import os
import json
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
import pydicom
import gradio as gr
import imageio

MODEL_PATH = "final_model.h5"
MAPPING_PATH = "class_mapping.json"
IMG_SIZE = (224, 224)

# Load model
model = load_model(MODEL_PATH)

# Load class mapping
with open(MAPPING_PATH, "r") as f:
    class_mapping = json.load(f)


def preprocess_file(file_path):
    ext = os.path.splitext(file_path)[1].lower()

    if ext == ".dcm":
        dicom = pydicom.dcmread(file_path)
        img = dicom.pixel_array.astype("float32")
    else:
        img = imageio.imread(file_path).astype("float32")
        if img.ndim == 3 and img.shape[2] == 4:
            img = img[..., :3]

        if img.ndim == 3:
            img = np.mean(img, axis=2)

    img -= np.min(img)
    mx = np.max(img)
    if mx > 0:
        img /= mx

    img = np.stack((img,) * 3, axis=-1)
    img = tf.image.resize(img, IMG_SIZE).numpy()

    return img


def predict(file):
    if file is None:
        return "No file uploaded.", {}

    image = preprocess_file(file)
    input_tensor = tf.expand_dims(image, axis=0)

    pred_target_prob, pred_class_probs = model.predict(input_tensor)

    target_prob = float(pred_target_prob[0][0])
    predicted_class_id = int(np.argmax(pred_class_probs[0]))
    predicted_class_name = class_mapping[str(predicted_class_id)]

    class_probs = {
        class_mapping[str(i)]: float(pred_class_probs[0][i])
        for i in range(pred_class_probs.shape[1])
    }

    return (
        f"Target: {'Abnormal' if target_prob >= 0.5 else 'Normal'} "
        f"({target_prob:.4f})",
        class_probs
    )


# Gradio Interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.File(label="Upload DICOM (.dcm) or Image File"),
    outputs=[
        gr.Textbox(label="Abnormality Prediction"),
        gr.Label(label="Class Probabilities")
    ],
    title="Lung Opacity Detector",
    description="Upload a DICOM or Image. Model predicts abnormality and class."
)


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)