# @title 4. Executar Treinamento (Fine-Tuning)
import os
import torch
import librosa
from torch.utils.data import Dataset
from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification, Trainer, TrainingArguments

# Configurações do Modelo
BASE_MODEL = "HyperMoon/wav2vec2-base-960h-finetuned-deepfake"
OUTPUT_DIR = "local_finetuned_model"

# Mapeamento Rígido de Labels para evitar conflitos (0=Real, 1=Fraude)
id2label = {0: "AUTHENTIC", 1: "FAKE"}
label2id = {"AUTHENTIC": 0, "FAKE": 1}

class DeepfakeDataset(Dataset):
    def __init__(self, root_dir, processor):
        self.files = []
        self.processor = processor
        
        # Carregamento explícito baseado em pastas
        for label_name, label_id in label2id.items():
            folder = "real" if label_name == "AUTHENTIC" else "fake"
            path = os.path.join(root_dir, folder)
            if os.path.exists(path):
                print(f"Carregando audios de: {folder}...")
                for f in os.listdir(path):
                    if f.lower().endswith(('.wav', '.mp3', '.flac')):
                        self.files.append({"path": os.path.join(path, f), "label": label_id})
            else:
                print(f"AVISO: Pasta {folder} não encontrada em {root_dir}")

    def __len__(self): return len(self.files)
    
    def __getitem__(self, idx):
        item = self.files[idx]
        try:
            speech, _ = librosa.load(item["path"], sr=16000)
            inputs = self.processor(speech, sampling_rate=16000, return_tensors="pt", padding="max_length", max_length=160000, truncation=True)
            return {"input_values": inputs.input_values[0], "labels": torch.tensor(item["label"])}
        except Exception as e:
            print(f"Erro ao processar {item['path']}: {e}")
            # Retorna o primeiro item como fallback para não quebrar o loop do Trainer
            return self.__getitem__(0)

print("Inicializando Processador e Modelo...")
try:
    processor = Wav2Vec2FeatureExtractor.from_pretrained(BASE_MODEL)
    # Adicionado id2label e label2id aqui para garantir consistência
    model = Wav2Vec2ForSequenceClassification.from_pretrained(
        BASE_MODEL, 
        num_labels=2, 
        id2label=id2label,
        label2id=label2id,
        ignore_mismatched_sizes=True
    )

    # Congelar base para focar no aprendizado das novas fraudes (Lógica Robusta)
    if hasattr(model, 'wav2vec2'):
        for param in model.wav2vec2.parameters():
            param.requires_grad = False
    print("Modelo carregado e camadas base congeladas com sucesso!")

    # Dataset (Aponte para a pasta onde você subiu os áudios no Colab)
    # Ex: /content/dataset_treino
    dataset_path = "/content/dataset" 
    train_data = DeepfakeDataset(dataset_path, processor)

    if len(train_data) == 0:
        print("ERRO: Nenhum dado encontrado. Verifique se as pastas 'real' e 'fake' existem dentro do caminho especificado.")
    else:
        training_args = TrainingArguments(
            output_dir=OUTPUT_DIR,
            num_train_epochs=3,
            per_device_train_batch_size=2,
            gradient_accumulation_steps=4,
            save_steps=50,
            logging_steps=10,
            learning_rate=2e-5,
            remove_unused_columns=False
        )

        trainer = Trainer(model=model, args=training_args, train_dataset=train_data)
        print("Iniciando Treinamento...")
        trainer.train()

        # Salva o resultado final
        model.save_pretrained(OUTPUT_DIR)
        processor.save_pretrained(OUTPUT_DIR)
        print(f"Sucesso! Modelo salvo em: {OUTPUT_DIR}")

except Exception as e:
    print(f"ERRO CRÍTICO: {e}")
    print("DICA: Se o erro for de conexão, tente rodar a célula novamente. O Hugging Face pode falhar ocasionalmente no download.")