"""
model.py — Degradation Function Estimation

Implements the model from the project PDF:

    dhi/dt = K * (∏_c I_c) * ∑_j (θ_ij,I * h_j + θ_ij,II * h_j * ln(h_j))

K is absorbed into theta_I and theta_II, so no separate input-scaling parameter is estimated.

Parameters:
    theta_I     : (N, N) — linear health coupling; theta_I[i,j] = θ_ij,I
    theta_II    : (N, N) — log-linear health coupling; theta_II[i,j] = θ_ij,II

Adjacency mask A (N x N, binary):
    A[i, j] = 1  iff component j is physically allowed to influence component i.
    Built from component_graph.COMPONENT_GRAPH so that the learned theta matrices
    can only be non-zero where a real physical coupling exists.  The mask is applied
    element-wise:   theta_I_eff = A * theta_I
    Positions where A[i,j] = 0 are zeroed on init and their gradients are zeroed
    during fitting — the model cannot learn phantom interactions.

Stochastic extension (§5): each Euler step subtracts Q_i * H_i where
    Q_i ~ N(0,1)²  (squared standard normal — event intensity, always ≥ 0)
    H_i ~ Poisson(λ_i * τ / n)  (event count per step; λ_i set per component)
Negative health values are valid in this mode and represent catastrophic failure.

Fitting: Euler-forward simulation + manual Jacobian recurrence → SGD + L1 regularisation.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import List, Optional, Tuple

import numpy as np


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------

@dataclass
class Sample:
    """One training example."""
    X: np.ndarray   # (C,) constant input vector over [0, tau]
    tau: float      # time horizon
    y: np.ndarray   # (N,) target health at tau
    M: np.ndarray   # (N,) mask: 1 = observed, 0 = ignored

    @classmethod
    def from_json(cls, path: str) -> List["Sample"]:
        """
        Load a list of samples from a JSON file.

        Expected format — a JSON array where each element has:
            "X"   : list of C floats  (inputs)
            "tau" : float             (time horizon)
            "y"   : list of N floats  (target health per component)
            "M"   : list of N ints    (mask: 1 = observed, 0 = ignored)
        """
        with open(path, "r") as f:
            records = json.load(f)
        return [
            cls(
                X   = np.array(r["X"],   dtype=float),
                tau = float(r["tau"]),
                y   = np.array(r["y"],   dtype=float),
                M   = np.array(r["M"],   dtype=float),
            )
            for r in records
        ]


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------

class DegradationModel:
    """
    Parametric ODE model for multi-component health degradation.

    State starts at h(0) = 1 (all components fully healthy).
    The ODE for component i couples to every component j via theta_I and theta_II,
    and to every input c via theta_input.
    """

    def __init__(
        self,
        N: int,
        C: int,
        lambda_rates: Optional[np.ndarray] = None,
        adjacency_mask: Optional[np.ndarray] = None,
        seed: int = 0,
    ) -> None:
        self.N = N
        self.C = C
        # Per-component Poisson rates λ_i for the stochastic shock term (§5).
        # Defaults to zeros — no randomness unless explicitly set.
        self.lambda_rates: np.ndarray = (
            np.asarray(lambda_rates, dtype=float)
            if lambda_rates is not None
            else np.zeros(N, dtype=float)
        )

        # Adjacency mask A[i,j] = 1 iff component j may influence component i.
        # If None, defaults to all-ones (fully connected — no structural constraint).
        # Pass build_adjacency_matrix() to enforce the physical graph topology.
        if adjacency_mask is not None:
            self.A: np.ndarray = np.asarray(adjacency_mask, dtype=float)
            if self.A.shape != (N, N):
                raise ValueError(
                    f"adjacency_mask must be ({N},{N}), got {self.A.shape}"
                )
        else:
            self.A = np.ones((N, N), dtype=float)

        rng = np.random.default_rng(seed)
        # Small negative diagonal drives self-degradation; off-diagonals start near 0.
        # Mask is applied immediately so forbidden positions start at exactly 0.
        self.theta_I: np.ndarray = -np.abs(rng.normal(0.0, 1e-3, (N, N))) * self.A
        np.fill_diagonal(self.theta_I, -1e-2)   # diagonal always in mask (self-coupling)
        self.theta_II: np.ndarray = np.zeros((N, N), dtype=float)

    # ------------------------------------------------------------------
    # Parameter vector helpers
    # ------------------------------------------------------------------

    @property
    def num_params(self) -> int:
        return 2 * self.N * self.N

    def _get_params(self) -> np.ndarray:
        return np.concatenate([
            self.theta_I.ravel(),
            self.theta_II.ravel(),
        ])

    def _set_params(self, p: np.ndarray) -> None:
        N = self.N
        self.theta_I  = np.minimum(p[: N * N].reshape(N, N), 0.0)
        self.theta_II = np.maximum(p[N * N :].reshape(N, N), 0.0)

    # ------------------------------------------------------------------
    # ODE
    # ------------------------------------------------------------------

    @staticmethod
    def _safe_hlog(h: np.ndarray) -> np.ndarray:
        """h * ln(h) with h clipped away from 0."""
        h_safe = np.clip(h, 1e-10, None)
        return h_safe * np.log(h_safe)

    def _P(self, I: np.ndarray) -> float:
        """Input product P = ∏_c I_c."""
        return float(np.prod(I))

    def f(self, h: np.ndarray, I: np.ndarray) -> np.ndarray:
        """Rate vector dh/dt, shape (N,).

        The mask A is applied element-wise before the matrix products so that
        forbidden couplings (A[i,j]=0) never contribute to dh/dt regardless of
        the current value of theta_I or theta_II.
        """
        P     = self._P(I)
        h_log = self._safe_hlog(h)
        g = (self.A * self.theta_I) @ h + (self.A * self.theta_II) @ h_log  # (N,)
        return P * g

    # ------------------------------------------------------------------
    # Forward simulation (Euler integration)
    # ------------------------------------------------------------------

    def simulate(
        self,
        X: np.ndarray,
        tau: float,
        n_steps: int = 100,
        stochastic: bool = False,
        seed: Optional[int] = None,
    ) -> np.ndarray:
        """
        Integrate from h(0)=1 to h(tau) using Euler steps.

        When stochastic=True, each step subtracts a random shock Q_i * H_i where
            Q_i ~ N(0,1)²   (squared standard normal — event intensity)
            H_i ~ Poisson(lambda_rates[i] * tau / n_steps)
        Negative health values are kept as-is; they represent catastrophic failure.

        Returns shape (n_steps + 1, N) — row 0 is h(0), row k is h(k * tau / n_steps).
        """
        h = np.ones(self.N, dtype=float)
        dt = tau / n_steps
        rng = np.random.default_rng(seed)
        trajectory = [h.copy()]

        for _ in range(n_steps):
            dh = dt * self.f(h, X)
            if stochastic:
                Q = rng.standard_normal(self.N) ** 2                    # N(0,1)²
                H = rng.poisson(self.lambda_rates * tau / n_steps)      # Poisson(λi*τ/n)
                dh -= Q * H
            h = h + dh
            if not stochastic:
                h = np.clip(h, 0.0, 1.0)
            trajectory.append(h.copy())

        return np.array(trajectory)  # (n_steps + 1, N)

    # ------------------------------------------------------------------
    # Loss
    # ------------------------------------------------------------------

    def loss(self, y_hat: np.ndarray, y: np.ndarray, M: np.ndarray) -> float:
        """Masked MSE: ∑_i (ŷi - yi)² * Mi."""
        return float(np.sum((y_hat - y) ** 2 * M))

    # ------------------------------------------------------------------
    # Jacobians
    # ------------------------------------------------------------------

    def _df_dtheta(self, h: np.ndarray, I: np.ndarray) -> np.ndarray:
        """
        ∂f/∂θ, shape (N, num_params).

        Columns correspond to [theta_I (flattened) | theta_II (flattened)].
        Gradient columns for positions where A[i,j]=0 are zeroed so those
        parameters receive no update signal during backprop.
        """
        N = self.N
        P = self._P(I)
        h_log = self._safe_hlog(h)

        jac = np.zeros((N, self.num_params))

        # ∂fi/∂θij,I = P * hj  — then zero out forbidden positions via A
        raw_I = np.kron(np.eye(N), P * h[np.newaxis, :])       # (N, N*N)
        jac[:, : N * N] = raw_I * self.A.ravel()[np.newaxis, :]

        # ∂fi/∂θij,II = P * hj * ln(hj) — same masking
        raw_II = np.kron(np.eye(N), P * h_log[np.newaxis, :])  # (N, N*N)
        jac[:, N * N :] = raw_II * self.A.ravel()[np.newaxis, :]

        return jac

    def _df_dh(self, h: np.ndarray, I: np.ndarray) -> np.ndarray:
        """
        ∂f/∂h (Jacobian of rate w.r.t. state), shape (N, N).

        ∂fi/∂hj = KP * (θij,I + θij,II * (1 + ln(hj)))
        """
        P = self._P(I)
        h_safe = np.clip(h, 1e-10, None)
        d_log = 1.0 + np.log(h_safe)   # d/dhj [hj ln hj] = 1 + ln hj
        return P * (self.theta_I + self.theta_II * d_log[np.newaxis, :])

    # ------------------------------------------------------------------
    # Gradient via Jacobian recurrence
    # ------------------------------------------------------------------

    def compute_gradient(
        self,
        sample: Sample,
        n_steps: int = 50,
        J_clip: float = 1e6,
    ) -> Tuple[np.ndarray, float]:
        """
        Gradient of the masked MSE loss for one sample, via:

            Jθŷ(t + dt) = Jθŷ(t) + dt * (∂f/∂θ + ∂f/∂ŷ · Jθŷ(t))
            Jθŷ(0) = 0

        where Jθŷ = ∂ŷ/∂θ has shape (N, num_params).

        Returns (gradient w.r.t. params, scalar loss).
        """
        X, tau, y, M = sample.X, sample.tau, sample.y, sample.M
        dt = tau / n_steps
        N, P = self.N, self.num_params

        h = np.ones(N, dtype=float)
        J = np.zeros((N, P))    # Jθŷ

        for _ in range(n_steps):
            df_dt = self._df_dtheta(h, X)  # (N, P)
            df_dh = self._df_dh(h, X)       # (N, N)
            J = np.clip(J + dt * (df_dt + df_dh @ J), -J_clip, J_clip)
            h = np.clip(h + dt * self.f(h, X), 0.0, 1.0)

        # ∂L/∂θr = 2 * ∑_i (ŷi − yi) * Mi * ∂ŷi/∂θr
        residual = (h - y) * M      # (N,)
        grad = 2.0 * (residual @ J) # (P,)
        loss_val = float(np.sum(residual ** 2))
        return grad, loss_val

    # ------------------------------------------------------------------
    # Summary
    # ------------------------------------------------------------------

    def summary(self, component_names: Optional[List[str]] = None) -> None:
        """Print a human-readable overview of the fitted parameters."""
        comp = component_names or [f"comp_{i}" for i in range(self.N)]
        w = max(len(n) for n in comp)   # column width

        print(f"DegradationModel  N={self.N}  C={self.C}  params={self.num_params}")
        print()

        print("Linear health coupling (theta_I[i,j])  -- row i influenced by col j:")
        header = " " * (w + 4) + "  ".join(f"{n:>{w}}" for n in comp)
        print(header)
        for i, row_name in enumerate(comp):
            vals = "  ".join(
                f"{self.theta_I[i, j]:+{w}.4f}" if self.A[i, j] else " " * (w + 1) + "-"
                for j in range(self.N)
            )
            print(f"  {row_name:<{w}}  {vals}")
        print()

        print("Log-linear health coupling (theta_II[i,j])  -- row i influenced by col j:")
        print(header)
        for i, row_name in enumerate(comp):
            vals = "  ".join(
                f"{self.theta_II[i, j]:+{w}.4f}" if self.A[i, j] else " " * (w + 1) + "-"
                for j in range(self.N)
            )
            print(f"  {row_name:<{w}}  {vals}")
        print()

        print("Stochastic shock rates (lambda_rates):")
        for i, (name, lam) in enumerate(zip(comp, self.lambda_rates)):
            print(f"  {name}: lambda={lam:.6f}")

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save(self, path: str) -> None:
        """Save all model arrays to a .npz file."""
        np.savez(
            path,
            N            = self.N,
            C            = self.C,
            lambda_rates = self.lambda_rates,
            A            = self.A,
            theta_I      = self.theta_I,
            theta_II     = self.theta_II,
        )

    @classmethod
    def load(cls, path: str) -> "DegradationModel":
        """Load a model saved with save()."""
        d = np.load(path)
        m = cls(
            N              = int(d["N"]),
            C              = int(d["C"]),
            lambda_rates   = d["lambda_rates"],
            adjacency_mask = d["A"],
        )
        m.theta_I  = np.minimum(d["theta_I"],  0.0)
        m.theta_II = np.maximum(d["theta_II"], 0.0)
        return m

    # ------------------------------------------------------------------
    # Fitting (SGD + L1)
    # ------------------------------------------------------------------

    def fit(
        self,
        dataset: List[Sample],
        lr: float = 1e-3,
        epochs: int = 100,
        lambda_l1: float = 1e-4,
        n_steps: int = 50,
        batch_size: Optional[int] = None,
        verbose: bool = True,
    ) -> List[float]:
        """
        Stochastic gradient descent with L1 regularisation.

        L1 promotes sparsity — zero parameters mean no coupling between components
        or inputs, letting the model discover the true dependency structure.

        Returns per-epoch average loss history.
        """
        rng = np.random.default_rng(0)
        loss_history: List[float] = []

        for epoch in range(epochs):
            indices = rng.permutation(len(dataset))
            if batch_size is not None:
                batches: List[np.ndarray] = [
                    indices[i : i + batch_size]
                    for i in range(0, len(indices), batch_size)
                ]
            else:
                batches = [indices]

            epoch_loss = 0.0
            for batch_idx in batches:
                results = [self.compute_gradient(dataset[i], n_steps) for i in batch_idx]
                grads, losses = zip(*results)
                grad = np.mean(grads, axis=0)
                batch_loss = float(np.mean(losses))

                # L1 subgradient
                params = self._get_params()
                grad = grad + lambda_l1 * np.sign(params)

                # Gradient clipping to prevent exploding updates
                grad_norm = float(np.linalg.norm(grad))
                if grad_norm > 1.0:
                    grad = grad / grad_norm

                self._set_params(params - lr * grad)

                # Re-apply mask after update: forbidden positions must stay at 0
                # even if numerical noise crept in through the L1 subgradient.
                self.theta_I  *= self.A
                self.theta_II *= self.A

                epoch_loss += batch_loss

            epoch_loss /= len(batches)
            loss_history.append(epoch_loss)

            if verbose and (epoch % max(1, epochs // 10) == 0 or epoch == epochs - 1):
                print(f"Epoch {epoch:4d}/{epochs}: loss = {epoch_loss:.6f}")

        return loss_history


COMPONENT_NAMES: List[str] = [
    "recoater_blade",
    "nozzle_plate",
    "heating_elements",
    "temperature_sensors",
    "insulation_panels",
    "firing_resistors",
    "cleaning_interface",
    "recoater_motor",
    "linear_rail",
]

INPUT_NAMES: List[str] = [
    "ambient_temperature_c",
    "build_chamber_temp_c",
    "ambient_humidity_pct",
    "powder_contamination_level",
    "print_hours",
    "build_volume_cm3",
    "recoating_speed_mm_s",
    "recoating_cycles",
    "maintenance_level",
]

if __name__ == "__main__":
    N, C = 9, 9
    lambda_rates = 1e-5 * np.ones(9)

    A = np.array([
        [1, 1, 1, 1, 0, 1, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 0, 0, 1, 1, 1, 1, 0, 0],
        [1, 0, 0, 1, 1, 1, 0, 0, 0],
        [0, 0, 0, 0, 1, 0, 1, 1, 1],
        [0, 0, 0, 0, 0, 0, 1, 1, 1],
        [0, 0, 0, 0, 0, 0, 1, 1, 1],
        ])   # fully connected; use build_adjacency_matrix() to enforce graph structure
    print("Adjacency mask A:")
    print(A)

    model = DegradationModel(N=N, C=C, lambda_rates=lambda_rates,
                             adjacency_mask=A, seed=42)
    print(f"\ntheta_I (masked, forbidden positions = 0):")
    print(model.theta_I.round(5))

    rng = np.random.default_rng(7)
    dataset: List[Sample] = Sample.from_json("samples.json")

    s = dataset[0]
    det = model.simulate(s.X, s.tau, n_steps=50, stochastic=False)
    sto = model.simulate(s.X, s.tau, n_steps=50, stochastic=True, seed=0)

    print("Deterministic simulation:")
    print(f"  y_hat = {det[-1]}")
    print(f"  loss  = {model.loss(det[-1], s.y, s.M):.4f}")
    print("Stochastic simulation (lambda_rates =", lambda_rates, "):")
    print(f"  y_hat = {sto[-1]}")
    print(f"  loss  = {model.loss(sto[-1], s.y, s.M):.4f}")


    model.summary(component_names=COMPONENT_NAMES)
    print("\nFitting:")
    learning_rates = [1e-5, 2e-5, 1e-4, 2e-4]
    min_loss = 10
    best_lr = 0
    for lr in learning_rates:
        import time
        model = DegradationModel(N=N, C=C, lambda_rates=lambda_rates,
                             adjacency_mask=A, seed=int(time.time()))
        final_loss = model.fit(dataset, lr=1e-3, epochs=300, lambda_l1=0, n_steps=20, verbose=True)[-1]
        if final_loss < min_loss:
            model.save("model.npz")
            min_loss = final_loss
            best_lr = lr
    print(best_lr, min_loss)
    model.summary(component_names=COMPONENT_NAMES)