"""Shared helpers for the fp_kiii study (Fundamental-Plane search for dimmed early-type galaxies)."""
import os

import numpy as np
from scipy import optimize, special, stats

HERE = os.path.dirname(os.path.abspath(__file__))
STUDY = os.path.dirname(HERE)
RAW = os.environ.get("FP_KIII_RAW", "fermi_raw/fp_kiii")
DATA = os.path.join(STUDY, "data")
FIGS = os.path.join(STUDY, "figures")
CACHE = os.path.join(RAW, "derived")       # large derived tables (re-creatable) are kept out of the shared folder
for _d in (DATA, FIGS, CACHE):
    os.makedirs(_d, exist_ok=True)


def sample_path(name):
    return os.path.join(CACHE if name == "sdss" else DATA, "sample_%s.csv" % name)


def resid_path(name):
    """Full per-galaxy table with residuals.  For SDSS (114k rows, ~60 MB) the full table lives in the cache;
    data/resid_sdss.csv is a slimmed copy with the columns needed to check the results."""
    return os.path.join(CACHE, "resid_sdss_full.csv") if name == "sdss" else os.path.join(DATA, "resid_%s.csv" % name)

C_KMS = 299792.458
H0 = 70.0          # km/s/Mpc, used for every physical length / volume quoted
OM = 0.3

# dimming fractions used throughout
F_GRID = np.array([0.3, 0.4, 0.5, 0.6, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95])


# ----------------------------------------------------------------------------
# cosmology (flat LCDM), fast interpolation
# ----------------------------------------------------------------------------
_zg = np.linspace(0.0, 0.3, 3001)
_Ez = np.sqrt(OM * (1 + _zg) ** 3 + 1 - OM)
_dc = np.concatenate([[0.0], np.cumsum(0.5 * (1 / _Ez[1:] + 1 / _Ez[:-1]) * np.diff(_zg))]) * C_KMS / H0


def comoving_distance(z):
    """Line-of-sight comoving distance in Mpc (H0=70, Om=0.3, flat)."""
    return np.interp(z, _zg, _dc)


def angular_diameter_distance(z):
    return comoving_distance(z) / (1.0 + np.asarray(z))


def luminosity_distance(z):
    return comoving_distance(z) * (1.0 + np.asarray(z))


def z_cmb_from_helio(z_helio, ra_deg, dec_deg):
    """Heliocentric -> CMB-frame redshift (Planck 2018 dipole: 369.82 km/s to l=264.021, b=48.253)."""
    # dipole apex in equatorial J2000: RA=167.942, Dec=-6.944
    ra0, de0, v = np.radians(167.942), np.radians(-6.944), 369.82
    ra, de = np.radians(ra_deg), np.radians(dec_deg)
    cosang = np.sin(de) * np.sin(de0) + np.cos(de) * np.cos(de0) * np.cos(ra - ra0)
    return (1.0 + z_helio) * (1.0 + v * cosang / C_KMS) - 1.0


def radec_to_xyz(ra_deg, dec_deg, dist):
    ra, de = np.radians(ra_deg), np.radians(dec_deg)
    return np.column_stack([dist * np.cos(de) * np.cos(ra), dist * np.cos(de) * np.sin(ra), dist * np.sin(de)])


def radec_to_galactic(ra_deg, dec_deg):
    """Equatorial J2000 -> Galactic (l, b) in degrees."""
    ra, de = np.radians(ra_deg), np.radians(dec_deg)
    ragp, degp, lncp = np.radians(192.85948), np.radians(27.12825), np.radians(122.93192)
    sinb = np.sin(de) * np.sin(degp) + np.cos(de) * np.cos(degp) * np.cos(ra - ragp)
    b = np.arcsin(np.clip(sinb, -1, 1))
    y = np.cos(de) * np.sin(ra - ragp)
    x = np.sin(de) * np.cos(degp) - np.cos(de) * np.sin(degp) * np.cos(ra - ragp)
    l = (lncp - np.arctan2(y, x)) % (2 * np.pi)
    return np.degrees(l), np.degrees(b)


# ----------------------------------------------------------------------------
# statistics
# ----------------------------------------------------------------------------
def poisson_ul(n, cl=0.95):
    """Classical one-sided Poisson upper limit on the mean given n observed events."""
    return stats.chi2.ppf(cl, 2 * (np.asarray(n) + 1)) / 2.0


def mad_sigma(x):
    x = np.asarray(x)
    return 1.4826 * np.median(np.abs(x - np.median(x)))


def fit_truncated_plane(y, X, ylim=None, yerr2=None, clip=3.5, n_iter=6, verbose=False):
    """Maximum-likelihood linear regression  y ~ N(X.c, s_int^2 + yerr2)  for a sample
    truncated at y > ylim (per-object lower limits; pass None for no truncation).

    Robustness: objects with |standardised residual| > clip are excluded from the *fit*
    (never from the sample) and the fit is iterated.  Returns dict with coefficients,
    intrinsic scatter, model mean and total sigma for every object, and the fit mask.
    """
    y = np.asarray(y, float)
    X = np.asarray(X, float)
    n, k = X.shape
    yerr2 = np.zeros(n) if yerr2 is None else np.asarray(yerr2, float)
    use = np.isfinite(y) & np.all(np.isfinite(X), axis=1)
    c = np.linalg.lstsq(X[use], y[use], rcond=None)[0]
    s = np.std(y[use] - X[use] @ c)
    p = np.concatenate([c, [np.log(max(s, 1e-3))]])

    def nll(p, m):
        mu = X[m] @ p[:k]
        var = np.exp(2 * p[k]) + yerr2[m]
        sd = np.sqrt(var)
        ll = -0.5 * ((y[m] - mu) / sd) ** 2 - np.log(sd)
        if ylim is not None:
            ll -= special.log_ndtr((mu - ylim[m]) / sd)
        return -np.sum(ll)

    mask = use.copy()
    for it in range(n_iter):
        res = optimize.minimize(nll, p, args=(mask,), method="BFGS")
        p = res.x
        mu = X @ p[:k]
        sd = np.sqrt(np.exp(2 * p[k]) + yerr2)
        zres = (y - mu) / sd
        newmask = use & (np.abs(zres) < clip)
        if verbose:
            print("  iter", it, "nfit", mask.sum(), "coef", np.round(p[:k], 4), "s_int", round(float(np.exp(p[k])), 4))
        if np.array_equal(newmask, mask):
            break
        mask = newmask
    # covariance from BFGS inverse Hessian (approximate)
    try:
        cov = res.hess_inv
        err = np.sqrt(np.diag(cov))
    except Exception:  # pragma: no cover
        err = np.full(k + 1, np.nan)
    return dict(coef=p[:k], coef_err=err[:k], s_int=float(np.exp(p[k])), mu=mu, sd=sd,
                resid=y - mu, zres=zres, fitmask=mask)


def weighted_orthogonal_fp(r, s, i):
    """Classical FP  r = a s + b i + c  by direct least squares on r and by orthogonal (PCA) fit,
    with 3-sigma iterative clipping used for the fit only.  Returns dict."""
    out = {}
    X = np.column_stack([s, i, np.ones_like(s)])
    m = np.isfinite(r) & np.isfinite(s) & np.isfinite(i)
    for _ in range(6):
        c = np.linalg.lstsq(X[m], r[m], rcond=None)[0]
        res = r - X @ c
        sd = mad_sigma(res[m])
        m2 = np.isfinite(res) & (np.abs(res) < 3 * sd)
        if np.array_equal(m, m2):
            break
        m = m2
    out["direct"] = dict(a=c[0], b=c[1], c=c[2], rms=float(np.std(res[m])), mad=float(sd))
    # orthogonal
    P = np.column_stack([r, s, i])
    mo = np.all(np.isfinite(P), axis=1)
    for _ in range(6):
        mean = P[mo].mean(0)
        w, v = np.linalg.eigh(np.cov((P[mo] - mean).T))
        nvec = v[:, 0]                 # smallest-variance direction
        d = (P - mean) @ nvec
        sd = mad_sigma(d[mo])
        m2 = np.all(np.isfinite(P), axis=1) & (np.abs(d) < 3 * sd)
        if np.array_equal(mo, m2):
            break
        mo = m2
    a = -nvec[1] / nvec[0]
    b = -nvec[2] / nvec[0]
    c0 = mean[0] - a * mean[1] - b * mean[2]
    res_o = r - (a * s + b * i + c0)
    out["orthogonal"] = dict(a=a, b=b, c=c0, rms_r=float(np.std(res_o[mo])), mad_r=float(mad_sigma(res_o[mo])),
                             rms_orth=float(np.std(d[mo])))
    out["resid_r_direct"] = res
    out["resid_r_orth"] = res_o
    return out