#!/usr/bin/env python3
"""Validate serialized results, exact trajectories, controls, and data schema."""
import argparse
from dataclasses import fields
from fractions import Fraction
from pathlib import Path
import json
import time
import numpy as np
from cell_auto.core import Model
from cell_auto.experiments import write_json
from cell_auto.tests import run_tests


def require(condition, message):
    if not bool(condition):
        raise AssertionError(message)


def verify(root, sample=False, output=None):
    root = Path(root)
    started = time.perf_counter()
    counts = dict(npz_files=0, arrays=0, integer_trajectories=0, exact_transitions=0)
    files = sorted(root.rglob('*.npz'))
    require(bool(files), 'No results found')
    allowed = {f.name for f in fields(Model)}
    max_register = 0
    min_physical_margin = float('inf')
    models = []
    for path in files:
        with np.load(path, allow_pickle=False) as z:
            for name in z.files:
                array = z[name]
                require(array.dtype.kind in 'biufc', (path, name, 'nonnumeric array'))
                require(np.all(np.isfinite(array)), (path, name, 'nonfinite'))
                counts['arrays'] += 1
            counts['npz_files'] += 1
            if 'Q' not in z.files:
                continue
            metadata = json.loads(path.with_suffix('.json').read_text())
            params = {k: v for k, v in metadata['model'].items() if k in allowed}
            params['delta_override'] = Fraction(metadata['model']['delta'])
            model = Model(**params)
            models.append(model.to_dict())
            q, p = z['Q'], z['P']
            counts['integer_trajectories'] += 1
            require(q.dtype == np.int64 and p.dtype == np.int64, (path, 'register dtype'))
            require(q.shape == p.shape and q.shape[1:] == (model.species, model.sites), (path, 'register shape'))
            require(np.array_equal(z['fields'], float(model.delta)*q.astype(float)), (path, 'physical fields'))
            require(np.array_equal(z['times'], np.arange(len(q))*float(model.tau)), (path, 'time grid'))
            # Convert extrema before abs, so a corrupted int64 minimum cannot overflow.
            absolute = max(abs(int(a.min())) for a in (q, p))
            absolute = max(absolute, *(abs(int(a.max())) for a in (q, p)))
            require(2*absolute < model.modulus, (path, 'no-wrap bound'))
            diagnostics = metadata['diagnostics']
            require(absolute == diagnostics['max_abs_register'], (path, 'register maximum'))
            require(diagnostics['nowrap_pass'] and diagnostics['paired_consistent'], (path, 'run diagnostics'))
            require(model.modulus//2-absolute == diagnostics['integer_margin'], (path, 'integer margin'))
            max_register = max(max_register, absolute)
            min_physical_margin = min(min_physical_margin, float((model.modulus//2-absolute)*model.delta))
            indices = range(len(q)-1)
            if sample:
                indices = sorted(set(np.linspace(0, len(q)-2, min(17, len(q)-1), dtype=int).tolist()))
            for n in indices:
                qn, pn = model.step(q[n], p[n], modular=False)
                require(np.array_equal(qn, q[n+1]) and np.array_equal(pn, p[n+1]), (path, n, 'forward'))
                qp, pp = model.inverse(q[n+1], p[n+1], modular=False)
                require(np.array_equal(qp, q[n]) and np.array_equal(pp, p[n]), (path, n, 'inverse'))
                counts['exact_transitions'] += 1
            if 'energy_times' in z.files:
                require(np.array_equal(z['energy_times'], z['times'][1:-1]), (path, 'energy time grid'))
                require(len(z['diagnostic_energy']) == len(z['energy_times']), (path, 'energy shape'))
    collision = json.loads((root / 'collision/summary.json').read_text())
    for row in collision['runs']:
        if row['epsilon'] == 0:
            with np.load(root / 'collision' / (row['label']+'.npz'), allow_pickle=False) as z:
                require(np.array_equal(z['Q'][:, 0], z['Q'][:, 1])
                        and np.array_equal(z['P'][:, 0], z['P'][:, 1]), (row['label'], 'symmetric control'))
    principal = next(row for row in collision['runs'] if row['label'] == collision['principal_label'])
    reflected = next(row for row in collision['runs'] if row['sites'] == principal['sites']
                     and row['epsilon'] == -principal['epsilon'] and row['content'] == 'both' and not row['linear'])
    with np.load(root / 'collision' / (principal['label']+'.npz'), allow_pickle=False) as positive, \
            np.load(root / 'collision' / (reflected['label']+'.npz'), allow_pickle=False) as negative:
        require(all(np.array_equal(positive[key][:, ::-1], negative[key]) for key in ('Q', 'P')),
                'Epsilon reflection does not interchange all species registers')
    require(all(m['linear_cfl_strict_margin'] > 0 for m in models), 'Nonpositive linear CFL margin')
    result = dict(status='passed', mode='sampled transitions' if sample else 'all saved integer transitions',
        counts=counts, maximum_saved_register=max_register, minimum_physical_register_margin=min_physical_margin,
        no_pickle_required=True, all_reference_arrays_finite=True, all_linear_cfl_margins_positive=True,
        epsilon_reflection_registers_exact=True, tests=run_tests(), runtime_seconds=time.perf_counter()-started)
    write_json(Path(output) if output else root / 'verification.json', result)
    return result


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--sample', action='store_true')
    parser.add_argument('--results', default='results')
    parser.add_argument('--output', help='Report path; defaults to verification.json in the results directory')
    args = parser.parse_args()
    here = Path(__file__).resolve().parent
    path = Path(args.results)
    if not path.is_absolute():
        path = here / path
    output = Path(args.output) if args.output else None
    if output is not None and not output.is_absolute():
        output = here / output
    print(json.dumps(verify(path, args.sample, output), indent=2))