"""
The purpose of this script is to compute the NPP coresponding to cbm's runs
"""

from functools import cached_property
from typing import Union, List
import warnings
import pandas
import numpy as np
from eu_cbm_hat.constants import eu_cbm_data_pathlib

def npp_components(df: pandas.DataFrame, groupby: Union[List[str], str]):
    """Compute the Net Annual Increment and Gross Annual Increment on climates, for npp-nai analysis

    Based on stock change and movements to the product pools as well as
    turnover and mouvements to air.
  
        >>> from eu_cbm_hat.core.continent import continent
        >>> from eu_cbm_hat.post_processor.nai import NAI_AGG_COLS
        >>> from eu_cbm_hat.post_processor.nai import compute_nai_gai
        >>> runner = continent.combos['reference'].runners['ZZ'][-1]
        >>> index = ["status"]
        >>> nai_st = runner.post_processor.npp.compute_npp

    """
   
    if isinstance(groupby, str):
        groupby = [groupby]
    # Order by groupby variables, then years
    df.sort_values(groupby + ["year"], inplace=True)

    # Check that there are no duplications over the groupby variables plus year
    #selector = df[["year"] + groupby].duplicated(keep=False)
    #if any(selector):
    #    msg = "The following rows have duplications along the groupby variables.\n"
    #    msg += f"{df.loc[selector, ['year'] + groupby ]}"
    #    msg += "\nPlease aggregate first along the groupby variables and year:\n"
    #    msg += f"{['year'] + groupby }\n Then run this function.\n"
    #    raise ValueError(msg)

    # exclude disturbances
    df = df[(df['status'].isin(["ForAWS", "ForNAWS"]))] #& (df['disturbance_type'] == 0)] # exclude all stands with antropogenic loss
    
    # compute stock in living biomass      
    df["biomass_stock"] = df["merch"] + df["other"]+ df["foliage"]+ df["roots"]

    # Compute NAI for the merchantable pool
    df["biomass_loss"] = df[[#"net_agb",
                    # biomass to production
                    "merch_prod", # this is zero for dist=0
                    "oth_prod",# this is zero for dist=0
                    # transfers to non living pools through disturbances
                    'disturbance_merch_litter_input',
                    'disturbance_oth_litter_input',
                    'disturbance_fol_litter_input',
                    'disturbance_coarse_litter_input',
                    'disturbance_fine_litter_input',        
                    'disturbance_merch_to_air',
                    'disturbance_oth_to_air',
                    'disturbance_fol_to_air',
                    'disturbance_coarse_to_air',
                    'disturbance_fine_to_air',
                    # transfers to non living pools through natural decay
                    'turnover_merch_litter_input',
                    'turnover_oth_litter_input',
                    'turnover_fol_litter_input',
                    'turnover_coarse_litter_input',
                    'turnover_fine_litter_input']].sum(axis=1)

    df =df[['year', 'status', 'forest_type', 'region',
            'mgmt_type', 'mgmt_strategy', 'climate', 'con_broad', 'site_index',
            'area', 'biomass_stock', 'biomass_loss']]

    # Group and sum first
    df_grouped = df.groupby(['year', 'status', 'climate', 'con_broad'])[['area', 'biomass_stock', 'biomass_loss']].sum().reset_index()
    
    # Sort by year to ensure proper chronological order
    df_grouped = df_grouped.sort_values(['year','status', 'climate', 'con_broad'])
    
    # Calculate net biomass change (difference in stock)
    df_grouped["net_biomass"] = df_grouped.groupby(['status', 'climate', 'con_broad'], observed=True)["biomass_stock"].diff()
    
    # Calculate NPP (Net Primary Productivity)
    df_grouped['npp'] = df_grouped["net_biomass"] + df_grouped["biomass_loss"]
    return df_grouped


class NPP:
    """Compute the net growth
    Usage:

        >>> # Net Annual Increment of the merchantable pool (nai_merch) and of
        >>> # all the above ground biomass (nai_agb) by status
        >>> nai_st = runner.post_processor.nai.df_agg(["status"])
        >>> selector = nai_st["status"] == 'ForAWS'
        >>> nai_st.loc[selector, ["year", "status", "area", "nai_merch", "nai_agb"]].head()
     
        """

    def __init__(self, parent):
        self.parent = parent
        self.runner = parent.runner
        self.combo_name = self.runner.combo.short_name

    @cached_property
    def compute_npp(self):
        """Merchantable pools and fluxes aggregated at the classifiers level"""
        df = self.parent.pools_fluxes_morf
        df_out = npp_components(df, groupby= ['year','status','climate', 'con_broad'])
        df_out = df_out.groupby(['year', 'climate', 'con_broad'])[['area','npp']].sum().reset_index()
        df_out['npp_ha_sim'] = df_out['npp']/df_out['area']
        df_out=df_out[['year', 'climate', 'con_broad', 'npp_ha_sim']]
        return df_out
