import json
import pandas as pd
import plotly.express as px
import streamlit as st
from config.config import CO2_UNIT, hidden_resource_names
from utils.data_utils import is_valid_scenario
from tabs.tab4_optimisation import parse_zip_scenario
from utils.ui_helpers import apply_chart_style, register_chart_data, show_bar


def panel_load_scenario():
    """Render the file uploader panel to load one or more scenarios for comparison.

    Accepts JSON and ZIP files. Adds valid scenarios to st.session_state.scenarios.

    Returns:
        None
    """
    st.markdown("""
        <p class="crra-sub-heading">Load a scenario from file</p>""",
        unsafe_allow_html=True,
    )
    st.markdown('<p class="caption">Accepts `.json` or `.zip` (CSV export)</p>', unsafe_allow_html=True)

    uploaded = st.file_uploader("Upload scenario", type=["json", "zip"], key="comparison_loader", accept_multiple_files=True)
    if uploaded is None:
        return
    for uploaded_file in uploaded:
        if uploaded_file is not None:
            try:
                if uploaded_file.name.endswith(".json"):
                    data = json.load(uploaded_file)
                    if not is_valid_scenario(data):
                        st.error(f"Invalid scenario format in {uploaded_file.name}.")
                        continue
                else:
                    data = parse_zip_scenario(uploaded_file)
                    if data is None:
                        st.error(f"Invalid ZIP format in {uploaded_file.name}. Expected: resource_caps.csv, method_constraints.csv, method_costs.csv, allocations.csv")
                        continue
                st.session_state.scenarios[uploaded_file.name] = data
            except Exception as e:
                st.error(f"Error loading {uploaded_file.name}: {e}")

def prepare_comparison_data(selected):
    """Build comparison DataFrames from the currently selected scenarios.

    Args:
        selected (list[str]): list of scenario names to include.

    Returns:
        tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
            df_summary — one row per scenario with total CO₂ and per-method values,
            df_methods — long-format method allocation data,
            df_resources — long-format resource usage data per method and scenario.
    """
    all_methods = sorted({
        m for name in selected
        for m in st.session_state.scenarios[name].get("allocations", {})
    })
    # Built from the stored resource_usage keys (not resource_caps): a custom
    # batch pooled with a standard resource (e.g. "Basalt A" complementing
    # "Basalt mineral") is never added to resource_caps — only resource_usage
    # tracks it — so sourcing the list from resource_caps would silently drop
    # every such batch from the comparison entirely.
    all_resources = sorted({
        r for name in selected
        for r in st.session_state.scenarios[name].get("resource_usage", {})
        if r not in hidden_resource_names
    })

    summary_rows, method_alloc_rows, resource_usage_rows = [], [], []

    for name in selected:
        scenario = st.session_state.scenarios[name]
        allocations = scenario.get("allocations", {})
        # Use the resource_usage already computed by the LP at save time —
        # NOT a recomputation from method_costs × allocations, which double-
        # counts standard resources that are pooled with a custom batch (the
        # LP only attributes the standard's own share to it, not the batch's).
        resource_usage = scenario.get("resource_usage", {})

        row = {
            "Scenario": name,
            f"CO₂ Removed ({CO2_UNIT})": scenario.get("summary", {}).get("total_co2_removed", 0),
        }

        for method in all_methods:
            val = allocations.get(method, 0)
            row[f"{method} ({CO2_UNIT})"] = val
            method_alloc_rows.append({"Scenario": name, "Method": method, f"{CO2_UNIT}": val})

        for resource in all_resources:
            per_method = resource_usage.get(resource, {})
            total_used = sum(per_method.values())
            for method, used in per_method.items():
                if used > 0:
                    resource_usage_rows.append({
                        "Scenario": name, "Method": method,
                        "Resource": resource, "Used": used,
                    })
            row[f"{resource} used"] = total_used

        summary_rows.append(row)

    return (
        pd.DataFrame(summary_rows),
        pd.DataFrame(method_alloc_rows),
        pd.DataFrame(resource_usage_rows),
    )


def show_method_allocations(df):
    """Render a grouped bar chart of method allocations across scenarios.

    Args:
        df (pd.DataFrame): long-format data with columns Scenario, Method, CO₂ value.

    Returns:
        None
    """
    if df.empty:
        return
    st.markdown("""<p class="crra-sub-heading">Method Allocations</p>""", unsafe_allow_html=True)
    fig = px.bar(df, x="Method", y=f"{CO2_UNIT}", color="Scenario", barmode="group")
    fig.update_layout(xaxis_tickangle=-45)
    apply_chart_style(fig)
    st.plotly_chart(fig, width='stretch')
    register_chart_data("Scenario comparison — method allocations", df)



def show_resource_usage(df):
    """Render grouped bar charts of resource usage per method and total per scenario.

    Args:
        df (pd.DataFrame): long-format data with columns Scenario, Method, Resource, Used.

    Returns:
        None
    """
    if df.empty:
        st.info("No resource usage data available (method_costs missing from scenarios?).", icon=":material/info:")
        return

    st.markdown("""<p class="crra-sub-heading">Resource Use per Method</p>""", unsafe_allow_html=True)
    for method in df["Method"].unique():
        df_m = df[df["Method"] == method]
        fig = px.bar(df_m, x="Resource", y="Used", color="Scenario", barmode="group", title=method)
        fig.update_layout(xaxis_tickangle=-45)
        apply_chart_style(fig)
        st.plotly_chart(fig, width='stretch')

    st.markdown("""<p class="crra-sub-heading">Total Resource Use per Scenario</p>""", unsafe_allow_html=True)
    df_total = df.groupby(["Scenario", "Resource"]).sum(numeric_only=True).reset_index()
    fig = px.bar(df_total, x="Scenario", y="Used", color="Resource", barmode="group")
    fig.update_layout(xaxis_tickangle=-45)
    apply_chart_style(fig)
    st.plotly_chart(fig, width='stretch')
    register_chart_data("Scenario comparison — resource usage", df_total)



def render_tab():
    """Render the Scenario Comparison tab.

    Returns:
        None
    """
    st.markdown('<h1 class="crra-heading">Scenario Comparison</h1>', unsafe_allow_html=True)
    show_bar()

    st.session_state.setdefault("scenarios", {})

    panel_load_scenario()

    if not st.session_state.scenarios:
        st.info("No scenarios saved yet. Run an optimisation in **Portfolio generation** tab and **save the scenario** or load a file above.", icon=":material/info:")
        return

    # Select & manage scenarios
    st.markdown("""<p class="crra-sub-heading">Select Scenarios to Compare</p>""", unsafe_allow_html=True)
    scenario_names = list(st.session_state.scenarios.keys())
    selected = st.multiselect("Select scenarios", scenario_names, default=scenario_names)

    for name in selected:
        col1, col2 = st.columns([10, 1])
        with col1:
            st.markdown(f"**{name}**")
        with col2:
            if st.button(label="", key=f"remove_{name}", icon=":material/delete:", help=f"Remove scenario {name} from comparison"):
                st.session_state.scenarios.pop(name, None)
                st.rerun()

    if not selected:
        st.warning("Please select at least one scenario.", icon=":material/warning:")
        return

    # Comparison
    df_summary, df_methods, df_resources = prepare_comparison_data(selected)

    st.markdown("""<p class="crra-sub-heading">Summary Table</p>""", unsafe_allow_html=True)
    st.dataframe(df_summary.style.format(na_rep="-", thousands=","), width='stretch')

    show_method_allocations(df_methods)
    show_resource_usage(df_resources)