import json
import os
import random
import subprocess
import sys

import numpy as np
import torch
from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    DataCollatorForSeq2Seq,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    set_seed,
)


# ============================================================
# CONFIGURAÇÕES
# ============================================================

MODEL_NAME = "google/byt5-small"

TRAIN_FILE = "data/train.jsonl"
VALIDATION_FILE = "data/validation.jsonl"

# Pasta diferente para não sobrescrever os testes anteriores
OUTPUT_DIR = "models/seeds-byt5-overfit-v2"
FINAL_MODEL_DIR = "models/seeds-byt5-overfit-v2-final"

# ------------------------------------------------------------
# TESTE CONTROLADO DE MEMORIZAÇÃO
# ------------------------------------------------------------

MAX_TRAIN_EXAMPLES = 20
MAX_VALIDATION_EXAMPLES = 20

# Os exemplos selecionados são curtos, então é possível aumentar
# o limite sem sofrer com truncamento.
MAX_SOURCE_LENGTH = 256
MAX_TARGET_LENGTH = 256

NUM_EPOCHS = 10
LEARNING_RATE = 1e-4

TRAIN_BATCH_SIZE = 1
EVAL_BATCH_SIZE = 1


# Cada exemplo gera uma atualização do otimizador.
GRADIENT_ACCUMULATION_STEPS = 1

SEED = 42


# ============================================================
# FUNÇÕES AUXILIARES
# ============================================================

def carregar_jsonl(caminho):
    with open(caminho, "r", encoding="utf-8") as arquivo:
        return [json.loads(linha) for linha in arquivo]


def garantir_dataset():

    if os.path.exists(TRAIN_FILE) and os.path.exists(VALIDATION_FILE):
        return

    print("Arquivos de dataset não encontrados.")
    print("Executando automaticamente o script 02_prepare_dataset.py...")

    resultado = subprocess.run(
        [sys.executable, "src/02_prepare_dataset.py"],
        capture_output=True,
        text=True,
    )

    print(resultado.stdout)

    if resultado.returncode != 0:
        print(resultado.stderr)

        raise RuntimeError(
            "Não foi possível preparar o dataset automaticamente."
        )


def preparar_exemplos():

    train_data = carregar_jsonl(TRAIN_FILE)
    validation_data = carregar_jsonl(VALIDATION_FILE)

    print(f"Treino completo: {len(train_data)} exemplos")
    print(f"Validação completa: {len(validation_data)} exemplos")

    # --------------------------------------------------------
    # Selecionado os exemplos mais curtos.
    #
    # Isso torna o teste de memorização mais limpo:
    # o modelo não precisa lidar com frases enormes nem
    # truncamento.
    # --------------------------------------------------------

    train_data = sorted(
        train_data,
        key=lambda item: len(item["source"]) + len(item["target"])
    )[:MAX_TRAIN_EXAMPLES]

    validation_data = sorted(
        validation_data,
        key=lambda item: len(item["source"]) + len(item["target"])
    )[:MAX_VALIDATION_EXAMPLES]

    print(f"Treino selecionado: {len(train_data)} exemplos")
    print(f"Validação selecionada: {len(validation_data)} exemplos")

    print("\nExemplos selecionados para treino:")

    for indice, exemplo in enumerate(train_data[:5], start=1):

        print(
            f"{indice}. "
            f"source_chars={len(exemplo['source'])}, "
            f"target_chars={len(exemplo['target'])}"
        )

        print(f"   EN: {exemplo['source']}")
        print(f"   MBC: {exemplo['target']}")
        print()

    return train_data, validation_data


def criar_dataset(exemplos):

    return Dataset.from_list(
        [
            {
                "source": exemplo["source"],
                "target": exemplo["target"],
            }
            for exemplo in exemplos
        ]
    )


# ============================================================
# PREPARAÇÃO
# ============================================================

print("=" * 70)
print("SEEDS — TESTE FINAL DE MEMORIZAÇÃO")
print("=" * 70)

garantir_dataset()

set_seed(SEED)

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

train_data, validation_data = preparar_exemplos()

train_dataset = criar_dataset(train_data)
validation_dataset = criar_dataset(validation_data)


# ============================================================
# TOKENIZER
# ============================================================

print("\nCarregando tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME
)


# ============================================================
# MODELO
# ============================================================

print("Carregando modelo...")

model = AutoModelForSeq2SeqLM.from_pretrained(
    MODEL_NAME
)

# Evita incompatibilidade entre cache e
# gradient checkpointing.
model.config.use_cache = False

# Evita que o Transformers tente amarrar novamente
# shared.weight e lm_head.weight.
model.config.tie_word_embeddings = False


# ============================================================
# TOKENIZAÇÃO
# ============================================================

def preprocess_function(exemplos):

    prompts = [
        f"English: {texto}\nMacushi:"
        for texto in exemplos["source"]
    ]

    targets = exemplos["target"]

    # Entrada
    model_inputs = tokenizer(
        prompts,
        max_length=MAX_SOURCE_LENGTH,
        truncation=True,
    )

    # Saída esperada
    labels = tokenizer(
        text_target=targets,
        max_length=MAX_TARGET_LENGTH,
        truncation=True,
    )

    model_inputs["labels"] = labels["input_ids"]

    return model_inputs


print("\nTokenizando dataset de treino...")

tokenized_train = train_dataset.map(
    preprocess_function,
    batched=True,
    remove_columns=train_dataset.column_names,
)


print("Tokenizando dataset de validação...")

tokenized_validation = validation_dataset.map(
    preprocess_function,
    batched=True,
    remove_columns=validation_dataset.column_names,
)


# ============================================================
# VERIFICAÇÃO DOS TAMANHOS
# ============================================================

print("\nComprimentos dos primeiros exemplos tokenizados:")

for indice in range(min(5, len(tokenized_train))):

    print(
        f"Exemplo {indice + 1}: "
        f"source={len(tokenized_train[indice]['input_ids'])}, "
        f"target={len(tokenized_train[indice]['labels'])}"
    )


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

data_collator = DataCollatorForSeq2Seq(
    tokenizer=tokenizer,
    model=model,
    padding=True,
)


# ============================================================
# ARGUMENTOS DE TREINAMENTO
# ============================================================

training_args = Seq2SeqTrainingArguments(

    output_dir=OUTPUT_DIR,

    # --------------------------------------------------------
    # ÉPOCAS
    # --------------------------------------------------------

    num_train_epochs=NUM_EPOCHS,

    # --------------------------------------------------------
    # BATCH
    # --------------------------------------------------------

    per_device_train_batch_size=TRAIN_BATCH_SIZE,
    per_device_eval_batch_size=EVAL_BATCH_SIZE,

    # --------------------------------------------------------
    # GRADIENT ACCUMULATION
    # --------------------------------------------------------

    # Cada exemplo gera uma atualização.
    gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,

    # --------------------------------------------------------
    # LEARNING RATE
    # --------------------------------------------------------

    learning_rate=LEARNING_RATE,

    # Mantem o learning rate constante durante o teste.
    lr_scheduler_type="constant",
    warmup_steps=0,

    # Não queremos regularização adicional nesse teste.
    weight_decay=0.0,

    # --------------------------------------------------------
    # LOGS
    # --------------------------------------------------------

    logging_steps=1,

    # --------------------------------------------------------
    # AVALIAÇÃO E CHECKPOINTS
    # --------------------------------------------------------

    eval_strategy="epoch",
    save_strategy="epoch",

    save_total_limit=2,

    # Não é preciso gerar durante cada avaliação.
    # A geração será feita manualmente no final.
    predict_with_generate=False,

    # --------------------------------------------------------
    # MEMÓRIA
    # --------------------------------------------------------

    gradient_checkpointing=True,

    use_cpu=True,

    # Não enviar métricas para serviços externos.
    report_to="none",

    seed=SEED,
)


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

trainer = Seq2SeqTrainer(

    model=model,

    args=training_args,

    train_dataset=tokenized_train,

    eval_dataset=tokenized_validation,

    processing_class=tokenizer,

    data_collator=data_collator,
)


# ============================================================
# TREINAMENTO
# ============================================================

print("\n" + "=" * 70)
print("INICIANDO TREINAMENTO")
print("=" * 70)

print(f"Exemplos de treino: {len(tokenized_train)}")

print(f"Épocas: {NUM_EPOCHS}")

print(
    "Atualizações por época: "
    f"aproximadamente {len(tokenized_train)}"
)

print(
    "Atualizações totais: "
    f"aproximadamente {len(tokenized_train) * NUM_EPOCHS}"
)

print(f"Learning rate: {LEARNING_RATE}")

print(
    "Gradient accumulation: "
    f"{GRADIENT_ACCUMULATION_STEPS}"
)


train_result = trainer.train()


print("\nTreinamento concluído.")

print("\nMétricas finais de treino:")

print(train_result.metrics)


# ============================================================
# SALVAMENTO
# ============================================================

print("\nSalvando modelo final...")

trainer.save_model(
    FINAL_MODEL_DIR
)

tokenizer.save_pretrained(
    FINAL_MODEL_DIR
)

print(
    f"Modelo salvo em: {FINAL_MODEL_DIR}"
)


# ============================================================
# RECARREGAMENTO
# ============================================================

print("\nRecarregando modelo salvo...")

final_tokenizer = AutoTokenizer.from_pretrained(
    FINAL_MODEL_DIR
)

final_model = AutoModelForSeq2SeqLM.from_pretrained(
    FINAL_MODEL_DIR
)

final_model.config.use_cache = True

print("Modelo recarregado com sucesso.")


# ============================================================
# GERAÇÃO
# ============================================================

def gerar_traducao(texto, num_beams):

    prompt = f"English: {texto}\nMacushi:"

    inputs = final_tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=MAX_SOURCE_LENGTH,
    )

    with torch.no_grad():

        output_ids = final_model.generate(
            **inputs,

            max_length=MAX_TARGET_LENGTH,

            num_beams=num_beams,

            early_stopping=True,

            # Evita repetições muito óbvias.
            no_repeat_ngram_size=3,
        )

    return final_tokenizer.decode(
        output_ids[0],
        skip_special_tokens=True,
    )


# ============================================================
# TESTE DE MEMORIZAÇÃO
# ============================================================

print("\n" + "=" * 70)
print("TESTE DE MEMORIZAÇÃO")
print("=" * 70)

print(
    "\nO modelo será testado nos próprios exemplos "
    "que foram utilizados durante o treinamento."
)

print(
    "Isso NÃO é uma avaliação de generalização."
)

print(
    "O objetivo aqui é verificar se o modelo consegue "
    "simplesmente aprender os pares apresentados."
)


for indice, exemplo in enumerate(
    train_data[:5],
    start=1
):

    print("\n" + "-" * 70)

    print(f"EXEMPLO {indice}")

    print(
        f"\nEN esperado:\n"
        f"{exemplo['source']}"
    )

    print(
        f"\nMBC esperado:\n"
        f"{exemplo['target']}"
    )


    # --------------------------------------------------------
    # GREEDY
    # --------------------------------------------------------

    traducao_gulosa = gerar_traducao(
        exemplo["source"],
        num_beams=1,
    )


    # --------------------------------------------------------
    # BEAM SEARCH
    # --------------------------------------------------------

    traducao_beam = gerar_traducao(
        exemplo["source"],
        num_beams=4,
    )


    print(
        f"\nGerado — greedy:\n"
        f"{traducao_gulosa}"
    )

    print(
        f"\nGerado — beam 4:\n"
        f"{traducao_beam}"
    )


# ============================================================
# FIM
# ============================================================

print("\n" + "=" * 70)
print("TESTE FINALIZADO")
print("=" * 70)

print(
    "\nSe o modelo ainda não conseguir reproduzir "
    "os próprios exemplos de treinamento, "
    "vamos considerar este teste suficiente "
    "para abandonar essa abordagem e partir "
    "para outra estratégia."
)