import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
import librosa
import numpy as np
import math
import os
from transformers import Wav2Vec2Model, Wav2Vec2Config, Wav2Vec2FeatureExtractor
from huggingface_hub import hf_hub_download

# ==========================================
# 1. මොඩලයේ ව්‍යුහය (Architecture)
# ==========================================
class SelfAttentionPooling(nn.Module):
    def __init__(self, input_dim):
        super(SelfAttentionPooling, self).__init__()
        self.W = nn.Linear(input_dim, 128)
        self.V = nn.Linear(128, 1)

    def forward(self, x, attention_mask=None):
        scores = self.V(torch.tanh(self.W(x)))
        if attention_mask is not None:
            indices = torch.linspace(0, attention_mask.size(1) - 1, steps=x.size(1)).long().to(x.device)
            mask = torch.index_select(attention_mask, 1, indices).unsqueeze(-1)
            scores = scores.masked_fill(mask == 0, -1e4)
        attn_weights = F.softmax(scores, dim=1)
        return torch.sum(x * attn_weights, dim=1), attn_weights

class SinhalaPhonoNet(nn.Module):
    # 🌟 num_classes=255 ලෙස සකසා ඇත
    def __init__(self, base_model="facebook/wav2vec2-xls-r-300m", embedding_dim=256, num_classes=255):
        super(SinhalaPhonoNet, self).__init__()
        self.config = Wav2Vec2Config.from_pretrained(base_model, output_hidden_states=True)
        self.backbone = Wav2Vec2Model.from_pretrained(base_model, config=self.config)
        self.layer_weights = nn.Parameter(torch.ones(self.config.num_hidden_layers + 1))
        self.attention = SelfAttentionPooling(self.config.hidden_size)
        self.fc = nn.Sequential(
            nn.Linear(self.config.hidden_size, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, embedding_dim),
            nn.BatchNorm1d(embedding_dim)
        )
        self.classifier = nn.Linear(embedding_dim, num_classes)

    def forward(self, input_values, attention_mask=None):
        outputs = self.backbone(input_values=input_values, attention_mask=attention_mask)
        stacked_hidden_states = torch.stack(outputs.hidden_states, dim=0)
        weights = F.softmax(self.layer_weights, dim=0).view(-1, 1, 1, 1)
        weighted_hidden_state = torch.sum(stacked_hidden_states * weights, dim=0)
        pooled, _ = self.attention(weighted_hidden_state, attention_mask)
        embeddings = self.fc(pooled)
        
        # 🌟 Training එකේ වගේම අගයන් 3ක් Return කරයි
        norm_embeddings = F.normalize(embeddings, p=2, dim=1)
        logits = self.classifier(norm_embeddings)
        return embeddings, norm_embeddings, logits

# ==========================================
# 2. මොඩලයන් පූරණය කිරීම (Hugging Face)
# ==========================================
DEVICE = torch.device("cpu")
BASE_MODEL_NAME = "facebook/wav2vec2-xls-r-300m"
PROCESSOR = Wav2Vec2FeatureExtractor.from_pretrained(BASE_MODEL_NAME)

REPO_ID = "TD-jayadeera/model_255" 
MODEL_FILENAME= "SinhalaPhonoNet_Final_Checkpoint_v4.pth"

try:
    print("⏳ Downloading & Loading Custom Model from Hugging Face...")
    model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
    custom_model = SinhalaPhonoNet(num_classes=255).to(DEVICE)
    
    # 🌟 Checkpoint එකෙන් මොළය පමණක් වෙන් කර ගැනීම
    checkpoint = torch.load(model_path, map_location=DEVICE, weights_only=False)
    custom_model.load_state_dict(checkpoint['model_state_dict'])
    custom_model.eval()
    print("✅ Custom Model Loaded Successfully!")
    
    print("⏳ Loading Base Model...")
    base_wav2vec2 = Wav2Vec2Model.from_pretrained(BASE_MODEL_NAME).to(DEVICE)
    base_wav2vec2.eval()
    print("✅ Base Model Loaded Successfully!")
    
except Exception as e:
    print(f"❌ Error loading models: {e}")

# ==========================================
# 3. ප්‍රධාන Analysis Logic
# ==========================================
def process_audio(teacher_audio, student_audio, use_custom=True):
    if teacher_audio is None or student_audio is None:
        return "කරුණාකර ශබ්ද ගොනු දෙකම ලබා දෙන්න.", {} 

    try:
        def get_emb(path):
            speech, _ = librosa.load(path, sr=16000)
            speech, _ = librosa.effects.trim(speech, top_db=25)
            inputs = PROCESSOR(speech, sampling_rate=16000, return_tensors="pt", padding=True)
            
            with torch.no_grad():
                if use_custom:
                    # 🌟 අගයන් 3න් මැද අගය (Norm Embeddings) පමණක් ලබාගැනීම
                    _, emb, _ = custom_model(inputs.input_values, inputs.attention_mask)
                else:
                    outputs = base_wav2vec2(inputs.input_values, attention_mask=inputs.attention_mask)
                    emb = torch.mean(outputs.last_hidden_state, dim=1)
                    emb = F.normalize(emb, p=2, dim=1)
            return emb.cpu().numpy()

        emb_t = get_emb(teacher_audio)
        emb_s = get_emb(student_audio)
        raw_dist = float(np.linalg.norm(emb_t - emb_s))

        # =========================================================
        # 🌟 අලුත් මොඩලයට ගැලපෙන සේ Calibration (Thresholds) වෙනස් කළා
        # =========================================================
        if use_custom:
            # 0.26 (Match) සහ 0.36 (Mismatch) අතර හරි මැද ලක්ෂ්‍යය
            center_point = 0.31  
            # පරතරය කුඩා නිසා Sigmoid curve එකේ බෑවුම වැඩි කිරීම
            steepness = 40       
        else:
            center_point = 0.85  
            steepness = 12
            
        accuracy = (1 / (1 + math.exp(steepness * (raw_dist - center_point)))) * 100
        # =========================================================

        if accuracy >= 85:
            verdict, color, msg = "EXCELLENT", "green", "ඉතාම නිවැරදියි! 🏆"
        elif accuracy >= 65:
            verdict, color, msg = "GOOD", "orange", "හොඳයි, තව උත්සාහ කරන්න! ⭐"
        else:
            verdict, color, msg = "INCORRECT", "red", "නැවත උත්සාහ කරන්න. ❌"

        results_labels = {
            "Excellent (ඉතා විශිෂ්ටයි)": 1.0 if verdict == "EXCELLENT" else 0.0,
            "Good (හොඳයි)": 1.0 if verdict == "GOOD" else 0.0,
            "Needs Work (නැවත උත්සාහ කරන්න)": 1.0 if verdict == "INCORRECT" else 0.0
        }

        model_type_str = "Custom SinhalaPhonoNet (255-Class)" if use_custom else "Base Wav2Vec2-300m"
        info_html = f"""
        <div style='text-align: center; padding: 20px; border-radius: 10px; background-color: #f0f2f6; border: 2px solid {color};'>
            <p style='color: #555; font-weight: bold;'>භාවිතා කළ මොඩලය: {model_type_str}</p>
            <h2 style='color: {color}; margin-top: 0;'>{verdict}</h2>
            <h3 style='color: #333;'>{msg}</h3>
            <p style='font-size: 1.4em;'>නිරවද්‍යතාවය: <b>{accuracy:.2f}%</b></p>
            <p style='font-size: 0.9em; color: #666;'>Raw Distance: {raw_dist:.4f}</p>
        </div>
        """
        return info_html, results_labels

    except Exception as e:
        return f"<p style='color:red;'>Error: {str(e)}</p>", {} 

def analyze_custom(t, s): return process_audio(t, s, use_custom=True)
def analyze_base(t, s): return process_audio(t, s, use_custom=False)

# ==========================================
# 4. Gradio UI
# ==========================================
with gr.Blocks() as demo: 
    gr.Markdown("# 🎙️ සිංහල මිතුරු (Sinhala Mithuru) - Pronunciation Lab")
    gr.Markdown("පර්යේෂණ අරමුණු සඳහා මොඩලයන් දෙකෙහි වෙනස මෙතැනින් පරීක්ෂා කරන්න.")
    
    with gr.Row():
        with gr.Column(scale=1):
            t_input = gr.Audio(type="filepath", label="ගුරුවරයාගේ ශබ්දය (Teacher)")
            s_input = gr.Audio(type="filepath", label="ඔබේ ශබ්දය (Student)")
            
            with gr.Row():
                btn_custom = gr.Button("Analyze (Custom Model)", variant="primary")
                btn_base = gr.Button("Analyze (Base Model)", variant="secondary")
            
        with gr.Column(scale=1):
            result_html = gr.HTML(label="Result Status")
            label_output = gr.Label(num_top_classes=1, label="Verdict Visualization")

    btn_custom.click(fn=analyze_custom, inputs=[t_input, s_input], outputs=[result_html, label_output])
    btn_base.click(fn=analyze_base, inputs=[t_input, s_input], outputs=[result_html, label_output])

demo.launch()