import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
import numpy as np
from huggingface_hub import hf_hub_download
from datasets import load_dataset, Audio # On garde l'import au cas où

# =============================================================================
# PARTIE 1 : DÉFINITION DE L'ARCHITECTURE (inchangée)
# =============================================================================

NUM_CLASSES = 2
N_MELS = 128

class AudioClassifier(nn.Module):
    """Réseau de Neurones Convolutionnels (CNN) que nous avons entraîné."""
    def __init__(self):
        super(AudioClassifier, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=(5, 5), padding=2)
        self.bn1 = nn.BatchNorm2d(32)
        self.pool1 = nn.MaxPool2d(kernel_size=(2, 2))
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(3, 3), padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.pool2 = nn.MaxPool2d(kernel_size=(2, 2))
        self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(3, 3), padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.pool3 = nn.MaxPool2d(kernel_size=(2, 2))
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc1 = nn.Linear(128 * 1 * 1, NUM_CLASSES)

    def forward(self, x):
        x = self.pool1(F.relu(self.bn1(self.conv1(x))))
        x = self.pool2(F.relu(self.bn2(self.conv2(x))))
        x = self.pool3(F.relu(self.bn3(self.conv3(x))))
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return self.fc1(x)


# =============================================================================
# PARTIE 2 : FONCTIONS DE PRÉPARATION POUR L'INFÉRENCE (inchangée)
# =============================================================================

SAMPLING_RATE = 16000
N_FFT = 400
HOP_LENGTH = 160
MAX_TIME_STEPS = 300

def prepare_spectrogram(audio_path):
    """
    Charge un fichier audio, calcule le Log-Mel Spectrogramme, et le prépare
    pour le modèle.
    """
    # 1. Charger et Rééchantillonner
    waveform, sr = torchaudio.load(audio_path)
    if sr != SAMPLING_RATE:
        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=SAMPLING_RATE)
        waveform = resampler(waveform)

    if waveform.shape[0] > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)

    # 2. Calculer le Log-Mel Spectrogramme
    mel_spectrogram_transform = torchaudio.transforms.MelSpectrogram(
        sample_rate=SAMPLING_RATE,
        n_fft=N_FFT,
        hop_length=HOP_LENGTH,
        n_mels=N_MELS,
    )
    mel_spectrogram = mel_spectrogram_transform(waveform.squeeze(0))
    log_mel_spectrogram = torchaudio.transforms.AmplitudeToDB()(mel_spectrogram)

    # 3. Tronquer
    if log_mel_spectrogram.shape[1] > MAX_TIME_STEPS:
        log_mel_spectrogram = log_mel_spectrogram[:, :MAX_TIME_STEPS]

    # 4. Ajouter les dimensions Batch et Channel : [1, 1, N_MELS, Time_Steps]
    input_tensor = log_mel_spectrogram.unsqueeze(0).unsqueeze(0)

    return input_tensor


def predict_audio(model, audio_tensor):
    """
    Effectue la prédiction et retourne l'étiquette.
    """
    model.eval()
    device = next(model.parameters()).device

    with torch.no_grad():
        audio_tensor = audio_tensor.to(device)
        outputs = model(audio_tensor)

        probabilities = F.softmax(outputs, dim=1)
        predicted_index = torch.argmax(probabilities, dim=1).item()

        # Décodeur des classes (assumant 0 = Calé, 1 = Décalé)
        class_labels = {0: "Calé (On Beat)", 1: "Décalé (Off Beat)"}

        return class_labels[predicted_index], probabilities[0].cpu().numpy()


# =============================================================================
# PARTIE 3 : CHARGEMENT DU MODÈLE ET EXÉCUTION DU TEST (CORRIGÉE)
# =============================================================================

# --- Configuration Hugging Face ---
REPO_ID = "Clemylia/Musica1"
MODEL_FILENAME = "pytorch_model.bin"

# 1. Télécharger les poids du modèle
print(f"1. Téléchargement des poids du modèle depuis {REPO_ID}...")
try:
    model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
except Exception as e:
    print(f"Erreur de téléchargement : {e}. Vérifiez le nom du dépôt et les permissions.")
    exit()

# 2. Charger le modèle
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AudioClassifier()

try:
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    print(f"2. Modèle chargé avec succès sur {device}.")
except Exception as e:
    print(f"Erreur lors du chargement des poids : {e}")
    exit()

# 3. Préparer une donnée de test

# *** CORRECTION MAJEURE ***
# Veuillez remplacer le chemin ci-dessous par un chemin valide sur votre système !
# Exemple : "C:/Users/Clemylia/Desktop/mes_sons/calé_test.wav" ou "./data/audio_test.mp3"
# ----------------------------------------------------------------------------------
AUDIO_TEST_FILE = input("Veuillez entrer le chemin complet d'un fichier audio (ex: /path/to/test.wav): ")
# ----------------------------------------------------------------------------------

print(f"\n3. Préparation d'un échantillon de test à partir de: {AUDIO_TEST_FILE}...")

try:
    input_tensor = prepare_spectrogram(AUDIO_TEST_FILE)
    audio_test_path = AUDIO_TEST_FILE # Pour l'affichage final

except Exception as e:
    print(f"Erreur lors de la préparation de l'échantillon de test (le fichier existe-t-il ? le format est-il pris en charge par torchaudio ?) : {e}")
    input_tensor = None


# 4. Exécuter la prédiction
if input_tensor is not None:
    print(f"\n4. Exécution de la prédiction sur l'échantillon...")

    prediction, probabilities = predict_audio(model, input_tensor)

    # Affichage des résultats
    print("\n-------------------------------------------")
    print(f"FICHIER TESTÉ: {audio_test_path}")
    print(f"PRÉDICTION: {prediction}")
    print(f"PROBABILITÉS: Calé={probabilities[0]:.4f}, Décalé={probabilities[1]:.4f}")
    print("-------------------------------------------")
else:
    print("Test annulé faute de pouvoir traiter le fichier audio.")