"""Neural prefix free energies and forward trajectory-balance learning."""

import json, time
from pathlib import Path
import numpy as np
import torch
from torch import nn
from .exact import sample, endpoint_distribution, normalize_rewards


class Sampler(nn.Module):
    """Value and edge-policy MLPs over fixed-width node features.

    The feature matrix has shape (nodes, feature_width). Edge logits use the
    source feature, destination feature, and scalar execution cost.
    """

    def __init__(self, graph, hidden=64):
        super().__init__()
        self.graph = graph
        widths = {len(n.features) for n in graph.nodes}
        if len(widths) != 1 or 0 in widths:
            raise ValueError("All nodes require equally sized nonempty features")
        x = torch.tensor([n.features for n in graph.nodes], dtype=torch.float32)
        self.register_buffer("x", x)
        self.register_buffer(
            "src", torch.tensor([graph.node_index[e.source] for e in graph.edges])
        )
        self.register_buffer(
            "dst", torch.tensor([graph.node_index[e.target] for e in graph.edges])
        )
        self.register_buffer(
            "cost", torch.tensor([e.cost for e in graph.edges], dtype=torch.float32)
        )
        d = x.shape[1]
        self.value = nn.Sequential(
            nn.Linear(d, hidden),
            nn.SiLU(),
            nn.Linear(hidden, hidden),
            nn.SiLU(),
            nn.Linear(hidden, 1),
        )
        self.policy = nn.Sequential(
            nn.Linear(2 * d + 1, hidden),
            nn.SiLU(),
            nn.Linear(hidden, hidden),
            nn.SiLU(),
            nn.Linear(hidden, 1),
        )
        self.log_z = nn.Parameter(torch.tensor(0.0))

    def values(self):
        """Evaluate prefix log partitions with the root value fixed to zero."""
        v = self.value(self.x).squeeze(-1)
        return v - v[self.graph.root_index]

    def forward_logs(self):
        """Evaluate edge log probabilities, normalized over outgoing edges."""
        features = torch.cat(
            [self.x[self.src], self.x[self.dst], self.cost[:, None]], dim=1
        )
        logits = self.policy(features).squeeze(-1)
        out = torch.empty_like(logits)
        for ids in self.graph.outgoing:
            if ids:
                out[ids] = torch.log_softmax(logits[ids], dim=0)
        return out

    def backward_logs(self, temperature):
        """Normalize cost-weighted predecessor values for each destination."""
        v = self.values()
        out = torch.empty_like(self.cost)
        for ids in self.graph.incoming:
            if ids:
                out[ids] = torch.log_softmax(
                    v[self.src[ids]] - self.cost[ids] / temperature, dim=0
                )
        return out

    def prefix_loss(self, temperature):
        """Mean fitted Bellman error over every nonroot node in the graph."""
        v = self.values()
        pred = []
        target = []
        # Detached targets implement fitted soft Bellman updates.
        for i in self.graph.order:
            ids = self.graph.incoming[i]
            if ids:
                pred.append(v[i])
                target.append(
                    torch.logsumexp(
                        v[self.src[ids]].detach() - self.cost[ids] / temperature, dim=0
                    )
                )
        return torch.mean((torch.stack(pred) - torch.stack(target)) ** 2)

    @torch.no_grad()
    def probabilities(self):
        """Return one float64 probability per edge, normalized by source node."""
        p = self.forward_logs().exp().cpu().numpy().astype(float)
        for ids in self.graph.outgoing:
            if ids:
                p[ids] /= p[ids].sum()
        return p


def train(
    graph,
    log_rewards,
    temperature=1.0,
    steps=2000,
    batch_size=64,
    hidden=64,
    lr=0.002,
    seed=0,
    exploration=0.15,
    output=None,
    backward="learned",
    resume=None,
    initialize=None,
):
    """Fit a sampler on a declared graph and optionally write a checkpoint.

    Log rewards have one finite value per canonical terminal outcome. Steps is
    the total update count, including any updates restored through resume.
    Route temperature and edge costs must use consistent units. Checkpoints
    include optimizer state and the exact graph needed for reproducible reload.
    """
    if not 0 <= exploration <= 1:
        raise ValueError("Exploration must lie in [0,1]")
    if not np.isfinite(temperature) or temperature <= 0:
        raise ValueError("Temperature must be positive")
    if steps < 1 or batch_size < 1:
        raise ValueError("Steps and batch size must be positive")
    normalize_rewards(graph, log_rewards)
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.set_num_threads(1)
    model = Sampler(graph, hidden)
    # A log-mean initialization avoids a long scalar-normalizer warmup.
    # The normalizer remains trainable throughout trajectory-balance fitting.
    with torch.no_grad():
        model.log_z.fill_(
            np.log(len(log_rewards)) + np.mean(list(log_rewards.values()))
        )
    if initialize:
        source, _ = load_model(initialize)
        if source.graph != graph:
            raise ValueError("Warm-start initialization requires the same graph")
        model.policy.load_state_dict(source.policy.state_dict())
        model.value.load_state_dict(source.value.state_dict())
    if initialize and resume:
        raise ValueError("Choose either resume or warm-start initialization")
    optp = torch.optim.Adam(list(model.policy.parameters()) + [model.log_z], lr=lr)
    optv = torch.optim.Adam(model.value.parameters(), lr=lr)
    from .exact import uniform_policy, prefix_values, backward_policy

    uniform = uniform_policy(graph)
    if backward == "exact":
        fixed = backward_policy(graph, prefix_values(graph, temperature), temperature)
    elif backward == "uniform":
        fixed = np.empty(len(graph.edges))
        for ids in graph.incoming:
            if ids:
                fixed[ids] = 1 / len(ids)
    elif backward in ["learned", "unnormalized"]:
        fixed = None
    else:
        raise ValueError("Unknown backward policy")
    history = []
    start_step = 0
    if resume:
        restored, state = load_model(resume)
        if restored.graph != graph:
            raise ValueError("Resume graph differs from the checkpoint graph")
        for key, value in [
            ("temperature", temperature),
            ("log_rewards", log_rewards),
            ("backward", backward),
            ("seed", seed),
            ("batch_size", batch_size),
            ("exploration", exploration),
            ("lr", lr),
        ]:
            if state.get(key) != value:
                raise ValueError(f"Resume setting differs for {key}")
        model.load_state_dict(restored.state_dict())
        optp.load_state_dict(state["policy_optimizer"])
        optv.load_state_dict(state["value_optimizer"])
        start_step = state["steps"]
        history = json.loads((Path(resume) / "training.json").read_text())
    start = time.perf_counter()
    for step in range(start_step, steps):
        if fixed is None:
            optv.zero_grad()
            vloss = model.prefix_loss(temperature)
            vloss.backward()
            torch.nn.utils.clip_grad_norm_(model.value.parameters(), 10)
            optv.step()
        else:
            vloss = torch.tensor(0.0)
        # Behavior probabilities are used only to collect trajectories.
        behavior = (1 - exploration) * model.probabilities() + exploration * uniform
        paths = sample(graph, behavior, batch_size, seed + step * 1009)
        logp = model.forward_logs()
        if backward == "unnormalized":
            v = model.values().detach()
            logq = v[model.src] - model.cost / temperature - v[model.dst]
        else:
            logq = (
                model.backward_logs(temperature).detach()
                if fixed is None
                else torch.tensor(np.log(fixed), dtype=torch.float32)
            )
        residual = torch.stack(
            [
                model.log_z
                + logp[p["edge_indices"]].sum()
                - float(log_rewards[p["outcome"]])
                - logq[p["edge_indices"]].sum()
                for p in paths
            ]
        )
        loss = (residual**2).mean()
        optp.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.policy.parameters(), 10)
        optp.step()
        if step % 100 == 0 or step == steps - 1:
            target, _ = normalize_rewards(graph, log_rewards)
            actual = endpoint_distribution(graph, model.probabilities())
            history.append(
                {
                    "step": step,
                    "balance_mse": float(loss.detach()),
                    "prefix_mse": float(vloss.detach()),
                    "endpoint_tv": sum(abs(actual[y] - target[y]) for y in target) / 2,
                    "seconds": time.perf_counter() - start,
                }
            )
    if output:
        out = Path(output)
        out.mkdir(parents=True, exist_ok=True)
        graph.save(out / "graph.json")
        torch.save(
            {
                "state_dict": model.state_dict(),
                "hidden": hidden,
                "temperature": temperature,
                "log_rewards": log_rewards,
                "seed": seed,
                "steps": steps,
                "backward": backward,
                "batch_size": batch_size,
                "exploration": exploration,
                "lr": lr,
                "policy_optimizer": optp.state_dict(),
                "value_optimizer": optv.state_dict(),
            },
            out / "model.pt",
        )
        (out / "training.json").write_text(json.dumps(history, indent=2))
    return model, history


def load_model(directory):
    """Reload the serialized graph, network weights, and training settings."""
    from .graph import Graph

    directory = Path(directory)
    g = Graph.load(directory / "graph.json")
    d = torch.load(directory / "model.pt", map_location="cpu", weights_only=True)
    model = Sampler(g, d["hidden"])
    model.load_state_dict(d["state_dict"])
    model.eval()
    return model, d