"""PyTorch Dataset classes.

We provide three:
  - ACSADataset           text only,  per-aspect labels   -> Baseline 3 (no meta)
  - MetaACSADataset       text + meta features            -> Proposed model
  - OverallSentimentDataset  text only, single 3-class    -> Baseline 2 (BERT-overall)
"""
from typing import Optional

import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset

from . import config as cfg
from .meta_encoder import MetaEncoder


class _BaseTextDataset(Dataset):
    def __init__(self, df: pd.DataFrame, tokenizer, max_length: int = cfg.MAX_LENGTH,
                 text_col: str = "full_text"):
        self.df = df.reset_index(drop=True)
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.text_col = text_col

    def __len__(self):
        return len(self.df)

    def _encode_text(self, text: str):
        enc = self.tokenizer(
            text,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )
        return enc["input_ids"].squeeze(0), enc["attention_mask"].squeeze(0)


class ACSADataset(_BaseTextDataset):
    """Per-aspect multi-head training, text-only. (Baseline 3.)"""

    def __init__(self, df, tokenizer, **kw):
        super().__init__(df, tokenizer, **kw)
        self.aspect_cols = [f"aspect_{a}" for a in cfg.ASPECTS]
        missing = [c for c in self.aspect_cols if c not in self.df.columns]
        if missing:
            raise KeyError(f"ACSADataset: missing aspect label columns {missing}. "
                           f"Did you run weak labeling (script 03)?")

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        input_ids, attn = self._encode_text(str(row[self.text_col]))
        labels = torch.tensor([int(row[c]) for c in self.aspect_cols], dtype=torch.long)
        return {"input_ids": input_ids, "attention_mask": attn, "labels": labels}


class MetaACSADataset(_BaseTextDataset):
    """Per-aspect multi-head training, text + metadata. (Proposed model.)

    The metadata is pre-encoded to a numpy matrix of shape (N, meta_dim) using
    a fit MetaEncoder, and indexed in lock-step with the text rows.

    If the df contains an 'overall_label' column (0=Neg, 1=Neu, 2=Pos), it is
    returned as part of each batch for the joint overall auxiliary loss.
    """

    def __init__(self, df, tokenizer, meta_encoder: MetaEncoder, **kw):
        super().__init__(df, tokenizer, **kw)
        self.aspect_cols = [f"aspect_{a}" for a in cfg.ASPECTS]
        missing = [c for c in self.aspect_cols if c not in self.df.columns]
        if missing:
            raise KeyError(f"MetaACSADataset: missing aspect label columns {missing}.")

        self.has_overall = "overall_label" in self.df.columns

        # Pre-encode metadata once, in the same row order as self.df.
        self.meta_matrix = meta_encoder.transform(self.df)
        assert self.meta_matrix.shape[0] == len(self.df), \
            "meta matrix and df length mismatch"

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        input_ids, attn = self._encode_text(str(row[self.text_col]))
        labels = torch.tensor([int(row[c]) for c in self.aspect_cols], dtype=torch.long)
        meta = torch.from_numpy(self.meta_matrix[idx]).float()
        item = {
            "input_ids": input_ids,
            "attention_mask": attn,
            "meta_features": meta,
            "labels": labels,
        }
        if self.has_overall:
            item["overall_labels"] = torch.tensor(int(row["overall_label"]), dtype=torch.long)
        return item


class OverallSentimentDataset(_BaseTextDataset):
    """Single 3-class overall sentiment label. (Baseline 2: BERT-overall.)"""

    def __init__(self, df, tokenizer, label_col: str = "overall_label", **kw):
        super().__init__(df, tokenizer, **kw)
        self.label_col = label_col
        if self.label_col not in self.df.columns:
            raise KeyError(f"OverallSentimentDataset: missing {self.label_col} column.")

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        input_ids, attn = self._encode_text(str(row[self.text_col]))
        return {
            "input_ids": input_ids,
            "attention_mask": attn,
            "labels": torch.tensor(int(row[self.label_col]), dtype=torch.long),
        }