#!/usr/bin/env python3
"""Reproduce the public87-row SH0ES teaching subsystem, Python3 standard library only.

Place shoes-subsystem.json and shoes-subsystem-benchmarks.json beside this file.
This code is an independent implementation of GLS, not the author's MCMC code.
"""
import argparse
import hashlib
import json
import math
from pathlib import Path

INPUT_SHA256 = '24b93d2ce7bdbcbfe9789ac35091dc5737e326beecedf0fa84a80d02ed287447'

def cholesky(matrix):
    size = len(matrix)
    factor = [[0.0] * size for _ in range(size)]
    for row in range(size):
        for column in range(row + 1):
            value = matrix[row][column] - sum(factor[row][k] * factor[column][k] for k in range(column))
            if row == column:
                if value <= 0:
                    raise ValueError('Covariance is not positive definite')
                factor[row][column] = math.sqrt(value)
            else:
                factor[row][column] = value / factor[column][column]
    return factor

def forward(factor, vector):
    result = []
    for row, value in enumerate(vector):
        result.append((value - sum(factor[row][k] * result[k] for k in range(row))) / factor[row][row])
    return result

def backward(factor, vector):
    size = len(vector)
    result = [0.0] * size
    for row in reversed(range(size)):
        result[row] = (vector[row] - sum(factor[k][row] * result[k] for k in range(row + 1, size))) / factor[row][row]
    return result

def dot(left, right):
    return sum(a*b for a, b in zip(left, right))

def rank(matrix):
    basis = []
    for column in zip(*matrix):
        residual = list(column)
        for _ in range(2):
            for vector in basis:
                projection = dot(residual, vector)
                residual = [a-projection*b for a, b in zip(residual, vector)]
        norm = math.sqrt(dot(residual, residual))
        if norm > 1e-10:
            basis.append([value/norm for value in residual])
    return len(basis)

def fit(data, anchor=0.0, zero_point=0.0, common=False, no_anchor=False, drop_metallicity=False, diagonal=False):
    selected = [i for i, row in enumerate(data['observations'])
                if not (no_anchor and row['kind'] == 'geometric-anchor')
                and not (drop_metallicity and row['sourceRowIndex'] == 3209)]
    design = [data['designMatrix'][i] for i in selected]
    values = []
    for i in selected:
        row = data['observations'][i]
        value = row['preparedMagnitude']
        if row['kind'] == 'geometric-anchor':
            value += anchor
        if row['kind'] == 'cepheid' and (common or row['host'] != 'N4258'):
            value += zero_point
        values.append(value)
    covariance = [[data['covariance'][i][j] if not diagonal or i == j else 0.0 for j in selected] for i in selected]
    actual_rank = rank(design)
    identified = actual_rank == 8
    columns = list(range(8)) if identified else [0, 1, 3, 4, 5, 6, 7]
    factor = cholesky(covariance)
    whitened_values = forward(factor, values)
    whitened_columns = [forward(factor, [row[column] for row in design]) for column in columns]
    normal = [[dot(a, b) for b in whitened_columns] for a in whitened_columns]
    normal_factor = cholesky(normal)
    solve = lambda vector: backward(normal_factor, forward(normal_factor, vector))
    parameters_reduced = solve([dot(column, whitened_values) for column in whitened_columns])
    parameters = [0.0] * 8
    for column, value in zip(columns, parameters_reduced):
        parameters[column] = value
    inverse_columns = [solve([float(i == j) for i in range(len(columns))]) for j in range(len(columns))]
    inverse = list(map(list, zip(*inverse_columns)))
    residuals = [value-dot(row, parameters) for row, value in zip(design, values)]
    whitened_residuals = forward(factor, residuals)
    relative_variance = inverse[0][0] + inverse[1][1] - 2*inverse[0][1]
    result = {'identified': identified, 'rank': actual_rank, 'parameters': parameters,
              'parameterInterpretation': 'identified' if identified else 'gauge-fixed for display only',
              'chiSquared': dot(whitened_residuals, whitened_residuals),
              'physicalSlope': data['basePeriodLuminositySlope'] + parameters[4],
              'hostModulusDifference': parameters[1]-parameters[0],
              'hostModulusDifferenceSigma': math.sqrt(relative_variance),
              'h0Median': None, 'h0Central68': None, 'covariance': inverse if identified else None}
    if identified:
        sigma = math.sqrt(inverse[7][7])
        result['h0Median'] = 10 ** (parameters[7]/5)
        result['h0Central68'] = [10 ** ((parameters[7]+sign*sigma)/5) for sign in (-1, 1)]
    return result

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--data', type=Path, default=Path(__file__).with_name('shoes-subsystem.json'))
    parser.add_argument('--reference', type=Path, default=None)
    parser.add_argument('--output', type=Path, default=Path('public-reproduction.json'))
    args = parser.parse_args()
    raw = args.data.read_bytes()
    if hashlib.sha256(raw).hexdigest() != INPUT_SHA256:
        raise ValueError('Input checksum differs from the frozen selected data')
    data = json.loads(raw)
    assert len(data['observations']) == 87
    original_columns = data['originalParameterIndices']
    for row, coefficients in zip(data['designMatrix'], data['observations']):
        for column, value in coefficients['originalNonzeroColumns']:
            assert column in original_columns, 'An original nonzero nuisance column was omitted'
            assert row[original_columns.index(column)] == value
    baseline = fit(data)
    variants = {'anchorPlus0.1': fit(data, anchor=0.1), 'hostZeroPointPlus0.1': fit(data, zero_point=0.1),
                'commonZeroPointPlus0.1': fit(data, zero_point=0.1, common=True), 'noAnchor': fit(data, no_anchor=True),
                'withoutMetallicitySummary': fit(data, drop_metallicity=True),
                'diagonalCovarianceCounterfactual': fit(data, diagonal=True)}
    assert variants['noAnchor']['rank'] == 7 and variants['noAnchor']['h0Median'] is None
    assert abs(variants['commonZeroPointPlus0.1']['h0Median'] / baseline['h0Median'] - 1) < 1e-10
    expected_ratio = 10**(-0.1/5)
    assert abs(variants['hostZeroPointPlus0.1']['h0Median']/baseline['h0Median'] - expected_ratio) < 1e-10
    if args.reference:
        reference = json.loads(args.reference.read_text())['default']
        assert abs(baseline['h0Median'] - reference['h0Median']) < 1e-8
        assert max(abs(a-b) for a, b in zip(baseline['parameters'], reference['parameters'])) < 1e-8
    result = {'sourceSha256': INPUT_SHA256, 'baseline': baseline, 'variants': variants,
              'scope': 'Conditional87-row measured subsystem, not a full SH0ES reanalysis or published precision reproduction.'}
    args.output.write_text(json.dumps(result, indent=2, allow_nan=False) + '\n')
    print(json.dumps({'rows': 87, 'h0': baseline['h0Median'], 'central68': baseline['h0Central68'],
                      'noAnchorRank': variants['noAnchor']['rank'], 'output': str(args.output)}))

if __name__ == '__main__':
    main()