import whisper
import gradio as gr
from docx import Document
import time
import librosa
import soundfile as sf
from scipy.cluster.hierarchy import dendrogram, fcluster
from sklearn.preprocessing import StandardScaler
import numpy as np

import ssl
import urllib.request

# 临时禁用证书验证
context = ssl._create_unverified_context()
urllib.request.urlopen("https://huggingface.co", context=context)

# Load the Whisper model (choose a model based on performance/accuracy requirements)
model = whisper.load_model("base")  # You can use "small", "medium", "large" for better accuracy

def extract_mfcc(audio_path, sr=None):
    """Extract MFCC features from audio file"""
    y, sr = librosa.load(audio_path, sr=sr)
    mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    return mfccs.T  # Transpose to get (time_steps, n_mfcc)

def perform_clustering(features):
    """Perform hierarchical clustering on MFCC features"""
    scaler = StandardScaler()
    scaled_features = scaler.fit_transform(features)
    
    # Compute pairwise distances
    distance_matrix = np.sqrt(((scaled_features - scaled_features[:, np.newaxis])**2).sum(axis=2))
    
    # Perform hierarchical clustering
    linkage = np.linalg.norm(distance_matrix[:, np.newaxis] - distance_matrix[np.newaxis, :], axis=2)
    
    # Generate dendrogram and get flat clusters
    clusters = fcluster(linkage, t=0.5, criterion='distance')  # Threshold needs tuning
    return clusters

def transcribe_and_generate_docx(audio_path, language):
    # Step 1: Basic transcription with Whisper
    audio_file = whisper.load_audio(audio_path)
    language_code = 'en' if language == "English" else 'zh'
    result = model.transcribe(audio_file, language=language_code)
    transcribed_text = result["text"]
    
    # Step 2: Extract MFCC features for clustering
    mfcc_features = extract_mfcc(audio_path)
    if len(mfcc_features) < 2:
        return "output_transcription.docx", transcribed_text
    
    # Step 3: Perform hierarchical clustering
    try:
        clusters = perform_clustering(mfcc_features)
    except Exception as e:
        log_message(f"Clustering error: {str(e)}")
        raise RuntimeError("Failed to cluster speakers") from e
    
    # Assign speaker labels
    unique_speakers = list(np.unique(clusters))
    speaker_map = {i+1: speaker for i, speaker in enumerate(unique_speakers)}
    
    # Split transcription text by speaker clusters
    time_segments = model.get_timestamps(audio_file)
    speaker_texts = {}
    
    for cluster_id, segment in zip(clusters, time_segments):
        speaker_id = speaker_map[cluster_id]
        start_time = segment['start']
        end_time = segment['end']
        
        # Find corresponding text片段 (approximation method)
        current_text = []
        for part in result["alternatives"][0]["transcript"].split(' '):
            part_duration = len(part) * 0.03  # Approximate word duration
            if start_time <= cumulative_time < end_time:
                current_text.append(part)
                cumulative_time += part_duration
        
        speaker_texts[speaker_id] = ' '.join(current_text).strip()
    
    # Create a docx document with speaker separation
    doc = Document()
    doc.add_heading("Speaker Transcription", level=1)
    
    for speaker_id, text in speaker_texts.items():
        doc.add_heading(f"Speaker {speaker_id}", level=2)
        doc.add_paragraph(text)
    
    # Save the document
    output_filename = f"speaker_transcription_{os.path.basename(audio_path)}.docx"
    doc.save(output_filename)
    
    return output_filename, "\n\n".join([f"Speaker {i}: {text}" for i, text in speaker_texts.items()])

# ... [Keep Gradio interface code unchanged] ...

# Launch the interface
iface.launch()