"""Paper-aligned two-layer streamfunction model-error correction."""

from __future__ import annotations

import json
import math
import os
import random
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Dict, Iterable, Tuple

import numpy as np
import torch
import torch.distributed as dist
import yaml
from torch import Tensor, nn
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, Dataset, DistributedSampler


@dataclass
class QGConfig:
    nx: int = 40
    ny: int = 20
    reference_dt_minutes: int = 10
    model_dt_minutes: int = 20
    observation_interval_minutes: int = 120
    window_batches: int = 12
    diffusion: float = 0.002
    truth_advection: float = 0.18
    model_advection: float = 0.15
    truth_coupling: float = 0.025
    model_coupling: float = 0.018
    truth_damping: float = 0.006
    model_damping: float = 0.009

    def __post_init__(self) -> None:
        if (self.nx, self.ny) != (40, 20):
            raise ValueError("The paper state grid is fixed at nx=40, ny=20")
        if self.model_dt_minutes != 20 or self.reference_dt_minutes != 10:
            raise ValueError("Paper time steps are fixed at 20 min (model) and 10 min (reference)")
        if self.observation_interval_minutes != 120 or self.window_batches != 12:
            raise ValueError("A DA window must contain 12 observation batches spaced by 2 h")

    @property
    def state_size(self) -> int:
        return 2 * self.ny * self.nx

    @property
    def model_steps_per_observation(self) -> int:
        return self.observation_interval_minutes // self.model_dt_minutes


def seed_all(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def load_yaml(path: str | Path) -> dict:
    with open(path, "r", encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def _fixed_y(field: Tensor) -> Tensor:
    field = field.clone()
    field[..., 0, :] = 0.0
    field[..., -1, :] = 0.0
    return field


class TwoLayerQG(nn.Module):
    """Executable channel dynamics for a two-layer streamfunction state.

    x differences wrap periodically. y derivatives are interior centered
    differences and both meridional streamfunction boundaries stay fixed.
    """

    def __init__(self, config: QGConfig, truth: bool = False):
        super().__init__()
        self.config = config
        self.truth = truth

    def tendency(self, psi: Tensor) -> Tensor:
        dx = 2.0 * math.pi / self.config.nx
        dy = 1.0 / (self.config.ny - 1)
        ddx = (torch.roll(psi, -1, -1) - torch.roll(psi, 1, -1)) / (2.0 * dx)
        ddy = torch.zeros_like(psi)
        ddy[..., 1:-1, :] = (psi[..., 2:, :] - psi[..., :-2, :]) / (2.0 * dy)
        lap = (torch.roll(psi, -1, -1) - 2.0 * psi + torch.roll(psi, 1, -1)) / dx**2
        lap[..., 1:-1, :] += (psi[..., 2:, :] - 2.0 * psi[..., 1:-1, :] + psi[..., :-2, :]) / dy**2
        advection = self.config.truth_advection if self.truth else self.config.model_advection
        coupling = self.config.truth_coupling if self.truth else self.config.model_coupling
        damping = self.config.truth_damping if self.truth else self.config.model_damping
        velocity_x = 0.35 + 0.15 * torch.tanh(-ddy)
        velocity_y = 0.08 * torch.tanh(ddx)
        other = psi.flip(1)
        tendency = -advection * (velocity_x * ddx + velocity_y * ddy)
        tendency = tendency + self.config.diffusion * lap + coupling * (other - psi) - damping * psi
        tendency[..., 0, :] = 0.0
        tendency[..., -1, :] = 0.0
        return tendency

    def _step(self, psi: Tensor, dt_minutes: int) -> Tensor:
        dt = dt_minutes / 120.0
        midpoint = _fixed_y(psi + 0.5 * dt * self.tendency(psi))
        return _fixed_y(psi + dt * self.tendency(midpoint))

    def forward(self, psi: Tensor) -> Tensor:
        if self.truth:
            state = self._step(psi, self.config.reference_dt_minutes)
            return self._step(state, self.config.reference_dt_minutes)
        return self._step(psi, self.config.model_dt_minutes)

    def advance_window(self, psi: Tensor) -> Tensor:
        state = psi
        for _ in range(self.config.window_batches * self.config.model_steps_per_observation):
            state = self(state)
        return state


class DModel(nn.Module):
    """Final paper D model: one 8-node linear hidden Dense layer."""

    def __init__(self, state_size: int = 1600, hidden_size: int = 8):
        super().__init__()
        if state_size != 1600:
            raise ValueError("D model requires the complete 1600-component state")
        if hidden_size < 8:
            raise ValueError("The final D model hidden layer cannot be smaller than 8")
        self.state_size = state_size
        self.hidden_size = hidden_size
        self.input_dense = nn.Linear(state_size, hidden_size)
        self.output_dense = nn.Linear(hidden_size, state_size)

    def forward(self, psi: Tensor) -> Tensor:
        shape = psi.shape
        return self.output_dense(self.input_dense(psi.reshape(shape[0], self.state_size))).reshape(shape)


class HybridSurrogate(nn.Module):
    def __init__(self, config: QGConfig, hidden_size: int = 8):
        super().__init__()
        self.knowledge = TwoLayerQG(config, truth=False)
        self.correction = DModel(config.state_size, hidden_size)

    def forward(self, analysis: Tensor) -> Tensor:
        return self.knowledge.advance_window(analysis) + self.correction(analysis)


def structured_wave_fields(count: int, config: QGConfig, device: torch.device) -> Tensor:
    """Create smooth channel waves with periodic x and zero fixed y edges."""
    x = torch.arange(config.nx, device=device) * (2.0 * math.pi / config.nx)
    y = torch.linspace(0.0, math.pi, config.ny, device=device)
    yy, xx = torch.meshgrid(y, x, indexing="ij")
    fields = []
    for _ in range(count):
        layers = []
        shared_phase = 2.0 * math.pi * torch.rand((), device=device)
        for layer in range(2):
            psi = torch.zeros_like(xx)
            for mode in range(1, 5):
                phase = shared_phase + 0.35 * layer + 2.0 * math.pi * torch.rand((), device=device)
                amplitude = (0.25 + 0.5 * torch.rand((), device=device)) / mode
                psi += amplitude * torch.sin((mode % 3 + 1) * yy) * torch.cos(mode * xx + phase)
            layers.append(psi)
        fields.append(_fixed_y(torch.stack(layers)))
    return torch.stack(fields)


def _sample_observation_geometry(samples: int, batches: int, count: int, config: QGConfig) -> Tuple[Tensor, Tensor, Tensor]:
    layer = torch.randint(0, 2, (samples, batches, count))
    x = torch.rand(samples, batches, count) * config.nx
    y = 1.0 + torch.rand(samples, batches, count) * (config.ny - 3)
    x0 = torch.floor(x).long() % config.nx
    y0 = torch.floor(y).long().clamp(0, config.ny - 2)
    x1 = (x0 + 1) % config.nx
    y1 = y0 + 1
    indices = torch.stack((layer, y0, x0, layer, y0, x1, layer, y1, x0, layer, y1, x1), -1)
    indices = indices.reshape(samples, batches, count, 4, 3)
    wx, wy = x - torch.floor(x), y - torch.floor(y)
    weights = torch.stack(((1 - wx) * (1 - wy), wx * (1 - wy), (1 - wx) * wy, wx * wy), -1)
    locations = torch.stack((layer.float(), y, x), -1)
    return locations, indices, weights


def bilinear_observe(states: Tensor, indices: Tensor, weights: Tensor) -> Tensor:
    values = []
    for corner in range(4):
        index = indices[..., corner, :]
        batch = torch.arange(states.shape[0], device=states.device)[:, None, None]
        time = torch.arange(states.shape[1], device=states.device)[None, :, None]
        values.append(states[batch, time, index[..., 0], index[..., 1], index[..., 2]])
    return (torch.stack(values, -1) * weights).sum(-1)


@torch.no_grad()
def create_dataset(path: str, samples: int, config: QGConfig, seed: int = 7,
                   observation_count: int = 50, observation_variance: float = 0.1,
                   analysis_gain: float = 0.35, format_version: str = "2.0") -> None:
    seed_all(seed)
    truth_model = TwoLayerQG(config, truth=True)
    model = TwoLayerQG(config, truth=False)
    start_truth = structured_wave_fields(samples, config, torch.device("cpu"))
    locations, indices, weights = _sample_observation_geometry(samples, config.window_batches, observation_count, config)
    truth_batches, state = [], start_truth
    for _ in range(config.window_batches):
        for _ in range(config.model_steps_per_observation):
            state = truth_model(state)
        truth_batches.append(state)
    truth_batches = torch.stack(truth_batches, 1)
    clean_observations = bilinear_observe(truth_batches, indices, weights)
    observations = clean_observations + math.sqrt(observation_variance) * torch.randn_like(clean_observations)

    # Executable sparse-observation analysis: sequential bilinear innovation spreading.
    analysis = start_truth + 0.08 * torch.randn_like(start_truth)
    analysis = _fixed_y(analysis)
    for batch in range(config.window_batches):
        for _ in range(config.model_steps_per_observation):
            analysis = model(analysis)
        for corner in range(4):
            idx = indices[:, batch, :, corner]
            predicted = bilinear_observe(analysis[:, None], indices[:, batch:batch + 1], weights[:, batch:batch + 1])[:, 0]
            innovation = observations[:, batch] - predicted
            for sample in range(samples):
                analysis[sample].index_put_(tuple(idx[sample].T), analysis_gain * weights[sample, batch, :, corner] * innovation[sample], accumulate=True)
        analysis = _fixed_y(analysis)
    next_analysis = analysis
    model_forecast = model.advance_window(start_truth)
    target_increment = next_analysis - model_forecast
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(path, analysis=start_truth.numpy(), next_analysis=next_analysis.numpy(),
                        model_forecast=model_forecast.numpy(), target_increment=target_increment.numpy(),
                        truth_window=truth_batches.numpy(), observations=observations.numpy(),
                        observation_locations=locations.numpy(), observation_indices=indices.numpy(),
                        observation_weights=weights.numpy(), observation_variance=np.array(observation_variance),
                        window_start_utc=np.array("01:00"), model_config=np.array(json.dumps(asdict(config))),
                        format_version=np.array(format_version))


class IncrementDataset(Dataset):
    def __init__(self, path: str, format_version: str):
        archive = np.load(path)
        actual = str(archive["format_version"]) if "format_version" in archive.files else None
        if actual != format_version:
            raise ValueError(f"Data format_version must be {format_version}, got {actual}")
        required = {"analysis", "next_analysis", "model_forecast", "target_increment", "observations", "observation_indices", "observation_weights"}
        if missing := required.difference(archive.files):
            raise ValueError(f"Dataset is missing fields: {sorted(missing)}")
        if archive["analysis"].shape[1:] != (2, 20, 40) or archive["observations"].shape[1:] != (12, 50):
            raise ValueError("Strict shapes are analysis [N,2,20,40] and observations [N,12,50]")
        expected = archive["next_analysis"] - archive["model_forecast"]
        if not np.allclose(archive["target_increment"], expected, rtol=1e-6, atol=1e-6):
            raise ValueError("Target must equal x_a_{k+1} - M_o(x_a_k)")
        self.analysis = torch.from_numpy(archive["analysis"]).float()
        self.target = torch.from_numpy(archive["target_increment"]).float()

    def __len__(self) -> int:
        return len(self.analysis)

    def __getitem__(self, index: int) -> Tuple[Tensor, Tensor]:
        return self.analysis[index], self.target[index]


def setup_distributed() -> Tuple[int, int, int, torch.device]:
    world_size, rank, local_rank = (int(os.environ.get(key, default)) for key, default in (("WORLD_SIZE", "1"), ("RANK", "0"), ("LOCAL_RANK", "0")))
    if world_size > 1 and not dist.is_initialized():
        dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
    device = torch.device("cuda", local_rank) if torch.cuda.is_available() else torch.device("cpu")
    if device.type == "cuda":
        torch.cuda.set_device(local_rank)
    return rank, world_size, local_rank, device


def train_model(data_path: str, checkpoint_path: str, metrics_path: str, config: QGConfig,
                model_settings: dict, training_settings: dict) -> Dict[str, Iterable[float]]:
    rank, world_size, local_rank, device = setup_distributed()
    seed_all(training_settings["seed"] + rank)
    dataset = IncrementDataset(data_path, training_settings["format_version"])
    sampler = DistributedSampler(dataset, shuffle=True) if world_size > 1 else None
    loader = DataLoader(dataset, batch_size=training_settings["batch_size"], shuffle=sampler is None, sampler=sampler)
    model = DModel(config.state_size, model_settings["hidden_size"]).to(device)
    trainable = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) if world_size > 1 else model
    history, stage_history = [], []
    for stage, (epochs, learning_rate) in enumerate(zip(training_settings["stage_epochs"], training_settings["stage_learning_rates"]), 1):
        optimizer = torch.optim.Adam(trainable.parameters(), lr=learning_rate)
        for epoch in range(epochs):
            if sampler is not None:
                sampler.set_epoch(len(history))
            total = 0.0
            for analysis, target in loader:
                analysis, target = analysis.to(device), target.to(device)
                optimizer.zero_grad(set_to_none=True)
                loss = nn.functional.mse_loss(trainable(analysis), target)
                loss.backward()
                optimizer.step()
                total += float(loss.detach())
            history.append(total / len(loader))
        stage_history.append({"stage": stage, "epochs": epochs, "learning_rate": learning_rate, "final_loss": history[-1]})
    if rank == 0:
        Path(checkpoint_path).parent.mkdir(parents=True, exist_ok=True)
        torch.save({"model": model.state_dict(), "qg_config": asdict(config), "hidden_size": model_settings["hidden_size"],
                    "format_version": training_settings["format_version"]}, checkpoint_path)
        Path(metrics_path).parent.mkdir(parents=True, exist_ok=True)
        with open(metrics_path, "w", encoding="utf-8") as handle:
            json.dump({"training_loss_by_epoch": history, "stages": stage_history, "world_size": world_size}, handle, indent=2)
    if dist.is_initialized():
        dist.barrier()
        dist.destroy_process_group()
    return {"loss": history}


def load_model(checkpoint_path: str, device: torch.device, format_version: str) -> Tuple[DModel, QGConfig, dict]:
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    if checkpoint.get("format_version") != format_version:
        raise ValueError("Checkpoint format version mismatch")
    config = QGConfig(**checkpoint["qg_config"])
    model = DModel(config.state_size, checkpoint["hidden_size"]).to(device)
    model.load_state_dict(checkpoint["model"])
    model.eval()
    return model, config, checkpoint


def run_inference(data_path: str, checkpoint_path: str, output_path: str, format_version: str) -> Dict[str, np.ndarray]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    correction, config, _ = load_model(checkpoint_path, device, format_version)
    archive = np.load(data_path)
    if str(archive["format_version"]) != format_version:
        raise ValueError("Data format version mismatch")
    analysis = torch.from_numpy(archive["analysis"]).float().to(device)
    target = torch.from_numpy(archive["next_analysis"]).float().to(device)
    knowledge = TwoLayerQG(config).to(device)
    with torch.no_grad():
        knowledge_prediction = knowledge.advance_window(analysis)
        predicted_increment = correction(analysis)
        hybrid_prediction = knowledge_prediction + predicted_increment
    result = {"format_version": np.array(format_version), "analysis": analysis.cpu().numpy(),
              "target_analysis": target.cpu().numpy(), "knowledge_prediction": knowledge_prediction.cpu().numpy(),
              "predicted_increment": predicted_increment.cpu().numpy(), "hybrid_prediction": hybrid_prediction.cpu().numpy(),
              "observations": archive["observations"], "observation_locations": archive["observation_locations"],
              "observation_indices": archive["observation_indices"], "observation_weights": archive["observation_weights"]}
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output_path, **result)
    return result


def validate_and_plot(npz_path: str, metrics_path: str, figure_path: str) -> Dict[str, float]:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    result = np.load(npz_path)
    required = {"analysis", "target_analysis", "knowledge_prediction", "predicted_increment", "hybrid_prediction", "observations", "observation_indices", "observation_weights"}
    if missing := required.difference(result.files):
        raise ValueError(f"Missing inference arrays: {sorted(missing)}")
    if result["analysis"].shape[1:] != (2, 20, 40) or result["observations"].shape[1:] != (12, 50):
        raise ValueError("Inference did not preserve the complete paper state/window dimensions")
    for key in required:
        if not np.isfinite(result[key]).all():
            raise ValueError(f"Non-finite values in {key}")
    target = result["target_analysis"]
    knowledge_rmse = float(np.sqrt(np.mean((result["knowledge_prediction"] - target) ** 2)))
    hybrid_rmse = float(np.sqrt(np.mean((result["hybrid_prediction"] - target) ** 2)))
    increment_rmse = float(np.sqrt(np.mean((result["predicted_increment"] - (target - result["knowledge_prediction"])) ** 2)))
    fig, axes = plt.subplots(1, 3, figsize=(12, 3.5))
    axes[0].bar(["model", "hybrid"], [knowledge_rmse, hybrid_rmse], color=["#67788a", "#c45b3c"])
    axes[0].set(ylabel="streamfunction RMSE", title="Next analysis")
    for axis, field, title in zip(axes[1:], [target[0, 0], result["hybrid_prediction"][0, 0]], ["target upper psi", "hybrid upper psi"]):
        image = axis.imshow(field, origin="lower", cmap="RdBu_r")
        axis.set_title(title)
        fig.colorbar(image, ax=axis, shrink=0.75)
    fig.tight_layout()
    Path(figure_path).parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(figure_path, dpi=140)
    plt.close(fig)
    summary = {"knowledge_analysis_rmse": knowledge_rmse, "hybrid_analysis_rmse": hybrid_rmse,
               "analysis_increment_rmse": increment_rmse, "samples": int(target.shape[0]),
               "state_shape": [2, 20, 40], "observation_shape_per_sample": [12, 50]}
    Path(metrics_path).parent.mkdir(parents=True, exist_ok=True)
    with open(metrics_path, "w", encoding="utf-8") as handle:
        json.dump(summary, handle, indent=2)
    return summary