'''
Script to test the TE mode converter from the fundamental mode to the second-order mode
'''

# Import necessary python libraries and SPINS-B libraries.
import os
import sys

import datetime
import logging
import os
import pickle

import matplotlib.pyplot as plt
import numpy as np

from spins import goos
from spins.goos import compat
from spins.goos_sim import maxwell
from spins.invdes.problem_graph import optplan, log_tools

'''
Set up Optimization Plan object.
'''
## Create folder where all spins-b results will be saved. ##

# Currently, the folder will be named "TEmc_{current_timestamp}" and will
# be located in the folder containing spins-b. To change this, change
# the following line and set `out_folder_name` somewhere else.
out_folder_name = "TEmc_" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
out_folder = os.path.join(os.getcwd(),out_folder_name)
if (not os.path.exists(out_folder)):
  os.makedirs(out_folder)

## Setup logging and Optimization Plan. ##
plan = goos.OptimizationPlan(save_path = out_folder)

'''
Define the constant background structures that will not be changed
during the design. In this case, these are the input and output waveguides.
'''
with plan:
  # Define input waveguide.
  wg_in = goos.Cuboid(pos=goos.Constant([0, -2000, 0]),
                            extents=goos.Constant([400, 3000, 220]),
                            material=goos.material.Material(index=3.45))
  # Define output waveguide.
  wg_out = goos.Cuboid(pos=goos.Constant([0, 2000, 0]),
                             extents=goos.Constant([400, 3000, 220]),
                             material=goos.material.Material(index=3.45))

  # Group these background structures together.
  eps_background_structures = goos.GroupShape([wg_in, wg_out])

'''
Visualize the constant background structures we just defined.
'''
with plan:
  eps_rendered = maxwell.RenderShape(
            eps_background_structures,
            region=goos.Box3d(center=[0, 0, 0], extents=[3000, 3000, 0]),
            mesh=maxwell.UniformMesh(dx=40),
            wavelength=1550,
        )

  goos.util.visualize_eps(eps_rendered.get().array[2])

'''
Define and initialize the design region.
'''
with plan:
  # Use random initialization, where each pixel is randomly assigned
  # a value in the range [0.3,0.7].
  def initializer(size):
            # Set the seed immediately before calling `random` to ensure
            # reproducibility.
            np.random.seed(247)
            return np.random.random(size) * 0.2 + 0.5

  # Define the design region as a pixelated continuous shape, which is composed
  # of voxels whose permittivities can take on any value between `material` and
  # `material2`.
  var, design = goos.pixelated_cont_shape(
                initializer=initializer,
                pos=goos.Constant([0, 0, 0]),
                extents=[2000, 2000, 220],
                material=goos.material.Material(index=1),
                material2=goos.material.Material(index=3.45),
                pixel_size=[40, 40, 220],
                var_name="var_cont")

  eps = goos.GroupShape([eps_background_structures, design])

'''
Visualize the design region as a sanity check.
'''
with plan:
  eps_rendered = maxwell.RenderShape(
            eps,
            region=goos.Box3d(center=[0, 0, 0], extents=[3000, 3000, 0]),
            mesh=maxwell.UniformMesh(dx=40),
            wavelength=1550,
        )
  goos.util.visualize_eps(eps_rendered.get().array[2])

'''
Setup EM solver - in this case, we use the FDFD solver that comes with
spins-b, maxwell.
'''
with plan:
  # Define wavelength and solver.
  my_wavelength = 1550

  sim_3d = True

  if sim_3d:
      sim_z_extent = 2500
      solver_info = maxwell.MaxwellSolver(solver="maxwell_cg",
                                          err_thresh=1e-2)
      pml_thickness = [400] * 6
  else:
      sim_z_extent = 40
      solver_info = maxwell.DirectSolver()
      pml_thickness = [400, 400, 400, 400, 0, 0]

  # Define simulation space.
  my_simulation_space = maxwell.SimulationSpace(
            mesh=maxwell.UniformMesh(dx=40),
            sim_region=goos.Box3d(
                center=[0, 0, 0],
                extents=[4000, 4000, sim_z_extent],
            ),
            pml_thickness=pml_thickness)

  # Define a waveguide mode source.
  my_sources = [maxwell.WaveguideModeSource(center=[0, -1400, 0],
                                        extents=[2500, 0, 1000],
                                        normal=[0, 1, 0],
                                        mode_num=0,
                                        power=1)]


  # Define simulation outputs.
  my_outputs=[ maxwell.Epsilon(name="eps"),
               maxwell.ElectricField(name="field"),
               maxwell.WaveguideModeOverlap(name="overlap",
                                         center=[0, 1400, 0],
                                         extents=[2500, 0, 1000],
                                         normal=[0, 1, 0],
                                         mode_num=2,
                                         power=1)]

  # Setup the simulation object.
  sim = maxwell.fdfd_simulation(
        name="sim_cont",
        wavelength= my_wavelength,
        background=goos.material.Material(index=1.0),
        eps=eps,
        simulation_space = my_simulation_space,
        solver_info = solver_info,
        sources = my_sources,
        outputs= my_outputs
    )

'''
Visualize simulation of initial structure, as a sanity check.
'''
with plan:
  initial_structure = np.squeeze(sim['eps'].get().array)
  initial_field = np.squeeze(sim['field'].get().array)
  initial_overlap = np.squeeze(sim['overlap'].get().array)

  plt.figure()
  plt.subplot(1,2,1)
  plt.imshow(
    np.abs(
      initial_structure[2]))
  plt.colorbar()

  plt.subplot(1,2,2)
  plt.imshow(
    np.abs(
      initial_field[2]))
  plt.colorbar()
  plt.show()

  print("Initial overlap is {}.".format(np.abs(initial_overlap)))

'''
Define objective function to maximize the transmission of the fundamental
waveguide mode.
'''
obj = goos.rename(goos.abs(sim["overlap"]), name="obj_cont")

'''
Setup the optimization strategy.
'''
with plan:
  goos.opt.scipy_maximize(
                obj,
                "L-BFGS-B",
                monitor_list=[sim["eps"], sim["field"], sim["overlap"], obj],
                max_iters=20,
                name="opt_cont")

'''
Save the plan.
'''
with plan:
  plan.save()

'''
Run the plan.
'''
# Setup logging.
LOG_FORMAT = "[%(asctime)-15s][%(levelname)s][%(module)s][%(funcName)s] %(message)s"
logging.basicConfig(stream=sys.stdout,format=LOG_FORMAT)
logging.getLogger("").setLevel(logging.INFO)
logger = logging.getLogger("")

# Run the plan.
with plan:
  plan.run()

'''
Set the folder for visualizing optimization results.
'''
# Currently, this will visualize results in the output folder defined
# at the beginning of this notebook in the cell where we define the
# OptimizationPlan object - but, you can change this variable
# to visualize results from other folders in your google drive.
folder_plt = out_folder

'''
Open a pkl file and look at what's inside, as a sanity check.
'''
with open(os.path.join(folder_plt, "step1.pkl"), "rb") as fp:
  data = pickle.load(fp)

print("Pkl dictionary: ", data.keys())
print("'monitor_data' dictionary: ", data['monitor_data'].keys())
print("'objective value': ", data['monitor_data']['obj_cont'])

'''
Visualize the objective function and overlap monitors as a function of iteration.
'''

# Load the logged monitor data.
df = log_tools.create_log_data_frame(log_tools.load_all_logs(folder_plt))
monitor_obj = log_tools.get_joined_scalar_monitors(df,"obj_cont", event_name = "optimizing", scalar_operation = None)
monitor_overlap = log_tools.get_joined_scalar_monitors(df,"sim_cont.overlap", event_name = "optimizing",scalar_operation = "magnitude")

# Plot the logged monitor data.
plt.figure()

plt.subplot(2, 1, 1)
plt.plot(monitor_obj.iterations,monitor_obj.data,'k-')
plt.xlabel("Iteration")
plt.ylabel("Objective")
plt.tight_layout()

plt.subplot(2, 1, 2)
plt.plot(monitor_overlap.iterations,monitor_overlap.data,'k-')
plt.xlabel("Iteration")
plt.ylabel("Overlap")
plt.tight_layout()


'''
Visualize the final structure and fields.
'''
step = goos.util.get_latest_log_step(folder_plt)

with open(os.path.join(folder_plt, "step{}.pkl".format(step)), "rb") as fp:
  data = pickle.load(fp)

eps = np.linalg.norm(data["monitor_data"]["sim_cont.eps"], axis=0)
field_norm = np.linalg.norm(data["monitor_data"]["sim_cont.field"], axis=0)

plt.figure()

plt.subplot(1, 2, 1)
plt.imshow(
    eps[:, :, eps.shape[2] // 2].squeeze())
plt.colorbar()
plt.title("|Eps_z|")
plt.tight_layout()

plt.subplot(1, 2, 2)
plt.imshow(
    field_norm[:, :, field_norm.shape[2] // 2].squeeze())
plt.colorbar()
plt.title("|E_z|")
plt.tight_layout()

plt.savefig(os.path.join(folder_plt,"fields_step"+str(step)+".png"))
plt.show()