dashboard/plots/__init__.py

"""
Plotting functions for the Mangrove ER Dashboard.
"""
# Imports will be updated as needed when moving functions from app.py 
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import numpy as np
import warnings
import traceback
from er_model_core.allometry import calculate_biomass
from er_model_core.er_model import ERModel
from pathlib import Path
from er_model_core.growth_models.chapman_richards import chapman_richards_growth
from er_model_core.growth_models.linear import linear_growth, linear_plateau_growth
from er_model_core.growth_models.declining_increment import declining_increment_growth, continuous_declining_increment_growth

# Helper for checking complex numbers (copied from app.py)
def check_complex(arr, label):
    if np.iscomplexobj(arr):
        complex_indices = np.where(np.iscomplex(arr))[0]
        warnings.warn(f"[ERROR] Complex values in {label}: indices={complex_indices}, values={arr[complex_indices]}")
        traceback.print_stack()

def create_growth_increment_plots(config, model_type=None):
    """
    Create a 2x2 grid of growth and increment plots for DBH and Height using Plotly.
    model_type: 'chapman_richards', 'linear', 'linear_plateau', or 'declining_increment'
    """
    if model_type is None:
        model_type = config.get('growth_model', 'Unknown Model')
    N = config["project"]["duration_years"]
    ages = np.arange(0, N + 1)  # 0 to N
    ages_inc = np.arange(1, N + 1)  # 1 to N
    fig = make_subplots(rows=2, cols=2, subplot_titles=("DBH Growth", "HEIGHT Growth", "DBH Annual Increment", "HEIGHT Annual Increment"))
    from er_model_core.er_model import ERModel
    SPECIES_DISPLAY_NAMES = {
        'species_A': 'Rhizophora spp.',
        'species_B': 'Avicennia germinans'
    }
    use_continuous = config.get('continuous_growth', False)
    for sp in config["species"]:
        name = SPECIES_DISPLAY_NAMES.get(sp["name"], sp["name"])
        initial_dbh = sp["initial_values"]["dbh"]
        initial_height = sp["initial_values"]["height"]
        if model_type == "linear":
            dbh = [linear_growth(t, sp["linear"]["dbh"], initial_dbh) for t in ages]
            height = [linear_growth(t, sp["linear"]["height"], initial_height) for t in ages]
        elif model_type == "linear_plateau":
            dbh = [linear_plateau_growth(t, sp["linear_plateau"]["dbh"], initial_dbh) for t in ages]
            height = [linear_plateau_growth(t, sp["linear_plateau"]["height"], initial_height) for t in ages]
        elif model_type == "declining_increment":
            if use_continuous:
                dbh = [continuous_declining_increment_growth(t, sp["declining_increment"]["dbh"], initial_dbh) for t in ages]
                height = [continuous_declining_increment_growth(t, sp["declining_increment"]["height"], initial_height) for t in ages]
            else:
                dbh = [declining_increment_growth(t, sp["declining_increment"]["dbh"], initial_dbh) for t in ages]
                height = [declining_increment_growth(t, sp["declining_increment"]["height"], initial_height) for t in ages]
        else:
            dbh = [chapman_richards_growth(t, sp["chapman_richards"]["dbh"], initial_dbh) for t in ages]
            height = [chapman_richards_growth(t, sp["chapman_richards"]["height"], initial_height) for t in ages]
        dbh = np.array(dbh)
        height = np.array(height)
        check_complex(dbh, f"{name} DBH (growth)")
        check_complex(height, f"{name} Height (growth)")
        dbh_inc = dbh[1:] - dbh[:-1]
        height_inc = height[1:] - height[:-1]
        check_complex(dbh_inc, f"{name} DBH Δ (increment)")
        check_complex(height_inc, f"{name} Height Δ (increment)")
        fig.add_trace(go.Scatter(x=ages, y=dbh, mode='lines', name=f"{name} DBH", legendgroup=name, line=dict(width=2)), row=1, col=1)
        fig.add_trace(go.Scatter(x=ages, y=height, mode='lines', name=f"{name} Height", legendgroup=name, line=dict(width=2)), row=1, col=2)
        fig.add_trace(go.Scatter(x=ages_inc, y=dbh_inc, mode='lines', name=f"{name} DBH Δ", legendgroup=name, showlegend=False, line=dict(width=2)), row=2, col=1)
        fig.add_trace(go.Scatter(x=ages_inc, y=height_inc, mode='lines', name=f"{name} Height Δ", legendgroup=name, showlegend=False, line=dict(width=2)), row=2, col=2)
    fig.update_layout(height=700, width=900, title_text="Growth and Increment Curves", hovermode="x unified")
    fig.update_xaxes(title_text="Age (years)", row=1, col=1)
    fig.update_xaxes(title_text="Age (years)", row=1, col=2)
    fig.update_xaxes(title_text="Age (years)", row=2, col=1)
    fig.update_xaxes(title_text="Age (years)", row=2, col=2)
    fig.update_yaxes(title_text="Size (cm)", row=1, col=1)
    fig.update_yaxes(title_text="Size (m)", row=1, col=2)
    fig.update_yaxes(title_text="Annual increment (cm/year)", row=2, col=1)
    fig.update_yaxes(title_text="Annual increment (m/year)", row=2, col=2)
    return fig


def create_all_plots(results, species_results, config):
    # Diagnostic: Check for complex numbers in results DataFrame columns used for plotting
    for col in ["gross_carbon", "buffer_carbon"]:
        arr = np.array(results[col])
        check_complex(arr, f"results['{col}']")
    # 1. Carbon curve (gross and buffer, with and without soil)
    fig1 = go.Figure()
    fig1.add_trace(go.Scatter(x=results["year"], y=results["gross_carbon"], mode="lines+markers", name="Gross Carbon", line=dict(width=2, color="blue")))
    fig1.add_trace(go.Scatter(x=results["year"], y=results["buffer_carbon"], mode="lines+markers", name="Buffer Carbon", line=dict(width=2, color="red")))
    if "gross_carbon_with_soil" in results.columns:
        # For backward compatibility, but we now sum soil carbon in summary, so recalc here for plot
        soil_cumsum = results["soil_carbon"].cumsum()
        gross_with_soil = results["gross_carbon"] + soil_cumsum
        buffer_with_soil = results["buffer_carbon"] + (soil_cumsum * (1 - config["carbon"]["buffer_percentage"] / 100))
        fig1.add_trace(go.Scatter(x=results["year"], y=gross_with_soil, mode="lines+markers", name="Gross Carbon (with Soil)", line=dict(width=2, dash="dot", color="green")))
        fig1.add_trace(go.Scatter(x=results["year"], y=buffer_with_soil, mode="lines+markers", name="Buffer Carbon (with Soil)", line=dict(width=2, dash="dot", color="#43c6ac")))
    fig1.update_layout(title="Carbon Sequestration Over Time", xaxis_title="Year", yaxis_title="Carbon (tCO2)", hovermode="x unified", template="plotly_white")

    # 2. Annual buffer carbon (ERs)
    annual_ers = results["buffer_carbon"].diff().fillna(results["buffer_carbon"].iloc[0])
    arr = np.array(annual_ers)
    check_complex(arr, "annual_ers")
    fig2 = go.Figure()
    fig2.add_trace(go.Bar(x=results["year"], y=annual_ers, name="Annual Buffer ERs", marker_color="#2ecc71", opacity=0.7))
    fig2.update_layout(title="Annual Emission Reductions (Post Buffer)", xaxis_title="Year", yaxis_title="Annual Buffer ERs (tCO2)", hovermode="x unified", template="plotly_white")

    # 3. Biomass per tree for all species (refactored to use model.species_metrics)
    years = results["year"]
    fig3 = go.Figure()
    import tempfile, yaml
    from er_model_core.er_model import ERModel
    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()
    df = model.species_metrics.copy()
    SPECIES_DISPLAY_NAMES = {
        'species_A': 'Rhizophora spp.',
        'species_B': 'Avicennia germinans'
    }
    for sp in config["species"]:
        name = SPECIES_DISPLAY_NAMES.get(sp["name"], sp["name"])
        # Get per-tree biomass for this species
        sp_biomass = df[df["Species"] == sp["name"]].set_index("Year")["Biomass per Tree (kg)"]
        # Ensure correct order and fill missing years with NaN
        sp_biomass = sp_biomass.reindex(years)
        fig3.add_trace(go.Scatter(x=years, y=sp_biomass, mode="lines+markers", name=name, line=dict(width=2)))
    fig3.update_layout(title="Total Biomass per Tree", xaxis_title="Year since planting", yaxis_title="Biomass (kg per tree)", hovermode="x unified", template="plotly_white")
    return (fig1, fig2, fig3)