#!/usr/bin/env python3
"""Public standalone SDSS exercise: Python 3 standard library, no site checkout.

Download this file, sdss-dr17.csv and photometry-grid.json into one directory.
Optionally download sdss-results.json for --reference verification.
"""
import argparse
import csv
import hashlib
import json
import math
from pathlib import Path


def logsumexp(values):
    maximum = max(values)
    return maximum + math.log(sum(math.exp(value - maximum) for value in values))


def cosmic_age(redshift):
    return 2 * (977.79222168 / 70) / (3 * math.sqrt(0.7)) * math.asinh(math.sqrt(0.7 / 0.3) / (1 + redshift) ** 1.5)


def amplitude_marginal(flux, sigma, shape):
    precision = 1 / 10000**2
    overlap = squared = 0
    for value, error, ratio in zip(flux, sigma, shape):
        precision += ratio**2 / error**2
        overlap += value * ratio / error**2
        squared += value**2 / error**2
    argument = overlap / math.sqrt(precision)
    log_cdf = math.log(0.5 * math.erfc(-argument / math.sqrt(2)))
    return -(squared - overlap**2 / precision) / 2 + math.log(2) + log_cdf - math.log(10000) - math.log(precision) / 2, max(0, overlap / precision)


def infer(flux, sigma, grid):
    log_density = []
    best_states = []
    age_count, dust_count = len(grid['agesGyr']), len(grid['dust'])
    for redshift_index, redshift in enumerate(grid['redshifts']):
        alternatives = []
        best = (-math.inf, 0, 0, 0, [])
        for age_index, age in enumerate(grid['agesGyr']):
            if age > cosmic_age(redshift):
                continue
            for dust_index in range(dust_count):
                start = ((redshift_index * age_count + age_index) * dust_count + dust_index) * 5
                shape = grid['ratios'][start:start + 5]
                weight, amplitude = amplitude_marginal(flux, sigma, shape)
                alternatives.append(weight)
                if weight > best[0]:
                    best = (weight, age_index, dust_index, amplitude, shape)
        log_density.append(logsumexp(alternatives) - math.log(len(alternatives)))
        best_states.append(best)
    maximum = max(log_density)
    relative = [math.exp(value - maximum) for value in log_density]
    z = grid['redshifts']
    norm = sum((relative[index] + relative[index - 1]) * (z[index] - z[index - 1]) / 2 for index in range(1, len(z)))
    density = [value / norm for value in relative]
    cumulative = [0]
    for index in range(1, len(z)):
        cumulative.append(cumulative[-1] + (density[index] + density[index - 1]) * 0.0025)
    def quantile(probability):
        right = next(index for index, value in enumerate(cumulative) if value >= probability)
        fraction = (probability - cumulative[right - 1]) / (cumulative[right] - cumulative[right - 1])
        return z[right - 1] + fraction * 0.005
    mode_index = log_density.index(maximum)
    best = best_states[mode_index]
    prediction = [best[3] * ratio for ratio in best[4]]
    return {'mode': z[mode_index], 'median': quantile(0.5), 'interval68': [quantile(0.16), quantile(0.84)],
            'density': density, 'bestAgeIndex': best[1], 'bestDustIndex': best[2], 'amplitude': best[3],
            'predictedFlux': prediction, 'chiSquared': sum(((value - model) / error)**2 for value, model, error in zip(flux, prediction, sigma))}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=Path, default=Path('sdss-dr17.csv'))
    parser.add_argument('--grid', type=Path, default=Path('photometry-grid.json'))
    parser.add_argument('--output', type=Path, default=Path('reproduced-sdss.json'))
    parser.add_argument('--reference', type=Path)
    args = parser.parse_args()
    raw = args.input.read_bytes()
    if hashlib.sha256(raw).hexdigest() != 'fd93102dd5ab47d780c470ff9687dccfab7370c46bb244744d373dd7222fb869':
        raise SystemExit('The frozen SDSS source checksum differs')
    grid_bytes = args.grid.read_bytes()
    if hashlib.sha256(grid_bytes).hexdigest() != 'e831838cfe343e0aeb4b18ad2c29dfff19711217341d53e36d560ac9bcc3fe56':
        raise SystemExit('The frozen photometry grid checksum differs')
    grid = json.loads(grid_bytes)
    rows = csv.DictReader(line for line in raw.decode().splitlines() if not line.startswith('#'))
    offsets = [-0.04, 0, 0, 0, 0.02]
    output = []
    for row in rows:
        flux, sigma = [], []
        for band, offset in zip('ugriz', offsets):
            factor = 10 ** (0.4 * float(row['extinction_' + band]) - 0.4 * offset)
            value = float(row['modelFlux_' + band]) * factor
            error = math.hypot(factor / math.sqrt(float(row['modelFluxIvar_' + band])), 0.03 * abs(value))
            flux.append(value)
            sigma.append(error)
        posterior = infer(flux, sigma, grid)
        truth = float(row['z'])
        output.append({'objectId': row['objID'], 'spectrumId': row['specObjID'], 'flux': flux, 'sigma': sigma,
                       'posterior': posterior, 'spectroscopicRedshift': truth,
                       'coveredByCentral68': posterior['interval68'][0] <= truth <= posterior['interval68'][1]})
    comparison = None
    if args.reference:
        expected = json.loads(args.reference.read_text())['results']
        maximum_difference = 0
        for result, reference in zip(output, expected):
            assert result['objectId'] == reference['objectId']
            assert result['posterior']['mode'] == reference['posterior']['mode']
            maximum_difference = max(maximum_difference, max(abs(left - right) for left, right in zip(result['posterior']['density'], reference['posterior']['density'])))
        assert len(output) == len(expected) == 24
        assert maximum_difference < 1e-6
        comparison = {'maximumAbsoluteDensityDifference': maximum_difference, 'modeAgreementCount': 24}
    document = {'source': 'Frozen SDSS DR17 + finite BC03/Calzetti/SDSS2010 grid', 'amplitudePrior': 'Half-normal, scale10000 nMgy; integrated analytically',
                'summary': {'count': len(output), 'central68CoverageCount': sum(result['coveredByCentral68'] for result in output)},
                'referenceComparison': comparison, 'results': output}
    args.output.write_text(json.dumps(document, indent=2, allow_nan=False) + '\n')
    print(json.dumps({'output': str(args.output), 'summary': document['summary'], 'referenceComparison': comparison}, indent=2))


if __name__ == '__main__':
    main()