#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
build_mock_cov.py
Build self-consistent Patchy P0+P2 (200x200) mock covariances for the 3
remaining BOSS DR12 bins (SGC_z1, NGC_z3, SGC_z3) and add them to
mock_cov_cache.npz, preserving the existing NGC_z1 entry.

Reuses dfd_pk_pipeline.py helpers EXACTLY:
  - read_ps1D  (13-col fast reader)
  - rebin_to_measured (Nmode-weighted rebin to 100 bins, dk=0.01)
  - make_kmask / K_MEAS / KMIN_FIT / KMAX_FIT for the validation sub-block
"""
import os
import sys
import glob
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

# Reuse the pipeline helpers/constants verbatim
from dfd_pk_pipeline import (read_ps1D, rebin_to_measured, make_kmask,
                             K_MEAS, KMIN_FIT, KMAX_FIT, N_MEAS)

CACHE = os.path.join(HERE, "mock_cov_cache.npz")

# bin -> local fast-disk dir holding the copied mocks
BINS = {
    "SGC_z1": "/tmp/patchy_SGC_z1",
    "NGC_z3": "/tmp/patchy_NGC_z3",
    "SGC_z3": "/tmp/patchy_SGC_z3",
}


def build_one(cap_zbin, localdir, max_mocks=300):
    files = sorted(glob.glob(os.path.join(localdir, "ps1D_patchy_*_renorm.dat")))
    files = files[:max_mocks]
    vecs = []
    bad = 0
    # The Patchy mock ps1D extends only to k~0.40, so the upper measured bins
    # (k>~0.40 in BOTH the P0 and P2 halves) are unpopulated -> NaN, with an
    # IDENTICAL pattern across every mock (same deterministic k-grid).  Those
    # bins lie far outside the k in [0.03,0.20] fit/validation sub-block and are
    # never indexed by the pipeline's mask, so we set them to 0 (not drop the
    # mock).  This matches the existing NGC_z1 cache convention.
    for fp in files:
        try:
            kf, P0f, P2f, Nm = read_ps1D(fp)
            P0, P2 = rebin_to_measured(kf, P0f, P2f, Nm)
            v = np.concatenate([P0, P2])  # 200-vec
            # require the FIT sub-block (k in [0.03,0.20], P0+P2) to be finite;
            # zero-fill only the always-empty high-k bins.
            fitm = make_kmask(use_p2=True)
            if np.any(~np.isfinite(v[fitm])):
                bad += 1
                continue
            v = np.where(np.isfinite(v), v, 0.0)
            vecs.append(v)
        except Exception:
            bad += 1
            continue
    V = np.array(vecs)                 # (Nmock, 200)
    nmock = V.shape[0]
    cov = np.cov(V, rowvar=False)      # (200,200)
    mean = V.mean(axis=0)              # (200,)
    return cov, mean, nmock, bad, V


def validate(cov, mean, V, nmock):
    """One mock vs mock mean over k in [0.03,0.20], P0+P2 (~34-dim), Hartlap."""
    mask = make_kmask(use_p2=True)
    p = int(mask.sum())
    Csub = cov[np.ix_(mask, mask)]
    hart = (nmock - p - 2.0) / (nmock - 1.0)
    Cinv = hart * np.linalg.inv(Csub)
    # average chi^2/dof over all mocks (each mock vs the mean) -> should be ~1
    chis = []
    for v in V:
        r = (v - mean)[mask]
        chis.append(float(r @ Cinv @ r))
    chis = np.array(chis)
    # single-mock check (first mock) as specified
    r0 = (V[0] - mean)[mask]
    chi2_single = float(r0 @ Cinv @ r0)
    return p, hart, chi2_single, chis.mean(), chis.mean() / p


def main():
    # Load existing cache (preserve NGC_z1)
    existing = {}
    if os.path.exists(CACHE):
        z = np.load(CACHE)
        for k in z.files:
            existing[k] = z[k]
    print(f"[cache] existing keys: {sorted(existing.keys())}")

    out = dict(existing)  # start from existing, add/overwrite the 3 new bins
    summary = []
    for name, localdir in BINS.items():
        cov, mean, nmock, bad, V = build_one(name, localdir)
        p, hart, chi2_single, chi2_mean, chi2_dof = validate(cov, mean, V, nmock)
        cond = np.linalg.cond(cov[np.ix_(make_kmask(True), make_kmask(True))])
        print(f"[{name}] nmock={nmock} (bad/skipped={bad})  Nbin(fit sub-block)={p}  "
              f"hartlap={hart:.4f}")
        print(f"        single-mock chi^2/dof = {chi2_single/p:.3f}   "
              f"avg-over-mocks chi^2/dof = {chi2_dof:.3f}   (sub-block cond={cond:.1e})")
        out[f"{name}_cov"] = cov
        out[f"{name}_mean"] = mean
        out[f"{name}_nmock"] = np.int64(nmock)
        summary.append((name, nmock, p, chi2_single / p, chi2_dof))

    np.savez(CACHE, **out)
    print(f"\n[cache] wrote {CACHE}")
    z2 = np.load(CACHE)
    print(f"[cache] new keys: {sorted(z2.files)}")
    print("\n=== VALIDATION SUMMARY (per new bin) ===")
    print(f"{'bin':<8}{'nmock':>7}{'Nbin':>6}{'single chi2/dof':>18}{'avg chi2/dof':>15}")
    for name, nmock, p, c1, cavg in summary:
        print(f"{name:<8}{nmock:>7}{p:>6}{c1:>18.3f}{cavg:>15.3f}")


if __name__ == "__main__":
    main()
