from enum import Enum
from typing import Tuple

from autograd.extend import primitive, defvjp
from autograd import numpy as anp
from autograd import tensor_jacobian_product
import matplotlib.pyplot as plt
import meep as mp
import meep.adjoint as mpa


RESOLUTION_UM = 100
WAVELENGTH_UM = 0.5
GRATING_PERIOD_UM = 2.5
PADDING_UM = 3.0

Polarization = Enum("Polarization", "S P")


def mapping(
        weights: anp.ndarray,
        eta: float,
        beta: float,
        filter_radius_um: float,
        design_region_size: mp.Vector3,
        design_region_resolution: float
) -> anp.ndarray:
    """A differentiable mapping function for the design weights."""

    # Note: weights should be a 2d array. If it is 1d, it will be reshaped                                        
    #       to 2d by conic_filter. There is no need to reshape it here.                                           
    filtered_field = mpa.conic_filter(
        weights,
        filter_radius_um,
        design_region_size.x,
        design_region_size.y,
        design_region_resolution,
    )

    if beta == 0:
        return filtered_field.flatten()
    else:
        projected_field = mpa.smoothed_projection(
            filtered_field,
            beta,
            eta,
            design_region_resolution,
        )
        return projected_field.flatten()


def design_region_to_meshgrid(
        design_region_size: mp.Vector3,
        nx: int,
        ny: int
) -> Tuple[anp.ndarray, anp.ndarray]:
    """Returns the meshgrid coordinate arrays for the design region."""

    xcoord = anp.linspace(
        -0.5*design_region_size.x,
        +0.5*design_region_size.x,
        nx
    )
    ycoord = anp.linspace(
        -0.5*design_region_size.y,
        +0.5*design_region_size.y,
        ny
    )
    xv, yv = anp.meshgrid(xcoord, ycoord, indexing='ij')

    return xv, yv


@primitive
def grating_levelset(
        grating_height_um: float,
        grating_duty_cycle: float,
        design_region_size: mp.Vector3,
        nx: int,
        ny: int
) -> anp.ndarray:
    """Returns the weights for the grating as a 1D array."""

    xv, yv = design_region_to_meshgrid(design_region_size, nx, ny)

    weights = anp.where(
        anp.abs(yv) <= 0.5 * grating_duty_cycle * GRATING_PERIOD_UM,
        anp.where(
            xv <= xv[0][0] + grating_height_um,
            1,
            0
        ),
        0
    )

    return weights.flatten()


def grating_levelset_vjp(
        ans,
        grating_height_um: float,
        grating_duty_cycle: float,
        design_region_size: mp.Vector3,
        nx: int,
        ny: int
) -> anp.ndarray:
    """Returns the vector-Jacobian product."""

    xv, yv = design_region_to_meshgrid(design_region_size, nx, ny)

    # pixel dimensions                                                                                            
    delta_x = design_region_size.x / (nx - 1)
    delta_y = design_region_size.y / (ny - 1)

    # Gradient of the level set with respect to the grating height.                                               
    # The gradient is obtained using a finite difference. The gradient is                                         
    # therefore a zero matrix with non-zero entries only at the locations                                         
    # of the height boundary. Since the height boundary is normal to the                                          
    # x axis, these non-zero matrix elements have a value of 1 / Δx.                                              
    weights = anp.where(
        anp.abs(yv) <= 0.5 * grating_duty_cycle * GRATING_PERIOD_UM,
        anp.where(
            xv <= xv[0][0] + grating_height_um + 0.5 * delta_x,
            anp.where(
                xv >= xv[0][0] + grating_height_um - 0.5 * delta_x,
                1 / delta_x,
                0
            ),
            0
        ),
        0
    )

    jacobian = weights.flatten()

    return lambda g: anp.tensordot(g, jacobian, axes=1)


def grating_1d(
        pol: Polarization,
        grating_height_um: float,
        design_region_size: mp.Vector3,
        nx: int,
        ny: int
) -> mpa.OptimizationProblem:
    """Sets up the adjoint optimization of a 1D grating."""

    frequency = 1 / WAVELENGTH_UM
    substrate_um = 3.0
    pml_um = 1.0
    pml_layers = [mp.PML(direction=mp.X, thickness=pml_um)]

    n_glass = 1.5
    glass = mp.Medium(index=n_glass)

    size_x_um = pml_um + substrate_um + grating_height_um + PADDING_UM + pml_um
    size_y_um = GRATING_PERIOD_UM
    cell_size = mp.Vector3(size_x_um, size_y_um, 0)

    k_point = mp.Vector3()

    if pol.name == "S":
        eig_parity = mp.ODD_Z
        src_cmpt = mp.Ez
    else:
        eig_parity = mp.EVEN_Z
        src_cmpt = mp.Hz

    src_pt = mp.Vector3(-0.5 * size_x_um + pml_um, 0, 0)
    sources = [
        mp.Source(
            mp.GaussianSource(frequency, fwidth=0.1 * frequency),
            component=src_cmpt,
            center=src_pt,
            size=mp.Vector3(0, size_y_um, 0),
        )
    ]

    matgrid = mp.MaterialGrid(
        mp.Vector3(nx, ny),
        mp.air,
        glass,
        weights=anp.ones((nx, ny)),
        do_averaging=False,
    )

    matgrid_region = mpa.DesignRegion(
        matgrid,
        volume=mp.Volume(
            center=mp.Vector3(
                (-0.5 * size_x_um + pml_um + substrate_um +
                 0.5 * design_region_size.x),
                0,
                0
            ),
            size=design_region_size,
        ),
    )

    geometry = [
        mp.Block(
            material=glass,
            size=mp.Vector3(pml_um + substrate_um, mp.inf, mp.inf),
            center=mp.Vector3(
                -0.5 * size_x_um + 0.5 * (pml_um + substrate_um), 0, 0
            ),
        ),
        mp.Block(
            material=matgrid,
            size=matgrid_region.size,
            center=matgrid_region.center,
        ),
    ]

    sim = mp.Simulation(
        resolution=RESOLUTION_UM,
        cell_size=cell_size,
        boundary_layers=pml_layers,
        k_point=k_point,
        sources=sources,
        geometry=geometry,
    )

    tran_pt = mp.Vector3(0.5 * size_x_um - pml_um, 0, 0)

    order_y = 1
    kdiff = mp.Vector3(
        (frequency**2 - (order_y / GRATING_PERIOD_UM)**2)**0.5,
        order_y / GRATING_PERIOD_UM,
        0
    )
    print(f"kdiff = ({kdiff.x:.5f}, {kdiff.y:.5f}, {kdiff.z:.5f})")

    obj_list = [
        mpa.EigenmodeCoefficient(
            sim,
            mp.Volume(
                center=tran_pt,
                size=mp.Vector3(0, size_y_um, 0),
            ),
            mode=1,
            kpoint_func=lambda *not_used: kdiff,
            eig_parity=eig_parity,
            eig_vol=mp.Volume(
                center=tran_pt,
                size=mp.Vector3(0, 1 / RESOLUTION_UM, 0)
            ),
        ),
    ]

    def J(tran_mon):
        return anp.abs(tran_mon)**2

    opt = mpa.OptimizationProblem(
        simulation=sim,
        objective_functions=J,
        objective_arguments=obj_list,
        design_regions=[matgrid_region],
        frequencies=[frequency],
    )

    return opt


if __name__ == "__main__":
    grating_height_um = 0.5
    grating_duty_cycle = 0.5
    height_perturbation_um = 1 / RESOLUTION_UM

    design_region_size = mp.Vector3(
        grating_height_um + PADDING_UM,
        GRATING_PERIOD_UM,
        0
    )
    design_region_resolution_um = int(2 * RESOLUTION_UM)
    nx = int(design_region_size.x * design_region_resolution_um) + 1
    ny = int(design_region_size.y * design_region_resolution_um) + 1

    # Properties of the convolution filter and projection operator.                                               
    beta = 0.2
    eta_intrinsic = 0.5
    eta_erosion = 0.75
    min_feature_size_um = 0.5
    filter_radius_um = mpa.get_conic_radius_from_eta_e(
        min_feature_size_um,
        eta_erosion
    )

    opt = grating_1d(
        Polarization.P,
        grating_height_um,
        design_region_size,
        nx,
        ny,
    )

    unperturbed_design_weights = grating_levelset(
        grating_height_um,
        grating_duty_cycle,
        design_region_size,
        nx,
        ny,
    )
    mapped_design_weights = mapping(
        unperturbed_design_weights,
        eta_intrinsic,
        beta,
        filter_radius_um,
        design_region_size,
        design_region_resolution_um,
    )
    obj_value_unperturbed, grad_unperturbed = opt(
        [mapped_design_weights],
        need_gradient=True,
    )

    # Backpropagate the adjoint gradient through two functions in the correct order.                              
    grad_backprop = tensor_jacobian_product(mapping, 0)(
        mapped_design_weights,
        eta_intrinsic,
        beta,
        filter_radius_um,
        design_region_size,
        design_region_resolution_um,
        grad_unperturbed,
    )

    defvjp(grating_levelset, grating_levelset_vjp)
    grad_backprop = tensor_jacobian_product(grating_levelset, 0)(
        grating_height_um,
        grating_duty_cycle,
        design_region_size,
        nx,
        ny,
        grad_backprop,
    )

    fig, ax = plt.subplots()
    opt.plot2D(init_opt=False, ax=ax)
    if mp.am_master():
        fig.savefig(
            'grating_1d_plot2D.png', dpi=150, bbox_inches='tight'
        )

    perturbed_design_weights = grating_levelset(
        grating_height_um + height_perturbation_um,
        grating_duty_cycle,
        design_region_size,
        nx,
        ny,
    )
    mapped_design_weights = mapping(
        perturbed_design_weights,
        eta_intrinsic,
        beta,
        filter_radius_um,
        design_region_size,
        design_region_resolution_um,
    )
    obj_value_perturbed, _ = opt(
        [mapped_design_weights],
        need_gradient=False
    )

    adj_dirderiv = height_perturbation_um * grad_backprop
    fnd_dirderiv = obj_value_perturbed[0] - obj_value_unperturbed[0]
    rel_err = abs(adj_dirderiv - fnd_dirderiv) / fnd_dirderiv
    print(
        f"dirderiv:, {fnd_dirderiv} (finite difference), "
        f"{adj_dirderiv} (adjoint), {rel_err:.6f} (error)"
    )

   # Obtain the grating as a level set.
    unperturbed_design_weights = grating_weights(
        grating_height_um,
        grating_duty_cycle,
        design_region_size,
        nx,
        ny,
    )
   # Apply a filter and projection to the level set.
    mapped_unperturbed_design_weights = mapping(
        unperturbed_design_weights,
        eta_intrinsic,
        npa.inf,
        filter_radius_um,
        design_region_size,
        design_region_resolution_um,
    )

    obj_value_unperturbed, grad_unperturbed = opt(
        [mapped_unperturbed_design_weights],
        need_gradient=True,
    )

   # Obtain the grating as a level set.
    height_perturbation_um = 1e-2
    perturbed_design_weights = grating_weights(
        grating_height_um + height_perturbation_um,
        grating_duty_cycle,
        design_region_size,
        nx,
        ny,
    )
   # Apply a filter and projection to the level set.
    mapped_perturbed_design_weights = mapping(
        perturbed_design_weights,
        eta_intrinsic,
        npa.inf,
        filter_radius_um,
        design_region_size,
        design_region_resolution_um,
    )

    obj_value_perturbed, _ = opt(
        [mapped_perturbed_design_weights],
        need_gradient=False
    )

def mapping(
        weights: npa.ndarray,
        eta: float,
        beta: float,
        filter_radius_um: float,
        design_region_size: mp.Vector3,
        design_region_resolution: float
) -> np.ndarray:
    """A differentiable mapping function for the design weights."""

    filtered_field = mpa.conic_filter(
        weights,
        filter_radius_um,
	design_region_size.x,
        design_region_size.y,
        design_region_resolution,
    )

    if beta == 0:
        return filtered_field.flatten()

    else:
        projected_field = mpa.smoothed_projection(
            filtered_field,
            beta,
            eta,
            design_region_resolution,
        )

        return projected_field.flatten()

def grating_levelset(
        grating_height_um: float,
        grating_duty_cycle: float,
        design_region_size: mp.Vector3,
        nx: int,
        ny: int
) -> npa.ndarray:
    """Returns the weights for the grating as a 1D array."""

    xcoord = npa.linspace(
        -0.5*design_region_size.x,
        +0.5*design_region_size.x,
        nx
    )
    ycoord = npa.linspace(
        -0.5*design_region_size.y,
        +0.5*design_region_size.y,
        ny
    )
    xv, yv = npa.meshgrid(xcoord, ycoord, indexing='ij')
    weights = npa.where(
        npa.abs(yv) <= 0.5 * grating_duty_cycle * GRATING_PERIOD_UM,
        npa.where(
            xv <= xcoord[0] + grating_height_um,
            1,
            0
        ),
        0
    )

    return weights.flatten()

def grating_levelset_vjp(
        ans,
        grating_height_um: float,
        grating_duty_cycle: float,
        design_region_size: mp.Vector3,
        nx: int,
        ny: int
) -> np.ndarray:
    """Returns the vector-Jacobian product."""

    xv, yv = design_region_to_meshgrid(design_region_size, nx, ny)

    # pixel dimensions
    delta_x = design_region_size.x / (nx - 1)
    delta_y = design_region_size.y / (ny - 1)

    # Gradient of the level set with respect to the grating height.
    # The gradient is obtained using a finite difference. The gradient is
    # therefore a zero matrix with non-zero entries only at the locations
    # of the height boundary. Since the height boundary is normal to the
    # x axis, these non-zero matrix elements have a value of 1 / Δx.
    weights = np.where(
        np.abs(yv) <= 0.5 * grating_duty_cycle * GRATING_PERIOD_UM,
        np.where(
            xv <= xv[0][0] + grating_height_um + 0.5 * delta_x,
            np.where(
                xv >= xv[0][0] + grating_height_um - 0.5 * delta_x,
                1 / delta_x,
                0
            ),
            0
        ),
        0
    )

    jacobian = weights.flatten()

    return lambda g: np.tensordot(g, jacobian, axes=1)

    matgrid = mp.MaterialGrid(
        mp.Vector3(nx, ny),
        mp.air,
        glass,
        weights=anp.ones((nx, ny)),
        do_averaging=True,
        beta=anp.inf
    )

def signed_distance(weights: anp.ndarray) -> anp.ndarray:
    """Maps the 2D weights using a signed-distance function."""

    # Create signed distance function.
    sd = skfmm.distance(weights - 0.5)

    # Linear interpolation of zero-levelset onto 0.5-levelset.
    xp = [anp.min(sd.flatten()), 0, anp.max(sd.flatten())]
    yp = [0, 0.5001, 1]

    return anp.interp(sd.flatten(), xp, yp)
  
  
@primitive
def levelset_and_smoothing(
        grating_height_um: float,
        grating_duty_cycle: float,
        design_region_size: mp.Vector3,
        nx: int,
        ny: int
) -> anp.ndarray:
    """Generates the grating as a levelset with signed-distance smoothing."""

    xv, yv = design_region_to_meshgrid(design_region_size, nx, ny)

    weights = anp.where(
        anp.abs(yv) <= 0.5 * grating_duty_cycle * GRATING_PERIOD_UM,
        anp.where(
            xv <= xv[0][0] + grating_height_um,
            1,
            0
        ),
        0
    )

    return signed_distance(weights)


def levelset_and_smoothing_vjp(
        ans: anp.ndarray,
        grating_height_um: float,
        grating_duty_cycle: float,
        design_region_size: mp.Vector3,
        nx: int,
        ny: int
) -> anp.ndarray:
    """Computes the vector-Jacobian product."""

    xv, yv = design_region_to_meshgrid(design_region_size, nx, ny)

    # pixel dimensions 
    delta_x = design_region_size.x / (nx - 1)

    weights = anp.where(
        anp.abs(yv) <= 0.5 * grating_duty_cycle * GRATING_PERIOD_UM,
        anp.where(
            xv <= xv[0][0] + grating_height_um + delta_x,
            1,
            0
        ),
        0
    )

    jacobian = (signed_distance(weights) - ans) / delta_x

    return lambda g: anp.tensordot(g, jacobian, axes=1)