import gradio as gr
import numpy as np
import onnxruntime as ort
import json
from cryptography.fernet import Fernet
import os

# Model load
key = os.getenv("ARTER_MODEL_KEY")
cipher = Fernet(key)

with open("model.bin", "rb") as f:
    encrypted = f.read()
    decrypted = cipher.decrypt(encrypted)
    ort_session = ort.InferenceSession(decrypted)

# Initialize ONNX session
input_name = ort_session.get_inputs()[0].name

def softmax(x):
    return np.exp(x) / np.sum(np.exp(x), axis=1)

# Load mappings
with open("species_to_idx.json", "r") as f:
    species_to_idx = json.load(f)

with open("idx_to_target.json", "r") as f:
    idx_to_target = json.load(f)
    
target_list = list(idx_to_target.keys())

def predict(species):
    
    if not species:
        return {}, {}, {}
    
    input = [species_to_idx[s] for s in species]
    input_np = np.array(input, dtype=np.int64).reshape(1, -1)

    output = ort_session.run([], {input_name: input_np})
    output = {target: softmax(output[i]) for i, target in enumerate(target_list)}
    
    format_output = []
    
    for target in target_list:
        probs = output[target][0]
        classes = idx_to_target[target]
        format_output.append({classes[str(i)]: float(p) for i, p in enumerate(probs)})
    
    return format_output

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Naturtype og §3 klassifikation baseret på artsfund")
    gr.Markdown("Dette værktøj analyserer artsregistreringer for at bestemme et områdes hovednaturtype, naturtilstand og om området er omfattet af Naturbeskyttelseslovens paragraf 3.")

    with gr.Row():
        choices = sorted(species_to_idx.keys())
        species_dropdown = gr.Dropdown(
            choices=choices,
            multiselect=True,
            label="Vælg arter",
        )
    
    with gr.Row():
        naturtype = gr.Label("", label="Hovednaturtype")
        p3 = gr.Label("", label="§3 område")
        tilstand = gr.Label("", label="Naturtilstand")

    gr.Markdown("App og model af Kenneth Thorø Martinsen (kenneth2810@gmail.com).")

    species_dropdown.change(
        predict,
        inputs=[species_dropdown],
        outputs=[naturtype, p3, tilstand]
    )

if __name__ == "__main__":
    demo.launch()