import numpy as np
import keras
from PIL import Image
import gradio as gr

MODEL_PATH = "small32cnn_mlbt_mmat.keras"
STATS_PATH = "jet_image_scale_and_stats.npz"

model = keras.models.load_model(MODEL_PATH, compile=False)

stats = np.load(STATS_PATH)
SCALE = float(stats["SCALE"])
MEAN  = float(stats["MEAN"])
STD   = float(stats["STD"])
print("Loaded SCALE/MEAN/STD:", SCALE, MEAN, STD, flush=True)

CLASS_NAMES = ["MMAT", "MLBT"]

# ---------------------------------------------------------
# PREPROCESSING
# ---------------------------------------------------------
def preprocess(img: np.ndarray) -> np.ndarray:
    if img.ndim == 3 and img.shape[2] == 3:
        img_gray = np.dot(img[..., :3], [0.2989, 0.5870, 0.1140])
    elif img.ndim == 3 and img.shape[2] == 1:
        img_gray = img[..., 0]
    elif img.ndim == 2:
        img_gray = img
    else:
        img_gray = img[..., 0]

    pil_img = Image.fromarray(img_gray.astype("uint8"))
    pil_img = pil_img.resize((32, 32), Image.BILINEAR)
    arr = np.array(pil_img).astype("float32")

    arr_unscaled = arr / SCALE
    arr_norm = (arr_unscaled - MEAN) / (STD + 1e-8)
    arr_norm = arr_norm[None, ..., None]
    return arr_norm


# ---------------------------------------------------------
# PREDICT FUNCTION
# ---------------------------------------------------------
def predict(img: np.ndarray):
    x = preprocess(img)
    raw = float(model.predict(x, verbose=0)[0, 0])
    prob_mlbt = raw
    prob_mmat = 1.0 - prob_mlbt
    return {"MLBT": prob_mlbt, "MMAT": prob_mmat}



# ---------------------------------------------------------
# ABOUT PAGE — Flip-cards recreated in HTML/CSS
# ---------------------------------------------------------
about_html = """
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.0/css/all.min.css">

<style>
  .grid {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(210px, 1fr));
    gap: 16px;
    margin-top: 20px;
  }

  .flip-card {
    background: transparent;
    width: 100%;
    height: 230px;
    perspective: 1400px;
    cursor: pointer;
  }

  .flip-inner {
    position: relative;
    width: 100%;
    height: 100%;
    transition: transform 0.6s;
    transform-style: preserve-3d;
  }

  .flip-card:hover .flip-inner {
    transform: rotateY(180deg);
  }

  .flip-face {
    position: absolute;
    inset: 0;
    border-radius: 12px;
    backface-visibility: hidden;
    display: flex;
    flex-direction: column;
    justify-content:center;
    align-items:center;
    padding: 14px;
    text-align:center;
    box-shadow: 0 6px 18px rgba(16,24,40,0.1);
  }

  .front {
    background:white;
    font-size:26px;
    font-weight:800;
    color:#0f172a;
  }

  .front small {
    margin-top:8px;
    font-size:12px;
    color:#64748b;
    font-weight:500;
  }

  .back {
    background: linear-gradient(135deg,#7c3aed,#2563eb);
    color:white;
    transform: rotateY(180deg);
    font-size:14px;
    line-height:1.35;
  }

  h1 {
    text-align:center;
    color:#5b21b6;
    font-size:32px;
  }

  h2 {
    text-align:center;
    color:#5b21b6;
    margin-top:25px;
  }
</style>


<h1> Learn more about high-energy physics! ⚛️ </h1>
<h2> 🔍 Key Concepts </h2>

<div class="grid">
  <!-- QGP -->
  <div class="flip-card">
    <div class="flip-inner">
      <div class="flip-face front">QGP <small>click to learn more</small></div>
      <div class="flip-face back">
        <b>Quark–Gluon Plasma (QGP)</b><br><br>
        A super-hot state of matter where quarks and gluons move freely.<br><br>
        Helps us understand conditions microseconds after the Big Bang.
      </div>
    </div>
  </div>

  <!-- Pb-Pb -->
  <div class="flip-card">
    <div class="flip-inner">
      <div class="flip-face front">Pb–Pb <small>click to learn more</small></div>
      <div class="flip-face back">
        <b>Lead–Lead Collisions</b><br><br>
        High-energy lead nuclei collisions that create QGP.<br><br>
        Lets us study how jets lose energy in extreme matter.
      </div>
    </div>
  </div>

  <!-- Jet -->
  <div class="flip-card">
    <div class="flip-inner">
      <div class="flip-face front">Jet <small>click to learn more</small></div>
      <div class="flip-face back">
        <b>Jets & Jet Quenching</b><br><br>
        Sprays of particles created during collisions.<br><br>
        Energy loss → reveals QGP properties.
      </div>
    </div>
  </div>

  <!-- αₛ -->
  <div class="flip-card">
    <div class="flip-inner">
      <div class="flip-face front">αₛ <small>click to learn more</small></div>
      <div class="flip-face back">
        <b>Strong Coupling Constant (αₛ)</b><br><br>
        Measures how strongly quarks bind together.<br><br>
        Different αₛ values help model jet energy loss.
      </div>
    </div>
  </div>

  <!-- Q0 -->
  <div class="flip-card">
    <div class="flip-inner">
      <div class="flip-face front">Q₀ <small>click to learn more</small></div>
      <div class="flip-face back">
        <b>Virtuality Scale (Q₀)</b><br><br>
        Determines how quantum or classical jet evolution is.<br><br>
        Controls which processes dominate in QGP.
      </div>
    </div>
  </div>

  <!-- MATTER -->
  <div class="flip-card">
    <div class="flip-inner">
      <div class="flip-face front">MATTER <small>click to learn more</small></div>
      <div class="flip-face back">
        <b>MATTER Model</b><br><br>
        Describes energy loss via radiation + elastic collisions.<br><br>
        Represents early jet evolution.
      </div>
    </div>
  </div>

  <!-- MATTER-LBT -->
  <div class="flip-card">
    <div class="flip-inner">
      <div class="flip-face front">MATTER–LBT <small>click to learn more</small></div>
      <div class="flip-face back">
        <b>MATTER–LBT Model</b><br><br>
        Hybrid model combining radiation + medium scattering.<br><br>
        Simulates realistic QGP interactions.
      </div>
    </div>
  </div>

</div>
"""


# ---------------------------------------------------------
# FINAL GRADIO APP WITH TABS
# ---------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("<h1 style='text-align:center;'>MLBT vs MMAT Jet Classifier</h1>")

    with gr.Tabs():
        with gr.Tab("🔬 Classifier"):
            input_component = gr.Image(type="numpy", label="Upload a jet image")
            output_component = gr.Label(num_top_classes=2, label="Predicted probabilities")
            input_component.change(predict, inputs=input_component, outputs=output_component)

        with gr.Tab("📘 About"):
            gr.HTML(about_html)

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