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

@author: richter bettina

Edited on 01.11.2022
@author: stephanie mayer
@author: richter bettina

"""


import numpy as np
import datetime
import pandas as pd

def read_profile(filename,timestamp=None,is3d=False,remove_soil=False,only_header=False):
    datafile=open(filename); content= datafile.readlines(); datafile.close()

    prof={}; is_data=False
    prof['info']={};prof['data']={}

    
    for line in content:
        
        if line.startswith('Altitude='): prof['info']['altitude']=float(line.strip().split('=')[1])
        elif line.startswith('Latitude='): prof['info']['latitude']=float(line.strip().split('=')[1])
        elif line.startswith('Longitude='): prof['info']['longitude']=float(line.strip().split('=')[1])
        elif line.startswith('SlopeAngle='): prof['info']['slopeAngle']=float(line.strip().split('=')[1])
        elif line.startswith('SlopeAzi='): prof['info']['slopeAzi']=float(line.strip().split('=')[1])
        elif line.startswith('StationName='): prof['info']['stationName']=line.strip().split('=')[1]

        elif line.startswith('[DATA]'):
            if only_header: break
            is_data=True        

        elif line.startswith('0500') and is_data: 
            ts = datetime.datetime.strptime( line.strip().split(',')[1], '%d.%m.%Y %H:%M:%S' )
            prof['data'][ts]={}
        elif line.startswith('0501') and is_data:
            height = np.array(line.strip().split(','))[2:].astype(float)
            if len(height) == 1 and height.item() == 0: continue
            else: prof['data'][ts]['height'] = height
#        elif line.startswith('0504') and is_data:
#            prof['data'][ts]['element_ID']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
        elif line.startswith('0502') and is_data:
            prof['data'][ts]['density']= np.array(line.strip().split(','))[2:].astype(float)
        elif line.startswith('0503') and is_data:
            prof['data'][ts]['temperature']= np.array(line.strip().split(','))[2:].astype(float)
        elif line.startswith('0506') and is_data:
            prof['data'][ts]['lwc']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
#        elif line.startswith('0508') and is_data:
#            prof['data'][ts]['dendricity']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
        elif line.startswith('0509') and is_data:
            prof['data'][ts]['sphericity']= np.array(line.strip().split(','))[2:].astype(float)
#        elif line.startswith('0510') and is_data:
#            prof['data'][ts]['coord_number']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
#        elif line.startswith('0511') and is_data:
#            prof['data'][ts]['bond_size']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
        elif line.startswith('0512') and is_data:
            prof['data'][ts]['grain_size']= np.array(line.strip().split(','))[2:].astype(float)
        elif line.startswith('0513') and is_data:
            prof['data'][ts]['grain_type']= np.array(line.strip().split(','))[2:-1].astype(float) #last entry is () and 772 is crust
#        elif line.startswith('0514') and is_data:        
#            prof['data'][ts]['sh_at_surface']= np.array(line.strip().split(','))[2:].astype(float)
#        elif line.startswith('0515') and is_data:
#            prof['data'][ts]['ice_content']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
#        elif line.startswith('0516') and is_data:
#            prof['data'][ts]['air_content']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
        elif line.startswith('0517') and is_data:
            prof['data'][ts]['stress']= np.array(line.strip().split(','))[2:].astype(float)
#        elif line.startswith('0518') and is_data:
#            prof['data'][ts]['viscosity']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
#        elif line.startswith('0520') and is_data:
#            prof['data'][ts]['temperature_gradient']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
        elif line.startswith('0523') and is_data:
            prof['data'][ts]['viscous_deformation_rate']= np.array(line.strip().split(','))[2:].astype(float)
#        elif line.startswith('0530') and is_data:
#            prof['data'][ts]['stab_indices']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
#        elif line.startswith('0531') and is_data:
#            prof['data'][ts]['stab_deformation_rate']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
        elif line.startswith('0532') and is_data:
            prof['data'][ts]['sn38']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
        elif line.startswith('0533') and is_data:
            prof['data'][ts]['sk38']= np.array(line.strip().split(','))[2:].astype(float)
        elif line.startswith('0534') and is_data:
            prof['data'][ts]['hand_hardness']= np.array(line.strip().split(','))[2:].astype(float)
#        elif line.startswith('0535') and is_data:
#            prof['data'][ts]['opt_equ_grain_size']= array([ float(x) for x in line.strip().split(',') ])[2:]
        elif line.startswith('0601') and is_data:
            prof['data'][ts]['shear_strength']= np.array(line.strip().split(','))[2:].astype(float)
#        elif line.startswith('0602') and is_data:
#            prof['data'][ts]['gs_difference']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
#        elif line.startswith('0603') and is_data:
#            prof['data'][ts]['hardness_difference']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
        elif line.startswith('0604') and is_data:
            prof['data'][ts]['ssi']= np.array([ float(x) for x in line.strip().split(',') ])[2:]
#        elif line.startswith('1501') and is_data:
#            prof['data'][ts]['height_nodes']= array([ float(x) for x in line.strip().split(',') ])[2:]
#        elif line.startswith('1532') and is_data:
#            prof['data'][ts]['sn38_nodes']= array([ float(x) for x in line.strip().split(',') ])[2:]
#        elif line.startswith('1533') and is_data:
#            prof['data'][ts]['sk38_nodes']= array([ float(x) for x in line.strip().split(',') ])[2:]
#        elif line.startswith('0540') and is_data:
#            prof['data'][ts]['date_of_birth']= np.array( line.strip().split(',') )[2:]
#        elif line.startswith('0607') and is_data:
#            prof['data'][ts]['accumulated_temperature_gradient']= array([ float(x) for x in line.strip().split(',') ])[2:]

    if is3d: getCoordinates(prof)
    if remove_soil:
        for profile in prof['data'].values():
            try:
                i_gr = np.where((profile['height']==0))[0].item()
                profile['height'] = profile['height'][i_gr+1:]
                profile['density'] = profile['density'][i_gr:]
#                profile['temperature'] = profile['temperature'][i_gr:]
#                profile['temperature_gradient'] = profile['temperature_gradient'][i_gr:]
                profile['stress'] = profile['stress'][i_gr:]
#                profile['ice_content'] = profile['ice_content'][i_gr:]
#                profile['air_content'] = profile['air_content'][i_gr:]
#                profile['viscosity'] = profile['viscosity'][i_gr:]
#                profile['element_ID'] = profile['element_ID'][i_gr:]
                profile['lwc'] = profile['lwc'][i_gr:]
#                profile['date_of_birth'] = profile['date_of_birth'][i_gr:] 
            except: continue

    for ts in prof['data'].keys():
        for var in prof['data'][ts].keys():
            data = prof['data'][ts][var]
            try: prof['data'][ts][var] = np.where((data==-999),np.nan,data)
            except: pass

    if timestamp: return prof['data'][timestamp]
    
    
    ts = sorted(prof['data'].keys())
    for dt in ts:
        try:
            pro=prof['data'][dt]
            rc=rc_flat(pro)
            prof['data'][dt]['critical_cut_length'] = rc
            
            pendepth = comp_pendepth(pro, prof['info']['slopeAngle'])
            prof['data'][dt]['penetration_depth'] = pendepth
        except:
            continue

    return prof

def getCoordinates(prof,yllcorner=186500,xllcorner=779000,gridsize=1,nx=600,ny=600,dem=None):
    ix,iy,name=prof['info']['stationName'].strip().split('_')
    ix=int(ix);iy=int(iy)
    prof['info']['ind_x'] = ix
    prof['info']['ind_y'] = ny-1-iy
    prof['info']['coord_x'] = xllcorner+ix*gridsize
    prof['info']['coord_y'] = yllcorner+iy*gridsize
    if dem:
        if dem[prof['info']['ind_y'],prof['info']['ind_x']] != prof['info']['altitude']: print( 'ACHTUNG')
        
#%%
    ##########################################################################################################################
    ########### Add improved and validated critical cut length for flat field according to Richter et al., 2019 (https://doi.org/10.5194/tc-13-3353-2019) ##############
    ########### Original formulation of Gaume et al., 2017 (https://doi.org/10.5194/tc-11-217-2017) resulted in negative values for slope simulations #################
    ########### In Mayer et al., 2022 (https://doi.org/10.5194/tc-16-4593-2022) flat rc values from slope simulation were ranked highly for computing probality of being unstable #########
    ########### Also add penetrations depth ##################################################################################
    ##########################################################################################################################

def rc_flat(pro):
    rho_ice = 917. #kg m-3
    gs_0 = 0.00125 #m
    thick = np.diff(np.concatenate((np.array([0]), pro['height'])))
    rho = pro['density']
    gs = pro['grain_size']*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 = pro['shear_strength']*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/rho_ice * gs/gs_0 )**b)*np.sqrt(2*tau_p*eprime*dsl_over_sigman)
    return rc_flat
        
def comp_pendepth(pro, slopeangle):
    top_crust = 0
    thick_crust = 0
    rho_Pk = 0
    dz_Pk = 1.e-12
    crust = False
    e_crust = -999
    
    layer_top = pro['height']
    ee = len(layer_top)-1
    thick = np.diff(np.concatenate((np.array([0]), layer_top)))
    rho = pro['density']
    HS = layer_top[-1]
    graintype = pro['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*100, (HS-top_crust)]) #NOTE  Pre-factor 0.8 introduced May 2006 by S. Bellaire , Pk = 34.6/rho_30
