import math
import pandas as pd
import gradio as gr
import plotly.graph_objects as go

MATERIALS = {
    "aluminum": {"name":"Aluminio", "family":"structural"},
    "pmc_cf_resin": {"name":"Compuesto (CF+Resina)", "family":"structural"},
    "eva_waste": {"name":"EVA/Tejidos sintéticos", "family":"textiles"},
    "foam_pack": {"name":"Espuma de empaque (Zotek/Plastazote)", "family":"foam"},
    "bubble_wrap": {"name":"Burbuja / Air Cushions", "family":"packaging"},
    "textiles": {"name":"Textiles / Wipes", "family":"textiles"},
    "food_pouches": {"name":"Pouches comida/bebida/termales", "family":"packaging"},
    "filters_mesh": {"name":"Filtros/Mallas lab", "family":"lab"},
    "carbon": {"name":"Carbono excedente (experimentos O₂)", "family":"lab"},
    "resealable_bags": {"name":"Bolsas resellables/antiestáticas", "family":"packaging"},
    "nitrile_gloves": {"name":"Guantes nitrilo", "family":"lab"},
}

PROCESS_EFF = {
    "shred": 0.95, "sort": 0.97, "compact": 0.98, "regolith_mix": 0.96, "form": 0.95
}
WATER = {"wash": {"use_L_per_kg": 0.6, "recovery": 0.93},
         "mix":  {"use_L_per_kg": 0.4, "recovery": 0.95}}
PRODUCT_YIELDS = {
    "Utensilios": {"kg_per_unit": 0.05},
    "Cajas de Almacen": {"kg_per_unit": 0.8},
    "Paneles Interiores": {"kg_per_unit": 2.0},
    "Herramientas": {"kg_per_unit": 0.2},
    "Pads Aislantes": {"kg_per_unit": 0.5},
}
REGOLITH_FRACTION = 0.20
SCENARIOS = {
    "Residence Renovations": {
        "desc": "Reuso del cubo 3D y empaques del hábitat.",
        "inventory_kg": {"aluminum": 120, "pmc_cf_resin": 60, "foam_pack": 80, "bubble_wrap": 20, "resealable_bags": 10}
    },
    "Cosmic Celebrations": {
        "desc": "Fiesta con materiales reutilizados dentro del hábitat.",
        "inventory_kg": {"textiles": 25, "eva_waste": 15, "bubble_wrap": 10, "food_pouches": 12, "resealable_bags": 5}
    },
    "Daring Discoveries": {
        "desc": "Post-experimentos de O₂: reutilizar equipos y carbono excedente.",
        "inventory_kg": {"filters_mesh": 18, "carbon": 35, "eva_waste": 12, "nitrile_gloves": 6, "resealable_bags": 4}
    }
}

def simulate_flow(inventory_kg: dict, eff=PROCESS_EFF, rego_frac=REGOLITH_FRACTION):
    inv = pd.Series(inventory_kg, dtype=float).reindex(MATERIALS.keys()).fillna(0.0)
    total_in = inv.sum()
    wash_kg = inv[["eva_waste","textiles","foam_pack","bubble_wrap","food_pouches","resealable_bags","nitrile_gloves","pmc_cf_resin"]].sum()
    water_wash_use = wash_kg * WATER["wash"]["use_L_per_kg"]
    water_wash_recovered = water_wash_use * WATER["wash"]["recovery"]
    water_wash_loss = water_wash_use - water_wash_recovered
    after_shred   = total_in * eff["shred"]
    after_sort    = after_shred  * eff["sort"]
    after_compact = after_sort   * eff["compact"]
    polymer_mass = after_compact
    regolith_added = polymer_mass * rego_frac
    mix_total = polymer_mass + regolith_added
    after_mix = mix_total * eff["regolith_mix"]
    final_useful = after_mix * eff["form"]
    water_mix_use = polymer_mass * WATER["mix"]["use_L_per_kg"]
    water_mix_recovered = water_mix_use * WATER["mix"]["recovery"]
    water_mix_loss = water_mix_use - water_mix_recovered
    kpi = {
        "Input total (kg)": round(total_in, 2),
        "After Shred (kg)": round(after_shred, 2),
        "After Sort (kg)": round(after_sort, 2),
        "After Compact (kg)": round(after_compact, 2),
        "Regolith added (kg)": round(regolith_added, 2),
        "After Mix (kg)": round(after_mix, 2),
        "Final Useful Output (kg)": round(final_useful, 2),
        "Water use wash (L)": round(water_wash_use, 2),
        "Water recovered wash (L)": round(water_wash_recovered, 2),
        "Water loss wash (L)": round(water_wash_loss, 2),
        "Water use mix (L)": round(water_mix_use, 2),
        "Water recovered mix (L)": round(water_mix_recovered, 2),
        "Water loss mix (L)": round(water_mix_loss, 2),
        "Water total loss (L)": round(water_wash_loss + water_mix_loss, 2),
        "Crew time (min/day)": 12
    }
    family_weights = {"structural": 0.35, "textiles": 0.15, "foam": 0.15, "packaging": 0.20, "lab": 0.15}
    family_mass = {f: final_useful*wt for f, wt in family_weights.items()}
    def units(m, kgpu): return int(max(0, m // kgpu))
    product_plan = {
        "Utensilios":        units(family_mass["packaging"]*0.5 + family_mass["textiles"]*0.2, PRODUCT_YIELDS["Utensilios"]["kg_per_unit"]),
        "Cajas de Almacen":  units(family_mass["structural"]*0.35 + family_mass["packaging"]*0.15, PRODUCT_YIELDS["Cajas de Almacen"]["kg_per_unit"]),
        "Paneles Interiores":units(family_mass["structural"]*0.4 + family_mass["foam"]*0.3, PRODUCT_YIELDS["Paneles Interiores"]["kg_per_unit"]),
        "Herramientas":      units(family_mass["lab"]*0.6 + family_mass["structural"]*0.1, PRODUCT_YIELDS["Herramientas"]["kg_per_unit"]),
        "Pads Aislantes":    units(family_mass["foam"]*0.7 + family_mass["textiles"]*0.1, PRODUCT_YIELDS["Pads Aislantes"]["kg_per_unit"]),
    }
    return kpi, product_plan, inv

def sankey_from_kpis(kpi):
    labels = ["Input","Shred","Sort","Compact","Regolith","Mix","Form (Final)"]
    idx = {lab:i for i,lab in enumerate(labels)}
    inp,a1,a2,a3,reg,mix,final = kpi["Input total (kg)"],kpi["After Shred (kg)"],kpi["After Sort (kg)"],kpi["After Compact (kg)"],kpi["Regolith added (kg)"],kpi["After Mix (kg)"],kpi["Final Useful Output (kg)"]
    loss_shred=max(0,inp-a1); loss_sort=max(0,a1-a2); loss_compact=max(0,a2-a3); loss_mix=max(0,(a3+reg)-mix); loss_form=max(0,mix-final)
    def l(s,t,v): return (s,t,round(v,2)) if v>0.0001 else None
    links=[l(idx["Input"],idx["Shred"],a1), l(idx["Input"],idx["Input"],loss_shred),
           l(idx["Shred"],idx["Sort"],a2), l(idx["Shred"],idx["Shred"],loss_sort),
           l(idx["Sort"],idx["Compact"],a3), l(idx["Sort"],idx["Sort"],loss_compact),
           l(idx["Compact"],idx["Regolith"],reg), l(idx["Regolith"],idx["Mix"],a3+reg),
           l(idx["Mix"],idx["Form (Final)"],final), l(idx["Mix"],idx["Mix"],loss_mix),
           l(idx["Form (Final)"],idx["Form (Final)"],loss_form)]
    links=[x for x in links if x]
    fig=go.Figure(data=[go.Sankey(
        node=dict(pad=20, thickness=18, line=dict(color="white", width=0.5), label=labels),
        link=dict(source=[s for s,_,__ in links], target=[t for _,t,__ in links], value=[v for __,__,v in links])
    )])
    fig.update_layout(margin=dict(l=10,r=10,t=10,b=10), height=430)
    return fig

def run_sim(scenario, aluminum, pmc, eva, foam, bubble, textiles, pouches, filters, carbon, bags, gloves, rego_frac):
    base = SCENARIOS[scenario]["inventory_kg"].copy()
    base.update({
        "aluminum": aluminum, "pmc_cf_resin": pmc, "eva_waste": eva, "foam_pack": foam, "bubble_wrap": bubble,
        "textiles": textiles, "food_pouches": pouches, "filters_mesh": filters, "carbon": carbon,
        "resealable_bags": bags, "nitrile_gloves": gloves
    })
    kpi, products, inv = simulate_flow(base, PROCESS_EFF, rego_frac)
    kpi_df  = pd.DataFrame(list(kpi.items()), columns=["Métrica","Valor"])
    prod_df = pd.DataFrame(list(products.items()), columns=["Producto","Unidades (aprox)"])
    inv_df  = pd.DataFrame({"Material":[MATERIALS[k]["name"] for k in inv.index], "kg":inv.values})
    fig = sankey_from_kpis(kpi)
    return kpi_df, prod_df, inv_df, fig

with gr.Blocks(title="MARS-LOOP — Jezero Crater") as demo:
    gr.Markdown("# 🪐 MARS-LOOP — Reciclaje circular en Marte\n**Convierte residuos inorgánicos + regolito MGS-1 en piezas útiles del hábitat.**  \n_Sin incinerar, minimizando agua, y con operación ligera (<15 min/día)._")
    with gr.Row():
        scenario = gr.Dropdown(choices=list(SCENARIOS.keys()), value="Residence Renovations", label="Escenario NASA")
        rego_frac = gr.Slider(0.0, 0.5, value=REGOLITH_FRACTION, step=0.05, label="Fracción de regolito (0–50%)")
    gr.Markdown("### Inventario (kg) — Ajusta por material")
    with gr.Row():
        with gr.Column():
            aluminum = gr.Slider(0, 300, value=SCENARIOS["Residence Renovations"]["inventory_kg"].get("aluminum",0), step=5, label="Aluminio (kg)")
            pmc      = gr.Slider(0, 300, value=SCENARIOS["Residence Renovations"]["inventory_kg"].get("pmc_cf_resin",0), step=5, label="Compuesto CF+Resina (kg)")
            eva      = gr.Slider(0, 200, value=SCENARIOS["Cosmic Celebrations"]["inventory_kg"].get("eva_waste",0), step=2, label="EVA/Textiles sintéticos (kg)")
            foam     = gr.Slider(0, 200, value=SCENARIOS["Residence Renovations"]["inventory_kg"].get("foam_pack",0), step=2, label="Espuma empaque (kg)")
            bubble   = gr.Slider(0, 100, value=SCENARIOS["Residence Renovations"]["inventory_kg"].get("bubble_wrap",0), step=2, label="Burbuja/Air Cushions (kg)")
            textiles = gr.Slider(0, 100, value=SCENARIOS["Cosmic Celebrations"]["inventory_kg"].get("textiles",0), step=2, label="Textiles/Wipes (kg)")
        with gr.Column():
            pouches = gr.Slider(0, 100, value=SCENARIOS["Cosmic Celebrations"]["inventory_kg"].get("food_pouches",0), step=2, label="Pouches comida/bebida/termal (kg)")
            filters = gr.Slider(0, 100, value=SCENARIOS["Daring Discoveries"]["inventory_kg"].get("filters_mesh",0), step=2, label="Filtros/Mallas (kg)")
            carbon  = gr.Slider(0, 100, value=SCENARIOS["Daring Discoveries"]["inventory_kg"].get("carbon",0), step=2, label="Carbono excedente (kg)")
            bags    = gr.Slider(0, 60,  value=SCENARIOS["Cosmic Celebrations"]["inventory_kg"].get("resealable_bags",0), step=1, label="Bolsas resellables (kg)")
            gloves  = gr.Slider(0, 40,  value=SCENARIOS["Daring Discoveries"]["inventory_kg"].get("nitrile_gloves",0), step=1, label="Guantes nitrilo (kg)")
    run = gr.Button("▶️ Ejecutar simulación")
    kpi_df  = gr.Dataframe(label="KPIs")
    prod_df = gr.Dataframe(label="Productos estimados")
    inv_df  = gr.Dataframe(label="Inventario usado (kg)")
    sankey  = gr.Plot(label="Flujo de materiales — Sankey")
    run.click(run_sim,
              inputs=[scenario, aluminum, pmc, eva, foam, bubble, textiles, pouches, filters, carbon, bags, gloves, rego_frac],
              outputs=[kpi_df, prod_df, inv_df, sankey])
demo.launch()