"""Audit the controlled output-only eight-branch experiment without hiding raw errors."""
import argparse
import json
from pathlib import Path

import numpy as np


def load(path):
    return json.loads(path.read_bytes())


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('directory', type=Path)
    args = parser.parse_args()
    out, result = args.directory, {}
    for name in ('full', 'noncyclic'):
        before, after = (out / stage / name for stage in ('before', 'after'))
        a, b = (load(p/'native/result.json') for p in (before, after))
        at, bt = (np.asarray(r['series']['time']) for r in (a, b))
        indices = np.searchsorted(bt, at)
        assert np.array_equal(bt[indices], at), 'Every old sample must remain at its real timestamp'
        differences = []
        for key, values in a['series'].items():
            old, new = np.asarray(values), np.asarray(b['series'][key])[indices]
            if not np.array_equal(old, new):
                differences.append(dict(key=key, maxAbsoluteDifference=float(np.max(np.abs(new-old)))))
        counters = ('acceptedSteps', 'rejectedSteps', 'nfev', 'njev', 'nlu', 'solverStarts', 'stateTransitions')
        counter_pairs = {key: [a[key], b[key]] for key in counters}
        paired = np.load(after/'curves.npz')
        contact_phase = {}
        for key in paired.files:
            if key.startswith('platform|') and key.endswith('.gap'):
                other = 'amesim|' + key.split('|', 1)[1]
                mismatch = paired['phaseMatched'] & ((paired[key]<0) != (paired[other]<0))
                contact_phase[key.split('|', 1)[1]] = dict(mismatchCount=int(mismatch.sum()),
                    times=paired['time'][mismatch].tolist())
        old_metrics = load(before/'raw-time-comparison.json')
        new_raw = load(after/'raw-time-comparison.json')
        aligned = load(after/'comparison.json')
        quantities = {}
        for quantity in ('signal', 'force', 'pressure', 'temperature', 'mass_flow', 'enthalpy_flow', 'velocity'):
            groups = []
            for rows in (old_metrics['curves'], new_raw['curves'], aligned['curves']):
                rows = [r for r in rows if r['quantity'] == quantity]
                worst = max(rows, key=lambda r: r['maxAbsoluteError'])
                groups.append(dict(maxAbsoluteError=worst['maxAbsoluteError'], key=worst['key'],
                    time=worst['worstTime'], above5PercentCount=sum(r['above5PercentCount'] for r in rows)))
            quantities[quantity] = dict(beforeRaw=groups[0], afterRaw=groups[1], afterSamePhase=groups[2])
        result[name] = dict(finalStateExactlyEqual=a['finalState']==b['finalState'],
            solverCountersEqual=all(x==y for x,y in counter_pairs.values()), counters=counter_pairs,
            originalSampleCount=len(at), newSampleCount=len(bt), extraSamples=len(bt)-len(at),
            allOriginalSamplesPreserved=True, changedOriginalOutputs=differences,
            solveSeconds=[a['solveSeconds'], b['solveSeconds']],
            processWallSeconds=[load(before/'native-summary.json')['processWallSeconds'],
                                load(after/'native-summary.json')['processWallSeconds']],
            phasePairing=load(after/'phase-pairing.json'), contactPhase=contact_phase,
            quantities=quantities,
            above5PercentCount=[old_metrics['above5PercentCount'], new_raw['above5PercentCount'], aligned['above5PercentCount']],
            above5PercentOutsideEvents=aligned['above5PercentOutsideEvents'],
            rawPeakContactForce=[load(before/'comparison.json')['extraEventPointMaxContactForce'], aligned['extraEventPointMaxContactForce']],
            sourceVerification={s: load(out/s/'source-verification.json') for s in ('before','after')})
    (out/'effect-summary.json').write_text(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False)+'\n', encoding='utf-8')
    for name, row in result.items():
        print(name, json.dumps({k:v for k,v in row.items() if k not in ('phasePairing','quantities','sourceVerification')}, ensure_ascii=False))
        print(json.dumps(row['quantities'], ensure_ascii=False))


if __name__ == '__main__':
    main()