import numpy as np
import gradio as gr
import joblib
import tensorflow as tf
import pandas as pd

MODEL_PATH = "./NN_12_feat.keras"
SCALER_PATH = "./scaler.pkl"

# 1) Cargar modelo y scaler
model = tf.keras.models.load_model(MODEL_PATH)
scaler = joblib.load(SCALER_PATH)

# 2) Definí el orden EXACTO de tus 12 features (como entrenaste)

FEATURE_NAMES = ['Angular Width (deg)',
 'MPA (deg)',
 'Linear Speed (km/s)',
 '2nd ord. Final Linear Speed (km/s)',
 'Mass (g)',
 'Bx (nT)',
 'Bz (nT)',
 'Proton Temperature (K)',
 'Flow Speed (km/s)',
 'Flow Longitude (deg)',
 'Alpha to Proton Ratio',
 'Plasma Pressure (nPa)']

# 3) Rangos/steps/valores iniciales para sliders (ajustalos a tu dataset)
#    Si no sabés rangos, poné inputs como gr.Number (más abajo te muestro opción).
'''RANGES = {
    name: {"min": 0.0, "max": 1.0, "step": 0.01, "value": 0.5}
    for name in FEATURE_NAMES
}'''

RANGES = {
    "Angular Width (deg)": {
        "min": 0, "max": 360, "step": 1, "value": 180, "unit": "deg"
    },
    "MPA (deg)": {
        "min": 0, "max": 360, "step": 1, "value": 180, "unit": "deg"
    },
    "Linear Speed (km/s)": {
        "min": 100, "max": 3000, "step": 10, "value": 800, "unit": "km/s"
    },
    "2nd ord. Final Linear Speed (km/s)": {
        "min": 100, "max": 3000, "step": 10, "value": 800, "unit": "km/s"
    },
    "Mass (g)": {
        "min": 1e12, "max": 1e17, "step": 1e12, "value": 5e14, "unit": "g"
    },
    "Bx (nT)": {
        "min": -20, "max": 20, "step": 0.5, "value": 0, "unit": "nT"
    },
    "Bz (nT)": {
        "min": -40, "max": 20, "step": 0.5, "value": -10, "unit": "nT"
    },
    "Proton Temperature (K)": {
        "min": 1e3, "max": 2e6, "step": 1e3, "value": 1e5, "unit": "K"
    },
    "Flow Speed (km/s)": {
        "min": 100, "max": 1000, "step": 10, "value": 500, "unit": "km/s"
    },
    "Flow Longitude (deg)": {
        "min": -10, "max": 10, "step": 1, "value": 0, "unit": "deg"
    },
    "Alpha to Proton Ratio": {
        "min": 0, "max": 10, "step": 0.005, "value": 0.02, "unit": ""
    },
    "Plasma Pressure (nPa)": {
        "min": 0, "max": 20, "step": 0.1, "value": 2, "unit": "nPa"
    }
}

# 4) Umbrales para feedback “dinámico” (en horas) - ajustá si querés
THRESH_FAST = 24
THRESH_MED = 48

def _predict_from_vector(x_raw: np.ndarray) -> float:
    """x_raw: shape (12,) sin escalar"""
    x_scaled = scaler.transform(x_raw.reshape(1, -1))
    y = model.predict(x_scaled, verbose=0).reshape(-1)[0]
    return float(y)

def predict(*inputs):
    x = np.array(inputs, dtype=float)

    y = _predict_from_vector(x)

    # Texto dinámico
    if y <= THRESH_FAST:
        label = "Rápida"
        msg = f"🟢 Predicción **rápida**: **{y:.2f} h** (≤ {THRESH_FAST} h)."
    elif y <= THRESH_MED:
        label = "Intermedia"
        msg = f"🟠 Predicción **intermedia**: **{y:.2f} h** (entre {THRESH_FAST} y {THRESH_MED} h)."
    else:
        label = "Lenta"
        msg = f"🔴 Predicción **lenta**: **{y:.2f} h** (> {THRESH_MED} h)."

    table = [[FEATURE_NAMES[i], float(x[i])] for i in range(len(FEATURE_NAMES))]
    #gauge = {"value": y, "min": 0, "max": max(120, y * 1.2), "label": f"Tiempo de propagación (h) — {label}"}

    return y, table, msg
'''
def what_if(feature_idx, *inputs):
    x0 = np.array(inputs, dtype=float)
    fname = FEATURE_NAMES[int(feature_idx)]
    r = RANGES[fname]

    xs = np.linspace(r["min"], r["max"], 60)
    ys = []
    for v in xs:
        x = x0.copy()
        x[int(feature_idx)] = v
        ys.append(_predict_from_vector(x))

    data = [{"x": float(a), "y": float(b)} for a, b in zip(xs, ys)]
    title = f"Curva what-if variando: **{fname}**"
    return data, title'''

def what_if(feature_idx, *inputs):
    x0 = np.array(inputs, dtype=float)

    idx = int(feature_idx)  # por si viene como string
    fname = FEATURE_NAMES[idx]
    r = RANGES[fname]

    xs = np.linspace(r["min"], r["max"], 60)
    ys = []
    for v in xs:
        x = x0.copy()
        x[idx] = v
        ys.append(_predict_from_vector(x))

    df = pd.DataFrame({"x": xs.astype(float), "y": np.array(ys, dtype=float)})

    title = f"Curva what-if variando: **{fname}**"
    return df, title

with gr.Blocks() as demo:
    gr.Markdown(
        "## Predicción del tiempo de propagación de una CME\n"
        "Ingresá los 12 atributos.\n"
        "Se puede explorar la sensibilidad de cada atributo mediante **what-if**."
    )

    with gr.Row():
        with gr.Column(scale=2):
            inputs = []
            for name in FEATURE_NAMES:
                r = RANGES[name]
                inputs.append(
                    gr.Slider(
                        minimum=r["min"], maximum=r["max"], step=r["step"],
                        value=r["value"],
                        label=name
                    )
                )
            btn = gr.Button("Predecir")

        with gr.Column(scale=1):
            out_value = gr.Number(label="Tiempo de propagación predicho (h)", precision=2)
            out_table = gr.Dataframe(headers=["Atributo", "Valor"], interactive=False, row_count=12, col_count=2)
            out_text = gr.Markdown()

    btn.click(predict, inputs=inputs, outputs=[out_value, out_table, out_text])

    gr.Markdown("### What-if (sensibilidad)")
    with gr.Row():
        feature_pick = gr.Dropdown(
            choices=[(f"{i} — {FEATURE_NAMES[i]}", i) for i in range(len(FEATURE_NAMES))],
            value=0,
            label="Elegí el atributo a variar"
        )
        whatif_btn = gr.Button("Generar curva")

    plot = gr.LinePlot(
        x="x", y="y",
        title="Predicción vs valor del atributo",
        x_title="Valor del atributo",
        y_title="Tiempo de propagación (h)",
        height=350
    )
    plot_title = gr.Markdown()

    whatif_btn.click(
        what_if,
        inputs=[feature_pick] + inputs,
        outputs=[plot, plot_title]
    )

demo.launch()