#!/usr/bin/env python3
"""Validate and reproduce the Atlas with Python's standard library.
Place this script and atlas.json in one directory; run: python reproduce.py
Writes atlas.sqlite and calculations-reproduced.csv beside the input.
"""
import csv
import json
import pathlib
import sqlite3
from decimal import Decimal


def validate(data):
    for collection in ('assets', 'entities', 'relationships', 'metrics', 'events', 'sources'):
        ids = [row['id'] for row in data[collection]]
        assert len(ids) == len(set(ids)), 'Duplicate IDs in ' + collection
    assets = {a['id'] for a in data['assets']}
    entities = {e['id'] for e in data['entities']}
    sources = {s['id'] for s in data['sources']}
    equity = {}
    for row in data['relationships']:
        assert row['asset_id'] in assets
        assert row['from_entity_id'] in entities
        assert row['to_id'] in (assets if row['to_type'] == 'asset' else entities)
        assert row['source_id'] in sources
        assert row['status'] in ('active', 'historical', 'terminated')
        assert row['snapshot_date']
        if row['equity_percent'] is not None:
            assert row['relationship_type'] == 'equity_interest'
            key = (row['to_id'], row['snapshot_date'])
            equity[key] = equity.get(key, Decimal('0')) + Decimal(str(row['equity_percent']))
    assert all(abs(v - 100) < Decimal('.001') for v in equity.values()), equity
    for row in data['metrics']:
        assert row['entity_id'] in entities
        assert row['source_id'] in sources
        assert row['value'] is not None and row['unit'] == 'USD'
        assert row['period_end'] and row['basis'] and row['locator']
    for row in data['assets']:
        assert row['geography']['type'] == 'regional_locator'
        assert row['geography']['route_geometry'] is None
        assert row['concession_end_precision'] == ('year' if row['concession_end_year'] else None)
    for row in data['events']:
        assert row['asset_id'] in assets and row['source_id'] in sources
    assert not any(m['metric'] == 'net_operating_cash' for m in data['metrics'])
    return equity


def reproduce(data, directory):
    results = []
    for row in data['same_basis']:
        assert row['revenue'] > 0
        assert row['remainder'] == row['revenue'] - row['operations'] - row['maintenance'] - row['debt_service']
        output = {'asset_id': row['asset_id'], 'period_end': row['period_end']}
        for key in ('operations', 'maintenance', 'debt_service', 'remainder'):
            output[key + '_per_100'] = str(Decimal(row[key]) / Decimal(row['revenue']) * 100)
        assert abs(sum(Decimal(output[k + '_per_100']) for k in ('operations', 'maintenance', 'debt_service', 'remainder')) - 100) < Decimal('1e-20')
        results.append(output)
    with (directory / 'calculations-reproduced.csv').open('w', newline='') as stream:
        writer = csv.DictWriter(stream, fieldnames=list(results[0]))
        writer.writeheader()
        writer.writerows(results)
    database = sqlite3.connect(directory / 'atlas.sqlite')
    database.execute('PRAGMA foreign_keys = ON')
    # A disposable, reproducible local output; publication never writes remote data.
    for table in ('relationships', 'metrics', 'events', 'assets', 'entities', 'sources'):
        database.execute(f'DROP TABLE IF EXISTS "{table}"')
    for table in ('assets', 'entities', 'sources', 'relationships', 'metrics', 'events'):
        rows = data[table]
        keys = list(dict.fromkeys(key for row in rows for key in row))
        columns = ', '.join('"' + key + '" ' + ('REAL' if key in ('value', 'equity_percent') else 'TEXT') + (' PRIMARY KEY' if key == 'id' else '') for key in keys)
        database.execute(f'CREATE TABLE "{table}" ({columns})')
        names = ','.join('"' + key + '"' for key in keys)
        sql = f'INSERT INTO "{table}" ({names}) VALUES ({",".join("?" for _ in keys)})'
        database.executemany(sql, [[json.dumps(row.get(key), ensure_ascii=False) if isinstance(row.get(key), (list, dict)) else row.get(key) for key in keys] for row in rows])
    database.commit()
    count = database.execute('SELECT count(*) FROM relationships').fetchone()[0]
    database.close()
    return count


if __name__ == '__main__':
    directory = pathlib.Path(__file__).resolve().parent
    data = json.loads((directory / 'atlas.json').read_text())
    equity = validate(data)
    count = reproduce(data, directory)
    print(f'Validated {len(data["assets"])} assets, {count} relationships, {len(data["metrics"])} observations and {len(equity)} complete equity snapshots.')
    print('Wrote atlas.sqlite and calculations-reproduced.csv. No statewide financial total calculated.')