import pytest
from numpy import pi
from numpy.testing import assert_allclose

from masque.builder.tools import AutoTool
from masque.builder.pather import Pather
from masque.library import Library
from masque.pattern import Pattern
from masque.ports import Port


def _make_transition_straight(length: float, width: float = 2, ptype: str = "wire") -> Pattern:
    pat = Pattern()
    pat.rect((1, 0), xmin=0, xmax=length, yctr=0, ly=width)
    pat.ports["in"] = Port((0, 0), 0, ptype=ptype)
    pat.ports["out"] = Port((length, 0), pi, ptype=ptype)
    return pat


@pytest.fixture
def autotool_setup() -> tuple[Pather, AutoTool, Library]:
    lib = Library()

    bend_pat = Pattern()
    bend_pat.ports["in"] = Port((0, 0), 0, ptype="wire")
    bend_pat.ports["out"] = Port((2, -2), pi / 2, ptype="wire")
    lib["bend"] = bend_pat
    lib.abstract("bend")

    via_pat = Pattern()
    via_pat.ports["m1"] = Port((0, 0), 0, ptype="wire_m1")
    via_pat.ports["m2"] = Port((1, 0), pi, ptype="wire_m2")
    lib["via"] = via_pat
    via_abs = lib.abstract("via")

    tool_m1 = AutoTool(
        straights=[
            AutoTool.Straight(ptype="wire_m1", fn=lambda length: _make_transition_straight(length, ptype="wire_m1"), in_port_name="in", out_port_name="out")
        ],
        bends=[],
        sbends=[],
        transitions={("wire_m2", "wire_m1"): AutoTool.Transition(via_abs, "m2", "m1")},
        default_out_ptype="wire_m1",
    )

    p = Pather(lib, tools=tool_m1)
    p.ports["start"] = Port((0, 0), pi, ptype="wire_m2")

    return p, tool_m1, lib

def test_autotool_transition(autotool_setup: tuple[Pather, AutoTool, Library]) -> None:
    p, _tool, _lib = autotool_setup

    p.straight("start", 10)

    # Via length is 1, so the remaining wire_m1 straight length is 9.
    assert_allclose(p.ports["start"].offset, [10, 0], atol=1e-10)
    assert p.ports["start"].ptype == "wire_m1"

def make_straight(length, width=2, ptype="wire"):
    pat = Pattern()
    pat.rect((1, 0), xmin=0, xmax=length, yctr=0, ly=width)
    pat.ports["A"] = Port((0, 0), 0, ptype=ptype)
    pat.ports["B"] = Port((length, 0), pi, ptype=ptype)
    return pat

def make_bend(R, width=2, ptype="wire", clockwise=True):
    pat = Pattern()
    # Rectangular approximation of a 90 degree bend.
    if clockwise:
        pat.rect((1, 0), xmin=0, xmax=R, yctr=0, ly=width)
        pat.rect((1, 0), xctr=R, lx=width, ymin=-R, ymax=0)
        pat.ports["A"] = Port((0, 0), 0, ptype=ptype)
        pat.ports["B"] = Port((R, -R), pi/2, ptype=ptype)
    else:
        pat.rect((1, 0), xmin=0, xmax=R, yctr=0, ly=width)
        pat.rect((1, 0), xctr=R, lx=width, ymin=0, ymax=R)
        pat.ports["A"] = Port((0, 0), 0, ptype=ptype)
        pat.ports["B"] = Port((R, R), -pi/2, ptype=ptype)
    return pat

@pytest.fixture
def multi_bend_tool():
    lib = Library()

    lib["b1"] = make_bend(2, ptype="wire")
    b1_abs = lib.abstract("b1")
    lib["b2"] = make_bend(5, ptype="wire")
    b2_abs = lib.abstract("b2")

    tool = AutoTool(
        straights=[
            AutoTool.Straight(ptype="wire", fn=make_straight, in_port_name="A", out_port_name="B", length_range=(0, 10)),
            AutoTool.Straight(ptype="wire", fn=lambda l: make_straight(l, width=4), in_port_name="A", out_port_name="B", length_range=(10, 1e8))
        ],
        bends=[
            AutoTool.Bend(b1_abs, "A", "B", clockwise=True, mirror=True),
            AutoTool.Bend(b2_abs, "A", "B", clockwise=True, mirror=True)
        ],
        sbends=[],
        transitions={},
        default_out_ptype="wire"
    )
    return tool, lib

@pytest.fixture
def asymmetric_transition_tool() -> AutoTool:
    lib = Library()

    bend_pat = Pattern()
    bend_pat.ports["in"] = Port((0, 0), 0, ptype="core")
    bend_pat.ports["out"] = Port((2, -2), pi / 2, ptype="core")
    lib["core_bend"] = bend_pat

    trans_pat = Pattern()
    trans_pat.ports["CORE"] = Port((0, 0), 0, ptype="core")
    trans_pat.ports["MID"] = Port((3, 1), pi, ptype="mid")
    lib["core_mid"] = trans_pat

    return AutoTool(
        straights=[
            AutoTool.Straight(
                ptype="core",
                fn=lambda length: make_straight(length, ptype="core"),
                in_port_name="A",
                out_port_name="B",
                length_range=(0, 3),
                ),
            AutoTool.Straight(
                ptype="mid",
                fn=lambda length: make_straight(length, ptype="mid"),
                in_port_name="A",
                out_port_name="B",
                length_range=(0, 1e8),
                ),
            ],
        bends=[
            AutoTool.Bend(lib.abstract("core_bend"), "in", "out", clockwise=True, mirror=True),
            ],
        sbends=[],
        transitions={
            ("mid", "core"): AutoTool.Transition(lib.abstract("core_mid"), "MID", "CORE"),
            },
        default_out_ptype="core",
        ).add_complementary_transitions()

def assert_trace_matches_plan(plan_port: Port, tree: Library, port_names: tuple[str, str] = ("A", "B")) -> None:
    pat = tree.top_pattern()
    out_port = pat[port_names[1]]
    dxy, rot = pat[port_names[0]].measure_travel(out_port)
    assert_allclose(dxy, plan_port.offset)
    assert rot is not None
    assert plan_port.rotation is not None
    assert_allclose(rot, plan_port.rotation)
    assert out_port.ptype == plan_port.ptype

def test_autotool_planL_selection(multi_bend_tool) -> None:
    tool, _ = multi_bend_tool

    p, data = tool.planL(True, 5)
    assert data.straight.length_range == (0, 10)
    assert data.straight_length == 3
    assert data.bend.abstract.name == "b1"
    assert_allclose(p.offset, [5, 2])

    p, data = tool.planL(True, 15)
    assert data.straight.length_range == (10, 1e8)
    assert data.straight_length == 13
    assert_allclose(p.offset, [15, 2])


@pytest.mark.parametrize("ccw", [False, True])
def test_autotool_traceL_matches_plan_with_post_bend_transition(ccw: bool) -> None:
    lib = Library()

    bend_pat = Pattern()
    bend_pat.ports["A"] = Port((0, 0), 0, ptype="core")
    bend_pat.ports["B"] = Port((2, -2), pi / 2, ptype="core")
    lib["core_bend"] = bend_pat

    trans_pat = Pattern()
    trans_pat.ports["CORE"] = Port((0, 0), 0, ptype="core")
    trans_pat.ports["EXT"] = Port((3, 1), pi, ptype="ext")
    lib["out_trans"] = trans_pat

    tool = AutoTool(
        straights=[
            AutoTool.Straight(
                ptype="core",
                fn=lambda length: make_straight(length, ptype="core"),
                in_port_name="A",
                out_port_name="B",
                length_range=(0, 1e8),
                ),
            ],
        bends=[
            AutoTool.Bend(lib.abstract("core_bend"), "A", "B", clockwise=True, mirror=True),
            ],
        sbends=[],
        transitions={
            ("ext", "core"): AutoTool.Transition(lib.abstract("out_trans"), "EXT", "CORE"),
            },
        default_out_ptype="core",
        )

    plan_port, data = tool.planL(ccw, 10, out_ptype="ext")

    assert data.out_transition is not None

    tree = tool.traceL(ccw, 10, out_ptype="ext")
    assert_trace_matches_plan(plan_port, tree)


def test_autotool_planU_consistency(multi_bend_tool) -> None:
    tool, lib = multi_bend_tool

    p, data = tool.planU(20, length=10)
    assert data.ldata0.straight_length == 7
    assert data.ldata0.bend.abstract.name == "b2"
    assert data.l2_length == 13
    assert data.ldata1.straight_length == 0
    assert data.ldata1.bend.abstract.name == "b1"

def test_autotool_traceU_matches_plan_with_asymmetric_transition(asymmetric_transition_tool: AutoTool) -> None:
    tool = asymmetric_transition_tool

    plan_port, data = tool.planU(12, length=0, in_ptype="core")

    assert data.ldata1.in_transition is not None
    assert data.ldata1.b_transition is not None

    tree = tool.traceU(12, length=0, in_ptype="core")
    assert_trace_matches_plan(plan_port, tree)

def test_autotool_planS_double_L(multi_bend_tool) -> None:
    tool, lib = multi_bend_tool

    p, data = tool.planS(20, 10)
    assert_allclose(p.offset, [20, 10])
    assert_allclose(p.rotation, pi)

    assert data.ldata0.straight_length == 16
    assert data.ldata1.straight_length == 0
    assert data.l2_length == 6

def test_autotool_traceS_double_l_matches_plan_with_asymmetric_transition(asymmetric_transition_tool: AutoTool) -> None:
    tool = asymmetric_transition_tool

    plan_port, data = tool.planS(4, 10, in_ptype="core")

    assert isinstance(data, AutoTool.UData)
    assert data.ldata1.in_transition is not None
    assert data.ldata1.b_transition is not None

    tree = tool.traceS(4, 10, in_ptype="core")
    assert_trace_matches_plan(plan_port, tree)

def test_autotool_planS_pure_sbend_with_transition_dx() -> None:
    lib = Library()

    def make_straight(length: float) -> Pattern:
        pat = Pattern()
        pat.ports["A"] = Port((0, 0), 0, ptype="core")
        pat.ports["B"] = Port((length, 0), pi, ptype="core")
        return pat

    def make_sbend(jog: float) -> Pattern:
        pat = Pattern()
        pat.ports["A"] = Port((0, 0), 0, ptype="core")
        pat.ports["B"] = Port((10, jog), pi, ptype="core")
        return pat

    trans_pat = Pattern()
    trans_pat.ports["EXT"] = Port((0, 0), 0, ptype="ext")
    trans_pat.ports["CORE"] = Port((5, 0), pi, ptype="core")
    lib["xin"] = trans_pat

    tool = AutoTool(
        straights=[
            AutoTool.Straight(
                ptype="core",
                fn=make_straight,
                in_port_name="A",
                out_port_name="B",
                length_range=(1, 1e8),
                )
            ],
        bends=[],
        sbends=[
            AutoTool.SBend(
                ptype="core",
                fn=make_sbend,
                in_port_name="A",
                out_port_name="B",
                jog_range=(0, 1e8),
                )
            ],
        transitions={
            ("ext", "core"): AutoTool.Transition(lib.abstract("xin"), "EXT", "CORE"),
            },
        default_out_ptype="core",
        )

    p, data = tool.planS(15, 4, in_ptype="ext")

    assert_allclose(p.offset, [15, 4])
    assert_allclose(p.rotation, pi)
    assert data.straight_length == 0
    assert data.jog_remaining == 4
    assert data.in_transition is not None