"""
Personality Probes for Interpretability Research

Linear and non-linear probes for predicting personality dimensions from activations.
Includes training utilities and analysis tools.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

from .collector import ActivationDataset


# Personality dimensions in canonical order
DIMENSIONS = [
    "openness",
    "conscientiousness",
    "extraversion",
    "agreeableness",
    "neuroticism",
    "assertiveness",
    "ambition",
    "resilience",
    "integrity",
    "curiosity",
]


class PersonalityProbe(nn.Module):
    """
    Probe for predicting personality dimensions from activations.

    Supports linear and MLP variants. Linear probes are preferred for
    interpretability (weights directly indicate feature importance).
    """

    def __init__(
        self,
        hidden_dim: int,
        num_dimensions: int = 10,
        probe_type: Literal["linear", "mlp"] = "linear",
        mlp_hidden: int = 256,
        dropout: float = 0.1,
    ):
        super().__init__()

        self.hidden_dim = hidden_dim
        self.num_dimensions = num_dimensions
        self.probe_type = probe_type

        if probe_type == "linear":
            self.probe = nn.Linear(hidden_dim, num_dimensions)
        else:
            self.probe = nn.Sequential(
                nn.Linear(hidden_dim, mlp_hidden),
                nn.ReLU(),
                nn.Dropout(dropout),
                nn.Linear(mlp_hidden, mlp_hidden // 2),
                nn.ReLU(),
                nn.Dropout(dropout),
                nn.Linear(mlp_hidden // 2, num_dimensions),
            )

    def forward(self, activations: torch.Tensor) -> torch.Tensor:
        """
        Predict personality scores from activations.

        Args:
            activations: [batch, hidden_dim] or [batch, seq_len, hidden_dim]

        Returns:
            Personality scores [batch, num_dimensions] in [0, 1]
        """
        # Handle sequence dimension
        if activations.dim() == 3:
            activations = activations.mean(dim=1)  # Mean pool over sequence

        logits = self.probe(activations)
        return torch.sigmoid(logits)

    def get_weights(self) -> torch.Tensor:
        """
        Get probe weights for interpretability analysis.

        Returns:
            Weight matrix [num_dimensions, hidden_dim] for linear probe
        """
        if self.probe_type != "linear":
            raise ValueError("Weight extraction only supported for linear probes")
        return self.probe.weight.detach()

    def get_dimension_weights(self, dimension: str | int) -> torch.Tensor:
        """
        Get weights for a specific personality dimension.

        Args:
            dimension: Dimension name or index

        Returns:
            Weight vector [hidden_dim]
        """
        if isinstance(dimension, str):
            dimension = DIMENSIONS.index(dimension)

        weights = self.get_weights()
        return weights[dimension]


@dataclass
class ProbeMetrics:
    """Metrics for probe evaluation."""

    mse: float
    mae: float
    per_dimension_r2: dict[str, float]
    per_dimension_correlation: dict[str, float]
    overall_r2: float
    overall_correlation: float

    def to_dict(self) -> dict[str, Any]:
        return {
            "mse": self.mse,
            "mae": self.mae,
            "per_dimension_r2": self.per_dimension_r2,
            "per_dimension_correlation": self.per_dimension_correlation,
            "overall_r2": self.overall_r2,
            "overall_correlation": self.overall_correlation,
        }


@dataclass
class ProbeTrainerConfig:
    """Configuration for probe training."""

    learning_rate: float = 1e-3
    weight_decay: float = 1e-4
    batch_size: int = 32
    num_epochs: int = 100
    early_stopping_patience: int = 10
    use_confidence_weighting: bool = True
    device: str = "cuda" if torch.cuda.is_available() else "cpu"


class ProbeTrainer:
    """
    Trainer for personality probes.

    Handles training loop, evaluation, and analysis utilities.
    """

    def __init__(
        self,
        probe: PersonalityProbe,
        config: ProbeTrainerConfig | None = None,
    ):
        self.probe = probe
        self.config = config or ProbeTrainerConfig()
        self.probe.to(self.config.device)

        self.optimizer = torch.optim.AdamW(
            self.probe.parameters(),
            lr=self.config.learning_rate,
            weight_decay=self.config.weight_decay,
        )
        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer, mode="min", factor=0.5, patience=5
        )

        self.train_losses: list[float] = []
        self.val_losses: list[float] = []

    def train(
        self,
        train_dataset: ActivationDataset,
        val_dataset: ActivationDataset | None = None,
    ) -> ProbeMetrics:
        """
        Train the probe on activation dataset.

        Args:
            train_dataset: Training data
            val_dataset: Optional validation data for early stopping

        Returns:
            Final metrics on validation set (or training if no val)
        """
        train_loader = DataLoader(
            train_dataset,
            batch_size=self.config.batch_size,
            shuffle=True,
            num_workers=0,
        )

        val_loader = None
        if val_dataset:
            val_loader = DataLoader(
                val_dataset,
                batch_size=self.config.batch_size,
                shuffle=False,
                num_workers=0,
            )

        best_val_loss = float("inf")
        patience_counter = 0
        best_state = None

        for epoch in range(self.config.num_epochs):
            # Training
            self.probe.train()
            train_loss = 0.0

            for batch in train_loader:
                activations = batch["activation"].to(self.config.device)
                labels = batch["labels"].to(self.config.device)
                confidence = batch["confidence"].to(self.config.device)

                self.optimizer.zero_grad()

                predictions = self.probe(activations)
                loss = self._compute_loss(predictions, labels, confidence)

                loss.backward()
                self.optimizer.step()

                train_loss += loss.item()

            train_loss /= len(train_loader)
            self.train_losses.append(train_loss)

            # Validation
            if val_loader:
                val_loss = self._validate(val_loader)
                self.val_losses.append(val_loss)

                self.scheduler.step(val_loss)

                # Early stopping
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    patience_counter = 0
                    best_state = {k: v.cpu().clone() for k, v in self.probe.state_dict().items()}
                else:
                    patience_counter += 1

                if patience_counter >= self.config.early_stopping_patience:
                    print(f"Early stopping at epoch {epoch}")
                    break

            if epoch % 10 == 0:
                val_str = f", val_loss={val_loss:.4f}" if val_loader else ""
                print(f"Epoch {epoch}: train_loss={train_loss:.4f}{val_str}")

        # Restore best model
        if best_state:
            self.probe.load_state_dict(best_state)
            self.probe.to(self.config.device)

        # Final evaluation
        eval_loader = val_loader or train_loader
        return self.evaluate(eval_loader)

    def _compute_loss(
        self,
        predictions: torch.Tensor,
        labels: torch.Tensor,
        confidence: torch.Tensor,
    ) -> torch.Tensor:
        """Compute MSE loss with optional confidence weighting."""
        mse = F.mse_loss(predictions, labels, reduction="none")
        mse = mse.mean(dim=1)  # Mean over dimensions

        if self.config.use_confidence_weighting:
            mse = mse * confidence

        return mse.mean()

    @torch.no_grad()
    def _validate(self, val_loader: DataLoader) -> float:
        """Compute validation loss."""
        self.probe.eval()
        total_loss = 0.0

        for batch in val_loader:
            activations = batch["activation"].to(self.config.device)
            labels = batch["labels"].to(self.config.device)
            confidence = batch["confidence"].to(self.config.device)

            predictions = self.probe(activations)
            loss = self._compute_loss(predictions, labels, confidence)
            total_loss += loss.item()

        return total_loss / len(val_loader)

    @torch.no_grad()
    def evaluate(self, loader: DataLoader) -> ProbeMetrics:
        """
        Evaluate probe and compute detailed metrics.

        Returns:
            ProbeMetrics with per-dimension and overall statistics
        """
        self.probe.eval()

        all_predictions = []
        all_labels = []

        for batch in loader:
            activations = batch["activation"].to(self.config.device)
            labels = batch["labels"]

            predictions = self.probe(activations).cpu()
            all_predictions.append(predictions)
            all_labels.append(labels)

        predictions = torch.cat(all_predictions, dim=0)
        labels = torch.cat(all_labels, dim=0)

        # Overall metrics
        mse = F.mse_loss(predictions, labels).item()
        mae = F.l1_loss(predictions, labels).item()

        # Per-dimension metrics
        per_dim_r2 = {}
        per_dim_corr = {}

        for i, dim in enumerate(DIMENSIONS):
            pred_dim = predictions[:, i]
            label_dim = labels[:, i]

            # R² score
            ss_res = ((pred_dim - label_dim) ** 2).sum()
            ss_tot = ((label_dim - label_dim.mean()) ** 2).sum()
            r2 = 1 - (ss_res / (ss_tot + 1e-8))
            per_dim_r2[dim] = r2.item()

            # Pearson correlation
            pred_centered = pred_dim - pred_dim.mean()
            label_centered = label_dim - label_dim.mean()
            corr = (pred_centered * label_centered).sum() / (
                (pred_centered.norm() * label_centered.norm()) + 1e-8
            )
            per_dim_corr[dim] = corr.item()

        overall_r2 = sum(per_dim_r2.values()) / len(per_dim_r2)
        overall_corr = sum(per_dim_corr.values()) / len(per_dim_corr)

        return ProbeMetrics(
            mse=mse,
            mae=mae,
            per_dimension_r2=per_dim_r2,
            per_dimension_correlation=per_dim_corr,
            overall_r2=overall_r2,
            overall_correlation=overall_corr,
        )

    def save(self, path: Path) -> None:
        """Save probe weights and training history."""
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)

        torch.save(self.probe.state_dict(), path / "probe.pt")

        config = {
            "hidden_dim": self.probe.hidden_dim,
            "num_dimensions": self.probe.num_dimensions,
            "probe_type": self.probe.probe_type,
            "train_losses": self.train_losses,
            "val_losses": self.val_losses,
        }
        with open(path / "config.json", "w") as f:
            json.dump(config, f, indent=2)

    @classmethod
    def load(cls, path: Path) -> tuple[PersonalityProbe, dict[str, Any]]:
        """Load probe from disk."""
        path = Path(path)

        with open(path / "config.json") as f:
            config = json.load(f)

        probe = PersonalityProbe(
            hidden_dim=config["hidden_dim"],
            num_dimensions=config["num_dimensions"],
            probe_type=config["probe_type"],
        )
        probe.load_state_dict(torch.load(path / "probe.pt", weights_only=True))

        return probe, config


class ProbeAnalyzer:
    """
    Analysis utilities for trained probes.

    Provides methods for understanding what the probe learned and
    how dimensions relate to model internals.
    """

    def __init__(self, probe: PersonalityProbe):
        if probe.probe_type != "linear":
            raise ValueError("ProbeAnalyzer requires a linear probe")
        self.probe = probe
        self.weights = probe.get_weights()

    def dimension_similarity_matrix(self) -> torch.Tensor:
        """
        Compute cosine similarity between dimension weight vectors.

        Returns:
            [num_dimensions, num_dimensions] similarity matrix
        """
        normalized = F.normalize(self.weights, dim=1)
        return normalized @ normalized.T

    def find_shared_circuits(self, threshold: float = 0.5) -> list[tuple[str, str, float]]:
        """
        Find dimension pairs with high weight similarity.

        Args:
            threshold: Minimum cosine similarity to report

        Returns:
            List of (dim1, dim2, similarity) tuples
        """
        sim_matrix = self.dimension_similarity_matrix()
        shared = []

        for i in range(len(DIMENSIONS)):
            for j in range(i + 1, len(DIMENSIONS)):
                sim = sim_matrix[i, j].item()
                if sim > threshold:
                    shared.append((DIMENSIONS[i], DIMENSIONS[j], sim))

        return sorted(shared, key=lambda x: -x[2])

    def top_features_for_dimension(
        self, dimension: str | int, k: int = 20
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Find top-k most important features for a dimension.

        Args:
            dimension: Dimension name or index
            k: Number of features to return

        Returns:
            (indices, weights) of top-k features
        """
        weights = self.probe.get_dimension_weights(dimension)
        values, indices = torch.abs(weights).topk(k)
        return indices, weights[indices]

    def feature_importance_overlap(
        self, dim1: str | int, dim2: str | int, k: int = 50
    ) -> float:
        """
        Compute overlap in top-k features between two dimensions.

        Returns:
            Jaccard similarity of top-k feature sets
        """
        idx1, _ = self.top_features_for_dimension(dim1, k)
        idx2, _ = self.top_features_for_dimension(dim2, k)

        set1 = set(idx1.tolist())
        set2 = set(idx2.tolist())

        intersection = len(set1 & set2)
        union = len(set1 | set2)

        return intersection / union
