# -*- coding: utf-8 -*-
# authors: emoisan (enora.moisan@ens-lyon.fr)
#          achatain (audrey.chatain@latmos.ipsl.fr)
# created on: 2023-01-25
# last modification: 2023-07-25

import numpy as np
import matplotlib.pyplot as plt

################################################################################
################################################################################
## Creates topography map, according to the namelist parameters.
## Needs to be adapted to read lake maps for 3D.
################################################################################
################################################################################

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
## Flags
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
plot_flag = 1 # creates a figure to visualize input_topo

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
## Functions definition
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 

def give_value(string): 
    '''
    Check for the data type in the string (int, float, boolean), 
    and returns the corresponding value.

    '''
    try:
        int(string)
        return int(string)
    except ValueError:
        try:
            float(string)
            return float(string)
        except ValueError:
            try:
                bool(string)
                return bool(string)
            except ValueError:
                print("Error: unknown value type." + string)

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
## Read namelist.input
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 

namelist_file_name = input("Enter name of namelist.input file. ")
namelist_file = open(namelist_file_name, 'r')
namelist = []
namelist = namelist_file.readlines()

# Read &domains section of the namelist
domains_dict = {}
i_domains_start = namelist.index(' &domains\n')
i_domains_end = namelist.index(' /\n', i_domains_start) # first ' /\n' after start of &domains
for i in range(i_domains_start+1, i_domains_end): # each line has the structure: ' name      = value,\n'
    i_eq = namelist[i].index('=') # get '=' index
    i_coma = namelist[i].index(',') # get ',' index
    key = namelist[i][:i_eq].strip() # get namelist variable name
    value = give_value(namelist[i][i_eq+1:i_coma]) # get namelist variable value
    domains_dict[key] = value # store variable name and value in the dictionary

# Read &dynamics section of the namelist
dynamics_dict = {}
i_dynamics_start = namelist.index(' &dynamics\n')
i_dynamics_end = namelist.index(' /\n', i_dynamics_start) # 1st ' /\n' after start of &dynamics
for i in range(i_dynamics_start+1,i_dynamics_end): # each line has the structure: ' name      = value,\n'
    i_eq = namelist[i].index('=') # get '=' index
    i_coma = namelist[i].index(',') # get ',' index
    key = namelist[i][:i_eq].strip() # get namelist variable name
    value = give_value(namelist[i][i_eq+1:i_coma]) # get namelist variable value
    dynamics_dict[key] = value # store variable name and value in the dictionary

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
## Create topography array
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

ishole = int(input("Type 1 if hole, 0 if rampart: "))

i_max = 2000
j_max = 1

## Get values from the namelist
s_we = domains_dict['s_we']
e_we = domains_dict['e_we']
s_sn = domains_dict['s_sn']
e_sn = domains_dict['e_sn']
dx   = domains_dict['dx']/1000 # in km
dy   = domains_dict['dy']/1000 # in km
lake_half_size = dynamics_dict['lake_half_size'] 
rampart_width  = dynamics_dict['rampart_width']
rampart_height = dynamics_dict['rampart_height']

dim_x = e_we-s_we
dim_y = e_sn-s_sn

lake_center = (e_we-1)*dx//2
x = np.arange(dx/2,dim_x*dx,dx) - lake_center # x coordinate in km

## Rampart topography
hp = rampart_height # 300 m
lm  = lake_half_size
rw  = rampart_width

## Assign topography to an array
topo = np.zeros((dim_y,dim_x)) # topo[y,x]

## for a rectangle shape (WRF does not like it)
# icm = e_we//2 # index of the center of the model
# for i_python in range(dim_x):
#     i_fortran = i_python +1 # python starts at 0, fortran starts at 1
#     if icm-lm-rw < i_fortran <= icm-lm or icm+lm < i_fortran <= icm+lm+rw:
#         topo[:,i_python] = hp # if rampart, topo=rampart_height. Else, 0.
        
## for a exp*cos**2 shape: double dump
# if ishole == 0:
#     Vx = lm      #km  valley floor half width
#     Sx = rw/2    #km  hill half width
#     Px = 0       #km  hill top width
#     a  = Sx/2.5  # expo smooth
#     a2 = a*10
#     for i in range(dim_x):
#         if abs(x[i]) <= Vx:
#             topo[:,i] = 0
#         elif abs(x[i]) < Vx+Sx:
#             topo[:,i] = (np.cos(np.pi*(abs(x[i])-Vx-Sx)/a2))**2*np.exp(-((abs(x[i])-Vx-Sx)/a)**2)
#         elif abs(x[i]) <= Vx+Sx+Px:
#             topo[:,i] = 1
#         elif abs(x[i]) < Vx+2*Sx+Px:
#             topo[:,i] = (np.cos(np.pi*(abs(x[i])-Vx-Sx-Px)/a2))**2*np.exp(-((abs(x[i])-Vx-Sx-Px)/a)**2)
#         else:
#             topo[:,i] = 0
#     topo = hp*topo

## for a exp*cos**2 hole
# if ishole == 1:
#     Vx = lm      #km  valley floor half width
#     Sx = rw/2    #km  hill half width
#     a  = Sx/2.5  # expo smooth
#     a2 = a*10
#     for i in range(dim_x):
#         if abs(x[i]) <= Vx:
#             topo[:,i] = 0
#         elif abs(x[i]) < Vx+Sx:
#             topo[:,i] = (np.cos(np.pi*(abs(x[i])-Vx-Sx)/a2))**2*np.exp(-((abs(x[i])-Vx-Sx)/a)**2)
#         else:
#             topo[:,i] = 1
#     topo = hp*topo
    
    
## for a double dump with smootherstep
if ishole == 0:
    ## Topo parameters
    hp = 300  #m   height
    Vx = lm   #km  valley floor half width
    Sx = rw/2 #km  hill step length
    Px = 0    #km  hill top width
    ## Smootherstep function
    def smootherstep(x):
        y=6*x**5-15*x**4+10*x**3
        return y
    ## Assign topography to an array
    topo = np.zeros((dim_y,dim_x)) # topo[y,x]
    for i in range(dim_x):
        if abs(x[i]) <= Vx:
            topo[:,i] = 0
        elif abs(x[i]) < Vx+Sx:
            topo[:,i] = smootherstep((abs(x[i])-Vx)/Sx)
        elif abs(x[i]) <= Vx+Sx+Px:
            topo[:,i] = 1
        elif abs(x[i]) < Vx+2*Sx+Px:
            topo[:,i] = 1-smootherstep((abs(x[i])-Vx-Sx-Px)/Sx)
        else:
            topo[:,i] = 0
    topo = hp*topo
    
    
## for a hole with smootherstep
if ishole == 1:
    ## Topo parameters
    hp = 300  #m   height
    Vx = lm   #km  valley floor half width
    Sx = rw #/2 #km  hill step length
    Px = 0    #km  hill top width
    ## Smootherstep function
    def smootherstep(x):
        y=6*x**5-15*x**4+10*x**3
        return y
    ## Assign topography to an array
    topo = np.zeros((dim_y,dim_x)) # topo[y,x]
    for i in range(dim_x):
        if abs(x[i]) <= Vx:
            topo[:,i] = 0
        elif abs(x[i]) < Vx+Sx:
            topo[:,i] = smootherstep((abs(x[i])-Vx)/Sx)
        else:
            topo[:,i] = 1
    topo = hp*topo


#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
## Write input_topo file
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 

file_input_topo = open('input_topo', 'w')
for j in range(dim_y-1,-1,-1): # from dim_y-1 to 0
    for i in range(dim_x): # from 0 to dim_x-1
        file_input_topo.write(f'{topo[j,i]:.2f}' + 4*' ')
    for i in range(dim_x,i_max): # from dim_x to i_max-1
        file_input_topo.write('0.00' + 4*' ')
    file_input_topo.write('\n')
file_input_topo.close()

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
## Plot
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
if plot_flag:
    fig = plt.figure(figsize=(8,6))  #figsize=(8,2)
    axes = fig.add_axes([0.09,0.0,0.9,0.9])
#   axes.axis('equal')
    aspect_ratio = 0.5 #1e-2
    vertical_exag = aspect_ratio/1e-3 # vertical exaggeration: aspect_ratio/1e-3 (10 here)
    axes.set_aspect(aspect_ratio)
    plt.plot(x, topo[0,:],'-+')
    plt.xlim([-100,100])
#   plt.grid()
    axes.spines['top'].set_visible(False)
    axes.spines['right'].set_visible(False)
    # axes.set_yticks([hp])
    plt.xlabel('Distance [km]')
    plt.ylabel('Topography [m]')
    axes.set_title('Topography at the center of the model (vertical exaggeration: %i)'%(vertical_exag),\
                   y=1, pad=50)
    plt.savefig('topo.jpg')
    plt.close()