import gradio as gr
import matplotlib.pyplot as plt
import librosa
import torch
import librosa.display
import numpy as np
import torch.nn as nn
import os

# ---- CR2NN Model ----
class CR2NNModel(nn.Module):
    def __init__(self, num_classes=2, lr=1e-3):
        super(CR2NNModel, self).__init__()
        self.lr = lr
        
        # Convolutional feature extractor:
        # Input shape: (batch, 1, n_mels, time)
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),  # (batch, 32, n_mels, time)
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(2, 2)),                        # (batch, 32, n_mels/2, time/2)
            
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),   # (batch, 64, n_mels/2, time/2)
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(2, 2))                         # (batch, 64, n_mels/4, time/4)
        )
        # For example, if n_mels=256 and fixed_length (time dimension) = 1400,
        # after pooling the output becomes roughly (batch, 64, 64, 350).
        # Averaging over the frequency dimension (n_mels/4) yields a sequence of length 350 with 64 features.
        
        # GRU: Now stacking two layers instead of one.
        # Using num_layers=2 applies dropout (if specified) between layers.
        self.gru = nn.GRU(input_size=64, hidden_size=128, num_layers=2,
                          batch_first=True, bidirectional=True, dropout=0.5)
        # The output of a bidirectional GRU with hidden_size=128 is 256 per time step.
        self.fc = nn.Linear(128 * 2, num_classes)  # Maps 256 features to num_classes

    def forward(self, x):
        # x: (batch, 1, n_mels, time)
        conv_out = self.conv(x)             # (batch, 64, n_mels/4, time/4)
        conv_out = conv_out.mean(dim=2)       # Average over the frequency axis: (batch, 64, time/4)
        conv_out = conv_out.permute(0, 2, 1)    # (batch, time/4, 64)
        gru_out, _ = self.gru(conv_out)       # (batch, time/4, 256)
        last_out = gru_out[:, -1, :]          # Use the output from the last time step: (batch, 256)
        logits = self.fc(last_out)            # (batch, num_classes)
        return logits

    
# Function to preprocess the audio file
def preprocess_audio(file_path, sr=22050, fixed_length=1400):
    # Load audio file
    # audio, _ = load_audio_with_skip(file_path, sr=sr, skip_seconds = 5)
    
    audio, _ = librosa.load(file_path, sr=sr)

    # Generate Mel spectrogram
    #mel_spec = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=128, fmax=8000)
    mel_spec = librosa.feature.melspectrogram(
        y=audio, 
        sr=sr, 
        n_fft=4096,        # Higher frequency resolution
        hop_length=256,    # Higher time resolution
        n_mels=256,        # More Mel bins for finer frequency detail
        fmax=8000
    )
    mel_spec_db = librosa.power_to_db(mel_spec, ref=np.max)

    # Normalize the spectrogram
    mel_spec_db = (mel_spec_db - mel_spec_db.min()) / (mel_spec_db.max() - mel_spec_db.min())

    # Pad or truncate to fixed length
    if mel_spec_db.shape[1] < fixed_length:
        pad_width = fixed_length - mel_spec_db.shape[1]
        mel_spec_db = np.pad(mel_spec_db, ((0, 0), (0, pad_width)), mode='constant')
    else:
        mel_spec_db = mel_spec_db[:, :fixed_length]

    # Add channel dimension and convert to PyTorch tensor
    mel_spec_db = np.expand_dims(mel_spec_db, axis=0)  # Add channel dimension
    mel_spec_db = torch.tensor(mel_spec_db, dtype=torch.float32).unsqueeze(0)  # Add batch dimension
    return mel_spec_db

# Function to classify the audio file
def classify_audio(file_path, model, device):
    # Preprocess the audio file
    input_tensor = preprocess_audio(file_path).to(device)

    # Perform inference
    model.eval()
    with torch.no_grad():
        logits = model(input_tensor)
        probabilities = torch.nn.functional.softmax(logits, dim=1)
    human_percent = probabilities[0,0].item() * 100
    ai_percent = probabilities[0,1].item() * 100

    # Map prediction to labels

    return {"Human" : human_percent, "AI" : ai_percent}

def load_model():
    # Construct the model file path using a relative path
    model_path = os.path.join(os.path.dirname(__file__), "model", "model.pth")

    # Load the model (adjust loading code as needed for your model type)
    model = CR2NNModel(num_classes=2)
    model.load_state_dict(torch.load(model_path, map_location="cpu"))
    
    model.eval()  # Set the model to evaluation mode
    return model

def analyze_mp3(file_path):
    # Load the MP3 file. This returns the audio data and sampling rate.
    try:
        audio, sr = librosa.load(file_path, sr=None)
        duration = librosa.get_duration(y=audio, sr=sr)
    except Exception as e:
        # If there is an error loading the file, assume a default duration.
        duration = 0

    # load the model
    model = load_model()

    
    # Replace this with your own logic/model to determine the percentages.
    result = classify_audio(file_path, model, "cpu")
    ai_percentage = result['AI']
    human_percentage = result['Human']

    # Create a pie chart with fixed labels: "AI" and "Human"
    labels = ['AI', 'Human']
    sizes = [ai_percentage, human_percentage]

    fig, ax = plt.subplots()
    ax.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=90)
    ax.axis('equal')  # Ensure that the pie is drawn as a circle.

    plt.close(fig)
    
    return fig

# Create the Gradio interface
iface = gr.Interface(
    fn=analyze_mp3,
    inputs=gr.Audio(type="filepath", label="Upload MP3"),
    outputs=gr.Plot(label="Analysis Pie Chart"),
    title="MP3 Analysis Pie Chart",
    description="Upload an MP3 file to generate a pie chart with AI and Human labels."
)

iface.launch()