"""Independent conservation identity check and completed-artifact audit."""
import csv
import json
from pathlib import Path
import sys

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model import PINNShallowWater
from common import arguments, device_for, errors, load_config, output_path, runtime_info, save_json, sha256
from train import check_formal_config


def conservation_identity(config, device):
    """Compare expanded R=A*r to derivatives of the conserved fluxes.

    This checks a separate mathematical representation, including mixed
    momentum fluxes; it catches dropped mass-residual terms and axis mistakes.
    """
    torch.manual_seed(config['seed'])
    model = PINNShallowWater(**config['model']).to(device=device, dtype=torch.float64)
    points = torch.rand(19, 3, dtype=torch.float64, device=device).requires_grad_(True)
    points = points * torch.tensor([20, 20, 0.8], device=device) - torch.tensor([10, 10, 0], device=device)
    state, expanded = model.residuals(points, config['physics']['gravity'], True)
    h, u, v = state.unbind(1)
    g = config['physics']['gravity']
    conserved = torch.stack((h, h * u, h * v), dim=1)
    flux_x = torch.stack((h * u, h * u.square() + 0.5 * g * h.square(), h * u * v), dim=1)
    flux_y = torch.stack((h * v, h * u * v, h * v.square() + 0.5 * g * h.square()), dim=1)
    expected = torch.zeros_like(expanded)
    for component in range(3):
        for value, axis in [(conserved, 2), (flux_x, 0), (flux_y, 1)]:
            expected[:, component] += torch.autograd.grad(value[:, component].sum(), points, retain_graph=True)[0][:, axis]
    maximum_error = (expanded - expected).abs().max().item()
    torch.testing.assert_close(expanded, expected, rtol=1e-9, atol=1e-10)
    (expanded.square().mean() + model.slope_recovery()).backward()
    assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters())
    assert model.slopes.numel() == 1500
    return {'flux_vs_expanded_max_abs': maximum_error, 'backward_all_parameters_finite': True,
            'neuron_slopes': model.slopes.numel(), 'double_precision_identity_passed': True}


def main():
    parser = arguments(__doc__)
    parser.add_argument('--preflight', action='store_true')
    args = parser.parse_args()
    config, root = load_config(args.config)
    check_formal_config(config)
    device = device_for(config)
    identity = conservation_identity(config, device)
    if args.preflight:
        report = {'status': 'passed', 'execution_readiness': 'ready', 'preflight_passed': True,
                  'evidence': {'preflight': {'status': 'passed'}},
                  'runtime': runtime_info(device), 'conservation_identity': identity,
                  'entry_config': 'formal dCDB9 checked', 'training_performed': False}
        save_json(root / '.onescience/preflight_compute.json', report)
        print('PREFLIGHT_COMPUTE_PASSED', report, flush=True)
        return
    results = output_path(root, config['paths']['results'])
    summary = json.loads((results / 'training_summary.json').read_text())
    metrics = json.loads((results / 'metrics.json').read_text())
    with (results / 'training_history.csv').open() as handle:
        history = list(csv.DictReader(handle))
    assert len(history) == 5000
    assert [int(row['iteration']) for row in history] == list(range(1, 5001))
    assert all(float(row['learning_rate']) == 0.001 for row in history)
    assert all(np.isfinite(float(row[key])) for row in history for key in row)
    best_path = output_path(root, config['paths']['best_checkpoint'])
    best = torch.load(best_path, map_location='cpu', weights_only=False)
    last = torch.load(output_path(root, config['paths']['last_checkpoint']), map_location='cpu', weights_only=False)
    assert last['iteration'] == best['completed_iterations'] == summary['completed_iterations'] == 5000
    assert best['iteration'] == summary['best_iteration']
    assert best['formal_run_complete'] and not summary['stopped_early']
    model = PINNShallowWater(**config['model']).to(device)
    model.load_state_dict(best['model_state_dict'], strict=True)
    assert all(torch.isfinite(p).all() for p in model.parameters())
    assert sha256(best_path) == summary['best_checkpoint_sha256'] == metrics['checkpoint_sha256']
    assert sha256(output_path(root, config['paths']['samples'])) == best['sampling_sha256']
    with (np.load(output_path(root, config['paths']['predictions'])) as prediction,
          np.load(output_path(root, config['paths']['hll_reference'])) as reference):
        fields, refs = prediction['fields'], reference['fields']
        assert fields.shape == refs.shape == (3, 500, 500, 3)
        assert np.isfinite(fields).all() and np.isfinite(refs).all()
        for k in range(3):
            recomputed = errors(fields[k].reshape(-1, 3), refs[k].reshape(-1, 3))
            for key in ['MAE', 'RMSE']:
                np.testing.assert_allclose(recomputed[key], metrics['errors'][k][key], rtol=1e-10, atol=1e-12)
    figures = ['dCDB9_fields_and_errors.png', 'dCDB9_fields_and_errors.pdf',
               'dCDB9_centerline_comparison.png', 'dCDB9_water_depth_surfaces.png', 'training_loss.png']
    for name in figures:
        assert (results / name).stat().st_size > 1000
    final = {'status': 'passed', 'formal_iterations': 5000, 'history_rows': len(history),
             'learning_rate_constant': True, 'checkpoint_strict_load': True,
             'checkpoint_finite': True, 'checkpoint_sha256': sha256(best_path),
             'conservative_identity': identity, 'predictions_shape': list(fields.shape),
             'metrics_independently_recomputed': True, 'figures_nonempty': figures,
             'runtime': runtime_info(device), 'accuracy_claim': 'actual metrics reported; no claim of paper accuracy parity'}
    save_json(results / 'verification.json', final)
    print('FINAL_VERIFICATION_PASSED', final, flush=True)


if __name__ == '__main__':
    main()