"""Transformer Fine-Tuning - fine-tune FinBERT, DeBERTa-v3, and ModernBERT for financial sentiment classification."""

import json
import time
import warnings

import evaluate
import matplotlib.pyplot as plt
import numpy as np
import polars as pl
import seaborn as sns
import torch
from datasets import Dataset, DatasetDict
from sklearn.metrics import confusion_matrix
from sklearn.model_selection import train_test_split
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    EarlyStoppingCallback,
    Trainer,
    TrainingArguments,
)
from transformers import (
    set_seed as set_transformers_seed,
)

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

warnings.filterwarnings("ignore")

SEED = 42
MAX_TRAIN_STEPS = -1  # -1 = train full epochs

# Reproducibility - set_global_seeds covers Python random / NumPy / Torch.
# transformers Trainer uses its own RNG that needs explicit seeding.
set_global_seeds(SEED)
set_transformers_seed(SEED)

CONFIG = {
    "random_seed": SEED,
    "dataset": {
        "name": "takala/financial_phrasebank",
        "subset": "sentences_allagree",
        "test_size": 0.15,
        "val_size": 0.15,
    },
    "models": {
        "finbert": {
            "model_id": "ProsusAI/finbert",
            "description": "FinBERT fine-tuned on Financial PhraseBank sentiment",
        },
        "deberta": {
            "model_id": "microsoft/deberta-v3-small",
            "description": "DeBERTa-v3 with disentangled attention",
        },
        "modernbert": {
            "model_id": "answerdotai/ModernBERT-base",
            "description": "ModernBERT with 8,192 token context length",
        },
    },
    "training": {
        "learning_rate": 2e-5,
        "batch_size": 16,
        "num_epochs": 3,
        "weight_decay": 0.01,
        "warmup_ratio": 0.1,
        "max_length": 128,
        "early_stopping_patience": 2,
    },
}

print("=" * 70)
print("EXPERIMENT CONFIGURATION")
print("=" * 70)
print(json.dumps(CONFIG, indent=2))
print("=" * 70)

# Check for GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"\nUsing device: {device}")

df = load_financial_phrasebank(agreement="100")
print(f"Dataset size: {len(df):,}")

# Prepare train/val/test splits (convert to pandas for sklearn, then back to Polars)
df_pd = df.to_pandas()
train_pd, temp_pd = train_test_split(
    df_pd, test_size=0.3, random_state=SEED, stratify=df_pd["label"]
)
val_pd, test_pd = train_test_split(
    temp_pd, test_size=0.5, random_state=SEED, stratify=temp_pd["label"]
)

train_df = pl.from_pandas(train_pd)
val_df = pl.from_pandas(val_pd)
test_df = pl.from_pandas(test_pd)

print(f"Train: {len(train_df)}, Val: {len(val_df)}, Test: {len(test_df)}")

def create_dataset_dict(train_df: pl.DataFrame, val_df: pl.DataFrame, test_df: pl.DataFrame):
    return DatasetDict(
        {
            "train": Dataset.from_pandas(
                train_df.select(["sentence", "label"]).to_pandas(), preserve_index=False
            ),
            "validation": Dataset.from_pandas(
                val_df.select(["sentence", "label"]).to_pandas(), preserve_index=False
            ),
            "test": Dataset.from_pandas(
                test_df.select(["sentence", "label"]).to_pandas(), preserve_index=False
            ),
        }
    )


dataset = create_dataset_dict(train_df, val_df, test_df)

# Model configurations
MODELS = {
    "FinBERT": "ProsusAI/finbert",
    "DeBERTa-v3": "microsoft/deberta-v3-small",
    "ModernBERT": "answerdotai/ModernBERT-base",
}

# Label mappings
id2label = {0: "negative", 1: "neutral", 2: "positive"}
label2id = {"negative": 0, "neutral": 1, "positive": 2}

def tokenize_function(examples, tokenizer, max_length=128):
    return tokenizer(
        examples["sentence"],
        truncation=True,
        max_length=max_length,
        # Note: Don't pad here; use DataCollatorWithPadding for dynamic padding
    )

accuracy_metric = evaluate.load("accuracy")
f1_metric = evaluate.load("f1")


def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    acc = accuracy_metric.compute(predictions=predictions, references=labels)
    f1 = f1_metric.compute(predictions=predictions, references=labels, average="macro")
    return {"accuracy": acc["accuracy"], "f1": f1["f1"]}

# Training function
def fine_tune_model(model_name: str, model_path: str, dataset: DatasetDict) -> dict:
    """Fine-tune a transformer model for sentiment classification."""
    print(f"\n{'=' * 60}")
    print(f"Fine-tuning: {model_name}")
    print(f"Model: {model_path}")
    print(f"{'=' * 60}")

    # Load tokenizer and model
    tokenizer = AutoTokenizer.from_pretrained(model_path)

    # Handle models without pad token (proper approach for encoder-only models)
    # Encoder models like BERT don't have eos_token; use [PAD] or add one
    if tokenizer.pad_token is None:
        if tokenizer.eos_token is not None:
            tokenizer.pad_token = tokenizer.eos_token
        else:
            # For BERT-like models, add [PAD] token
            tokenizer.add_special_tokens({"pad_token": "[PAD]"})

    model = AutoModelForSequenceClassification.from_pretrained(
        model_path,
        num_labels=3,
        id2label=id2label,
        label2id=label2id,
        ignore_mismatched_sizes=True,
    )

    # Resize embeddings if we added a new token
    model.resize_token_embeddings(len(tokenizer))

    # Handle pad token in model config
    if model.config.pad_token_id is None:
        model.config.pad_token_id = tokenizer.pad_token_id

    # Create data collator for dynamic padding (more efficient than max_length padding)
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    # Tokenize dataset
    tokenized = dataset.map(
        lambda x: tokenize_function(x, tokenizer),
        batched=True,
        remove_columns=["sentence"],
    )

    # Training arguments - save checkpoints under chapter output directory
    chapter_dir = get_chapter_dir(10)
    output_dir = chapter_dir / "output" / "bert_finetuning" / model_name.lower().replace("-", "_")
    output_dir.mkdir(parents=True, exist_ok=True)

    # Use CONFIG values consistently
    train_config = CONFIG["training"]
    num_epochs = train_config["num_epochs"]
    batch_size = train_config["batch_size"]
    train_size = len(tokenized["train"])

    max_steps = MAX_TRAIN_STEPS

    # Build training arguments dict with version-compatible parameter names
    # transformers 4.36+ uses eval_strategy, older versions use evaluation_strategy
    eval_strat_key = "eval_strategy"  # Default to newer API
    save_strat_key = "save_strategy"
    try:
        # Test if the newer API is available
        import inspect

        sig = inspect.signature(TrainingArguments)
        if "evaluation_strategy" in sig.parameters and "eval_strategy" not in sig.parameters:
            eval_strat_key = "evaluation_strategy"
            save_strat_key = "save_strategy"
    except Exception:
        pass  # Use defaults

    # Compute eval_steps and save_steps ensuring save_steps is a multiple of eval_steps
    # (required by load_best_model_at_end)
    if max_steps > 0:
        eval_steps = max(10, max_steps // 5)
        # save_steps must be a multiple of eval_steps
        save_steps = eval_steps * 2  # evaluate twice, save once
    else:
        eval_steps = None
        save_steps = None

    training_kwargs = {
        "output_dir": str(output_dir),
        "num_train_epochs": num_epochs,
        "max_steps": max_steps,
        "per_device_train_batch_size": batch_size,
        "per_device_eval_batch_size": batch_size * 2,
        "warmup_steps": min(50, max(10, train_size // batch_size // 4)),
        "weight_decay": train_config["weight_decay"],
        "logging_steps": max(10, train_size // batch_size // 3),
        eval_strat_key: "steps" if max_steps > 0 else "epoch",
        "eval_steps": eval_steps,
        save_strat_key: "steps" if max_steps > 0 else "epoch",
        "save_steps": save_steps,
        "load_best_model_at_end": True,
        "metric_for_best_model": "f1",
        "greater_is_better": True,
        "report_to": "none",  # Disable wandb/tensorboard
        "fp16": torch.cuda.is_available(),
    }

    training_args = TrainingArguments(**training_kwargs)

    # Create trainer with data collator for dynamic padding
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized["train"],
        eval_dataset=tokenized["validation"],
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
    )

    # Train
    start_time = time.time()
    train_result = trainer.train()
    train_time = time.time() - start_time

    # Evaluate on test set
    test_results = trainer.evaluate(tokenized["test"])

    # Get predictions for confusion matrix
    predictions = trainer.predict(tokenized["test"])
    y_pred = np.argmax(predictions.predictions, axis=-1)
    y_true = predictions.label_ids

    return {
        "model_name": model_name,
        "accuracy": test_results["eval_accuracy"],
        "f1": test_results["eval_f1"],
        "train_time": train_time,
        "num_params": sum(p.numel() for p in model.parameters()),
        "y_pred": y_pred,
        "y_true": y_true,
        "train_loss": train_result.training_loss,
    }

results = {name: fine_tune_model(name, path, dataset) for name, path in MODELS.items()}

# Summary table
summary_df = pl.DataFrame(
    [
        {
            "Model": r["model_name"],
            "Accuracy": f"{r['accuracy']:.1%}",
            "F1 (macro)": f"{r['f1']:.3f}",
            "Parameters": f"{r['num_params'] / 1e6:.1f}M",
            "Train Time": f"{r['train_time']:.0f}s",
        }
        for r in results.values()
    ]
)

summary_df

# Visualization: Performance comparison
fig, axes = plt.subplots(1, 2, figsize=(12, 5))

# Accuracy/F1 comparison
ax = axes[0]
models = list(results.keys())
accuracies = [results[m]["accuracy"] for m in models]
f1_scores = [results[m]["f1"] for m in models]

x = np.arange(len(models))
width = 0.35

ax.bar(x - width / 2, accuracies, width, label="Accuracy", color="#0a1628")
ax.bar(x + width / 2, f1_scores, width, label="F1 (macro)", color="#D4A84B")

ax.set_ylabel("Score")
ax.set_title("Fine-Tuned Model Performance on Financial PhraseBank")
ax.set_xticks(x)
ax.set_xticklabels(models)
ax.legend()
ax.set_ylim(0, 1)

# Training time comparison
ax = axes[1]
times = [results[m]["train_time"] for m in models]
bars = ax.bar(models, times, color="#1a2d4a")
ax.set_ylabel("Training Time (seconds)")
ax.set_title("Training Time on RTX 3090")

for bar in bars:
    height = bar.get_height()
    ax.annotate(
        f"{height:.0f}s",
        xy=(bar.get_x() + bar.get_width() / 2, height),
        xytext=(0, 3),
        textcoords="offset points",
        ha="center",
        va="bottom",
    )

plt.show()

# Confusion matrices for each model
n_models = len(results)
fig, axes = plt.subplots(1, n_models, figsize=(5 * n_models, 4))
fig.suptitle("Fine-Tuned Model Confusion Matrices", fontsize=12, y=1.02)

labels = ["negative", "neutral", "positive"]

for ax, (name, r) in zip(axes, results.items(), strict=False):
    cm = confusion_matrix(r["y_true"], r["y_pred"])
    sns.heatmap(
        cm,
        annot=True,
        fmt="d",
        cmap="Blues",
        xticklabels=labels,
        yticklabels=labels,
        ax=ax,
    )
    ax.set_title(f"{name}\nAcc: {r['accuracy']:.1%}")
    ax.set_xlabel("Predicted")
    ax.set_ylabel("Actual")

plt.show()

# Structured per-model summary for chapter cross-reference
for r in results.values():
    print(f"\n{r['model_name']}:")
    print(f"  Accuracy: {r['accuracy']:.1%}")
    print(f"  F1 (macro): {r['f1']:.3f}")
    print(f"  Parameters: {r['num_params'] / 1e6:.1f}M")
    print(f"  Train time: {r['train_time']:.0f}s")