"""Train a small Physics-Informed Neural Network (PINN) on the 1D heat equation.

PDE: u_t = alpha*u_xx on 0 <= x <= 1 and 0 <= t <= 1.
Initial: u(x, 0) = sin(pi*x). Boundaries: u(0,t) = u(1,t) = 0.
Exact reference: exp(-alpha*pi**2*t)*sin(pi*x).

Requires Python 3.10+, PyTorch, NumPy, and Matplotlib.
Example: python pinn_heat_equation.py --epochs 2500
"""
from __future__ import annotations

import argparse
import csv
import json
import math
from pathlib import Path

import matplotlib
matplotlib.use("Agg")  # Save figures on machines without a display.
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn


def exact_solution(x: torch.Tensor, t: torch.Tensor, alpha: float) -> torch.Tensor:
    """Q: equally shaped coordinates; alpha > 0. R: exact heat solution at each pair."""
    return torch.exp(-alpha * math.pi**2 * t) * torch.sin(math.pi * x)


def generate_data(path: Path, alpha: float) -> None:
    """Q: writable path, alpha > 0. R: CSV with 121 exact (x,t,u) rows."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["x", "t", "u"])
        # Invariant: all preceding grid points were written once.
        # Bound: 121 - number of written rows.
        for index in range(121):
            t, x = (index // 11) / 10, (index % 11) / 10
            u = math.exp(-alpha * math.pi**2 * t) * math.sin(math.pi * x)
            writer.writerow([f"{x:.1f}", f"{t:.1f}", f"{u:.12f}"])


def load_data(path: Path, device: torch.device) -> tuple[torch.Tensor, ...]:
    """Q: nonempty x,t,u CSV. R: finite float32 columns of shape (N,1), x,t in [0,1]."""
    with path.open(newline="") as handle:
        reader = csv.DictReader(handle)
        if reader.fieldnames != ["x", "t", "u"]:
            raise ValueError("CSV header must be x,t,u in that order")
        rows = [[float(row[key]) for key in ("x", "t", "u")] for row in reader]
    if not rows:
        raise ValueError("CSV has no observations")
    values = torch.tensor(rows, dtype=torch.float32, device=device)
    if not torch.isfinite(values).all():
        raise ValueError("CSV contains NaN or infinity")
    if ((values[:, :2] < 0) | (values[:, :2] > 1)).any():
        raise ValueError("x and t must be within [0,1]")
    return values[:, 0:1], values[:, 1:2], values[:, 2:3]


class PINN(nn.Module):
    """Q: width > 0. R: smooth, independent-per-row map (N,2) -> (N,1)."""

    def __init__(self, width: int = 32) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, width), nn.Tanh(),
            nn.Linear(width, width), nn.Tanh(),
            nn.Linear(width, width), nn.Tanh(),
            nn.Linear(width, 1),
        )
        # Invariant: preceding Linear layers have Xavier weights and zero biases.
        # Bound: remaining layers in self.net.
        for layer in self.net:
            if isinstance(layer, nn.Linear):
                nn.init.xavier_normal_(layer.weight)
                nn.init.zeros_(layer.bias)

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """Q: x,t are float columns on the model's device. R: temperature column."""
        coordinates = torch.cat((x, t), dim=1)
        return self.net(2.0 * coordinates - 1.0)


def derivative(y: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
    """Q: differentiable, independent rows; x requires grad. R: per-row dy/dx."""
    return torch.autograd.grad(
        y, x, grad_outputs=torch.ones_like(y), create_graph=True
    )[0]


def pde_residual(model: nn.Module, x: torch.Tensor, t: torch.Tensor,
                 alpha: float) -> torch.Tensor:
    """Q: grad-enabled coordinate columns. R: u_t - alpha*u_xx at each point."""
    u = model(x, t)
    u_t = derivative(u, t)
    u_x = derivative(u, x)
    u_xx = derivative(u_x, x)
    return u_t - alpha * u_xx


def physics_points(count: int, device: torch.device) -> tuple[torch.Tensor, ...]:
    """Q: count > 0. R: fresh grad-enabled columns sampled uniformly in [0,1)."""
    x = torch.rand(count, 1, device=device, requires_grad=True)
    t = torch.rand(count, 1, device=device, requires_grad=True)
    return x, t


def condition_points(count: int, device: torch.device) -> tuple[torch.Tensor, ...]:
    """Q: count >= 2. R: initial and both-end boundary coordinates and targets."""
    x_i = torch.linspace(0, 1, count, device=device).reshape(-1, 1)
    t_i, u_i = torch.zeros_like(x_i), torch.sin(math.pi * x_i)
    times = torch.linspace(0, 1, count, device=device).reshape(-1, 1)
    x_b = torch.cat((torch.zeros_like(times), torch.ones_like(times)))
    t_b = torch.cat((times, times))
    return x_i, t_i, u_i, x_b, t_b


def loss_terms(model: PINN, physics: tuple, conditions: tuple, data: tuple,
               alpha: float) -> dict[str, torch.Tensor]:
    """Q: valid point tuples. R: four scalar, nonnegative mean-squared losses."""
    x_f, t_f = physics
    x_i, t_i, u_i, x_b, t_b = conditions
    x_d, t_d, u_d = data
    return {
        "physics": pde_residual(model, x_f, t_f, alpha).square().mean(),
        "initial": (model(x_i, t_i) - u_i).square().mean(),
        "boundary": model(x_b, t_b).square().mean(),
        "data": (model(x_d, t_d) - u_d).square().mean(),
    }


def train(args: argparse.Namespace, device: torch.device) -> tuple[PINN, list]:
    """Q: validated args. R: exactly epochs Adam steps, finite parameters and logs."""
    path = Path(args.data)
    if not path.exists():
        generate_data(path, args.alpha)
        print(f"Generated synthetic observations: {path}")
    x, t, u = load_data(path, device)
    # Select observations strictly inside x and after initial time.
    mask = ((x > 0) & (x < 1) & (t > 0)).reshape(-1)
    x, t, u = x[mask], t[mask], u[mask]
    if len(x) == 0:
        raise ValueError("Data must include at least one interior observation at t > 0")
    indices = torch.randperm(len(x), device=device)[:args.data_points]
    data = x[indices], t[indices], u[indices]
    model = PINN(args.width).to(device)
    conditions = condition_points(args.n_conditions, device)
    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
    history = []
    print(f"Device: {device}; training observations: {len(indices)}; physics points/step: {args.n_physics}")
    # Invariant: epoch-1 complete optimizer updates; history has epoch-1 rows.
    # Bound: args.epochs - epoch + 1. This guarantees steps, not convergence.
    for epoch in range(1, args.epochs + 1):
        optimizer.zero_grad(set_to_none=True)
        terms = loss_terms(model, physics_points(args.n_physics, device),
                           conditions, data, args.alpha)
        loss = (args.physics_weight * terms["physics"]
                + 10 * terms["initial"] + 10 * terms["boundary"]
                + args.data_weight * terms["data"])
        if not torch.isfinite(loss):
            raise RuntimeError("Nonfinite loss; reduce the learning rate")
        loss.backward()
        if any(p.grad is None or not torch.isfinite(p.grad).all() for p in model.parameters()):
            raise RuntimeError("Nonfinite or missing gradient; inspect data and learning rate")
        optimizer.step()
        if any(not torch.isfinite(p).all() for p in model.parameters()):
            raise RuntimeError("Nonfinite model parameter after update; reduce learning rate")
        row = {"epoch": epoch, "total": loss.item(),
               **{name: term.item() for name, term in terms.items()}}
        history.append(row)
        if epoch == 1 or epoch % args.log_every == 0 or epoch == args.epochs:
            print(" | ".join([f"epoch {epoch:5d}"] +
                            [f"{name} {row[name]:.3e}" for name in row if name != "epoch"]))
    return model, history


def evaluate(model: PINN, device: torch.device, alpha: float) -> tuple[dict, list]:
    """Q: trained model, alpha > 0. R: dense-grid errors plus new-point residual RMSE."""
    axis = torch.linspace(0, 1, 101, device=device)
    t_grid, x_grid = torch.meshgrid(axis, axis, indexing="ij")
    x, t = x_grid.reshape(-1, 1), t_grid.reshape(-1, 1)
    model.eval()
    with torch.no_grad():
        predicted, exact = model(x, t), exact_solution(x, t, alpha)
        difference = predicted - exact
        relative_l2 = torch.linalg.vector_norm(difference) / torch.linalg.vector_norm(exact)
        metrics = {"relative_l2": relative_l2.item(),
                   "max_absolute_error": difference.abs().max().item()}
        arrays = [v.reshape(101, 101).cpu().numpy()
                  for v in (exact, predicted, difference.abs())]
    # Physics evaluation requires autograd, even though model is in eval mode.
    residual = pde_residual(model, *physics_points(1000, device), alpha)
    metrics["residual_rmse"] = residual.square().mean().sqrt().item()
    return metrics, arrays


def save_results(model: PINN, history: list, metrics: dict, arrays: list,
                 args: argparse.Namespace) -> None:
    """Q: completed training/evaluation. R: checkpoint, metrics, CSV, and two PNGs."""
    output = Path(args.output_dir)
    output.mkdir(parents=True, exist_ok=True)
    torch.save({"state_dict": model.state_dict(), "width": args.width, "alpha": args.alpha},
               output / "pinn_heat_equation.pt")
    (output / "metrics.json").write_text(json.dumps({"config": vars(args), **metrics}, indent=2) + "\n")
    with (output / "loss_history.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(history[0]))
        writer.writeheader()
        writer.writerows(history)
    fig, axes = plt.subplots(1, 3, figsize=(13, 4), constrained_layout=True)
    # Invariant: preceding axes contain correct arrays; bound: remaining 3-index.
    for index, title in enumerate(("Exact", "PINN", "Absolute error")):
        image = axes[index].imshow(arrays[index], origin="lower", extent=[0, 1, 0, 1],
                                   aspect="auto", cmap="viridis" if index < 2 else "magma",
                                   vmin=0, vmax=1 if index < 2 else None)
        axes[index].set(xlabel="position x", ylabel="time t", title=title)
        fig.colorbar(image, ax=axes[index])
    fig.savefig(output / "pinn_results.png", dpi=160)
    plt.close(fig)
    fig, plot = plt.subplots(figsize=(8, 4), constrained_layout=True)
    # Invariant: earlier loss series are drawn; bound: remaining keys.
    for name in ("total", "physics", "initial", "boundary", "data"):
        plot.semilogy([row["epoch"] for row in history],
                      [max(row[name], 1e-16) for row in history], label=name)
    plot.set(xlabel="optimizer step (epoch)", ylabel="loss", title="Training history")
    plot.legend()
    fig.savefig(output / "loss_history.png", dpi=160)
    plt.close(fig)


def predict(args: argparse.Namespace, device: torch.device) -> None:
    """Q: compatible checkpoint and coordinates in [0,1]. R: print predicted u."""
    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=True)
    model = PINN(checkpoint["width"]).to(device)
    model.load_state_dict(checkpoint["state_dict"])
    model.eval()
    x, t = (torch.tensor([[value]], device=device) for value in args.predict)
    with torch.no_grad():
        value = model(x, t).item()
    print(f"u(x={args.predict[0]}, t={args.predict[1]}) = {value:.8f}")


def self_test() -> int:
    """Check the analytic residual, tensor shapes, parameters, and full loss gradients."""
    device = torch.device("cpu")
    class Exact(nn.Module):
        def forward(self, x, t):
            return exact_solution(x, t, 0.1)
    x, t = physics_points(32, device)
    if pde_residual(Exact(), x, t, 0.1).abs().max().item() > 1e-6:
        raise AssertionError("analytic solution does not satisfy the residual")
    model = PINN()
    if model(x, t).shape != (32, 1):
        raise AssertionError("prediction shape must be (N,1)")
    if sum(p.numel() for p in model.parameters()) != 2241:
        raise AssertionError("unexpected parameter count")
    data = x.detach(), t.detach(), exact_solution(x.detach(), t.detach(), 0.1)
    terms = loss_terms(model, (x, t), condition_points(8, device), data, 0.1)
    sum(terms.values()).backward()
    if any(p.grad is None or not torch.isfinite(p.grad).all() for p in model.parameters()):
        raise AssertionError("loss must backpropagate finite gradients")
    return 4


def parse_args() -> argparse.Namespace:
    """Q: command-line values. R: finite, positive counts/rates and valid coordinates."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--epochs", type=int, default=2500)
    parser.add_argument("--width", type=int, default=32)
    parser.add_argument("--n-physics", type=int, default=1000)
    parser.add_argument("--n-conditions", type=int, default=64)
    parser.add_argument("--data-points", type=int, default=16)
    parser.add_argument("--learning-rate", type=float, default=0.001)
    parser.add_argument("--alpha", type=float, default=0.1)
    parser.add_argument("--physics-weight", type=float, default=1.0)
    parser.add_argument("--data-weight", type=float, default=1.0)
    parser.add_argument("--data", default="heat_equation_data.csv")
    parser.add_argument("--output-dir", default="output")
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--log-every", type=int, default=250)
    parser.add_argument("--threads", type=int, default=1)
    parser.add_argument("--predict", type=float, nargs=2, metavar=("X", "T"))
    parser.add_argument("--checkpoint", default="output/pinn_heat_equation.pt")
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test and args.predict is not None:
        parser.error("--self-test and --predict are mutually exclusive")
    # Invariant: previous fields passed validation; bound: remaining field names.
    for name in ("epochs", "width", "n_physics", "n_conditions", "data_points", "log_every", "threads"):
        if getattr(args, name) <= 0:
            parser.error(f"--{name.replace('_', '-')} must be positive")
    if args.n_conditions < 2:
        parser.error("--n-conditions must be at least 2")
    for name in ("alpha", "learning_rate", "physics_weight", "data_weight"):
        value = getattr(args, name)
        if not math.isfinite(value) or value < 0 or (name in ("alpha", "learning_rate") and value == 0):
            parser.error(f"--{name.replace('_', '-')} has an invalid value")
    if args.predict and any(not math.isfinite(v) or v < 0 or v > 1 for v in args.predict):
        parser.error("--predict coordinates must be in [0,1]")
    return args


def main() -> None:
    """Q: valid CLI. R: train and save results, or load and predict one value."""
    args = parse_args()
    torch.set_num_threads(args.threads)
    torch.manual_seed(args.seed)
    if args.self_test:
        print(f"{self_test()} checks passed")
        return
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if args.predict is not None:
        predict(args, device)
        return
    model, history = train(args, device)
    metrics, arrays = evaluate(model, device, args.alpha)
    save_results(model, history, metrics, arrays, args)
    print(json.dumps(metrics, indent=2))
    print(f"Saved model, figures, and logs in {args.output_dir}/")


if __name__ == "__main__":
    main()