#!/usr/bin/env python3
"""Six-path smoke checks and independent derivative/gradient equivalence checks."""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import torch

from scripts.common import (ROOT, build_model, gather, load_dns, predict, read_config,
                            relative_l2, sample_indices, save_json, setup, tensor)
from scripts.train import objective


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--benchmark", action="store_true")
    args = parser.parse_args()
    cfg = read_config()
    device, dtype = setup(cfg, args.device)
    dns = load_dns(cfg)
    index = sample_indices(cfg, dns)
    bq, buv = gather(dns, index["boundary"])
    rq, _ = gather(dns, index["residual"])
    model = build_model(cfg, device)
    q = tensor(rq[:32], device, dtype)
    checks = {}
    output = model(q)
    assert output.shape == (32, 2) and torch.isfinite(output).all()
    checks["forward_pass"] = "passed"
    values, residuals, div = model.residual_autograd(q)
    tvalues, tresiduals, _ = model.residual_taylor(q)
    torch.testing.assert_close(tvalues, values, atol=1e-10, rtol=1e-9)
    torch.testing.assert_close(tresiduals, residuals, atol=1e-10, rtol=1e-9)
    params = tuple(model.parameters())
    aloss = residuals.square().mean() + values[:, :2].square().mean() + 0.25 * model.slope_recovery()
    tloss = tresiduals.square().mean() + tvalues[:, :2].square().mean() + 0.25 * model.slope_recovery()
    ag = torch.autograd.grad(aloss, params)
    tg = torch.autograd.grad(tloss, params)
    for left, right in zip(ag, tg):
        torch.testing.assert_close(left, right, atol=2e-10, rtol=2e-8)
    assert float(div.abs().max()) < 1e-10
    checks["backward_pass"] = "passed: finite parameter gradients; Taylor/reverse-mode equivalence"
    bq_t, buv_t, rq_t = (tensor(x, device, dtype) for x in (bq[:64], buv[:64], rq[:64]))
    opt = torch.optim.Adam(model.parameters(), lr=cfg["training"]["learning_rate"], eps=1e-7)
    losses = []
    for _ in range(3):
        opt.zero_grad(set_to_none=True)
        value = objective(model, bq_t, buv_t, rq_t, cfg, backward=True)
        opt.step()
        losses.append(value["loss"])
    assert np.isfinite(losses).all() and losses[-1] < losses[0], losses
    checks["train_loop_dry_run"] = {"status": "passed", "losses": losses}
    valq, valuv = gather(dns, index["validation"])
    valp = predict(model, valq[:128], device, dtype, 64)
    error = relative_l2(valp[:, :2], valuv[:128])
    assert np.isfinite(error).all()
    checks["validation_loop_dry_run"] = {"status": "passed", "relative_l2": error.tolist()}
    assert np.allclose(relative_l2(dns["uv"].reshape(-1, 2), dns["uv"].reshape(-1, 2)), 0)
    assert np.allclose(relative_l2(np.zeros_like(dns["uv"].reshape(-1, 2)), dns["uv"].reshape(-1, 2)), 1)
    checks["domain_metric_dry_run"] = "passed: L2 self=0, zero predictor=1; streamfunction divergence <1e-10"
    assert len(model.frequencies) == 4 and sum(f.numel() for f in model.frequencies) == 200
    assert [l.out_features for l in model.layers] == [50, 50, 50, 50, 2]
    assert index["boundary"].shape == (200, 400) and index["residual"].shape == (200, 1000)
    checks["config_consistency"] = "passed: 4x50, 200 frequencies, (x,y,t)->(psi,p), 4000 configured updates, per-time counts"
    benchmark = None
    if args.benchmark:
        model = build_model(cfg, device)
        fullb, fulluv, fullr = (tensor(x, device, dtype) for x in (bq, buv, rq))
        started = time.perf_counter()
        model.zero_grad(set_to_none=True)
        value = objective(model, fullb, fulluv, fullr, cfg, backward=True)
        if device.type == "cuda":
            torch.cuda.synchronize()
        benchmark = {"seconds_one_full_objective_backward": time.perf_counter()-started, "score": value,
                     "peak_device_memory_bytes": torch.cuda.max_memory_allocated() if device.type == "cuda" else 0}
    result = {"status": "passed", "attempt": 1, "checks": checks, "benchmark": benchmark,
              "parameter_count": sum(p.numel() for p in model.parameters()), "torch": torch.__version__,
              "device": torch.cuda.get_device_name() if device.type == "cuda" else "CPU"}
    save_json(ROOT / ".onescience/smoke_result.json", result)
    print(json.dumps(result, ensure_ascii=False, indent=2), flush=True)


if __name__ == "__main__":
    main()