import torch
import gradio as gr
import models as MOD
import process_data as PD

model_master = {
    "SSL-AASIST (Trained on ASV-Spoof5)": {"eer_threshold": 3.3330237865448,
                                           "data_process_func": "process_ssl_assist_input",
                                           "note": "This model is trained only on ASVSpoof 2024 training data.",
                                           "model_class": "Model",
                                           "model_checkpoint": "ssl_aasist_epoch_7.pth"},
    "AASIST": {"eer_threshold": 1.8018419742584229,
               "data_process_func": "process_assist_input",
               "note": "This model is trained on ASVSpoof 2024 training data.",
               "model_class":"AASIST_Model",
               "model_checkpoint": "orig_aasist_epoch_1.pth"}
}

model = MOD.Model(None, "cpu")
model.load_state_dict(torch.load("ssl_aasist_epoch_7.pth", map_location="cpu"))
model.eval()
loaded_model = "SSL-AASIST (Trained on ASV-Spoof5)"


def process(file, type):
    global model
    global loaded_model
    inp = getattr(PD, model_master[type]["data_process_func"])(file)
    if not loaded_model == type:
        model = getattr(MOD, model_master[type]["model_class"])(None, "cpu")
        model.load_state_dict(torch.load(model_master[type]["model_checkpoint"], map_location="cpu"))
        model.eval()
        loaded_model = type

    op = model(inp).detach().squeeze()[1].item()

    response_text = "Decision score: {} \nDecision threshold: {} \nNotes: 1. Any score below threshold is indicative of fake. \n2. {} ".format(
        str(op), str(model_master[type]["eer_threshold"]), model_master[type]["note"])
    return response_text


demo = gr.Blocks()
file_proc = gr.Interface(
    fn=process,
    inputs=[
        gr.Audio(sources=["upload"], label="Audio file", type="filepath"),
        gr.Radio(["SSL-AASIST (Trained on ASV-Spoof5)", "AASIST"], label="Select Model", type="value"),
    ],
    outputs="text",
    title="Find the Fake: Analyze 'Real' or 'Fake'.",
    description=(
        "Analyze fake or real with a click of a button. Upload a .wav or .flac file."
    ),
    examples=[
        ["./bonafide.flac", "SSL-AASIST (Trained on ASV-Spoof5)"],
        ["./fake.flac", "SSL-AASIST (Trained on ASV-Spoof5)"],
        ["./bonafide.flac", "AASIST"],
        ["./fake.flac", "AASIST"],
    ],
    cache_examples=True,
    allow_flagging="never",
)

with demo:
    gr.TabbedInterface([file_proc], ["Analyze Audio File"])
demo.queue(max_size=10)
demo.launch(share=True)