"""
Data preparation: combine 5 HF datasets into unified train/test splits.
Labels: 0=negative, 1=neutral, 2=positive
"""
import numpy as np
from datasets import load_dataset, Dataset, DatasetDict, concatenate_datasets, Features, Value


def load_financial_phrasebank():
    ds = load_dataset("nickmuchi/financial-classification")
    label_map = {0: 0, 1: 1, 2: 2}  # already negative/neutral/positive
    def process(ex):
        return {"text": ex["text"], "label": label_map[ex["labels"]]}
    train = ds["train"].map(process, remove_columns=ds["train"].column_names)
    test = ds["test"].map(process, remove_columns=ds["test"].column_names)
    return train, test


def load_twitter_financial():
    ds = load_dataset("zeroshot/twitter-financial-news-sentiment")
    label_map = {0: 0, 1: 2, 2: 1}  # 0=Bear→neg, 1=Bull→pos, 2=Neutral→neutral
    def process(ex):
        return {"text": ex["text"], "label": label_map[ex["label"]]}
    train = ds["train"].map(process, remove_columns=ds["train"].column_names)
    test = ds["validation"].map(process, remove_columns=ds["validation"].column_names)
    return train, test


def load_auditor_sentiment():
    ds = load_dataset("FinanceInc/auditor_sentiment")
    label_map = {0: 0, 1: 1, 2: 2}
    def process(ex):
        return {"text": ex["sentence"], "label": label_map[ex["label"]]}
    train = ds["train"].map(process, remove_columns=ds["train"].column_names)
    test = ds["test"].map(process, remove_columns=ds["test"].column_names)
    return train, test


def load_fiqa():
    ds = load_dataset("pauri32/fiqa-2018")
    def to_label(score):
        if score < -0.15: return 0
        elif score > 0.15: return 2
        else: return 1
    def process(ex):
        return {"text": ex["sentence"], "label": to_label(ex["sentiment_score"])}
    train = ds["train"].map(process, remove_columns=ds["train"].column_names)
    val = ds["validation"].map(process, remove_columns=ds["validation"].column_names)
    test = ds["test"].map(process, remove_columns=ds["test"].column_names)
    train = concatenate_datasets([train, val])
    return train, test


def load_climate_sentiment():
    ds = load_dataset("climatebert/climate_sentiment")
    # ClassLabel: 0=risk→neg, 1=neutral, 2=opportunity→pos
    def process(ex):
        return {"text": ex["text"], "label": int(ex["label"])}
    train = ds["train"].map(process, remove_columns=ds["train"].column_names)
    test = ds["test"].map(process, remove_columns=ds["test"].column_names)
    return train, test


def load_combined_dataset():
    """Load and combine all 5 datasets."""
    loaders = [
        ("financial_phrasebank", load_financial_phrasebank),
        ("twitter_financial", load_twitter_financial),
        ("auditor_sentiment", load_auditor_sentiment),
        ("fiqa", load_fiqa),
        ("climate_sentiment", load_climate_sentiment),
    ]
    
    all_train, all_test = [], []
    feat = Features({"text": Value("string"), "label": Value("int64")})
    for name, loader in loaders:
        print(f"Loading {name}...")
        train, test = loader()
        # Cast to uniform schema
        train = train.cast(feat)
        test = test.cast(feat)
        print(f"  {name}: {len(train)} train, {len(test)} test")
        all_train.append(train)
        all_test.append(test)
    
    combined_train = concatenate_datasets(all_train)
    combined_test = concatenate_datasets(all_test)
    
    # Shuffle
    combined_train = combined_train.shuffle(seed=42)
    combined_test = combined_test.shuffle(seed=42)
    
    print(f"\nCombined: {len(combined_train)} train, {len(combined_test)} test")
    
    # Label distribution
    from collections import Counter
    train_dist = Counter(combined_train["label"])
    print(f"Train distribution: {dict(train_dist)}")
    
    return DatasetDict({"train": combined_train, "test": combined_test})


if __name__ == "__main__":
    ds = load_combined_dataset()
    print(ds)