#!/localstorage/peruzzetto/anaconda3/bin/python
# -*- coding: utf-8 -*-
"""
Created on Thu Feb  6 14:24:10 2020

@author: peruzzetto

MIT License, Copyright (c) 2024 Marc PERUZZETTO
"""

import numpy as np

# No data value
NDV=-10


def make_topo():
    # Grid extent
    xmin=0
    xmax=2
    ymin=0
    ymax=1
    
    #Number of cells
    nx=150
    ny=75
    
    #Slope angle
    theta=20
    
    x=np.linspace(xmin,xmax,nx)
    y=np.linspace(ymin,ymax,ny)
    
    xmesh,ymesh=np.meshgrid(x,y,indexing='ij')
    
    z=-np.tan(np.deg2rad(theta))*xmesh
    z=z-np.min(z)
    
    np.savetxt('z.d',np.reshape(z,(nx*ny),order='F'),fmt='%.10e')
    
    return xmesh, ymesh, z

def paraboloid(xmesh, ymesh, xc=0.5,yc=0.5,xw=0.15,yw=0.15,hmax=0.15):
    """
    Create a paraboloid mass

    Parameters
    ----------
    xc : float, optional
        x-position of mass center. The default is 0.5.
    yc : float, optional
        y-position of mass center. The default is 0.5.
    xw : float, optional
        x-width of mass. The default is 0.15.
    yw : float, optional
        y-width of mass. The default is 0.15.
    hmax : float, optional
        maximum thickness of mass, at position (xc, yc). The default is 0.15.

    Returns
    -------
    h : TYPE
        DESCRIPTION.

    """
    h=hmax*(1-(xmesh-xc)**2/xw**2-(ymesh-yc)**2/yw**2)
    return h

def make_init_mass(xmesh, ymesh, hmax=0.25):
    h=paraboloid(xmesh, ymesh, xc=0.5,yc=0.5,xw=0.15,yw=0.15,hmax=hmax)
    ind=h<0
    h[ind]=0
    np.savetxt('m.d',np.reshape(h,(xmesh.shape[0]*xmesh.shape[1]),
                                order='F'),fmt='%.10e')
    return ~ind
    
def make_init_vel(xmesh, ymesh, u0=0, ut0=0, ind=None):
    nx = xmesh.shape[0]
    ny = xmesh.shape[1]

    uinit=np.zeros((nx,ny))
    utinit=np.zeros((nx,ny))
    
    if ind is not None:
        uinit[ind]=u0
        utinit[ind]=ut0
        
    np.savetxt('u.d',np.reshape(uinit,(nx*ny),order='F'),fmt='%.10e')
    np.savetxt('ut.d',np.reshape(utinit,(nx*ny),order='F'),fmt='%.10e')

def make_flux(xmesh, ymesh, flux=0.02, u0=0, ut0=0, flux_mask=True,
              xc=0.5, yc=0.5):

    nx = xmesh.shape[0]
    ny = xmesh.shape[1]
    
    rho_source=paraboloid(xmesh, ymesh, xc=xc, yc=yc, xw=0.15, yw=0.15,
                          hmax=flux)
    ind=rho_source<=0
    rho_source[ind]=NDV
    if flux_mask:
        rho_source[~ind]=1

    usource=np.zeros((nx,ny))
    utsource=np.zeros((nx,ny))
    usource[ind]=NDV
    utsource[ind]=NDV
    usource[~ind]=u0
    utsource[~ind]=ut0
    
    np.savetxt('u_source.d',np.reshape(usource,(nx*ny),order='F'),
               fmt='%.10e')
    np.savetxt('ut_source.d',np.reshape(utsource,(nx*ny),order='F'),
               fmt='%.10e')
    np.savetxt('rho_source.d',np.reshape(rho_source,(nx*ny),order='F'),
               fmt='%.10e')

def make_params(xmesh, ymesh, flux=0):
    
    nx = xmesh.shape[0]
    ny = xmesh.shape[1]
    params = dict()
    params['nx']=nx
    params['ny']=ny
    params['per']=xmesh[-1, 0]-xmesh[0, 0]
    params['pery']=ymesh[0, -1]-ymesh[0, 0]
    params['tmax']=1
    params['nbim']=10
    
    params['initz']=0
    params['file_z_init']='z.d'
    params['ipr']=0
    
    params['icomp']=1
    params['delta1']=30
    
    return params

def write_params(params, file='params.txt'):

    with open('params.txt','w') as file_params:
        for name in params:
            val=params[name]
            if type(val)==int or type(val)==np.int64:
                file_params.write('{:s} {:d}\n'.format(name,val))
            if type(val)==float or type(val)==np.float64:
                file_params.write('{:s} {:.4G}\n'.format(name,val))
            if type(val)==str:
                file_params.write('{:s} {:s}\n'.format(name,val))

if __name__ == '__main__':
    
    xmesh, ymesh, z = make_topo()
    ind_mass = make_init_mass(xmesh, ymesh)
    
    # Uncomment section to run it
    
    # %% Simple simulation with mass on inclined plane
    # params = make_params(xmesh, ymesh)
    # write_params(params)
    
    # %% Add initial velocity from file
    # make_init_vel(xmesh, ymesh, u0=-0.5, ut0=-0.5, ind=ind_mass)
    # params = make_params(xmesh, ymesh)
    # params['initu']=0
    # params['initut']=0
    # write_params(params)
    
    # %% Use input flux from mass file, mass value become local flux
    # make_init_mass(xmesh, ymesh, hmax=4e-4)
    # params = make_params(xmesh, ymesh)
    # params['ipr']=2
    # params['zv']=0
    # params['i_rho_source']=2
    # params['file_rho_source']='m.d'
    # params['tmax'] = 2
    # params['nbim'] = 20
    # write_params(params)
    
    # %% Flux location given by a mask, total flux given as parameter
    # make_flux(xmesh, ymesh)
    # params = make_params(xmesh, ymesh)
    # params['ndv']=NDV
    # params['ipr']=2
    # params['zv']=0
    # params['i_rho_source']=2
    # params['file_rho_source']='rho_source.d'
    # params['source_flux']=0.2
    # params['tmax'] = 2
    # params['nbim'] = 20
    # write_params(params)
    
    # %% Flux location given by a mask, total flux and velocity given as   
    # # parameters
    # make_flux(xmesh, ymesh)
    # params = make_params(xmesh, ymesh)
    # params['ndv']=NDV
    # params['ipr']=2
    # params['zv']=0
    # params['i_rho_source']=2
    # params['file_rho_source']='rho_source.d'
    # params['source_flux']=0.15
    # params['tmax'] = 1
    # params['nbim'] = 10
    # params['i_vel_source'] = 2 #0: nothing, 1: given by file, 2: uniform
    # params['file_u_source']='rho_source.d' # Mask where velocity is imposed
    # params['vel_source_norm'] = 0.1 # Velocity norm (for i_vel_source=2)
    # params['vel_source_dir'] = 60 # Direction of velocity (NDV: steepest slope)
    # write_params(params)
    
    # %% Use initial mass and input flux between two given times
    make_flux(xmesh, ymesh, xc=0.75, yc=0.5)
    params = make_params(xmesh, ymesh)
    params['ndv']=NDV
    params['ipr']=0
    params['i_rho_source']=2
    params['file_rho_source']='rho_source.d'
    params['source_flux']=0.05 #-0.001 for "hole"
    params['tmax'] = 1.5
    params['nbim'] = 15
    params['t_source_start']=0.6
    params['t_source_end']=1.5
    write_params(params)
    
