#!/usr/bin/env python3
"""Atlas UMMO — Janus / DESI DR2, 2026-09-09.

Run: python reproduire.py [--data DIR] [--output DIR]
Offline, versioned inputs included. No downloaded source code is executed.
Two separate experiments: per-bin scale-profiled geometry, then common-ruler BAO.
Editorial calculation assisted by AI; not an independently reviewed replication.
"""
import os
os.environ.setdefault('OPENBLAS_NUM_THREADS', '1')
os.environ.setdefault('OMP_NUM_THREADS', '1')
import argparse
import csv
import hashlib
import json
import platform
from pathlib import Path
import numpy as np
import scipy
from numpy.polynomial.legendre import leggauss
from scipy.integrate import quad
from scipy.linalg import cho_factor, cho_solve
from scipy.optimize import minimize_scalar, differential_evolution, brentq, least_squares
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

FILES = {
    'desi_gaussian_bao_ALL_GCcomb_mean.txt': '9ac154ab583ce759c0f7eef3c978c7c70a6ead2d18774caceadf1a350a640585',
    'desi_gaussian_bao_ALL_GCcomb_cov.txt': '252a143274c8a07c78694c119617d36594f6d7965d00319ca611c6ffb886e509',
}
TRACERS = {0.295: 'BGS', 0.510: 'LRG1', 0.706: 'LRG2', 0.934: 'LRG3+ELG1', 1.321: 'ELG2', 1.484: 'QSO', 2.330: 'Lya'}
LABELS = {'janus': 'Janus · branche sans pression', 'lcdm': 'ΛCDM plat', 'cpl': 'w₀wₐCDM plat · exploratoire'}
GX, GW = leggauss(96)


def distances(z, model, theta):
    """Dimensionless (H0/c) D_M and (H0/c) D_H, fixed c, z=a0/a-1."""
    z = np.asarray(z, dtype=float)
    if model == 'janus':
        q = float(theta[0])
        if q > 0 or np.any(1 + 2*q*z <= 0):
            raise ValueError('Outside expanding accelerating dust branch')
        root = np.sqrt(1+2*q*z)
        # Rationalized Mattig distance, finite at q=0. Omega_k = 1-2*q.
        dl = z + z*z*(1-q)/(1+q*z+root)
        return dl/(1+z), 1/((1+z)*root)
    om = float(theta[0])
    w0, wa = (-1., 0.) if model == 'lcdm' else (float(theta[1]), float(theta[2]))
    def einv(t):
        de = np.exp(3*(1+w0+wa)*np.log1p(t) - 3*wa*t/(1+t))
        return 1/np.sqrt(om*(1+t)**3 + (1-om)*de)
    t = z[..., None]*(GX+1)/2
    return z/2*np.sum(GW*einv(t), axis=-1), einv(z)


class Experiment:
    def __init__(self, z, y, quantities, cov, geometry=False):
        self.z, self.y, self.quantities, self.cov = z, y, quantities, cov
        self.geometry = geometry
        self.cf = cho_factor(cov, lower=True)
        self.groups = np.unique(z)

    def design(self, model, theta):
        dm, dh = distances(self.z, model, theta)
        v = np.where(self.quantities == 'DM_over_rs', dm,
                     np.where(self.quantities == 'DH_over_rs', dh, np.cbrt(self.z*dm*dm*dh)))
        if self.geometry:
            # One independent length per redshift; only D_M/D_H constrains shape.
            return v[:, None]*(self.z[:, None] == self.groups[None, :])
        return v[:, None]

    def objective(self, model, theta, details=False):
        x = self.design(model, theta)
        ci_x = cho_solve(self.cf, x)
        coeff = np.linalg.solve(x.T@ci_x, ci_x.T@self.y)
        if np.any(coeff <= 0):
            raise ValueError('Unphysical negative profiled scale')
        pred = x@coeff
        res = self.y-pred
        value = float(res@cho_solve(self.cf, res))
        return (value, coeff, pred) if details else value

    def fit(self, model):
        if model == 'cpl':
            # Fixed exploration box, not claimed to be a DESI posterior analysis.
            bounds = [(0.01, 0.99), (-3., 1.), (-3., 2.)]
            opts = [differential_evolution(lambda p: self.objective(model, p), bounds,
                    seed=s, tol=1e-10, popsize=18, maxiter=1000, polish=True) for s in [41, 73]]
            assert all(o.success for o in opts)
            assert abs(opts[0].fun-opts[1].fun) < 1e-5
            opt = min(opts, key=lambda o: o.fun)
            theta = list(map(float, opt.x))
            intervals = None
            agreement = float(abs(opts[0].fun-opts[1].fun))
        else:
            bounds = [(-0.5/self.z.max()+1e-8, 0.)] if model == 'janus' else [(0., 1.)]
            fun = lambda p: self.objective(model, [p])
            grid = np.linspace(*bounds[0], 801)
            vals = np.array([fun(p) for p in grid])
            idx = int(np.argmin(vals))
            opt = minimize_scalar(fun, bounds=(grid[max(0, idx-1)], grid[min(len(grid)-1, idx+1)]),
                                  method='bounded', options={'xatol': 1e-13})
            assert opt.success
            candidates = [(float(opt.x), float(opt.fun)), (bounds[0][0], fun(bounds[0][0])), (bounds[0][1], fun(bounds[0][1]))]
            par, value = min(candidates, key=lambda v: v[1])
            assert value <= vals.min()+1e-7
            theta = [float(par)]
            intervals = [float(brentq(lambda p: fun(p)-value-1, min(par, edge), max(par, edge)))
                         if fun(edge) > value+1 else None for edge in bounds[0]]
            agreement = float(value-vals.min())
        value, coeff, pred = self.objective(model, theta, True)
        k = len(theta)+len(coeff)
        boundary = [name for name, p, (lo, hi) in zip(['q0' if model == 'janus' else 'Omega_m', 'w0', 'wa'], theta, bounds)
                    if min(p-lo, hi-p) < 1e-5]
        return {'theta': theta, 'parameter_names': ['q0'] if model == 'janus' else (['Omega_m'] if model == 'lcdm' else ['Omega_m', 'w0', 'wa']),
                'bounds': bounds, 'profile_delta_chi2_1': intervals, 'scales': coeff.tolist(),
                'chi2': value, 'n': len(self.y), 'k': k, 'dof_nominal': len(self.y)-k,
                'aic': value+2*k, 'boundary_parameters': boundary,
                'optimizer_check': agreement, 'predicted': pred.tolist()}


def load(data):
    for name, sha in FILES.items():
        assert hashlib.sha256((data/name).read_bytes()).hexdigest() == sha, name+' changed'
    raw = np.loadtxt(data/list(FILES)[0], dtype=str, skiprows=1)
    z, y, quantities = raw[:, 0].astype(float), raw[:, 1].astype(float), raw[:, 2]
    cov = np.loadtxt(data/list(FILES)[1])
    assert len(y) == 13 and cov.shape == (13, 13)
    assert len(set(zip(z, quantities))) == 13
    assert np.array_equal(np.unique(z), np.array(list(TRACERS)))
    np.testing.assert_allclose(cov, cov.T, rtol=0, atol=1e-12)
    cho_factor(cov)
    assert np.isfinite(cov).all() and np.isfinite(y).all()
    # Current source is block diagonal between redshift bins.
    assert np.all(cov[z[:, None] != z[None, :]] == 0)
    assert quantities[11] == 'DH_over_rs' and quantities[12] == 'DM_over_rs'
    return z, y, quantities, cov


def subset(data, geometry=False, omit=None):
    z, y, qty, cov = data
    keep = np.ones(len(z), dtype=bool)
    if geometry: keep &= qty != 'DV_over_rs'
    if omit is not None: keep &= z != omit
    return Experiment(z[keep], y[keep], qty[keep], cov[np.ix_(keep, keep)], geometry)


def checks(data, fits):
    zz = np.linspace(.001, 2.33, 40)
    max_error = 0.
    for q in [-.2, -.087, -.001, 0.]:
        dm, _ = distances(zz, 'janus', [q])
        k = 1-2*q
        numeric = np.array([np.sinh(np.sqrt(k)*quad(lambda t: 1/((1+t)*np.sqrt(1+2*q*t)), 0, z,
                            epsabs=1e-12, epsrel=1e-12)[0])/np.sqrt(k) for z in zz])
        max_error = max(max_error, float(np.max(np.abs(dm/numeric-1))))
    assert max_error < 1e-10
    dm, dh = distances(zz, 'janus', [0.])
    np.testing.assert_allclose(dm/dh, zz+zz*zz/2, rtol=1e-13)
    np.testing.assert_allclose(distances(zz, 'lcdm', [1.])[0], 2*(1-1/np.sqrt(1+zz)), rtol=1e-12)
    np.testing.assert_allclose(distances(zz, 'lcdm', [.3]), distances(zz, 'cpl', [.3, -1, 0]), rtol=1e-13)
    for model, truth in [('janus', -.08), ('lcdm', .30)]:
        ex = subset(data)
        pred = ex.design(model, [truth])[:, 0]*30.
        syn = Experiment(ex.z, pred, ex.quantities, ex.cov)
        recovered = syn.fit(model)
        assert abs(recovered['theta'][0]-truth) < 1e-6 and abs(recovered['scales'][0]-30) < 1e-5
    # Expected published BAO-only LCDM benchmark ~10.2. Use the actual pinned
    # covariance, whose errors differ from Table 4; omit radiation/neutrino detail.
    assert abs(fits['joint']['lcdm']['chi2']-10.2) < .6
    # Check exact profiling vs a direct 2D distance-to-line residual at each bin.
    ex = subset(data, geometry=True)
    for model in ['janus', 'lcdm']:
        theta = fits['geometry'][model]['theta']
        total = 0.
        for z in ex.groups:
            im = int(np.flatnonzero((ex.z == z)&(ex.quantities == 'DM_over_rs'))[0])
            ih = int(np.flatnonzero((ex.z == z)&(ex.quantities == 'DH_over_rs'))[0])
            dm, dh = distances(np.array([z]), model, theta)
            f = float(dm[0]/dh[0])
            total += (ex.y[im]-f*ex.y[ih])**2/(ex.cov[im, im]+f*f*ex.cov[ih, ih]-2*f*ex.cov[im, ih])
        assert abs(total-fits['geometry'][model]['chi2']) < 1e-8
    # A separate optimizer and quadrature implementation of the JOINT fit.
    direct = {}
    ex = subset(data)
    lc = np.linalg.cholesky(ex.cov)
    for model, x0, bounds in [('janus', [-.08, 30.], ([-.2145922, 0.], [0., 100.])),
                              ('lcdm', [.3, 30.], ([0., 0.], [1., 100.]))]:
        def residual(p):
            t0, amplitude = p
            e = (lambda t: (1+t)*np.sqrt(1+2*t0*t)) if model == 'janus' else (lambda t: np.sqrt(t0*(1+t)**3+1-t0))
            dc = np.array([quad(lambda t: 1/e(t), 0, z)[0] for z in ex.z])
            dh = 1/e(ex.z)
            dm = np.sinh(np.sqrt(1-2*t0)*dc)/np.sqrt(1-2*t0) if model == 'janus' else dc
            v = np.where(ex.quantities == 'DM_over_rs', dm, np.where(ex.quantities == 'DH_over_rs', dh, np.cbrt(ex.z*dm**2*dh)))
            return np.linalg.solve(lc, ex.y-amplitude*v)
        fit = least_squares(residual, x0, bounds=bounds, gtol=1e-10, ftol=1e-12, xtol=1e-12)
        assert fit.success
        direct[model] = float(fit.fun@fit.fun)
        assert abs(direct[model]-fits['joint'][model]['chi2']) < 1e-6
    return {'quadrature_max_relative_error': max_error, 'lcdm_published_benchmark_approx': 10.2,
            'lcdm_recalculated': fits['joint']['lcdm']['chi2'],
            'direct_quadrature_joint_chi2': direct,
            'checks': 'input SHA256; 13 unique rows; ordering Lyα DH then DM; Cholesky; block independence; quadrature; Milne / Einstein-de Sitter / CPL-LCDM limits; noiseless recovery; profiling identity; optimizer scans and two seeds; direct joint least-squares with independent quadrature'}


def products(data, fits, out):
    z, y, qty, cov = data
    plt.rcParams.update({'font.family': 'DejaVu Sans', 'font.size': 11, 'svg.fonttype': 'none',
                         'svg.hashsalt': 'atlas-janus-desi-v1', 'axes.spines.top': False, 'axes.spines.right': False})
    colors = {'janus': '#007e91', 'lcdm': '#a94813', 'cpl': '#7762a8'}
    rows = []
    for zz in np.unique(z[qty != 'DV_over_rs']):
        im = int(np.flatnonzero((z == zz)&(qty == 'DM_over_rs'))[0])
        ih = int(np.flatnonzero((z == zz)&(qty == 'DH_over_rs'))[0])
        f = y[im]/y[ih]
        sigma = np.sqrt(cov[im, im]+f*f*cov[ih, ih]-2*f*cov[im, ih])/y[ih]
        row = {'tracer': TRACERS[zz], 'z': float(zz), 'F_AP': float(f), 'sigma_F_AP_linear': float(sigma)}
        for model in ['janus', 'lcdm']:
            dm, dh = distances(np.array([zz]), model, fits['geometry'][model]['theta'])
            row[model] = float(dm[0]/dh[0])
        rows.append(row)
    with (out/'rapports.csv').open('w') as fp:
        writer = csv.DictWriter(fp, fieldnames=list(rows[0])); writer.writeheader(); writer.writerows(rows)
    fig, ax = plt.subplots(figsize=(10, 5.5), layout='constrained')
    zs = np.linspace(.01, 2.33, 350)
    ax.errorbar([r['z'] for r in rows], [r['F_AP'] for r in rows], yerr=[r['sigma_F_AP_linear'] for r in rows],
                fmt='o', color='#253044', capsize=3, label='DESI DR2 · erreurs propagées pour l’affichage')
    for model in ['janus', 'lcdm']:
        dm, dh = distances(zs, model, fits['geometry'][model]['theta'])
        ax.plot(zs, dm/dh, color=colors[model], label=LABELS[model], lw=2)
    ax.set(xlabel='Décalage vers le rouge z', ylabel='F_AP = Dₘ / Dₕ', title='La géométrie sans fixer la taille de la règle acoustique')
    ax.legend(frameon=False); ax.grid(alpha=.15)
    fig.savefig(out/'geometrie.svg', metadata={'Date': None}); plt.close(fig)
    residual_rows = []
    for i in range(len(y)):
        row = {'tracer': TRACERS[z[i]], 'z': float(z[i]), 'observable': qty[i], 'measured': float(y[i]), 'sigma_marginal': float(np.sqrt(cov[i, i]))}
        for model in ['janus', 'lcdm', 'cpl']:
            pred = fits['joint'][model]['predicted'][i]
            row[model+'_predicted'] = pred
            row[model+'_residual_sigma_marginal'] = float((y[i]-pred)/np.sqrt(cov[i, i]))
        residual_rows.append(row)
    with (out/'residus.csv').open('w') as fp:
        writer = csv.DictWriter(fp, fieldnames=list(residual_rows[0])); writer.writeheader(); writer.writerows(residual_rows)
    fig, ax = plt.subplots(figsize=(10, 6), layout='constrained')
    x = np.arange(13)
    for model, shift in [('janus', -.12), ('lcdm', .12)]:
        label = 'Janus · limite q₀ = 0' if model == 'janus' else LABELS[model]
        ax.errorbar(x+shift, [r[model+'_residual_sigma_marginal'] for r in residual_rows], yerr=1,
                    fmt='o', capsize=2, ms=4, color=colors[model], label=label)
    labels = [TRACERS[zz].replace('LRG3+ELG1', 'LRG3+\nELG1')+'\n'+qq[:2] for zz, qq in zip(z, qty)]
    ax.set_xticks(x, labels, fontsize=9)
    ax.axhline(0, color='#333333', lw=1); ax.grid(axis='y', alpha=.15)
    ax.set(ylabel='(Mesure − prédiction) / incertitude marginale', title='Distances BAO : résidus après ajustement d’une échelle commune')
    ax.legend(frameon=False)
    fig.savefig(out/'residus.svg', metadata={'Date': None}); plt.close(fig)
    return rows


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--data', type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument('--output', type=Path, default=Path(__file__).resolve().parent)
    args = parser.parse_args(); args.output.mkdir(parents=True, exist_ok=True)
    data = load(args.data)
    fits = {mode: {model: subset(data, geometry=(mode == 'geometry')).fit(model)
                  for model in ['janus', 'lcdm']} for mode in ['geometry', 'joint']}
    fits['joint']['cpl'] = subset(data).fit('cpl')
    # Defined sensitivity: remove each whole tracer bin, refit both baseline models.
    fits['leave_one_bin_out'] = [{'omitted': name, 'z': z, **{model: subset(data, omit=z).fit(model)
                                   for model in ['janus', 'lcdm']}} for z, name in TRACERS.items()]
    fits['checks'] = checks(data, fits)
    fits['ratios'] = products(data, fits, args.output)
    fits['environment'] = {'python': platform.python_version(), 'numpy': np.__version__, 'scipy': scipy.__version__, 'matplotlib': matplotlib.__version__}
    fits['date'] = '2026-09-09'
    fits['interpretation'] = 'Conditional Gaussian compressed-distance test; not raw spectra, not an early-universe calculation, not a probability for a model. Geometry and joint chi2 must never be added.'
    (args.output/'resultats.json').write_text(json.dumps(fits, indent=2, ensure_ascii=False)+'\n')
    print(json.dumps({mode: {model: {k: v for k, v in result.items() if k not in ['predicted', 'optimizer_check']}
                               for model, result in fits[mode].items()} for mode in ['geometry', 'joint']}, indent=2))
    print('Checks:', fits['checks'])
    print('Leave-one-bin-out delta chi2:', [(r['omitted'], r['janus']['chi2']-r['lcdm']['chi2']) for r in fits['leave_one_bin_out']])


if __name__ == '__main__': main()