import os
from pathlib import Path
from typing import Any, Generator, Iterator


from . import lattice
from . import elements


def flatten(iterable: Iterator[Any]) -> Generator[Any, None, None]:
    """Flatten arbitrarily nested iterable.
    Special case for strings that avoids infinite recursion.  Non iterables passed
    as arguments are yielded.

    :param iterable: Any iterable to be flattened.
    :raises: RecursionError

    """

    def _flatten(iterable):
        try:
            for item in iterable:
                try:
                    iter(item)
                except TypeError:
                    yield item
                else:
                    yield from _flatten(item)
        except TypeError:  # If iterable isn't actually iterable, then yield it.
            yield iterable

    try:
        yield from _flatten(iterable)
    except RecursionError:
        raise RecursionError(
            "Maximum recusion reached.  Possibly trying"
            " to flatten an infinitely nested iterable."
        )


class UnknownBeamlineDescription(RuntimeError):
    pass


class UnknownElementType(RuntimeError):
    pass


def from_ocelot(ocelot_lattice: Any) -> lattice.Beamline:
    seq = ocelot_lattice

    try:
        seq = ocelot_lattice.sequence
    except AttributeError:
        pass

    beamline = lattice.Beamline(_loop_lattice_from_ocelot(seq))
    return beamline


def coerce_beamline(some_beamline: Any) -> lattice.Beamline:
    if isinstance(some_beamline, lattice.Beamline):
        return some_beamline
    try:
        some_beamline = Path(some_beamline)
    except TypeError:
        pass

    try:
        import ocelot
    except ImportError:
        pass
    else:
        return from_ocelot(some_beamline)

    raise UnknownBeamlineDescription


def _loop_lattice_from_ocelot(
    ocelot_sequence: Any,
) -> Generator[elements.Element, None, None]:
    assert not isinstance(ocelot_sequence, str)

    try:
        from ocelot.cpbd import elements as ole
    except ImportError:
        pass

    for ele in flatten(ocelot_sequence):
        name = ele.id
        l = ele.l
        if isinstance(ele, ole.Marker):
            yield elements.Marker(name)
        elif isinstance(ele, ole.Monitor):
            yield elements.Monitor(name)
        elif isinstance(ele, ole.Drift):
            yield elements.Drift(name, l)
        elif isinstance(ele, ole.RBend):
            yield elements.RBend(name, l, ele.angle)
        elif isinstance(ele, ole.SBend):
            yield elements.SBend(name, l, ele.angle)
        elif isinstance(ele, ole.Quadrupole):
            yield elements.Quadrupole(name, l, ele.k1)
        elif isinstance(ele, ole.Sextupole):
            yield elements.Sextupole(name, l, ele.k2)
        elif isinstance(ele, ole.TDCavity):
            yield elements.TransverseDeflectingCavity(name, l, ele.v)
        elif isinstance(ele, ole.Vcor):
            yield elements.VKicker(name, l, ele.angle)
        elif isinstance(ele, ole.Hcor):
            yield elements.HKicker(name, l, ele.angle)
        elif isinstance(ele, ole.Cavity):
            yield elements.RFCavity(name, l, ele.v, ele.phi)
        elif isinstance(ele, ole.Solenoid):
            yield elements.Solenoid(name, l, ele.k)
        elif isinstance(ele, ole.Undulator):
            undulator_length = ele.nperiods * ele.lperiod
            yield elements.Undulator(name, undulator_length, kx=ele.Kx, ky=ele.Ky)
        elif isinstance(ele, ole.Octupole):
            yield elements.Octupole(name, l, ele.k3)
        else:
            raise UnknownElementType(name, type(ele))
