import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), 'modules')))

import modules.util_stackup_reader as stackup_reader
import modules.util_gds_reader as gds_reader
import modules.util_utilities as utilities
import modules.util_simulation_setup as simulation_setup
import modules.util_meshlines as util_meshlines

from openEMS import openEMS
import numpy as np
import matplotlib.pyplot as plt


# Model comments
# 
# Example how layout can be created from code, with or without loading other layout from GDSII file
# See code lines with  allpolygons.add_rectangle() and allpolygons.add_polygon()


# ======================== workflow settings ================================

# preview model/mesh only?
# postprocess existing data without re-running simulation?
preview_only = False
postprocess_only = False

# ===================== input files and path settings =======================

# gds_filename = "line_simple_viaport.gds"   # geometries
XML_filename = "SG13G2_nosub.xml"               # stackup

# preprocess GDSII for safe handling of cutouts/holes?
preprocess_gds = False

# get path for this simulation file
script_path = utilities.get_script_path(__file__)

# use script filename as model basename
model_basename = utilities.get_basename(__file__)

# set and create directory for simulation output
sim_path = utilities.create_sim_path (script_path,model_basename)
print('Simulation data directory: ', sim_path)

# change current path to model script path
os.chdir(os.path.dirname(os.path.abspath(__file__)))

# ======================== simulation settings ================================

unit   = 1e-6  # geometry is in microns
margin = 50    # distance in microns from GDSII geometry boundary to simulation boundary 

fstart =  0e9
fstop  = 110e9
numfreq = 401

refined_cellsize = 0.5  # mesh cell size in conductor region

# choices for boundary: 
# 'PEC' : perfect electric conductor (default)
# 'PMC' : perfect magnetic conductor, useful for symmetries
# 'MUR' : simple MUR absorbing boundary conditions
# 'PML_8' : PML absorbing boundary conditions
Boundaries = ['PEC', 'PEC', 'PEC', 'PEC', 'PEC', 'PEC']

cells_per_wavelength = 20   # how many mesh cells per wavelength, must be 10 or more
energy_limit = -40          # end criteria for residual energy (dB)

# ports from GDSII Data, polygon geometry from specified special layer
# note that for multiport simulation, excitations are switched on/off in simulation_setup.createSimulation below

simulation_ports = simulation_setup.all_simulation_ports()
# instead of in-plane port specified with target_layername, we here use via port specified with from_layername and to_layername
simulation_ports.add_port(simulation_setup.simulation_port(portnumber=1, voltage=1, port_Z0=50, source_layernum=201, from_layername='Metal1', to_layername='TopMetal2', direction='z'))
simulation_ports.add_port(simulation_setup.simulation_port(portnumber=2, voltage=1, port_Z0=50, source_layernum=202, from_layername='Metal1', to_layername='TopMetal2', direction='z'))

# ======================== simulation ================================

# get technology stackup data
materials_list, dielectrics_list, metals_list = stackup_reader.read_substrate (XML_filename)
# get list of layers from technology
layernumbers = metals_list.getlayernumbers()
layernumbers.extend(simulation_ports.portlayers)

# Here, we do NOT load a GDSII file, and create an empty all_polygons_list() instead
# allpolygons = gds_reader.read_gds(gds_filename, layernumbers, purposelist=[0], metals_list=metals_list, preprocess=preprocess_gds)
allpolygons = gds_reader.all_polygons_list()

# Add rectangles manually
# microstrip line
allpolygons.add_rectangle(x1=0, y1=-1.5, x2=100, y2=1.5, layernum=metals_list.getbylayername('TopMetal2').layernum)
# port rectangles, mapped to via ports defined above by their layer number. Parameter is_port=True leads to mesh refinement for this polygon.
allpolygons.add_rectangle(x1=0,  y1=-1.5, x2=1,   y2=1.5, layernum=201, is_port=True)
allpolygons.add_rectangle(x1=99, y1=-1.5, x2=100, y2=1.5, layernum=202, is_port=True)
# ground plane on Metal1 
allpolygons.add_rectangle(x1=-20,  y1=-20, x2=120, y2=20, layernum=metals_list.getbylayername('Metal1').layernum)
# test adding polygon 
# allpolygons.add_polygon( xy=[[-20,-20],[-20,20],[120,20],[120,-20]], layernum=metals_list.getbylayername('Metal1').layernum)


# calculate maximum cellsize from wavelength in dielectric
wavelength_air = 3e8/fstop / unit
max_cellsize = (wavelength_air)/(np.sqrt(materials_list.eps_max)*cells_per_wavelength) 

# define excitation and stop criteria and boundaries
FDTD = openEMS(EndCriteria=np.exp(energy_limit/10 * np.log(10)))
FDTD.SetGaussExcite( (fstart+fstop)/2, (fstop-fstart)/2 )
FDTD.SetBoundaryCond( Boundaries )


########### create model, run and post-process ###########

# Create simulation for port 1 and 2 excitation, return value is data path for that excitation
excite_ports = [1]  # list of ports that are excited for this simulation run
FDTD = simulation_setup.setupSimulation (excite_ports, 
                                         simulation_ports, 
                                         FDTD, 
                                         materials_list, 
                                         dielectrics_list, 
                                         metals_list, 
                                         allpolygons, 
                                         max_cellsize, 
                                         refined_cellsize, 
                                         margin, 
                                         unit, 
                                         xy_mesh_function=util_meshlines.create_xy_mesh_from_polygons)


# create dump box in frequency domain, type 12 for electric current frequency-domain dump
CSX = FDTD.GetCSX()

dump_frequency = 100e9
Jf = CSX.AddDump('Jf', file_type=0, dump_type=12, frequency=dump_frequency, sub_sampling=[1,1,1])

TopMetal2 = metals_list.getbylayername('TopMetal2')
start = [0, -5, TopMetal2.zmin]
stop  = [100, 5, TopMetal2.zmax]
Jf.AddBox(start, stop)


sub1_data_path = simulation_setup.runSimulation (excite_ports, FDTD, sim_path, model_basename, preview_only, postprocess_only)


########## evaluation of results with composite GSG ports ###########

if preview_only==False:

    # define dB function for S-parameters
    def dB(value):
        return 20.0*np.log10(np.abs(value))        

    # define phase function for S-parameters
    def phase(value):
        return np.angle(value, deg=True) 


    f = np.linspace(fstart,fstop,numfreq)

    # get results, CSX port definition is read from simulation ports object
    s11 = utilities.calculate_Sij (1, 1, f, sim_path, simulation_ports)
    s21 = utilities.calculate_Sij (2, 1, f, sim_path, simulation_ports)

    # S12, S22 is NOT available because we have NOT simulated port2 excitation
    # fake it by assuming symmetry
    s22 = s11
    s12 = s21

    # write Touchstone S2P file
    s2p_name = os.path.join(sim_path, model_basename + '.s2p')
    utilities.write_snp (np.array([[s11, s21],[s12,s22]]),f, s2p_name)

    fig, axis = plt.subplots(num='Return Loss', tight_layout=True)
    axis.plot(f/1e9, dB(s11), 'k-',  linewidth=2, label='S11 (dB)')
    axis.grid()
    axis.set_xmargin(0)
    axis.set_xlabel('Frequency (GHz)')
    axis.set_title('Return Loss')
    axis.legend()
    
    fig, axis = plt.subplots(num='Insertion Loss', tight_layout=True)
    axis.plot(f/1e9, dB(s21), 'k-',  linewidth=2, label='S21 (dB)')
    axis.grid()
    axis.set_xmargin(0)
    axis.set_xlabel('Frequency (GHz)')
    axis.set_title('Insertion Loss')
    axis.legend()

    fig, axis = plt.subplots(num='Transmission Phase', tight_layout=True)
    axis.plot(f/1e9, phase(s21), 'k-',  linewidth=2, label='S21 (dB)')
    axis.grid()
    axis.set_xmargin(0)
    axis.set_xlabel('Frequency (GHz)')
    axis.set_title('Transmission Phase')
    axis.legend()

    # show all plots
    plt.show()

