ml/models/geometry/sequence_vae.py

"""
ml/models/geometry/sequence_vae.py
1D-CNN Variational Autoencoder (VAE) for Novel Geometry Anomaly Detection (tc.v1 SOTA).
Features:
  1. Dilated 1D-CNN Encoder (captures 1-3m, 3-10m, and 10-30m wavelengths)
  2. Dual-Path Anomaly Scoring (Reconstruction MSE + Latent Mahalanobis Distance)
  3. Sigmoid Distance Calibration for [0.0, 1.0] confidence mapping
"""

import os
from pathlib import Path
from typing import Dict, Tuple, Optional, Union, Any, List
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from ml.core.schema import CalibratedSignal, SignalType, DefectClass
from ml.core.registry import register_model
from ml.calibration.patchcore_scale import SigmoidDistanceCalibrator
from ml.models.geometry.sequence_vae_enhanced import EnhancedSequenceVAE


class DilatedEncoder(nn.Module):
    """
    Multi-Scale Dilated 1D-CNN Encoder capturing short (1-3m), medium (3-10m),
    and long (10-30m) spatial track wavelengths simultaneously.
    """

    def __init__(self, n_features: int = 5, seq_len: int = 80, latent_dim: int = 16):
        super().__init__()
        self.n_features = n_features
        self.seq_len = seq_len
        self.latent_dim = latent_dim

        # Parallel multi-scale dilated convolutions
        self.branch_short = nn.Conv1d(n_features, 32, kernel_size=3, dilation=1, padding=1)   # 1-3m wavelengths
        self.branch_medium = nn.Conv1d(n_features, 32, kernel_size=3, dilation=4, padding=4)  # 3-10m wavelengths
        self.branch_long = nn.Conv1d(n_features, 32, kernel_size=3, dilation=10, padding=10) # 10-30m wavelengths

        self.fusion = nn.Sequential(
            nn.Conv1d(96, 64, kernel_size=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool1d(20),  # Downsample sequence to 20 spatial bins
        )

        self.flat_size = 64 * 20  # 1280

        # Latent Space Projections
        self.fc_mu = nn.Linear(self.flat_size, latent_dim)
        self.fc_logvar = nn.Linear(self.flat_size, latent_dim)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        # x: [batch, n_features, seq_len]
        short = F.relu(self.branch_short(x))
        medium = F.relu(self.branch_medium(x))
        long = F.relu(self.branch_long(x))

        combined = torch.cat([short, medium, long], dim=1)  # [B, 96, seq_len]
        h = self.fusion(combined).flatten(1)               # [B, 1280]

        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar


class SequenceVAE(nn.Module):
    """
    1D-CNN Variational Autoencoder with Dilated Multi-Scale Encoder
    and Transposed Conv1D Decoder for track geometry waveform reconstruction.
    """

    def __init__(
        self,
        seq_len: int = 80,
        n_features: int = 5,
        latent_dim: int = 16,
    ):
        super().__init__()
        self.seq_len = seq_len
        self.n_features = n_features
        self.latent_dim = latent_dim

        # Multi-scale dilated encoder
        self.encoder = DilatedEncoder(n_features=n_features, seq_len=seq_len, latent_dim=latent_dim)
        self.flat_size = self.encoder.flat_size

        # Decoder: Reconstructs 5-channel sequence back to 80 spatial bins
        self.decoder_input = nn.Linear(latent_dim, self.flat_size)
        self.decoder = nn.Sequential(
            nn.Unflatten(1, (64, 20)),
            nn.ConvTranspose1d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1),  # 20 -> 40
            nn.ReLU(),
            nn.ConvTranspose1d(32, n_features, kernel_size=3, stride=2, padding=1, output_padding=1),  # 40 -> 80
        )

    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if x.dim() == 2:
            x = x.unsqueeze(0)

        # Transpose [batch, seq_len, features] -> [batch, features, seq_len] for Conv1d
        x_trans = x.transpose(1, 2)
        mu, logvar = self.encoder(x_trans)
        z = self.reparameterize(mu, logvar)

        h = F.relu(self.decoder_input(z))
        recon_trans = self.decoder(h)
        recon_x = recon_trans.transpose(1, 2)  # [batch, seq_len, features]

        return recon_x, mu, logvar

    def loss_function(
        self,
        recon_x: torch.Tensor,
        x: torch.Tensor,
        mu: torch.Tensor,
        logvar: torch.Tensor,
        beta: float = 0.01,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if recon_x.shape != x.shape:
            min_len = min(recon_x.shape[1], x.shape[1])
            recon_x = recon_x[:, :min_len, :]
            x = x[:, :min_len, :]

        recon_loss = F.mse_loss(recon_x, x, reduction="mean")
        kld_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
        total_loss = recon_loss + beta * kld_loss
        return total_loss, recon_loss, kld_loss

    def compute_reconstruction_error(self, x: torch.Tensor) -> torch.Tensor:
        self.eval()
        with torch.no_grad():
            recon_x, _, _ = self.forward(x)
            if recon_x.shape != x.shape:
                min_len = min(recon_x.shape[1], x.shape[1])
                recon_x = recon_x[:, :min_len, :]
                x = x[:, :min_len, :]
            error = torch.mean((recon_x - x) ** 2, dim=(1, 2))
        return error

    def compute_anomaly_score(self, x: torch.Tensor) -> torch.Tensor:
        """Alias for compute_reconstruction_error."""
        return self.compute_reconstruction_error(x)


@register_model("geometry_sequence_vae")
class SequenceVAEDetector:
    """
    Production detector wrapper for SequenceVAE featuring:
      - Enhanced multi-scale dilated convolutions matching trained SOTA checkpoint
      - Dual-Path Anomaly Scoring (Reconstruction MSE + Latent Mahalanobis Distance)
      - EVT & Sigmoid Calibration for calibrated [0.0, 1.0] novelty probabilities
    """

    def __init__(
        self,
        weights_path: Optional[Union[str, Path]] = "artifacts/checkpoints/geometry/sequence_vae_enhanced.pt",
        calibrator_path: Optional[Union[str, Path]] = "artifacts/calibration/vae_calibration.json",
        device: str = "cpu",
        seq_len: int = 80,
        n_features: int = 5,
        latent_dim: int = 16,
        alpha: float = 0.7,
        threshold: float = 0.50,
    ):
        self.device = torch.device(device if torch.cuda.is_available() or device == "cpu" else "cpu")
        self.seq_len = seq_len
        self.n_features = n_features
        self.latent_dim = latent_dim
        self.alpha = float(alpha)
        self.threshold = threshold

        # Default model is EnhancedSequenceVAE matching trained SOTA checkpoint ([128, 192, 1] fusion)
        self.model: Union[EnhancedSequenceVAE, SequenceVAE] = EnhancedSequenceVAE(
            seq_len=seq_len, n_features=n_features, latent_dim=latent_dim
        ).to(self.device)

        # Resolve candidate weight paths in order of preference
        candidate_weights = [weights_path] if weights_path else []
        candidate_weights.extend([
            "artifacts/checkpoints/geometry/sequence_vae_enhanced.pt",
            "artifacts/checkpoints/geometry/sequence_vae.pt",
            "ml/models/geometry/weights/sequence_vae.pt",
            "ml/models/geometry/weights/sequence_vae_enhanced.pt",
        ])

        loaded_checkpoint = False
        for wp in candidate_weights:
            if wp and os.path.exists(str(wp)):
                try:
                    ckpt = torch.load(wp, map_location=self.device)
                    state = ckpt["model_state_dict"] if isinstance(ckpt, dict) and "model_state_dict" in ckpt else ckpt
                    # Auto-detect architecture from state dict
                    if isinstance(state, dict) and "encoder.fusion.0.weight" in state:
                        in_channels = state["encoder.fusion.0.weight"].shape[1]
                        if in_channels == 96:
                            self.model = SequenceVAE(seq_len=seq_len, n_features=n_features, latent_dim=latent_dim).to(self.device)
                        else:
                            self.model = EnhancedSequenceVAE(seq_len=seq_len, n_features=n_features, latent_dim=latent_dim).to(self.device)
                    self.model.load_state_dict(state, strict=True)
                    loaded_checkpoint = True
                    break
                except Exception as e:
                    if weights_path is not None and str(wp) == str(weights_path):
                        raise RuntimeError(
                            f"VAE weight load FAILED for {weights_path} — architecture mismatch. "
                            f"Do NOT fall back to a random model. Error: {e}"
                        )

        self.model.eval()

        # Latent distribution baseline for Mahalanobis distance
        self.normal_latent_mean: Optional[np.ndarray] = None
        self.normal_latent_cov_inv: Optional[np.ndarray] = None

        # Sigmoid / EVT threshold calibrator
        self.calibrator = SigmoidDistanceCalibrator(threshold=1.65, steepness_k=2.0)
        candidate_calibs = [calibrator_path] if calibrator_path else []
        candidate_calibs.extend([
            "artifacts/calibration/vae_calibration.json",
            "artifacts/calibration/sequence_vae_calibration.json",
        ])
        loaded_calib = False
        for cp in candidate_calibs:
            if cp and os.path.exists(str(cp)):
                try:
                    with open(cp, "r", encoding="utf-8") as f:
                        cdata = json.load(f)
                    thresh = float(cdata.get("threshold_evt", cdata.get("threshold_p99", 1.65)))
                    k_val = float(cdata.get("steepness_k", cdata.get("steepness", 2.0)))
                    self.calibrator = SigmoidDistanceCalibrator(threshold=thresh, steepness_k=k_val)
                    if isinstance(self.model, EnhancedSequenceVAE):
                        self.model.threshold_evt = thresh
                        self.model.steepness_k = k_val
                    loaded_calib = True
                    break
                except Exception:
                    pass

        if not loaded_calib:
            # Auto-initialize baseline distribution on nominal noise
            synth_baseline = np.random.normal(0.0, 0.1, (30, self.seq_len, self.n_features)).astype(np.float32)
            self.fit_latent_distribution(synth_baseline)
            synth_scores = [self.compute_anomaly_score(s) for s in synth_baseline]
            self.calibrator.fit(synth_scores, percentile=99.0)
            self.calibrator.threshold = max(self.calibrator.threshold, 0.5)

    def _format_input(self, geometry_window: Union[np.ndarray, Dict[str, np.ndarray], torch.Tensor]) -> torch.Tensor:
        if isinstance(geometry_window, dict):
            keys = [
                ("twist_3m", "twist_3m_mm"),
                ("versine_10m", "versine_10m_mm"),
                ("versine_20m", "versine_20m_mm"),
                ("unevenness_10m", "unevenness_10m_mm"),
                ("cant", "cant_mm"),
            ]
            cols = []
            for k1, k2 in keys:
                if k1 in geometry_window:
                    cols.append(np.asarray(geometry_window[k1]))
                elif k2 in geometry_window:
                    cols.append(np.asarray(geometry_window[k2]))
                else:
                    cols.append(np.zeros(self.seq_len))
            arr = np.column_stack(cols)
        elif isinstance(geometry_window, torch.Tensor):
            arr = geometry_window.cpu().numpy()
        else:
            arr = np.asarray(geometry_window)

        if arr.ndim == 2:
            arr = np.expand_dims(arr, axis=0)  # [1, seq_len, 5]

        b, t, c = arr.shape
        if t != self.seq_len:
            fixed = np.zeros((b, self.seq_len, c), dtype=np.float32)
            copy_len = min(t, self.seq_len)
            fixed[:, :copy_len, :] = arr[:, :copy_len, :]
            arr = fixed

        return torch.tensor(arr, dtype=torch.float32).to(self.device)

    def fit_latent_distribution(self, normal_sequences: Union[np.ndarray, torch.Tensor, List[np.ndarray]]):
        """Fits baseline latent Gaussian distribution (mean and regularized inverse covariance)."""
        if isinstance(self.model, EnhancedSequenceVAE):
            if isinstance(normal_sequences, list):
                normal_tensor = torch.stack([torch.tensor(s, dtype=torch.float32) for s in normal_sequences])
            elif isinstance(normal_sequences, np.ndarray):
                normal_tensor = torch.tensor(normal_sequences, dtype=torch.float32)
            else:
                normal_tensor = normal_sequences
            self.model.fit_latent_distribution(normal_tensor)
            self.normal_latent_mean = self.model.latent_mean
            self.normal_latent_cov_inv = self.model.latent_cov_inv
            return

        self.model.eval()
        latents = []

        if isinstance(normal_sequences, list):
            normal_tensor = torch.stack([torch.tensor(s, dtype=torch.float32) for s in normal_sequences])
        elif isinstance(normal_sequences, np.ndarray):
            normal_tensor = torch.tensor(normal_sequences, dtype=torch.float32)
        else:
            normal_tensor = normal_sequences

        if normal_tensor.ndim == 2:
            normal_tensor = normal_tensor.unsqueeze(0)

        with torch.no_grad():
            normal_tensor = normal_tensor.to(self.device)
            _, mu, _ = self.model(normal_tensor)
            latents = mu.cpu().numpy()

        self.normal_latent_mean = np.mean(latents, axis=0)
        cov = np.cov(latents.T)
        if cov.ndim == 0:
            cov = np.array([[cov]])
        cov_reg = cov + np.eye(cov.shape[0]) * 1e-6
        self.normal_latent_cov_inv = np.linalg.inv(cov_reg)

    def compute_anomaly_score(self, sequence: Union[np.ndarray, Dict[str, np.ndarray], torch.Tensor]) -> float:
        """
        Computes Dual-Path Anomaly Score:
            Score = alpha * MSE_recon + (1 - alpha) * D_mahalanobis
        """
        if isinstance(self.model, EnhancedSequenceVAE):
            if isinstance(sequence, dict):
                tensor_in = self._format_input(sequence)
                res = self.model.compute_anomaly_score(tensor_in)
            else:
                res = self.model.compute_anomaly_score(sequence)
            return float(res.get("combined_score", res.get("recon_error", 0.0)))

        self.model.eval()
        tensor_in = self._format_input(sequence)

        with torch.no_grad():
            recon_x, mu, _ = self.model(tensor_in)
            mse = float(F.mse_loss(recon_x, tensor_in).item())

            # Path 2: Latent Mahalanobis Distance
            if self.normal_latent_mean is not None and self.normal_latent_cov_inv is not None:
                mu_np = mu.cpu().numpy()[0]
                diff = mu_np - self.normal_latent_mean
                mahal_dist = float(np.sqrt(np.clip(diff @ self.normal_latent_cov_inv @ diff.T, 0.0, None)))
            else:
                mahal_dist = float(torch.norm(mu).item())

        combined_score = float((self.alpha * mse) + ((1.0 - self.alpha) * mahal_dist))
        return combined_score

    def fit_calibration(self, normal_windows: Union[np.ndarray, torch.Tensor], percentile: float = 99.0):
        """Fits latent distribution and Sigmoid distance calibrator on nominal track validation sequences."""
        self.fit_latent_distribution(normal_windows)

        scores = []
        if isinstance(normal_windows, torch.Tensor):
            arr = normal_windows.cpu().numpy()
        else:
            arr = np.asarray(normal_windows)

        for seq in arr:
            scores.append(self.compute_anomaly_score(seq))

        self.calibrator.fit(scores, percentile=percentile)
        return self.calibrator.threshold

    def predict(self, geometry_window: Union[np.ndarray, Dict[str, np.ndarray], torch.Tensor]) -> CalibratedSignal:
        """
        Runs the 1D-CNN Dual-Path VAE on the 20m geometry window.
        Returns a calibrated CalibratedSignal indicating structural novelty.
        """
        if isinstance(self.model, EnhancedSequenceVAE):
            if isinstance(geometry_window, dict):
                tensor_in = self._format_input(geometry_window)
                pred = self.model.predict(tensor_in)
            else:
                pred = self.model.predict(geometry_window)
            raw_score = float(pred["raw_score"])
            calibrated_prob = float(pred["calibrated_prob"])
            is_anomaly = bool(pred["is_anomaly"])
            thresh = float(pred.get("threshold", self.threshold))
        else:
            raw_score = self.compute_anomaly_score(geometry_window)
            calibrated_prob = float(self.calibrator.scale(raw_score))
            is_anomaly = bool(calibrated_prob >= self.threshold)
            thresh = self.threshold

        explanation = {
            "combined_anomaly_score": round(raw_score, 4),
            "reconstruction_error": round(raw_score, 4),
            "calibrated_prob": round(calibrated_prob, 4),
            "p99_threshold": round(self.calibrator.threshold, 4),
            "anomaly_detected": is_anomaly,
        }

        return CalibratedSignal(
            name="sequence_vae_anomaly",
            stream_name="geometry_vae",
            model_version="0.1.0",
            signal_type=SignalType.GEOMETRY_NOVEL,
            value=calibrated_prob,
            raw_score=raw_score,
            calibrated_prob=calibrated_prob,
            threshold=thresh,
            fired=is_anomaly,
            is_anomaly=is_anomaly,
            predicted_class=DefectClass.GEOMETRY_ANOMALY if is_anomaly else DefectClass.NORMAL,
            bbox=None,
            explanation=explanation,
            metadata={
                "model_name": "sequence_vae_dual_path",
                "latent_dim": self.model.latent_dim,
                "combined_score": raw_score,
                "calibrator_threshold": self.calibrator.threshold,
            },
        )