import gradio as gr
from PIL import Image
import torch

from uei_core.models import ModelPortfolio
from uei_core.uncertainty import UncertaintyEstimator
from uei_core.energy import EnergyProfiler
from uei_core.policy import UEIPolicy
from visuals import get_visual_ui
device = "cpu"  # M1 CPU is optimized enough; GPU optional

models = ModelPortfolio(device=device)
unc = UncertaintyEstimator()
energy = EnergyProfiler()
policy = UEIPolicy()


def uei_infer(img):

    x = models.preprocess(img)

    # Step 1: Small model
    logits_s, E_s = energy.measure(models.infer_small, x)
    U_s = unc.estimate(logits_s)

    # Step 2: Large model (to measure marginal utility)
    logits_l, E_l = energy.measure(models.infer_large, x)
    U_l = unc.estimate(logits_l)

    decision = policy.decide(U_s, U_l, E_s, E_l)

    if decision == "small":
        final_logits = logits_s
        energy_used = E_s
        unc_used = U_s
        model_name = "Low-Energy Small Model"
    else:
        final_logits = logits_l
        energy_used = E_l
        unc_used = U_l
        model_name = "High-Energy Large Model"

    probs = torch.softmax(final_logits, dim=1).squeeze()
    top_idx = torch.argmax(probs).item()
    confidence = float(probs[top_idx])

    return {
        "Predicted Class Index": top_idx,
        "Confidence": confidence,
        "Selected Model": model_name,
        "Uncertainty": float(unc_used),
        "Energy (proxy units)": energy_used
    }


demo = gr.Interface(
    fn=uei_infer,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(),
    title="Uncertainty-Elastic Inference (UEI)",
    description="Energy-efficient inference by dynamically selecting models based on uncertainty-energy tradeoffs."
)

# ---------------------------
# NEW UI: Add Visualization Tab
# ---------------------------
visual_tab = get_visual_ui()

app = gr.TabbedInterface(
    [
        demo,            # your original UEI inference app
        visual_tab       # new visualization dashboard
    ],
    [
        "UEI Inference",
        "Visual Dashboard"
    ]
)

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