"""
Path C: Compute kappa_pred(x,y) from baryonic inputs using the framework's
three-layer K_act composite pipeline.

Uses JWST NIRCam F277W combined mosaic as stellar+ICL surface brightness tracer.
Scales to Sigma_baryonic via stellar-to-total-baryon ratio appropriate for
galaxy clusters (gas ~85%, stars ~15% -> stellar mass x ~6.5 = total baryon).

Applies the locked UM/FUM framework constants:
  phi = 1.6180339887498949
  alpha_struct = 0.0073032157
  Eidolon = (1-alpha_struct)/alpha_struct
  K_struct = 4*phi^3*Eidolon/3 = 767.70
  g_critical = c*H0/(2*omega_C1) ~= 1.042e-10 m/s^2

Pipeline per pixel:
  g_bar(x,y) = 2*pi*G*Sigma_baryonic(x,y)
  K_act,internal(x,y) = max(1, sqrt(g_critical/g_bar(x,y)))
  Lambda(x,y) = 1/[1 + (K_int-1)/K_struct]
  K_act,observed(x,y) = K_int * F_env * F_ML^-1
  kappa_baryonic(x,y) = Sigma_baryonic / Sigma_crit_lens
  kappa_pred(x,y) = K_act,observed * kappa_baryonic
"""
import os
import sys
import warnings
warnings.filterwarnings('ignore')

import numpy as np
from astropy.io import fits
from astropy.wcs import WCS
from astropy.cosmology import Planck18
import astropy.units as u
import astropy.constants as const
from astropy.coordinates import SkyCoord
from reproject import reproject_interp

print("=" * 70)
print("Bullet Cluster kappa_pred Pipeline -- Path C")
print("=" * 70)

# ===== Locked framework constants =====
PHI = 1.6180339887498949
ALPHA_STRUCT = 0.0073032157
EIDOLON = (1 - ALPHA_STRUCT) / ALPHA_STRUCT
K_STRUCT = 4 * PHI**3 * EIDOLON / 3
G_CRITICAL = 1.042e-10  # m/s^2

# Cluster-locus environmental parameters (cluster core is the host)
EPS_EXT = 0.0
ETA_TID = 0.0
DELTA_ML = 0.05
F_ENV = 1.0 / ((1.0 + EPS_EXT) * (1.0 + ETA_TID**2))
F_ML_INV = 1.0 / (1.0 + DELTA_ML)

print(f"K_struct = 4*phi^3*Eidolon/3 = {K_STRUCT:.4f}")
print(f"g_critical = {G_CRITICAL:.4e} m/s^2")
print(f"F_env = {F_ENV:.4f}, F_ML^-1 = {F_ML_INV:.4f}")

# ===== Cosmology / lensing geometry =====
Z_LENS = 0.296
Z_SOURCE = 2.0
COSMO = Planck18

D_L = COSMO.angular_diameter_distance(Z_LENS)
D_S = COSMO.angular_diameter_distance(Z_SOURCE)
D_LS = COSMO.angular_diameter_distance_z1z2(Z_LENS, Z_SOURCE)
G = const.G.to(u.m**3 / (u.kg * u.s**2))
c_speed = const.c.to(u.m / u.s)

Sigma_crit_lens = (c_speed**2 * D_S / (4 * np.pi * G * D_L * D_LS)).to(u.kg / u.m**2).value
print(f"D_L = {D_L:.2f}, D_S = {D_S:.2f}, D_LS = {D_LS:.2f}")
print(f"Sigma_crit_lens = {Sigma_crit_lens:.4e} kg/m^2")

# ===== Load kappa_observed_baseline (target grid) =====
TARGET_FITS = 'bullet_kappa_observed_SLbaseline.fits'
print(f"\n--- Loading target grid: {TARGET_FITS} ---")
with fits.open(TARGET_FITS) as hdul:
    kappa_obs = hdul[0].data.astype(np.float64)
    target_header = hdul[0].header
    target_wcs = WCS(target_header)
print(f"Target shape: {kappa_obs.shape}")
print(f"Target pixel scale: {target_wcs.proj_plane_pixel_scales()[0].to('arcsec').value:.3f} arcsec/pixel")

# ===== Load JWST F277W combined mosaic (stellar+ICL tracer) =====
F277W_PATH = 'JWST_GO4598_NIRCAM/jw04598-o001_t001_nircam_clear-f277w_i2d.fits'
if not os.path.exists(F277W_PATH):
    print(f"ERROR: {F277W_PATH} not found")
    sys.exit(1)

print(f"\n--- Loading JWST F277W: {F277W_PATH} ---")
with fits.open(F277W_PATH, memmap=True) as hdul:
    sci = None
    for i, h in enumerate(hdul):
        if h.name == 'SCI':
            sci = i
            break
    if sci is None:
        sci = 1
    flux = hdul[sci].data.astype(np.float64)
    src_header = hdul[sci].header
    src_wcs = WCS(src_header)
print(f"F277W shape: {flux.shape}, units: {src_header.get('BUNIT', '?')}")
pix_scale_arcsec_source = src_wcs.proj_plane_pixel_scales()[0].to('arcsec').value
print(f"F277W pixel scale: {pix_scale_arcsec_source:.4f} arcsec/pixel")

# ===== Convert F277W flux (MJy/sr) -> stellar surface mass density =====
# F277W observes rest-frame near-IR at z=0.296: lambda_obs=2.77 um -> lambda_rest=2.14 um
# Rest-frame ~2.1 um is excellent stellar mass tracer (M*/L ~ 0.5-1 M_sun/L_sun)
#
# Conversion path:
#   flux [MJy/sr] -> surface brightness mu_AB [mag/arcsec^2]
#   AB mag -> L per unit solid angle (using D_L) -> stellar luminosity surface density
#   M*/L ratio -> stellar mass surface density [M_sun/Mpc^2]
#
# For simplicity, use a single conversion factor calibrated to give typical
# cluster-core stellar surface densities ~100 M_sun/pc^2 = 1e8 M_sun/Mpc^2
#
# 1 MJy/sr at z=0.296 with M/L ~ 1 corresponds to roughly:
#   L_nu(observed) per sr -> L_nu(rest)*(1+z) due to bandpass
#   Distance: D_L = 1614 Mpc (Planck18 at z=0.296)
#   M*_per_pixel = flux_per_pixel * conversion_factor

# Conversion factor calibrated: 1 MJy/sr -> 1e9 M_sun/Mpc^2 (order of magnitude
# for cluster cores). The framework's pixel-level test is most sensitive to
# spatial topology rather than absolute mass, so we set the factor to a
# physically motivated baseline.
# Note: this is a baseline conversion; production-grade work would use full
# stellar population synthesis SED fitting.

F277W_TO_SIGMA_STAR = 1.0e9  # M_sun/Mpc^2 per (MJy/sr)
# Account for non-stellar baryon (gas, ICL): total baryon ~ 6.5x stellar
BARYON_OVER_STAR = 6.5
F277W_TO_SIGMA_BARYON = F277W_TO_SIGMA_STAR * BARYON_OVER_STAR

# Convert M_sun/Mpc^2 to kg/m^2
M_SUN = 1.989e30  # kg
MPC = 3.0857e22   # m
sigma_baryon = flux * F277W_TO_SIGMA_BARYON  # M_sun/Mpc^2
sigma_baryon_kgm2 = sigma_baryon * M_SUN / MPC**2  # kg/m^2

# Replace NaN/Inf with 0
sigma_baryon_kgm2 = np.nan_to_num(sigma_baryon_kgm2, nan=0.0, posinf=0.0, neginf=0.0)
print(f"\nSigma_baryon stats (kg/m^2):")
print(f"  min = {sigma_baryon_kgm2.min():.3e}")
print(f"  max = {sigma_baryon_kgm2.max():.3e}")
print(f"  mean = {sigma_baryon_kgm2.mean():.3e}")
print(f"Sigma_baryon stats (M_sun/pc^2):")
sb_mpc2 = sigma_baryon_kgm2 / M_SUN * (3.0857e16)**2  # to M_sun/pc^2
print(f"  min = {sb_mpc2.min():.3e}")
print(f"  max = {sb_mpc2.max():.3e}")
print(f"  mean = {sb_mpc2.mean():.3e}")

# ===== Reproject Sigma_baryon to target grid =====
print("\n--- Reprojecting Sigma_baryon to kappa_observed grid ---")
# Build a temporary HDU for reprojection
temp_hdu = fits.PrimaryHDU(data=sigma_baryon_kgm2.astype(np.float32),
                            header=src_header)
sigma_target, footprint = reproject_interp(temp_hdu, target_wcs,
                                            shape_out=kappa_obs.shape,
                                            order='bilinear')
sigma_target = np.nan_to_num(sigma_target, nan=0.0, posinf=0.0, neginf=0.0)
print(f"Reprojected coverage: {(footprint > 0).sum()/footprint.size*100:.1f}% of target grid")
print(f"Sigma_target (kg/m^2): min={sigma_target.min():.3e} max={sigma_target.max():.3e}")

# ===== Apply framework pipeline =====
print("\n--- Applying three-layer K_act composite law ---")

# g_bar = 2 * pi * G * Sigma_baryon
G_NEWTON = 6.674e-11  # m^3/kg/s^2
g_bar = 2 * np.pi * G_NEWTON * sigma_target
print(f"g_bar (m/s^2): min={g_bar.min():.3e} max={g_bar.max():.3e}")

# K_act,internal = max(1, sqrt(g_critical/g_bar))
safe_g = np.where(g_bar > 1e-30, g_bar, 1e-30)
K_int = np.maximum(1.0, np.sqrt(G_CRITICAL / safe_g))
print(f"K_act,internal: min={K_int.min():.3f} max={K_int.max():.3f}")

# Lambda = 1/[1 + (K_int-1)/K_struct]
Lambda = 1.0 / (1.0 + (K_int - 1.0) / K_STRUCT)
print(f"Lambda: min={Lambda.min():.5f} max={Lambda.max():.5f}")

# K_act,observed = K_int * F_env * F_ML^-1
K_obs = K_int * F_ENV * F_ML_INV

# kappa_baryonic = Sigma / Sigma_crit
kappa_baryon = sigma_target / Sigma_crit_lens
print(f"kappa_baryonic: min={kappa_baryon.min():.4e} max={kappa_baryon.max():.4e}")

# kappa_pred = K_act_observed * kappa_baryonic
kappa_pred = K_obs * kappa_baryon
print(f"kappa_pred: min={kappa_pred.min():.4e} max={kappa_pred.max():.4e} mean={kappa_pred.mean():.4e}")

# ===== kappa_residual = kappa_observed - kappa_pred =====
print("\n--- Computing kappa_residual ---")
kappa_residual = kappa_obs - kappa_pred
print(f"kappa_residual: min={kappa_residual.min():.4f} max={kappa_residual.max():.4f}")
print(f"  RMS = {np.sqrt(np.mean(kappa_residual**2)):.4f}")
print(f"  abs mean = {np.mean(np.abs(kappa_residual)):.4f}")

# Centroid comparison
def centroid(arr):
    arr = np.maximum(arr, 0)
    total = arr.sum()
    if total <= 0:
        return None
    iy, ix = np.indices(arr.shape)
    cx = (ix * arr).sum() / total
    cy = (iy * arr).sum() / total
    return cx, cy

c_obs = centroid(kappa_obs)
c_pred = centroid(kappa_pred)
print(f"\nkappa_observed centroid (pixel): {c_obs}")
print(f"kappa_pred centroid (pixel):     {c_pred}")
if c_obs and c_pred:
    pix_scale = target_wcs.proj_plane_pixel_scales()[0].to('arcsec').value
    sep_pix = np.sqrt((c_obs[0]-c_pred[0])**2 + (c_obs[1]-c_pred[1])**2)
    sep_arcsec = sep_pix * pix_scale
    print(f"Centroid separation: {sep_pix:.2f} pixels = {sep_arcsec:.2f} arcsec")

# ===== Save outputs =====
print("\n--- Saving outputs ---")
def make_hdu(data, base_header, bunit, comment):
    h = base_header.copy()
    h['BUNIT'] = (bunit, 'Unit')
    h['HISTORY'] = comment
    return fits.PrimaryHDU(data=data.astype(np.float32), header=h)

out_pred = 'bullet_kappa_pred.fits'
out_resid = 'bullet_kappa_residual.fits'
out_sigma = 'bullet_sigma_baryon_proxy.fits'

make_hdu(kappa_pred, target_header, 'dimensionless',
         'kappa_pred from UM/FUM three-layer composite').writeto(out_pred, overwrite=True)
print(f"Wrote {out_pred}  ({os.path.getsize(out_pred)/1024:.1f} KB)")

make_hdu(kappa_residual, target_header, 'dimensionless',
         'kappa_residual = kappa_observed - kappa_pred').writeto(out_resid, overwrite=True)
print(f"Wrote {out_resid}  ({os.path.getsize(out_resid)/1024:.1f} KB)")

make_hdu(sigma_target, target_header, 'kg/m^2',
         'Sigma_baryon proxy from JWST F277W via stellar M/L conversion').writeto(out_sigma, overwrite=True)
print(f"Wrote {out_sigma}  ({os.path.getsize(out_sigma)/1024:.1f} KB)")

# Summary
with open('bullet_kappa_pipeline_summary.txt', 'w') as f:
    f.write("Bullet Cluster Path C Pipeline Summary\n")
    f.write("Charles Anthony Hyatt Battiste UM/FUM\n")
    f.write("USPTO 19/640,364 anchors\n\n")
    f.write("Locked framework constants:\n")
    f.write(f"  phi = {PHI}\n")
    f.write(f"  alpha_struct = {ALPHA_STRUCT}\n")
    f.write(f"  Eidolon = {EIDOLON:.6f}\n")
    f.write(f"  K_struct = 4*phi^3*Eidolon/3 = {K_STRUCT:.4f}\n")
    f.write(f"  g_critical = {G_CRITICAL:.4e} m/s^2\n\n")
    f.write("Cluster-locus environmental parameters:\n")
    f.write(f"  eps_ext = {EPS_EXT}\n")
    f.write(f"  eta_tid = {ETA_TID}\n")
    f.write(f"  Delta_ML = {DELTA_ML}\n")
    f.write(f"  F_env = {F_ENV:.4f}\n")
    f.write(f"  F_ML^-1 = {F_ML_INV:.4f}\n\n")
    f.write(f"Lensing geometry:\n")
    f.write(f"  z_L = {Z_LENS}, z_S = {Z_SOURCE}\n")
    f.write(f"  Sigma_crit_lens = {Sigma_crit_lens:.4e} kg/m^2\n\n")
    f.write(f"Sigma_baryon proxy:\n")
    f.write(f"  Source: JWST NIRCam F277W combined mosaic (Cha 2025 ICL filter)\n")
    f.write(f"  Conversion: 1 MJy/sr -> {F277W_TO_SIGMA_BARYON:.2e} M_sun/Mpc^2 (stars * 6.5 baryon factor)\n")
    f.write(f"  Sigma range (kg/m^2): {sigma_target.min():.3e} to {sigma_target.max():.3e}\n\n")
    f.write(f"kappa_pred stats:\n")
    f.write(f"  min = {kappa_pred.min():.4e}\n")
    f.write(f"  max = {kappa_pred.max():.4e}\n")
    f.write(f"  mean = {kappa_pred.mean():.4e}\n\n")
    f.write(f"kappa_residual = kappa_observed - kappa_pred:\n")
    f.write(f"  min = {kappa_residual.min():.4f}\n")
    f.write(f"  max = {kappa_residual.max():.4f}\n")
    f.write(f"  RMS = {np.sqrt(np.mean(kappa_residual**2)):.4f}\n")
    f.write(f"  abs mean = {np.mean(np.abs(kappa_residual)):.4f}\n")
    if c_obs and c_pred:
        f.write(f"\nCentroid separation: {sep_arcsec:.2f} arcsec\n")

print("\n=== Path C complete ===")