import numpy as np
import gradio as gr

from core import predict_stats, gradcam_overlay, CONF_TAU_MONTHS
from preprocess_infer import preprocess_image


def _fmt_readable(months: float) -> str:
    return f"{int(months // 12)} ปี {int(round(months % 12))} เดือน"


# ==========================================
# API 1: Bone Age (Decimal, เดือน)
# ==========================================
def api_bone_age(image: np.ndarray, sex: str):
    if image is None:
        return {"error": "no image provided"}
    mean, _sd, conf = predict_stats(image, sex)
    return {
        "bone_age_months": round(mean, 2),
        "bone_age_years": round(mean / 12.0, 2),
        "readable": _fmt_readable(mean),
        "confidence_pct": round(conf, 1),
        "sex": "female" if str(sex).lower().startswith("f") else "male",
    }


# ==========================================
# API 2: Confidence Interval (Decimal ± SD, เดือน)
# ==========================================
def api_confidence_interval(image: np.ndarray, sex: str, k: float = 1.0):
    if image is None:
        return {"error": "no image provided"}
    mean, sd, conf = predict_stats(image, sex)
    return {
        "mean_months": round(mean, 2),
        "sd_months": round(sd, 2),
        "confidence_pct": round(conf, 1),
        "confidence_note": f"มั่นใจ {round(conf, 1)}% ว่าอายุกระดูกอยู่ในช่วง ±{int(CONF_TAU_MONTHS)} เดือนของค่าที่ทำนาย",
        "k": k,
        "lower_months": round(mean - k * sd, 2),
        "upper_months": round(mean + k * sd, 2),
        "interval_str": f"{mean:.2f} ± {sd:.2f} เดือน",
    }


# ==========================================
# API 3: Grad-CAM (Image)
# ==========================================
def api_gradcam(image: np.ndarray, sex: str):
    if image is None:
        return None
    overlay, _mean = gradcam_overlay(image, sex)
    return overlay


# ==========================================
# API 4 (debug): ดูภาพที่โมเดลเห็นจริง (หลัง crop + histogram match + pad)
# ==========================================
def api_preprocess(image: np.ndarray):
    if image is None:
        return None
    return preprocess_image(image)  # (512,512) uint8


# ==========================================
# UI (Blocks) — แต่ละปุ่มผูก api_name แยกกัน
# ==========================================
with gr.Blocks(title="Bone Age Prediction") as demo:
    gr.Markdown("# 🦴 Bone Age Prediction (ConvNeXtV2 Tiny)")
    gr.Markdown("อัปโหลด X-ray มือ + เลือกเพศ แล้วเรียกแต่ละฟังก์ชันได้แยกกัน")

    with gr.Row():
        inp_img = gr.Image(type="numpy", image_mode="L", label="Hand X-ray")
        inp_sex = gr.Radio(choices=["male", "female"], value="male", label="Sex")

    with gr.Tab("Bone Age"):
        btn1 = gr.Button("Predict Bone Age", variant="primary")
        out1 = gr.JSON(label="Bone Age (months)")
        btn1.click(api_bone_age, [inp_img, inp_sex], out1, api_name="bone_age")

    with gr.Tab("Confidence Interval"):
        inp_k = gr.Slider(0.5, 3.0, value=1.0, step=0.5, label="k (จำนวน SD)")
        btn2 = gr.Button("Compute Interval", variant="primary")
        out2 = gr.JSON(label="Confidence Interval (± SD, months)")
        btn2.click(api_confidence_interval, [inp_img, inp_sex, inp_k], out2,
                   api_name="confidence_interval")

    with gr.Tab("Grad-CAM"):
        btn3 = gr.Button("Generate Grad-CAM", variant="primary")
        out3 = gr.Image(type="numpy", label="Grad-CAM overlay")
        btn3.click(api_gradcam, [inp_img, inp_sex], out3, api_name="gradcam")

    with gr.Tab("Debug: ภาพหลัง preprocess"):
        gr.Markdown("ดูว่าโมเดลเห็นอะไรจริง — ถ้ายังเห็น marker/พื้นหลัง = crop ไม่ทำงาน")
        btn4 = gr.Button("Show preprocessed", variant="secondary")
        out4 = gr.Image(type="numpy", image_mode="L", label="Preprocessed (512×512)")
        btn4.click(api_preprocess, [inp_img], out4, api_name="preprocess")


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