import gradio as gr
import onnxruntime
from transformers import AutoTokenizer
import torch
import os
from transformers import pipeline
import subprocess
import moviepy.editor as mp
import base64

token  = AutoTokenizer.from_pretrained('distilroberta-base')

inf_session = onnxruntime.InferenceSession('classifier-quantized2.onnx')
input_name = inf_session.get_inputs()[0].name
output_name = inf_session.get_outputs()[0].name

classes = ['Art', 'Astrology', 'Biology', 'Chemistry', 'Economics', 'History', 'Literature', 'Philosophy', 'Physics', 'Politics', 'Psychology', 'Sociology']

### --- Audio/Video to txt ---###
device = "cuda:0" if torch.cuda.is_available() else "cpu"
pipe = pipeline("automatic-speech-recognition", 
                    model="openai/whisper-tiny.en",
                    chunk_length_s=30, device=device)

### --- Text Summary --- ###
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=device)
     

def video_identity(video):
    transcription = pipe(video)["text"]
    return transcription

def summary(text):
    text = text.split('.')
    max_chunk = 500
    current_chunk = 0
    chunks = []
    
    
    for t in text:
        if len(chunks) == current_chunk + 1:
            if len(chunks[current_chunk]) + len(t.split(' ')) <= max_chunk:
                chunks[current_chunk].extend(t.split(' '))
            else:
                current_chunk += 1
                chunks.append(t.split(' '))
        else:
          chunks.append(t.split(' '))
    
    for chunk in range(len(chunks)):
        chunks[chunk] =' '.join(chunks[chunk])
    
    summ = summarizer(chunks,max_length = 100)

    return summ

def classify(video_file,encoded_video):
    
    if encoded_video != "":
    
        decoded_file_data = base64.b64decode(encoded_video)

        with open("temp_video.mp4", "wb") as f:
            f.write(decoded_file_data)
        
        video_file = "temp_video.mp4"

    clip = mp.VideoFileClip(video_file)
    clip.audio.write_audiofile(r"audio.wav")
    
    full_text = video_identity(r"audio.wav")
    sum = summary(full_text)[0]['summary_text']

    
    input_ids = token(sum)['input_ids'][:512]
    logits = inf_session.run([output_name],{input_name : [input_ids]})[0]
    logits = torch.FloatTensor(logits)
    probs = torch.sigmoid(logits)[0]
    probs = list(probs)
    label = classes[probs.index(max(probs))]

    final = {
    'text':full_text,
    'summary':sum,
    'label':label,
}
    return final

text1 = gr.Textbox(label="Text")
text2 =  gr.Textbox(label="Summary")



    
iface = gr.Interface(fn=classify,
                     inputs=['video','text'],
                     outputs = ['json'])
iface.launch(inline=False)