import math
import os
import json
import numpy as np
from scipy.integrate import quad

C_LIGHT = 299792.458  # km/s


def _omega_gamma_h2(Tcmb: float) -> float:
    # Photon density today (approx)
    return 2.469e-5 * (Tcmb / 2.7255) ** 4


def _omega_nu_h2(Tcmb: float, N_eff: float) -> float:
    return _omega_gamma_h2(Tcmb) * 0.2271 * N_eff


def _ez(z: float, om: float, orad: float, ok: float, w0: float, wa: float) -> float:
    a = 1.0 / (1.0 + z)
    # CPL dark energy evolution
    f_de = (a ** (-3.0 * (1.0 + w0 + wa))) * math.exp(-3.0 * wa * (1.0 - a))
    ode = 1.0 - om - orad - ok
    return math.sqrt(
        orad * (1 + z) ** 4
        + om * (1 + z) ** 3
        + ok * (1 + z) ** 2
        + ode * f_de
    )


def _quad_integral(func, a, b):
    val, err = quad(func, a, b, epsabs=1e-10, epsrel=1e-10, limit=400)
    return val


def comoving_distance(z: float, H0: float, om: float, orad: float, ok: float, w0: float, wa: float) -> float:
    invE = lambda zp: 1.0 / _ez(zp, om, orad, ok, w0, wa)
    integral = _quad_integral(invE, 0.0, z)
    return (C_LIGHT / H0) * integral  # Mpc


def z_star(ombh2: float, ommh2: float) -> float:
    g1 = 0.0783 * (ombh2 ** (-0.238)) / (1.0 + 39.5 * (ombh2 ** 0.763))
    g2 = 0.560 / (1.0 + 21.1 * (ombh2 ** 1.81))
    return 1048.0 * (1.0 + 0.00124 * (ombh2 ** (-0.738))) * (1.0 + g1 * (ommh2 ** g2))


def z_drag(ombh2: float, ommh2: float) -> float:
    b1 = 0.313 * (ommh2 ** (-0.419)) * (1.0 + 0.607 * (ommh2 ** 0.674))
    b2 = 0.238 * (ommh2 ** 0.223)
    return 1291.0 * (ommh2 ** 0.251) / (1.0 + 0.659 * (ommh2 ** 0.828)) * (1.0 + b1 * (ombh2 ** b2))


def sound_horizon_rd(
    H0: float,
    omb: float,
    om: float,
    orad: float,
    ok: float,
    w0: float,
    wa: float,
    Tcmb: float,
    N_eff: float,
) -> float:
    """Precision drag-horizon approximation for the active flat-LCDM slice.

    The previous v23 implementation combined the Eisenstein-Hu drag-redshift
    fit with a sound-horizon integral truncated at z=1e5.  At Row-20 precision
    those two approximations compensate numerically and are not suitable for
    the final likelihood.

    Use the Aubourg et al. (2015) numerical approximation

        r_d = 55.154 exp[-72.3 (omega_nu + 0.0006)^2]
              / (omega_cb^0.25351 omega_b^0.12807) Mpc.

    The current UTFANSWF background treats the N_eff neutrino contribution as
    radiation and instantiates no massive-neutrino matter density, therefore
    omega_nu=0 and omega_cb=omega_b+omega_c=Omega_m h^2.
    No fitted calibration or BAO target value enters this formula.
    """
    h = H0 / 100.0
    omega_b = omb * h * h
    omega_cb = om * h * h
    omega_nu = 0.0
    if omega_b <= 0.0 or omega_cb <= 0.0:
        raise ValueError("Physical baryon and cb densities must be positive for r_d.")
    return (
        55.154
        * math.exp(-72.3 * (omega_nu + 0.0006) ** 2)
        / (omega_cb ** 0.25351 * omega_b ** 0.12807)
    )


def sound_horizon_rs(
    z_start: float,
    H0: float,
    omb: float,
    om: float,
    orad: float,
    ok: float,
    w0: float,
    wa: float,
    Tcmb: float,
    N_eff: float = 3.046,
) -> float:
    # Kept for completeness / debugging only.
    og_h2 = _omega_gamma_h2(Tcmb)
    h = H0 / 100.0
    og = og_h2 / (h * h)

    a_start = 1.0 / (1.0 + z_start)

    def E_of_a(a):
        ode = 1.0 - om - orad - ok
        f_de = (a ** (-3.0 * (1.0 + w0 + wa))) * math.exp(-3.0 * wa * (1.0 - a))
        return math.sqrt(orad / (a ** 4) + om / (a ** 3) + ok / (a ** 2) + ode * f_de)

    def integrand(a):
        R = (3.0 / 4.0) * (omb / og) * a
        cs = 1.0 / math.sqrt(3.0 * (1.0 + R))
        return cs / (a * a * E_of_a(a))

    integral = _quad_integral(integrand, 1.0e-8, a_start)
    return (C_LIGHT / H0) * integral


def luminosity_distance(z: float, H0: float, om: float, orad: float, ok: float, w0: float, wa: float) -> float:
    Dc = comoving_distance(z, H0, om, orad, ok, w0, wa)
    return (1.0 + z) * Dc


def distance_modulus(
    z: float,
    H0: float,
    om: float,
    orad: float,
    ok: float,
    w0: float,
    wa: float,
    M_shift: float = 0.0,
) -> float:
    dL_Mpc = luminosity_distance(z, H0, om, orad, ok, w0, wa)
    mu = 5.0 * math.log10(dL_Mpc) + 25.0
    return mu + M_shift


def _infer_gw_length(dataset_dir: str, default_n: int = 3) -> int:
    gw_path = os.path.join(dataset_dir, "gw_data.csv")
    try:
        if os.path.exists(gw_path):
            data = np.loadtxt(gw_path, delimiter=",")
            data = np.atleast_1d(data).reshape(-1)
            return int(data.shape[0])
    except Exception:
        pass
    return int(default_n)


def generate_theory_vectors(params: dict, dataset_dir: str = None) -> dict:
    H0 = float(params.get("H0", 67.4))
    h = H0 / 100.0
    ombh2 = float(params.get("ombh2", 0.02237))
    omch2 = float(params.get("omch2", 0.1200))
    ok = float(params.get("ok", 0.0))
    w0 = float(params.get("w0", -1.0))
    wa = float(params.get("wa", 0.0))
    Tcmb = float(params.get("T_CMB", 2.7255))
    N_eff = float(params.get("N_eff", 3.046))

    omb = ombh2 / (h * h)
    omc = omch2 / (h * h)
    om = omb + omc
    orad = (_omega_gamma_h2(Tcmb) + _omega_nu_h2(Tcmb, N_eff)) / (h * h)

    zst = z_star(ombh2, ombh2 + omch2)
    Dc = comoving_distance(zst, H0, om, orad, ok, w0, wa)
    rd = sound_horizon_rd(H0, omb, om, orad, ok, w0, wa, Tcmb, N_eff)

    # Keep rs available in metadata/debugging only.
    rs = sound_horizon_rs(zst, H0, omb, om, orad, ok, w0, wa, Tcmb)

    # CMB distance priors.  The acoustic angular scale is defined with the
    # sound horizon at last scattering, r_s(z_*), not the drag horizon r_d.
    # Older bundle revisions used r_d multiplied by an empirical 1.022 bridge;
    # v23 derives the conversion directly by using r_s(z_*) in the denominator.
    legacy_cmb_calibration = params.get("legacy_cmb_calibration_ignored", params.get("cmb_calibration", None))
    R = math.sqrt(om) * (H0 / C_LIGHT) * Dc
    lA = math.pi * (Dc / rs)

    # BAO points
    z_bao_dv = params.get("bao_z_dv", None)
    z_bao_dm_h = params.get("bao_z_dm_h", None)

    if z_bao_dv is None:
        z_bao_dv = [0.38, 0.51, 0.61]
    z_bao_dv = [float(z) for z in z_bao_dv]

    if z_bao_dm_h is None:
        legacy = params.get("bao_z", [0.38, 0.51, 0.61])
        z_bao_dm_h = [float(z) for z in legacy]
    else:
        z_bao_dm_h = [float(z) for z in z_bao_dm_h]

    dv_rd = []
    dm_h_vec = []

    rd_fid = float(params.get("rd_fid_Mpc", 147.78))

    for z in z_bao_dv:
        Dc_z = comoving_distance(z, H0, om, orad, ok, w0, wa)
        Hz = H0 * _ez(z, om, orad, ok, w0, wa)
        DA = Dc_z / (1.0 + z)
        DV = (((1.0 + z) ** 2 * DA * DA * (C_LIGHT * z / Hz))) ** (1.0 / 3.0)
        dv_rd.append(DV / rd)

    for z in z_bao_dm_h:
        Dc_z = comoving_distance(z, H0, om, orad, ok, w0, wa)
        Hz = H0 * _ez(z, om, orad, ok, w0, wa)

        dm_obs = Dc_z * (rd_fid / rd)
        h_obs = Hz * (rd / rd_fid)

        dm_h_vec.extend([dm_obs, h_obs])

    # Full DESI DR2 Gaussian BAO likelihood in its native observable basis.
    # Ordering follows desi_gaussian_bao_ALL_GCcomb_mean.txt exactly.
    # No H-like conversion or fiducial-rd rescaling is applied here.
    def native_bao(z):
        DM = comoving_distance(z, H0, om, orad, ok, w0, wa)
        Hz = H0 * _ez(z, om, orad, ok, w0, wa)
        DH = C_LIGHT / Hz
        return DM / rd, DH / rd

    full_desi = []
    dm, dh = native_bao(0.295)
    DV = ((0.295 * (dm * rd) ** 2 * (dh * rd))) ** (1.0 / 3.0)
    full_desi.append(DV / rd)
    for z in (0.510, 0.706, 0.934, 1.321, 1.484):
        dm, dh = native_bao(z)
        full_desi.extend([dm, dh])
    dm, dh = native_bao(2.330)
    full_desi.extend([dh, dm])

    # SN compressed points
    z_sn = params.get("sn_z", [0.10, 0.35, 0.70])
    z_sn = [float(z) for z in z_sn]

    sn_M_shift = float(params.get("sn_M_shift", 0.0))
    sn_mu_pivot = float(params.get("sn_mu_pivot", np.mean(z_sn)))
    legacy_sn_mu_slope = params.get("legacy_sn_mu_slope_ignored", params.get("sn_mu_slope", None))

    # Generate the raw distance-modulus prediction only.  The v23 Row-20 gate
    # profiles the intercept and linear-in-redshift calibration slope jointly as
    # declared covariance-weighted nuisance parameters; no sweep-selected slope
    # is inserted into the physical theory vector.
    sn_vec = []
    for z in z_sn:
        mu = distance_modulus(z, H0, om, orad, ok, w0, wa, M_shift=sn_M_shift)
        sn_vec.append(mu)

    # GW must match gw_data.csv length
    gw_n = _infer_gw_length(dataset_dir) if dataset_dir else int(params.get("gw_n", 3))
    gw_vec = np.full(gw_n, H0, dtype=float)

    return {
        "cmb": np.array([R, lA, ombh2], dtype=float),
        "bao_dv": np.array(dv_rd, dtype=float),
        "bao_dm_h": np.array(dm_h_vec, dtype=float),
        "bao_desi_dr2_full": np.array(full_desi, dtype=float),
        "sn": np.array(sn_vec, dtype=float),
        "gw": gw_vec,
        "meta": {
            "H0": H0,
            "h": h,
            "ombh2": ombh2,
            "omch2": omch2,
            "Omega_m": om,
            "Omega_b": omb,
            "Omega_r": orad,
            "z_star": zst,
            "z_star_method": "Hu-Sugiyama fitting relation already implemented in UTFANSWF",
            "r_d_Mpc": rd,
            "r_d_method": "Aubourg2015_precision_fit__omega_nu_massive=0",
            "r_s_zstar_Mpc": rs,
            "r_d_fid_Mpc": rd_fid,
            "bao_z_dv": z_bao_dv,
            "bao_z_dm_h": z_bao_dm_h,
            "sn_z": z_sn,
            "sn_M_shift": sn_M_shift,
            "sn_mu_pivot": sn_mu_pivot,
            "sn_nuisance_model": "profile GLS intercept + linear redshift slope in Row 20",
            "legacy_sn_mu_slope_ignored": legacy_sn_mu_slope,
            "gw_n": gw_n,
            "cmb_acoustic_scale_definition": "lA = pi * D_M(z_star) / r_s(z_star)",
            "derived_r_d_over_r_s_zstar": rd / rs,
            "legacy_cmb_calibration_ignored": legacy_cmb_calibration,
        }
    }


def write_vector_csv(path: str, vec: np.ndarray):
    with open(path, "w", encoding="utf-8") as f:
        for v in vec.reshape(-1):
            f.write(f"{float(v):.10g}\n")


def maybe_generate(dataset_dir: str, params_path: str = None) -> dict:
    if params_path is None:
        params_path = os.path.join(dataset_dir, "cosmo_params.json")
    if os.path.exists(params_path):
        with open(params_path, "r", encoding="utf-8") as f:
            params = json.load(f)
    else:
        params = {}

    out = generate_theory_vectors(params, dataset_dir=dataset_dir)
    write_vector_csv(os.path.join(dataset_dir, "cmb_theory.csv"), out["cmb"])
    write_vector_csv(os.path.join(dataset_dir, "bao_theory_DV.csv"), out["bao_dv"])
    write_vector_csv(os.path.join(dataset_dir, "bao_theory_DM_H.csv"), out["bao_dm_h"])
    write_vector_csv(os.path.join(dataset_dir, "bao_theory_DESI_DR2_FULL.csv"), out["bao_desi_dr2_full"])
    write_vector_csv(os.path.join(dataset_dir, "sn_theory.csv"), out["sn"])
    write_vector_csv(os.path.join(dataset_dir, "gw_theory.csv"), out["gw"])
    return out