"""Sequential reservoir experiment: stop after a failed positive case."""
import argparse
import json
import os
from pathlib import Path
import sys
from unittest.mock import patch

import numpy as np
import CelerisLab
from CelerisLab.cuda import compiler_v2 as compiler
from run_j1 import ROOT, make_sim, off_path_check
from run_j2_slot import NX, WIDTH, GATES, run_case, restart_check


def old_velocity_check(baseline, width, top_u):
    states = []
    for source in (baseline, Path(compiler.kernel_path(''))):
        with patch.object(compiler, '_KERNEL_DIR', str(source)):
            with make_sim(.002, nx=NX, slot_width=width, top_u=top_u) as sim:
                sim.initialize()
                initial = sim.get_ddf().copy()
                sim.run(64)
                states.append((initial, sim.get_ddf().copy(), sim.get_flags().copy()))
    return bool(all(np.array_equal(a, b) for a, b in zip(*states)))


def checkpoint_check():
    with make_sim(.002, nx=NX, top_u=0., slot_width=WIDTH, top_mode='reservoir') as sim:
        sim.initialize()
        sim.run(11)
        return restart_check(sim, .002, top_mode='reservoir')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--baseline-kernels', type=Path, required=True)
    args = parser.parse_args()
    assert Path.cwd() == ROOT == Path('/home/frank14f/CelerisLab')
    assert sys.executable == '/home/frank14f/anaconda3/envs/pycuda_3_10/bin/python'
    assert Path(CelerisLab.__file__).resolve() == ROOT/'src/CelerisLab/__init__.py'
    assert Path(compiler.kernel_path('kernel_v2.cu')).resolve() == ROOT/'src/CelerisLab/lbm/kernels/kernel_v2.cu'
    assert os.environ.get('CUDA_VISIBLE_DEVICES') in ('0', '1')
    report = dict(python=sys.executable, package=CelerisLab.__file__,
                  cuda_source=compiler.kernel_path('kernel_v2.cu'), physical_gpu=os.environ['CUDA_VISIBLE_DEVICES'],
                  gates=dict(GATES, top_density_error=5e-7), regressions={}, cases=[],
                  status=dict(positive='not_tested', negative='not_tested', zero='not_tested',
                              checkpoint='not_tested'))

    def save():
        args.output.write_text(json.dumps(report, indent=2)+'\n')

    for name, result in (
        ('j1_velocity_top', old_velocity_check(args.baseline_kernels, 0, .02)),
        ('j2_velocity_top', old_velocity_check(args.baseline_kernels, WIDTH, 0.)),
        ('default_off', off_path_check(args.baseline_kernels)),
    ):
        report['regressions'][name] = 'passed' if result else 'failed'
    save()
    print(json.dumps(report['regressions']), flush=True)
    if 'failed' in report['regressions'].values():
        return 1
    for name, amplitude in (('positive', .002), ('negative', -.002), ('zero', 0.)):
        case = run_case(amplitude, top_mode='reservoir', check_restart=False)
        report['cases'].append(case)
        report['status'][name] = 'passed' if case['passed'] else 'failed'
        save()
        print(json.dumps({k: case[k] for k in ('amplitude', 'steps', 'checks', 'mass_offset',
                                             'mass_window_span', 'flux_error', 'top_density_error',
                                             'top_velocity_range', 'symmetry')}), flush=True)
        if not case['passed']:
            return 1
    report['status']['checkpoint'] = 'passed' if checkpoint_check() else 'failed'
    save()
    print(json.dumps(report['status']), flush=True)
    return 0 if all(status == 'passed' for status in report['status'].values()) else 1


if __name__ == '__main__':
    raise SystemExit(main())