"""OneScience installer preflight and independent structural checks; no formal weights."""
import ast
import importlib
import json
import os
import shutil
import time

import torch
from torch.utils.data import DataLoader

from common import (ROOT, Trajectories, Windows, build_model, device, load_config,
                    output_dir, parser, provenance, save_json, setup, validate)
from train import train_batch


def main():
    args = parser(__doc__).parse_args()
    cfg = load_config(args.config)
    setup(cfg["seed"])
    target_device = device()
    for module in ("onescience", "torch", "numpy", "h5py", "yaml", "matplotlib"):
        importlib.import_module(module)
        print("DEPENDENCY_OK " + module, flush=True)
    for path in ROOT.rglob("*.py"):
        ast.parse(path.read_text())
    for command in ("sbatch", "squeue", "sacct"):
        assert shutil.which(command), command
    datasets = {split: Trajectories(cfg, split) for split in cfg["data"]["expected_counts"]}
    out = output_dir(cfg)
    save_json(out / "dataset_validation.json", {s: d.metadata for s, d in datasets.items()})
    model = build_model(cfg, target_device)
    sample = datasets["train"].values[0, :6].unsqueeze(0).to(target_device)
    model.eval()
    with torch.no_grad():
        initial = sample[:, 0]
        predicted = model(initial)
        assert predicted.shape == initial.shape
        assert torch.isfinite(predicted).all()
        assert predicted[:, 0].min() >= 0
        drift = (predicted.double().sum((-2, -1)) - initial.double().sum((-2, -1))).abs()
        scale = initial.double().abs().sum((-2, -1)).clamp_min(1.0)
        assert (drift / scale).max() < 2e-6
        # A uniform state remains uniform for all learned transport plans.
        uniform = torch.ones_like(initial)
        assert torch.allclose(model(uniform), uniform, atol=2e-6)
        # Every channel conserves its sum for arbitrary signed momentum.
        random_state = torch.rand_like(initial)
        random_state[:, 1:] -= 0.5
        result = model(random_state)
        random_drift = (result.double().sum((-2, -1)) - random_state.double().sum((-2, -1))).abs()
        assert (random_drift / random_state.double().abs().sum((-2, -1)).clamp_min(1)).max() < 2e-6
    print("STRUCTURE_OK forward, positivity, uniform state, per-channel conservation", flush=True)
    model.train()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
    first, terminal = train_batch(model, sample, optimizer, 5)
    assert any(p.grad is not None for p in model.parameters())
    validation = validate(model, DataLoader(Windows(datasets["val"], 5), batch_size=2), target_device)
    assert validation["mse"] >= 0
    # Independent stop-gradient check on an explicit two-state feedback graph.
    x = sample[:, 0].detach().requires_grad_(True)
    intermediate = model(x)
    final = model(intermediate.detach())
    assert torch.autograd.grad(final.sum(), x, allow_unused=True)[0] is None
    props = torch.cuda.get_device_properties(target_device)
    result = {"status": "passed", "preflight_passed": True, "execution_readiness": "ready",
              "evidence": {"preflight": {"status": "passed"}}, "installer_reason": "preflight_validation",
              "execution_channel": "local_slurm", "provenance": provenance(),
              "gpu_memory_gb": props.total_memory / 2**30,
              "model_parameters": sum(p.numel() for p in model.parameters()),
              "structural_checks": ["shape", "finite", "positivity", "uniform_state", "conservation", "detached_feedback"],
              "real_data_train_batch_loss": {"step1": first, "step5": terminal},
              "real_data_validation": validation,
              "formal_training_completed": False, "formal_weights_written": False}
    save_json(out / "preflight_validation.json", result)
    print("PREFLIGHT_PASSED " + json.dumps(result), flush=True)


if __name__ == "__main__":
    main()