import warnings
warnings.filterwarnings("ignore")

import os
import re
import librosa
import webrtcvad
import nbimporter
import torchaudio
import numpy as np
import gradio as gr
import scipy.signal
import soundfile as sf
from transformers import pipeline
from transformers import AutoProcessor
from pyctcdecode import build_ctcdecoder
from transformers import Wav2Vec2ProcessorWithLM

from text2int import text_to_int
from isNumber import is_number
from Text2List import text_to_list
from convert2list import convert_to_list
from processDoubles import process_doubles
from replaceWords import replace_words
from applyVad import apply_vad
from wienerFilter import wiener_filter
from highPassFilter import high_pass_filter



transcriber_hindi_new = pipeline(task="automatic-speech-recognition", model="cdactvm/w2v-bert-2.0-hindi_new")
transcriber_hindi_old = pipeline(task="automatic-speech-recognition", model="cdactvm/huggingface-hindi_model")
processor = AutoProcessor.from_pretrained("cdactvm/w2v-bert-2.0-hindi_new")
vocab_dict = processor.tokenizer.get_vocab()
sorted_vocab_dict = {k.lower(): v for k, v in sorted(vocab_dict.items(), key=lambda item: item[1])}
decoder = build_ctcdecoder(
    labels=list(sorted_vocab_dict.keys()),
    kenlm_model_path="lm.binary",
    )
processor_with_lm = Wav2Vec2ProcessorWithLM(
    feature_extractor=processor.feature_extractor,
    tokenizer=processor.tokenizer,
    decoder=decoder
    )
processor.feature_extractor._processor_class = "Wav2Vec2ProcessorWithLM"
transcriber_hindi_lm = pipeline("automatic-speech-recognition", model="cdactvm/w2v-bert-2.0-hindi_new", tokenizer=processor_with_lm, feature_extractor=processor_with_lm.feature_extractor, decoder=processor_with_lm.decoder)


def transcribe_hindi_new(audio):
    # # Process the audio file
    transcript = transcriber_hindi_new(audio)
    text_value = transcript['text']
    processd_doubles=process_doubles(text_value)
    replaced_words = replace_words(processd_doubles)
    converted_text=text_to_int(replaced_words)
    return converted_text
    
def transcribe_hindi_lm(audio):
    # # Process the audio file
    transcript = transcriber_hindi_lm(audio)
    text_value = transcript['text']
    processd_doubles=process_doubles(text_value)
    replaced_words = replace_words(processd_doubles)
    converted_text=text_to_int(replaced_words)
    return converted_text

def transcribe_hindi_old(audio):
    # # Process the audio file
    transcript = transcriber_hindi_old(audio)
    text_value = transcript['text']
    cleaned_text=text_value.replace("<s>","")
    processd_doubles=process_doubles(cleaned_text)
    replaced_words = replace_words(processd_doubles)
    converted_text=text_to_int(replaced_words)
    return converted_text

## implementation of noise reduction techniques.
###############################################
def noise_reduction_pipeline(filepath):
    # Your existing noise reduction code
    audio, sr = librosa.load(filepath, sr=None)
    audio_hp = high_pass_filter(audio, sr, cutoff=100, order=5)
    audio_wiener = wiener_filter(audio_hp)
    audio_vad = apply_vad(audio_wiener, sr)
    output_filepath = "processed_output.wav"
    sf.write(output_filepath, audio_vad, sr)
    return output_filepath

# Hugging Face ASR function uses the pre-loaded model
def transcribe_with_huggingface(filepath):
    result = transcriber_hindi_lm(filepath)
    text_value = result['text']
    cleaned_text = text_value.replace("<s>", "")
    converted_to_list = convert_to_list(cleaned_text, text_to_list())
    processed_doubles = process_doubles(converted_to_list)
    replaced_words = replace_words(processed_doubles)
    converted_text = text_to_int(replaced_words)
    print("Transcription: ", converted_text)
    return converted_text

# Combined function to process and transcribe audio
def process_audio_and_transcribe(audio):
    # Step 1: Preprocess (Noise Reduction)
    try:
        processed_filepath = noise_reduction_pipeline(audio)
    except webrtcvad.Error as e:
        return f"Error in processing audio for VAD: {str(e)}"
    
    # Step 2: Transcription
    try:
        transcription = transcribe_with_huggingface(processed_filepath)
    except Exception as e:
        return f"Transcription failed: {str(e)}"
    
    return transcription
#################################################

def sel_lng(lng, mic=None, file=None):
    if mic is not None:
        audio = mic
    elif file is not None:
        audio = file
    else:
        return "You must either provide a mic recording or a file"
    
    if lng == "model_1":
        return transcribe_hindi_old(audio)
    elif lng == "model_2":
        return transcribe_hindi_new(audio)
    elif lng== "model_3":
        return transcribe_hindi_lm(audio)
    elif lng== "model_4":
        return process_audio_and_transcribe(audio)
            
        
# demo=gr.Interface(
#     transcribe,
#     inputs=[
#         gr.Audio(sources=["microphone","upload"], type="filepath"),
#     ],
#     outputs=[
#         "textbox"
#     ],
#     title="Automatic Speech Recognition",
#     description = "Demo for Automatic Speech Recognition. Use microphone to record speech. Please press Record button. Initially it will take some time to load the model. The recognized text will appear in the output textbox",
#       ).launch()

demo=gr.Interface(
    fn=sel_lng, 
      
    inputs=[
        gr.Dropdown([
            "model_1","model_2","model_3","model_4"],label="Select Model"),
        gr.Audio(sources=["microphone","upload"], type="filepath"),
    ],
    outputs=[
        "textbox"
    ],
    title="Automatic Speech Recognition",
    description = "Demo for Automatic Speech Recognition. Use microphone to record speech. Please press Record button. Initially it will take some time to load the model. The recognized text will appear in the output textbox",
      ).launch()