import xarray as xr
import numpy as np
import pandas as pd
from . import utils


class DynsoilModel:
    def __init__(self, params:pd.DataFrame, temp:xr.DataArray=None, runoff:xr.DataArray=None, slope:xr.DataArray=None):
        self.Rgas = 8.314

        self.params = params
        try:
            self.ke = self.params['ke']
        except:
            self.ke = self.params['# ke']
        
        for v in self.params.columns[1:]:
            self.__dict__[v] = self.params[v]

        self.temp = temp
        self.runoff = runoff
        self.slope = slope

    def run(self, idx=0, verbose=True):
        # optimal regolith production rate:
        self.RPopt = self.reg_prod_opt(idx=idx)
        self.RPopt.name = 'RPopt'
        self.RPopt.attrs['units'] = 'N/A'  # TODO
        self.RPopt.attrs['long_name'] = 'Optimal Regolith Production Rate'
        if verbose: utils.p_success('>>> DynsoilModel.RPopt created')

        # regolith erosion rate:
        self.E = self.erosion(idx=idx)
        self.E.name = 'E'
        self.E.attrs['units'] = 'N/A'  # TODO
        self.E.attrs['long_name'] = 'Regolith Erosion Rate'
        if verbose: utils.p_success('>>> DynsoilModel.E created')

        # regolith thickness
        self.h = self.eq_reg_thick(self.RPopt, self.E, idx=idx)
        self.h.name = 'h'
        self.h.attrs['units'] = 'N/A'  # TODO
        self.h.attrs['long_name'] = 'Regolith Thickness'
        if verbose: utils.p_success('>>> DynsoilModel.h created')

        # dissolution constant:
        self.Kmain = self.dissolution_constant(idx=idx)
        self.Kmain.name = 'Kmain'
        self.Kmain.attrs['units'] = 'N/A'  # TODO
        self.Kmain.attrs['long_name'] = 'Dissolution Constant'
        if verbose: utils.p_success('>>> DynsoilModel.Kmain created')

        # Primary phases proportion at surface:
        self.xs = self.eq_x_p_surface(self.h, self.E, self.Kmain, idx=idx)
        self.xs.name = 'xs'
        self.xs.attrs['units'] = 'N/A'  # TODO
        self.xs.attrs['long_name'] = 'Primary Phases Proportion at Surface'
        if verbose: utils.p_success('>>> DynsoilModel.xs created')

        # Weathering rate:
        # self.W = xr.where(self.runoff > 0, self.steady_state_weathering(self.E, self.xs), 0)
        self.W = self.steady_state_weathering(self.E, self.xs)
        self.W.name = 'W'
        self.W.attrs['units'] = 'N/A'  # TODO
        self.W.attrs['long_name'] = 'Weathering Rate'
        if verbose: utils.p_success('>>> DynsoilModel.W created')


    def erosion(self, runoff_only=True, idx=0):
        """
        Calculate erosion rate.

        Parameters:
        temp : float
            Temperature in °C.
        runoff : float
            Runoff in m/y.
        slope : float
            Slope.

        Returns:
        float
            Erosion rate.
        """
        if runoff_only:
            return self.ke[idx] * (self.runoff**self.a[idx]) * (self.slope**self.b[idx])
        else:
            return self.ke[idx] * (self.runoff**self.a[idx]) * (self.slope**self.b[idx]) * np.maximum(2, self.temp)  # BQART exponents


    def reg_prod_opt(self, idx=0):
        """
        Calculate optimal regolith production rate.

        Parameters:
        temp : float
            Temperature in °C.
        runoff : float
            Runoff in m/y.

        Returns:
        float
            Optimal regolith production rate.
        """
        return self.krp[idx] * self.runoff * np.exp(-(self.Ea_rp[idx] / self.Rgas) * (1.0 / (self.temp + 273.15) - 1.0 / self.T0_rp[idx]))


    def soil_prod_func(self, h_soil, idx=0):
        """
        Calculate soil production function.

        Parameters:
        h_soil : float
            Soil thickness.

        Returns:
        float
            Soil production rate.
        """
        return self.h0[idx] / h_soil  # Inverse function


    def eq_reg_thick(self, RPopt, E, idx=0):
        """
        Calculate equilibrium regolith thickness.

        Parameters:
        RPopt : float
            Optimal regolith production rate.
        E : float
            Erosion rate.

        Returns:
        float
            Equilibrium regolith thickness.
        """
        return np.maximum(self.h0[idx] * np.log(RPopt / E), 0.0)  # Exponential function


    def dissolution_constant(self, idx=0):
        """
        Calculate rock dissolution constant.

        Parameters:
        temp : float
            Temperature in °C.
        runoff : float
            Runoff in m/y.

        Returns:
        float
            Rock dissolution constant.
        """
        return self.kd[idx] * (1 - np.exp(-self.kw[idx] * self.runoff)) * np.exp((self.Ea[idx] / self.Rgas) * (1.0 / self.T0[idx] - 1.0 / (self.temp + 273.15)))


    def eq_x_p_surface(self, h, E, K, idx=0):
        """
        Calculate primary phases proportion at surface at steady-state.

        Parameters:
        h : float
            Soil thickness.
        E : float
            Erosion rate.
        K : float
            Dissolution constant.

        Returns:
        float
            Primary phases proportion at surface.
        """
        return np.exp(-K * ((h / E) ** (self.sigma[idx] + 1)) / (self.sigma[idx] + 1))

    def steady_state_weathering(self, E, x_p_surf):
        return E * (1 - x_p_surf)
