import contextlib
import pathlib
from collections import OrderedDict
from typing import Optional, Tuple

import numpy as np
from omegaconf import DictConfig

from gdsfactory import Component
from gdsfactory.config import CONFIG

c = 2.9979e8
pi = np.pi
um = 1e-6


def install_design_kit(
    session: object,
    install_dir: pathlib.Path = CONFIG["interconnect"],
    overwrite: bool = False,
):
    from gdsfactory.pdk import get_interconnect_cml_path

    cml_path = get_interconnect_cml_path()
    session.installdesignkit(str(cml_path), str(install_dir), overwrite)


def set_named_settings(session: object, simulation_settings: dict, element: str):
    for param, val in zip(simulation_settings.keys(), simulation_settings.values()):
        session.setnamed(element, param, val)


# TODO: Add custom s-parameter models using compound elements


def add_interconnect_element(
    session: object,
    label: str,
    model: str,
    loc: Tuple[float, float] = (200.0, 200.0),
    flip_vert: bool = False,
    flip_horiz: bool = False,
    rotation: float = 0.0,
    simulation_props: Optional[OrderedDict] = None,
):
    """Add an element to the Interconnect session.

    TODO: Need to connect this to generated s-parameters and add them to the model as well

    Args:
        session: Interconnect session.
        label: label for Interconnect component.
        model:
        loc:
        flip_vert:
        flip_horiz:
        rotation:
        extra_props:

    """
    props = OrderedDict(
        [
            ("name", label),
            ("x position", loc[0]),
            ("y position", loc[1]),
            ("horizontal flipped", float(flip_horiz)),
            ("vertical flipped", float(flip_vert)),
            ("rotated", rotation),
        ]
    )
    if simulation_props:
        if "library" in simulation_props.keys():
            _ = simulation_props.pop("library")
        if "properties" in simulation_props.keys():
            props.update(simulation_props["properties"])
        else:
            props.update(simulation_props)
    return session.addelement(model, properties=props)


def get_interconnect_settings(instance):
    info = instance.info.copy()
    if "interconnect" not in info.keys():
        return dict()
    settings = info["interconnect"]
    if "properties" not in settings:
        settings["properties"] = []
    if "layout_model_property_pairs" in settings.keys():
        pairs = settings.pop("layout_model_property_pairs")
        for inc_name, (layout_name, scale) in pairs.items():
            settings["properties"][inc_name] = info[layout_name] * scale
    return settings


def send_to_interconnect(
    component: Component,
    session: object,
    ports_in: dict = None,
    ports_out: dict = None,
    placements: dict = None,
    simulation_settings: OrderedDict = None,
    drop_port_prefix: str = None,
    component_distance_scaling_x: float = 1,
    component_distance_scaling_y: float = 1,
    setup_mc: bool = False,
    exclude_electrical: bool = True,
    **settings,
) -> object:
    """Send netlist components to Interconnect and connect them according to netlist.

    Args:
        component: component from which to extract netlist.
        session: Interconnect session.
        placements: x,y pairs for where to place the components in the Interconnect GUI.
        simulation_settings: global settings for Interconnect simulation.
        drop_port_prefix: if components are written with some prefix, drop up to and including
            the prefix character.  (i.e. "c1_input" -> "input").
        component_distance_scaling: scaling factor for component distances when
            laying out Interconnect schematic.

    """
    if not session:
        import lumapi

        session = lumapi.INTERCONNECT()

    install_design_kit(session=session)

    session.switchtolayout()
    session.deleteall()

    # Make compound element for circuit
    compound_element = session.addelement(
        "Compound Element", properties=OrderedDict([("name", str(component.name))])
    )
    if setup_mc:
        MC_param_element = f"::Root Element::{component.name}"
        # Add Monte-Carlo params
        session.addproperty(
            MC_param_element, "MC_uniformity_thickness", "wafer", "Matrix"
        )
        session.addproperty(MC_param_element, "MC_uniformity_width", "wafer", "Matrix")
        session.addproperty(MC_param_element, "MC_non_uniform", "wafer", "Number")
        session.addproperty(MC_param_element, "MC_grid", "wafer", "Number")
        session.addproperty(MC_param_element, "MC_resolution_x", "wafer", "Number")
        session.addproperty(MC_param_element, "MC_resolution_y", "wafer", "Number")

    # Switch groupscope to compound element so that anything added will go into it
    session.groupscope(compound_element.name)

    c = component

    netlist = c.get_netlist()

    instances: DictConfig = netlist["instances"]
    connections: DictConfig = netlist["connections"]
    placements: DictConfig = placements or netlist["placements"]
    ports: DictConfig = netlist["ports"]

    relay_count = 1
    excluded = []
    for instance in instances:
        if exclude_electrical:
            # Exclude if purely electrical
            with contextlib.suppress(Exception):
                port_types = instances[instance].full["cross_section"].port_types
                if ("electrical" in port_types) and ("optical" not in port_types):
                    excluded.append(instance)
                    continue

        loc = (
            component_distance_scaling_x * placements[instance].x,
            component_distance_scaling_y * placements[instance].y,
        )

        sim_props = get_interconnect_settings(instances[instance])

        if "model" not in sim_props.keys():
            raise KeyError(f"Please specify an interconnect model for {instance!r}")
        model = sim_props.pop("model")

        if "layout_model_port_pairs" in sim_props.keys():
            _ = sim_props.pop("layout_model_port_pairs")

        add_interconnect_element(
            session=session,
            label=instance,
            loc=loc,
            rotation=float(placements[instance].rotation),
            model=model,
            simulation_props=sim_props,
        )
        if instance in ports_in:
            # Add input port and connect to the compound element input port
            session.addport(
                f"::Root Element::{compound_element.name}",
                f"{instance}.{ports_in[instance]}",
                "input",
                "Optical Signal",
                "left",
            )
            session.connect(
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "output",
                instance,
                ports_in[instance],
            )
            session.setnamed(
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "y position",
                loc[1] - 50,
            )
            session.setnamed(
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "x position",
                loc[0] - 50,
            )
            relay_count += 1
        elif instance in ports_out:
            session.addport(
                f"::Root Element::{compound_element.name}",
                f"{instance}.{ports_out[instance]}",
                "output",
                "Optical Signal",
                "right",
            )
            session.connect(
                instance,
                ports_out[instance],
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "input",
            )
            session.setnamed(
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "y position",
                loc[1] + 50,
            )
            session.setnamed(
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "x position",
                loc[0] + 50,
            )
            relay_count += 1

    if ports:
        for port in ports_in:
            input_instance, instance_port = ports[port].split(",")
            info = get_interconnect_settings(instances[input_instance])
            if "layout_model_port_pairs" in info.keys():
                instance_port = info["layout_model_port_pairs"][instance_port]
            session.addport(
                f"::Root Element::{compound_element.name}",
                str(port),
                "input",
                "Optical Signal",
                "left",
            )
            session.connect(
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "output",
                input_instance,
                instance_port,
            )
            session.setnamed(
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "y position",
                session.getnamed(input_instance, "y position") - 50,
            )
            session.setnamed(
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "x position",
                session.getnamed(input_instance, "x position") - 50,
            )
            relay_count += 1
        for port in ports_out:
            output_instance, instance_port = ports[port].split(",")
            info = get_interconnect_settings(instances[output_instance])
            if "layout_model_port_pairs" in info.keys():
                instance_port = info["layout_model_port_pairs"][instance_port]
            session.addport(
                f"::Root Element::{compound_element.name}",
                str(port),
                "output",
                "Optical Signal",
                "right",
            )
            session.connect(
                output_instance,
                instance_port,
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "input",
            )
            session.setnamed(
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "y position",
                session.getnamed(output_instance, "y position") + 50,
            )
            session.setnamed(
                f"::Root Element::{compound_element.name}::RELAY_{relay_count}",
                "x position",
                session.getnamed(output_instance, "x position") + 50,
            )
            relay_count += 1

    for connection in connections:
        element2, port2 = connection.split(",")
        element1, port1 = connections[connection].split(",")
        if (element1 in excluded) or (element2 in excluded):
            continue
        if drop_port_prefix:
            # a bad way to autodetect which ports need to have prefixes dropped..
            with contextlib.suppress(Exception):
                port1 = port1[port1.index(drop_port_prefix) + 1 :]
            with contextlib.suppress(Exception):
                port2 = port2[port2.index(drop_port_prefix) + 1 :]

        # EBeam ports are not named consistently between Klayout and Interconnect..
        element1_info = get_interconnect_settings(instances[element1])
        if "layout_model_port_pairs" in element1_info.keys():
            port1 = element1_info["layout_model_port_pairs"][port1]
        element2_info = get_interconnect_settings(instances[element2])
        if "layout_model_port_pairs" in element2_info.keys():
            port2 = element2_info["layout_model_port_pairs"][port2]

        session.connect(element1, port1, element2, port2)

    if simulation_settings:
        set_named_settings(
            session, simulation_settings, element=f"::Root Element::{component.name}"
        )
    # Return to highest-level element
    session.groupscope("::Root Element")

    # Auto-distribute ports on the compound element
    session.autoarrange(component.name)
    return session


def run_wavelength_sweep(
    component: Component,
    session: Optional[object] = None,
    setup_simulation: bool = True,
    is_top_level: bool = False,
    ports_in: Optional[dict] = None,
    ports_out: Optional[dict] = None,
    mode: int = 1,
    wavelength_range: Tuple[float, float] = (1.500, 1.600),
    n_points: int = 1000,
    results: Tuple[str, ...] = ("transmission",),
    extra_ona_props: Optional[dict] = None,
    **kwargs,
) -> dict:
    """Args are the following.

    component:
    session:
    setup_simulation: whether to send the component to interconnect before running the sweep.
    ports_in: specify the port in the Interconnect model to attach the ONA output to.
    ports_out: specify the ports in the Interconnect models to attach the ONA input to.
    wavelength_range:
    n_points:
    results:
    extra_ona_props:
    kwargs:
    """
    if len(ports_in) > 1:
        raise ValueError("Only 1 input port is supported at this time")

    import lumapi

    if not session:
        session = lumapi.INTERCONNECT()

    install_design_kit(session=session)

    if setup_simulation:
        session = send_to_interconnect(
            component=component,
            session=session,
            ports_in=ports_in,
            ports_out=ports_out,
            **kwargs,
        )

    ona_props = OrderedDict(
        [
            ("number of input ports", len(ports_out)),
            ("number of points", n_points),
            ("input parameter", "start and stop"),
            ("start frequency", (c / (wavelength_range[1] * um))),
            ("stop frequency", (c / (wavelength_range[0] * um))),
            ("plot kind", "wavelength"),
            ("relative to center", float(False)),
        ]
    )
    if extra_ona_props:
        ona_props.update(extra_ona_props)

    ona = add_interconnect_element(
        session=session,
        model="Optical Network Analyzer",
        label="ONA_1",
        loc=(0, -50),
        simulation_props=ona_props,
    )
    for port in ports_in.keys():
        name = port if is_top_level else f"{port}.{ports_in[port]}"
        session.connect(ona.name, "output", component.name, name)
    for i, port in enumerate(ports_out.keys()):
        name = port if is_top_level else f"{port}.{ports_out[port]}"
        session.connect(ona.name, f"input {i+1}", component.name, name)

    session.run()

    # inc.close()
    return {
        result: {
            port: session.getresult(ona.name, f"input {i+1}/mode {mode}/{result}")
            for i, port in enumerate(ports_out)
        }
        for result in results
    }


def plot_wavelength_sweep(
    ports_out, results, result_name: str = "TE Transmission", show: bool = True
):
    import matplotlib.pyplot as plt

    for port in ports_out:
        wl = results["transmission"][port]["wavelength"] / um
        T = 10 * np.log10(np.abs(results["transmission"][port][result_name]))
        plt.plot(wl, T, label=str(port))
    plt.legend()
    plt.xlabel(r"Wavelength ($\mu$m)")
    plt.ylabel(f"{result_name} (dB)")

    if show:
        plt.show()


if __name__ == "__main__":

    import ubcpdk.components as pdk

    mzi = pdk.mzi()

    netlist = mzi.get_netlist()

    ports_in = {"o1": "o1"}
    ports_out = {"o2": "o2"}

    simulation_settings = OrderedDict(
        [
            ("MC_uniformity_thickness", np.array([200, 200])),
            ("MC_uniformity_width", np.array([200, 200])),
            ("MC_non_uniform", 0),
            ("MC_grid", 1e-5),
            ("MC_resolution_x", 200),
            ("MC_resolution_y", 0),
        ]
    )
    results = run_wavelength_sweep(
        component=mzi,
        ports_in=ports_in,
        ports_out=ports_out,
        results=("transmission",),
        component_distance_scaling=50,
        simulation_settings=simulation_settings,
        setup_mc=True,
        is_top_level=True,
    )

    plot_wavelength_sweep(ports_out=ports_out, results=results)
