import os
# --- METHOD 1: PERSISTENT CACHE IMPLEMENTATION ---
# Forces Hugging Face to cache weights locally within your Space directory
os.environ["HF_HOME"] = "/app/.hf_cache"
os.environ["XDG_CACHE_HOME"] = "/app/.hf_cache"

import torch
import torchaudio
import gradio as gr
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC

MODEL_OPTIONS = [
    "facebook/wav2vec2-large-960h-lv60-self", # Highly accurate, crisp character mapping (~1.2 GB)
    "facebook/mms-1b-fl102",                  # 1-Billion parameters, ultimate acoustic accuracy (~3.8 GB)
    "facebook/wav2vec2-xlsr-53-espeak-cv-ft", # Keep for your English-IPA mapper 
    "facebook/wav2vec2-base-960h"              # Your original base model
]

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# --- COMPLETE ENGLISH ALPHABET MAPPER FOR THE IPA MODEL ---
IPA_TO_ENGLISH_MAP = {
    # Vowels & Diphthongs
    'ɑː': 'ah', 'ɑ': 'ah', 'ɒ': 'o', 'æ': 'a', 'ʌ': 'uh', 'ɔː': 'aw', 'ɔ': 'o',
    'ɛ': 'e', 'ɜː': 'er', 'ə': 'uh', 'ɪ': 'i', 'iː': 'ee', 'ʊ': 'oo', 'uː': 'oo',
    'aɪ': 'ye', 'aʊ': 'ow', 'eɪ': 'ay', 'oʊ': 'oh', 'ɔɪ': 'oy',
    # Consonants 
    'b': 'b', 'd': 'd', 'f': 'f', 'ɡ': 'g', 'g': 'g', 'h': 'h', 'j': 'y', 'k': 'k', 'l': 'l', 'm': 'm',
    'n': 'n', 'ŋ': 'ng', 'p': 'p', 'ɹ': 'r', 'r': 'r', 's': 's', 'ʃ': 'sh', 't': 't', 'tʃ': 'ch',
    'θ': 'th', 'ð': 'th', 'v': 'v', 'w': 'w', 'z': 'z', 'ʒ': 'zh', 'dʒ': 'j', 'ɲ': 'ny', 'ɾ': 't'
}

def clean_ipa_to_english(ipa_str):
    """Translates the linguistic IPA symbols into readable English letters"""
    text = ipa_str.lower()
    # Sort keys by length so multi-character symbols match first
    for ipa_char in sorted(IPA_TO_ENGLISH_MAP.keys(), key=len, reverse=True):
        text = text.replace(ipa_char, IPA_TO_ENGLISH_MAP[ipa_char])
    return text

def load_model_and_processor(model_choice):
    processor = Wav2Vec2Processor.from_pretrained(model_choice)
    model = Wav2Vec2ForCTC.from_pretrained(model_choice).to(device)
    model.eval()
    return model, processor

model_name = MODEL_OPTIONS[0]
model, processor = load_model_and_processor(model_name)

def transcribe_audio(audio_path, model_choice):
    try:
        global processor, model, model_name

        if model_choice != model_name:
            model, processor = load_model_and_processor(model_choice)
            model_name = model_choice  
    
        if not audio_path:
            return "Please provide an audio file!"

        # Audio Load & Preprocessing
        waveform, sample_rate = torchaudio.load(audio_path)
        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)
        if sample_rate != 16000:
            resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
            waveform = resampler(waveform)
    
        input_values = processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_values.to(device)
        
        with torch.no_grad():
            logits = model(input_values).logits
        predicted_ids = torch.argmax(logits, dim=-1)
        
        raw_tokens = []
        previous_token_id = None
        silence_counter = 0
        
        pad_id = processor.tokenizer.pad_token_id
        word_delimiter_id = getattr(processor.tokenizer, "word_delimiter_token_id", None)
        
        for i in range(predicted_ids.shape[-1]):
            token_id = predicted_ids[0, i].item()
            
            # If model is predicting empty space or padding
            if token_id == pad_id or token_id == word_delimiter_id:
                silence_counter += 1
                # Inject a period if a notable pause occurs between characters
                if silence_counter > 12 and raw_tokens and raw_tokens[-1] != ".":
                    raw_tokens.append(".")
                previous_token_id = token_id
                continue
                
            silence_counter = 0
            
            # CTC character collapse
            if token_id != previous_token_id:
                token_char = processor.tokenizer.decode([token_id], clean_up_tokenization_spaces=False)
                raw_tokens.append(token_char)
                previous_token_id = token_id

        # Combine text and clean up edge boundaries
        output_text = "".join(raw_tokens).strip(".")
        
        # Clean consecutive duplicate periods down to single ones (e.g. 'abc...def' -> 'abc.def')
        cleaned_splits = []
        for char in output_text:
            if char == "." and cleaned_splits and cleaned_splits[-1] == ".":
                continue
            cleaned_splits.append(char)
        output_text = "".join(cleaned_splits)

        if "espeak" in model_choice:
            # 1. Translate the raw IPA array into English alphabetical groupings
            readable_english_gibberish = clean_ipa_to_english(output_text)
            return readable_english_gibberish if readable_english_gibberish else "[Unintelligible Sounds]"
        else:
            # For the base model, enforce lowercase formatting
            return output_text.lower() if output_text else "[Unintelligible Sounds]"
    
    except Exception as e:
        return f"Error: {str(e)}"
    
# Gradio Interface setup
iface = gr.Interface(
    fn=transcribe_audio,
    inputs=[
        gr.Audio(type="filepath", label="Record or Upload Audio"),
        gr.Dropdown(MODEL_OPTIONS, label="Select Transcription Strategy", value=MODEL_OPTIONS[0]),
    ],
    outputs=gr.Textbox(label="Phonetic 'Gibberish' Transcription Output"),
    title="Phonetic Gibberish Transcriber",
    description="Speak into the microphone. Conversational pauses will generate period splits (.), and phonetic outputs are displayed entirely in standard English characters.",
)

if __name__ == "__main__":
    iface.launch()