#!/usr/bin/env python3
"""What do peculiar velocities and survey systematics do to the sphere statistic?

 1. pv_expectation():  linear-theory rms line-of-sight bulk velocity of a top-hat sphere of radius R
    (BBKS transfer function, sigma_8 = 0.81) and the FP-residual shift it causes,
        d(delta_I) = -c_r * log10(1 + v/cz)      (c_r ~ -1.1: a receding sphere looks too bright).
 2. finger_of_god():  the SDSS sample uses individual redshifts.  Inside a cluster, galaxies with negative
    line-of-sight velocities are assigned too small a distance, hence too small an R_e, hence look too DIM for
    the FP.  A redshift-space sphere on the near side of a cluster therefore shows a coherent deficit that has
    nothing to do with luminosity.  For the members of each extreme SDSS sphere that have a group redshift in
    the Howlett+22 catalogue the residual is recomputed with the group redshift.
 3. groupz_scan():  the same sphere scan on the Howlett+22 sample (group redshifts throughout).
 4. env_table():  mean FP residual versus local galaxy over-density for the three samples.
"""
import json
import os

import numpy as np
import pandas as pd
from scipy import integrate
from scipy.spatial import cKDTree

import analysis_regions as G
import fpk_common as F

SIGMA8, NS, OB, H = 0.81, 0.965, 0.048, F.H0 / 100.0


def _pk_unnorm(k):
    """BBKS (Bardeen et al. 1986) shape with the Sugiyama (1995) baryon correction; k in h/Mpc."""
    gam = F.OM * H * np.exp(-OB * (1 + np.sqrt(2 * H) / F.OM))
    q = k / gam
    t = np.log(1 + 2.34 * q) / (2.34 * q) * (1 + 3.89 * q + (16.1 * q) ** 2 + (5.46 * q) ** 3 + (6.71 * q) ** 4) ** -0.25
    return k ** NS * t ** 2


def _w(x):
    return 3 * (np.sin(x) - x * np.cos(x)) / x ** 3


def pv_expectation(c_r):
    lk = np.linspace(np.log(1e-4), np.log(50.0), 6000)
    k = np.exp(lk)
    p = _pk_unnorm(k)
    s8 = np.sqrt(integrate.trapezoid(k ** 3 * p * _w(8.0 * k) ** 2, lk) / (2 * np.pi ** 2))
    p *= (SIGMA8 / s8) ** 2
    fgrow = F.OM ** 0.55
    rows = []
    for R in G.R_LIST:
        rh = R * H                                           # Mpc/h
        v3 = 100.0 * fgrow * np.sqrt(integrate.trapezoid(k * p * _w(rh * k) ** 2, lk) / (2 * np.pi ** 2))
        v1 = v3 / np.sqrt(3)
        for z in (0.02, 0.04, 0.08, 0.12):
            rows.append(dict(R_mpc=R, z=z, v_los_rms_kms=v1,
                             dI_rms_dex=abs(c_r) * np.log10(1 + v1 / (F.C_KMS * z))))
    return pd.DataFrame(rows)


def _scan(df, R, rng, n_shuffle=200):
    lam, sv, a = G.standardise(df)
    zeta = G.zeta_of(df.dI.values, df.dI_sd.values, lam, sv)
    zeta = zeta - G.z_trend(df.zdist.values, zeta)[0]
    xyz = F.radec_to_xyz(df.ra.values, df["dec"].values, F.comoving_distance(df.zdist.values))
    tree, Aij = G.adjacency(xyz, R)
    N = np.asarray(Aij.sum(axis=1)).ravel()
    good = N >= G.N_MIN
    S = np.where(good, (Aij @ zeta) / np.sqrt(np.maximum(N, 1)), np.nan)
    null = np.array([((Aij @ rng.permutation(zeta))[good] / np.sqrt(N[good])).min() for _ in range(n_shuffle)])
    return S, N, xyz, Aij, null


def groupz_scan(rng):
    """Sphere scan of the Howlett+22 sample, which uses group redshifts (no finger-of-god effect)."""
    d = pd.read_csv(F.resid_path("howlett"))
    d = d[(d.zdist >= 0.02) & (d.zdist <= 0.10)].reset_index(drop=True)
    rows, pat = [], []
    for R in G.R_LIST:
        S, N, xyz, Aij, null = _scan(d, R, rng)
        smin = float(np.nanmin(S))
        rows.append(dict(sample="howlett (group z)", R_mpc=R, n_gal=len(d), S_min=smin,
                         p_shuffle=float((np.sum(null <= smin) + 1) / (len(null) + 1)),
                         S_min_shuffle_median=float(np.median(null)), S_std=float(np.nanstd(S))))
        for idx in G.greedy_top(S, N, xyz, R, k=5, sign=-1):
            nb = Aij[idx].indices
            pat.append(dict(sample="howlett (group z)", R_mpc=R, S=float(S[idx]), N=int(N[idx]),
                            mean_dI_dex=float(d.dI.values[nb].mean()), ra=float(d.ra[idx]),
                            dec=float(d["dec"][idx]), z=float(d.zdist[idx])))
    return pd.DataFrame(rows), pd.DataFrame(pat)


def finger_of_god(c_r, c_z):
    """Group-redshift version of the residual for the members of the extreme SDSS spheres."""
    s = pd.read_csv(F.resid_path("sdss"))
    s = s[(s.zdist >= 0.02) & (s.zdist <= 0.12)].reset_index(drop=True)
    h = pd.read_csv(os.path.join(F.DATA, "sample_howlett.csv"))[["id", "zdist", "Ngroup"]] \
        .rename(columns={"zdist": "z_group"})
    s = s.merge(h, on="id", how="left")
    zg = s.z_group.values
    # R_e scales with the angular-diameter distance; I_e carries (1+z)^4; the plane has a redshift term
    d_logr = np.log10(F.angular_diameter_distance(zg) / F.angular_diameter_distance(s.zdist.values))
    d_logi = 4 * np.log10((1 + zg) / (1 + s.zdist.values))
    s["dI_groupz"] = s.dI.values + d_logi - c_r * d_logr - c_z * (zg - s.zdist.values)
    xyz = F.radec_to_xyz(s.ra.values, s["dec"].values, F.comoving_distance(s.zdist.values))
    tree = cKDTree(xyz)
    pt = pd.read_csv(os.path.join(F.DATA, "region_extreme_patches.csv"))
    pt = pt[(pt["sample"] == "sdss") & (pt.variant == "raw") & pt.kind.str.contains("negative")]
    rows = []
    for _, p in pt.iterrows():
        c = F.radec_to_xyz(np.array([p.ra]), np.array([p["dec"]]), F.comoving_distance(np.array([p.z])))[0]
        nb = np.array(tree.query_ball_point(c, p.R_mpc), dtype=int)
        m = np.isfinite(s.dI_groupz.values[nb])
        mm = nb[m]
        rows.append(dict(R_mpc=p.R_mpc, ra=p.ra, dec=p["dec"], z=p.z, S=p.S, N=len(nb),
                         mean_dI_all=float(s.dI.values[nb].mean()), n_with_group_z=int(m.sum()),
                         mean_dI_ownz_subset=float(s.dI.values[mm].mean()) if m.any() else np.nan,
                         mean_dI_groupz_subset=float(s.dI_groupz.values[mm].mean()) if m.any() else np.nan,
                         sem_subset=float(s.dI_groupz.values[mm].std() / np.sqrt(m.sum())) if m.sum() > 1 else np.nan,
                         mean_z_own_subset=float(s.zdist.values[mm].mean()) if m.any() else np.nan,
                         mean_z_group_subset=float(zg[mm].mean()) if m.any() else np.nan))
    return pd.DataFrame(rows)


def env_table():
    """Mean I-direction FP residual in bins of local over-density (neighbours within 5 Mpc in redshift space,
    relative to the mean at that redshift)."""
    rows = []
    for name, lab in (("sdss", "SDSS own (individual z)"), ("howlett", "Howlett+22 (group z)"),
                      ("6df", "6dFGSv (group z)")):
        d = pd.read_csv(F.resid_path(name))
        d = d[d.zdist >= 0.01].reset_index(drop=True)
        xyz = F.radec_to_xyz(d.ra.values, d["dec"].values, F.comoving_distance(d.zdist.values))
        tree = cKDTree(xyz)
        n5 = np.array([len(v) for v in tree.query_ball_point(xyz, 5.0)], float)
        zb = np.quantile(d.zdist.values, np.linspace(0, 1, 11))
        k = np.clip(np.searchsorted(zb, d.zdist.values, side="right") - 1, 0, 9)
        nbar = np.array([n5[k == i].mean() for i in range(10)])
        od = n5 / nbar[k]
        ok = np.abs(d.dI.values) < 0.35                       # robust mean: ignore the far tails
        for lo, hi in ((0, 0.5), (0.5, 1), (1, 2), (2, 4), (4, 8), (8, 1e9)):
            m = (od >= lo) & (od < hi) & ok
            if m.sum() >= 20:
                rows.append(dict(sample=lab, overdensity_lo=lo, overdensity_hi=hi, n=int(m.sum()),
                                 mean_dI_dex=float(d.dI.values[m].mean()),
                                 sem_dex=float(d.dI.values[m].std() / np.sqrt(m.sum()))))
    return pd.DataFrame(rows)


def survey_volume():
    """Footprint area from the occupancy of equal-area sky cells, and the comoving volume it encloses."""
    rows = []
    for name, lab, zmin, zmax, cell in (("sdss", "SDSS own", 0.02, 0.12, 4.0),
                                        ("6df", "6dFGSv", 3000 / F.C_KMS, 16120 / F.C_KMS, 25.0)):
        d = pd.read_csv(F.resid_path(name))
        d = d[(d.zdist >= zmin) & (d.zdist <= zmax)]
        nz = int(round(41252.96 / cell / 2 / np.sqrt(41252.96 / cell / np.pi))) * 2      # ~square cells
        iz = np.floor((np.sin(np.radians(d["dec"].values)) + 1) / 2 * nz).astype(int)
        nra = int(round(41252.96 / cell / nz))
        ira = np.floor(d.ra.values / 360.0 * nra).astype(int)
        occ = len(np.unique(iz * 100000 + ira))
        area = occ * 41252.96 / (nz * nra)
        vol = area / 41252.96 * 4 * np.pi / 3 * (F.comoving_distance(zmax) ** 3 - F.comoving_distance(zmin) ** 3)
        rows.append(dict(sample=lab, n_gal=len(d), zmin=zmin, zmax=zmax, cell_deg2=41252.96 / (nz * nra),
                         footprint_deg2=area, volume_mpc3=float(vol),
                         d_min_mpc=float(F.comoving_distance(zmin)), d_max_mpc=float(F.comoving_distance(zmax))))
    return pd.DataFrame(rows)


def main():
    rng = np.random.default_rng(G.SEED + 1)
    fits = json.load(open(os.path.join(F.DATA, "fp_fits.json")))
    c_r, c_z = fits["sdss"]["fpI_coef"][2], fits["sdss"]["fpI_coef"][3]
    pv = pv_expectation(c_r)
    pv.to_csv(os.path.join(F.DATA, "pv_expectation.csv"), index=False, float_format="%.4g")
    print(pv.to_string(index=False))
    fog = finger_of_god(c_r, c_z)
    fog.to_csv(os.path.join(F.DATA, "region_patches_groupz.csv"), index=False, float_format="%.5g")
    pd.set_option("display.width", 250)
    print(fog.to_string(index=False))
    rows, pat = groupz_scan(rng)
    rows.to_csv(os.path.join(F.DATA, "region_scan_howlett_groupz.csv"), index=False, float_format="%.5g")
    pat.to_csv(os.path.join(F.DATA, "region_patches_howlett_groupz.csv"), index=False, float_format="%.5g")
    print(rows.to_string(index=False))
    print(pat.to_string(index=False))
    sv = survey_volume()
    sv.to_csv(os.path.join(F.DATA, "survey_volume.csv"), index=False, float_format="%.5g")
    print(sv.to_string(index=False))
    env = env_table()
    env.to_csv(os.path.join(F.DATA, "env_trend.csv"), index=False, float_format="%.4g")
    print(env.to_string(index=False))
    return True


if __name__ == "__main__":
    main()