import pathlib
import pprint
from typing import List, Union

import lark
import pytest

from ...version4 import (
    Drift,
    DuplicatedLineItem,
    LatticeSequenceConst,
    LatticeSequencePolynom,
    LatticeSequencePower,
    LatticeSequenceRandom,
    PositionedLineItem,
    Genesis4,
    Lattice,
    Line,
    Quadrupole,
    Undulator,
)
from ...version4.input.parsers import new_lattice_parser
from ...version4.types import BeamlineElement

test_v4_root = pathlib.Path(__file__).resolve().parent
genesis4_lattices = tuple(test_v4_root.glob("*.lat"))


@pytest.fixture(scope="module")
def lattice_parser() -> lark.Lark:
    return new_lattice_parser()


element_sources = pytest.mark.parametrize(
    ("element_source", "expected_elements"),
    [
        pytest.param("D1: DRIFT = { l = 0.44};", Drift(label="D1", L=0.44)),
        pytest.param("D2: DRIFT = { l = 0.24};", Drift(label="D2", L=0.24)),
        pytest.param(
            "QF: QUADRUPOLE = { l = 0.080000, k1= 2.000000 };",
            Quadrupole(label="QF", L=0.080000, k1=2.000000),
        ),
        pytest.param(
            "QD: QUADRUPOLE = { l = 0.080000, k1= -2.000000 };",
            Quadrupole(label="QD", L=0.080000, k1=-2.000000),
        ),
        pytest.param(
            "QD: QUAD= { l = 0.080000, k1= -2.000000 };",
            Quadrupole(label="QD", L=0.080000, k1=-2.000000),
        ),
        pytest.param(
            "UND: UNDULATOR = { lambdau=0.015000, nwig=266, aw=0.84853, helical= True};",
            Undulator(
                label="UND", lambdau=0.015000, nwig=266, aw=0.84853, helical=True
            ),
        ),
        pytest.param(
            "FODO: LINE={UND,D1,QF,D2,UND,D1,QD,D2};",
            Line(
                label="FODO",
                elements=["UND", "D1", "QF", "D2", "UND", "D1", "QD", "D2"],
            ),
        ),
        pytest.param(
            "FEL: LINE={6*FODO};",
            Line(label="FEL", elements=[DuplicatedLineItem(label="FODO", count=6)]),
        ),
        pytest.param(
            "FODO: Line = {UND,F@0.2,D@1.2,M@2.0}; ",
            Line(
                label="FODO",
                elements=[
                    "UND",
                    PositionedLineItem(label="F", position=0.2),
                    PositionedLineItem(label="D", position=1.2),
                    PositionedLineItem(label="M", position=2.0),
                ],
            ),
        ),
        pytest.param(
            "Line1: LINE= {UND}; Line2: Line = {Line1};",
            [
                Line(label="Line1", elements=["UND"]),
                Line(label="Line2", elements=["Line1"]),
            ],
        ),
    ],
)


@element_sources
def test_element_parsing(
    lattice_parser: lark.Lark,
    element_source: str,
    expected_elements: Union[BeamlineElement, List[BeamlineElement]],
) -> None:
    print(lattice_parser.parse(element_source))
    assert BeamlineElement.from_contents(element_source) == expected_elements


@element_sources
def test_elements_roundtrip(
    element_source: str,
    expected_elements: Union[BeamlineElement, List[BeamlineElement]],
) -> None:
    if not isinstance(expected_elements, list):
        expected_elements = [expected_elements]

    for ele in expected_elements:
        assert BeamlineElement.from_contents(ele.to_genesis()) == ele


sequence_types = pytest.mark.parametrize(
    ("element_source", "expected_element"),
    [
        pytest.param(
            "VAL: SEQUENCE = {type = const, c0 = 3.5};",
            LatticeSequenceConst(label="VAL", c0=3.5),
        ),
        pytest.param(
            "VAL: SEQUENCE = {type = polynom, c0 = 3.5, c1=0.0, c2=0.0, c3=0.0, c4=0.0};",
            LatticeSequencePolynom(label="VAL", c0=3.5, c1=0.0, c2=0.0, c3=0.0, c4=0.0),
        ),
        pytest.param(
            "VAL: SEQUENCE = {type = power, c0 = 3.5};",
            LatticeSequencePower(label="VAL", c0=3.5),
        ),
        pytest.param(
            "VAL: SEQUENCE = {type = random, c0 = 3.5, dc=0.001};",
            LatticeSequenceRandom(label="VAL", c0=3.5, dc=0.001),
        ),
    ],
)


@sequence_types
def test_sequence_types_parse(
    lattice_parser: lark.Lark, element_source: str, expected_element: BeamlineElement
) -> None:
    print(lattice_parser.parse(element_source))
    assert BeamlineElement.from_contents(element_source) == expected_element


@sequence_types
def test_sequence_types_round_trip(
    element_source: str, expected_element: BeamlineElement
) -> None:
    assert (
        BeamlineElement.from_contents(expected_element.to_genesis()) == expected_element
    )


@pytest.mark.parametrize(
    "filename", [pytest.param(fn, id=fn.name) for fn in genesis4_lattices]
)
def test_load_file(lattice_parser: lark.Lark, filename: pathlib.Path) -> None:
    with open(filename) as fp:
        contents = fp.read()
    print(f"\n\nLattice ({filename})")
    print(contents)

    print("\n\nAs a lark Tree:")
    print(lattice_parser.parse(contents))

    print("\n\nAs dataclasses:")
    lattice = Lattice.from_file(filename)
    pprint.pprint(lattice, width=1)

    print("\n\nSerialized back to a file:")
    round_tripped = lattice.to_genesis()
    print(round_tripped)

    print("\n\nChecking file output vs initial dataclasses..")
    second_lattice = Lattice.from_contents(round_tripped)
    assert lattice.to_genesis() == second_lattice.to_genesis()


@pytest.mark.parametrize(
    "filename", [pytest.param(fn, id=fn.name) for fn in genesis4_lattices[:1]]
)
def test_load_file_with_input_parser_helper(filename: pathlib.Path) -> None:
    assert isinstance(Genesis4.lattice_parser(filename), Lattice)


@pytest.mark.parametrize(
    "filename", [pytest.param(fn, id=fn.name) for fn in genesis4_lattices]
)
def test_serialize(filename: pathlib.Path) -> None:
    lattice = Lattice.from_file(filename)

    print("This lattice:")
    pprint.pprint(lattice)

    print("Is serialized as follows:")
    serialized = lattice.model_dump()
    pprint.pprint(serialized)

    print("Deserialized back to dataclasses:")
    deserialized = Lattice.model_validate(serialized)
    assert lattice == deserialized
