#!/usr/bin/env python3
import argparse
import json
import os
from pathlib import Path
import shutil
import subprocess
import tempfile

import numpy as np

root = Path(__file__).resolve().parents[1]
parser = argparse.ArgumentParser(description='Compare upstream and current Legendre kernels at the selected precision.')
parser.add_argument('--resolutions', nargs='+', type=int, choices=[21, 42, 85], default=[21, 42, 85])
parser.add_argument('--repeats', type=int, default=20)
parser.add_argument('--timeout', type=int, default=180)
parser.add_argument('--work', type=Path)
parser.add_argument('--reference', default='exoplasim-upstream-import')
parser.add_argument('--precision', type=int, choices=[4, 8], default=8)
args = parser.parse_args()
if args.repeats < 1:
    parser.error('--repeats must be positive')
(root / 'tmp').mkdir(exist_ok=True)
work = args.work.resolve() if args.work else Path(tempfile.mkdtemp(prefix='transform-benchmark-', dir=root / 'tmp'))
work.mkdir(parents=True, exist_ok=True)
baseline = subprocess.check_output(['git', 'show', args.reference + ':exoplasim/plasim/src/legmod.f90'], cwd=root, timeout=20)
candidate = (root / 'exoplasim/plasim/src/legmod.f90').read_bytes()
stub, program = (root / 'scripts/transform-benchmark.f90').read_text().split('program transform_benchmark', 1)
results = []
for truncation in args.resolutions:
    for partition in [1, 2]:
        for mode, flags in [('checked', ['-fcheck=all', '-DEXOPLASIM_CHECKED']), ('native', ['-march=native', '-ffp-contract=off'])]:
            outputs = {}
            timings = {}
            for label, source in [('baseline', baseline), ('candidate', candidate)]:
                directory = work / f't{truncation}-p{partition}-{mode}-{label}'
                directory.mkdir(exist_ok=True)
                (directory / 'pumamod.f90').write_text(stub)
                (directory / 'driver.f90').write_text('program transform_benchmark' + program)
                (directory / 'legmod.f90').write_bytes(source)
                command = [os.environ.get('FC', 'gfortran'), '-O3', '-cpp', f'-DTRUNC={truncation}',
                           f'-DPARTITION={partition}', *(['-fdefault-real-8'] if args.precision == 8 else []), '-ffree-line-length-none',
                           '-ffpe-trap=invalid,zero,overflow', *flags, 'pumamod.f90', 'legmod.f90',
                           'driver.f90', '-o', 'benchmark']
                with (directory / 'build.log').open('w') as log:
                    subprocess.run(command, cwd=directory, check=True, timeout=args.timeout,
                                   stdout=log, stderr=subprocess.STDOUT)
                output = subprocess.check_output(['./benchmark', str(args.repeats)], cwd=directory, timeout=args.timeout)
                (directory / 'timings.txt').write_bytes(output)
                timings[label] = list(map(float, output.decode().split()))
                outputs[label] = np.fromfile(directory / 'output.bin', dtype=np.float64 if args.precision == 8 else np.float32)
            reference, current = outputs['baseline'], outputs['candidate']
            if current.shape != reference.shape or not np.isfinite(current).all() or not np.isfinite(reference).all():
                raise AssertionError('Invalid kernel output')
            delta = float(np.max(np.abs(current - reference)))
            scaled = delta / max(1.0, float(np.max(np.abs(reference))))
            record = dict(truncation=truncation, partition=partition, mode=mode, precision=args.precision,
                          max_absolute_error=delta, scaled_error=scaled,
                          bitwise_identical=reference.tobytes() == current.tobytes(),
                          baseline_seconds=timings['baseline'], candidate_seconds=timings['candidate'],
                          speedup=[a / b for a, b in zip(timings['baseline'], timings['candidate'])])
            results.append(record)
            (work / 'results.json').write_text(json.dumps(results, indent=2) + '\n')
            print(json.dumps(record), flush=True)
            if scaled > 1e-12:
                raise AssertionError(f'Kernel numerical mismatch: {record}')
print('Results:', work / 'results.json')