import xarray as xr
import numpy as np
import pandas as pd
from tqdm import tqdm
from typing import List
from . import utils
from .dynsoil import DynsoilModel

class GDSSModel:
    def __init__(self, temp:xr.DataArray=None, runoff:xr.DataArray=None, slope:xr.DataArray=None, lith_frac:xr.DataArray=None, land_area:xr.DataArray=None):
        self.temp = temp
        self.runoff = runoff
        self.slope = slope
        self.slope.attrs['long_name'] = 'slope'
        self.land_area = land_area
        self.lith_frac = lith_frac

    def interp_climate(self, CO2:float=None, vns:List[str]=None):
        ''' Interpolate the climate variable to a given CO2 level
        '''
        for vn in vns:
            clim = self.__dict__[vn]
            CO2_levels = clim.lvl
            nlev = len(CO2_levels)
            kinf = 0
            ksup = nlev - 1

            # Find by dichotomy (binary search):
            while ksup - kinf > 1:
                k = kinf + (ksup - kinf) // 2
                if CO2 >= CO2_levels[k]:
                    kinf = k
                else:
                    ksup = k

            # Interpolation coefficient:
            xi = (CO2 - CO2_levels[kinf]) / (CO2_levels[ksup] - CO2_levels[kinf])
            clim_interp = (1 - xi)*clim.isel(lvl=kinf) + xi*clim.isel(lvl=ksup)
            clim_interp.name = self.__dict__[vn].name
            clim_interp.attrs = self.__dict__[vn].attrs
            self.__dict__[vn] = clim_interp
            utils.p_success(f'>>> GDSSModel.{vn} interpolated to the CO2 level: {CO2}')

    def run_dynsoil(self, params:pd.DataFrame=None):
        self.dynsoil_mdl = DynsoilModel(params=params, temp=self.temp, runoff=self.runoff, slope=self.slope)
        utils.p_success('>>> GDSSModel.dynsoil_mdl created')

        self.CaMg_rock = self.dynsoil_mdl.params.filter(like='CaMg')
        type_idx = [int(col.split('_')[-1]) for col in self.CaMg_rock.columns if col.startswith('CaMg_')]
        lith_frac_sum = self.lith_frac.isel(lith=type_idx).sum(dim='lith')
        self.lith_frac_norm = self.lith_frac.isel(lith=type_idx) / lith_frac_sum
        utils.p_success('>>> GDSSModel.lith_frac_norm created')

        self.nruns = len(params)
        utils.p_header(f'>>> GDSSModel.nruns = {self.nruns}')
        self.WCaMg = []
        self.Fsilwth = []
        for i in tqdm(range(self.nruns)):
            self.dynsoil_mdl.run(idx=i, verbose=False)

            WCaMg = np.sum(self.dynsoil_mdl.W * self.lith_frac_norm * self.CaMg_rock.values[i], axis=-1)
            WCaMg.name = 'WCaMg'
            WCaMg.attrs['units'] = 'N/A'  # TODO
            WCaMg.attrs['long_name'] = 'Lithology-dependent CaMg weathering'
            # WCaMg = WCaMg.where(WCaMg != 0, np.nan)
            self.WCaMg.append(WCaMg)
            # self.land_area = self.land_area.where(self.land_area != 0, np.nan)
            Fsilwth = np.sum(self.land_area * WCaMg)
            self.Fsilwth.append(Fsilwth)
        
        self.Fsilwth = np.array(self.Fsilwth)
