import logging
import os
import shutil
import sys
from functools import partial
from pathlib import Path
from typing import Any, Generator

import matplotlib.pyplot as plt
import numpy as np
import pytest
from component_model.model import Model
from component_model.utils.controls import Control, Controls

from py_crane.animation import AnimateCrane
from py_crane.boom import Boom, Wire
from py_crane.crane import Crane

sys.path.insert(0, str(Path(__file__).parent.parent / "examples"))

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)  # DEBUG)


def do_plot(time: list[float], traces: tuple[str, ...], data: dict[str, list[float]], title: str = "CraneOnSpring"):
    """Plot selected 'traces' from the 'data' repository, with 'title'."""
    fig, ax = plt.subplots()
    for k, v in data.items():
        if k in traces:
            ax.plot(time, v, label=k)
    ax.legend()
    plt.title(title)
    plt.show()


@pytest.fixture(scope="session")
def mobile_crane_fmu():
    return _get_fmu("MobileCrane.fmu")


def make_fmu(build_path: Path, source: Path, resource: Path, newargs: dict[str, Any] | None = None):
    """Make single FMU. Non-default arguments allowed."""
    build_path.mkdir(exist_ok=True)
    fmu = Model.build(
        str(source),
        project_files=[resource],
        dest=build_path,
        newargs=newargs,
    )
    return fmu


def make_fmus():
    """Re-make the FMUs. Only default arguments used here."""
    make_fmu(
        Path(__file__).parent.parent / "examples",
        Path(__file__).parent.parent / "src" / "py_crane" / "mobile_crane.py",
        Path(__file__).parent.parent / "src" / "py_crane",
    )
    make_fmu(
        Path(__file__).parent.parent / "examples",
        Path(__file__).parent.parent.parent / "component-model" / "examples" / "oscillator_6dof_fmu.py",
        Path(__file__).parent.parent.parent / "component-model" / "src" / "component_model",
    )
    shutil.copy(
        Path(__file__).parent.parent.parent / "component-model" / "examples" / "HarmonicOscillator6D.fmu",
        Path(__file__).parent.parent / "examples",
    )


def make_mobile_crane_straight():
    fmu = make_fmu(
        Path(__file__).parent.parent / "examples",
        Path(__file__).parent.parent / "src" / "py_crane" / "mobile_crane.py",
        Path(__file__).parent.parent / "src" / "py_crane",
        newargs={"pedestalCoM": (0.5, 0, 0), "boomAngle": "0 deg", "wire_length": 5},
    )
    shutil.copy(fmu, fmu.parent / "MobileCraneStraight.fmu")


@pytest.fixture(scope="session")
def oscillator_fmu():
    return _get_fmu("HarmonicOscillator6D.fmu")


def _get_fmu(fmu_file: str) -> Path:
    # First assume that the FMU is in the current working directory, which is expected to be the 'test_working_directory' folder.
    # If not found there, look in the 'examples' folder.
    fmu = Path.cwd() / fmu_file
    if not fmu.exists():
        fmu = Path(__file__).parent.parent / "examples" / fmu_file
    assert fmu.exists(), f"{fmu_file} file expected at {fmu}. Not found."
    return fmu


def _boom(b: Boom, idx: int, x: float | None = None) -> float:
    """Move boom 'b', or return current setting. Used as partial() to define Control."""
    if x is not None:
        arg = [None if i != idx else x for i in range(3)]
        b.boom_setter(arg)
    return b.boom[idx]


# def test_visual_simulation_1():
#     sim = VisualSimulator()
#     sim.start(
#         points_3d=[
#             (
#                 ("mobileCrane", "pedestal_end[0]"),
#                 ("mobileCrane", "pedestal_end[1]"),
#                 ("mobileCrane", "pedestal_end[2]"),
#             ),
#             (
#                 ("mobileCrane", "boom_end[0]"),
#                 ("mobileCrane", "boom_end[1]"),
#                 ("mobileCrane", "boom_end[2]"),
#             ),
#             (
#                 ("mobileCrane", "wire_end[0]"),
#                 ("mobileCrane", "wire_end[1]"),
#                 ("mobileCrane", "wire_end[2]"),
#             ),
#         ],
#         osp_system_structure="OspSystemStructure.xml",
#     )


#    def test_visual_simulation_2():
#     simulator = VisualSimulator()
#     sim.start(
#         points_3d=[
#             (
#                 ("mobileCrane", "pedestal_end[0]"),
#                 ("mobileCrane", "pedestal_end[1]"),
#                 ("mobileCrane", "pedestal_end[2]"),
#             ),
#             (
#                 ("mobileCrane", "wire_end[0]"),
#                 ("mobileCrane", "wire_end[1]"),
#                 ("mobileCrane", "wire_end[2]"),
#             ),
#         ],
#         osp_system_structure="OspSystemStructure.xml",
#     )


def ensure_subpath(pkg: str, folder: str) -> bool:
    """Ensure that the path 'pkg/folder' is in sys.path.
    It is expected that pkg is loaded, such that its path is known
    and that 'folder' is a sub-path below pkg, like 'examples'."""
    pkg_path = ""
    for p in sys.path:
        if p.endswith(os.path.join(pkg, folder)):
            return True
        elif p.endswith(pkg):
            pkg_path = p
    if pkg_path == "":
        # raise ValueError(f"Package {pkg} is not loaded and we cannot include a sub-folder in sys.path") from None
        return False
    sys.path.insert(0, os.path.join(pkg_path, folder))
    return True


component_model_examples_loaded: bool = ensure_subpath("component_model", "examples")


@pytest.mark.skipif(not component_model_examples_loaded, reason="component_model examples not loaded")
def make_crane_on_spring(
    pM: float = 10000.0,  # in kg
    pCoM: tuple[float, ...] = (0.5, -1.0, 0.8),  # in m
    pH: float = 3.0,  # in m
    bM: float = 1000.0,  # in kg
    bL: float = 8.0,  # in m
    bA: float = np.radians(90.0),  # in radians
    wM: float = 50.0,  # in kg
    wL: float = 1e-6,
    wQ: float = 50.0,  # dimensionless quality factor
    k: tuple[float, ...] = (1e4,) * 6,
    c: tuple[float, ...] = (0,) * 6,
    m: float = 1e4,
    x0: tuple[float, ...] = (0.0,) * 6,
    v0: tuple[float, ...] = (0.0,) * 6,
):
    """Initialize and return crane, force and oscillator."""

    def calc_force(t: float, x: np.ndarray, v: np.ndarray, crane: Any, dt: float | None = None):
        """Connect the crane to the oscillator.
        Transfer 6D position and 6D speed of oscillator to crane, step crane and return updated force and torque.
        """
        crane.position, crane.angular = np.split(x, 2)  # x[:3], x[3:]
        crane.velocity, crane.d_angular = np.split(v, 2)  # v[:3], v[3:]
        crane.do_step(t, dt)  # calc_statics_dynamics( dt)
        # print(f"Crane z:{x[2]}, v:{v[2]}, a:{crane.boom0.acceleration[2]} => f:{crane.force[2]}")
        return np.append(crane.force, crane.torque)

    crane = Crane()
    _ = crane.add_boom(
        "pedestal",
        description="The crane base, on one side fixed to the vessel and on the other side the first crane boom is fixed to it. The mass should include all additional items fixed to it, like the operator's cab",
        mass=pM,
        mass_center=pCoM,
        boom=(pH, 0, 0),
    )
    _ = crane.add_boom(
        "boom",
        description="The first boom. Can be lifted",
        mass=bM,
        mass_center=0.5,
        boom=(bL, bA, 0),
    )
    _ = crane.add_boom(
        "wire",
        description="The wire fixed to the last boom. Flexible connection",
        mass=wM,  # so far basically the hook
        mass_rng=(50, 2000),
        mass_center=1.0,
        boom=(wL, np.radians(90), 0),
        q_factor=wQ,
    )

    assert ensure_subpath("component-model", "examples")
    from oscillator_xd import Force, OscillatorXD  # type: ignore  ## it works for pytest!

    force = Force(dim=6, func=partial(calc_force, crane=crane))
    osc = OscillatorXD(dim=6, k=k, c=c, m=m, force=force)
    assert osc.force is not None
    for i in range(len(x0)):
        osc.x[i] = x0[i]
    for i in range(len(v0)):
        osc.v[i] = v0[i]

    return (crane, force, osc)


@pytest.mark.skipif(not component_model_examples_loaded, reason="component_model examples not loaded")
def test_crane_on_spring(show: bool = False):
    """Test the crane object as force on a 6D oscillator. No FMUs involved."""

    def do_experiment(
        pM: float = 10000.0,
        pCoM: tuple[float, ...] = (0.5, -1.0, 0.8),
        pH: float = 3.0,
        bM: float = 1000.0,
        bL: float = 8.0,
        bA: float = np.radians(90.0),
        wM: float = 50.0,
        wL: float = 1e-6,
        wQ: float = 50.0,
        k: tuple[float, ...] = (1e4,) * 6,
        c: tuple[float, ...] = (0,) * 6,
        m: float = 1e4,
        x0: tuple[float, ...] = (0.0,) * 6,
        v0: tuple[float, ...] = (0.0,) * 6,
        title: str = "Experiment",
        show: bool = show,
    ):
        crane, force, osc = make_crane_on_spring(pM, pCoM, pH, bM, bL, bA, wM, wL, wQ, k, c, m, x0, v0)

        results: dict[str, list[float]] = {}
        for i in range(3):
            results.update({f"boom.end[{i}]": [], f"f[{i}]": [], f"v[{i}]": []})
        times: list[float] = []
        t = 0.0
        dt = 0.01
        _boom = crane.boom_by_name("boom")
        assert _boom is not None, "Expect 'boom' to exist"
        while t <= 3.145:  # 10.0:
            t += dt
            times.append(t)
            osc.do_step(t, dt)  # takes updated force, x and v and calculates updated x and v
            for i in range(3):
                results[f"boom.end[{i}]"].append(_boom.end[i])
                results[f"f[{i}]"].append(osc.force.out[i])
                results[f"v[{i}]"].append(osc.v[i])

        for b in crane.booms():
            print(f"{b.name}: {b.direction}")
        print(f"Force: {osc.force.out}")
        print(f"Speed: {osc.v}")

        if show:
            do_plot(times, ("boom.end[0]", "boom.end[1]", "boom.end[2]", "f[1]", "v[1]"), results, title)

    do_experiment(m=1e4, c=(0.5,) * 6, pCoM=(0.5, 0, 0), bA=np.radians(180), title="Straight crane", show=show)
    # do_experiment(m=1e4, c=(0.5,) * 6, pCoM=(0.5, 0, 0), bA=np.radians(90), title="90deg boom crane", show=show)


@pytest.mark.skipif(not component_model_examples_loaded, reason="component_model examples not loaded")
def test_controlled_crane_on_spring(show: bool = False):
    """Move the crane-on-spring, which should lead to wobling."""

    def movement(crane: Crane, dt: float, t_end: float) -> Generator[tuple[float, Crane], None, None]:
        """Create movement of the crane through definition and usage of Controls.
        Generaor function which yields updated crane objects.
        time is defined global as a simple way to draw the current time together with the title.
        """
        # initial definition of controls and start values
        controls = Controls(limit_err=logging.WARNING)  # CRITICAL)
        f, p, b1, w = list(crane.booms())
        assert isinstance(w, Wire)
        w.additional_checks = True
        controls.extend(
            (
                Control("turn", (None, (-0.31, 0.31), (-1, 1)), rw=partial(_boom, p, 2)),
                Control("luff", ((0, 1.58), (-0.18, 0.09), (-0.09, 0.05)), rw=partial(_boom, b1, 1)),
                Control("boom", ((8, 50), (-0.7, 0.5), (-0.1, 0.05)), rw=partial(_boom, b1, 0)),
                Control("wire", ((0.5, 50), (-0.1, 1.0), (-0.4, 0.5)), rw=partial(_boom, w, 0)),
            )
        )

        # From time 0 we set three goals
        controls["wire"].setgoal(0, 7.0)  # wire length increase to 9m
        controls["turn"].setgoal(0, np.radians(90))  # turn pedestal 90 deg
        controls["luff"].setgoal(0, np.radians(45))  # luff boom to 45 deg
        controls["boom"].setgoal(0, 12.0)  # increase length  to 15m (0.1m/s)
        for time in np.linspace(0.0, t_end, int(t_end / dt) + 1):
            # if abs(time-10.0)<0.5*dt:  # Stop wire length increase
            #    controls["wire"].setgoal(1, 0.0) # wire length increase with 1m/s
            controls.step(time, dt)
            # if time > 7.2:
            #    print(f"w.additional:{w.additional_checks}")
            crane.do_step(time, dt)  # takes updated force, x and v and calculates updated x and v
            # crane.boom0.translate(-crane.boom0.origin)
            yield (time + dt, crane)

    if not show:  # nothing to do in this case
        return

    crane, force, osc = make_crane_on_spring(
        m=1e7, bL=8.0, wL=1.0, wM=1000.0, wQ=10000
    )  # define a default crane on spring

    ani = AnimateCrane(crane, movement=movement, dt=0.01, t_end=15, interval=10)  # type: ignore  ## It is a Generator!
    ani.do_animation()


if __name__ == "__main__":
    """Run the tests defined here.

    Note: The FMUs are not produced here. Only loaded. To change the FMUs use
    'test_mobile_crane_*.py' for the crane
    'test_oscillator_6dof_fmu.py' from the component-model package for the oscillator
    """
    retcode = 0  # pytest.main(["-rA", "-v", "--rootdir", "../", "--show", "False", __file__])
    assert retcode == 0, f"Non-zero return code {retcode}"
    make_fmus()
    make_mobile_crane_straight()
    # test_crane_on_spring( show=True)
    test_controlled_crane_on_spring(show=True)
