"""Compare real native/browser results to the frozen 72-curve Amesim reference.

Run with Python + NumPy. --plot additionally requires Matplotlib, available in
this workstation's system Python. Outputs stay in the requested artifact folder.
"""
from __future__ import annotations
import argparse
import csv
import hashlib
import json
from pathlib import Path
import numpy as np

ROOT = Path(__file__).resolve().parents[2]
LIMITS = {'pressure': 250., 'temperature': .015, 'displacement': 2e-6,
          'velocity': 1e-5, 'mass_flow': 3e-5}
UNITS = {'pressure': 'Pa', 'temperature': 'K', 'displacement': 'm',
         'velocity': 'm/s', 'mass_flow': 'kg/s'}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('result', type=Path)
    parser.add_argument('--reference', type=Path, default=ROOT/'tests/baselines/simulation/test_mql_4/test-mql-4-amesim-reference.json')
    parser.add_argument('--output-dir', type=Path, required=True)
    parser.add_argument('--plot', action='store_true')
    args = parser.parse_args()
    ref = json.loads(args.reference.read_text())
    data = json.loads(args.result.read_text())
    if 'snapshot' in data:
        data = data['snapshot']['result']
    series = data['series']; times = np.array(series['time']); ref_times = np.array(ref['times'])
    if not data['success'] or times[-1] < ref_times[-1] or times[0] > ref_times[0]:
        raise ValueError('A completed result covering all reference times is required.')
    if np.any(np.diff(times) < 0):
        raise ValueError('Native sample times are not ordered.')
    finite = all(np.all(np.isfinite(values)) for values in series.values())
    rows = []; grouped = {}
    for key, curve in ref['series'].items():
        actual = np.interp(ref_times, times, series[key]); error = actual - curve['values']
        worst = int(np.argmax(abs(error))); quantity = curve['quantity']
        row = {'key': key, 'quantity': quantity, 'unit': UNITS[quantity],
               'maxAbsoluteError': float(abs(error[worst])), 'rmsError': float(np.sqrt(np.mean(error**2))),
               'worstTime': float(ref_times[worst]), 'limit': LIMITS[quantity],
               'passed': bool(np.max(abs(error)) <= LIMITS[quantity])}
        rows.append(row)
        if quantity not in grouped or row['maxAbsoluteError'] > grouped[quantity]['maxAbsoluteError']:
            grouped[quantity] = row
    masses = np.array([v for k, v in series.items() if k.rsplit('.', 1)[-1] in ('m', 'm1', 'm2')])
    total = np.sum(masses, axis=0)
    mass_drift = float(np.max(abs(total-total[0])))
    force_peaks = sorted(({'key': k, 'peakAbsolute': max(map(abs, v))} for k, v in series.items()
                          if k.startswith('amesim_lstp00a_') and k.endswith('.f')),
                         key=lambda row: row['peakAbsolute'], reverse=True)
    summary = {'resultPath': str(args.result.resolve()),
               'resultSha256': hashlib.sha256(args.result.read_bytes()).hexdigest(),
               'referencePath': str(args.reference.resolve()),
               'referenceSha256': hashlib.sha256(args.reference.read_bytes()).hexdigest(),
               'referenceSource': ref['source'], 'freshAmesimRun': False,
               'referenceCurveCount': len(rows), 'referenceTimeCount': len(ref_times),
               'nativeSampleCount': len(times), 'allSamplesFinite': bool(finite),
               'massStateCount': len(masses), 'maxTotalMassDriftKg': mass_drift,
               'groupedWorstErrors': grouped, 'curves': rows, 'contactForcePeaksOutsideReferenceScope': force_peaks,
               'passed': bool(finite and mass_drift <= 1e-10 and all(row['passed'] for row in rows))}
    args.output_dir.mkdir(parents=True, exist_ok=True)
    (args.output_dir/'comparison.json').write_text(json.dumps(summary, ensure_ascii=False, indent=2)+'\n')
    with (args.output_dir/'curve-errors.csv').open('w', newline='') as output:
        writer = csv.DictWriter(output, fieldnames=list(rows[0])); writer.writeheader(); writer.writerows(rows)
    if args.plot:
        import matplotlib
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt
        fig, axes = plt.subplots(5, 2, figsize=(13, 15), constrained_layout=True)
        for (quantity, row), (left, right) in zip(grouped.items(), axes):
            key = row['key']; values = ref['series'][key]['values']
            actual = np.interp(ref_times, times, series[key])
            left.plot(times, series[key], linewidth=1, label='Native, rtol=1e-8')
            left.plot(ref_times, values, '.', markersize=3, label='Saved Amesim reference')
            left.set_title(key, fontsize=9); left.set_ylabel(UNITS[quantity]); left.legend(fontsize=7)
            right.plot(ref_times, actual-values, '.-', linewidth=.8, markersize=3)
            right.axhline(row['limit'], color='gray', linestyle='--', linewidth=.6)
            right.axhline(-row['limit'], color='gray', linestyle='--', linewidth=.6)
            right.set_title(f'{quantity}: sampled difference and acceptance limits', fontsize=9)
            right.set_ylabel(UNITS[quantity])
            for ax in (left, right): ax.grid(alpha=.2); ax.set_xlabel('Time [s]')
        fig.suptitle('Four-branch model: 72 curves checked at 57 stored Amesim times\nWorst curve per quantity; contact force is outside this reference', fontsize=12)
        fig.savefig(args.output_dir/'comparison.svg'); fig.savefig(args.output_dir/'comparison.png', dpi=140)
    print(json.dumps({k: summary[k] for k in ('passed', 'nativeSampleCount', 'maxTotalMassDriftKg', 'groupedWorstErrors')}, indent=2))
    if not summary['passed']: raise SystemExit(1)


if __name__ == '__main__': main()