#!/usr/bin/env python3
"""
helix_cdc/training/plan_critic.py

Plan Critic model for evaluating plans and returning risk flags.
WO-ECHO-MIRROR-TRAIN-01: Phase 2 - Critic Model

Architecture:
    - Plan encoder: Linear projection of plan embeddings
    - Evidence encoder: Linear projection of evidence embeddings
    - Classifier: MLP that outputs PASS/AMBER/FAIL probabilities
    - Risk head: Multi-label classifier for risk flags

Input:
    - Plan text (embedded to fixed dimension)
    - Top-K evidence chunks (embedded and pooled)

Output:
    - Label: PASS (0), AMBER (1), FAIL (2)
    - Risk flags: 10 binary flags
    - Confidence: softmax probability of predicted label
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False


# Risk flag definitions
RISK_FLAG_NAMES = [
    "missing_receipt",
    "missing_prerequisite",
    "low_cosine_similarity",
    "accuracy_below_threshold",
    "gate_failing",
    "gqa_not_faithful",
    "training_divergence",
    "hardware_failure",
    "vram_exceeded",
    "regression_risk",
]

# Label mapping
# WO-ECHO-MEMORY-PIPELINE-03: Added UNKNOWN for unparseable receipts
# Model outputs 3 classes (PASS/AMBER/FAIL). UNKNOWN is filtered during training.
LABEL_NAMES = ["PASS", "AMBER", "FAIL", "UNKNOWN"]
TRAINABLE_LABELS = ["PASS", "AMBER", "FAIL"]  # Labels used for model training
NUM_TRAINABLE_LABELS = len(TRAINABLE_LABELS)

LABEL_TO_IDX = {name: idx for idx, name in enumerate(TRAINABLE_LABELS)}
IDX_TO_LABEL = {idx: name for idx, name in enumerate(TRAINABLE_LABELS)}


@dataclass
class CriticOutput:
    """Output from Plan Critic model."""
    label: str  # PASS / AMBER / FAIL
    label_idx: int
    confidence: float
    label_probs: List[float]  # [P(PASS), P(AMBER), P(FAIL)]
    risk_flags: List[str]
    risk_probs: List[float]  # Probabilities for each risk flag
    raw_logits: Optional[List[float]] = None
    raw_risk_logits: Optional[List[float]] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "label": self.label,
            "label_idx": self.label_idx,
            "confidence": self.confidence,
            "label_probs": self.label_probs,
            "risk_flags": self.risk_flags,
            "risk_probs": self.risk_probs,
        }


if TORCH_AVAILABLE:

    class PlanCritic(nn.Module):
        """
        Plan Critic model: MLP classifier over plan and evidence embeddings.

        Architecture:
            plan_embed (embed_dim) --> plan_encoder (hidden_dim) --+
                                                                   |--> concat --> classifier --> 3 classes
            evidence_embed (embed_dim) --> evidence_encoder (hidden_dim) --+
                                                                           |
                                                                           +--> risk_head --> 10 flags

        Args:
            embed_dim: Input embedding dimension (default: 384 for sentence-transformers)
            hidden_dim: Hidden layer dimension
            num_labels: Number of output labels (default: 3 for PASS/AMBER/FAIL)
            num_risk_flags: Number of risk flags to predict
            dropout: Dropout rate
        """

        def __init__(
            self,
            embed_dim: int = 384,
            hidden_dim: int = 256,
            num_labels: int = 3,
            num_risk_flags: int = 10,
            dropout: float = 0.1,
        ):
            super().__init__()

            self.embed_dim = embed_dim
            self.hidden_dim = hidden_dim
            self.num_labels = num_labels
            self.num_risk_flags = num_risk_flags

            # Plan encoder
            self.plan_encoder = nn.Sequential(
                nn.Linear(embed_dim, hidden_dim),
                nn.LayerNorm(hidden_dim),
                nn.ReLU(),
                nn.Dropout(dropout),
            )

            # Evidence encoder
            self.evidence_encoder = nn.Sequential(
                nn.Linear(embed_dim, hidden_dim),
                nn.LayerNorm(hidden_dim),
                nn.ReLU(),
                nn.Dropout(dropout),
            )

            # Classifier: takes concatenated plan + evidence features
            self.classifier = nn.Sequential(
                nn.Linear(hidden_dim * 2, hidden_dim),
                nn.ReLU(),
                nn.Dropout(dropout),
                nn.Linear(hidden_dim, num_labels),
            )

            # Risk head: multi-label classification
            self.risk_head = nn.Sequential(
                nn.Linear(hidden_dim * 2, hidden_dim),
                nn.ReLU(),
                nn.Dropout(dropout),
                nn.Linear(hidden_dim, num_risk_flags),
            )

            # Initialize weights
            self._init_weights()

        def _init_weights(self):
            """Initialize weights with Xavier uniform."""
            for module in self.modules():
                if isinstance(module, nn.Linear):
                    nn.init.xavier_uniform_(module.weight)
                    if module.bias is not None:
                        nn.init.zeros_(module.bias)

        def forward(
            self,
            plan_embed: torch.Tensor,
            evidence_embed: torch.Tensor,
            labels: Optional[torch.Tensor] = None,
            risk_labels: Optional[torch.Tensor] = None,
        ) -> Dict[str, torch.Tensor]:
            """
            Forward pass.

            Args:
                plan_embed: Plan embedding [batch, embed_dim]
                evidence_embed: Pooled evidence embedding [batch, embed_dim]
                labels: Optional ground truth labels [batch] for training
                risk_labels: Optional ground truth risk flags [batch, num_risk_flags]

            Returns:
                Dict with:
                    - logits: Classification logits [batch, num_labels]
                    - risk_logits: Risk flag logits [batch, num_risk_flags]
                    - loss: Total loss (if labels provided)
                    - label_loss: Cross-entropy loss (if labels provided)
                    - risk_loss: BCE loss (if risk_labels provided)
            """
            # Encode plan and evidence
            plan_features = self.plan_encoder(plan_embed)  # [batch, hidden]
            evidence_features = self.evidence_encoder(evidence_embed)  # [batch, hidden]

            # Concatenate
            combined = torch.cat([plan_features, evidence_features], dim=-1)  # [batch, hidden*2]

            # Classify
            logits = self.classifier(combined)  # [batch, num_labels]
            risk_logits = self.risk_head(combined)  # [batch, num_risk_flags]

            output = {
                "logits": logits,
                "risk_logits": risk_logits,
            }

            # Compute losses if labels provided
            if labels is not None:
                label_loss = F.cross_entropy(logits, labels)
                output["label_loss"] = label_loss
                output["loss"] = label_loss

            if risk_labels is not None:
                risk_loss = F.binary_cross_entropy_with_logits(risk_logits, risk_labels.float())
                output["risk_loss"] = risk_loss
                if "loss" in output:
                    output["loss"] = output["loss"] + 0.5 * risk_loss  # Weight risk loss lower
                else:
                    output["loss"] = risk_loss

            return output

        def predict(
            self,
            plan_embed: torch.Tensor,
            evidence_embed: torch.Tensor,
            risk_threshold: float = 0.5,
        ) -> CriticOutput:
            """
            Make a prediction for a single sample.

            Args:
                plan_embed: Plan embedding [1, embed_dim] or [embed_dim]
                evidence_embed: Evidence embedding [1, embed_dim] or [embed_dim]
                risk_threshold: Threshold for risk flag prediction

            Returns:
                CriticOutput with prediction details
            """
            # Ensure batch dimension
            if plan_embed.dim() == 1:
                plan_embed = plan_embed.unsqueeze(0)
            if evidence_embed.dim() == 1:
                evidence_embed = evidence_embed.unsqueeze(0)

            self.eval()
            with torch.no_grad():
                output = self.forward(plan_embed, evidence_embed)

                # Get label prediction
                logits = output["logits"][0]
                probs = F.softmax(logits, dim=-1)
                label_idx = torch.argmax(probs).item()
                confidence = probs[label_idx].item()

                # Get risk flags
                risk_logits = output["risk_logits"][0]
                risk_probs = torch.sigmoid(risk_logits)
                risk_flags = [
                    RISK_FLAG_NAMES[i]
                    for i, prob in enumerate(risk_probs)
                    if prob > risk_threshold
                ]

            return CriticOutput(
                label=IDX_TO_LABEL[label_idx],
                label_idx=label_idx,
                confidence=confidence,
                label_probs=probs.tolist(),
                risk_flags=risk_flags,
                risk_probs=risk_probs.tolist(),
                raw_logits=logits.tolist(),
                raw_risk_logits=risk_logits.tolist(),
            )

        @classmethod
        def from_config(cls, config: Dict[str, Any]) -> "PlanCritic":
            """Create model from config dict."""
            return cls(
                embed_dim=config.get("embed_dim", 384),
                hidden_dim=config.get("hidden_dim", 256),
                num_labels=config.get("num_labels", 3),
                num_risk_flags=config.get("num_risk_flags", 10),
                dropout=config.get("dropout", 0.1),
            )

        def get_config(self) -> Dict[str, Any]:
            """Get model config dict."""
            return {
                "embed_dim": self.embed_dim,
                "hidden_dim": self.hidden_dim,
                "num_labels": self.num_labels,
                "num_risk_flags": self.num_risk_flags,
            }


    class SimpleTextEncoder:
        """
        Simple text encoder using mean pooling of character embeddings.
        For production, replace with sentence-transformers or similar.
        """

        def __init__(self, embed_dim: int = 384, vocab_size: int = 256):
            self.embed_dim = embed_dim
            self.vocab_size = vocab_size

            # Create random but deterministic embeddings
            torch.manual_seed(42)
            self.char_embeds = torch.randn(vocab_size, embed_dim) / (embed_dim ** 0.5)

        def encode(self, text: str) -> torch.Tensor:
            """Encode text to embedding via mean pooling of char embeddings."""
            # Convert to bytes and look up embeddings
            bytes_arr = text.encode("utf-8", errors="replace")[:1000]  # Truncate
            indices = torch.tensor([b % self.vocab_size for b in bytes_arr], dtype=torch.long)

            # Mean pool
            embeds = self.char_embeds[indices]  # [seq_len, embed_dim]
            return embeds.mean(dim=0)  # [embed_dim]

        def encode_batch(self, texts: List[str]) -> torch.Tensor:
            """Encode batch of texts."""
            return torch.stack([self.encode(t) for t in texts])


    class CriticDataset(torch.utils.data.Dataset):
        """
        Dataset for Plan Critic training.

        Loads from JSONL file with format:
        {
            "plan_text": str,
            "evidence_chunks": List[str],
            "label": "PASS" | "AMBER" | "FAIL" | "UNKNOWN",
            "risk_flags": List[str],
            ...
        }

        WO-ECHO-MEMORY-PIPELINE-03: Added filter_unknown option to exclude
        UNKNOWN labels from training (default: True).
        """

        def __init__(
            self,
            data_path: str,
            encoder: Optional[SimpleTextEncoder] = None,
            max_evidence_chunks: int = 5,
            filter_unknown: bool = True,
        ):
            import json
            from pathlib import Path

            self.encoder = encoder or SimpleTextEncoder()
            self.max_evidence_chunks = max_evidence_chunks
            self.filter_unknown = filter_unknown
            self.examples = []

            # Load data
            path = Path(data_path)
            unknown_count = 0
            with open(path) as f:
                for line in f:
                    if line.strip():
                        example = json.loads(line)
                        # Filter UNKNOWN labels if requested
                        if filter_unknown and example.get("label") == "UNKNOWN":
                            unknown_count += 1
                            continue
                        self.examples.append(example)

            if unknown_count > 0:
                print(f"CriticDataset: Filtered {unknown_count} UNKNOWN examples")

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

        def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
            example = self.examples[idx]

            # Encode plan
            plan_embed = self.encoder.encode(example["plan_text"])

            # Encode and pool evidence
            evidence_chunks = example.get("evidence_chunks", [])[:self.max_evidence_chunks]
            if evidence_chunks:
                evidence_embeds = self.encoder.encode_batch(evidence_chunks)
                evidence_embed = evidence_embeds.mean(dim=0)  # Mean pool
            else:
                evidence_embed = torch.zeros(self.encoder.embed_dim)

            # Get label
            label_str = example.get("label", "AMBER")
            # Map UNKNOWN to AMBER for training (shouldn't happen if filter_unknown=True)
            if label_str == "UNKNOWN":
                label_str = "AMBER"
            label = LABEL_TO_IDX.get(label_str, 1)  # Default to AMBER

            # Get risk flags
            risk_flags = example.get("risk_flags", [])
            risk_labels = torch.zeros(len(RISK_FLAG_NAMES))
            for flag in risk_flags:
                if flag in RISK_FLAG_NAMES:
                    risk_labels[RISK_FLAG_NAMES.index(flag)] = 1.0

            return {
                "plan_embed": plan_embed,
                "evidence_embed": evidence_embed,
                "labels": torch.tensor(label, dtype=torch.long),
                "risk_labels": risk_labels,
            }


else:
    # Stubs for when torch is not available

    class PlanCritic:
        def __init__(self, *args, **kwargs):
            raise ImportError("torch is required for PlanCritic")

    class SimpleTextEncoder:
        def __init__(self, *args, **kwargs):
            raise ImportError("torch is required for SimpleTextEncoder")

    class CriticDataset:
        def __init__(self, *args, **kwargs):
            raise ImportError("torch is required for CriticDataset")


# Export
__all__ = [
    "PlanCritic",
    "CriticOutput",
    "SimpleTextEncoder",
    "CriticDataset",
    "RISK_FLAG_NAMES",
    "LABEL_NAMES",
    "TRAINABLE_LABELS",
    "NUM_TRAINABLE_LABELS",
    "LABEL_TO_IDX",
    "IDX_TO_LABEL",
]
