tabs/tab4_optimization.py

import io
import json
import unicodedata
import zipfile

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import streamlit as st

from config.config import method_order, resource_order, resource_units
from optimization.optimization import run_optimization
from utils.ui_helpers import format_results_summary


def render_tab():
    st.header("Portfolio generation")
    st.caption("Tab 4 - Seaborn fallback for vertical histograms")

    st.info("""
    **Objective:**
    Run the optimization tool, explore results, and save scenarios for comparison or export.
    """)

    # Validate necessary inputs
    if "constraints" not in st.session_state or "costs" not in st.session_state or "resource_caps" not in st.session_state:
        st.error("Missing inputs. Please complete previous steps.")
        return

    if "scenarios" not in st.session_state:
        st.session_state.scenarios = {}

    if sum(st.session_state.resource_caps.values()) <= 0:
        st.warning("\u26a0\ufe0f All resources are currently set to 0. Please increase availability.")
        return

    # Run optimization
    if st.button("Run Optimization"):
        success, result = run_optimization(
            resource_caps=st.session_state.resource_caps,
            method_constraints=st.session_state.constraints,
            method_costs=st.session_state.costs,
        )

        if success:
            allocations = result["method_usage"]
            total_removed = result["total_removed"]

            if total_removed <= 0:
                st.warning("\u274c Optimization removed 0 tCO₂. Relax constraints or increase resources.")
                return

            total_cost = sum(
                allocations[m] * sum(st.session_state.costs[m].values())
                for m in allocations
            )

            st.session_state.latest_result = {
                "summary": {
                    "total_co2_removed": total_removed,
                    "total_cost": total_cost
                },
                "allocations": allocations,
                "resource_caps": st.session_state.resource_caps.copy(),
                "method_constraints": st.session_state.constraints,
                "method_costs": st.session_state.costs,
            }

            st.success(f"\u2705 Optimization successful: {total_removed:,.0f} tCO₂ removed")
        else:
            st.error("\u274c Optimization failed")
            if "message" in result:
                st.error(result["message"])

    if "latest_result" not in st.session_state:
        return

    latest = st.session_state.latest_result
    allocations = latest["allocations"]
    resource_caps = latest["resource_caps"]
    method_costs = latest["method_costs"]

    if not allocations or sum(allocations.values()) == 0:
        return

    # Metrics
    col1, col2 = st.columns(2)
    col1.metric("Total CO₂ Removed", f"{latest['summary']['total_co2_removed']:,.0f} t")
    col2.metric("Methods Used", len(allocations))

    # Summary table
    st.subheader("Optimization Summary")
    df_result = pd.DataFrame.from_dict(allocations, orient="index", columns=["tCO₂ removed"])
    df_result["tCO₂ removed"] = pd.to_numeric(df_result["tCO₂ removed"], errors="coerce")
    total = df_result["tCO₂ removed"].sum()
    if total > 0:
        df_result["Share (%)"] = (df_result["tCO₂ removed"] / total * 100).round(2)
    df_result = df_result.sort_values(by="tCO₂ removed", ascending=False)
    st.table(format_results_summary(df_result))

    # Donut chart
    st.markdown("### Method allocation breakdown")
    labels = df_result.index.tolist()
    values = df_result["tCO₂ removed"].tolist()

    import plotly.graph_objects as go
    fig_donut = go.Figure(data=[go.Pie(
        labels=labels,
        values=values,
        hole=0.4,
        textinfo="label+percent",
        textposition="outside",
        marker=dict(line=dict(color="#000000", width=1)),
        showlegend=False
    )])
    fig_donut.update_layout(margin=dict(t=10, b=10))
    st.plotly_chart(fig_donut, use_container_width=True)

    # Heatmap
    st.markdown("### Resource usage heatmap (% of cap)")
    heatmap_matrix, heatmap_y = [], []
    available_resources = [r for r in resource_order if resource_caps.get(r, 0) > 0]

    for method in df_result.index:
        row = []
        for r in available_resources:
            usage = method_costs.get(method, {}).get(r, 0.0) * allocations.get(method, 0.0)
            row.append((usage / resource_caps[r]) * 100 if resource_caps[r] > 0 else 0)
        heatmap_matrix.append(row)
        heatmap_y.append(method)

    fig_heatmap = go.Figure(data=go.Heatmap(
        z=heatmap_matrix,
        x=available_resources,
        y=heatmap_y,
        colorscale="Blues",
        hovertemplate="Method: %{y}<br>Resource: %{x}<br>% Used: %{z:.2f}%<extra></extra>",
    ))
    fig_heatmap.update_layout(
        title="Resource Usage vs. Cap",
        xaxis_title="Resource",
        yaxis_title="Method",
        height=400,
    )
    st.plotly_chart(fig_heatmap, use_container_width=True)

    # Absolute usage (fallback with seaborn)
    st.markdown("### Absolute resource usage by unit (Seaborn)")
    df_abs = []
    for m, co2_removed in allocations.items():
        for unit, resources in resource_units.items():
            for r in resources:
                if r in method_costs.get(m, {}):
                    used = method_costs[m][r] * co2_removed
                    if used > 0:
                        df_abs.append({
                            "Method": m,
                            "Resource": r,
                            "Used": used,
                            "Unit": unit
                        })

    if df_abs:
        df_abs = pd.DataFrame(df_abs)
        for unit in df_abs["Unit"].unique():
            st.markdown(f"#### Resources in {unit}")
            df_unit = df_abs[df_abs["Unit"] == unit]

            fig, ax = plt.subplots(figsize=(8, 4 + 0.3 * df_unit["Resource"].nunique()))
            sns.barplot(data=df_unit, x="Resource", y="Used", hue="Method", ax=ax)
            ax.set_ylabel(f"Used ({unit})")
            ax.set_xlabel("Resource")
            ax.tick_params(axis='x', rotation=45)
            ax.legend(title="Method", bbox_to_anchor=(1.05, 1), loc='upper left')
            st.pyplot(fig)

    # Save scenario
    st.markdown("### Save scenario")
    st.text_input("Scenario name", key="scenario_name")

    if st.button("Save Scenario"):
        name = st.session_state.get("scenario_name", "").strip()
        if not name:
            st.warning("Please enter a name before saving.")
        elif name in st.session_state.scenarios:
            st.warning("Scenario name already exists.")
        else:
            st.session_state.scenarios[name] = latest
            st.success(f"Scenario '{name}' saved.")

    # Preview saved scenario
    current_name = st.session_state.get("scenario_name", "").strip()
    if current_name and current_name in st.session_state.scenarios:
        scenario = st.session_state.scenarios[current_name]
        st.markdown(f"#### Scenario: {current_name}")
        df_saved = pd.DataFrame.from_dict(
            scenario["allocations"], orient="index", columns=["tCO₂ removed"]
        )
        df_saved["tCO₂ removed"] = pd.to_numeric(df_saved["tCO₂ removed"], errors="coerce")
        total = df_saved["tCO₂ removed"].sum()
        if total > 0:
            df_saved["Share (%)"] = (df_saved["tCO₂ removed"] / total * 100).round(2)
        df_saved = df_saved.sort_values(by="tCO₂ removed", ascending=False)
        st.table(format_results_summary(df_saved))

        # Export buttons
        safe_name = unicodedata.normalize("NFKD", current_name).encode("ascii", "ignore").decode("ascii")

        st.download_button(
            label="Export JSON",
            data=json.dumps(scenario, indent=2),
            file_name=f"{safe_name}_scenario.json",
            mime="application/json"
        )

        inputs_df = pd.DataFrame.from_dict(scenario["resource_caps"], orient="index", columns=["Available"])
        constraints_df = pd.DataFrame.from_dict({m: v.get("resources", {}) for m, v in scenario["method_constraints"].items()}, orient="index").fillna(0)
        allocations_df = pd.DataFrame.from_dict(scenario["allocations"], orient="index", columns=["tCO₂ removed"])

        csv_buffer = io.BytesIO()
        with zipfile.ZipFile(csv_buffer, "w") as zipf:
            zipf.writestr("resource_caps.csv", inputs_df.to_csv())
            zipf.writestr("method_constraints.csv", constraints_df.to_csv())
            zipf.writestr("allocations.csv", allocations_df.to_csv())

        st.download_button(
            label="Export CSV",
            data=csv_buffer.getvalue(),
            file_name=f"{safe_name}_scenario.zip",
            mime="application/zip"
        )