"""Recompute the HERA H1C IDR2 upper limits and a window-weighted model comparison.

Python standard library only. Inputs, in the same folder or given on the command line:
  hera-h1c-idr2-window-extract.json  derived extract of the public HERA Phase I files
  model-spectra.json                 density-only model on the HERA window grid (ideal reference)

Usage:
  python3 -I recompute_window_comparison.py [extract.json] [model.json] [output.json]

Units: k in h Mpc^-1; Delta^2, sigma and upper limits in mK^2; covariance in mK^4.
"""
import json
import math
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
FIRST_DECIMATED_K = {0: 0.256, 1: 0.192}  # Band 1 starts at 0.256, Band 2 at 0.192
K_MAX = 0.6


def upper_limit(estimate, sigma, keep_negative=False, sigma_multiple=2.0):
    """HERA rule: max(estimate, 0) + 2 sigma. Options reproduce two common mistakes."""
    central = estimate if keep_negative else max(estimate, 0.0)
    return central + sigma_multiple * sigma


def windowed(window_row, model):
    if len(window_row) != len(model):
        raise ValueError('window row and model grid differ in length')
    return sum(weight * value for weight, value in zip(window_row, model))


def decimated_rows(band):
    start = FIRST_DECIMATED_K[band['spectralWindow']]
    rows = [row for row in band['rows'] if row['kHPerMpc'] >= start - 1e-9 and row['kHPerMpc'] <= K_MAX + 1e-9]
    return rows[::2]


def compare(extract, model):
    field = next(entry for entry in extract['fields'] if entry['field'] == 1)
    model_bands = {band['spectralWindow']: band for band in model['bands']}
    output = []
    for band in field['bands']:
        reference = model_bands[band['spectralWindow']]
        if any(abs(a - b) > 1e-9 for a, b in zip(reference['kHPerMpc'], band['windowKHPerMpc'])):
            raise ValueError('model grid does not match the window grid')
        rows = []
        for row in decimated_rows(band):
            weights = row['windowRow']
            if abs(sum(weights) - 1.0) > 1e-6:
                raise ValueError('window row is not normalised to unit sum')
            limit = upper_limit(row['estimateMk2'], row['oneSigmaMk2'])
            entry = {
                'kHPerMpc': round(row['kHPerMpc'], 6),
                'estimateMk2': row['estimateMk2'], 'oneSigmaMk2': row['oneSigmaMk2'],
                'upperLimitMk2': limit, 'upperLimitMk': math.sqrt(limit),
                'failureCases': {
                    'keepNegativeEstimateMk2': upper_limit(row['estimateMk2'], row['oneSigmaMk2'], keep_negative=True),
                    'oneSigmaOnlyMk2': upper_limit(row['estimateMk2'], row['oneSigmaMk2'], sigma_multiple=1.0),
                },
                'models': {},
            }
            index = row['index'] - 1  # window columns start at k index 1
            for name in ('heated', 'cold'):
                curve = reference['deltaSquaredMk2'][name]
                with_window = windowed(weights, curve)
                entry['models'][name] = {
                    'windowedMk2': with_window, 'atBinCentreMk2': curve[index],
                    'windowedOverLimit': with_window / limit,
                    'hUnitMistakeWindowedMk2': windowed(weights, reference['hUnitMistakeDeltaSquaredMk2'][name]),
                    'excluded': with_window > limit,
                }
            rows.append(entry)
        # Neighbouring native rows overlap through the taper; report their correlation coefficient.
        correlations = []
        for first, second in zip(band['rows'][:-1], band['rows'][1:]):
            covariance = first['covarianceWithNextSelectedMk4']
            correlations.append({'kPairHPerMpc': [round(first['kHPerMpc'], 6), round(second['kHPerMpc'], 6)],
                                 'correlation': covariance / (first['oneSigmaMk2'] * second['oneSigmaMk2'])})
        output.append({'spectralWindow': band['spectralWindow'], 'redshift': band['bandCentreRedshift'], 'rows': rows,
                       'adjacentNativeRowCorrelation': correlations})
    return output


def main():
    extract_path = Path(sys.argv[1]) if len(sys.argv) > 1 else HERE / 'hera-h1c-idr2-window-extract.json'
    model_path = Path(sys.argv[2]) if len(sys.argv) > 2 else HERE / 'model-spectra.json'
    output_path = Path(sys.argv[3]) if len(sys.argv) > 3 else None
    extract = json.loads(extract_path.read_text())
    model = json.loads(model_path.read_text())
    result = {'formatVersion': 1, 'python': sys.version.split()[0],
              'labels': {'limits': extract['label'], 'models': model['label']},
              'upperLimitRule': 'max(estimate, 0) + 2 sigma', 'bands': compare(extract, model)}
    text = json.dumps(result, indent=1) + '\n'
    if output_path:
        output_path.write_text(text)
    for band in result['bands']:
        print(f"redshift {band['redshift']:.2f}")
        for row in band['rows']:
            heated, cold = row['models']['heated'], row['models']['cold']
            print(f"  k={row['kHPerMpc']:.3f}  U=({row['upperLimitMk']:.2f})^2 mK^2  heated/U={heated['windowedOverLimit']:.4f}  cold/U={cold['windowedOverLimit']:.3f}")


if __name__ == '__main__':
    main()