import pytest
import gym_electric_motor as gem
from gym_electric_motor.physical_systems import (
    PolynomialStaticLoad,
    MechanicalLoad,
    ConstantSpeedLoad,
    ExternalSpeedLoad,
)
from gymnasium.spaces import Box
import numpy as np
from scipy import signal
import math

# The load parameter values used in the test
load_parameter1 = {
    "j_load": 0.2,
    "state_names": ["omega"],
    "j_rot_load": 0.25,
    "omega_range": (0, 1),
    "parameter": dict(a=0.12, b=0.13, c=0.4, j_load=0.2),
}
# The different initializers used in test
test_const_initializer = {
    "states": {"omega": 15.0},
    "interval": None,
    "random_init": None,
    "random_params": (None, None),
}
# todo random init as testcase ?
test_rand_initializer = {
    "interval": None,
    "random_init": None,
    "random_params": (None, None),
}
# profile_parameter
test_amp = 20
test_bias = 10
test_freq = 2


def speed_profile_(t, amp, freq, bias):
    return amp * signal.sawtooth(2 * np.pi * freq * t, width=0.5) + bias


@pytest.fixture
def defaultMechanicalLoad():
    """
     pytest fixture that returns a default MechanicalLoad object
    :return: MechanicalLoad object initialized with default values
    """
    return MechanicalLoad()


@pytest.fixture
def concreteMechanicalLoad():
    """
     pytest fixture that returns a  MechanicalLoad object with concrete parameters
    :return: MechanicalLoad object initialized with concrete values
    """
    # Parameters picked from the parameter dict above
    state_names = load_parameter1["state_names"]
    j_load = load_parameter1["j_load"]
    test_initializer = test_const_initializer
    return MechanicalLoad(state_names, j_load, load_initializer=test_initializer)


@pytest.fixture
def defaultPolynomialLoad():
    """
    pytest fixture that returns a default PolynomialLoad object
    :return: PolynomialLoad object initialized with default values
    """
    return PolynomialStaticLoad()


@pytest.fixture
def concretePolynomialLoad():
    """
    pytest fixture that returns a  PolynomialLoad object with concrete parameters
    :return: PolynomialLoad object initialized with concrete values
    """
    test_load_params = dict(a=0.01, b=0.05, c=0.1, j_load=0.1)
    test_initializer = test_const_initializer
    # x, y are random kwargs
    return PolynomialStaticLoad(
        load_parameter=test_load_params, load_initializer=test_initializer
    )


def test_InitMechanicalLoad(defaultMechanicalLoad, concreteMechanicalLoad):
    """
    Test whether the MechanicalLoad object attributes are initialized with the correct values
    :param defaultMechanicalLoad:
    :param concreteMechanicalLoad:
    :return:
    """
    assert defaultMechanicalLoad.state_names == load_parameter1["state_names"]
    assert defaultMechanicalLoad.j_total == 0
    assert defaultMechanicalLoad.limits == {}
    assert defaultMechanicalLoad.initializer == {}
    assert concreteMechanicalLoad.state_names == load_parameter1["state_names"]
    assert concreteMechanicalLoad.j_total == load_parameter1["j_load"]
    assert concreteMechanicalLoad.limits == {}
    assert concreteMechanicalLoad.initializer == test_const_initializer


def test_MechanicalLoad_set_j_rotor(concreteMechanicalLoad):
    """
    Test the set_j_rotor() function
    :param concreteMechanicalLoad:
    :return:
    """
    test_val = 1
    concreteMechanicalLoad.set_j_rotor(test_val)
    assert concreteMechanicalLoad.j_total == test_val + load_parameter1["j_load"]


def test_MechanicalLoad_MechanicalOde(concreteMechanicalLoad):
    """
    test the mechanical_ode() function
    :param concreteMechanicalLoad:
    :return:
    """
    test_args = 0.5, [2, 3, 4], 0.99
    with pytest.raises(NotImplementedError):
        concreteMechanicalLoad.mechanical_ode(*test_args)


def test_MechanicalLoad_reset(concreteMechanicalLoad):
    """
    Test the reset() function
    :param concreteMechanicalLoad:
    :return:
    """
    # random testcase for the necessary parameters needed for initialization
    test_positions = {"omega": 0}
    test_nominal = np.array([80])
    # gymnasium.Box state space with random size
    test_space = Box(low=-1.0, high=1.0, shape=(3,))
    # set additional random kwargs
    resetVal = concreteMechanicalLoad.reset(
        a=7,
        b=9,
        state_positions=test_positions,
        nominal_state=test_nominal,
        state_space=test_space,
    )
    testVal = np.asarray(list(test_const_initializer["states"].values()))
    assert (resetVal == testVal).all()


def test_MechanicalLoad_get_state_space(concreteMechanicalLoad):
    """
    test the get_state_space() function
    :param concreteMechanicalLoad:
    :return:
    """
    test_omega_range = 0, 5
    test_omega_range_2 = 1, 4, 6, 7
    retVal = concreteMechanicalLoad.get_state_space(test_omega_range)
    assert retVal[0]["omega"] == test_omega_range[0]
    assert retVal[1]["omega"] == test_omega_range[1]

    retVal2 = concreteMechanicalLoad.get_state_space(
        test_omega_range_2
    )  # additional negative test
    assert retVal2[0]["omega"] == test_omega_range_2[0]
    assert retVal2[1]["omega"] == test_omega_range_2[1]


def test_InitPolynomialStaticLoad(defaultPolynomialLoad):
    """
    Test the __init__() function of PolynomialStaticLoad class with default parameters
    :param defaultPolynomialLoad:
    :return:
    """
    test_default_load_parameterVal = {"a": 0.01, "b": 0.05, "c": 0.1, "j_load": 0.1}
    assert defaultPolynomialLoad.load_parameter == test_default_load_parameterVal


def test_InitPolynomialStaticLoad(concretePolynomialLoad):
    """
    Test the __init__() function of PolynomialStaticLoad class with concrete parameters
    :param concretePolynomialLoad:
    :return:
    """
    test_concrete_load_parameterVal = {"a": 0.01, "b": 0.05, "c": 0.1, "j_load": 0.1}
    assert concretePolynomialLoad.load_parameter == test_concrete_load_parameterVal


@pytest.mark.parametrize(
    "omega, expected_result", [(-3, 23400), (0, 20000), (5, 11400)]
)  # to verify all 3 branches
def test_PolynomialStaticLoad_MechanicalOde(
    concretePolynomialLoad, omega, expected_result
):
    """
    test the mechanical_ode() function of the PolynomialStaticLoad class
    :param concretePolynomialLoad:
    :param omega:
    :return:
    """

    test_mechanical_state = np.array([omega])
    test_t = 1
    test_torque = 2
    load_parameter = dict(j_load=1e-4, a=0.01, b=0.02, c=0.03)
    op = PolynomialStaticLoad(load_parameter=load_parameter)
    output_val = op.mechanical_ode(test_t, test_mechanical_state, test_torque)[0]
    # output_val = concretePolynomialLoad.mechanical_ode(test_t, test_mechanical_state, test_torque)
    assert math.isclose(expected_result, output_val, abs_tol=1e-6)


class TestMechanicalLoad:
    key = ""
    class_to_test = MechanicalLoad
    kwargs = {}


class TestConstSpeedLoad(TestMechanicalLoad):
    key = "ConstSpeedLoad"
    class_to_test = ConstantSpeedLoad

    @pytest.fixture
    def const_speed_load(self):
        return ConstantSpeedLoad(omega_fixed=60)

    def test_initialization(self):
        load = ConstantSpeedLoad(omega_fixed=60)
        assert load.omega_fixed == 60

    def test_mechanical_ode(self, const_speed_load):
        assert all(const_speed_load.mechanical_ode() == 0)

    def test_reset(self, const_speed_load):
        test_positions = {"omega": 0}
        test_nominal = np.array([80])
        # gymnasium.Box state space with random size
        test_space = Box(low=-1.0, high=1.0, shape=(3,))
        reset_val = const_speed_load.reset(
            state_positions=test_positions,
            nominal_state=test_nominal,
            state_space=test_space,
        )
        # set additional random kwargs
        assert all(reset_val == np.array([const_speed_load.omega_fixed]))

    @pytest.mark.parametrize(
        "omega, omega_fixed, expected",
        [(-0.5, 1000, (0, 0)), (0, 0, (0, 0)), (2, -523, (0, 0)), (2, 0.2, (0, 0))],
    )
    def test_jacobian(self, omega, omega_fixed, expected):
        test_object = self.class_to_test(omega_fixed)

        # 2 Runs to test independence on time and torque
        result0 = test_object.mechanical_jacobian(0.456, np.array([omega]), 0.385)
        result1 = test_object.mechanical_jacobian(5.345, np.array([omega]), -0.654)

        assert result0[0] == result1[0] == expected[0]
        assert result0[1] == result1[1] == expected[1]


class TestExtSpeedLoad(TestMechanicalLoad):
    key = "ExtSpeedLoad"
    class_to_test = ExternalSpeedLoad
    kwargs = dict(
        speed_profile=speed_profile_,
        speed_profile_kwargs=dict(amp=test_amp, bias=test_bias, freq=test_freq),
    )

    @pytest.fixture
    def ext_speed_load(self):
        return ExternalSpeedLoad(
            speed_profile=speed_profile_,
            speed_profile_kwargs=dict(amp=test_amp, bias=test_bias, freq=test_freq),
        )

    def test_initialization(self):
        load = ExternalSpeedLoad(
            speed_profile=speed_profile_,
            speed_profile_kwargs=dict(amp=test_amp, bias=test_bias, freq=test_freq),
        )
        assert load._speed_profile == speed_profile_
        assert load.omega == speed_profile_(
            t=0, amp=test_amp, bias=test_bias, freq=test_freq
        )
        for key in ["amp", "bias", "freq"]:
            assert key in load.speed_profile_kwargs

    # to verify all 3 branches
    @pytest.mark.parametrize(
        "omega, expected_result", [(-3, -69840.0), (0, -99840.0), (5, -149840.0)]
    )
    def test_mechanical_ode(self, ext_speed_load, omega, expected_result):
        test_mechanical_state = np.array([omega])
        test_t = 1
        op = ext_speed_load
        output_val = op.mechanical_ode(test_t, test_mechanical_state)[0]
        assert math.isclose(expected_result, output_val, abs_tol=1e-6)

    def test_reset(self, ext_speed_load):
        test_positions = {"omega": 0}
        test_nominal = np.array([80])
        # gymnasium.Box state space with random size
        test_space = Box(low=-1.0, high=1.0, shape=(3,))
        reset_var = ext_speed_load.reset(
            state_positions=test_positions,
            nominal_state=test_nominal,
            state_space=test_space,
        )
        assert all(reset_var == np.array([ext_speed_load._omega_initial]))

    @pytest.mark.parametrize(
        "omega, omega_initial, expected",
        [
            (-0.5, 1000, (None, None)),
            (0, 0, (None, None)),
            (2, -523, (None, None)),
            (2, 0.2, (None, None)),
        ],
    )
    def test_jacobian(self, omega, omega_initial, expected):
        test_object = self.class_to_test(
            speed_profile_,
            speed_profile_kwargs=dict(amp=test_amp, bias=test_bias, freq=test_freq),
        )

        # 2 Runs to test independence on time and torque
        result0 = test_object.mechanical_jacobian(0.456, np.array([omega]), 0.385)
        result1 = test_object.mechanical_jacobian(5.345, np.array([omega]), -0.654)

        assert result0 == result1 == expected[0]
        assert result0 == result1 == expected[1]


class TestPolyStaticLoad(TestMechanicalLoad):
    class_to_test = PolynomialStaticLoad
    key = "PolyStaticLoad"

    @pytest.mark.parametrize(
        "omega, load_parameter, expected",
        [
            (
                -0.5,
                dict(a=12, b=1, c=0, j_load=1),
                (np.array([[-1.0]]), np.array([[1.0]])),
            ),
            (0, dict(j_load=0.5), (np.array([[-1000.0]]), np.array([[2.0]]))),
            (2, dict(a=20, b=0, c=2, j_load=0.25), (-32, 4)),
            (2, dict(a=20, b=0.125, c=2, j_load=0.25), (-32.5, 4)),
        ],
    )
    def test_jacobian(self, omega, load_parameter, expected):
        test_object = self.class_to_test(load_parameter)

        # 2 Runs to test independence on time and torque
        result0 = test_object.mechanical_jacobian(0.456, np.array([omega]), 0.385)
        result1 = test_object.mechanical_jacobian(5.345, np.array([omega]), -0.654)

        assert result0[0] == result1[0] == expected[0]
        assert result0[1] == result1[1] == expected[1]
