import matplotlib.pyplot as plt
import numpy as np
import tidy3d as td
import tidy3d.web as web

lda0 = 0.78  # Simulation wavelength
freq0 = td.C_0 / lda0  # Simulation frequency
fwidth = freq0 / 10  # Source frequency width

h_wg = 0.25  # Waveguide thickness
w_wg = 0.45  # Waveguide width
r = 1.8  # Radius of the design region
h_box = 2  # Box layer thickness

# Define the materials
SiN = td.Medium.from_nk(n=2.0255, k=0, freq=freq0)
Si = td.Medium.from_nk(n=3.4945, k=0, freq=freq0)
SiO2 = td.Medium.from_nk(n=1.4537, k=0, freq=freq0)

inf_eff = 1e2  # Effective infinity

# Define the SiN waveguides connected to the design region
waveguide_left = td.Structure(
    geometry=td.Box.from_bounds(rmin=(-inf_eff, -w_wg / 2, 0), rmax=(-r, w_wg / 2, h_wg)),
    medium=SiN,
)

waveguide_right = td.Structure(
    geometry=td.Box.from_bounds(rmin=(r, -w_wg / 2, 0), rmax=(inf_eff, w_wg / 2, h_wg)),
    medium=SiN,
)

# Define the SiO2 box and substrate
box = td.Structure(
    geometry=td.Box.from_bounds(rmin=(-inf_eff, -inf_eff, -h_box), rmax=(inf_eff, inf_eff, 0)),
    medium=SiO2,
)

substrate = td.Structure(
    geometry=td.Box.from_bounds(
        rmin=(-inf_eff, -inf_eff, -inf_eff), rmax=(inf_eff, inf_eff, -h_box)
    ),
    medium=Si,
)

# Define the mode monitors
mode_spec = td.ModeSpec(num_modes=1, target_neff=2)
mode_monitor_left = td.ModeMonitor(
    center=(-r - lda0 / 2, 0, h_wg / 2),
    size=(0, 4 * w_wg, 5 * h_wg),
    freqs=[freq0],
    mode_spec=mode_spec,
    name="mode_left",
)

mode_monitor_right = td.ModeMonitor(
    center=(r + lda0 / 2, 0, h_wg / 2),
    size=(0, 4 * w_wg, 5 * h_wg),
    freqs=[freq0],
    mode_spec=mode_spec,
    name="mode_right",
)

def make_sims(design_params: np.ndarray):
    """
    Creates two Tidy3D simulations.
    One simulation corresponds to right circularly polarized (RCP) light,
    and the other to left circularly polarized (LCP) light.
    """

    # Create the design region structure
    design_region = td.Structure(geometry=create_design_region(design_params), medium=SiN)

    # Define the X-polarized component of the Gaussian beam source
    gaussian_beam_x = td.GaussianBeam(
        center=(0, 0, h_wg + lda0 / 5),
        size=(td.inf, td.inf, 0),
        source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth),
        direction="-",
        waist_radius=r,
    )

    # Define the Y-polarized component for Right Circular Polarization (RCP)
    gaussian_beam_y_right = td.GaussianBeam(
        center=(0, 0, h_wg + lda0 / 5),
        size=(td.inf, td.inf, 0),
        source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth, phase=np.pi / 2),
        direction="-",
        waist_radius=r,
        pol_angle=np.pi / 2,
    )

    # Define the Y-polarized component for Left Circular Polarization (LCP)
    gaussian_beam_y_left = gaussian_beam_y_right.updated_copy(
        source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth, phase=-np.pi / 2)
    )

    # Define the simulation domain box
    sim_box = td.Box.from_bounds(
        rmin=(-r - lda0, -r - lda0, -h_box - lda0 / 2),
        rmax=(r + lda0, r + lda0, h_wg + lda0),
    )

    # Create the simulation object for the RCP case
    sim_right = td.Simulation(
        center=sim_box.center,
        size=sim_box.size,
        grid_spec=td.GridSpec.auto(min_steps_per_wvl=12),
        run_time=2e-13,
        structures=[design_region, waveguide_left, waveguide_right, box, substrate],
        sources=[gaussian_beam_x, gaussian_beam_y_right],
        monitors=[mode_monitor_left, mode_monitor_right],
    )

    # Create the simulation object for the LCP case
    sim_left = sim_right.updated_copy(sources=[gaussian_beam_x, gaussian_beam_y_left])

    return sim_right, sim_left

def postprocess(sim_data_right, sim_data_left):
    """
    Post-processes simulation data to calculate the objective function value.

    Args:
        sim_data_right (td.SimulationData): Simulation data for right-circularly polarized (RCP) input.
        sim_data_left (td.SimulationData): Simulation data for left-circularly polarized (LCP) input.


    Returns:
        float: The calculated objective function value.
    """
    # Calculate the transmitted power for RCP input, detected as left-circularly polarized (LCP).
    T_RCP_L = np.abs(sim_data_right["mode_left"].amps.sel(direction="-").item()) ** 2
    # Calculate the transmitted power for RCP input, detected as right-circularly polarized (RCP).
    T_RCP_R = np.abs(sim_data_right["mode_right"].amps.sel(direction="+").item()) ** 2
    # Calculate the transmitted power for LCP input, detected as left-circularly polarized (LCP).
    T_LCP_L = np.abs(sim_data_left["mode_left"].amps.sel(direction="-").item()) ** 2
    # Calculate the transmitted power for LCP input, detected as right-circularly polarized (RCP).
    T_LCP_R = np.abs(sim_data_left["mode_right"].amps.sel(direction="+").item()) ** 2

    # Calculate the final metric based on the transmitted powers.
    return 0.5 * ((T_RCP_L - T_RCP_R) / T_RCP_L + (T_LCP_R - T_LCP_L) / T_LCP_R)

n_designs = 20  # Number of initial designs

# Generate random design parameters
rng = np.random.default_rng(seed=42)
design_params_history = rng.integers(0, 2, size=(n_designs, n_pixel))

# Create simulations
sims = {
    f"{i}_{pol}": sim
    for i, design in enumerate(design_params_history)
    for pol, sim in zip(("right", "left"), make_sims(design))
}

# Run simulations
batch = web.Batch(simulations=sims, verbose=False)
batch_results = batch.run(path_dir="data")

def obj_func(design_params):
    """
    Objective function for optimization.

    Args:
        design_params (np.array): The design parameters.

    Returns:
        float: The objective value.
    """

    sim_right, sim_left = make_sims(design_params)
    batch = web.Batch(simulations={"right": sim_right, "left": sim_left}, verbose=False)
    batch_results = batch.run(path_dir="data")
    return postprocess(batch_results["right"], batch_results["left"])


def direct_binary_search(obj_func, initial_params, termination, max_passes=3):
    """
    Performs Direct Binary Search (DBS) to maximize the objective function.

    Args:
        obj_func (callable): The objective function to maximize.
        initial_design (np.array): The starting binary design vector.
        termination (float): Optimization termination condition
        max_passes (int): Maximum number of full passes over all pixels.

    Returns:
        best_design (np.array): The optimized binary array.
        best_objective (float): The corresponding function value.
    """

    current_design = initial_params.copy().astype(int)

    # Calculate initial objective
    current_objective = obj_func(current_design)
    best_design = current_design.copy()
    best_objective = current_objective

    n_params = len(current_design)

    print(f"Starting DBS. Parameters: {n_params}, Initial objective: {current_objective:.4f}")

    np.random.seed(123)
    for pass_num in range(max_passes):
        improved_in_this_pass = False

        # Create a random order of indices to visit (stochastic traversal prevents bias)
        indices = np.random.permutation(n_params)

        for idx in indices:
            # Flip the bit (0->1 or 1->0)
            current_design[idx] = 1 - current_design[idx]

            # Evaluate new design
            new_objective = obj_func(current_design)

            # Decision Logic
            if new_objective > current_objective:
                # Keep the change
                current_objective = new_objective
                best_design = current_design.copy()
                best_objective = new_objective
                improved_in_this_pass = True

                print(
                    f"  [Pass {pass_num + 1}] Improved! Idx {idx}: {1 - current_design[idx]}->{current_design[idx]} | Highest objective value: {new_objective:.4f}"
                )

                if new_objective >= termination:
                    print(
                        f"Termination condition reached (objective={new_objective:.4f} ≥ 0.99). Stopping optimization."
                    )
                    return best_design, best_objective

            else:
                # Revert the change
                current_design[idx] = 1 - current_design[idx]

        # Convergence check
        if not improved_in_this_pass:
            print(f"Converged at pass {pass_num + 1}. No further improvements found.")
            break

    return best_design, best_objective