#!/usr/bin/env python3
"""Check force evaluation, saved trajectories, normalization, grids, and frequency fits.

This script does not rerun the experiment suite.  It performs the five checks
already used during the mathematical/code review: an independent Fraction
force oracle, exact replay of selected saved trajectories, Parseval power
normalization, saved-grid time alignment, and frequency-fit sensitivity to
three initial search intervals.  Paths are resolved relative to this file,
so no files outside the delivered project are required.

Usage, from any working directory:
    python /path/to/project/experiment/verify_results.py
"""
from __future__ import annotations

import argparse
from fractions import Fraction
import hashlib
import json
import os
from math import ceil, floor
from pathlib import Path
import platform

import numpy as np
import scipy

from cell_auto.core import Model
from cell_auto.experiments import mode_frequency, temporal_frequency


HERE = Path(__file__).resolve().parent
PROJECT_ROOT = HERE.parent


def require(condition, message):
    """Keep checks active even when Python is invoked with optimization."""
    if not bool(condition):
        raise AssertionError(message)


def relative_name(path):
    return Path(os.path.relpath(Path(path).resolve(), PROJECT_ROOT)).as_posix()


def sha256(path):
    digest = hashlib.sha256()
    with Path(path).open('rb') as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b''):
            digest.update(block)
    return digest.hexdigest()


def fraction_force_oracle():
    # Use floor/ceil on Fraction values instead of the production code's
    # shared-denominator arithmetic and vectorized integer-rounding helper.
    def nearest_odd(value):
        if value >= 0:
            return floor(value + Fraction(1, 2))
        return ceil(value - Fraction(1, 2))

    model = Model(
        7, Fraction(8), courant=Fraction(2, 5), species=2,
        mass_phi_sq=Fraction(2, 3), mass_chi_sq=Fraction(7, 11),
        eta=Fraction(3, 7), lambda_phi=Fraction(5, 13),
        lambda_chi=Fraction(2, 9), g=Fraction(-1, 17),
    )
    generator = np.random.default_rng(191)
    fields = generator.integers(-10**6, 10**6, (2, 7)).astype(object)
    fields[0, 1] = 10**30 + 1
    fields[1, 4] = -10**31 + 3
    actual = model.force(fields)
    largest_digits = 0
    for component in range(2):
        mass = (model.mass_phi_sq, model.mass_chi_sq)[component]
        coupling = (model.lambda_phi, model.lambda_chi)[component]
        for site in range(7):
            value = model.delta * int(fields[component, site])
            other = model.delta * int(fields[1-component, site])
            onsite = (mass * value + model.eta * (value-other)
                      + coupling * value**3 + model.g * value * other**2)
            edge_sum = sum(
                nearest_odd(model.nu_sq * (
                    int(fields[component, (site+direction) % 7])
                    - int(fields[component, site])))
                for direction in (-1, 1)
            )
            expected = edge_sum - nearest_odd(model.tau**2 / model.delta * onsite)
            require(actual[component, site] == expected,
                    f'Fraction force mismatch at component {component}, site {site}')
            largest_digits = max(largest_digits, len(str(abs(expected))))
    return dict(
        passed=True, compared_components=14, random_seed=191,
        largest_force_result_digits=largest_digits,
        large_register_values=[str(10**30+1), str(-10**31+3)],
        model=model.to_dict(),
        note='Exact arithmetic check only; these large registers are not a continuum-limit experiment.',
    )


def model_from_saved_json(path):
    metadata = json.loads(Path(path).read_text())['model']
    keys = ('sites', 'length', 'courant', 'delta_factor', 'mass_phi_sq',
            'mass_chi_sq', 'eta', 'lambda_phi', 'lambda_chi', 'g',
            'species', 'modulus')
    return Model(**{key: metadata[key] for key in keys},
                 delta_override=Fraction(metadata['delta']))


def replay_saved_trajectory(base_path):
    json_path = base_path.with_suffix('.json')
    data_path = base_path.with_suffix('.npz')
    model = model_from_saved_json(json_path)
    with np.load(data_path, allow_pickle=False) as stored:
        fields, momenta = stored['Q'], stored['P']
    q, p = fields[0].astype(object), momenta[0].astype(object)
    for level in range(1, len(fields)):
        q, p = model.step(q, p, modular=True)
        require(np.array_equal(q, fields[level])
                and np.array_equal(p, momenta[level]),
                f'Forward replay mismatch: {relative_name(data_path)}, level {level}')
    for level in range(len(fields)-2, -1, -1):
        q, p = model.inverse(q, p, modular=True)
        require(np.array_equal(q, fields[level])
                and np.array_equal(p, momenta[level]),
                f'Inverse replay mismatch: {relative_name(data_path)}, level {level}')
    return dict(
        passed=True, input=relative_name(data_path), metadata=relative_name(json_path),
        stored_levels=len(fields), forward_steps=len(fields)-1,
        inverse_steps=len(fields)-1, all_integer_states_exact=True,
        note='Replay checks saved data against the production update; force correctness is checked separately by the Fraction oracle.',
    )


def parseval_check(base_path):
    model = model_from_saved_json(base_path.with_suffix('.json'))
    with np.load(base_path.with_suffix('.npz'), allow_pickle=False) as stored:
        fields = stored['Q']
    h, delta = float(model.h), float(model.delta)
    branch = delta * (fields[:, 0]-fields[:, 1]) / np.sqrt(2)
    direct_power = h * np.sum(branch * branch, axis=-1)
    fourier_power = (h / model.sites
                     * np.sum(abs(np.fft.fft(branch, axis=-1))**2, axis=-1))
    discrepancy = float(np.max(abs(direct_power-fourier_power)))
    tolerance = 64 * np.finfo(float).eps * max(1., float(direct_power.max()))
    require(discrepancy <= tolerance, 'Parseval normalization check failed')
    return dict(
        passed=True, input=relative_name(base_path.with_suffix('.npz')),
        max_absolute_discrepancy=discrepancy, absolute_tolerance=tolerance,
        direct_normalization='h sum_j |v_j|^2',
        fourier_normalization='(h/J) sum_k |FFT(v)_k|^2',
        note='This verifies the Parseval factor used by the saved Fourier/UV power diagnostics; it is not an energy-conservation claim.',
    )


def grid_alignment_check(coarse_path, fine_path):
    with np.load(coarse_path, allow_pickle=False) as coarse:
        coarse_times, coarse_shape = coarse['times'], coarse['fields'].shape
    with np.load(fine_path, allow_pickle=False) as fine:
        fine_times = fine['times']
        restricted_shape = fine['fields'][::2, :, ::2].shape
    require(np.array_equal(coarse_times, fine_times[::2]),
            'Saved coarse and restricted-fine time arrays differ')
    require(coarse_shape == restricted_shape,
            'Space-time restricted array shapes differ')
    return dict(
        passed=True, coarse_input=relative_name(coarse_path),
        fine_input=relative_name(fine_path), refinement_ratio=2,
        time_arrays_exactly_equal=True, coarse_shape=list(coarse_shape),
        restricted_fine_shape=list(restricted_shape),
    )


def timefit_sensitivity(results_root):
    rows = []
    factors = (0.95, 1., 1.05)
    for branch, sites, mode in (('u', 128, 1), ('v', 512, 0), ('v', 128, 4)):
        input_path = results_root / 'spectroscopy' / f'{branch}_grid_{sites}.npz'
        with np.load(input_path, allow_pickle=False) as stored:
            times = stored['times']
            positions = np.flatnonzero(stored['mode_numbers'] == mode)
            require(len(positions) == 1, f'Missing or duplicated mode {mode}')
            series = stored['projected_modes'][:, int(positions[0])]
        initial, recurrence_residual = mode_frequency(series, times[1]-times[0])
        fits = []
        for factor in factors:
            estimate, residual = temporal_frequency(times, series, initial * factor)
            fits.append(dict(initial_guess_factor=factor,
                             fitted_frequency=estimate,
                             relative_fit_residual=residual))
        frequencies = [row['fitted_frequency'] for row in fits]
        spread = max(frequencies)-min(frequencies)
        tolerance = 1e-6 * max(abs(value) for value in frequencies)
        require(spread <= tolerance,
                f'Temporal-fit search-interval sensitivity too large: {branch}, {sites}, {mode}')
        rows.append(dict(
            passed=True, input=relative_name(input_path), branch=branch,
            sites=sites, mode=mode, recurrence_initial_guess=initial,
            recurrence_relative_residual=recurrence_residual, fits=fits,
            absolute_frequency_spread=spread, absolute_spread_tolerance=tolerance,
        ))
    return dict(
        passed=True, rows=rows,
        largest_absolute_frequency_spread=max(row['absolute_frequency_spread'] for row in rows),
        note='No theoretical frequency enters these fits. The recurrence supplies only a data-derived search center. Search-interval sensitivity is not a statistical confidence interval.',
    )


def run_checks(results_root):
    replay_bases = [results_root / 'convergence/grid_32',
                   results_root / 'collision/n256_eps1over10_both']
    collision_base = replay_bases[1]
    coarse_path = collision_base.with_suffix('.npz')
    fine_path = results_root / 'collision/n512_eps1over10_both.npz'
    checks = dict(
        fraction_force_oracle=fraction_force_oracle(),
        saved_trajectory_replay=[replay_saved_trajectory(path) for path in replay_bases],
        parseval_power_normalization=parseval_check(collision_base),
        saved_grid_alignment=grid_alignment_check(coarse_path, fine_path),
        timefit_initial_search_sensitivity=timefit_sensitivity(results_root),
    )
    inputs = {path.with_suffix(extension)
              for path in replay_bases for extension in ('.json', '.npz')}
    inputs.add(fine_path)
    for branch, sites in (('u', 128), ('v', 512), ('v', 128)):
        inputs.add(results_root / 'spectroscopy' / f'{branch}_grid_{sites}.npz')
    sources = [Path(__file__), HERE / 'cell_auto/core.py',
               HERE / 'cell_auto/reference.py', HERE / 'cell_auto/experiments.py']
    return dict(
        passed=True,
        scope='Independent force oracle; selected exact saved-state replays; Parseval normalization; saved time-grid alignment; temporal-fit search sensitivity.',
        checks=checks,
        environment=dict(python=platform.python_version(), numpy=np.__version__, scipy=scipy.__version__),
        source_sha256={relative_name(path): sha256(path) for path in sources},
        input_sha256={relative_name(path): sha256(path) for path in sorted(inputs)},
        limitations='This is a limited reproducibility/code audit, not a second execution of the entire experiment suite or a substitute for the analytical convergence proof.',
    )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--results', default='results',
                        help='Input results directory, relative to this script unless absolute')
    parser.add_argument('--output', default='results/result_checks.json',
                        help='Output JSON, relative to this script unless absolute')
    args = parser.parse_args()
    results_root, output = Path(args.results), Path(args.output)
    if not results_root.is_absolute():
        results_root = HERE / results_root
    if not output.is_absolute():
        output = HERE / output
    report = run_checks(results_root)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(report, indent=2, ensure_ascii=False, allow_nan=False) + '\n')
    print(f'Result checks passed: {relative_name(output)}')
    levels = ', '.join(str(row['stored_levels']) for row in report['checks']['saved_trajectory_replay'])
    print(f'Fraction oracle: 14 components; saved-state replays: {levels} levels, forward and inverse.')
    print('Largest time-fit search-interval spread: '
          f'{report["checks"]["timefit_initial_search_sensitivity"]["largest_absolute_frequency_spread"]:.12g}')


if __name__ == '__main__':
    main()