"""Phase 2 data cleaning and splitting.

Loads raw questions, removes duplicates, and writes train/val/test splits.
"""

import json
import os
import sys

import pandas as pd
from sklearn.model_selection import train_test_split

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import config  # noqa: E402


class DatasetProcessor:
    """Clean, deduplicate, and split the raw dataset."""

    def __init__(self, processing_config):
        """Store processing config."""
        self.cfg = processing_config

    def load_raw(self):
        """Load raw JSONL into a DataFrame."""
        records = []
        with open(self.cfg.raw_path, "r", encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if line:
                    records.append(json.loads(line))
        frame = pd.DataFrame(records)
        print(f"Loaded {len(frame)} raw questions.")
        return frame

    @staticmethod
    def _normalise(text):
        """Normalize text for exact duplicate checks."""
        return " ".join(text.lower().split())

    def remove_exact_duplicates(self, frame):
        """Drop exact duplicates and keep the first one."""
        before = len(frame)
        frame = frame.copy()
        frame["_norm"] = frame["question"].map(self._normalise)
        frame = frame.drop_duplicates(subset="_norm", keep="first")
        frame = frame.drop(columns="_norm").reset_index(drop=True)
        print(f"Exact dedup: {before} -> {len(frame)} "
              f"({before - len(frame)} removed).")
        return frame

    def remove_near_duplicates(self, frame):
        """Remove near-duplicate questions with embedding similarity."""
        from sentence_transformers import SentenceTransformer, util

        before = len(frame)
        questions = frame["question"].tolist()

        print(f"Encoding {before} questions for near-dup detection "
              f"(model: {self.cfg.embed_model})...")
        model = SentenceTransformer(self.cfg.embed_model)
        embeddings = model.encode(
            questions, convert_to_tensor=True, show_progress_bar=True
        )

        # Returns [score, i, j] sorted by score.
        pairs = util.paraphrase_mining_embeddings(embeddings)

        removed = set()
        for score, i, j in pairs:
            if score < self.cfg.near_dup_threshold:
                break
            low, high = (i, j) if i < j else (j, i)
            if low not in removed and high not in removed:
                removed.add(high)

        kept_mask = [idx not in removed for idx in range(before)]
        frame = frame[kept_mask].reset_index(drop=True)
        print(f"Near-dup removal (threshold {self.cfg.near_dup_threshold}): "
              f"{before} -> {len(frame)} ({before - len(frame)} removed).")
        return frame

    def split_in_domain(self, frame):
        """Split in-domain rows into train/val/test."""
        in_domain = frame[frame["split"] == "train"].copy().reset_index(drop=True)
        ood = frame[frame["split"] == "ood"].copy().reset_index(drop=True)

        strat_key = in_domain["domain"] + "|" + in_domain["bloom_class"]

        # Split off test first.
        train_val, test = train_test_split(
            in_domain,
            test_size=self.cfg.test_size,
            stratify=strat_key,
            random_state=self.cfg.random_state,
        )
        # Then split validation from the rest.
        strat_key_tv = train_val["domain"] + "|" + train_val["bloom_class"]
        relative_val = self.cfg.val_size / (1.0 - self.cfg.test_size)
        train, val = train_test_split(
            train_val,
            test_size=relative_val,
            stratify=strat_key_tv,
            random_state=self.cfg.random_state,
        )

        splits = {
            "train": train.reset_index(drop=True),
            "val": val.reset_index(drop=True),
            "test": test.reset_index(drop=True),
            "ood_test": ood,
        }
        return splits

    def save_splits(self, splits):
        """Save each split as a JSONL file."""
        os.makedirs(self.cfg.processed_dir, exist_ok=True)
        keep_cols = [
            "question", "bloom_sublevel", "bloom_class",
            "domain", "topic", "source",
        ]
        for name, frame in splits.items():
            path = os.path.join(self.cfg.processed_dir, f"{name}.jsonl")
            with open(path, "w", encoding="utf-8") as handle:
                for _, row in frame.iterrows():
                    record = {col: row[col] for col in keep_cols if col in row}
                    handle.write(json.dumps(record, ensure_ascii=False) + "\n")
            print(f"  wrote {len(frame):4d} -> {path}")

    @staticmethod
    def print_breakdown(splits):
        """Print split sizes and class/domain counts."""
        print("\nSplit sizes:")
        for name, frame in splits.items():
            print(f"  {name:9s} {len(frame)}")

        print("\nClass balance per split:")
        for name, frame in splits.items():
            counts = frame["bloom_class"].value_counts().to_dict()
            print(f"  {name:9s} {counts}")

        print("\nPer-(domain x class) counts (all splits combined):")
        combined = pd.concat(splits.values())
        table = combined.groupby(["domain", "bloom_class"]).size().unstack(fill_value=0)
        print(table.to_string())

    def run(self):
        """Run the full Phase 2 pipeline."""
        frame = self.load_raw()
        frame = self.remove_exact_duplicates(frame)
        frame = self.remove_near_duplicates(frame)
        splits = self.split_in_domain(frame)
        self.save_splits(splits)
        self.print_breakdown(splits)
        print("\nPhase 2 complete.")


def main():
    """Run Phase 2 processing."""
    processor = DatasetProcessor(config.ProcessingConfig())
    processor.run()


if __name__ == "__main__":
    main()