# This is version 2 updated on 17th Sept 2024.
# Uses the Whiper Medium model ( on RTX 4070 with 8GB vram)
#Beep done changed and beepify_segments function not used instead now using audio_to_beep.overlay
# Please change beep sound wave filepath according to your local dir  in   "Beeped_Audio_Path":  line 254
#output audio stored in "pii_beep_audio_uploads" in local dir where this file located


import gradio as gr
import os
import random
import whisper_timestamped as whisper
from pydub import AudioSegment
import numpy as np
import spacy
import torch
import threading
import zipfile
import shutil
from pathlib import Path
from werkzeug.utils import secure_filename
import time
from gradio_rich_textbox import RichTextbox
import re

# Worker class to process the audio file and load models
class Worker(threading.Thread):
    def __init__(self, audio_file_path, model_directory, callback):
        threading.Thread.__init__(self)
        self._AudiofileName = audio_file_path
        self._ModelDirectory = model_directory
        self._BeepAudiofileName = "beep2.wav"
        self.callback = callback

        self._PII_text_and_Timestamp =""
        self._Transcribe_Text_With_Entities =""
        self._Metrics =""
        self._BeepedAudiofileName =""

        print(f"Audio File: {self._AudiofileName}")
        print(f"Model Directory: {self._ModelDirectory}")
        print(f"Beep Audio File: {self._BeepAudiofileName}")

    def run(self):
        try:
            print("loading SpaCy model with custom model ",str(self._ModelDirectory))
            # Load spaCy model from directory or a known model name
            self.nlp = spacy.load(str(self._ModelDirectory))
            print("SpaCy model loaded.")

            # Load Whisper model
            devices = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
            print(devices)
            time.sleep(0.2)
            self.model = whisper.load_model("medium", device=devices)
            
            print("Whisper model loaded.")

            self.processData()
            self.callback("callback Processing complete!")

        except Exception as e:
            print(f"Error during processing: {str(e)}")

    def count_entities(self,entities):
        entity_counts = {}  # Initialize an empty dictionary to store counts

        for _, entity_type in entities:
            # Increment the count for each entity type
            entity_counts[entity_type] = entity_counts.get(entity_type, 0) + 1

        return entity_counts

    def colorize_entities(self, data, entities):
        # Define color mappings (you can customize these)
        color_map = {
            'PERSON': 'blue',
            'GPE': 'green',
            'LOC': 'purple',
            'PHONE': 'orange',
            'EMAIL': 'blue',
            'CAR_PLATE':'red', 
            'ORG':'purple',              
            'NRIC': 'red',
            'PASSPORT_NUM':'green'
        }

        print("entities",entities)
        # Replace entities with colored versions
        for entity, entity_type in entities:
            #print("before update data",data)
            color = color_map.get(entity_type, 'blue')  # Default to blue if type not found
            colored_entity = f'<span style="color: {color};">{entity} {entity_type}</span>'
            data = data.replace(entity, colored_entity)
            #print("after update data",data)

        return data

    def processData(self):
        # Transcribe audio and extract entities
        try:
            # Load audio
            audio = whisper.load_audio(self._AudiofileName)
            output = whisper.transcribe(self.model, audio, beam_size=5, best_of=5, temperature=(0.0, 0.2, 0.4, 0.6, 0.8, 1.0),vad=True, language="en", remove_punctuation_from_words=True,refine_whisper_precision=0.6,min_word_duration=0.01)
            #output = whisper.transcribe(self.model, audio, language="en", task='transcribe', temperature=(0.0, 0.2, 0.4, 0.6, 0.8, 1.0), best_of=5, beam_size=5)""
            transcription_text = output['text']
            transcription_text = re.sub(r"\.(?!\S)", " ", transcription_text)
            print("~~~~~~~~~~~~~~~~")
            print(transcription_text)

            #append text
            self._PII_text_and_Timestamp += (transcription_text)+"\n"
            # Run NER with spaCy
            doc = self.nlp(transcription_text)
            entities = [(ent.text, ent.label_) for ent in doc.ents]
            uniqueentities = list(set(entities))
            entity_counts = self.count_entities(entities)

            for entity_type, count in entity_counts.items():
                #append to metrics
                self._Metrics += (entity_type+ " : "+ str(count))+"\n"

            transcribeWithEntities = self.colorize_entities(transcription_text, uniqueentities)

            #append to transcribeWithEntities
            self._Transcribe_Text_With_Entities = transcribeWithEntities

            print(f"Transcription: {transcription_text}")
            print(f"Entities: {entities}")

            # Beepify audio segments containing PII entities
            audio_to_beep = AudioSegment.from_file(self._AudiofileName)

            # Process the audio file to beepify words (remaining unchanged)
            # Extract segments to be beeped
            self.segments_to_beep = []

            pii_Text_TimeStamp = []

            for ent in doc.ents:
                self.segments_to_beep.append((ent.start_char, ent.end_char))
                pii_Text_TimeStamp.append((ent.text,ent.start_char*200,ent.end_char*200))
                print("=======")
                print("ent.text",ent.text)
                print("ent.start",ent.start_char)
                print("ent.end",ent.end_char)
                
                print(pii_Text_TimeStamp)
            for ent in pii_Text_TimeStamp:
                self._PII_text_and_Timestamp += ("Timestamp: "+str(ent[1]/1000)+ " --- "+str(ent[2]/1000)+" sec")+"\n"
                self._PII_text_and_Timestamp  += ("Text: "+ent[0])+"\n"


            # Convert character offsets to time (assuming 1 character = 20 ms)
            segments_in_ms = [(start*200, end*200) for start, end in self.segments_to_beep]
            print("Segments:", segments_in_ms)


            
            words_to_beepify =[]

            # append the all text in the doc the words_to_beepify array
            for word in doc.ents:
                # words_to_beepify.append(word.text)
                words_to_beepify.append(word.text.replace('.', ''))
               
            print(words_to_beepify)

            # New list to store individual words
            individual_words_to_beepify = []
            
            # Split each phrase into individual words and append to the new list
            for phrase in words_to_beepify:
                individual_words_to_beepify.extend(phrase.split())

            # Remove duplicates by converting the list to a set and then back to a list
            #individual_words_to_beepify = list(set(individual_words_to_beepify))
            individual_words_to_beepify = list(dict.fromkeys(individual_words_to_beepify))
            
            print(individual_words_to_beepify)

            # Load the beep sound
            beep_sound = AudioSegment.from_file(self._BeepAudiofileName)



            # Iterate over the words array in segment array of the output
            for segment in output["segments"]:
                for word in segment["words"]:
                    
                    # Check if the word is in the list of words to beepify
                    if word["text"] in individual_words_to_beepify:
                        # Get the start and end time of the word
                        print("*******")
                        print(word)
                        
                        start_time = word["start"]
                        end_time = word["end"]
                        
                        # Get the start and end indices of the word
                        start_index = float(start_time * 1000)
                        end_index = float(end_time * 1000 + 100) # Add 100ms buffer

                        # Calculate the duration of the word segment
                        word_duration = (end_index - start_index)
                        print(word_duration)   
                        # Create a silent segment with the same duration as the word
                        silent_segment = AudioSegment.silent(duration=word_duration)
                        

                        # Replace the word segment with the silent segment in the original audio
                        audio_to_beep = audio_to_beep[:int(start_index)] + silent_segment + audio_to_beep[int(end_index):]

                        # Get the start and end indices of the beep sound to match the word's duration
                        beep_start_index = 0
                        beep_end_index = word_duration + 200 # Add 200ms 
                        #beep_end_index = word_duration 
                        
                        
                        # Trim the beep sound to match the word's duration
                        beep_sound = beep_sound[beep_start_index:beep_end_index]
                        
                        """ if word_duration > len(beep_sound):
                            beep_sound = beep_sound + AudioSegment.silent(duration=word_duration - len(beep_sound))
                        else:
                            beep_sound = beep_sound[:word_duration] """

                        #Overlay the beep sound on the silent segment
                        audio_to_beep = audio_to_beep.overlay(beep_sound, position=int(start_index))

            # Save the beeped audio file
            random_filename = str(random.getrandbits(32)) + secure_filename(Path(self._AudiofileName).name)
            output_path = os.path.join("pii_beep_audio_uploads", f"new_{random_filename}")
            os.makedirs("pii_beep_audio_uploads", exist_ok=True)
            
            
            audio_to_beep.export(output_path)
            #audio_to_beep.export(output_path, format="wav")
            self._BeepedAudiofileName =output_path

            print(f"Beeped audio file saved at: {output_path}")
            self.callback({
                "PII_text_and_Timestamp": self._Transcribe_Text_With_Entities,
                "Transcribe_Text_With_Entities": self._PII_text_and_Timestamp,
                "Metrics": self._Metrics,
                "Beeped_Audio_Path": self._BeepedAudiofileName
            })
        except Exception as e:
            print(f"An error occurred during transcription: {str(e)}")

       

    
# Callback function for Gradio
def start_worker(audio_file_path, model_directory):
    result = {
        "PII_text_and_Timestamp": "Processing...",
        "Transcribe_Text_With_Entities": "Processing...",
        "Metrics": "Processing...",
        #"Beeped_Audio_Path": "/home/prema/Documents/Audio/beep2.wav"
        "Beeped_Audio_Path": "/content/drive/MyDrive/2024_Project/Pipeline/NER/beep2.wav"
    }

    def update_result(message):
        if isinstance(message, dict):
            result.update({
                "PII_text_and_Timestamp": str(message.get("PII_text_and_Timestamp")),
                "Transcribe_Text_With_Entities": message.get("Transcribe_Text_With_Entities"),
                "Metrics":  str(message.get('Metrics')),
                "Beeped_Audio_Path":  str(message.get('Beeped_Audio_Path'))

            })
        print("Processing complete.")

    if not audio_file_path or os.stat(audio_file_path).st_size == 0:
        return gr.update(visible=True), "Error: No input provided. Please upload a audio file"

    if not model_directory or os.stat(model_directory).st_size == 0:
        return gr.update(visible=True), "Error: No input provided. Please upload model(.zip)file"

   
    # Start worker in a separate thread
    worker = Worker(audio_file_path, model_directory, update_result)
    worker.start()

    # Wait for the worker to finish
    worker.join()

    #returning result to called function
    return result["PII_text_and_Timestamp"], result["Transcribe_Text_With_Entities"], result["Metrics"], result["Beeped_Audio_Path"]

def reset():
    return None, None, None, None, None

def get_audio_file_path(audio):
    return audio



def load_model(files):
    if files:
        # Assume the uploaded file is a zip file representing the directory
        zip_file_path = files.name

        # Define a directory to extract the zip
        extract_dir = "extracted_model"

        # Clean the directory if it already exists
        if os.path.exists(extract_dir):
            shutil.rmtree(extract_dir)

        os.makedirs(extract_dir, exist_ok=True)

        # Extract the zip file contents
        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
            zip_ref.extractall(extract_dir)

        # Debug output: List the contents of the extracted directory
        extracted_files = []
        for root, dirs, files in os.walk(extract_dir):
            for file in files:
                extracted_files.append(os.path.join(root, file))

        print("Extracted files:")
        for file in extracted_files:
            print(file)

        # Determine the base directory inside the extracted directory
        base_dir = None
        for root, dirs, files in os.walk(extract_dir):
            if files and 'meta.json' in files:
                base_dir = root
                break

        # Check if meta.json was found and construct the path
        if base_dir:
            meta_path = os.path.join(base_dir, "meta.json")
            if os.path.exists(meta_path):
                return base_dir
            else:
                directory_message = "Invalid model directory: meta.json not found"
        else:
            directory_message = "Invalid model directory: meta.json not found"

    else:
        directory_message = "No directory selected"

    return directory_message

# Function to load and return the audio file path
def load_audio(beep_audio_file_output):
    if beep_audio_file_output is not None:
        return beep_audio_file_output.name  # Return the path to the uploaded file
    return None
    
# Gradio UI
with gr.Blocks(css="""
    .centered {
        display: flex;
        justify-content: center;
        align-items: center;    }

        .custom-label {
            font-size: 14px;
            font-weight: bold;
            text-align: left;
            height: 100px;
            border: 0px solid black;
        }
""") as demo:

    gr.Markdown("# Speech De-Identification Framework ver-2.0", elem_classes="centered")

    with gr.Column():

        with gr.Row():
        
            audio_input = gr.Audio(label="Upload Audio File", type="filepath")
            audio_output = gr.Textbox(label="Audio File Path", interactive=False, visible = False)
            audio_input.change(fn=get_audio_file_path, inputs=audio_input, outputs=audio_output)



        # Model directory input (as a zip file)
            model_dir_input = gr.File(label="Select ML Model as zip file", file_count="single")
            model_output_path = gr.Textbox(label="Model Load Status", interactive=False, visible = False)
            model_dir_input.change(fn=load_model, inputs=model_dir_input, outputs=model_output_path)

        with gr.Row():
            gr.Markdown("")
            gr.Markdown("")
            gr.Markdown("")
            gr.Markdown("")
            gr.Markdown("")

            reset_button = gr.Button("Reset")
            submit_button = gr.Button("Submit")

       

        
        gr.Markdown("### Transcribe Text and Entities:")
        pii_text_output = RichTextbox(show_label=False , interactive=False)
        gr.Markdown("### PII Text and Time Stamps:")
        transcribe_text_output = gr.Textbox(show_label=False , interactive=False)
        gr.Markdown("### Metrics:")
        metrics_output = gr.Textbox(show_label=False , interactive=False)

        with gr.Row():
            # Audio component to display the audio file in the interface
            beep_audio_file_output = gr.File(label="Download Beeped Audio", interactive=False)
        
            # Audio player component to play the selected audio file
            audio_player = gr.Audio(label="Play Beeped Audio", type="filepath")
        
            # Automatically update the audio player when the file component changes
            beep_audio_file_output.change(load_audio, inputs=beep_audio_file_output, outputs=audio_player)


    # Event Handlers
    reset_button.click(reset, [], [audio_input, model_dir_input, pii_text_output, transcribe_text_output, metrics_output])
    submit_button.click(start_worker, [audio_output, model_output_path], [pii_text_output, transcribe_text_output, metrics_output,beep_audio_file_output])

demo.launch(inbrowser=True, show_error=True,share = True)