'''@author: Junkai Ge (gejunkai@mail.sdu.edu.cn) Shandong university'''
#generate gprMax in file for 3D model             
#GPrMax needs to be pre installed

import numpy as np
from utils.layers_generator import generate_layers_1d,assign_permittivity,assign_permittivity_with_smooth_transition,assign_integer_values
import pandas as pd
import os
import matplotlib.pyplot as plt
from config import Forward_Model_Config as cfg
from config import Path_Config as pcfg

# #Initial modeling parameters
model_num=cfg.model_num                         #Number of models
depth=cfg.depth                                 #Model depth unit:metre
air_depth=cfg.air_depth                         #Air layer thickness, unit:metre
grid_length=cfg.grid_length                     #grid length
surface_width=cfg.surface_width                 #surface width:metre
layers_range=cfg.layers_range                   #layers range
smooth_cell=cfg.smooth_cell                     #Model smoothing parameters
min_layer_thickness=cfg.min_layer_thickness     #<=grid length
first_layer_minlenth=cfg.first_layer_minlenth   #first layer minlenth
first_layer_maxlenth=cfg.first_layer_maxlenth   #first layer maxlenth
permittivity_range=cfg.permittivity_range       #the range of permittivity
first_layer_eps_range=cfg.first_layer_eps_range #the range of permittivity of first layer
Time=cfg.Time                                   #Collection time window (within the medium)
static_time=cfg.static_time                     #Collection time window (in the air)
frequency=cfg.frequency                         #Antenna frequency
Twindows=cfg.Twindows                           #total time
root = cfg.root                                 #current path


#Calculate the number of grids
grid_x=int(surface_width/grid_length)
grid_y=int(surface_width/grid_length)
grid_z=int(depth/grid_length)
graid_layer_min=int(min_layer_thickness/grid_length)
first_layer_minsize=int(first_layer_minlenth/grid_length)
first_layer_maxsize=int(first_layer_maxlenth/grid_length)

#Modeling
for index in range(model_num):

    layers_num=np.random.randint(layers_range[0],layers_range[1])#Number of generated layers
    layer=generate_layers_1d(grid_z, layers_num, graid_layer_min,smooth_cell,first_layer_minsize,first_layer_maxsize)#Generated layer thickness and arrangement
    eps_model=assign_permittivity_with_smooth_transition(layer, layers_num, first_layer_eps_range, permittivity_range, smooth_cell)#Generate random permittivity and smooth layer
    layer_id=assign_integer_values(eps_model)#assign id

    #Initialization of velocity field, calculation of velocity field
    vel_model=np.zeros(grid_z)
    for element in range(grid_z):
        vel_model[element]=0.3/np.sqrt(eps_model[element])

    z = np.zeros(grid_z)+grid_length
    v = vel_model
    t = np.cumsum(2*grid_length / v)
    Time_time = Time*1e9  
    new_time = np.arange(0, Time_time, Time_time/grid_z)  
    v_interpolated = np.interp(new_time, t, v) 

    # Cut the first grid-z points, and if there are not enough grid-z points, fill them with the last value
    if len(v_interpolated) < grid_z:
        v_interpolated = np.pad(v_interpolated, (0, grid_z - len(v_interpolated)), 'edge')
    else:
        v_interpolated = v_interpolated[:grid_z] 

    eps_interpolated=(0.3/v_interpolated)**2

    # Draw a distribution diagram of permittivity/velocity over time
    plt.figure()
    plt.plot(new_time[:grid_z], eps_interpolated)
    plt.xlabel('Time (s)')
    plt.ylabel('Permittivity')
    plt.title('Permittivity over time')
    plt.grid(True)
    plt.savefig('./dataset/eps_model_in_time/' + str(index + 1).zfill(5) + '.png')

    #permittivity model over grid
    plt.figure()
    plt.plot(eps_model, 'o-', markersize=5)
    plt.title("permittivity model over grid")
    plt.ylim(0, 30)
    plt.savefig('./dataset/eps_model_in_depth/' + str(index + 1).zfill(5) + '.png')

    #One dimensional dielectric constant over storage time
    output_file='./dataset/eps_label_in_time/'+str(index+1).zfill(5)+'.csv'
    model_txtDF = pd.DataFrame(eps_interpolated)
    model_txtDF.to_csv(output_file,index=False)

    #Write 'in' files
    fl = open(pcfg.in_data_dir+str(index+1).zfill(5)+'_'+str(int(eps_model[1])).zfill(2)+'.in','w')
    fl.write('#title: '+str(index+1).zfill(5)+'\n')
    fl.write('#domain: '+str(surface_width)+' '+str(surface_width)+' '+str(depth+air_depth)+'\n')
    fl.write('#dx_dy_dz: '+str(grid_length)+' '+str(grid_length)+' '+str(grid_length)+'\n')
    fl.write('#pml_cells: 20 20 20 20 20 20\n')
    fl.write('#time_window: '+str(Twindows)+'\n')
    fl.write('#waveform: ricker 1 '+str(frequency)+' my_ricker\n')
    fl.write('#hertzian_dipole: z '+str(surface_width/2)+' '+str(surface_width/2)+' '+str(depth)+' my_ricker\n')
    fl.write('#rx: '+str(surface_width/2)+' '+str(surface_width/2)+' '+str(depth)+'\n')
    fl.write('#material: 1 0 1 0 kongqi\n')
    fl.write('#box: 0 0 0 '+str(surface_width)+' '+str(surface_width)+' '+str(depth+air_depth)+' kongqi n \n')
    
    #Assign material properties and modeling
    current_eps = eps_model[0]
    end_z = depth
    material_idx=0
    for i in range(1, grid_z):
        if eps_model[i] != current_eps:
            material_idx=material_idx+1
            start_z = depth - i*grid_length
            fl.write('#material: '+str(eps_model[i-1])+' 0 1 0 '+'material'+str(material_idx).zfill(5)+'\n')
            fl.write('#box: 0 0 '+ str(start_z)+' '+str(surface_width)+' '+str(surface_width)+' '+str(end_z)+' '+'material'+str(material_idx).zfill(5)+' n \n')
            current_eps = eps_model[i]
            end_z = start_z 
        
    material_idx=material_idx+1
    fl.write('#material: '+str(eps_model[-1])+' 0 1 0 '+'material'+str(material_idx).zfill(5)+'\n')
    fl.write('#box: 0 0 0 '+str(surface_width)+' '+str(surface_width)+' '+str(end_z)+' '+'material'+str(material_idx).zfill(5)+' n \n')
    fl.close()