dashboard/app.py

"""
Mangrove ER Model Dashboard

- Each tab corresponds to a different model/config.
- To add a new model, add its config to configs/ and add an entry to MODEL_CONFIGS below.
- Each tab shows the results for the base config (no parameter editing).
"""

import gradio as gr
from pathlib import Path
import sys
import os

# Add the parent directory to Python path so we can import er_model_core
sys.path.insert(0, str(Path(__file__).parent.parent))

from er_model_core.er_model import ERModel
import pandas as pd
import numpy as np
import plotly.graph_objs as go
from plotly.subplots import make_subplots
import warnings
import traceback
from er_model_core.allometry import calculate_biomass
from dashboard.plots import check_complex, create_growth_increment_plots, create_all_plots
from dashboard.tables import create_survival_table, create_biomass_table, create_project_results_table
import yaml
import tempfile
from dashboard.inputs import get_model_inputs

MODEL_CONFIGS = {
    "Declining Increment": "configs/declining_increment.yaml"
}

# Add a mapping for display names
SPECIES_DISPLAY_NAMES = {
    'species_A': 'Rhizophora spp.',
    'species_B': 'Avicennia germinans'
}

def to_native(obj):
    if isinstance(obj, pd.DataFrame):
        obj = obj.applymap(lambda x: to_native(x))
        obj.index = obj.index.map(to_native)
        obj.columns = obj.columns.map(to_native)
        return obj
    elif isinstance(obj, pd.Series):
        obj = obj.map(to_native)
        obj.index = obj.index.map(to_native)
        return obj
    elif isinstance(obj, dict):
        return {to_native(k): to_native(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [to_native(x) for x in obj]
    elif isinstance(obj, tuple):
        return tuple(to_native(x) for x in obj)
    elif isinstance(obj, (np.generic, np.ndarray)):
        return obj.item() if obj.size == 1 else obj.tolist()
    else:
        return obj

# Helper to update planting schedule in a config file
def update_planting_schedule(config_path, year_areas):
    with open(config_path) as f:
        config = yaml.safe_load(f)
    for i, area in enumerate(year_areas, 1):
        config["project"]["planting_schedule"][f"year_{i}"] = area
    return config

def create_survival_table(model):
    """
    Returns a DataFrame with years as rows and columns for each species and total surviving trees.
    Uses model.species_metrics for all values.
    """
    df = model.species_metrics.copy()
    # Pivot to wide format: years as rows, species as columns
    surv = df.pivot(index="Year", columns="Species", values="Surviving Trees")
    # Rename columns to display names
    surv.columns = [SPECIES_DISPLAY_NAMES.get(sp, sp) for sp in surv.columns]
    # Add total surviving trees column
    surv["Total Surviving Trees"] = surv.sum(axis=1)
    # Format numbers
    for col in surv.columns:
        surv[col] = surv[col].apply(lambda x: f"{x:,.0f}")
    surv = surv.reset_index()
    return surv

def run_model_from_config(config):
    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as tmp:
        yaml.dump(config, tmp)
        tmp_path = tmp.name
    model = ERModel(Path(tmp_path))
    results, species_results = model.run()
    plots = create_all_plots(results, species_results, config)
    summary = create_summary(results, species_results, config, model)
    survival_table = create_survival_table(model)
    # Use modularized table formatting
    results_fmt = create_project_results_table(results)
    species_results_fmt = species_results.copy()
    for col in species_results_fmt.columns:
        if species_results_fmt[col].dtype in [float, int]:
            species_results_fmt[col] = species_results_fmt[col].apply(lambda x: f"{x:,.2f}" if isinstance(x, float) else f"{x:,}")
    return (*plots, summary, results_fmt, species_results_fmt, survival_table)

def run_model(config_path):
    with open(config_path) as f:
        config = yaml.safe_load(f)
    model = ERModel(Path(config_path))
    results, species_results = model.run()
    plots = create_all_plots(results, species_results, config)
    summary = create_summary(results, species_results, config, model)
    survival_table = create_survival_table(model)
    # Use modularized table formatting
    results_fmt = create_project_results_table(results)
    species_results_fmt = species_results.copy()
    for col in species_results_fmt.columns:
        if species_results_fmt[col].dtype in [float, int]:
            species_results_fmt[col] = species_results_fmt[col].apply(lambda x: f"{x:,.2f}" if isinstance(x, float) else f"{x:,}")
    return (*plots, summary, results_fmt, species_results_fmt, survival_table)

def create_summary(results, species_results, config, model):
    total_area = sum(config["project"]["planting_schedule"].values())
    final_gross = results["gross_carbon"].iloc[-1]
    final_buffer = results["buffer_carbon"].iloc[-1]
    # Sum soil carbon over all years
    total_soil_carbon = results["soil_carbon"].sum()
    final_gross_soil = final_gross + total_soil_carbon
    final_buffer_soil = final_buffer + (total_soil_carbon * (1 - config["carbon"]["buffer_percentage"] / 100))
    years = len(results)
    gross_carbon_per_ha = final_gross / total_area if total_area > 0 else 0
    buffer_carbon_per_ha = final_buffer / total_area if total_area > 0 else 0
    gross_carbon_per_ha_soil = final_gross_soil / total_area if total_area > 0 else 0
    buffer_carbon_per_ha_soil = final_buffer_soil / total_area if total_area > 0 else 0
    annual_buffer_per_ha = buffer_carbon_per_ha / years if years > 0 else 0
    annual_buffer_per_ha_soil = buffer_carbon_per_ha_soil / years if years > 0 else 0
    soil_carbon_val = config["carbon"].get("soil_carbon_per_ha_per_year", 0)
    # Milestone years
    def get_val(col, idx):
        return results[col].iloc[idx] if len(results) > idx else 0
    year_5_gross = get_val("gross_carbon", 4)
    year_5_buffer = get_val("buffer_carbon", 4)
    year_10_gross = get_val("gross_carbon", 9)
    year_10_buffer = get_val("buffer_carbon", 9)
    year_20_gross = get_val("gross_carbon", 19)
    year_20_buffer = get_val("buffer_carbon", 19)
    avg_annual_gross = final_gross / years if years > 0 else 0
    avg_annual_buffer = final_buffer / years if years > 0 else 0
    avg_annual_gross_soil = final_gross_soil / years if years > 0 else 0
    avg_annual_buffer_soil = final_buffer_soil / years if years > 0 else 0
    # Surviving trees at milestones
    milestones = [5, 10, 20, 30]
    survival_lines = []
    for m in milestones:
        if m <= years:
            surv = model.calculate_total_surviving_trees(m)
            surv_str = ", ".join([f"{SPECIES_DISPLAY_NAMES.get(k, k)}: {v:,.0f}" for k, v in surv.items()])
            survival_lines.append(f"Year {m} surviving trees: {surv_str}")
    # Per hectare (final year)
    surv_30 = model.calculate_total_surviving_trees(years)
    per_ha_lines = []
    for k, v in surv_30.items():
        per_ha = v / total_area if total_area > 0 else 0
        per_ha_lines.append(f"{SPECIES_DISPLAY_NAMES.get(k, k)}: {per_ha:,.0f} trees/ha")
    summary = (
        f"Project Overview:\n----------------\nDuration: {years} years\nTotal Area Planted: {total_area:,.0f} ha\nBuffer Pool: {config['carbon']['buffer_percentage']}%\n\n"
        f"Soil Carbon:\n-----------\nSoil carbon added: {soil_carbon_val} tCO2/ha/year\nTotal soil carbon added over project: {total_soil_carbon:,.0f} tCO2\n\n"
        f"Carbon Sequestration (Biomass Only):\n-------------------\nTotal Gross Carbon (Year {years}): {final_gross:,.0f} tCO2\nTotal Buffer Carbon (Year {years}): {final_buffer:,.0f} tCO2\nAverage Annual Gross: {avg_annual_gross:,.0f} tCO2/yr\nAverage Annual Buffer: {avg_annual_buffer:,.0f} tCO2/yr\n\n"
        f"Carbon Sequestration (With Soil Carbon):\n-------------------\nTotal Gross Carbon (Year {years}): {final_gross_soil:,.0f} tCO2\nTotal Buffer Carbon (Year {years}): {final_buffer_soil:,.0f} tCO2\nAverage Annual Gross: {avg_annual_gross_soil:,.0f} tCO2/yr\nAverage Annual Buffer: {avg_annual_buffer_soil:,.0f} tCO2/yr\n\n"
        f"Milestone Years (Biomass Only):\n--------------\nYear 5 Gross: {year_5_gross:,.0f} tCO2\nYear 5 Buffer: {year_5_buffer:,.0f} tCO2\nYear 10 Gross: {year_10_gross:,.0f} tCO2\nYear 10 Buffer: {year_10_buffer:,.0f} tCO2\nYear 20 Gross: {year_20_gross:,.0f} tCO2\nYear 20 Buffer: {year_20_buffer:,.0f} tCO2\n\n"
        "Surviving Trees (Milestones):\n----------------------------\n"
        + "\n".join(survival_lines)
        + "\n\nPer Hectare Surviving Trees (Final Year):\n----------------------------------------\n"
        + "\n".join(per_ha_lines)
        + f"\n\nPer Hectare Metrics (Year {years}):\n-----------------------------\nGross Carbon per ha (biomass only): {gross_carbon_per_ha:,.0f} tCO2/ha\nBuffer Carbon per ha (biomass only): {buffer_carbon_per_ha:,.0f} tCO2/ha\nGross Carbon per ha (with soil): {gross_carbon_per_ha_soil:,.0f} tCO2/ha\nBuffer Carbon per ha (with soil): {buffer_carbon_per_ha_soil:,.0f} tCO2/ha\nAverage Annual Buffer per ha (biomass only): {annual_buffer_per_ha:,.2f} tCO2/ha/yr\nAverage Annual Buffer per ha (with soil): {annual_buffer_per_ha_soil:,.2f} tCO2/ha/yr\n"
    )
    return summary

def hex_to_rgba(hex_color, alpha):
    hex_color = hex_color.lstrip('#')
    r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
    return f'rgba({r},{g},{b},{alpha})'

def get_all_model_inputs(config):
    # Use get_model_inputs to get all widgets and defaults
    model_inputs, model_defaults = get_model_inputs(config)
    # Unpack widgets for layout
    planting_density_A = model_inputs['planting_density_A']
    r0_dbh_A = model_inputs['r0_dbh_A']
    tm_dbh_A = model_inputs['tm_dbh_A']
    r0_height_A = model_inputs['r0_height_A']
    tm_height_A = model_inputs['tm_height_A']
    planting_density_B = model_inputs['planting_density_B']
    r0_dbh_B = model_inputs['r0_dbh_B']
    tm_dbh_B = model_inputs['tm_dbh_B']
    r0_height_B = model_inputs['r0_height_B']
    tm_height_B = model_inputs['tm_height_B']
    mort_vars = [model_inputs[f'year_{yr}'] for yr in range(1, 11)]
    mort_sub = model_inputs['subsequent']
    project_duration = model_inputs['project_duration']
    buffer_pct = model_inputs['buffer_pct']
    soil_carbon = model_inputs['soil_carbon']
    return (planting_density_A, r0_dbh_A, tm_dbh_A, r0_height_A, tm_height_A,
            planting_density_B, r0_dbh_B, tm_dbh_B, r0_height_B, tm_height_B,
            mort_vars, mort_sub, project_duration, buffer_pct, soil_carbon, model_defaults)

# Gradio UI code starts here
with gr.Blocks() as demo:
    gr.Markdown("""
    <div style='display: flex; align-items: center; gap: 16px;'>
        <img src='dashboard/sequestr_logo.png' alt='Sequestr Logo' style='height:48px;'>
        <div>
            <b>This tool is provided by <a href='https://www.sequestr.io/' target='_blank'>Sequestr</a>.</b><br>
            <span style='color: #d9534f;'>Please do not share this page.</span>
        </div>
    </div>
    """, elem_id="sequestr-header")
    gr.Markdown(r"""
# Mangrove Emissions Reduction (ER) Model & Dashboard

Welcome to the Mangrove ER Model Dashboard, a tool designed by [Sequestr](https://www.sequestr.io/) for simulating and analyzing carbon sequestration in our mangrove restoration project. This dashboard allows users to explore the "Declining Increment" growth model, adjust key project parameters, and visualize the resulting carbon and emission reduction outcomes over time.
""")
    with gr.Accordion("How the Model Works", open=False):
        gr.Markdown(r"""
## How the Model Works

The core of this dashboard is an emissions reduction (ER) model that projects mangrove growth, biomass accumulation, and carbon sequestration based on user-defined parameters and a chosen growth model. Currently, the active model is the **Declining Increment** model.

### Declining Increment Growth Model

This model assumes that the annual growth increment (for both Diameter at Breast Height - DBH, and Height) declines linearly over time, eventually reaching zero. The total size at any given age is the sum of all non-negative annual increments up to that age. This ensures that tree size never decreases.

**1. Discrete Declining Increment:** The size at a given age $t$ (in years) is calculated by summing annual increments:

$$
\mathrm{size}(t) = \mathrm{initial\_size} + \sum_{i=1}^{t} \mathrm{increment}(i)
$$

Where the increment for year $i$ is:

$$
\mathrm{increment}(i) = \max \left(0, r_0 \cdot \left(1 - \frac{i - 0.5}{T_m}\right) \right)
$$

- $\mathrm{initial\_size}$: The initial DBH (cm) or Height (m) at planting (age 0).
- $r_0$: The initial annual growth increment (cm/year for DBH, m/year for Height).
- $T_m$: The time (in years) at which the annual growth increment becomes zero.

**2. Continuous Declining Increment (Analytical Formula):** For a smoother representation, the continuous version uses an analytical formula:

$$
\mathrm{size}(t) = \mathrm{initial\_size} + r_0 \cdot \left(t - \frac{t^2}{2T_m}\right), \quad 0 \leq t \leq T_m
$$

If $t > T_m$, then $\mathrm{size}(t) = \mathrm{size}(T_m)$. This version may yield slightly different results from the discrete sum, especially for short time periods or small $T_m$ values. The dashboard configuration specifies which version is active.

### Survival and Mortality
The number of surviving trees in a cohort is calculated annually. Mortality can be specified using either:

1. **Annual Mortality Rates:** A percentage mortality rate applied each year, potentially varying for the first few years and then stabilizing.
2. **DBH-Dependent Mortality (Optional, if configured):** Mortality rate $m$ is a function of tree DBH:

$$
m = m_\mathrm{ref} \cdot \left( \frac{\mathrm{DBH}_\mathrm{ref}}{\mathrm{DBH}} \right)^p
$$

- $m_\mathrm{ref}$: Reference mortality rate at $\mathrm{DBH}_\mathrm{ref}$.
- $\mathrm{DBH}_\mathrm{ref}$: Reference DBH.
- $p$: Exponent controlling the sensitivity of mortality to DBH. The calculated $m$ is capped (e.g., between 0 and 0.99).
""")
    with gr.Accordion("How Emission Reductions (ERs) are Calculated", open=False):
        gr.Markdown(r"""
## How Emission Reductions (ERs) are Calculated

The model follows these steps to estimate net CO2 emission reductions:

**1. Biomass Calculation:** Total above-ground and below-ground biomass per tree is calculated using allometric equations, which typically relate DBH and Height to biomass. For example, using equations from Zanvo et al. (2023):
- For *Rhizophora racemosa* (Species A):
  $$
  \mathrm{Biomass}_{\mathrm{total}} = 2.0738 \times (\mathrm{DBH}^2 \cdot H)^{0.67628}
  $$
- For *Avicennia germinans* (Species B):
  $$
  \mathrm{Biomass}_{\mathrm{total}} = 1.5595 \times (\mathrm{DBH}^2 \cdot H)^{0.55864}
  $$
  *(Note: The specific equations are defined in the model configuration.)*

**2. Carbon Stock Calculation:** The total carbon stock in living biomass per hectare is calculated:

$$
\mathrm{Carbon\ Stock\ (tC/ha)} = \frac{\mathrm{Biomass\ per\ tree\ (kg)} \times \mathrm{Surviving\ Trees\ per\ ha} \times \mathrm{Biomass\ to\ Carbon\ ratio}}{1000}
$$

- The Biomass to Carbon ratio (e.g., 0.47) converts biomass to carbon mass.

**3. Gross CO2 Sequestration:** The carbon stock is then converted to tons of CO2 equivalent:

$$
\mathrm{Gross\ CO2eq\ (tCO2/ha)} = \mathrm{Carbon\ Stock\ (tC/ha)} \times \mathrm{Carbon\ to\ CO2\ ratio}
$$

- The Carbon to CO2 ratio (typically 3.67) is based on molecular weights.

**4. Soil Carbon (Optional):** If configured, annual soil carbon sequestration is added:

$$
\mathrm{Soil\ Carbon\ CO2eq\ (tCO2/ha/yr)} = \mathrm{User\ Defined\ Value\ (e.g.\ 1.0\ tCO2/ha/yr)}
$$

This is added to the gross CO2 sequestration from biomass.

**5. Net Emission Reductions (ERs):** Adjustments are made to the gross CO2 sequestration (including soil carbon, if applicable) to determine net ERs eligible for crediting:

$$
\mathrm{Net\ ERs} = (\mathrm{Gross\ CO2eq}_{\mathrm{biomass + soil}}) \times (1 - \mathrm{Buffer\ \%}) - (\mathrm{Gross\ CO2eq}_{\mathrm{biomass + soil}} \times \mathrm{Leakage\ \%}) - (\mathrm{Baseline\ Emissions\ per\ ha} \times \mathrm{Area})
$$

- **Buffer Pool**: A percentage deduction to account for risks like project failure or natural disturbances.
- **Leakage**: Emissions occurring outside the project boundary due to project activities (often assumed to be 0% for mangrove projects if activities are self-contained).
- **Baseline Emissions**: Emissions that would have occurred in the absence of the project (e.g., from degrading land).

This dashboard visualizes these values annually over the project duration, providing insights into the project's carbon sequestration potential.
""")
    # --- Key Input Params Card ---
    config = yaml.safe_load(open(MODEL_CONFIGS["Declining Increment"]))
    species_md = "| Species | Planting Density (trees/ha) | r0 (DBH, cm/yr) | Tm (DBH, yrs) | r0 (Height, m/yr) | Tm (Height, yrs) |\n|---|---|---|---|---|---|\n"
    for sp in config["species"]:
        species_md += f"| {sp['name']} | {sp['planting_density']} | {sp['declining_increment']['dbh']['r0']} | {sp['declining_increment']['dbh']['T_m']} | {sp['declining_increment']['height']['r0']} | {sp['declining_increment']['height']['T_m']} |\n"
    # Project details with proper line breaks
    project_md = (
        f"**Project Duration:** {config['project']['duration_years']} years  \n"
        f"**Buffer %:** {config['carbon']['buffer_percentage']}  \n"
        f"**Soil Carbon per ha per year:** {config['carbon']['soil_carbon_per_ha_per_year']} tCO2"
    )
    # Mortality table
    mortality_schedule = config['project'].get('annual_mortality_schedule', {})
    mortality_md = "| Year | Mortality Rate (%) |\n|---|---|\n"
    if mortality_schedule:
        for year, rate in mortality_schedule.items():
            mortality_md += f"| {year} | {rate} |\n"
    else:
        mortality_md += "| (no data) | (no data) |\n"
    # Carbon details
    carbon = config["carbon"]
    carbon_md = (
        f"**Biomass to Carbon:** {carbon['biomass_to_carbon']}  \n"
        f"**Carbon to CO2:** {carbon['carbon_to_co2']}  \n"
        f"**Buffer %:** {carbon['buffer_percentage']}  \n"
        f"**Leakage %:** {carbon['leakage_percentage']}  \n"
        f"**Baseline Emissions:** {carbon['baseline_emissions']} tCO2/ha/year  \n"
        f"**Soil Carbon per ha per year:** {carbon['soil_carbon_per_ha_per_year']} tCO2"
    )
    gr.Markdown("""
### Key Model Inputs
""" + species_md)
    gr.Markdown("""
### Carbon Details
""" + carbon_md)
    with gr.Tabs():
        with gr.Tab("Declining Increment"):
            config = yaml.safe_load(open(MODEL_CONFIGS["Declining Increment"]))
            use_continuous = config.get('continuous_growth', False)
            if use_continuous:
                gr.Markdown("**Mode:** Continuous declining increment growth")
            else:
                gr.Markdown("**Mode:** Discrete declining increment growth (sum of annual increments)")
            gr.Markdown("## Model Inputs")
            with gr.Row():
                # Column 1: Planting Schedule (as gr.Number)
                with gr.Column(scale=1, min_width=220):
                    with gr.Accordion("Planting Schedule (ha per year)", open=True):
                        year_1_default = config['project']['planting_schedule']['year_1']
                        year_2_default = config['project']['planting_schedule']['year_2']
                        year_3_default = config['project']['planting_schedule']['year_3']
                        year_4_default = config['project']['planting_schedule']['year_4']
                        year_5_default = config['project']['planting_schedule']['year_5']
                        year_1 = gr.Number(value=year_1_default, label="Area to be planted in year 1 (ha)", interactive=True)
                        year_2 = gr.Number(value=year_2_default, label="Area to be planted in year 2 (ha)", interactive=True)
                        year_3 = gr.Number(value=year_3_default, label="Area to be planted in year 3 (ha)", interactive=True)
                        year_4 = gr.Number(value=year_4_default, label="Area to be planted in year 4 (ha)", interactive=True)
                        year_5 = gr.Number(value=year_5_default, label="Area to be planted in year 5 (ha)", interactive=True)
                # Column 2: Planting Density (as gr.Slider)
                with gr.Column(scale=1, min_width=220):
                    with gr.Accordion("Planting Density", open=True):
                        pdA = config['species'][0]['planting_density']
                        pdB = config['species'][1]['planting_density']
                        planting_density_A = gr.Slider(minimum=pdA*0.75, maximum=pdA*1.25, value=pdA, step=1, label="Initial planting density for Species A")
                        planting_density_B = gr.Slider(minimum=pdB*0.75, maximum=pdB*1.25, value=pdB, step=1, label="Initial planting density for Species B")
                # Columns 3-4: Mortality Rates (as gr.Slider, split across two half-width columns)
                with gr.Column(scale=2, min_width=300):
                    with gr.Accordion("Mortality Rates (%)", open=True):
                        mort_defaults = [config['species'][0]['mortality_rates'][f'year_{yr}'] for yr in range(1, 11)]
                        mort_sub_default = config['species'][0]['mortality_rates']['subsequent']
                        with gr.Row():
                            with gr.Column():
                                mort_vars_1 = [gr.Slider(minimum=val*0.75, maximum=val*1.25, value=val, step=0.1, label=f"Annual mortality rate for year {i+1}") for i, val in enumerate(mort_defaults[:6])]
                            with gr.Column():
                                mort_vars_2 = [gr.Slider(minimum=val*0.75, maximum=val*1.25, value=val, step=0.1, label=f"Annual mortality rate for year {i+7}") for i, val in enumerate(mort_defaults[6:])]
                                mort_sub = gr.Slider(minimum=mort_sub_default*0.75, maximum=mort_sub_default*1.25, value=mort_sub_default, step=0.1, label="Annual mortality rate for years 11+")
                        mort_vars = mort_vars_1 + mort_vars_2
            # After the mortality columns, add:
            soil_carbon_default = config['carbon']['soil_carbon_per_ha_per_year']

            soil_carbon = gr.Number(value=soil_carbon_default, label="Soil Carbon per ha per year (tCO2)", interactive=True, info="Annual soil carbon addition per hectare")

            # Add a Reset to Defaults button
            reset_btn = gr.Button("Reset to Defaults", variant="secondary")

            def reset_to_defaults():
                # Reload config and return all default values in widget order
                config = yaml.safe_load(open(MODEL_CONFIGS["Declining Increment"]))
                pdA = config['species'][0]['planting_density']
                pdB = config['species'][1]['planting_density']
                mort_defaults = [config['species'][0]['mortality_rates'][f'year_{yr}'] for yr in range(1, 11)]
                mort_sub_default = config['species'][0]['mortality_rates']['subsequent']
                return [
                    config['project']['planting_schedule']['year_1'],
                    config['project']['planting_schedule']['year_2'],
                    config['project']['planting_schedule']['year_3'],
                    config['project']['planting_schedule']['year_4'],
                    config['project']['planting_schedule']['year_5'],
                    pdA, pdB,
                    *mort_defaults, mort_sub_default,
                    config['carbon']['soil_carbon_per_ha_per_year']
                ]
            reset_btn.click(
                reset_to_defaults,
                inputs=[],
                outputs=[year_1, year_2, year_3, year_4, year_5,
                         planting_density_A, planting_density_B,
                         *mort_vars, mort_sub, soil_carbon]
            )

            update_btn = gr.Button("Update Results", variant="primary")
            with gr.Row():
                carbon_plot = gr.Plot()
                biomass_plot = gr.Plot()
                annual_plot = gr.Plot()
            growth_plot = gr.Plot(label="Growth & Increment Curves (DBH/Height)")
            summary_box = gr.Textbox(label="Summary", lines=12)
            # --- Tabbed tables section ---
            with gr.Tabs():
                with gr.Tab("Project Results (Annual)"):
                    results_box = gr.Dataframe(label="Project Results (Annual)")
                    download_results_file = gr.File(label="Download CSV", visible=False)
                    download_results_btn = gr.Button("📥 Download CSV", variant="secondary")
                with gr.Tab("Species Results (Annual)"):
                    species_box = gr.Dataframe(label="Species Results (Annual)")
                    download_species_file = gr.File(label="Download CSV", visible=False)
                    download_species_btn = gr.Button("📥 Download CSV", variant="secondary")
                with gr.Tab("Surviving Trees Table"):
                    survival_box = gr.Dataframe(label="Surviving Trees Table")
                    download_survival_file = gr.File(label="Download CSV", visible=False)
                    download_survival_btn = gr.Button("📥 Download CSV", variant="secondary")
                with gr.Tab("Biomass Table"):
                    biomass_debug_table = gr.Dataframe(label="Biomass Table (inputs & outputs per year)")
                    download_biomass_file = gr.File(label="Download CSV", visible=False)
                    download_biomass_btn = gr.Button("📥 Download CSV", variant="secondary")
            # --- End tabbed tables ---
            # Update the update_declining_increment callback to use these new inputs
            def update_declining_increment(y1, y2, y3, y4, y5,
                                           planting_density_A, planting_density_B,
                                           *mortality_inputs_and_soil_carbon):
                # Unpack the last one from the end
                *mortality_inputs, soil_carbon = mortality_inputs_and_soil_carbon
                config = update_planting_schedule(MODEL_CONFIGS["Declining Increment"], [y1, y2, y3, y4, y5])
                # Update species_A params
                config['species'][0]['planting_density'] = planting_density_A
                # Update species_B params
                config['species'][1]['planting_density'] = planting_density_B
                # Update mortality for BOTH species from the same UI values (years 1-10 + subsequent)
                for i in [0, 1]:
                    for yr_idx, yr in enumerate(range(1, 11)):
                        config['species'][i]['mortality_rates'][f'year_{yr}'] = mortality_inputs[yr_idx]
                    config['species'][i]['mortality_rates']['subsequent'] = mortality_inputs[10]
                # Update project/carbon params
                config['carbon']['soil_carbon_per_ha_per_year'] = soil_carbon
                with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as tmp:
                    yaml.dump(config, tmp)
                    tmp_path = tmp.name
                model = ERModel(Path(tmp_path))
                results, species_results = model.run()
                plots = create_all_plots(results, species_results, config)
                summary = create_summary(results, species_results, config, model)
                survival_table = create_survival_table(model)
                growth_fig = create_growth_increment_plots(config, model_type="declining_increment")
                biomass_debug_df = create_biomass_table(config)
                results_fmt = create_project_results_table(results)
                species_results_fmt = species_results.copy()
                for col in species_results_fmt.columns:
                    if species_results_fmt[col].dtype in [float, int]:
                        species_results_fmt[col] = species_results_fmt[col].apply(lambda x: f"{x:,.2f}" if isinstance(x, float) else f"{x:,}")
                # --- Ensure all DataFrames are native types ---
                results_fmt = to_native(results_fmt)
                species_results_fmt = to_native(species_results_fmt)
                survival_table = to_native(survival_table)
                biomass_debug_df = to_native(biomass_debug_df)
                
                # --- Ensure all other values are native types ---
                if hasattr(summary, 'item'):
                    summary = summary.item()
                return (plots[0], plots[1], plots[2], growth_fig, summary, results_fmt, species_results_fmt, survival_table, biomass_debug_df)
            update_btn.click(
                update_declining_increment,
                inputs=[year_1, year_2, year_3, year_4, year_5,
                        planting_density_A, planting_density_B,
                        *mort_vars, mort_sub, soil_carbon],
                outputs=[carbon_plot, biomass_plot, annual_plot, growth_plot, summary_box, results_box, species_box, survival_box, biomass_debug_table]
            )

            # Download functionality for tables
            def create_download_file(df, filename):
                """Create a temporary CSV file for download"""
                import tempfile
                import os
                
                # Create a temporary file
                temp_dir = tempfile.gettempdir()
                filepath = os.path.join(temp_dir, filename)
                
                # Convert formatted strings back to numbers where possible
                if df is not None and not df.empty:
                    df_clean = df.copy()
                    for col in df_clean.columns:
                        if df_clean[col].dtype == 'object' and col.lower() != 'year':
                            try:
                                # Try to convert formatted strings back to numbers
                                df_clean[col] = df_clean[col].astype(str).str.replace(',', '').astype(float)
                            except:
                                pass  # Keep as string if conversion fails
                    
                    df_clean.to_csv(filepath, index=False)
                    return filepath
                return None
            
            def download_results_csv(y1, y2, y3, y4, y5, planting_density_A, planting_density_B, *mortality_inputs_and_soil_carbon):
                """Generate and download project results CSV"""
                try:
                    # Recreate the results using current parameters
                    *mortality_inputs, soil_carbon = mortality_inputs_and_soil_carbon
                    config = update_planting_schedule(MODEL_CONFIGS["Declining Increment"], [y1, y2, y3, y4, y5])
                    config['species'][0]['planting_density'] = planting_density_A
                    config['species'][1]['planting_density'] = planting_density_B
                    for i in [0, 1]:
                        for yr_idx, yr in enumerate(range(1, 11)):
                            config['species'][i]['mortality_rates'][f'year_{yr}'] = mortality_inputs[yr_idx]
                        config['species'][i]['mortality_rates']['subsequent'] = mortality_inputs[10]
                    config['carbon']['soil_carbon_per_ha_per_year'] = soil_carbon
                    
                    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as tmp:
                        yaml.dump(config, tmp)
                        tmp_path = tmp.name
                    model = ERModel(Path(tmp_path))
                    results, _ = model.run()
                    
                    filepath = create_download_file(results, "project_results.csv")
                    return gr.File(value=filepath, visible=True) if filepath else gr.File(visible=False)
                except Exception as e:
                    print(f"Error creating download: {e}")
                    return gr.File(visible=False)
                
            def download_species_csv(y1, y2, y3, y4, y5, planting_density_A, planting_density_B, *mortality_inputs_and_soil_carbon):
                """Generate and download species results CSV"""
                try:
                    *mortality_inputs, soil_carbon = mortality_inputs_and_soil_carbon
                    config = update_planting_schedule(MODEL_CONFIGS["Declining Increment"], [y1, y2, y3, y4, y5])
                    config['species'][0]['planting_density'] = planting_density_A
                    config['species'][1]['planting_density'] = planting_density_B
                    for i in [0, 1]:
                        for yr_idx, yr in enumerate(range(1, 11)):
                            config['species'][i]['mortality_rates'][f'year_{yr}'] = mortality_inputs[yr_idx]
                        config['species'][i]['mortality_rates']['subsequent'] = mortality_inputs[10]
                    config['carbon']['soil_carbon_per_ha_per_year'] = soil_carbon
                    
                    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as tmp:
                        yaml.dump(config, tmp)
                        tmp_path = tmp.name
                    model = ERModel(Path(tmp_path))
                    _, species_results = model.run()
                    
                    filepath = create_download_file(species_results, "species_results.csv")
                    return gr.File(value=filepath, visible=True) if filepath else gr.File(visible=False)
                except Exception as e:
                    print(f"Error creating download: {e}")
                    return gr.File(visible=False)
                
            def download_survival_csv(y1, y2, y3, y4, y5, planting_density_A, planting_density_B, *mortality_inputs_and_soil_carbon):
                """Generate and download survival table CSV"""
                try:
                    *mortality_inputs, soil_carbon = mortality_inputs_and_soil_carbon
                    config = update_planting_schedule(MODEL_CONFIGS["Declining Increment"], [y1, y2, y3, y4, y5])
                    config['species'][0]['planting_density'] = planting_density_A
                    config['species'][1]['planting_density'] = planting_density_B
                    for i in [0, 1]:
                        for yr_idx, yr in enumerate(range(1, 11)):
                            config['species'][i]['mortality_rates'][f'year_{yr}'] = mortality_inputs[yr_idx]
                        config['species'][i]['mortality_rates']['subsequent'] = mortality_inputs[10]
                    config['carbon']['soil_carbon_per_ha_per_year'] = soil_carbon
                    
                    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as tmp:
                        yaml.dump(config, tmp)
                        tmp_path = tmp.name
                    model = ERModel(Path(tmp_path))
                    model.run()
                    survival_table = create_survival_table(model)
                    
                    # Convert survival table to raw numbers
                    df_raw = model.species_metrics.copy()
                    surv_raw = df_raw.pivot(index="Year", columns="Species", values="Surviving Trees")
                    surv_raw.columns = [SPECIES_DISPLAY_NAMES.get(sp, sp) for sp in surv_raw.columns]
                    surv_raw["Total Surviving Trees"] = surv_raw.sum(axis=1)
                    surv_raw = surv_raw.reset_index()
                    
                    filepath = create_download_file(surv_raw, "surviving_trees.csv")
                    return gr.File(value=filepath, visible=True) if filepath else gr.File(visible=False)
                except Exception as e:
                    print(f"Error creating download: {e}")
                    return gr.File(visible=False)
                
            def download_biomass_csv(y1, y2, y3, y4, y5, planting_density_A, planting_density_B, *mortality_inputs_and_soil_carbon):
                """Generate and download biomass table CSV"""
                try:
                    *mortality_inputs, soil_carbon = mortality_inputs_and_soil_carbon
                    config = update_planting_schedule(MODEL_CONFIGS["Declining Increment"], [y1, y2, y3, y4, y5])
                    config['species'][0]['planting_density'] = planting_density_A
                    config['species'][1]['planting_density'] = planting_density_B
                    for i in [0, 1]:
                        for yr_idx, yr in enumerate(range(1, 11)):
                            config['species'][i]['mortality_rates'][f'year_{yr}'] = mortality_inputs[yr_idx]
                        config['species'][i]['mortality_rates']['subsequent'] = mortality_inputs[10]
                    config['carbon']['soil_carbon_per_ha_per_year'] = soil_carbon
                    
                    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as tmp:
                        yaml.dump(config, tmp)
                        tmp_path = tmp.name
                    model = ERModel(Path(tmp_path))
                    model.run()
                    
                    filepath = create_download_file(model.species_metrics, "biomass_table.csv")
                    return gr.File(value=filepath, visible=True) if filepath else gr.File(visible=False)
                except Exception as e:
                    print(f"Error creating download: {e}")
                    return gr.File(visible=False)
            
            # Connect download buttons to functions
            download_results_btn.click(
                download_results_csv,
                inputs=[year_1, year_2, year_3, year_4, year_5,
                        planting_density_A, planting_density_B,
                        *mort_vars, mort_sub, soil_carbon],
                outputs=[download_results_file]
            )
            
            download_species_btn.click(
                download_species_csv,
                inputs=[year_1, year_2, year_3, year_4, year_5,
                        planting_density_A, planting_density_B,
                        *mort_vars, mort_sub, soil_carbon],
                outputs=[download_species_file]
            )
            
            download_survival_btn.click(
                download_survival_csv,
                inputs=[year_1, year_2, year_3, year_4, year_5,
                        planting_density_A, planting_density_B,
                        *mort_vars, mort_sub, soil_carbon],
                outputs=[download_survival_file]
            )
            
            download_biomass_btn.click(
                download_biomass_csv,
                inputs=[year_1, year_2, year_3, year_4, year_5,
                        planting_density_A, planting_density_B,
                        *mort_vars, mort_sub, soil_carbon],
                outputs=[download_biomass_file]
            )
            
            # Show initial results
            c, b, a, g, summary, r, s, surv, biomass_debug_df = update_declining_increment(
                year_1_default, year_2_default, year_3_default, year_4_default, year_5_default,
                pdA, pdB,
                *mort_defaults, mort_sub_default,
                soil_carbon_default
            )
            carbon_plot.value = c
            biomass_plot.value = b
            annual_plot.value = a
            growth_plot.value = g
            summary_box.value = summary
            results_box.value = r
            species_box.value = s
            survival_box.value = surv
            biomass_debug_table.value = biomass_debug_df
    demo.launch(share=True)