"""
dataset_builder.py
Creates a stratified 'Golden Dataset' from the full MBFC evaluation data.
Targets 50-60 samples with uniform distribution across bias labels AND factuality levels.
"""

import json
import random
import logging
from typing import List, Dict, Any
from collections import defaultdict

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Input: the full scraped MBFC dataset (645 items currently)
RAW_DATA_PATH = "evaluation_dataset_full.json"
OUTPUT_PATH = "evaluation_dataset.json"

# 7 canonical bias categories
BIAS_CATEGORIES = [
    "Extreme Left", "Left", "Left-Center", "Least Biased",
    "Right-Center", "Right", "Extreme Right"
]

# 5 factuality levels (from best to worst)
FACTUALITY_LEVELS = ["VERY HIGH", "HIGH", "MOSTLY FACTUAL", "MIXED", "LOW", "VERY LOW"]

# Target: ~8 per bias category = 56 total
SAMPLES_PER_BIAS = 43
TARGET_TOTAL = SAMPLES_PER_BIAS * len(BIAS_CATEGORIES)  # 56


def normalize_bias_label(raw_label: str, bias_score: float = None) -> str:
    """Normalize messy MBFC bias labels to 7 canonical categories.
    Falls back to score-based bucketing if label is non-standard."""
    if raw_label is None and bias_score is None:
        return None

    # Try label-based mapping first
    if raw_label:
        upper = raw_label.upper().strip()

        # Direct matches
        label_map = {
            "EXTREME LEFT": "Extreme Left",
            "FAR LEFT": "Extreme Left",
            "FAR-LEFT": "Extreme Left",
            "FAR LEFT BIAS": "Extreme Left",
            "LEFT": "Left",
            "LEFT-CENTER": "Left-Center",
            "LEAST BIASED": "Least Biased",
            "CENTER": "Least Biased",
            "PRO-SCIENCE": "Least Biased",
            "RIGHT-CENTER": "Right-Center",
            "RIGHT": "Right",
            "FAR RIGHT": "Extreme Right",
            "FAR-RIGHT": "Extreme Right",
            "EXTREME RIGHT": "Extreme Right",
        }

        # Check direct match
        if upper in label_map:
            return label_map[upper]

        # Check partial matches for compound labels (e.g., "RIGHT CONSPIRACY-PSEUDOSCIENCE")
        for key, val in label_map.items():
            if upper.startswith(key):
                return val

    # Score-based fallback
    if bias_score is not None:
        return get_bias_label_from_score(bias_score)

    return None


def get_bias_label_from_score(score: float) -> str:
    if score is None:
        return None
    if -10.0 <= score <= -8.0:
        return "Extreme Left"
    if -7.9 <= score <= -5.0:
        return "Left"
    if -4.9 <= score <= -2.0:
        return "Left-Center"
    if -1.9 <= score <= 1.9:
        return "Least Biased"
    if 2.0 <= score <= 4.9:
        return "Right-Center"
    if 5.0 <= score <= 7.9:
        return "Right"
    if 8.0 <= score <= 10.0:
        return "Extreme Right"
    return None


def normalize_factuality(raw: str) -> str:
    """Normalize factuality labels."""
    if not raw:
        return None
    upper = raw.strip().upper()
    # Fix typos like "L OW"
    upper = " ".join(upper.split())
    if upper in FACTUALITY_LEVELS:
        return upper
    return None


def build_dataset():
    logger.info(f"Loading raw data from {RAW_DATA_PATH}...")
    try:
        with open(RAW_DATA_PATH, 'r', encoding='utf-8') as f:
            raw_data = json.load(f)
    except FileNotFoundError:
        logger.error(f"File {RAW_DATA_PATH} not found. "
                     f"Rename your full dataset to '{RAW_DATA_PATH}' first.")
        return

    # Bucket: bias_label -> factuality_label -> [items]
    buckets = defaultdict(lambda: defaultdict(list))

    valid_count = 0
    for item in raw_data:
        bias_score = item.get('bias_score')
        if bias_score is None:
            continue

        bias_cat = normalize_bias_label(item.get('bias_rating'), bias_score)
        if not bias_cat:
            continue

        fact_cat = normalize_factuality(item.get('factual_reporting'))
        if not fact_cat:
            continue

        # Store normalized labels back
        item['_norm_bias'] = bias_cat
        item['_norm_factuality'] = fact_cat

        buckets[bias_cat][fact_cat].append(item)
        valid_count += 1

    logger.info(f"Categorized {valid_count} valid entries into {len(buckets)} bias categories.")

    # Log distribution
    for bias in BIAS_CATEGORIES:
        facts = {f: len(buckets[bias][f]) for f in FACTUALITY_LEVELS if buckets[bias][f]}
        total = sum(facts.values())
        logger.info(f"  {bias}: {total} total | {facts}")

    # Stratified sampling: for each bias category, sample uniformly across factuality levels
    final_dataset = []
    random.seed(42)  # Reproducibility

    for bias in BIAS_CATEGORIES:
        available_facts = [f for f in FACTUALITY_LEVELS if buckets[bias][f]]
        if not available_facts:
            logger.warning(f"No data for {bias}, skipping.")
            continue

        selected = []
        # Round-robin across factuality levels
        target = SAMPLES_PER_BIAS
        round_robin_idx = 0

        while len(selected) < target:
            added = False
            for fact in available_facts:
                if len(selected) >= target:
                    break
                source = buckets[bias][fact]
                if source:
                    idx = random.randrange(len(source))
                    item = source.pop(idx)
                    # Clean up temp keys
                    item.pop('_norm_bias', None)
                    item.pop('_norm_factuality', None)
                    selected.append(item)
                    added = True

            if not added:
                break  # Exhausted all items for this bias

        logger.info(f"  {bias}: selected {len(selected)} samples")
        final_dataset.extend(selected)

    # Shuffle final dataset
    random.shuffle(final_dataset)

    with open(OUTPUT_PATH, 'w', encoding='utf-8') as f:
        json.dump(final_dataset, f, indent=2)

    # Print summary
    from collections import Counter
    bias_dist = Counter(normalize_bias_label(i.get('bias_rating'), i.get('bias_score')) for i in final_dataset)
    fact_dist = Counter(normalize_factuality(i.get('factual_reporting')) for i in final_dataset)
    has_fc = sum(1 for i in final_dataset
                 if i.get('failed_fact_checks')
                 and len(i['failed_fact_checks']) > 0
                 and i['failed_fact_checks'][0] not in ['None in the last 5 years', 'None Found', 'None'])

    logger.info(f"\nDataset created at {OUTPUT_PATH} with {len(final_dataset)} total samples.")
    logger.info(f"Bias distribution: {dict(bias_dist)}")
    logger.info(f"Factuality distribution: {dict(fact_dist)}")
    logger.info(f"Samples with failed fact checks: {has_fc}/{len(final_dataset)}")


if __name__ == "__main__":
    build_dataset()