"""
External dataset loaders for personality training.

Supports loading from public personality datasets:
- Big5-Chat: https://huggingface.co/datasets/yutong-wu/Big5-Chat
- PANDORA: https://huggingface.co/datasets/Nofear/PANDORA

These datasets provide Big Five personality labels which we map to
Lindr's 10-dimension personality space.
"""

from __future__ import annotations

from typing import Iterator, Any

import torch
from torch.utils.data import Dataset


# Mapping from Big Five to Lindr's 10 dimensions
# Big Five: Openness, Conscientiousness, Extraversion, Agreeableness, Neuroticism
# Extended: Assertiveness, Ambition, Resilience, Integrity, Curiosity
#
# Derivation rules for extended dimensions:
# - Assertiveness: High Extraversion + Low Agreeableness
# - Ambition: High Conscientiousness + High Extraversion
# - Resilience: Low Neuroticism
# - Integrity: High Conscientiousness + High Agreeableness
# - Curiosity: High Openness (direct mapping)


def _derive_extended_dimensions(big5: dict[str, float]) -> dict[str, float]:
    """
    Derive extended personality dimensions from Big Five scores.

    Args:
        big5: Dictionary with keys: openness, conscientiousness, extraversion,
              agreeableness, neuroticism (values in [0, 1])

    Returns:
        Dictionary with all 10 Lindr personality dimensions
    """
    o = big5.get("openness", 0.5)
    c = big5.get("conscientiousness", 0.5)
    e = big5.get("extraversion", 0.5)
    a = big5.get("agreeableness", 0.5)
    n = big5.get("neuroticism", 0.5)

    return {
        # Direct mappings
        "openness": o,
        "conscientiousness": c,
        "extraversion": e,
        "agreeableness": a,
        "neuroticism": n,
        # Derived dimensions
        "assertiveness": min(1.0, (e + (1 - a)) / 2),  # High E, Low A
        "ambition": min(1.0, (c + e) / 2),  # High C and E
        "resilience": 1 - n,  # Inverse of neuroticism
        "integrity": min(1.0, (c + a) / 2),  # High C and A
        "curiosity": o,  # Direct mapping to openness
    }


def _scores_to_list(scores: dict[str, float]) -> list[float]:
    """Convert scores dict to ordered list matching PERSONALITY_DIMENSIONS."""
    dimensions = [
        "openness",
        "conscientiousness",
        "extraversion",
        "agreeableness",
        "neuroticism",
        "assertiveness",
        "ambition",
        "resilience",
        "integrity",
        "curiosity",
    ]
    return [scores.get(d, 0.5) for d in dimensions]


class Big5ChatLoader:
    """
    Loader for the Big5-Chat dataset.

    Dataset: https://huggingface.co/datasets/yutong-wu/Big5-Chat
    Contains personality-labeled text samples from chat-like conversations.
    """

    HF_DATASET = "yutong-wu/Big5-Chat"

    def __init__(self, split: str = "train", max_samples: int | None = None) -> None:
        """
        Initialize Big5-Chat loader.

        Args:
            split: Dataset split to load (train, validation, test)
            max_samples: Maximum number of samples to load (None for all)
        """
        try:
            from datasets import load_dataset
        except ImportError:
            raise ImportError(
                "datasets library is required for loading external datasets. "
                "Install with: pip install datasets"
            )

        self.dataset = load_dataset(self.HF_DATASET, split=split)
        self.max_samples = max_samples

    def to_lindr_format(self) -> Iterator[dict]:
        """
        Convert Big5-Chat samples to Lindr training format.

        Yields:
            Dictionaries with 'text' and 'scores' keys
        """
        count = 0
        for sample in self.dataset:
            if self.max_samples and count >= self.max_samples:
                break

            # Extract text and Big Five scores
            text = sample.get("text", "")
            if not text:
                continue

            # Big5-Chat uses different score key names
            # Adapt based on actual dataset structure
            big5 = {
                "openness": self._normalize_score(sample.get("openness", sample.get("O", 0.5))),
                "conscientiousness": self._normalize_score(sample.get("conscientiousness", sample.get("C", 0.5))),
                "extraversion": self._normalize_score(sample.get("extraversion", sample.get("E", 0.5))),
                "agreeableness": self._normalize_score(sample.get("agreeableness", sample.get("A", 0.5))),
                "neuroticism": self._normalize_score(sample.get("neuroticism", sample.get("N", 0.5))),
            }

            # Derive all 10 dimensions
            full_scores = _derive_extended_dimensions(big5)

            yield {
                "text": text,
                "scores": _scores_to_list(full_scores),
                "source": "big5-chat",
            }
            count += 1

    def _normalize_score(self, value: Any) -> float:
        """Normalize score to [0, 1] range."""
        if value is None:
            return 0.5
        try:
            v = float(value)
            # If scores are in different ranges (e.g., 1-5, 0-100), normalize
            if v > 1:
                if v <= 5:
                    return (v - 1) / 4  # 1-5 scale
                elif v <= 100:
                    return v / 100  # 0-100 scale
            return max(0.0, min(1.0, v))
        except (TypeError, ValueError):
            return 0.5


class PANDORALoader:
    """
    Loader for the PANDORA Reddit personality dataset.

    Dataset: https://huggingface.co/datasets/Nofear/PANDORA
    Contains Reddit comments with Big Five personality annotations.
    """

    HF_DATASET = "Nofear/PANDORA"

    def __init__(self, split: str = "train", max_samples: int | None = None) -> None:
        """
        Initialize PANDORA loader.

        Args:
            split: Dataset split to load
            max_samples: Maximum number of samples to load
        """
        try:
            from datasets import load_dataset
        except ImportError:
            raise ImportError(
                "datasets library is required for loading external datasets. "
                "Install with: pip install datasets"
            )

        self.dataset = load_dataset(self.HF_DATASET, split=split)
        self.max_samples = max_samples

    def to_lindr_format(self) -> Iterator[dict]:
        """
        Convert PANDORA samples to Lindr training format.

        Yields:
            Dictionaries with 'text' and 'scores' keys
        """
        count = 0
        for sample in self.dataset:
            if self.max_samples and count >= self.max_samples:
                break

            # Extract text (Reddit comment)
            text = sample.get("text", sample.get("body", ""))
            if not text or len(text) < 20:  # Skip very short texts
                continue

            # Extract Big Five scores
            big5 = {
                "openness": self._normalize_score(sample.get("openness")),
                "conscientiousness": self._normalize_score(sample.get("conscientiousness")),
                "extraversion": self._normalize_score(sample.get("extraversion")),
                "agreeableness": self._normalize_score(sample.get("agreeableness")),
                "neuroticism": self._normalize_score(sample.get("neuroticism")),
            }

            # Derive all 10 dimensions
            full_scores = _derive_extended_dimensions(big5)

            yield {
                "text": text,
                "scores": _scores_to_list(full_scores),
                "source": "pandora",
            }
            count += 1

    def _normalize_score(self, value: Any) -> float:
        """Normalize score to [0, 1] range."""
        if value is None:
            return 0.5
        try:
            v = float(value)
            return max(0.0, min(1.0, v))
        except (TypeError, ValueError):
            return 0.5


class ExternalDataset(Dataset):
    """
    PyTorch Dataset wrapper for external personality datasets.

    Computes embeddings on-the-fly from text samples.
    """

    def __init__(
        self,
        samples: list[dict],
        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
    ) -> None:
        """
        Initialize external dataset.

        Args:
            samples: List of samples with 'text' and 'scores' keys
            embedding_model: Model name for computing embeddings
        """
        self.samples = samples
        self.embedding_model_name = embedding_model
        self._embedder = None

    @property
    def embedder(self):
        """Lazy-load sentence transformer."""
        if self._embedder is None:
            from sentence_transformers import SentenceTransformer
            self._embedder = SentenceTransformer(self.embedding_model_name)
        return self._embedder

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

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        sample = self.samples[idx]

        # Compute embedding
        text = sample.get("text", "")
        embedding_np = self.embedder.encode(
            text,
            convert_to_numpy=True,
            normalize_embeddings=True,
        )
        embedding = torch.tensor(embedding_np, dtype=torch.float32)

        # Get scores
        scores = sample.get("scores", [0.5] * 10)
        scores_tensor = torch.tensor(scores, dtype=torch.float32)

        return embedding, scores_tensor


def load_external_dataset(
    name: str,
    split: str = "train",
    max_samples: int | None = None,
    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
) -> ExternalDataset:
    """
    Load an external personality dataset by name.

    Args:
        name: Dataset name ('big5-chat' or 'pandora')
        split: Dataset split to load
        max_samples: Maximum number of samples
        embedding_model: Model for computing embeddings

    Returns:
        ExternalDataset ready for training
    """
    loaders = {
        "big5-chat": Big5ChatLoader,
        "pandora": PANDORALoader,
    }

    if name.lower() not in loaders:
        raise ValueError(f"Unknown dataset: {name}. Available: {list(loaders.keys())}")

    loader = loaders[name.lower()](split=split, max_samples=max_samples)
    samples = list(loader.to_lindr_format())

    return ExternalDataset(samples, embedding_model=embedding_model)
