#!/usr/bin/env python3
"""Fit the competing distance-redshift laws to the Pantheon+ SN Ia sample.

Diagonal (statistical) errors only -- the full 1701x1701 covariance is not used,
so the chi-squared differences below are UPPER bounds on the real significance.
The absolute magnitude / H0 combination is marginalised analytically, so only
the SHAPE of the distance-redshift relation is being tested.
"""
import math
from models import lcdm, milne, lisle, lisle_fixed, tired, mu

ZMIN = 0.023          # standard Pantheon+ cut: below this, peculiar velocities dominate

rows = []
with open("pantheon.dat") as fh:
    hdr = fh.readline().split()
    ix = {n: i for i, n in enumerate(hdr)}
    for line in fh:
        p = line.split()
        if not p:
            continue
        if int(p[ix["IS_CALIBRATOR"]]) == 1:
            continue
        z = float(p[ix["zHD"]])
        if z < ZMIN:
            continue
        rows.append((z,
                     float(p[ix["MU_SH0ES"]]),
                     float(p[ix["MU_SH0ES_ERR_DIAG"]])))

print("supernovae used: %d   z range %.3f - %.3f" %
      (len(rows), min(r[0] for r in rows), max(r[0] for r in rows)))

MODELS = [("LCDM (Om=0.3)", lcdm), ("Milne / empty", milne),
          ("Lisle Doppler", lisle), ("Lisle + reciprocity", lisle_fixed),
          ("Tired light", tired)]

results = {}
for name, f in MODELS:
    # analytic marginalisation over a constant offset
    sw = sd = 0.0
    pred = []
    for z, m, e in rows:
        _da, dl = f(z)
        pm = mu(dl)
        pred.append(pm)
        w = 1.0 / e ** 2
        sw += w
        sd += w * (m - pm)
    off = sd / sw
    chi2 = 0.0
    for (z, m, e), pm in zip(rows, pred):
        chi2 += ((m - pm - off) / e) ** 2
    results[name] = (chi2, off)

base = results["LCDM (Om=0.3)"][0]
print("\n%-22s %12s %10s %10s" % ("model", "chi2", "chi2/dof", "dchi2"))
for name, _f in MODELS:
    c, off = results[name]
    print("%-22s %12.1f %10.3f %10.1f" %
          (name, c, c / (len(rows) - 1), c - base))

print("\nsqrt(dchi2) ~ nominal sigma (diagonal errors only, so an over-estimate):")
for name, _f in MODELS:
    c, _ = results[name]
    if c > base:
        print("  %-22s %.1f" % (name, math.sqrt(c - base)))

# binned residuals against LCDM, for the article table
print("\nbinned Hubble-diagram residuals, mag  (model - LCDM, after offset fit)")
edges = [0.023, 0.05, 0.1, 0.2, 0.3, 0.5, 0.7, 1.0, 1.5, 2.3]
print("%-14s" % "z bin" + "".join("%16s" % n for n, _ in MODELS))
for a, b in zip(edges[:-1], edges[1:]):
    sel = [r for r in rows if a <= r[0] < b]
    if not sel:
        continue
    zc = sum(r[0] for r in sel) / len(sel)
    line = "%-14s" % ("%.3f-%.3f" % (a, b))
    for name, f in MODELS:
        off = results[name][1]
        off0 = results["LCDM (Om=0.3)"][1]
        _da, dl = f(zc)
        _da0, dl0 = lcdm(zc)
        line += "%16.3f" % ((mu(dl) + off) - (mu(dl0) + off0))
    print(line + "   (N=%d, <z>=%.3f)" % (len(sel), zc))