#!/usr/bin/env python3
"""
UVCS Bright-Dim Permutation Test
--------------------------------
Splits each (day, Rbin) group temporally (by observation timestamp),
computes early/late intensity contrast, and tests significance via
permutation of frame order.

The "bright/dim" labeling is applied AFTER the temporal split for sign
convention only - it does not affect the statistical test.

Input: uvcs_lya_summary.csv in the working directory -- a per-frame Ly-alpha
fit summary table reduced from archival SOHO/UVCS observations (not bundled
with this deposit; obtain the archival UVCS Ly-alpha data and reduce to one
row per frame). Required columns: day, date_obs (e.g. "12-Apr-1997
01:23:45.000"), I_rate (Ly-alpha intensity rate), lambda_centroid (Angstrom),
posx1 (heliocentric position, R_sun); optional: roll_deg, EXPTIME.

Dependencies: numpy, pandas.
"""
import os, numpy as np, pandas as pd

BIN_W = 0.75
N_PERM = 20000  # Number of permutations
LAM0  = 1215.67
C_KMS = 299792.458

def build_obs_pairs(df):
    """
    For each (day, Rbin) group:
    1. Sort frames by observation timestamp
    2. Split at temporal midpoint into early/late halves
    3. Label whichever half is brighter as "bright" (sign convention only)
    """
    rows = []
    for (day, rbin), g in df.groupby(["day", "Rbin"], sort=False):
        # TEMPORAL SPLIT: sort by observation time
        h = g.sort_values("_t")
        k = len(h) // 2
        early, late = h.iloc[:k], h.iloc[k:]
        
        if len(early) == 0 or len(late) == 0:
            continue
            
        I_early = early["I_rate"].median()
        I_late = late["I_rate"].median()
        L_early = early["lambda_centroid"].median()
        L_late = late["lambda_centroid"].median()
        
        if not (np.isfinite(I_early) and np.isfinite(I_late)):
            continue
            
        # Label brighter half as "bright" (sign convention only)
        if I_early >= I_late:
            I_bright, I_dim = I_early, I_late
            L_bright, L_dim = L_early, L_late
        else:
            I_bright, I_dim = I_late, I_early
            L_bright, L_dim = L_late, L_early
            
        contrast = I_bright / I_dim - 1.0
        delta_lambda = L_bright - L_dim
        rows.append([day, rbin, contrast, delta_lambda])
        
    obs = pd.DataFrame(rows, columns=["day", "Rbin", "obs_delta_I_frac", "obs_delta_lambda"])
    if not obs.empty:
        obs["obs_delta_v_kms"] = C_KMS * obs["obs_delta_lambda"] / LAM0
    return obs

def perm_null(df, day, rbin, nperm=N_PERM, rng_seed=42):
    """
    Generate null distribution by randomly permuting frame order.
    This breaks any temporal structure while preserving intensity values.
    """
    g = df[(df["day"] == day) & (df["Rbin"] == rbin)]
    if len(g) < 2:
        return np.array([]), np.array([])
        
    rng = np.random.default_rng(rng_seed)
    di = np.empty(nperm, dtype=float)
    dl = np.empty(nperm, dtype=float)
    vals_I = g["I_rate"].to_numpy()
    vals_L = g["lambda_centroid"].to_numpy()
    idx = np.arange(len(g))
    k = len(g) // 2
    
    for i in range(nperm):
        perm = rng.permutation(idx)
        A, B = perm[:k], perm[k:]
        Ia, Ib = np.nanmedian(vals_I[A]), np.nanmedian(vals_I[B])
        La, Lb = np.nanmedian(vals_L[A]), np.nanmedian(vals_L[B])
        
        # Random bright/dim assignment under null
        if rng.random() < 0.5:
            sI, cI, sL, cL = Ia, Ib, La, Lb
        else:
            sI, cI, sL, cL = Ib, Ia, Lb, La
            
        di[i] = (sI / cI - 1.0) if (np.isfinite(sI) and np.isfinite(cI) and cI != 0) else np.nan
        dl[i] = sL - cL
        
    return di[np.isfinite(di)], dl[np.isfinite(dl)]

def main():
    df = pd.read_csv("uvcs_lya_summary.csv")
    
    # Coerce numerics
    for c in ["I_rate", "lambda_centroid", "posx1", "roll_deg", "EXPTIME"]:
        if c in df:
            df[c] = pd.to_numeric(df[c], errors="coerce")

    # Parse observation time
    t = pd.to_datetime(df["date_obs"], format="%d-%b-%Y %H:%M:%S.%f", errors="coerce")
    t = t.fillna(pd.to_datetime(df["date_obs"], errors="coerce"))
    df["_t"] = t

    # Ly-alpha sanity band
    df = df[(df["lambda_centroid"] > 1180) & (df["lambda_centroid"] < 1250)].copy()

    # Radial binning
    r = df["posx1"]
    df["Rbin"] = np.where(np.isfinite(r), np.round(r / BIN_W) * BIN_W, 0.0)

    # Keep (day, Rbin) with >=2 frames
    kc = df.groupby(["day", "Rbin"])["I_rate"].size()
    keep = kc[kc >= 2].index
    df = df.set_index(["day", "Rbin"]).loc[keep].reset_index()

    print(f"Groups kept: {len(keep)}  (unique days: {df['day'].nunique()})")
    if df.empty:
        print("No (day, Rbin) with >=2 frames.")
        return

    obs = build_obs_pairs(df)
    if obs.empty:
        print("No observed pairs formed.")
        return

    rows = []
    keys = list(obs[["day", "Rbin"]].itertuples(index=False, name=None))
    
    for j, (day, rbin) in enumerate(keys, 1):
        di, dl = perm_null(df, day, rbin, N_PERM, rng_seed=1234 + j)
        if di.size == 0:
            continue
            
        obs_row = obs[(obs["day"] == day) & (obs["Rbin"] == rbin)].iloc[0]
        obs_di = obs_row["obs_delta_I_frac"]
        obs_dl = obs_row["obs_delta_lambda"]
        
        # Two-sided p-value: median-centered comparison
        # Tests whether |Δ_obs - median(null)| is extreme relative to null spread
        p_di = (np.sum(np.abs(di - np.median(di)) >= np.abs(obs_di - np.median(di))) + 1) / (len(di) + 1)
        p_dl = (np.sum(np.abs(dl - np.median(dl)) >= np.abs(obs_dl - np.median(dl))) + 1) / (len(dl) + 1)
        
        rows.append(dict(
            day=day, Rbin=rbin,
            obs_delta_I_frac=float(obs_di),
            obs_delta_lambda=float(obs_dl),
            obs_delta_v_kms=float(C_KMS * obs_dl / LAM0),
            p_two_sided_contrast=float(p_di),
            p_two_sided_dlambda=float(p_dl),
            n_null=int(di.size)
        ))
        
        if j % 25 == 0 or j == len(keys):
            print(f"...permuted {j}/{len(keys)} groups (N_PERM={N_PERM})")

    res = pd.DataFrame(rows).reset_index(drop=True)
    res.to_csv("uvcs_perm_results.csv", index=False)
    print(f"\nSaved: uvcs_perm_results.csv ({len(res)} groups)")

if __name__ == "__main__":
    main()
