from fastapi import FastAPI
from pydantic import BaseModel, Field
from typing import Literal
import json
import numpy as np
import onnxruntime as ort
from typing_extensions import Annotated
import gradio as gr
from cryptography.fernet import Fernet
import os

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

VERSION = "0.0.1"
TITLE = f"DVPI beregnings API (version {VERSION})"
DESCRIPTION = "Beregn Dansk Vandløbs Plante Indeks (DVPI) fra dækningsgrad af plantearter. Beregningen er baseret på en model som efterligner DVPI beregningsmetoden og er dermed ikke eksakt, usikkerheden er i gennemsnit **±0.05 EQR-enheder**."
URL = "https://huggingface.co/spaces/KennethTM/dvpi"

# Load ONNX model and species mappings
with open("model.bin", "rb") as f:
    encrypted = f.read()
    decrypted = cipher.decrypt(encrypted)
    ort_session = ort.InferenceSession(decrypted)

with open("spec2idx.json", "r") as f:
    spec2idx = json.load(f)

# Define types
valid_species = tuple(spec2idx.keys())

class SpeciesCover(BaseModel):
    species: dict[Literal[valid_species], Annotated[float, Field(ge=0, le=100)]]
    
    model_config = {
        "json_schema_extra": {
            "examples": [{
                "species": {
                    "Potamogeton alpinus": 25.0,
                    "Berula erecta": 15.5,
                    "Calamagrostis canescens": 10.0
                }
            }]
        }
    }


class EQRResult(BaseModel):
    EQR: float  # Round to 2 decimals
    DVPI: int
    version: str = VERSION

# Create FastAPI app
app = FastAPI(title=TITLE,
              description=DESCRIPTION)

def eqr_to_dvpi(eqr: float) -> int:
    if eqr < 0.20:
        return 1
    elif eqr < 0.35:
        return 2
    elif eqr < 0.50:
        return 3
    elif eqr < 0.70:
        return 4
    else:
        return 5

# FastAPI routes
@app.post("/dvpi")
def predict(cover_data: SpeciesCover) -> EQRResult:
    """Predict EQR and DVPI from species cover data"""
    # Initialize input vector with zeros
    input_vector = np.zeros((1, len(spec2idx)))

    print(cover_data.species)

    # Fill values from input
    for species, cover in cover_data.species.items():
        idx = spec2idx[species]
        input_vector[0, idx] = cover
    
    # Get prediction
    input_name = ort_session.get_inputs()[0].name
    ort_inputs = {input_name: input_vector.astype(np.float32)}
    ort_output = ort_session.run(None, ort_inputs)

    eqr = float(ort_output[0][0])
    dvpi = eqr_to_dvpi(eqr)
    
    return EQRResult(EQR=round(eqr, 2), DVPI=dvpi)

@app.get("/arter")
def list_species() -> dict:
    """Return list of valid species names"""
    return {"species": list(spec2idx.keys())}

# Gradio app
def add_entry(species, cover, current_dict) -> tuple[SpeciesCover, str]:
    
    current_dict[species] = cover

    return current_dict, current_dict

def gradio_predict(cover_data: dict):

    if len(cover_data) == 0:
        return {}

    data = SpeciesCover(species=cover_data)
    result = predict(data)    

    return result.model_dump()
    
with gr.Blocks() as io:

    gr.Markdown(f"# {TITLE}")
    gr.Markdown(DESCRIPTION)

    with gr.Tab(label = "Beregner"):

        gr.Markdown("Beregning er baseret på samfund af plantearter og deres dækningsgrad. Dækningsgraden angives i procent som summen af scoren for dækningsgraden (1-5) divideret med det samlede antal undersøgte kvadrater gange 5, og til sidste konverteret til procent. Eksempel: Potamogeton alpinus findes 3 felter med scorerne 2, 3 og 5 ud af 50 undersøgte kvadrater. Dækningsgraden for Potamogeton alpinus er derfor (2+3+5)/(50*5)*100 = 4%.")

        current_dict = gr.State({})
        
        with gr.Row():
            species_input = gr.Dropdown(choices=valid_species, label="Vælg art")
            cover_input = gr.Number(label="Dækningsgrad (%)", minimum=0, maximum=100)
        
        with gr.Row():
            add_btn = gr.Button("Tilføj")
            reset_btn = gr.Button("Nulstil")
        
        list_display = gr.JSON(label="Artsliste")
        
        calc_btn = gr.Button("Beregn")
        results = gr.JSON(label="Resultater")
            
        def reset_dict():
            return {}, {}, {}
        
        add_btn.click(
            add_entry,
            inputs=[species_input, cover_input, current_dict],
            outputs=[current_dict, list_display]
        )
        
        reset_btn.click(
            reset_dict,
            inputs=[],
            outputs=[current_dict, list_display, results]
        )
        
        calc_btn.click(
            gradio_predict,
            inputs=[current_dict],
            outputs=results
        )

        gr.Markdown("App og model af Kenneth Thorø Martinsen.")

    with gr.Tab(label="Dokumentation"):

        # Add markdown description with code to call the api in python
        gr.Markdown("## Eksempel på brug af API")
        gr.Markdown(f"API dokumentation kan findes på [{URL}/docs]({URL}/docs)")
        gr.Markdown("### Python")
        gr.Code(f"""
import requests
import json
            
data = {{
    "species": {{
        "Potamogeton alpinus": 25.0,
        "Berula erecta": 15.5,
        "Calamagrostis canescens": 10.0
    }}
}}

response = requests.post("{URL}/dvpi", json=data)
print(response.json())
        """)

        gr.Markdown("### R")
        gr.Code(f"""
library(httr)
library(jsonlite)

data <- list(species = list(
    "Potamogeton alpinus" = 25.0,
    "Berula erecta" = 15.5,
    "Calamagrostis canescens" = 10.0
))

response <- POST("{URL}/dvpi",
                body = toJSON(data, auto_unbox = TRUE),
                content_type("application/json"))

print(fromJSON(rawToChar(response$content)))
        """)
    
# Mount Gradio app
app = gr.mount_gradio_app(app, io, path="/")