#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

Created on Tue Apr 13 16:30:42 2021

@author: mayer stephanie

"""



import pandas as pd
import numpy as np
from scipy import stats as stats

################################################################  necessary functions   ###############################################################################

### to calculate critical crack length (Betti version) rcflat
###########################################################################################################
#####
#####   Name:       rc_flat
#####
#####   Purpose:     calculate critical crack length [m] (Richter 2019) for the flat field 
#####               given profile with arrays of layer tops, density rho, grain size gs, shear strength strength_s [kPa]   
#####
#####   Remarks:  rc is calculated for the flat, even if input profile is given at slope angle. 
#####             the formula is slope angle independent since SNOWPACK output refers to "gravitational layer height". 
#####               (--> normal stress in the flat is calculated exactly with this gravitational height)
#####               if formula was to be used for slopes, some cos, sin are needed!)
###########################################################################################################    
    
def rc_flat(layer_top, rho, gs, strength_s):
    rho_ice = 917. #kg m-3
    gs_0 = 0.00125 #m
    thick = np.diff(np.concatenate((np.array([0]), layer_top)))
    rho_wl = rho
    gs_wl = gs*0.001 #[m]
    rho_sl = np.append(np.flip(np.cumsum(rho[::-1]*thick[::-1])/np.cumsum(thick[::-1]))[1:len(rho)],np.nan)
    #rho_sl_old = np.array([ np.sum( rho[i+1:]* thick[i+1:])/np.sum(thick[i+1:]) if i<len(rho)-1 else np.nan for i in range(len(rho)) ] )
    tau_p = strength_s*1000. #[Pa]
    eprime = 5.07e9*(rho_sl/rho_ice)**5.13 / (1-0.2**2) #Eprime = E' = E/(1-nu**2) ; poisson ratio nu=0.2
    dsl_over_sigman = 1. / (9.81 * rho_sl) #D_sl/sigma_n = D_sl / (rho_sl*9.81*D_sl) = 1/(9.81*rho_sl)
    a = 4.6e-9
    b = -2.
    rc_flat = np.sqrt(a*( rho_wl/rho_ice * gs_wl/gs_0 )**b)*np.sqrt(2*tau_p*eprime*dsl_over_sigman)
    return rc_flat


###########################################################################################################
#####
#####   Name:       comp_features
#####
#####   Purpose:     given SNOWPACK profile in readProfile format, calculate all necessary features for every single layer (from bottom to top of profile)
#####               
#####   Remarks:  don't change order of features in output (the RF model needs exactly this order)
#####
###########################################################################################################
    
def comp_features(prof, slopeangle):
    
    strength_s = prof['shear_strength']  
    gs = prof['grain_size']
    rho = prof['density']
    layer_top = prof['height']
    nlayers = len(rho)
    
    # 1. viscous deformation rate
    viscdefrate = prof['viscous_deformation_rate']
    
    # 2. critical crack length
    rcflat = rc_flat(layer_top, rho, gs, strength_s)
    
    # 3. sphericity

    # 4. grainsize
  
    
    # 5. penetration depth
    
    # use simple trick: penetration depth = snowdepth - height of the top of the layer where sk38 is for the first time not equal to 6
    # does not always give same result as SP calculation 
#    depths = layer_top[-1]-layer_top
#    xx = np.where(prof['sk38']<6)[0] 
#    xx_min = np.argmin(depths[xx])
##    pen_depth_sk = depths[xx[xx_min]+1]/100
#    pen_depth_sk = depths[xx[xx_min]]/100
#    pen_depth_rep_sk = np.repeat(pen_depth_sk, nlayers) #repeat pen_depth to match output format (same value for every layer) 
    
    # calculation as in SP
    pen_depth_new = comp_pendepth(prof,slopeangle)
    pen_depth_rep = np.repeat(pen_depth_new, nlayers)
    
    # 6. mean of slab density divided by slab grain size <rho/gs>_{slab}
    thick = np.diff(np.concatenate((np.array([0]), layer_top)))
    #slab_rhogs_old = np.array([np.sum( rho[i+1:]* thick[i+1:]/gs[i+1:])/np.sum(thick[i+1:]) if i<len(rho)-1 else np.nan for i in range(len(rho))] )
    rhogs = rho*thick/gs
    slab_rhogs = np.append(np.flip(np.cumsum(rhogs[::-1])/np.cumsum(thick[::-1]))[1:len(rho)],np.nan)

    # put all features together into one dataframe. 
    d = {'viscdefrate': viscdefrate,
         'rcflat': rcflat,
         'sphericity': prof['sphericity'],
         'grainsize': gs,
         'penetrationdepth': pen_depth_rep,
         'slab_rhogs': slab_rhogs}  #
    
    features = pd.DataFrame(data = d)
    
    return features 


###########################################################################################################
#####
#####   Name:       comp_rf_probability
#####
#####   Purpose:     given SNOWPACK profile in readProfile format and RF model, calculate RF probability for each layer
#####               
#####   Remarks:  
#####
###########################################################################################################
    
def comp_rf_probability(features, model):
    # make sure features are in right order for model
    features = features[['viscdefrate', 'rcflat', 'sphericity', 'grainsize', 'penetrationdepth','slab_rhogs']]
    P_unstable = np.repeat(np.nan,len(features))    
    
    # for i, row in features.iterrows():
        # if i < len(features)-1: #all layers except the upper one
            # P_unstable[i] = model.predict_proba(features.loc[[i]])[:,0]
    if len(P_unstable)> 1:# more than one element is needed to compute P_unstable, wl + slab
        P_unstable[:-1] = model.predict_proba(features.iloc[:-1])[:,0] 
    else:
        if len(P_unstable) ==1:
            P_unstable = [np.nan]
        else:
            P_unstable = []
        

    return P_unstable

###########################################################################################################
#####
#####   Name:       comp_prof_dataframe
#####
#####   Purpose:     given SNOWPACK profile in readProfile format and RF model, calculate RF probability for each layer
#####               and save relevant properties in dataframe
#####   Remarks:  
#####
###########################################################################################################

def create_RFprof(prof, model,slopeangle=0):
    #get features for RF model
    features = comp_features(prof, slopeangle)
    #get P_unstable
    df = features
    df['P_unstable'] = comp_rf_probability(features, model)
    #get some additional features for plotting
    df['density'] = prof['density']
    df['height'] = prof['height']
    df['hand_hardness'] = prof['hand_hardness']
    df['grain_type'] = prof['grain_type']    
    return df
        



###########################################################################################################
#####
#####   Name:       comp_pendepth
#####
#####   Purpose:     given SNOWPACK profile in readProfile format, calculate skier penetration depth
#####               
#####   Remarks: might slightly differ from SNOWPACK source code, but is correct with regard to publications Jamieson Johnston (1998)
#####               and Bellaire (2006)
###########################################################################################################
def comp_pendepth(prof, slopeangle):
    top_crust = 0
    thick_crust = 0
    rho_Pk = 0
    dz_Pk = 1.e-12
    crust = False
    e_crust = -999
    
    layer_top = prof['height']
    ee = len(layer_top)-1
    thick = np.diff(np.concatenate((np.array([0]), layer_top)))
    rho = prof['density']
    HS = layer_top[-1]
    graintype = prof['grain_type']  
    min_thick_crust = 3 #cm
    
    while (ee >= 0) & ((HS-layer_top[ee])<30):
        
        rho_Pk = rho_Pk + rho[ee]*thick[ee]
        dz_Pk = dz_Pk + thick[ee]
        
        if crust == False:
        ##Test for strong mf-crusts MFcr.
        ## Look for the first (from top) with thickness perp to slope > 3cm
            if (graintype[ee] == 772) & (rho[ee] >500.): ## very high density threshold, but implemented as this in SP
                if e_crust == -999:
                   e_crust = ee
                   top_crust = layer_top[ee]
                   thick_crust = thick_crust + thick[ee]
                elif (e_crust - ee) <2:
                   thick_crust = thick_crust + thick[ee]
                   e_crust = ee
            elif e_crust > 0:
               if thick_crust*np.cos(np.deg2rad(slopeangle)) > min_thick_crust:
                   crust = True
               else:
                   e_crust = -999
                   top_crust = 0
                   thick_crust = 0

        ee = ee-1
                         

    
    rho_Pk = rho_Pk/dz_Pk        #average density of the upper 30 cm slab
    return np.min([0.8*43.3/rho_Pk, (HS-top_crust)/100.]) #NOTE  Pre-factor 0.8 introduced May 2006 by S. Bellaire , Pk = 34.6/rho_30 #/100 to get upper crust in meters
#original regression by Jamieson Johnston (1998)
