"""FinBERT Cross-Dataset Evaluation - measure distribution shift between Financial PhraseBank and FinMarBa."""

# 1. Setup and Configuration
import os
import warnings

import matplotlib.pyplot as plt
import numpy as np
import polars as pl
import torch
from datasets import load_dataset
from sklearn.metrics import accuracy_score, confusion_matrix, f1_score
from transformers import pipeline

from utils.paths import get_chapter_dir
from utils.reproducibility import set_global_seeds

# Polars display configuration

os.environ["TOKENIZERS_PARALLELISM"] = "false"
warnings.filterwarnings("ignore")

SEED = 42
SAMPLE_SIZE = 10000

# Reproducibility - set_global_seeds covers Python random / NumPy / Torch.
set_global_seeds(SEED)

device = 0 if torch.cuda.is_available() else -1
print(f"Using device: {'GPU' if device == 0 else 'CPU'}")

def load_finmarba_dataset(sample_size: int = 10000) -> pl.DataFrame:
    """Load FinMarBa dataset from Hugging Face.

    Args:
        sample_size: Number of samples to use (dataset has 60K+ samples).
            Set to 0 for all samples.

    Raises:
        RuntimeError: If dataset cannot be loaded (network error, etc.)
    """
    try:
        # Try 'test' split first, fall back to 'train' (dataset structure changed)
        try:
            ds = load_dataset(
                "baptle/financial_headlines_market_based",
                split="test",
            )
            print("Loaded FinMarBa test split")
        except ValueError:
            # Dataset may only have 'train' split now
            split_str = f"train[:{sample_size}]" if sample_size > 0 else "train"
            ds = load_dataset(
                "baptle/financial_headlines_market_based",
                split=split_str,
            )
            print(f"Loaded FinMarBa train split (sampled {sample_size:,} samples)")

        df = pl.DataFrame(ds.to_pandas())

        # Handle schema changes: newer versions have different column names
        # Map 'Title' -> 'text' and 'Global Sentiment' -> 'label'
        if "Title" in df.columns and "text" not in df.columns:
            df = df.rename({"Title": "text"})
        if "Global Sentiment" in df.columns and "label" not in df.columns:
            # Global Sentiment is -1/0/1, map to 0/1/2 for consistency
            # -1 (negative) -> 0, 0 (neutral) -> 1, 1 (positive) -> 2
            df = df.with_columns((pl.col("Global Sentiment") + 1).cast(pl.Int64).alias("label"))

        print(f"Loaded FinMarBa dataset: {len(df):,} samples")
        return df
    except Exception as e:
        raise RuntimeError(
            f"\n"
            f"{'=' * 70}\n"
            f"DATASET NOT AVAILABLE: FinMarBa\n"
            f"{'=' * 70}\n"
            f"\n"
            f"Error: {e}\n"
            f"\n"
            f"This notebook requires the FinMarBa dataset from Hugging Face:\n"
            f"  baptle/financial_headlines_market_based\n"
            f"\n"
            f"Possible causes:\n"
            f"  - No internet connection\n"
            f"  - HuggingFace servers unavailable\n"
            f"  - Dataset has been moved or renamed\n"
            f"\n"
            f"To run this notebook, ensure you have internet access.\n"
            f"{'=' * 70}\n"
        ) from e


df = load_finmarba_dataset(sample_size=SAMPLE_SIZE)
print("\nLabel distribution:")
print(df.group_by("label").len().sort("label"))

# Map labels to text
LABEL_MAP = {0: "negative", 1: "neutral", 2: "positive"}

# Load FinBERT pipeline
model_name = "ProsusAI/finbert"
print(f"Loading {model_name}...")

classifier = pipeline(
    "sentiment-analysis",
    model=model_name,
    tokenizer=model_name,
    device=device,
    truncation=True,
    max_length=512,
)

# ProsusAI/finbert label mapping (lowercase labels)
FINBERT_LABEL_MAP = {"negative": 0, "neutral": 1, "positive": 2}

def get_finbert_predictions(texts: list[str], batch_size: int = 32) -> np.ndarray:
    """Get predictions from FinBERT."""
    predictions = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i : i + batch_size]
        results = classifier(batch)
        for r in results:
            predictions.append(FINBERT_LABEL_MAP[r["label"]])
    return np.array(predictions)

# Get the text column (may be 'headline' or 'text' depending on dataset)
text_col = "headline" if "headline" in df.columns else "text"
texts = df[text_col].to_list()
true_labels = df["label"].to_numpy()

print(f"Running FinBERT on {len(texts)} samples...")
predictions = get_finbert_predictions(texts)

# Compute metrics
accuracy = accuracy_score(true_labels, predictions)
f1 = f1_score(true_labels, predictions, average="macro")

print("\n" + "=" * 60)
print("FinBERT on FinMarBa (OUT-OF-SAMPLE)")
print("=" * 60)
print(f"Accuracy: {accuracy:.1%}")
print(f"F1 (macro): {f1:.3f}")
print("=" * 60)

print("\n" + "=" * 70)
print("FinMarBa cross-domain accuracy for ProsusAI/finbert (zero-shot)")
print("=" * 70)
print(f"  Accuracy:    {accuracy:.1%}")
print(f"  F1 (macro):  {f1:.3f}")
print(f"  Test samples: {len(true_labels)}")
print("=" * 70)
print(
    "Reference for context (not measured here): Araci (2019) reports ~87% "
    "accuracy / ~0.85 macro F1 on Financial PhraseBank held-out test."
)

# Confusion matrix for FinMarBa
fig, ax = plt.subplots(figsize=(6, 5))
labels = ["negative", "neutral", "positive"]

cm = confusion_matrix(true_labels, predictions)
im = ax.imshow(cm, cmap="Blues")

# Add annotations; use light text on the dark high-count cells for legibility.
threshold = cm.max() / 2.0
for i in range(3):
    for j in range(3):
        ax.text(
            j,
            i,
            str(cm[i, j]),
            ha="center",
            va="center",
            fontsize=10,
            color="white" if cm[i, j] > threshold else "#0a1628",
        )

ax.set_xticks(range(3))
ax.set_yticks(range(3))
ax.set_xticklabels(labels)
ax.set_yticklabels(labels)
ax.set_xlabel("Predicted")
ax.set_ylabel("Actual")
ax.set_title(f"ProsusAI/finbert Cross-Dataset Evaluation\nFinMarBa Accuracy: {accuracy:.1%}")

plt.tight_layout()
plt.show()

# Structured output
print("=" * 70)
print("KEY INSIGHTS")
print("=" * 70)
print("\nDataset: FinMarBa")
print(f"Samples: {len(df)}")
print("\nFinBERT Performance (ProsusAI/finbert, zero-shot):")
print(f"  Accuracy:   {accuracy:.1%}")
print(f"  F1 (macro): {f1:.3f}")
print(
    "\nReference for context (not measured in this notebook): "
    "Araci (2019) reports ~87% accuracy / ~0.85 macro F1 on Financial "
    "PhraseBank held-out test."
)

# Save results
output_dir = get_chapter_dir(10) / "output" / "finbert_cross_dataset"
output_dir.mkdir(parents=True, exist_ok=True)
results_file = output_dir / "results.md"
with open(results_file, "w") as f:
    f.write("# FinBERT Cross-Dataset Evaluation Results\n\n")
    f.write("## Performance on FinMarBa (Out-of-Sample)\n\n")
    f.write(f"- **Accuracy**: {accuracy:.1%}\n")
    f.write(f"- **F1 (macro)**: {f1:.3f}\n")
    f.write(f"- **Samples**: {len(df):,}\n\n")
    f.write("## Reference for context (not measured in this notebook)\n\n")
    f.write(
        "Araci (2019) reports ~87% accuracy / ~0.85 macro F1 for "
        "ProsusAI/finbert on the held-out Financial PhraseBank test split.\n"
    )

print(f"\nResults saved to: {results_file}")