"""Reproduce ProteinIQ's genome sequencing cost analyses (September 24, 2026).

1. Halving time of NHGRI's sequencing cost by era, and inflation adjustment.
2. Survey of US university core facility prices for a 30x human genome.

Run from this directory: python analyze.py   (needs pandas and xlrd)
Writes summary.json, nhgri-halving-times.csv and core-facility-results.csv.
"""
import csv
import glob
import json
import statistics as st

import numpy as np
import pandas as pd

GENOME_GB = 90  # NHGRI's 30x genome: 3,000 Mb x 30 (cost per genome / cost per Mb = 90,000 Mb after 2008)
GB_PER_MILLION_PAIRS = 0.3  # 2 x 150 bp
ERAS = [
    ('Sanger sequencing', '2001-09-01', '2007-12-31'),
    ('Early next-generation sequencing', '2008-01-01', '2011-12-31'),
    ('Maturing short-read sequencing', '2012-01-01', '2015-12-31'),
    ('Recent short-read sequencing', '2016-01-01', '2022-05-31'),
]


def r(x, n=2):
    return None if x is None else round(float(x), n)


# ---------- 1. NHGRI series ----------
nh = pd.read_excel('raw/Sequencing_Cost_Data_Table_May2022.xls')
nh.columns = ['date', 'cost_per_mb', 'cost_per_genome']
nh['date'] = pd.to_datetime(nh['date'])
nh['implied_genome_mb'] = nh['cost_per_genome'] / nh['cost_per_mb']

cpi_rows = []
for f in sorted(glob.glob('raw/bls-CUUR0000SA0-*.json')):
    for d in json.load(open(f))['Results']['series'][0]['data']:
        if d['period'].startswith('M') and d['period'] != 'M13' and d['value'] != '-':  # BLS marks unpublished months with '-'
            cpi_rows.append((pd.Timestamp(int(d['year']), int(d['period'][1:]), 1), float(d['value'])))
cpi = pd.Series(dict(cpi_rows)).sort_index()
cpi_latest_month = cpi.index.max()
cpi_latest = cpi.iloc[-1]


def cpi_at(ts):
    return cpi[pd.Timestamp(ts.year, ts.month, 1)]


nh['cost_per_genome_aug2026_usd'] = [c * cpi_latest / cpi_at(d) for c, d in zip(nh['cost_per_genome'], nh['date'])]
t0 = nh['date'].min()
nh['months'] = (nh['date'] - t0).dt.days / (365.2425 / 12)

halving = []
for name, a, b in ERAS:
    s = nh[(nh['date'] >= a) & (nh['date'] <= b)]
    slope = np.polyfit(s['months'], np.log2(s['cost_per_mb']), 1)[0]  # log2 units per month
    halving.append({
        'era': name,
        'start': s['date'].min().date().isoformat(),
        'end': s['date'].max().date().isoformat(),
        'data_points': len(s),
        'halving_time_months': r(-1 / slope, 1),
        'fold_decline_per_year': r(2 ** (-12 * slope), 2),
        'start_cost_per_genome_usd': r(s['cost_per_genome'].iloc[0]),
        'end_cost_per_genome_usd': r(s['cost_per_genome'].iloc[-1]),
        'end_cost_per_30x_genome_usd': r(s['cost_per_mb'].iloc[-1] * GENOME_GB * 1000),
    })
with open('nhgri-halving-times.csv', 'w', newline='') as fh:
    w = csv.DictWriter(fh, fieldnames=list(halving[0]))
    w.writeheader()
    w.writerows(halving)

first, last = nh.iloc[0], nh.iloc[-1]
nhgri = {
    'source_file': 'raw/Sequencing_Cost_Data_Table_May2022.xls',
    'genome_definition_mb': {
        'sanger_era': r(nh.loc[nh['date'] < '2008-01-01', 'implied_genome_mb'].median(), 0),
        'ngs_era': r(nh.loc[nh['date'] >= '2008-01-01', 'implied_genome_mb'].median(), 0),
    },
    'first': {'date': first['date'].date().isoformat(), 'cost_per_genome_usd': r(first['cost_per_genome']),
              'cost_per_mb_usd': r(first['cost_per_mb'], 4)},
    'last': {'date': last['date'].date().isoformat(), 'cost_per_genome_usd': r(last['cost_per_genome']),
             'cost_per_mb_usd': r(last['cost_per_mb'], 6)},
    'minimum': {'date': nh.loc[nh['cost_per_genome'].idxmin(), 'date'].date().isoformat(),
                'cost_per_genome_usd': r(nh['cost_per_genome'].min())},
    'fold_decline_per_genome_nominal': r(first['cost_per_genome'] / last['cost_per_genome'], 0),
    'fold_decline_per_mb_nominal': r(first['cost_per_mb'] / last['cost_per_mb'], 0),
    'cpi_series': 'BLS CUUR0000SA0 (CPI-U, US city average, not seasonally adjusted)',
    'cpi_reference_month': cpi_latest_month.strftime('%Y-%m'),
    'first_cost_per_genome_in_reference_dollars': r(first['cost_per_genome_aug2026_usd'], 0),
    'last_cost_per_genome_in_reference_dollars': r(last['cost_per_genome_aug2026_usd'], 0),
    'fold_decline_per_genome_real': r(first['cost_per_genome_aug2026_usd'] / last['cost_per_genome_aug2026_usd'], 0),
    'moores_law_halving_months': 24,
    'halving_times': halving,
}

# ---------- 2. Core facility survey ----------
cores = list(csv.DictReader(open('core-facility-prices.csv', encoding='utf8')))
results = []
for c in cores:
    row = {'institution': c['institution'], 'sequencing_location': c['sequencing_location'], 'seq_option': c['seq_option']}
    if c['seq_price_internal_usd']:
        gb = float(c['seq_units_read_pairs_millions']) * GB_PER_MILLION_PAIRS
        for tier, col in (('internal', 'seq_price_internal_usd'), ('external', 'seq_price_external_academic_usd')):
            if c[col]:
                per_gb = float(c[col]) / gb
                row[f'usd_per_gb_{tier}'] = r(per_gb, 3)
                row[f'sequencing_per_genome_{tier}'] = r(per_gb * GENOME_GB)
    for tier, col in (('internal', 'library_internal_usd'), ('external', 'library_external_academic_usd')):
        if c[col]:
            row[f'library_{tier}'] = float(c[col])
    for tier in ('internal', 'external'):
        s, lib = row.get(f'sequencing_per_genome_{tier}'), row.get(f'library_{tier}')
        if s is not None and lib is not None:
            row[f'total_per_genome_{tier}'] = r(s + lib)
            row[f'library_share_{tier}'] = r(lib / (s + lib), 3)
    if c['seq_price_internal_usd'] and c['seq_price_external_academic_usd']:
        row['external_to_internal_sequencing_ratio'] = r(float(c['seq_price_external_academic_usd']) / float(c['seq_price_internal_usd']), 3)
    results.append(row)

fields = ['institution', 'sequencing_location', 'seq_option', 'usd_per_gb_internal', 'sequencing_per_genome_internal',
          'library_internal', 'total_per_genome_internal', 'library_share_internal', 'usd_per_gb_external',
          'sequencing_per_genome_external', 'library_external', 'total_per_genome_external', 'library_share_external',
          'external_to_internal_sequencing_ratio']
with open('core-facility-results.csv', 'w', newline='', encoding='utf8') as fh:
    w = csv.DictWriter(fh, fieldnames=fields)
    w.writeheader()
    w.writerows(results)


def describe(key, rows=results):
    v = sorted(x[key] for x in rows if x.get(key) is not None)
    if not v:
        return None
    q = np.percentile(v, [25, 75])
    return {'n': len(v), 'median': r(st.median(v), 3), 'q1': r(q[0], 3), 'q3': r(q[1], 3), 'min': r(v[0], 3), 'max': r(v[-1], 3)}


# Lane-level options that one laboratory can buy; the per-read shared run is summarized separately.
lane_rows = [x for x in results if 'shared run per million' not in x['seq_option']]
survey = {
    'retrieved': '2026-09-24',
    'genome_gb': GENOME_GB,
    'facilities': len(cores),
    'facilities_with_sequencing_price': sum(1 for x in results if x.get('sequencing_per_genome_internal') is not None),
    'sequencing_per_genome_internal_lane_level': describe('sequencing_per_genome_internal', lane_rows),
    'sequencing_per_genome_internal_all': describe('sequencing_per_genome_internal'),
    'sequencing_per_genome_external': describe('sequencing_per_genome_external'),
    'library_internal': describe('library_internal'),
    'library_external': describe('library_external'),
    'total_per_genome_internal_lane_level': describe('total_per_genome_internal', lane_rows),
    'library_share_internal_lane_level': describe('library_share_internal', lane_rows),
    'sequencing_per_gb_internal_lane_level': describe('usd_per_gb_internal', lane_rows),
    'external_to_internal_sequencing_ratio': describe('external_to_internal_sequencing_ratio'),
    'shared_run_per_read_example': next(x for x in results if x not in lane_rows),
    'sensitivity_sequencing_median_at_100_gb': r(describe('sequencing_per_genome_internal', lane_rows)['median'] * 100 / GENOME_GB),
}

json.dump({'nhgri': nhgri, 'core_facility_survey': survey}, open('summary.json', 'w'), indent=2)
print(json.dumps({'nhgri': nhgri, 'core_facility_survey': survey}, indent=2))
for x in results:
    print(x)