#!/usr/bin/env python3
"""Run one deterministic Adam step for legacy or rewritten ROOT-GNN."""

from __future__ import annotations

import argparse
from pathlib import Path

import dgl
import numpy as np
import torch

from validation.artifacts import load_artifact
from validation.forward import _graph


def _weighted_loss(logits, labels, weights):
    elementwise = torch.nn.functional.cross_entropy(logits, labels, reduction="none")
    result = logits.new_zeros(())
    for label in torch.unique(labels):
        mask = labels == label
        result = (
            result + (weights[mask] * elementwise[mask]).sum() / weights[mask].sum()
        )
    return result / len(torch.unique(labels))


def _flat_state(model, *, gradients: bool = False):
    values = []
    for name, parameter in sorted(model.named_parameters()):
        if gradients:
            if parameter.grad is None:
                raise RuntimeError(f"missing gradient for {name}")
            values.append(parameter.grad.detach().cpu().reshape(-1))
        else:
            values.append(parameter.detach().cpu().reshape(-1))
    return torch.cat(values).numpy()


def _flat_optimizer_state(model, optimizer, key):
    values = []
    steps = []
    for name, parameter in sorted(model.named_parameters()):
        state = optimizer.state[parameter]
        values.append(state[key].detach().cpu().reshape(-1))
        steps.append(float(state["step"]))
    return torch.cat(values).numpy(), np.asarray(steps, dtype=np.float64)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--artifact", type=Path, required=True)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--events", type=int, default=8)
    parser.add_argument("--steps", type=int, default=1)
    args = parser.parse_args()
    artifact = load_artifact(args.artifact)
    count = min(args.events, artifact.event_count)
    graphs = [_graph(artifact, index) for index in range(count)]
    graph = dgl.batch(graphs)
    empty_globals = torch.empty((count, 0), dtype=torch.float32)
    from gnn4colliders.models.root_gnn import (
        EdgeNetwork,
        load_legacy_edge_network_state_dict,
    )

    model = EdgeNetwork(graphs[0], empty_globals[:1], 64, 12, 4, 4, dropout=0.0)
    payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    load_legacy_edge_network_state_dict(model, payload)
    model.train()
    labels = torch.from_numpy(artifact.labels[:count]).to(torch.long)
    weights = torch.from_numpy(artifact.weights[:count]).to(torch.float32)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    histories = []
    for _ in range(args.steps):
        optimizer.zero_grad()
        logits = model(graph, None)
        loss = _weighted_loss(logits, labels, weights)
        loss.backward()
        gradients = _flat_state(model, gradients=True)
        optimizer.step()
        histories.append(float(loss.detach()))
    exp_avg, optimizer_steps = _flat_optimizer_state(model, optimizer, "exp_avg")
    exp_avg_sq, _ = _flat_optimizer_state(model, optimizer, "exp_avg_sq")
    np.savez_compressed(
        args.output,
        logits=logits.detach().cpu().numpy(),
        loss=np.asarray(float(loss.detach())),
        history=np.asarray(histories, dtype=np.float64),
        gradients=gradients,
        parameters=_flat_state(model),
        optimizer_steps=optimizer_steps,
        exp_avg=exp_avg,
        exp_avg_sq=exp_avg_sq,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())