"""Step 12: are underluminous Tully-Fisher residuals spatially clustered in 3D?

A multi-galaxy engineered region (radius R) would dim every disc inside it coherently.  For R = 2, 5, 10,
20 Mpc we (i) tile the survey with disjoint groups of >= NMIN sample galaxies lying within R of a seed
galaxy, (ii) take the MEDIAN detrended residual of each group, (iii) flag groups whose median is below
MED_CUT (a group in which every member lost >= 50% of its light has median <= -0.30 dex on average),
(iv) calibrate the false-positive rate by permuting residuals among galaxies in the same distance bin and
catalogue, and (v) measure the detection efficiency by injecting 50% dimming into each group in turn.
In the injection a dimmed member that would fall outside its catalogue's magnitude / surface-brightness range
(selection weight w_sel_0.5 of s07) is REMOVED from the sample, and the group only counts as recovered if at
least NMIN members remain and their median is below MED_CUT (revision; referee point E9).
What is limited is the fraction of galaxy-occupied, redshift-space regions that contain >= NMIN sample
galaxies -- not a volume-filling fraction: voids and sparsely sampled regions are not tested.  R = 2 Mpc is
kept for completeness only; it is smaller than the redshift-space smearing of groups (fingers of god).
Distance-error control: a coherent distance error moves luminosity AND HI mass residuals together, while
star harvesting leaves HI mass untouched, so for each flagged group we also report the median residual of
the HI-mass--velocity relation.
Outputs: data/clustering_summary.csv, data/clustering_flagged_groups.csv, data/pair_correlation.csv
"""
import os
import numpy as np
import pandas as pd
from scipy.spatial import cKDTree
from common import DATA, poisson_ul, robust_polyfit, radec_to_xyz, mad_sigma

RADII = (2.0, 5.0, 10.0, 20.0)
NMIN = 5
MED_CUT = -0.20
F_INJ = 0.5
NPERM = 500
SEED = 7
DMIN = 60.0     # Mpc (H0=70): beyond this a 300 km/s peculiar velocity changes luminosities by < 0.065 dex


def load(clean_only):
    u = pd.read_csv(os.path.join(DATA, "union_sample.csv"))
    if clean_only:
        u = u[u.auto_clean].copy()
    # one distance scale (H0 = 70): CF4 distances were computed with H0 = 75
    u["d70"] = np.where(u.src == "cf4", u.Dist * 75.0 / 70.0, u.Dist)
    u = u[u.d70 >= DMIN].reset_index(drop=True)
    # detrend D against distance separately per catalogue (HI-selection makes distant galaxies brighter at fixed V)
    u["dbin"] = np.digitize(u.d70, [80, 100, 120, 140, 170, 200])
    u["Dd"] = u.D - u.groupby(["src", "dbin"]).D.transform("median")
    # HI-mass--velocity residual (distance-error control)
    p, _ = robust_polyfit(u.logV.values - 2.2, u.logMHI.values, deg=2)
    u["r_HI"] = u.logMHI - np.polyval(p, u.logV - 2.2)
    u["r_HI"] = u.r_HI - u.groupby(["src", "dbin"]).r_HI.transform("median")
    xyz = radec_to_xyz(u.ra, u.dec) * u.d70.values[:, None]
    # would the galaxy stay in its catalogue if dimmed by 50%?  (fractional weights -> one seeded draw)
    u["keep_dim"] = np.random.default_rng(SEED + 1).random(len(u)) < u[f"w_sel_{F_INJ}"].values
    return u, xyz


def tile(xyz, R, rng):
    """Greedy tiling into disjoint groups: each group = unassigned galaxies within R of a seed galaxy."""
    tree = cKDTree(xyz)
    free = np.ones(len(xyz), bool)
    groups = []
    for i in rng.permutation(len(xyz)):
        if not free[i]:
            continue
        mem = np.array([j for j in tree.query_ball_point(xyz[i], R) if free[j]])
        if len(mem) >= NMIN:
            groups.append(mem)
            free[mem] = False
    return groups


def group_medians(vals, groups):
    return np.array([np.median(vals[g]) for g in groups])


def injection_efficiency(Dd, keep, groups, shift, nmin):
    """Fraction of groups recovered when all members are dimmed and members leaving the catalogue are dropped."""
    hits = []
    for g in groups:
        stay = g[keep[g]]
        hits.append(len(stay) >= nmin and np.median(Dd[stay] + shift) < MED_CUT)
    return float(np.mean(hits)) if hits else np.nan


def run(clean_only, label):
    u, xyz = load(clean_only)
    rng = np.random.default_rng(SEED)
    Dd = u.Dd.values
    keep = u.keep_dim.values
    strata = (u.src.astype(str) + "_" + u.dbin.astype(str)).values
    strata_idx = [np.where(strata == s)[0] for s in np.unique(strata)]
    shift = np.log10(1 - F_INJ)
    print(f"  clustering sample [{label}]: {len(u)} galaxies beyond {DMIN:.0f} Mpc; robust sigma(D) = {mad_sigma(Dd):.3f} dex")

    rows, flagged = [], []
    for R in RADII:
        groups = tile(xyz, R, rng)
        if not groups:
            continue
        sizes = np.array([len(g) for g in groups])
        med = group_medians(Dd, groups)
        nflag = int((med < MED_CUT).sum())
        # permutation null: shuffle residuals within (catalogue, distance-bin) strata
        null_counts = np.empty(NPERM, int)
        null_min = np.empty(NPERM)
        for k in range(NPERM):
            perm = Dd.copy()
            for idx in strata_idx:
                perm[idx] = Dd[rng.permutation(idx)]
            m = group_medians(perm, groups)
            null_counts[k] = (m < MED_CUT).sum()
            null_min[k] = m.min()
        # injection: dim every member of one group by F_INJ, one group at a time
        eff_nosel = np.mean([(np.median(Dd[g] + shift) < MED_CUT) for g in groups])   # first-pass definition
        eff = injection_efficiency(Dd, keep, groups, shift, NMIN)
        eff_lenient = injection_efficiency(Dd, keep, groups, shift, 3)
        ngal = int(sizes.sum())
        ul_regions = poisson_ul(nflag) / (len(groups) * eff)
        # galaxy-weighted filling fraction: fraction of tiled galaxies living in a flagged-like region
        ul_fill = min(1.0, ul_regions)
        big = sizes >= 10
        nflag_big = int((med[big] < MED_CUT).sum())
        eff_big = injection_efficiency(Dd, keep, [g for g, b in zip(groups, big) if b], shift, NMIN) if big.any() else np.nan
        rows.append(dict(sample=label, n_sample=len(u), R_Mpc=R, n_groups=len(groups), n_gal_in_groups=ngal, median_group_size=float(np.median(sizes)),
                         n_flagged=nflag, null_mean=null_counts.mean(), null_p95=float(np.percentile(null_counts, 95)),
                         p_value_count=float((null_counts >= nflag).mean()),
                         obs_min_median=float(med.min()), null_min_median_p5=float(np.percentile(null_min, 5)),
                         p_value_min=float((null_min <= med.min()).mean()),
                         injection_efficiency=float(eff), injection_efficiency_no_selection=float(eff_nosel),
                         injection_efficiency_ge3_survivors=float(eff_lenient),
                         frac_members_kept_when_dimmed=float(np.mean([keep[g].mean() for g in groups])),
                         UL95_fraction_of_regions=float(ul_fill),
                         UL95_no_selection=float(min(1.0, poisson_ul(nflag) / (len(groups) * eff_nosel))),
                         n_groups_ge10=int(big.sum()), n_flagged_ge10=nflag_big, injection_efficiency_ge10=float(eff_big),
                         UL95_fraction_of_regions_ge10=float(min(1.0, poisson_ul(nflag_big) / (big.sum() * eff_big)))
                         if big.any() else np.nan,
                         UL95_bkg_subtracted=float(min(1.0, max(poisson_ul(nflag) - null_counts.mean(), 0)
                                                       / (len(groups) * eff))),
                         caveat="below redshift-space smearing scale; not meaningful" if R < 5 else ""))
        for g, m in zip(groups, med):
            if m < MED_CUT:
                gg = u.iloc[g]
                flagged.append(dict(sample=label, R_Mpc=R, n=len(g), median_Dd=m, median_rHI=float(gg.r_HI.median()),
                                    frac_below_m03=float((gg.Dd < -0.301).mean()),
                                    ra=float(gg.ra.mean()), dec=float(gg.dec.mean()), d70=float(gg.d70.mean()),
                                    cz_spread=float(gg.cz.std()), members=" ".join(gg.uid)))
        r = rows[-1]
        print(f"  R={R:4.0f} Mpc: {r['n_groups']:4d} groups ({ngal} galaxies, median size {r['median_group_size']:.0f}); "
              f"flagged {nflag} (null {r['null_mean']:.1f}, p={r['p_value_count']:.2f}); efficiency {eff:.2f} "
              f"(without selection loss {eff_nosel:.2f}); "
              f"95% UL on fraction of such regions dimmed >=50%: {ul_fill:.3f}; groups with >=10 members: "
              f"{r['n_groups_ge10']}, flagged {r['n_flagged_ge10']}, UL {r['UL95_fraction_of_regions_ge10']:.3f}")
    summary = pd.DataFrame(rows)
    fl = pd.DataFrame(flagged)
    if len(fl):
        print("  flagged groups (median residual, median HI-mass residual, fraction of members below -0.3 dex):")
        print(fl[["R_Mpc", "n", "median_Dd", "median_rHI", "frac_below_m03", "ra", "dec", "d70"]].round(3).to_string(index=False))

    if clean_only:
        return summary, fl, None
    # ---- pairwise residual correlation as a function of 3D separation
    tree = cKDTree(xyz)
    edges = [0, 1, 2, 5, 10, 20, 40]
    prs = tree.query_pairs(edges[-1], output_type="ndarray")
    sep = np.linalg.norm(xyz[prs[:, 0]] - xyz[prs[:, 1]], axis=1)
    prod = Dd[prs[:, 0]] * Dd[prs[:, 1]]
    low = (Dd < -0.301)
    both_low = low[prs[:, 0]] & low[prs[:, 1]]
    high = (Dd > 0.301)                       # control: the over-luminous tail cannot be produced by harvesting
    both_high = high[prs[:, 0]] & high[prs[:, 1]]
    rHI = u.r_HI.values
    okh = np.isfinite(rHI[prs[:, 0]]) & np.isfinite(rHI[prs[:, 1]])
    var = np.var(Dd)
    plow, phigh = low.mean(), high.mean()
    rows = []
    for lo, hi in zip(edges[:-1], edges[1:]):
        m = (sep >= lo) & (sep < hi)
        n = int(m.sum())
        # permutation error bars
        perm_c, perm_l, perm_h = [], [], []
        for k in range(100):
            perm = Dd.copy()
            for idx in strata_idx:
                perm[idx] = Dd[rng.permutation(idx)]
            perm_c.append(np.mean(perm[prs[m, 0]] * perm[prs[m, 1]]) / var)
            pl = perm < -0.301
            perm_l.append(np.mean(pl[prs[m, 0]] & pl[prs[m, 1]]) / plow ** 2)
            ph = perm > 0.301
            perm_h.append(np.mean(ph[prs[m, 0]] & ph[prs[m, 1]]) / phigh ** 2)
        rows.append(dict(s_lo=lo, s_hi=hi, n_pairs=n, corr=float(prod[m].mean() / var), corr_null_sd=float(np.std(perm_c)),
                         low_pair_excess=float(both_low[m].mean() / plow ** 2), low_pair_null_mean=float(np.mean(perm_l)),
                         low_pair_null_sd=float(np.std(perm_l)),
                         high_pair_excess=float(both_high[m].mean() / phigh ** 2), high_pair_null_mean=float(np.mean(perm_h)),
                         high_pair_null_sd=float(np.std(perm_h)),
                         corr_HI_mass_resid=float(np.mean(rHI[prs[m & okh, 0]] * rHI[prs[m & okh, 1]]) / np.nanvar(rHI))))
        print(f"    s={lo:2d}-{hi:2d} Mpc: pairs={n:7d}  corr(D_i,D_j)={rows[-1]['corr']:+.4f} +/- {rows[-1]['corr_null_sd']:.4f}"
              f"   P(both<-0.3)/P^2 = {rows[-1]['low_pair_excess']:.3f} (null {rows[-1]['low_pair_null_mean']:.3f} +/- {rows[-1]['low_pair_null_sd']:.3f})"
              f"   P(both>+0.3)/P^2 = {rows[-1]['high_pair_excess']:.3f} (null {rows[-1]['high_pair_null_mean']:.3f} +/- {rows[-1]['high_pair_null_sd']:.3f})"
              f"   corr(rHI) = {rows[-1]['corr_HI_mass_resid']:+.4f}")
    return summary, fl, pd.DataFrame(rows)


def main():
    s1, f1, pc = run(False, "all sample galaxies")
    s2, f2, _ = run(True, "catalogue-clean only")
    pd.concat([s1, s2], ignore_index=True).to_csv(os.path.join(DATA, "clustering_summary.csv"), index=False)
    pd.concat([f1, f2], ignore_index=True).to_csv(os.path.join(DATA, "clustering_flagged_groups.csv"), index=False)
    pc.to_csv(os.path.join(DATA, "pair_correlation.csv"), index=False)


if __name__ == "__main__":
    main()