"""Training data assembly for dual-head fine-tuning.

Merges Stage 1 consensus labels with paragraph text, filters holdout and
truncated filings, assigns quality-tier sample weights, tokenizes, and
produces a train/val DatasetDict ready for the Trainer.
"""

import json
from pathlib import Path

from datasets import Dataset, DatasetDict
from transformers import PreTrainedTokenizerBase

# Canonical label maps — alphabetical for reproducibility
CATEGORIES = [
    "Board Governance",       # 0
    "Incident Disclosure",    # 1
    "Management Role",        # 2
    "None/Other",             # 3
    "Risk Management Process",# 4
    "Strategy Integration",   # 5
    "Third-Party Risk",       # 6
]
CAT2ID = {c: i for i, c in enumerate(CATEGORIES)}
NUM_CATEGORIES = len(CATEGORIES)

# Specificity: ordinal 1-4 → 0-indexed 0-3
NUM_SPECIFICITY = 4

# Quality tier → sample weight (STATUS.md: clean/headed/minor 1.0, degraded 0.5)
TIER_WEIGHT = {
    "clean": 1.0,
    "headed": 1.0,
    "minor": 1.0,
    "degraded": 0.5,
}


def _load_jsonl(path: str | Path) -> list[dict]:
    records = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


def load_finetune_data(
    *,
    paragraphs_path: str,
    consensus_path: str,
    quality_path: str,
    holdout_path: str,
    max_seq_length: int,
    validation_split: float,
    tokenizer: PreTrainedTokenizerBase,
    seed: int = 42,
) -> DatasetDict:
    """Load, merge, filter, tokenize, and split training data.

    Returns a DatasetDict with 'train' and 'test' splits, each containing:
        input_ids, attention_mask, category_labels, specificity_labels, sample_weight
    """
    print("  Loading data files...")

    # 1. Load all sources
    paragraphs = {p["id"]: p for p in _load_jsonl(paragraphs_path)}
    consensus = {c["paragraphId"]: c for c in _load_jsonl(consensus_path)}
    quality = {q["id"]: q for q in _load_jsonl(quality_path)}

    with open(holdout_path) as f:
        holdout_ids = set(json.load(f))

    print(f"    Paragraphs: {len(paragraphs):,}")
    print(f"    Consensus:  {len(consensus):,}")
    print(f"    Quality:    {len(quality):,}")
    print(f"    Holdout:    {len(holdout_ids):,}")

    # 2. Identify individually truncated paragraphs
    truncated_ids: set[str] = set()
    for q in quality.values():
        if "truncated" in q.get("issues", []):
            truncated_ids.add(q["id"])
    print(f"    Truncated paragraphs: {len(truncated_ids)}")

    # 3. Merge and filter
    records: list[dict] = []
    skipped_holdout = 0
    skipped_truncated = 0
    skipped_no_consensus = 0

    for pid, para in paragraphs.items():
        # Exclude holdout
        if pid in holdout_ids:
            skipped_holdout += 1
            continue

        # Exclude individually truncated paragraphs
        if pid in truncated_ids:
            skipped_truncated += 1
            continue

        # Must have consensus label
        cons = consensus.get(pid)
        if cons is None:
            skipped_no_consensus += 1
            continue

        label = cons["finalLabel"]
        cat = label["content_category"]
        spec = label["specificity_level"]  # 1-4

        # Quality tier weight
        q = quality.get(pid, {})
        tier = q.get("quality_tier", "clean")
        weight = TIER_WEIGHT.get(tier, 1.0)

        # Specificity confidence: 1.0 if all 3 runs agreed, 0.0 otherwise
        spec_agreed = cons.get("specificityAgreement", {}).get("agreed", True)

        records.append({
            "text": para["text"],
            "category_labels": CAT2ID[cat],
            "specificity_labels": spec - 1,  # 0-indexed
            "sample_weight": weight,
            "spec_mask": 1.0 if spec_agreed else 0.0,
        })

    print(f"\n    Skipped holdout:      {skipped_holdout:,}")
    print(f"    Skipped truncated:    {skipped_truncated:,}")
    print(f"    Skipped no consensus: {skipped_no_consensus:,}")
    print(f"    Training samples:     {len(records):,}")

    # 4. Distributions
    from collections import Counter
    cat_counts = Counter(r["category_labels"] for r in records)
    spec_counts = Counter(r["specificity_labels"] for r in records)
    spec_agreed = sum(1 for r in records if r["spec_mask"] == 1.0)
    print(f"    Spec high-confidence: {spec_agreed:,} ({spec_agreed/len(records)*100:.1f}%)")
    print(f"\n    Category distribution:")
    for i, name in enumerate(CATEGORIES):
        print(f"      {name}: {cat_counts[i]:,} ({cat_counts[i]/len(records)*100:.1f}%)")
    print(f"    Specificity distribution:")
    for i in range(NUM_SPECIFICITY):
        print(f"      L{i+1}: {spec_counts[i]:,} ({spec_counts[i]/len(records)*100:.1f}%)")

    # 5. Create HuggingFace dataset
    dataset = Dataset.from_list(records)

    # 6. Tokenize
    print(f"\n    Tokenizing to max_seq_length={max_seq_length}...")

    def tokenize_fn(examples):
        return tokenizer(
            examples["text"],
            truncation=True,
            max_length=max_seq_length,
            padding=False,  # dynamic padding in collator
        )

    dataset = dataset.map(
        tokenize_fn,
        batched=True,
        remove_columns=["text"],
        desc="Tokenizing",
    )

    # 7. Stratified train/val split (using sklearn since HF requires ClassLabel)
    from sklearn.model_selection import StratifiedShuffleSplit

    sss = StratifiedShuffleSplit(n_splits=1, test_size=validation_split, random_state=seed)
    train_idx, val_idx = next(sss.split(range(len(dataset)), dataset["category_labels"]))
    split = DatasetDict({
        "train": dataset.select(train_idx),
        "test": dataset.select(val_idx),
    })
    print(f"    Train: {len(split['train']):,} | Val: {len(split['test']):,}")

    return split


def compute_class_weights(dataset: Dataset) -> dict[str, list[float]]:
    """Compute inverse-frequency class weights from training data."""
    from collections import Counter
    import math

    n = len(dataset)

    # Category weights
    cat_counts = Counter(dataset["category_labels"])
    cat_weights = []
    for i in range(NUM_CATEGORIES):
        count = cat_counts.get(i, 1)
        # Inverse frequency, normalized so mean weight = 1.0
        cat_weights.append(n / (NUM_CATEGORIES * count))
    # Cap extreme weights (ID category is very rare)
    max_weight = 10.0
    cat_weights = [min(w, max_weight) for w in cat_weights]

    # Specificity weights
    spec_counts = Counter(dataset["specificity_labels"])
    spec_weights = []
    for i in range(NUM_SPECIFICITY):
        count = spec_counts.get(i, 1)
        spec_weights.append(n / (NUM_SPECIFICITY * count))

    return {
        "category_weights": cat_weights,
        "specificity_weights": spec_weights,
    }
