#!/usr/bin/env python3
"""Atlas UMMO: recalcul JLA/Janus 2018, 2026-09-09.

Run: python reproduire.py [--archive /path/jla_likelihood_v6.tgz] [--output DIR]
Without --archive the official JLA archive is downloaded and SHA-256 checked.
Only numeric text members are read; no upstream code/binary is executed.
The main comparison fixes alpha/beta/host step and uses the SAME full C.
The joint refit follows the published JLA chi-square criterion (no log det C).
This is an editorial numerical study assisted by AI, not a peer-reviewed replication.
"""
import os
os.environ.setdefault('OPENBLAS_NUM_THREADS', '1')
os.environ.setdefault('OMP_NUM_THREADS', '1')
import argparse
import hashlib
import io
import json
from pathlib import Path
import platform
import tarfile
import urllib.request
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, solve
from scipy.optimize import brentq, minimize, minimize_scalar
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

DATA_URL = 'https://supernovae.in2p3.fr/sdss_snls_jla/jla_likelihood_v6.tgz'
DATA_SHA = '2f06277628ab53ca6590c7da679714b39b2b5a0973c2f8978cebdb1626e762f4'
PAPER_URL = 'https://www.jp-petit.org/papers/cosmo/2018-AstrophysSpaceSci.pdf'
PAPER_SHA = 'f68fece1b103e0ef7757cd0e4c509425562cf5055285fbfd7358a2b4561f26c9'
A0, B0, STEP0 = .141, 3.101, -.070
GX, GW = leggauss(64)


def janus(z, q):
    """Dimensionless D_L = H0*d_L/c, Eq. 5/25; expanding dust branch."""
    z = np.asarray(z)
    if q > 0 or np.any(1 + 2*q*z <= 0):
        raise ValueError('Janus requires q <= 0 and 1 + 2*q*z > 0')
    return z + z*z*(1-q)/(1+q*z+np.sqrt(1+2*q*z))


def lcdm(z, omega, radiation=0.):
    """Flat FLRW distance; radiation=0 for the common SN-only comparison."""
    z = np.asarray(z)
    t = z[..., None]*(GX+1)/2
    e = np.sqrt(omega*(1+t)**3 + radiation*(1+t)**4 + 1-omega-radiation)
    return (1+z)*z/2*np.sum(GW/e, axis=-1)


def read_data(blob):
    if hashlib.sha256(blob).hexdigest() != DATA_SHA:
        raise ValueError('JLA archive hash differs: review the upstream revision before use.')
    with tarfile.open(fileobj=io.BytesIO(blob), mode='r:gz') as tf:
        def raw(name):
            return tf.extractfile('jla_likelihood_v6/data/'+name).read()
        data = np.genfromtxt(io.BytesIO(raw('jla_lcparams.txt')), names=True, dtype=None, encoding='utf8')
        matrices = []
        for name in ['v0', 'va', 'vb', 'v0a', 'v0b', 'vab']:
            v = np.loadtxt(io.BytesIO(raw(f'jla_{name}_covmatrix.dat')))
            assert int(v[0]) == len(data) and len(v) == len(data)**2+1
            matrices.append(v[1:].reshape(len(data),len(data)))
    assert len(data) == 740 and len(set(data['name'])) == 740
    return data, matrices


class Experiment:
    def __init__(self, data, matrices):
        self.d, self.m = data, matrices
        self.z = data['zcmb']
        self.host = (data['3rdvar'] > 10).astype(float)
        self.ones = np.ones(len(data))
        self.X = np.column_stack([self.ones, self.host])
        self.frame = 5*np.log10((1+data['zhel'])/(1+self.z))
        self.bounds = {'janus': (-.5/self.z.max()+1e-7, 0.), 'lcdm': (0.,1.)}

    def covariance(self, a, b):
        m,d = self.m,self.d
        c = m[0]+a*a*m[1]+b*b*m[2]+2*a*m[3]-2*b*m[4]-2*a*b*m[5]
        # dmb already contains intrinsic/lensing/redshift dispersion: do not add twice.
        diagonal = d['dmb']**2+(a*d['dx1'])**2+(b*d['dcolor'])**2+2*a*d['cov_m_s']-2*b*d['cov_m_c']-2*a*b*d['cov_s_c']
        c = c + np.diag(diagonal)
        assert np.isfinite(c).all() and np.allclose(c,c.T,atol=1e-10,rtol=0)
        return c

    def shape(self, model, theta, use_frames=True, z=None):
        zz = self.z if z is None else z
        mu = 5*np.log10(janus(zz,theta) if model=='janus' else lcdm(zz,theta))
        if use_frames:
            if z is not None: raise ValueError('Observed frame correction needs the catalogue rows')
            mu = mu + self.frame
        return mu

    def conditional(self, model, mode='full', host_step=STEP0, use_frames=True):
        c = self.covariance(A0,B0)
        if mode == 'diagonal': c = np.diag(np.diag(c))
        cf = cho_factor(c,lower=True)
        ci1 = cho_solve(cf,self.ones)
        y = self.d['mb']+A0*self.d['x1']-B0*self.d['color']+19.05-host_step*self.host
        def objective(theta, values=False):
            r = y-self.shape(model,theta,use_frames)
            offset = ci1@r/(ci1@self.ones)
            residual = r-offset
            value = float(residual@cho_solve(cf,residual))
            return (value,float(offset),residual) if values else value
        fit = minimize_scalar(objective,bounds=self.bounds[model],method='bounded',options={'xatol':1e-11})
        assert fit.success
        # Numerical scan catches a missed lower basin or a preferred boundary.
        scan = [objective(t) for t in np.linspace(*self.bounds[model],101)]
        assert fit.fun <= min(scan)+1e-7
        value,offset,residual = objective(fit.x,True)
        interval = []
        for endpoint in self.bounds[model]:
            interval.append(float(brentq(lambda t:objective(t)-value-1, min(endpoint,fit.x),max(endpoint,fit.x))) if objective(endpoint)>value+1 else None)
        return {'parameter':float(fit.x),'profile_delta_chi2_1':interval,'offset':offset,'chi2':value,'n':740,'k':2,'dof':738,'alpha':A0,'beta':B0,'host_step':host_step,'covariance':mode,'redshift_frames':use_frames},residual,c

    def joint(self, model, use_frames=True):
        def objective(v, values=False):
            theta,a,b=v
            cf=cho_factor(self.covariance(a,b),lower=True)
            y=self.d['mb']+a*self.d['x1']-b*self.d['color']+19.05-self.shape(model,theta,use_frames)
            ciX=cho_solve(cf,self.X)
            linear=solve(self.X.T@ciX,ciX.T@y,assume_a='pos')
            r=y-self.X@linear
            val=float(r@cho_solve(cf,r))
            return (val,linear) if values else val
        starts = [[-.067 if model=='janus' else .295,A0,B0],[-.2 if model=='janus' else .6,.18,2.7]]
        opts=[minimize(objective,s,method='Nelder-Mead',bounds=[self.bounds[model],(.01,.3),(1,5)],options={'xatol':1e-8,'fatol':1e-8,'maxiter':800}) for s in starts]
        assert all(o.success for o in opts) and abs(opts[0].fun-opts[1].fun)<1e-6
        fit=min(opts,key=lambda o:o.fun)
        for v,(lo,hi) in zip(fit.x,[self.bounds[model],(.01,.3),(1,5)]): assert lo+1e-5<v<hi-1e-5
        val,linear=objective(fit.x,True)
        return {'parameter':float(fit.x[0]),'alpha':float(fit.x[1]),'beta':float(fit.x[2]),'offset':float(linear[0]),'host_step':float(linear[1]),'chi2':val,'n':740,'k':5,'dof':735,'redshift_frames':use_frames,'optimizer_agreement':float(abs(opts[0].fun-opts[1].fun))}


def numerical_checks(ex):
    z=np.geomspace(.00001,1.299106,50)
    np.testing.assert_allclose(janus(z,0),z+z*z/2,rtol=1e-14)
    # Independent quadrature using E(z)=(1+z)*sqrt(1+2*q*z), Omega_k=1-2*q.
    maxerr=0.
    for q in [-.3,-.087,-.001]:
        k=1-2*q
        expected=np.array([(1+zz)/np.sqrt(k)*np.sinh(np.sqrt(k)*quad(lambda t:1/((1+t)*np.sqrt(1+2*q*t)),0,zz,epsabs=1e-12)[0]) for zz in z])
        maxerr=max(maxerr,float(np.max(np.abs(janus(z,q)/expected-1))))
    assert maxerr<1e-9
    np.testing.assert_allclose(lcdm(z,1),2*(1+z)*(1-1/np.sqrt(1+z)),rtol=1e-9)
    # Recover the shape of noiseless synthetic distances with an unknown offset.
    for model,truth in [('janus',-.087),('lcdm',.295)]:
        y=ex.shape(model,truth)+43.2
        def f(t):
            r=y-ex.shape(model,t);r-=r.mean();return float(r@r)
        fit=minimize_scalar(f,bounds=ex.bounds[model],method='bounded',options={'xatol':1e-11})
        assert abs(fit.x-truth)<1e-6
    # Public JLA test.cc benchmark: zcmb only, fixed rounded Table 10 parameters.
    y=ex.d['mb']+A0*ex.d['x1']-B0*ex.d['color']+19.05-STEP0*ex.host
    mu=5*np.log10(lcdm(ex.z,.295)*299792.458/70)+25
    r=y-mu;ref=float(r@cho_solve(cho_factor(ex.covariance(A0,B0)),r))
    assert abs(ref-682.9)<.05 # Radiation omitted, rounded nuisance parameters.
    return {'JLA_reference_chi2':ref,'JLA_reference_expected_approx':682.9,'Janus_quadrature_max_relative_error':maxerr,'checks':'740 unique rows; covariance symmetry/Cholesky; Milne and Einstein-de Sitter limits; noiseless parameter recovery; quadrature; JLA reference; two-start joint fits'}


def plots(ex,out,results,residuals,cov):
    plt.rcParams.update({'font.family':'DejaVu Sans','font.size':11,'svg.fonttype':'none','svg.hashsalt':'atlas-janus-v1','axes.spines.top':False,'axes.spines.right':False})
    colors={'janus':'#007e91','lcdm':'#a94813'}
    zz=np.geomspace(ex.z.min(),ex.z.max(),400)
    fig,axes=plt.subplots(3,1,figsize=(10.5,10.5),height_ratios=[1.4,1,1],sharex=True,layout='constrained')
    y=ex.d['mb']+A0*ex.d['x1']-B0*ex.d['color']+19.05-STEP0*ex.host-ex.frame
    err=np.sqrt(np.diag(cov))
    axes[0].errorbar(ex.z,y,yerr=err,fmt='.',ms=2,color='#6c7882',alpha=.32,elinewidth=.5,label='740 SNe JLA · erreurs marginales')
    curve_columns=[zz]
    for i,model in enumerate(['janus','lcdm']):
        f=results['conditional'][model]
        curve=ex.shape(model,f['parameter'],False,zz)+f['offset']
        curve_columns.append(curve)
        label='Janus 2018' if model=='janus' else 'ΛCDM plat'
        axes[0].plot(zz,curve,color=colors[model],lw=2,label=label)
        axes[i+1].errorbar(ex.z,residuals[model],yerr=err,fmt='.',ms=2,color=colors[model],alpha=.35,elinewidth=.5)
        axes[i+1].axhline(0,color='#222222',lw=.8)
        axes[i+1].set_ylabel(f'Résidu {label}\n(magnitude)')
        axes[i+1].set_ylim(-1.,1.)
    axes[0].set_ylabel('Module de distance\nnormalisation conventionnelle')
    axes[0].legend(loc='lower right',frameon=False,fontsize=10)
    axes[0].set_title('Même catalogue, même covariance, deux paramètres ajustés',loc='left',pad=16,fontweight='bold')
    axes[-1].set_xlabel('Décalage vers le rouge z (référentiel CMB, échelle logarithmique)')
    for ax in axes:ax.set_xscale('log');ax.grid(alpha=.15)
    fig.savefig(out/'comparaison.svg',metadata={'Date':None})
    plt.close(fig)
    np.savetxt(out/'courbes.csv',np.column_stack(curve_columns),delimiter=',',header='zcmb,mu_janus,mu_lcdm',comments='',fmt='%.12g')
    # The direct difference uses individually fitted offsets, no imposed common H0.
    fig,ax=plt.subplots(figsize=(10.5,4.5),layout='constrained')
    delta=curve_columns[1]-curve_columns[2]
    ax.plot(zz,delta,color=colors['janus'],lw=2)
    ax.axhline(0,color='#222222',lw=.8);ax.grid(alpha=.15)
    ax.set_xlabel('Décalage vers le rouge z');ax.set_ylabel('μ Janus − μ ΛCDM (magnitude)')
    ax.set_title('L’écart entre les deux courbes apparaît après soustraction',loc='left',pad=14,fontweight='bold')
    fig.savefig(out/'difference.svg',metadata={'Date':None});plt.close(fig)
    return {'maximum_absolute_curve_difference_mag':float(np.max(np.abs(delta))),'at_redshift':float(zz[np.argmax(np.abs(delta))])}


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--archive',type=Path)
    parser.add_argument('--output',type=Path,default=Path(__file__).resolve().parent)
    args=parser.parse_args();out=args.output;out.mkdir(parents=True,exist_ok=True)
    if args.archive:blob=args.archive.read_bytes()
    else:
        print('Downloading official JLA data (24 MB)...',flush=True)
        with urllib.request.urlopen(DATA_URL,timeout=120) as response:blob=response.read()
    data,matrices=read_data(blob);ex=Experiment(data,matrices)
    result={'analysis_date':'2026-09-09','protocol':'Atlas JLA v1','data_url':DATA_URL,'data_sha256':DATA_SHA,'paper_url':PAPER_URL,'paper_sha256':PAPER_SHA,'paper_doi':'10.1007/s10509-018-3365-3','published_2018':{'q0':-.087,'error_reported':.015,'chi2':657,'dof':738},'versions':{'python':platform.python_version(),'numpy':np.__version__,'scipy':scipy.__version__,'matplotlib':matplotlib.__version__},'zmin':float(ex.z.min()),'zmax':float(ex.z.max()),'checks':numerical_checks(ex),'conditional':{},'joint':{},'sensitivity':{}}
    result['revision_2024'] = {
        'doi':'10.1140/epjc/s10052-024-13569-w',
        'publication_date':'2024-11-26',
        'article_number':1226,
        'pdf_sha256':'f7d507f78f952d0bf5727e7e27945c178c77a7ab2ac5df10552026e3b6a6811c',
        'scope':'Documentary comparison: section 10 figure 10 cites reference 4 (2018). No new SN fit is assigned a separate chi-square here. Revised field equations are not fully rederived by this script.'
    }
    residuals={}
    for model in ['janus','lcdm']:
        fit,residual,cov=ex.conditional(model);result['conditional'][model]=fit;residuals[model]=residual
        result['joint'][model]=ex.joint(model)
        result['sensitivity'][model]={
            'diagonal_only':ex.conditional(model,mode='diagonal')[0],
            'no_host_step':ex.conditional(model,host_step=0.)[0],
            'zcmb_only':ex.conditional(model,use_frames=False)[0]
        }
    for key in ['conditional','joint']:
        result[key]['delta_chi2_janus_minus_lcdm']=result[key]['janus']['chi2']-result[key]['lcdm']['chi2']
    result['plots']=plots(ex,out,result,residuals,cov)
    import csv
    with (out/'residus.csv').open('w',newline='',encoding='utf8') as f:
        writer=csv.writer(f);writer.writerow(['name','zcmb','zhel','mu_standardized','sigma_marginal','mu_janus','mu_lcdm','residual_janus','residual_lcdm'])
        yy=data['mb']+A0*data['x1']-B0*data['color']+19.05-STEP0*ex.host
        for i in range(740):writer.writerow([data['name'][i],ex.z[i],data['zhel'][i],yy[i],np.sqrt(cov[i,i]),yy[i]-residuals['janus'][i],yy[i]-residuals['lcdm'][i],residuals['janus'][i],residuals['lcdm'][i]])
    (out/'resultats.json').write_text(json.dumps(result,ensure_ascii=False,indent=2)+'\n',encoding='utf8')
    print(json.dumps({k:result[k] for k in ['checks','conditional','joint','plots']},ensure_ascii=False,indent=2))

if __name__=='__main__':main()