from difflib import unified_diff

import pytest

from finesse import Model
from finesse.components.replace import MissingReplacementError


def compare_kat(a: str, b: str):
    if a != b:
        raise AssertionError("\n".join(unified_diff(a.splitlines(), b.splitlines())))


def test_explicit_space():
    # CASE 1: simple linear model using explicit spaces

    model = Model(
        """\
l l1
s s1 l1.p1 m1.p1
m m1
s s2 m1.p2 m2.p1
m m2"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        optical_ports=["m_miron.p1", "m_miron.p2"],
        verbose=False,
    )
    expected = """\
l l1
s s1 l1.p1 m_miron.p1
m m_miron R=0.5 T=0.5 L=0.0
s s2 m_miron.p2 m2.p1
m m2 R=0.5 T=0.5 L=0.0"""

    compare_kat(new_kat, expected)


def test_explicit__leading_space():
    # CASE 1: simple linear model using explicit spaces

    model = Model(
        """\
l l1
s s1 l1.p1 m1.p1
    m m1
s s2 m1.p2 m2.p1
m m2"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        optical_ports=["m_miron.p1", "m_miron.p2"],
        verbose=False,
    )
    expected = """\
l l1
s s1 l1.p1 m_miron.p1
m m_miron R=0.5 T=0.5 L=0.0
s s2 m_miron.p2 m2.p1
m m2 R=0.5 T=0.5 L=0.0"""

    compare_kat(new_kat, expected)


def test_optical_ports_too_short():
    # CASE 1: simple linear model using explicit spaces

    model = Model(
        """\
l l1
s s1 l1.p1 m1.p1
m m1
s s2 m1.p2 m2.p1
m m2"""
    )

    with pytest.raises(MissingReplacementError):
        model.m1.replace(
            "m m_miron R=0.5 T=0.5 L=0.0",
            optical_ports=[
                "m_miron.p1",
            ],
            verbose=False,
        )


def test_optical_ports_none():
    # CASE 1: simple linear model using explicit spaces

    model = Model(
        """\
l l1
s s1 l1.p1 m1.p1
m m1
s s2 m1.p2 m2.p1
m m2"""
    )

    with pytest.raises(MissingReplacementError):
        model.m1.replace(
            "m m_miron R=0.5 T=0.5 L=0.0",
            optical_ports=None,
            verbose=False,
        )


def test_optical_ports_none_allowed():
    model = Model("m m1 R=0.5 T=0.5 L=0.0")
    new_kat = model.m1.replace("m m_miron R=0.5 T=0.5 L=0.0", verbose=False)
    assert new_kat == "m m_miron R=0.5 T=0.5 L=0.0"


# CASE 2: using mirror aliases
def test_aliases():
    model = Model(
        """\
l l1
s s1 l1.p1 m1.p1
mirror m1 R=0.5 T=0.5 L=0.0
s s2 m1.p2 m2.p1
m m2 R=0.5 T=0.5 L=0.0"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        optical_ports=["m_miron.p1", "m_miron.p2"],
        verbose=False,
    )
    expected = """\
l l1
s s1 l1.p1 m_miron.p1
m m_miron R=0.5 T=0.5 L=0.0
s s2 m_miron.p2 m2.p1
m m2 R=0.5 T=0.5 L=0.0"""

    compare_kat(new_kat, expected)


# CASE 3: using arbitrary whitespace in katscript (seems to be removed by the unparsers?)
def test_whitespace():
    model = Model(
        """\
l l1
s s1 l1.p1 m1.p1
mirror           m1 R=0.5 T=0.5 L=0.0
s s2 m1.p2 m2.p1
m m2 R=0.5 T=0.5 L=0.0"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        optical_ports=["m_miron.p1", "m_miron.p2"],
        verbose=False,
    )
    expected = """\
l l1
s s1 l1.p1 m_miron.p1
m m_miron R=0.5 T=0.5 L=0.0
s s2 m_miron.p2 m2.p1
m m2 R=0.5 T=0.5 L=0.0"""

    compare_kat(new_kat, expected)


# CASE 4: replacing an end mirror
def test_single_connection():
    model = Model(
        """\
l l1
s s1 l1.p1 m1.p1
mirror m1 R=0.5 T=0.5 L=0.0"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        optical_ports=["m_miron.p1", "m_miron.p2"],
        verbose=False,
    )
    expected = """\
l l1
s s1 l1.p1 m_miron.p1
m m_miron R=0.5 T=0.5 L=0.0"""

    compare_kat(new_kat, expected)


# CASE 5: using the 'link' command with explicit ports
# unparsers turns links into spaces
def test_link_explicit():
    model = Model(
        """\
l l1
link(l1.p1, m1.p1)
mirror m1 R=0.5 T=0.5 L=0.0
link(m1.p2, m2.p1)
m m2 R=0.5 T=0.5 L=0.0"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        optical_ports=["m_miron.p1", "m_miron.p2"],
        verbose=False,
    )
    expected = """\
l l1

m m_miron R=0.5 T=0.5 L=0.0

m m2 R=0.5 T=0.5 L=0.0

space l1_p1__m1_p1 portA=l1.p1 portB=m_miron.p1
space m1_p2__m2_p1 portA=m_miron.p2 portB=m2.p1"""

    compare_kat(new_kat, expected)


# CASE 6: using the 'link' command without explicit ports
def test_link_implicit():
    model = Model(
        """\
l l1
link(l1, m1)
mirror m1 R=0.5 T=0.5 L=0.0
link(m1, m2)
m m2 R=0.5 T=0.5 L=0.0"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        optical_ports=["m_miron.p1", "m_miron.p2"],
        verbose=False,
    )
    expected = """\
l l1

m m_miron R=0.5 T=0.5 L=0.0

m m2 R=0.5 T=0.5 L=0.0

space l1_p1__m1_p1 portA=l1.p1 portB=m_miron.p1
space m1_p2__m2_p1 portA=m_miron.p2 portB=m2.p1"""

    compare_kat(new_kat, expected)


def test_detectors():
    model = Model(
        """\
m m1 R=0.5 T=0.5 L=0.0
ad amp1 m1.p1.i f=0
ad amp2 m1.p1.o f=0
pd pow1 m1.p2.i
pd pow2 m1.p2.o
"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        optical_ports=["m_miron.p1", "m_miron.p2"],
        verbose=False,
    )
    expected = """\
m m_miron R=0.5 T=0.5 L=0.0
ad amp1 m_miron.p1.i f=0
ad amp2 m_miron.p1.o f=0
pd pow1 m_miron.p2.i
pd pow2 m_miron.p2.o"""

    compare_kat(new_kat, expected)


def test_component():
    model = Model(
        """\
m m1 R=0.5 T=0.5 L=0.0
pendulum pend2 m1"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        component="m_miron",
        optical_ports=None,
        verbose=False,
    )

    expected = """\
m m_miron R=0.5 T=0.5 L=0.0
pendulum pend2 m_miron"""

    compare_kat(new_kat, expected)


def test_component_spaces():
    model = Model(
        """\
m m1 R=0.5 T=0.5 L=0.0
pendulum pend2 m1 mass=4"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        component="m_miron",
        optical_ports=None,
        verbose=False,
    )

    expected = """\
m m_miron R=0.5 T=0.5 L=0.0
pendulum pend2 m_miron mass=4"""

    compare_kat(new_kat, expected)


def test_component_comment():
    model = Model(
        """\
m m1 R=0.5 T=0.5 L=0.0
pendulum pend2 m1#comment"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        component="m_miron",
        optical_ports=None,
        verbose=False,
    )

    expected = """\
m m_miron R=0.5 T=0.5 L=0.0
pendulum pend2 m_miron#comment"""

    compare_kat(new_kat, expected)


def test_component_keyword():
    model = Model(
        """\
m m1 R=0.5 T=0.5 L=0.0
pendulum pend2 connected_to=m1 mass=4"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        component="m_miron",
        optical_ports=None,
        verbose=False,
    )

    expected = """\
m m_miron R=0.5 T=0.5 L=0.0
pendulum pend2 connected_to=m_miron mass=4"""

    compare_kat(new_kat, expected)


def test_component_keyword_comment():
    model = Model(
        """\
m m1 R=0.5 T=0.5 L=0.0
pendulum pend2 connected_to=m1#comment"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        component="m_miron",
        optical_ports=None,
        verbose=False,
    )

    expected = """\
m m_miron R=0.5 T=0.5 L=0.0
pendulum pend2 connected_to=m_miron#comment"""

    compare_kat(new_kat, expected)


def test_mechanical_ports():
    model = Model(
        """\
l l1
s s1 l1.p1 m1.p1
m m1
s s2 m1.p2 m2.p1
m m2
ad motion m1.mech.z f=0
"""
    )

    new_kat = model.m1.replace(
        "m m_miron R=0.5 T=0.5 L=0.0",
        optical_ports=["m_miron.p1", "m_miron.p2"],
        mechanical_ports=["m_miron.mech"],
        verbose=False,
    )

    expected = """\
l l1
s s1 l1.p1 m_miron.p1
m m_miron R=0.5 T=0.5 L=0.0
s s2 m_miron.p2 m2.p1
m m2 R=0.5 T=0.5 L=0.0
ad motion m_miron.mech.z f=0"""

    compare_kat(new_kat, expected)


def test_mechanical_ports_missing():
    model = Model(
        """\
l l1
s s1 l1.p1 m1.p1
m m1
s s2 m1.p2 m2.p1
m m2
ad motion m1.mech.z f=0
"""
    )

    with pytest.raises(MissingReplacementError):
        model.m1.replace(
            "m m_miron R=0.5 T=0.5 L=0.0",
            optical_ports=["m_miron.p1", "m_miron.p2"],
            mechanical_ports=None,
            verbose=False,
        )


def test_complicated():
    model = Model(
        """\
l l1

s LA l1.p1 PLACEHOLDER.p1
nothing PLACEHOLDER
s LB PLACEHOLDER.p2 central_bs.p1

bs central_bs R=0.5 T=0.5
s LW central_bs.p2 WE.p1 L=100
m WE R=0.9 T=0.1
s LN central_bs.p3 NE.p1 L=100
m NE R=0.9 T=0.1
pd pd1 central_bs.p4.o"""
    )

    IMC_KATSCRIPT = """\
bs bs1 R=0.9 T=0.1
s s1 bs1.p2 bs2.p1
bs bs2 R=0.9 T=0.1
s s2 bs2.p2 bs3.p2
bs bs3 R=0.9 T=0.1
s s3 bs1.p3 bs3.p1
"""

    new_kat = model.PLACEHOLDER.replace(
        IMC_KATSCRIPT, optical_ports=["bs1.p1", "bs3.p3"]
    )

    expected = """\
l l1

s LA l1.p1 bs1.p1
bs bs1 R=0.9 T=0.1
s s1 bs1.p2 bs2.p1
bs bs2 R=0.9 T=0.1
s s2 bs2.p2 bs3.p2
bs bs3 R=0.9 T=0.1
s s3 bs1.p3 bs3.p1
s LB bs3.p3 central_bs.p1

bs central_bs R=0.5 T=0.5 L=(1-(central_bs.R+central_bs.T))
s LW central_bs.p2 WE.p1 L=100
m WE R=0.9 T=0.1 L=(1-(WE.R+WE.T))
s LN central_bs.p3 NE.p1 L=100
m NE R=0.9 T=0.1 L=(1-(NE.R+NE.T))
pd pd1 central_bs.p4.o"""

    compare_kat(new_kat, expected)
