"""A real PINN for normalized exponential decay: u'(t) + u(t) = 0, u(0) = 1.

No sensors and no time/value training labels. The exact curve is evaluation only.
Run: python decay_pinn.py --output-dir runs/decay
"""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

import torch
from torch import nn


class DecayPINN(nn.Module):
    """A smooth independent-per-row map from time [N,1] to value [N,1]."""

    def __init__(self) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, 16), nn.Tanh(), nn.Linear(16, 16), nn.Tanh(), nn.Linear(16, 1)
        )

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        return self.net(t)


def derivative(u: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """Require grad-enabled independent coordinate rows. Return per-row du/dt."""
    return torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u), create_graph=True)[0]


def losses(model: nn.Module, t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Return ODE residual loss and initial-condition loss. No exact labels enter."""
    u = model(t)
    residual = derivative(u, t) + u
    physics = residual.square().mean()
    initial = (model(torch.zeros(1, 1)) - 1).square().mean()
    return physics, initial


def train(*, epochs: int = 1500, seed: int = 7) -> tuple[DecayPINN, list[dict]]:
    """Perform epochs finite Adam steps. Epoch count is not a convergence proof."""
    if epochs <= 0:
        raise ValueError("Epoch count must be positive")
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    model = DecayPINN()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.003)
    history = []
    for epoch in range(epochs):  # Bound: remaining epochs.
        t = (2 * torch.rand(64, 1)).requires_grad_()
        optimizer.zero_grad(set_to_none=True)
        physics, initial = losses(model, t)
        total = physics + 10 * initial
        if not torch.isfinite(total):
            raise RuntimeError("Nonfinite PINN loss")
        total.backward()
        if any(p.grad is None or not torch.isfinite(p.grad).all() for p in model.parameters()):
            raise RuntimeError("Missing or nonfinite PINN gradient")
        optimizer.step()
        if epoch % 100 == 0 or epoch == epochs - 1:
            history.append(
                {"epoch": epoch + 1, "physics": physics.item(), "initial": initial.item()}
            )
    return model.eval(), history


def evaluate(model: DecayPINN) -> tuple[dict, list[list[float]]]:
    """Evaluate new grid points against the exact solution. They never train the model."""
    t = torch.linspace(0, 2, 101).reshape(-1, 1).requires_grad_()
    prediction = model(t)
    expected = torch.exp(-t)
    residual = derivative(prediction, t) + prediction
    errors = {
        "rmse": (prediction - expected).square().mean().sqrt().item(),
        "max_absolute_error": (prediction - expected).abs().max().item(),
        "residual_rmse": residual.square().mean().sqrt().item(),
    }
    rows = torch.cat((t, expected, prediction), dim=1).detach().tolist()
    return errors, rows


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--epochs", type=int, default=1500)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--output-dir", type=Path, default=Path("runs/decay"))
    args = parser.parse_args()
    model, history = train(epochs=args.epochs, seed=args.seed)
    errors, rows = evaluate(model)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    with (args.output_dir / "evaluation.csv").open("w", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["t", "reference_evaluation_only", "prediction"])
        writer.writerows(rows)
    torch.save(model.state_dict(), args.output_dir / "model.pt")
    report = {
        "epochs": args.epochs,
        "seed": args.seed,
        "parameter_count": sum(p.numel() for p in model.parameters()),
        "training_value_labels": 0,
        "initial_condition": 1,
        "errors": errors,
        "history": history,
    }
    (args.output_dir / "results.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()