import pandas as pd
import numpy as np
from collections import OrderedDict
from scipy.integrate import solve_ivp
from scipy.interpolate import interp1d, pchip
import matplotlib.pyplot as plt
import matplotlib as mpl

mpl.rcParams['lines.linewidth']=2
mpl.rcParams['font.size']=16

# This is the disequilibrium melting model described in Elkins and Spiegelman (2021),
# applied here to stable element concentrations.

class Disequilibrium:
    ''' A class for calculating element concentrations using scaled chemical reactivity and
    disequilibrium scenarios, as described in Elkins and Spiegelman (2021).

    Interface:
    ----------
        model(C0s,D,F,dFdz,rho_f=2800.,rho_s=3300.,method=method,Da=inf)

    Parameters:
    -----------
        C0s     :   numpy array of initial concentrations
        D       :  Function D(z) -- returns an array of partition coefficents at scaled height z
        F       :  Function F(z) -- returns the degree of melting F
        dFdz    :  Function dFdz(z) --  returns the derivative of F
        rho_f   :  float -- melt density
        rho_s   :  float --  solid density
        method  :  string -- ODE time-stepping scheme to be passed to solve_ivp (one of 'RK45', 'Radau', 'BDF')
        Da      :  float -- Dahmköhler Number (defaults to \inf, unused in equilibrium model)

    Required Method:
    ----------------
        model.solve(): returns depth and log concentration numpy arrays  z, Us, Uf
    '''
    def __init__(self,C0s,D,F,dFdz,rho_f=2800.,rho_s=3300.,method='Radau',Da=0.):
        self.C0s = C0s
        self.N = len(C0s)
        self.D = lambda zp: np.array([D[i][0](zp) for i in range(self.N)])

        self.D0 = self.D(0.)
        self.F  = F
        self.dFdz = dFdz
        self.Da = Da
        self.rho_f = rho_f
        self.rho_s = rho_s
        self.method = method


    def rhs(self,z,U):
        ''' Returns right hand side of decay chain problem for the log of concentration of the solid
        U^s = U[:N]  where N=number of elements
        and
        U^f = U[N:]

        The full equations for dU/dz are given in Elkins and Spiegelman (2021), and here are calculated
        without radiaoctive terms for stable elements.
        Log concentrations are calculated similar to methods in Spiegelman (2000).

        C0, D, D0 are set by the UserCalc driver routine.
        '''

        # Determine F(z) and D(z):
        F = self.F(z)
        dFdz = self.dFdz(z)
        D = np.array(self.D(z))

        # Identify the initial concentration values:
        C0s = self.C0s

        # Identify the initial values for the partition coefficients:
        D0 = self.D0

        # Define the Dahmköhler number:
        Da = self.Da

        # Separate U into solid and melt components:
        N = self.N  # number of elements
        Us = U[:N]
        Uf = U[N:]

        # Calculate the change in log solid concentrations:
        dUs = ((1. - 1./D)*dFdz + Da/D*(D/D0*np.exp(Uf-Us)-1.)) / (1. - F)

        # Calculate the change in log liquid concentrations:
        if F == 0.:
            dUf = (dUs*(dFdz + Da)/dFdz) / (2. + Da/dFdz)
        else:
            dUf = (D0/D*np.exp(Us-Uf)-1.)*(dFdz + Da)/F

        # Return the full right-hand side of the equation:
        dU=np.zeros(2*N)
        dU[:N] = dUs
        dU[N:] = dUf

        return dU


    def solve(self,z_eval=None):
        ''' Solves the concentration equation as an ODE initial value problem.
        If z_eval = None, every point is saved
        Else will save output at every z_eval depth
        '''
        # Set initial conditions and solve the ODE:
        N = self.N  # number of elements
        U0 = np.zeros(2*N)
        try:
            sol = solve_ivp(self.rhs,(0.,1.),U0,t_eval=z_eval,method=self.method)
            z = sol.t
            U = sol.y
        except ValueError as err:
            print('Warning:  solver did not complete, returning NaN: {}'.format(err))
            z = np.array([1.])
            U = np.NaN*np.ones((2*N,len(z)))

        # Return the log concentrations with depth:
        Us = U[:N,:]
        Uf = U[N:,:]
        return z,Us,Uf


# This is the main stableCalc driver class:
class stableCalc:
    ''' A class for constructing solutions for 1-D, steady-state, open-system  transport calculations
        as in Spiegelman (2000) and Elkins and Spiegelman (2021), here for stable elements.

        Usage:
        ------

            us = stableCalc(df,dPdz = 0.32373,tol=1.e-6,model=Disequilibrium,Da=0.1,method='Radau')

        Parameters:
        -----------
            df     : A pandas dataframe with columns ['P','F', Kr','Di']
            dPdz   : float -- Pressure gradient, to convert pressure P to depth z
            tol    : float -- Tolerance for the ODE solver
            Da     : float -- Optional Da number for disequilibrium transport model
            method : string
                ODE time-stepping method to pass to solve_ivp (usually one of 'Radau', 'BDF', or 'RK45')

        Principal Methods:
        --------
            solve_1D              :  1D column solution for a single Decay chain
                                        with arbitrary D, lambda, alpha_0
            solve_all_1D          : Solves a single column model for both 238U and 235U chains.
                                        returns a pandas dataframe
            solve_grid            : Solves multiple column models for a grid of porosities and upwelling rates
                                        returns a 3-D array of activity ratios
    '''

    def __init__(self,df,dPdz = 0.32373,tol = 1.e-6,model=Disequilibrium,Da=0.1,method='Radau'):
        self.model = model
        self.Da = Da
        self.method = method

        # Check that the model selection and Da number are compatible:
        if self.Da is None or not isinstance(self.Da, float):
            raise ValueError('Da must be set to a floating point number.')

        self.df = df
        self.dPdz = dPdz
        self.tol = 1.e-6

        # Set the depth scale h:
        self.zmin = df['P'].min()/dPdz
        self.zmax = df['P'].max()/dPdz
        self.h = self.zmax - self.zmin

        # Lambda function to define scaled column height zprime:
        self.zp = lambda P: (self.zmax - P/dPdz)/self.h

        # Set interpolants for F, Kr, and P:
        self.F = pchip(self.zp(df['P']),df['F'])
        self.dFdz = self.F.derivative(nu=1)
        self.P = interp1d(self.zp(df['P']),df['P'],kind='cubic')

        # Set maximum degree of melting:
        self.Fmax = self.df['F'].max()

        # Set reference densities (assuming a mantle composition):
        self.rho_s = 3300.
        self.rho_f = 2800.

        # Set interpolation functions for partition coefficients for each element:
        self.Dnames = [col for col in df.columns if 'D' in col]
        self.D = []
        for name in self.Dnames:
            interpolated_func = interp1d(self.zp(df['P']), self.df[name], kind='cubic')
            self.D.append([interpolated_func])

        # Lambda function to retrieve partition coefficients at zprime = 0:
        self.get_D0 = lambda D, zp: np.array([D[i][0](zp) for i in range(len(D))])


    def solve_1D(self,D,C0,elements,z_eval = None):
        ''' Solves final concentrations for 1-D problem

        Usage:  z, Us, Uf = solve_1D(D,C0,z_eval)

        Parameters:
        -----------
            D      :  Bulk partition coefficients for each nuclide
            C0     :  Initial concentrations
            elements : List of elements to be analyzed
            z_eval :  Dimensionless column heights where solution is returned

        Returns:
        ---------
            z:   numpy array -- Coordinates where evaluated
            Cs:  numpy array -- Solid concentrations of each element
            Cf:  numpy array -- Liquid concentrations of each element
            Us:  numpy array -- Solid nuclide log concentration
            Uf:  numpy array -- Fluid nuclide log concentration
        '''

        # If z_eval is not set, use initial input values:
        if z_eval is None:
            z_eval = self.zp(self.df['P'])
        elif np.isscalar(z_eval):
            z_eval = np.array([z_eval])

        us = self.model(self.C0,self.D,self.F,self.dFdz,self.rho_f,self.rho_s,method=self.method,Da=self.Da)

        D0 = self.get_D0(D,z_eval)
        z, Us, Uf = us.solve(z_eval)

        # Calculate concentrations:
        Cs = [C0[i]*np.exp(Us[i]) for i in range(len(D0))]
        Cf = [C0[i]/D0[i]*np.exp(Uf[i]) for i in range(len(D0))]

        return z, Cs, Cf, Us, Uf


    def solve_all_1D(self,C0,elements,z_eval = None):
        ''' Solves 1-D column model for stable elements, and saves at all depths

        Parameters:
        -----------
            C0: float
                list of initial concentrations (ppm)
            elements: string
                list of element names extracted from input table headers

        Returns:
        --------
            df:  pandas dataframe containing physical and chemical parameters
        '''
        # Evaluate at input depths if not specified:
        if z_eval is None:
            z_eval = self.zp(self.df['P'])

        # Set initial concentrations for element i:
        self.elements = elements
        self.C0 = C0

        # Solve for element concentrations:
        z, Cs, Cf, Us, Uf = self.solve_1D(self.D,self.C0,self.elements,z_eval=z_eval)

        # Start building output dataframe:
        z = z_eval

        df = pd.DataFrame()
        df['P'] = self.P(z)
        df['depth'] = self.zmax - self.h*z
        df['F'] = self.F(z)

        for i,element in enumerate(elements):
            df['Cf({})'.format(element)] = Cf[i]
            df['Cs({})'.format(element)] = Cs[i]
            df['Uf({})'.format(element)] = Uf[i]
            df['Us({})'.format(element)] = Us[i]

        return df


### Utility plotting functions:
def plot_inputs(df,figsize=(8,6),show_grids=True):
    ''' pretty plots input data from pandas dataframe df

    Parameters:
    -----------
        df: pandas dataframe
            dataframe of input parameters ([P, K_r, Di's])
        figsize: tuple
            size of figure
        show_grids: bool
            if true, plot grids on all figures

    Returns:
    --------
        fig:  matplotlib figure object
        axes: tuple of matplotlib axes objects for replotting

    '''
    # columns = list(df)
    names = [col for col in df.columns if 'D' in col]

    fig, (ax1, ax2) = plt.subplots(1,2, sharey=True,figsize=figsize)
    ax1.plot(df['F']*100,df['P'])
    ax1.invert_yaxis()
    ax1.set_xlabel('F (%)')
    ax1.set_ylabel('Pressure (kbar)')
    xticks = np.linspace(0,max(df['F']),10)
    if show_grids:
        ax1.grid()
    ax1.set_xticks(xticks,minor=True)
    for s in names:
        ax2.semilogx(df[s],df['P'],label=s)
    ax2.set_xlabel(r'$D_i$')
    if show_grids:
        ax2.grid()
    ax2.legend(loc='best',bbox_to_anchor=(1.1,1))
    return fig, (ax1, ax2)

def plot_1Dcolumn(df,figsize=(8,6),show_grids=False):
    ''' pretty plots output data from dataframe of output

    Parameters:
    -----------
        df: pandas dataframe
            dataframe of model output
        figsize: tuple
            size of figure
        show_grids: bool
            if true, plot grids on all figures
            if false, just plot grid on activities plot

    Returns:
    --------
        fig:  matplotlib figure object
        axes: tuple of matplotlib axes objects for replotting
    '''
    Cf_all = [col for col in df.columns if 'Cf' in col]
    Cs_all = [col for col in df.columns if 'Cs' in col]
    elemlist = [e[3:] for e in Cf_all]
    elemlist = list(map(lambda i: i[: -1], elemlist))

    fig, (ax1, ax2, ax3) = plt.subplots(1,3, sharey=True,figsize=figsize)
    ax1.plot(df['F']*100,df['P'],'b',label='$F$')
    ax1.set_xlabel('Degree of melting (%)',color='b')
    ax1.set_ylabel('Pressure (kbar)')
    ax1.invert_yaxis()
    if show_grids:
        ax1.grid(color='b')

    for f in Cf_all:
        ax2.plot(df[f],df['P'],label=f)
    ax2.set_xlabel('$C_i^f$ (ppm)')
    ax2.grid()
    ax2.legend(elemlist,loc='best',bbox_to_anchor=(1.1,1))

    for s in Cs_all:
        ax3.plot(df[s],df['P'],label=s)
    ax3.set_xlabel('$C_i^s$ (ppm)')
    ax3.grid()
    ax3.legend(elemlist,loc='best',bbox_to_anchor=(1.1,1))
    return fig,(ax1,ax2,ax3)


#---------------------------------------
# Utility functions for calculating smoothed steps and ramps and the two-layer input generator

def _step(x, x_0, x_lambda, f):
    """ function to return smoothed step function between values f_min and f_max

        f_min + (f_max - f_min)*H((x-x0)/x_lambda)

        where H(x) is a smoothed Heaviside function

        Parameters:
        -----------

        x: array
            numpy array of locations x
        x_0: float
            location of step
        x_lambda: float
            smoothing lengthscale
        f:  array [f_min f_max]
            critical values in step function
            f_min value of step at x[0]
            f_max value of step at x[1]

    """
    return f[0] + (f[1] - f[0])*0.5*(np.tanh((x - x_0)/x_lambda) + 1)


def _ramp(x, x_0, x_lambda, f ):
    """ function to return smoothed broken ramp function  between values f_min, f_0 and f_max

        This function is just the integral of step(x, df1, df2, x_0, x_lambda) with values given by the slopes between

            df1 = (f_0 - f_min)/(x_0 - x_min)
            df2 = (f_max - f_0)/(x_max - x_0)

        Parameters:
        -----------

        x: array
            numpy array of locations x
        x_0: float
            location of step
        x_lambda: float
            smoothing lengthscale
        f: array [f_min f_0 f_max ]
            critical values in ramp function
            f_min value of ramp at x[0]
            f_0 value ramp at x_0
            f_max value of ramp at x[-1]

    """
    # set slopes of the two segments of the ramp far from the inflection point
    df1 = (f[1] - f[0])/(x_0 - x[0])
    df2 = (f[-1] - f[1])/(x[-1] - x_0)

    # calculate the scaled difference from the inflection point
    x_scaled = (x-x_0)/x_lambda

    # find all points > xscaled = -10
    indx = x_scaled >= -10
    istart =  np.where(indx)[0][0]

    # smoothed heaviside function
    H = (np.tanh(x_scaled) + 1.)/2.

    # initialized bottom part of ramp
    ramp = f[0] + df1*x
    # calculate integral of step function for x_scaled > -10
    intdf = df1*x[indx] + (df2 - df1)*(x[indx] - 0.5*x_lambda*np.log(H[indx]))
    # patch up ramp function
    ramp[indx]  = ramp[istart] - intdf[0] + intdf
    return ramp


def inputmodel(P, F, elements, D_lower, D_upper,  N=100, P_lambda = 1.,):
    """ create a pandas dataframe for the stableCalc sample two-layer model defined by a column that spans pressures from
        P_bottom to P_top with a possible change in layer properties at pressure P_boundary that is smoothed over a pressure range given by P_lambda
        Each layer can have its own bulk partition coefficient D.
        The degree of melting is given by a smoothed ramp function defined by three values of F, F(P_bottom), F(P_boundary) and F(P_top)
        Usage: df = stableCalc.inputmodel(P, F, elements, D_lower, D_upper, N, P_lambda)

        Parameters:
        -----------

        P: array [ P_bottom, P_boundary, P_top ]
            pressure bounds in the column (assumed kb)
            P[0] = P_bottom is the pressure at the bottom of the column
            P[1] = P_boundary is the pressure at the layer boundaries
            P[2] = P_top is the pressure at the top of the column
        F: array [ F_bottom, F_boundary, F_top]
            Degree of melting at the bottom, layer boundary and top of the column
        D_lower:  array [ D_i, D_j, D_k, etc. ]
            bulk partition coefficients for elements in the bottom layer
        D_upper:  array [ D_i, D_j, D_k, etc. ]
            bulk partition coefficients for elements in the upper layer
        elements: list of elements
        N: int
            Number of rows in dataframe
        P_lambda: float
            Pressure smoothing parameter.  Controls the width of smooth steps and ramps between layers (defaults to 1kb)

        Returns:
        --------

        df: pandas dataframe
            with columns P, F, Di

    """
    # basic check on array lengths
    assert(len(P) == 3)
    assert(len(F) == 3)
    assert(len(D_lower) == len(elements))
    assert(len(D_lower) == len(elements))

    # construct pressure array and reordered pressure array
    P_array = np.linspace(P[0], P[-1], N)
    x = P_array[0] - P_array
    x_0 = P_array[0] - P[1]

    Ds = list(zip(D_lower, D_upper))
    D_labels = ['D' + elem for elem in elements]

    _dict = OrderedDict()
    _dict['P'] = P_array
    _dict['F'] = _ramp(x, x_0, P_lambda, F)
    for i,D in enumerate(Ds):
        _dict[D_labels[i]] = _step(x, x_0, P_lambda, D)

    return pd.DataFrame.from_dict(_dict)


def smoothing(New_ID, P, F, Kr, D_lower, D_upper, D_labels, N=100, P_lambda = 1):
    """ Same as the inputmodel code above, but generates a dataframe with indices based on an existing input file,
        for smoothing of sharp discontinuities in that input file. The new dataframe is designed to update an old one.
        Takes the additional argument New_ID for indexing:

        New_ID: array [ID_bottom, ID_top]
            Index locations of the base and top of the section to be smoothed.

        Returns:
        --------

        df: pandas dataframe
            with columns New_ID, P, F, Kr (if already present), Di

    """
    # basic check on array lengths
    assert(len(P) == 3)
    assert(len(F) == 3)
    assert(len(New_ID) == 2)
    assert(len(D_lower) == len(D_upper))

    # construct pressure array and reordered pressure array
    P_array = np.linspace(P[0], P[-1], N)
    x = P_array[0] - P_array
    x_0 = P_array[0] - P[1]

    Ds = list(zip(D_lower, D_upper))

    _dict = OrderedDict()
    _dict['New_ID'] = np.arange(New_ID[0],New_ID[1]+1)
    _dict['P'] = P_array
    _dict['F'] = _ramp(x, x_0, P_lambda, F)
    if Kr != 'None':
        _dict['Kr'] = _step(x, x_0, P_lambda, Kr)
    for i,D in enumerate(Ds):
        _dict[D_labels[i]] = _step(x, x_0, P_lambda, D)

    return pd.DataFrame.from_dict(_dict)
