# ============================================================
# RESUMABLE WHISPER TRAINING SCRIPT WITH TIMESTAMP SUPPORT
# ============================================================

from dataclasses import dataclass
from typing import Any, Dict, List, Union
import os
import gc
import torch
import pandas as pd
import evaluate

from datasets import Dataset, Audio
from transformers import (
    WhisperForConditionalGeneration,
    WhisperProcessor,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
)

# ============================================================
# CONFIG
# ============================================================

MODEL_SIZE = "medium"
BASE_MODEL = f"neurlang/ipa-whisper-{MODEL_SIZE}"
OUTPUT_DIR = f"whisper-{MODEL_SIZE}-finetuned"

RESUME_TRAINING = False        # 🔁 flip to True to resume
RESUME_CHECKPOINT = "checkpoint-1840000"       # e.g. "checkpoint-40000"
RESUME_CHECKPOINT_TARGET = 1880000             # e.g.             80000

# don't tune this, it's auto tuned on training start/resume
lr = 1.251564455569462e-07 # 1e-5

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ============================================================
# LOAD DATA
#
# FORMAT:
# foo.mp3,hello world
#
# ============================================================

train_df = pd.read_csv("train.csv")
eval_df = pd.read_csv("test.csv")

train_df.columns = ["audio", "sentence"]
eval_df.columns = ["audio", "sentence"]

train_dataset = Dataset.from_pandas(train_df)
eval_dataset = Dataset.from_pandas(eval_df)

train_dataset = train_dataset.cast_column("audio", Audio(sampling_rate=16000))
eval_dataset = eval_dataset.cast_column("audio", Audio(sampling_rate=16000))

# Shuffle the dataset with exact seed control
train_dataset = train_dataset.shuffle(seed=42)  # Default shuffles all

# ============================================================
# PROCESSOR (TOKENIZER + FEATURE EXTRACTOR)
# ============================================================

if RESUME_TRAINING:
    print(f"🔁 Loading processor from {OUTPUT_DIR}")
    processor = WhisperProcessor.from_pretrained(OUTPUT_DIR)
else:
    print("🆕 Creating new processor")
    processor = WhisperProcessor.from_pretrained(
        BASE_MODEL,
        language="english",
        task="transcribe",
        predict_timestamps=True,
    )

    os.makedirs(OUTPUT_DIR, exist_ok=True)
    processor.save_pretrained(OUTPUT_DIR)  # 🔒 critical

# ============================================================
# DATA PREPARATION
# ============================================================

def prepare_dataset(batch):
    audio = batch["audio"]

    batch["input_features"] = processor.feature_extractor(
        audio["array"],
        sampling_rate=16000
    ).input_features[0]

    text = batch["sentence"] if batch["sentence"] else ""
    batch["labels"] = processor.tokenizer(
        text,
        return_tensors="pt"
    ).input_ids[0]

    del batch["audio"]
    del batch["sentence"]
    return batch

train_dataset = train_dataset.map(prepare_dataset, num_proc=1)
eval_dataset = eval_dataset.map(prepare_dataset, num_proc=1)

# ============================================================
# DATA COLLATOR
# ============================================================

@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    processor: Any

    def __call__(self, features):
        inputs = [{"input_features": f["input_features"]} for f in features]
        batch = self.processor.feature_extractor.pad(
            inputs, return_tensors="pt"
        )

        labels = [{"input_ids": f["labels"]} for f in features]
        labels_batch = self.processor.tokenizer.pad(
            labels, return_tensors="pt"
        )

        labels = labels_batch["input_ids"].masked_fill(
            labels_batch.attention_mask.ne(1), -100
        )

        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all():
            labels = labels[:, 1:]

        batch["labels"] = labels
        return batch

data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor)

# ============================================================
# METRICS
# ============================================================

cer_metric = evaluate.load("cer")

def compute_metrics(pred):
    pred_ids = pred.predictions
    label_ids = pred.label_ids

    label_ids[label_ids == -100] = processor.tokenizer.pad_token_id

    pred_str = processor.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    label_str = processor.tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    return {
        "cer": 100 * cer_metric.compute(
            predictions=pred_str,
            references=label_str
        )
    }

# ============================================================
# TRAINING ARGUMENTS
# ============================================================



training_args = Seq2SeqTrainingArguments(
    output_dir=OUTPUT_DIR,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=1,
    learning_rate=lr,
    warmup_steps=1000,
    max_steps=RESUME_CHECKPOINT_TARGET,
    evaluation_strategy="steps",
    save_strategy="steps",
    logging_steps=10*100,
    eval_steps=10*100,
    save_steps=10*100,
    save_total_limit=3,
    predict_with_generate=True,
    generation_max_length=225,
    fp16=False,
    report_to=["tensorboard"],
    load_best_model_at_end=False,
    metric_for_best_model="cer",
    greater_is_better=False,
    save_safetensors=True,   # 🔒 ensure safetensors
)

# ============================================================
# LOAD MODEL
# ============================================================

if RESUME_TRAINING:
    assert RESUME_CHECKPOINT is not None, "RESUME_CHECKPOINT must be set"

    checkpoint_path = os.path.join(OUTPUT_DIR, RESUME_CHECKPOINT)
    print(f"🔁 Loading model from {checkpoint_path}")

    model = WhisperForConditionalGeneration.from_pretrained(
        checkpoint_path,
        torch_dtype=torch.float32,
    )
else:
    print("🆕 Loading base model")
    model = WhisperForConditionalGeneration.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.float32,
    )

# 🔒 Modified safety check for Transformers version
# Transformers' Whisper uses different parameter naming
print(f"✅ Model loaded successfully")
print(f"   Model type: {type(model)}")
print(f"   Device: {next(model.parameters()).device}")

model.to(DEVICE)

# ============================================================
# TRAINER
# ============================================================

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

# ============================================================
# TRAIN
# ============================================================

torch.cuda.empty_cache()
gc.collect()

if RESUME_TRAINING:
    trainer.train(resume_from_checkpoint=checkpoint_path)
else:
    trainer.train()

# ============================================================
# SAVE FINAL
# ============================================================

trainer.save_model(OUTPUT_DIR)
processor.save_pretrained(OUTPUT_DIR)

print("✅ Training completed successfully")