import os
import torch
import pickle
import pandas as pd
import torch.nn.functional as F
import streamlit as st

from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download

from model import MedBERTClassifier


BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_NAME = "Charangan/MedBERT"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MAX_LEN = 128

# --- Load resources (cached) --- #
@st.cache_resource
def load_resources():

    # Download model from HF model repo
    model_path = hf_hub_download(
        repo_id="ilhamst/rgai_medbert_icd10",
        filename="medbert_epoch_11.pt"
    )

    # Load label encoder
    with open(os.path.join(BASE_DIR, "label_encoder.pkl"), "rb") as f:
        label_encoder = pickle.load(f)

    num_classes = len(label_encoder.classes_)

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    # Load model
    model = MedBERTClassifier(MODEL_NAME, num_classes).to(DEVICE)

    checkpoint = torch.load(
        model_path,
        map_location=DEVICE
    )

    model.load_state_dict(
        checkpoint["model_state_dict"]
    )

    model.eval()

    # ICD lookup
    icd_lookup = pd.read_csv(os.path.join(BASE_DIR, "icd_lookup.csv"))
    icd_dict = dict(zip(icd_lookup.dxcode, icd_lookup.longdesc))

    return model, tokenizer, label_encoder, icd_dict


# Load once
model, tokenizer, label_encoder, icd_dict = load_resources()


# --- Prediction function --- #
def predict_icd(text):

    inputs = tokenizer(
        text,
        padding="max_length",
        truncation=True,
        max_length=MAX_LEN,
        return_tensors="pt"
    )

    input_ids = inputs["input_ids"].to(DEVICE)
    attention_mask = inputs["attention_mask"].to(DEVICE)

    with torch.no_grad():

        logits = model(input_ids, attention_mask)

        probs = torch.softmax(logits, dim=1)

    probs = probs.cpu().numpy()[0]

    top3_idx = probs.argsort()[-3:][::-1]

    results = []

    for idx in top3_idx:

        code = label_encoder.inverse_transform([idx])[0]

        desc = icd_dict.get(code, "Unknown")

        conf = probs[idx]

        results.append((code, desc, conf))

    return results