import json
from pathlib import Path
import numpy as np
import pandas as pd
import xgboost as xgb
import gradio as gr
import spaces

MODEL_PATH = Path("surrogate_peak_temp.json")
BENCHMARK_JSON = Path("benchmark_results.json")

def load_surrogate():
    if not MODEL_PATH.exists():
        return None
    model = xgb.XGBRegressor()
    model.load_model(str(MODEL_PATH))
    return model

model = load_surrogate()

@spaces.GPU
def predict_thermal_performance(wall_t, roof_alb, shading, wwr, vent_ach):
    if model is None:
        return "Model missing", "Model missing"
    # Order: [wall_t, vent_ach, roof_alb, shading, wwr]
    x_in = np.array([[wall_t, vent_ach, roof_alb, shading, wwr]])
    pred_peak = float(model.predict(x_in)[0])
    peak_suppression = 45.52 - pred_peak
    return f"{pred_peak:.2f} °C", f"{peak_suppression:.2f} °C"

def get_benchmark_table(mode_choice):
    if not BENCHMARK_JSON.exists():
        return pd.DataFrame([{"Message": "Run `python cli.py --benchmark` to generate ablation results."}])
    
    with open(BENCHMARK_JSON) as f:
        bench_data = json.load(f)

    mode_key = "as_built" if "Mode A" in mode_choice else "shared_interior"
    table_rows = []

    # Ambient row
    table_rows.append({
        "Archetype": bench_data["ambient"]["label"],
        "Peak Temp [°C]": f"{bench_data['ambient']['peak']:.2f} °C",
        "Mean Temp [°C]": f"{bench_data['ambient']['mean']:.2f} °C",
        "DDH (>30°C)": f"{bench_data['ambient']['ddh_30']:.1f}",
        "Daily Decrement (f)": f"{bench_data['ambient']['decrement_f']:.2f}"
    })

    for k, v in bench_data[mode_key].items():
        table_rows.append({
            "Archetype": v["label"],
            "Peak Temp [°C]": f"{v['peak']:.2f} °C",
            "Mean Temp [°C]": f"{v['mean']:.2f} °C",
            "DDH (>30°C)": f"{v['ddh_30']:.1f}",
            "Daily Decrement (f)": f"{v['decrement_f']:.2f}"
        })

    return pd.DataFrame(table_rows)

def get_image_path(filename):
    path = Path(filename)
    return str(path) if path.exists() else None

with gr.Blocks(title="HeatBee - Vernacular Thermal Surrogate") as demo:
    gr.Markdown("# 🏛️ HeatBee: Vernacular Thermal Physics & AI Surrogate")
    gr.Markdown("**Surrogate Modeling of Moroccan Rammed Earth (*Pisé*) Architecture**\n*EnergyPlus 26.2 1D Conduction Finite Difference | Marrakesh Menara EPW (Hottest Outdoor Hour: 45.52°C)*")
    
    if model is None:
        gr.Warning("Surrogate model not found. Please run `python cli.py --train` first.")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🛠️ Design Envelope Controls")
            wall_t = gr.Slider(0.25, 0.60, value=0.50, step=0.05, label="Pisé Wall Thickness [m]", info="Analytical diurnal penetration depth is ~11 cm. Beyond 45 cm, mass yields diminishing returns.")
            roof_alb = gr.Slider(0.20, 0.85, value=0.60, step=0.05, label="Roof Solar Reflectance (Albedo)", info="Solar reflectance of horizontal roof finish. Dominant thermal parameter.")
            shading = gr.Slider(0.10, 0.80, value=0.40, step=0.05, label="Mashrabiya Glazing SHGC", info="Solar Heat Gain Coefficient. Lower value represents denser shading.")
            wwr = gr.Slider(0.20, 0.50, value=0.25, step=0.05, label="Courtyard Window-to-Wall Ratio")
            vent_ach = gr.Slider(1.0, 12.0, value=4.0, step=0.5, label="Ventilation Rate [ACH]", info="Free cooling ventilation rate when outdoor air is cooler than indoors.")

        with gr.Column():
            gr.Markdown("### ⚡ Real-Time Physics Inference (< 1 ms)")
            pred_out = gr.Textbox(label="Predicted Peak Operative Temp")
            supp_out = gr.Textbox(label="Peak Suppression vs Ambient", info="Difference between peak outdoor ambient (45.52°C) and indoor peak operative temperature.")
            
            gr.Markdown("---")
            gr.Markdown("### 🔬 Benchmark Comparison")
            mode_choice = gr.Radio(
                choices=["Mode A: As-Built Typologies (Archetype Partitions)", "Mode B: Shared Interior Mass (15cm Concrete Slab for All)"],
                value="Mode A: As-Built Typologies (Archetype Partitions)",
                label="Benchmark Mode:"
            )
            bench_table = gr.Dataframe(interactive=False)
            gr.Markdown("*Shared boundary conditions: Roof albedo 0.60, external wall solar absorptance 0.50, glazing SHGC 0.40, 4.0 ACH free cooling, adiabatic floor and party walls.*")

    gr.Markdown("---")
    gr.Markdown("### 📊 Visualizations")
    with gr.Tabs():
        with gr.Tab("1-to-1 Heatwave Ablation Profile"):
            gr.Image(value=get_image_path("rigorous_ablation_heatwave_profile.png"), show_label=False)
        with gr.Tab("Direct Physical Sweeps"):
            gr.Image(value=get_image_path("direct_physics_sweeps.png"), show_label=False)
        with gr.Tab("Surrogate Parity"):
            gr.Image(value=get_image_path("surrogate_parity_plot.png"), show_label=False)

    # Reactive Bindings
    inputs = [wall_t, roof_alb, shading, wwr, vent_ach]
    outputs = [pred_out, supp_out]

    for component in inputs:
        component.change(fn=predict_thermal_performance, inputs=inputs, outputs=outputs)
    
    demo.load(fn=predict_thermal_performance, inputs=inputs, outputs=outputs)
    mode_choice.change(fn=get_benchmark_table, inputs=[mode_choice], outputs=[bench_table])
    demo.load(fn=get_benchmark_table, inputs=[mode_choice], outputs=[bench_table])

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