"""bioai.models.pinn_fate -- Physics-Informed Neural Network for dsRNA environmental fate.

Predicts the first-order degradation rate ``k`` of a dsRNA in the field, given
8 environmental + sequence features. The physics constraint is the standard
exponential decay ``C(t) = C0 * exp(-k * t)`` so half-life is ``ln(2) / k``.

Wired into the pipeline (Task D deliverable #5): the ranker uses
``DegradationPINN.half_life(features)`` to penalise candidates whose
half-life is under 6 hours (too short to be effective in the field).

Specs:
    Input features (8-dim):
        temperature_C, pH, UV_index, GC_content, length,
        salinity_ppt, soil_clay_pct, humidity_pct
    Output: positive rate k (Softplus ensures positivity)
    Network: Linear(8,64) -> Tanh -> Linear(64,32) -> Tanh -> Linear(32,1) -> Softplus
"""

from __future__ import annotations

import math

import torch
import torch.nn as nn


# Feature names in the expected order. Used by callers to build feature vectors.
PINN_FEATURE_NAMES = (
    "temperature_C",   # ambient temperature in Celsius
    "pH",              # soil/leaf pH
    "UV_index",        # 0..15 typical daylight UV index
    "GC_content",      # 0..1 fraction G+C of the dsRNA
    "length",          # dsRNA length in nt (e.g. 200)
    "salinity_ppt",    # salinity in parts per thousand
    "soil_clay_pct",   # 0..100 soil clay percentage
    "humidity_pct",    # 0..100 relative humidity
)


class DegradationPINN(nn.Module):
    """PINN for dsRNA environmental fate.

    Forward pass returns the rate ``k`` (positive). Use :meth:`predict_rate`,
    :meth:`predict_concentration`, or :meth:`half_life` for the physics-wrapped
    quantities.
    """

    LN2 = math.log(2.0)

    def __init__(self, feature_dim: int = 8):
        super().__init__()
        self.feature_dim = feature_dim
        self.net = nn.Sequential(
            nn.Linear(feature_dim, 64), nn.Tanh(),
            nn.Linear(64, 32), nn.Tanh(),
            nn.Linear(32, 1), nn.Softplus(),  # ensures k > 0
        )

    # ------------------------------------------------------------------ #
    def forward(self, features: torch.Tensor) -> torch.Tensor:
        """Raw forward returns the rate ``k`` (positive). Shape ``(B, 1)``."""
        if features.dim() == 1:
            features = features.unsqueeze(0)
        return self.net(features)

    # ------------------------------------------------------------------ #
    def predict_rate(self, features: torch.Tensor) -> torch.Tensor:
        """Positive degradation rate ``k`` (1/hours). Shape ``(B, 1)``."""
        return self.forward(features)

    # ------------------------------------------------------------------ #
    def predict_concentration(
        self,
        C0: torch.Tensor | float,
        t_points: torch.Tensor,
        features: torch.Tensor,
    ) -> torch.Tensor:
        """Physics-informed concentration trajectory.

        Parameters
        ----------
        C0:
            Initial concentration (scalar or ``(B,)`` or ``(B, 1)``).
        t_points:
            1-D tensor of time points (hours), shape ``(T,)``.
        features:
            ``(B, feature_dim)`` env/sequence features.

        Returns
        -------
        torch.Tensor
            Concentration ``C(t) = C0 * exp(-k * t)`` for each sample and
            each time point. Shape ``(B, T)``.
        """
        k = self.predict_rate(features)                 # (B, 1)
        if not torch.is_tensor(C0):
            C0 = torch.tensor(float(C0), dtype=k.dtype, device=k.device)
        C0 = C0.view(-1, 1) if C0.dim() == 1 else C0.view(-1, 1)
        # Broadcast: (B, 1) * (1, T) -> (B, T)
        t = t_points.to(k.device).unsqueeze(0)
        return C0 * torch.exp(-k * t)

    # ------------------------------------------------------------------ #
    def half_life(self, features: torch.Tensor) -> torch.Tensor:
        """Half-life (hours) = ln(2) / k. Shape ``(B, 1)``.

        Softplus guarantees ``k > 0`` so the division is well-defined.
        """
        k = self.predict_rate(features)
        # Clamp k from below to avoid Inf half-lives from a near-zero rate.
        k_safe = torch.clamp(k, min=1e-6)
        return torch.tensor(self.LN2, dtype=k.dtype, device=k.device) / k_safe

    # ------------------------------------------------------------------ #
    def physics_consistency_loss(
        self,
        C0: torch.Tensor | float,
        t_points: torch.Tensor,
        features: torch.Tensor,
        C_observed: torch.Tensor,
    ) -> torch.Tensor:
        """MSE between the PINN-predicted concentration and observed C(t).

        Use this as the "physics-consistency" loss term in training. Combined
        with a direct MSE on ``k`` (when rate labels are available), it
        enforces that the network both fits the data *and* respects the
        exponential-decay ODE.
        """
        C_pred = self.predict_concentration(C0, t_points, features)
        return nn.functional.mse_loss(C_pred, C_observed)