import gradio as gr
import pymarc
from marcai.process import process
from marcai.utils.parsing import record_dict
import pandas as pd
from marcai.predict import predict_onnx
from marcai.utils import load_config

def compare(file1, file2):
    record1 = pymarc.parse_xml_to_array(file1)[0]
    record2 = pymarc.parse_xml_to_array(file2)[0]

    df1 = pd.DataFrame.from_dict([record_dict(record1)])
    df2 = pd.DataFrame.from_dict([record_dict(record2)])

    df = process(df1, df2)

    # Load model config
    config = load_config("config.yaml")
    model_onnx = "model.onnx"

    # Run ONNX model
    input_df = df[config["model"]["features"]]
    prediction = predict_onnx(model_onnx, input_df)

    prediction = prediction.item()

    return {"match": prediction, "not match": 1 - prediction}


interface = gr.Interface(
    fn=compare,
    inputs=[
        gr.File(label="MARC XML File 1"),
        gr.File(label="MARC XML File 2")
    ],
    outputs=gr.Label(label="Classification"),
    title="MARC Record Matcher",
    description="Upload two MARC XML files with one record each.",
    allow_flagging="never"
)
interface.launch()