#!/usr/bin/env python3
"""Extract the fixed 87-row SH0ES teaching subsystem, retaining full covariance.

Requires NumPy. The 48 MB source is cached outside the repository. No source code
from the unlicensed author repository is copied or executed by this script.
"""
import argparse
import hashlib
import json
from pathlib import Path
import urllib.request

import numpy as np

_source_parents = Path(__file__).resolve().parents
ROOT = _source_parents[3] if len(_source_parents) > 3 and (_source_parents[3] / 'src/features/labs').exists() else Path.cwd()
OUTPUT = ROOT / 'src/features/labs/distance-ladder/data'
DOCUMENTS = ROOT / 'docs/astronomy/ac08'
COMMIT = 'c447f0fea703fcd0fff57de5000947b5ca81286b'
N = 3492
BASE_SLOPE = -3.285
PARAMETER_INDICES = [0, 1, 37, 38, 41, 42, 43, 46]
PARAMETER_NAMES = ['muM101', 'muM1337', 'deltaMuN4258', 'cepheidAbsoluteW', 'slopeCorrection', 'supernovaAbsoluteB', 'metallicityCoefficient', 'fiveLogH0']
FILES = {
    'ally_shoes_ceph_topantheonwt6.0_112221.fits': '10bb034ca3fff53f6625c9fe47ebd054b37ef45d1140af1c766831cfa6871433',
    'alll_shoes_ceph_topantheonwt6.0_112221.fits': '9a2ce872dd20ed4fdf5005ce62a805d3eefbbb0b20b0c614013e7c0094009db8',
    'allc_shoes_ceph_topantheonwt6.0_112221.fits': 'a42778672d25df7a559bd2949b1e412b99b2e35dd9a36895f6b38828be019172',
}


def source_url(name):
    host = 'https://media.githubusercontent.com/media/' if name.startswith('allc_') else 'https://raw.githubusercontent.com/'
    return host + f'PantheonPlusSH0ES/DataRelease/{COMMIT}/SH0ES_Data/{name}'


def source_image(path, shape):
    content = path.read_bytes()
    assert b'BITPIX  =                  -32' in content[:2880]
    return np.frombuffer(content, dtype='>f4', count=int(np.prod(shape)), offset=2880).reshape(shape).astype(float)


def gls(design, values, covariance):
    factor = np.linalg.cholesky(covariance)
    whitened_design = np.linalg.solve(factor, design)
    whitened_values = np.linalg.solve(factor, values)
    parameters = np.linalg.lstsq(whitened_design, whitened_values, rcond=None)[0]
    parameter_covariance = np.linalg.inv(whitened_design.T @ whitened_design)
    residuals = whitened_values - whitened_design @ parameters
    sigma_five_log_h0 = np.sqrt(parameter_covariance[-1, -1])
    return {'parameters': parameters.tolist(), 'covariance': parameter_covariance.tolist(),
            'chiSquared': float(residuals @ residuals), 'rank': int(np.linalg.matrix_rank(whitened_design)),
            'h0Median': float(10 ** (parameters[-1] / 5)),
            'h0Central68': (10 ** ((parameters[-1] + np.array([-1, 1]) * sigma_five_log_h0) / 5)).tolist(),
            'physicalSlope': float(BASE_SLOPE + parameters[4]),
            'hostDistanceCorrelation': float(parameter_covariance[0, 1] / np.sqrt(parameter_covariance[0, 0] * parameter_covariance[1, 1]))}


def save_json(path, value):
    path.write_text(json.dumps(value, separators=(',', ':'), allow_nan=False) + '\n')
    return {'path': str(path.relative_to(ROOT)) if path.is_relative_to(ROOT) else path.name, 'bytes': path.stat().st_size,
            'sha256': hashlib.sha256(path.read_bytes()).hexdigest()}


def main():
    global OUTPUT, DOCUMENTS
    parser = argparse.ArgumentParser()
    parser.add_argument('--source-cache', type=Path, default=Path('/tmp/sandboxphysics-ac08-source'))
    parser.add_argument('--fetch', action='store_true')
    parser.add_argument('--output-directory', type=Path, help='Write the selected data, benchmark and provenance together in this directory')
    parser.add_argument('--full-reference', action='store_true', help='Also solve the complete author system once offline')
    args = parser.parse_args()
    if args.output_directory:
        OUTPUT = DOCUMENTS = args.output_directory.resolve()
    OUTPUT.mkdir(parents=True, exist_ok=True)
    DOCUMENTS.mkdir(parents=True, exist_ok=True)
    args.source_cache.mkdir(parents=True, exist_ok=True)
    metadata = []
    for name, expected_hash in FILES.items():
        path = args.source_cache / name
        if not path.exists():
            if not args.fetch:
                raise SystemExit(f'Missing {path}; --fetch downloads the pinned input')
            path.write_bytes(urllib.request.urlopen(source_url(name)).read())
        content = path.read_bytes()
        assert hashlib.sha256(content).hexdigest() == expected_hash
        metadata.append({'filename': name, 'url': source_url(name), 'bytes': len(content), 'sha256': expected_hash})
    table_path = args.source_cache / 'table2.tex'
    if not table_path.exists():
        if not args.fetch:
            raise SystemExit('Missing metadata table; add --fetch')
        table_path.write_bytes(urllib.request.urlopen(source_url('table2.tex')).read())
    table = []
    for line in table_path.read_text().splitlines():
        row = line.split('&')
        if len(row) >= 10 and row[4].strip().replace('.', '', 1).isdigit():
            table.append(row)
    assert len(table) == 3130
    y = source_image(args.source_cache / next(name for name in FILES if name.startswith('ally_')), (N,))
    stored_l = source_image(args.source_cache / next(name for name in FILES if name.startswith('alll_')), (47, N))
    full_c = np.memmap(args.source_cache / next(name for name in FILES if name.startswith('allc_')), dtype='>f4', offset=2880, shape=(N, N))
    for constraint_row in [3209, 3212, 3213]:
        cross_covariance = np.asarray(full_c[constraint_row], dtype=float).copy()
        cross_covariance[constraint_row] = 0
        assert np.count_nonzero(cross_covariance) == 0, 'Released constraint has cross covariance that must be retained'
    indices = sorted(set(np.linspace(0, 258, 24).round().astype(int).tolist()
                         + list(range(259, 274))
                         + np.linspace(2150, 2592, 24).round().astype(int).tolist()
                         + list(range(3130, 3135)) + [3209, 3212, 3213]
                         + np.linspace(3215, 3491, 16).round().astype(int).tolist()))
    assert len(indices) == 87
    original_design = stored_l[:, indices].T
    omitted_columns = [index for index in range(47) if index not in PARAMETER_INDICES]
    assert np.count_nonzero(original_design[:, omitted_columns]) == 0, 'No omitted parameter may be silently fixed'
    design = original_design[:, PARAMETER_INDICES]
    values = y[indices]
    covariance = np.asarray(full_c[np.ix_(indices, indices)], dtype=float)
    observations = []
    for local_index, source_index in enumerate(indices):
        record = {'sourceRowIndex': source_index, 'observationId': f'SH0ES_y[{source_index}]', 'preparedMagnitude': float(values[local_index]),
                  'originalNonzeroColumns': [[int(column), float(original_design[local_index, column])] for column in np.flatnonzero(original_design[local_index])]}
        if source_index < 3130:
            row = table[source_index]
            period = 10 ** (1 + stored_l[41, source_index])
            assert abs(period / float(row[4]) - 1) < 0.00051
            host = row[0].strip()
            assert host in ['M101', 'M1337', 'N4258']
            record.update({'kind': 'cepheid', 'host': host, 'starId': row[3].strip(),
                           'raDegrees': float(row[1]), 'decDegrees': float(row[2]), 'periodDays': float(period),
                           'tablePeriodDaysRounded': float(row[4]), 'f160wMagnitudeRounded': float(row[7]),
                           'logPeriodOffset': float(stored_l[41, source_index]),
                           'relativeMetallicityDex': float(stored_l[43, source_index])})
        elif source_index < 3207:
            record.update({'kind': 'calibrator-supernova', 'host': 'M101' if source_index < 3132 else 'M1337',
                           'supernova': '2011fe' if source_index < 3132 else '2006D'})
        elif source_index == 3213:
            assert np.count_nonzero(full_c[source_index, :]) == 1
            record.update({'kind': 'geometric-anchor', 'host': 'N4258', 'constraint': 'released maser distance-modulus offset'})
        elif source_index in [3209, 3212]:
            assert np.count_nonzero(full_c[source_index, :]) == 1
            record.update({'kind': 'released-summary-constraint', 'constraint': 'released metallicity summary' if source_index == 3209 else 'broad slope-correction regularizer'})
        else:
            record.update({'kind': 'hubble-flow-supernova', 'constraint': 'prepared magnitude-redshift combination; no unverified CID mapping'})
        observations.append(record)
    selected = {'source': {'repository': 'PantheonPlusSH0ES/DataRelease', 'commit': COMMIT, 'orientation': 'Author L is 47 by 3492; stored design is selected L transposed',
                'originalObservationCount': N, 'originalParameterCount': 47},
                'parameterNames': PARAMETER_NAMES, 'originalParameterIndices': PARAMETER_INDICES,
                'basePeriodLuminositySlope': BASE_SLOPE, 'anchorReferenceModulusMagnitude': 29.398,
                'anchorReferenceNote': 'Display reference from Riess et al. 2112.04510v3 section4.7; source prepared anchor vector is used unchanged',
                'observations': observations, 'designMatrix': design.tolist(), 'covariance': covariance.tolist(),
                'conditioning': 'Retains released gamma summary row3209 and broad slope row3212 with their supplied diagonal covariance; not an independently recalibrated new ladder',
                'selection': '24 evenly spaced matrix rows from M101, all15 M1337, 24 evenly spaced N4258; five host SN light-curve rows;16 evenly spaced HF rows; three released constraints',
                'omittedNonzeroParameterCount': 0}
    default_fit = gls(design, values, covariance)
    variants = {}
    for name, predicate in [('withoutMetallicitySummary', lambda row: row != 3209), ('withoutBroadSlopeRegularizer', lambda row: row != 3212)]:
        mask = [index for index, row in enumerate(indices) if predicate(row)]
        variants[name] = gls(design[mask], values[mask], covariance[np.ix_(mask, mask)])
    variants['diagonalCovarianceCounterfactual'] = gls(design, values, np.diag(np.diag(covariance)))
    for name, predicate in [('anchorPlus0p1', lambda row: row == 3213), ('commonCepheidPlus0p1', lambda row: row < 3130), ('hostCepheidPlus0p1', lambda row: row < 2150)]:
        shifted = values + np.array([0.1 if predicate(row) else 0 for row in indices])
        variants[name] = gls(design, shifted, covariance)
    no_anchor_mask = [index for index, row in enumerate(indices) if row != 3213]
    null_direction = np.array([1, 1, 1, -1, 0, -1, 0, -1])
    assert np.max(np.abs(design[no_anchor_mask] @ null_direction)) == 0
    no_anchor = {'rank': int(np.linalg.matrix_rank(design[no_anchor_mask])), 'parameterCount': 8,
                 'nullDirection': null_direction.tolist(), 'absoluteScaleIdentified': False,
                 'interpretation': 'All moduli shift together, luminosity zero points and fiveLogH0 shift oppositely; no finite H0 inference'}
    assert no_anchor['rank'] == 7
    benchmarks = {'numpyVersion': np.__version__, 'algorithm': 'Cholesky whitening followed by SVD least squares; parameter covariance from whitened normal matrix',
                  'default': default_fit, 'variants': variants, 'withoutAnchor': no_anchor}
    if args.full_reference:
        benchmarks['completeAuthorSystemReference'] = gls(stored_l.T, y, np.asarray(full_c, dtype=float))
        # The general helper's physicalSlope index applies to the subset only.
        benchmarks['completeAuthorSystemReference']['physicalSlope'] = float(BASE_SLOPE + benchmarks['completeAuthorSystemReference']['parameters'][41])
        benchmarks['completeAuthorSystemReference'].pop('hostDistanceCorrelation')
        benchmarks['completeAuthorSystemReference']['note'] = 'Offline full source check; not the 87-row browser inference or the additional systematic variant error'
    artifacts = [save_json(OUTPUT / 'shoes-subsystem.json', selected), save_json(OUTPUT / 'shoes-subsystem-benchmarks.json', benchmarks)]
    manifest = {'retrievedDate': '2026-10-28', 'sourceCommit': COMMIT, 'sources': metadata,
                'metadataTable': {'url': source_url('table2.tex'), 'bytes': table_path.stat().st_size, 'sha256': hashlib.sha256(table_path.read_bytes()).hexdigest(),
                'use': 'Only object IDs, coordinates, periods and display H-band magnitudes; never its diagonal errors for the fit'},
                'artifacts': artifacts, 'releasedConstraintCrossCovariance': {'sourceRows': [3209, 3212, 3213], 'otherSourceRowsChecked': 3491, 'nonzeroCrossEntries': 0, 'interpretation': 'Numerical zeros in the source covariance; summaries remain conditional on the released global analysis'}, 'publishedArticleLicense': 'CC-BY4.0 according to publisher-registered DOI metadata in published-license.json',
                'repositoryCodeLicense': 'No separate license found; source MCMC code was not copied. Numerical data subset credited to authors and official public release.'}
    (DOCUMENTS / 'provenance.json').write_text(json.dumps(manifest, indent=2) + '\n')
    print(json.dumps({'rows': len(indices), 'default': {key: value for key, value in default_fit.items() if key not in ['parameters', 'covariance']},
                      'withoutAnchor': no_anchor, 'artifacts': artifacts}, indent=2))


if __name__ == '__main__':
    main()