import gradio as gr
from transformers import pipeline
from pydub import AudioSegment

pipe = pipeline("automatic-speech-recognition", "facebook/wav2vec2-base-960h")
#morse code dictionary
dictionary = {' ': '/', 
	"'": '.----.', 
	'0': '-----', 
	'1': '.----', 
	'2': '..---', 
	'3': '...--', 
	'4': '....-', 
	'5': '.....', 
	'6': '-....', 
	'7': '--...', 
	'8': '---..', 
	'9': '----.', 
	'?': '..--..', 
	'A': '.-', 
	'B': '-...', 
	'C': '-.-.', 
	'D': '-..', 
	'E': '.', 
	'F': '..-.', 
	'G': '--.', 
	'H': '....', 
	'I': '..', 
	'J': '.---', 
	'K': '-.-', 
	'L': '.-..', 
	'M': '--', 
	'N': '-.', 
	'O': '---', 
	'P': '.--.', 
	'Q': '--.-', 
	'R': '.-.', 
	'S': '...', 
	'T': '-', 
	'U': '..-', 
	'V': '...-', 
	'W': '.--', 
	'X': '-..-', 
	'Y': '-.--', 
	'Z': '--..', 
}



def asr(speech):
  dummy = AudioSegment.from_file("dummy.mp3", format="mp3") #load audio files
  dot = AudioSegment.from_file("dot.mp3", format="mp3")
  dash = AudioSegment.from_file("dash.mp3", format="mp3")
  space = AudioSegment.from_file("space.mp3", format="mp3")
  i=0
  transcript = pipe(speech)['text'] #convert input speech to text
  morseCode = ""
  for character in transcript:
    morseCode += dictionary[character] + " "
  while i<len(morseCode): #append dash and dot sounds
    if morseCode[i] == ".":
      dummy = dummy + dot
      i = i+1
    elif morseCode[i] == "-":
      dummy = dummy + dash
      i = i+1
    elif morseCode[i] == " ":
      dummy = dummy + space
      i = i+1
    else:
      i = i+1
  morseAudio = dummy.export("output.mp3", format="mp3") #export final audio
  return transcript, morseCode, "output.mp3"
  
gr.Interface(fn=asr,
             inputs = gr.inputs.Audio(source="microphone", type="filepath", optional=False, label="Please record your voice"),
             #inputs = gr.inputs.Audio(source="upload", type="filepath", label="Upload your audio file here"),
			 outputs = [gr.outputs.Textbox(type="str", label="Text Translation"),
						gr.outputs.Textbox(type="str", label="Morse Code"),
						gr.outputs.Audio(type="file", label="Audio Translation")],
             title = "Speech to Morse Code",
             description = "This app will translate your speech into text, morse code, and audio using wav2vec2-base-960h.",
			 examples = [["helloworld.mp3"], ["SOS.mp3"], ["SMS.mp3"]],
			 article = "Model: <a href=\"https://huggingface.co/facebook/wav2vec2-base-960h\">Wav2Vec2-Base-960h</a>"
             ).launch()