
#!/usr/bin/env python3
from __future__ import annotations
import argparse, json
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from rbft_lib import ensure_dir

def gaussian(X,Y,x0,y0,sx,sy=None,amp=1.0):
    if sy is None: sy=sx
    return amp*np.exp(-0.5*(((X-x0)/sx)**2 + ((Y-y0)/sy)**2))

def get_point(obs, comp, typ):
    r = obs[(obs.component==comp)&(obs.type==typ)].iloc[0]
    return np.array([r.x_kpc, r.y_kpc], dtype=float)

def deflection_from_kappa(kappa, dx, dy):
    k = kappa - np.mean(kappa)
    ny, nx = k.shape
    kx = 2*np.pi*np.fft.fftfreq(nx, d=dx)
    ky = 2*np.pi*np.fft.fftfreq(ny, d=dy)
    KX, KY = np.meshgrid(kx, ky)
    K2 = KX**2 + KY**2
    kfft = np.fft.fft2(k)
    psi_fft = np.zeros_like(kfft, dtype=complex)
    mask = K2 > 0
    psi_fft[mask] = -2*kfft[mask]/K2[mask]
    ax = np.real(np.fft.ifft2(1j*KX*psi_fft))
    ay = np.real(np.fft.ifft2(1j*KY*psi_fft))
    return ax, ay

def norm_total(a):
    s = np.sum(a)
    return a/s if s > 0 else a

def corr(a,b,mask=None):
    if mask is not None:
        a=a[mask]; b=b[mask]
    a=a.ravel(); b=b.ravel()
    if np.std(a)==0 or np.std(b)==0:
        return np.nan
    return float(np.corrcoef(a,b)[0,1])

def main():
    ap = argparse.ArgumentParser(description="Run Bullet Cluster centroid and deflection-field proxy tests.")
    ap.add_argument("--root", default=".")
    args = ap.parse_args()
    root = Path(args.root)
    outdir = ensure_dir(root / "outputs" / "06_bullet")
    figdir = ensure_dir(root / "figures")

    kpc_per_pix = 15.4
    obs = pd.read_csv(root / "data" / "external" / "bullet_cluster_peak_coordinates.csv")
    cx, cy = obs["x_pix"].mean(), obs["y_pix"].mean()
    obs["x_kpc"] = (obs["x_pix"]-cx)*kpc_per_pix
    obs["y_kpc"] = (obs["y_pix"]-cy)*kpc_per_pix
    obs.to_csv(outdir / "bullet_observed_coordinates_used.csv", index=False)

    # Centroid audit
    rows = []
    for comp in ["main","sub"]:
        lens = get_point(obs, comp, "lensing_kappa_peak")
        gas = get_point(obs, comp, "gas_sigma_peak")
        gal = get_point(obs, comp, "compact_galaxy_peak")
        dgas = float(np.linalg.norm(lens-gas))
        dgal = float(np.linalg.norm(lens-gal))
        rows.append({
            "component": comp,
            "d_lens_to_gas_kpc": dgas,
            "d_lens_to_compact_galaxy_kpc": dgal,
            "galaxy_vs_gas_distance_improvement": 1-dgal/dgas,
            "pass_centroid_condition": dgal < dgas
        })
    centroid = pd.DataFrame(rows)
    centroid.to_csv(outdir / "bullet_centroid_distance_audit.csv", index=False)

    # Build map grid
    xmin, xmax = obs.x_kpc.min()-650, obs.x_kpc.max()+650
    ymin, ymax = obs.y_kpc.min()-520, obs.y_kpc.max()+520
    N=384
    x=np.linspace(xmin,xmax,N); y=np.linspace(ymin,ymax,N)
    dx=x[1]-x[0]; dy=y[1]-y[0]
    X,Y=np.meshgrid(x,y)

    target=np.zeros_like(X)
    gas_map=np.zeros_like(X)
    star_map=np.zeros_like(X)
    for comp, amp, sx, sy in [("main",1.0,145,100),("sub",0.9,125,90)]:
        lx,ly = get_point(obs, comp, "lensing_kappa_peak")
        target += gaussian(X,Y,lx,ly,sx,sy,amp=amp)
    for comp in ["main","sub"]:
        gx,gy = get_point(obs, comp, "gas_sigma_peak")
        sx0,sy0 = get_point(obs, comp, "compact_galaxy_peak")
        gas_map += gaussian(X,Y,gx,gy,278,190,amp=6.0)
        star_map += gaussian(X,Y,sx0,sy0,85,70,amp=1.0)

    C_star=0.584; C_gas_sparc=0.095; C_gas_shock=0.087
    maps = {
        "target_observed_kappa_proxy": norm_total(target),
        "mass_only_baryonic": norm_total(gas_map+star_map),
        "RBFT_SPARC_locked": norm_total((C_star**2)*star_map + (C_gas_sparc**2)*gas_map),
        "RBFT_shock_prelock": norm_total((C_star**2)*star_map + (C_gas_shock**2)*gas_map),
    }

    fields={}
    for name,k in maps.items():
        ax,ay=deflection_from_kappa(k,dx,dy)
        fields[name]={"kappa":k,"alpha_x":ax,"alpha_y":ay}

    target_field=fields["target_observed_kappa_proxy"]
    amag=np.sqrt(target_field["alpha_x"]**2+target_field["alpha_y"]**2)
    lens_strength=230/(np.percentile(amag[amag>0],99)+1e-30)
    mask90 = target_field["kappa"] > np.percentile(target_field["kappa"],90)
    metrics=[]
    for name in ["mass_only_baryonic","RBFT_SPARC_locked","RBFT_shock_prelock"]:
        f=fields[name]
        terr=np.sqrt(((f["alpha_x"]-target_field["alpha_x"])*lens_strength)**2+
                     ((f["alpha_y"]-target_field["alpha_y"])*lens_strength)**2)
        dot=f["alpha_x"]*target_field["alpha_x"]+f["alpha_y"]*target_field["alpha_y"]
        mag1=np.sqrt(f["alpha_x"]**2+f["alpha_y"]**2)
        mag2=np.sqrt(target_field["alpha_x"]**2+target_field["alpha_y"]**2)
        cosang=dot/(mag1*mag2+1e-30)
        metrics.append({
            "model":name,
            "kappa_shape_corr_all":corr(f["kappa"],target_field["kappa"]),
            "kappa_shape_corr_strong_region":corr(f["kappa"],target_field["kappa"],mask90),
            "deflection_RMSE_kpc_strong90":float(np.sqrt(np.mean(terr[mask90]**2))),
            "median_deflection_error_kpc_strong90":float(np.median(terr[mask90])),
            "mean_alpha_direction_cosine_strong90":float(np.mean(cosang[mask90]))
        })
    pd.DataFrame(metrics).to_csv(outdir / "bullet_deflection_field_metrics.csv", index=False)

    # Simple source/lensed image proxy
    def source_image(bx,by):
        return (gaussian(bx,by,-120,40,75,35,1.0)+
                gaussian(bx,by,-40,-35,45,95,0.65)+
                gaussian(bx,by,80,10,30,25,0.35))
    def lensed(f):
        return source_image(X-lens_strength*f["alpha_x"], Y-lens_strength*f["alpha_y"])
    Iobs = lensed(target_field)
    bright = Iobs > np.percentile(Iobs,75)
    img_rows=[]
    for name in ["mass_only_baryonic","RBFT_SPARC_locked","RBFT_shock_prelock"]:
        im=lensed(fields[name])
        res=im-Iobs
        img_rows.append({
            "model":name,
            "image_RMSE_bright75":float(np.sqrt(np.mean(res[bright]**2))),
            "image_corr_bright75":corr(im,Iobs,bright)
        })
    pd.DataFrame(img_rows).to_csv(outdir / "bullet_synthetic_image_metrics.csv", index=False)

    # Figures
    plt.figure(figsize=(7,5))
    colors={"gas_sigma_peak":"red","lensing_kappa_peak":"blue","compact_galaxy_peak":"orange"}
    markers={"gas_sigma_peak":"+","lensing_kappa_peak":"o","compact_galaxy_peak":"x"}
    for typ,g in obs.groupby("type"):
        plt.scatter(g.x_kpc,g.y_kpc,c=colors[typ],marker=markers[typ],s=100,label=typ)
    plt.gca().set_aspect("equal", adjustable="box")
    plt.xlabel("x [kpc centered]"); plt.ylabel("y [kpc centered]")
    plt.title("Bullet centroids used")
    plt.legend()
    plt.tight_layout()
    plt.savefig(figdir / "bullet_centroid_offsets.png", dpi=160)
    plt.close()

    fig, axes = plt.subplots(1,4,figsize=(17,4),constrained_layout=True)
    for ax,(name,f) in zip(axes,fields.items()):
        im=ax.imshow(f["kappa"],origin="lower",extent=[xmin,xmax,ymin,ymax],aspect="auto")
        ax.set_title(name.replace("_"," "))
        fig.colorbar(im,ax=ax,fraction=0.046,pad=0.04)
    plt.savefig(figdir / "bullet_kappa_proxy_maps.png", dpi=160, bbox_inches="tight")
    plt.close()

    print("Centroid:")
    print(centroid.to_string(index=False))
    print("\nDeflection metrics:")
    print(pd.DataFrame(metrics).to_string(index=False))

if __name__ == "__main__":
    main()
