# For any of the general bots, the test cases are generated by the following:
# First, loop through all the kinematics of the robot (h and p matrices) in the test_cases folder
# For each of these, we generate 20 random positions for the robot to be in (q values)
# We run forward kinematics for each of these, then run inverse kinematics on the result
# We then run forward kinematics on the result to make sure it is the same as the forward kinematics of the input

# For general bots, we don't have access to forward kinematics, so we can't do the same test
# Instead, we will just run inverse kinematics on all of the given test cases in the test_cases folder and make sure that we actually return a nonzero result
# THIS IS NOT A TEST FOR CORRECTNESS ON HARD-CODED BOTS, ONLY ON GENERAL BOTS

# Requires numpy to be installed

# Import the test case csv files

from collections import namedtuple
import numpy as np
import unittest
from ik_geo import Robot
from math import pi

# Import the csv files
import csv
import os

TestBot = namedtuple("TestBot", ["casename", "robot", "testcases"])
TestCaseGeneral = namedtuple("TestCaseGeneral", ["hVals", "pVals"])
TestCaseHardCoded = namedtuple(
    "TestCaseHardCoded", ["rotationMatrix", "positionVector"]
)


generalFilenames = [
    "IK_2_intersecting",
    "IK_2_parallel",
    "IK_3_parallel_2_intersecting",
    "IK_3_parallel",
    "IK_spherical",
    "IK_spherical_2_parallel",
    "IK_spherical_2_intersecting",
    "IK_gen_6_dof",
]
hardcodedFilenames = [
    "IRB_6640",
    "spherical_bot",
    "three_parallel_bot",
    "two_parallel_bot",
    "ur5"
]

epsilon = 1e-3


# Test the general robots
class TestGeneralRobots(unittest.TestCase):
    def check_general_robot(self, bot):
        np.random.seed(0)
        for configNum, testCase in enumerate(bot.testcases):
            hMatrix = np.reshape(testCase.hVals, (6, 3))
            pMatrix = np.reshape(testCase.pVals, (7, 3))
            robot = bot.robot(hMatrix, pMatrix)

            # Generate 20 random robot configurations
            qVals = np.random.rand(20, 6) * 2 * pi - pi
            # Generate fewer for the 6 dof bot because it is slower
            if bot.casename == "IK_gen_6_dof":
                qVals = np.random.rand(5, 6) * 2 * pi - pi
            for i, q in enumerate(qVals):
                # Get the forward kinematics result and then run inverse to see if we get the same thing
                forward_kinematics = robot.forward_kinematics(q)
                rotation = np.array(forward_kinematics[0])
                translation = np.array(forward_kinematics[1])

                # Get the inverse kinematics
                results = robot.get_ik_sorted(
                    forward_kinematics[0], forward_kinematics[1]
                )

                self.assertFalse(
                    len(results) == 0,
                    msg=f"Failed on test {i + 1} of {bot.casename} configuration {configNum + 1} with \nqVals: {str(q)} \
                                    \nRotation matrix expected: \n{str(rotation)}\nTranslation expected: {str(translation)} \
                                    \nhMatrix: \n{str(hMatrix)}\npMatrix: \n{str(pMatrix)} \
                                    \nNo result was returned.",
                )
                result = results[0]

                # Run forward kinematics on the result to make sure it is the same as the input
                resultForward = robot.forward_kinematics(result[0])
                resultRotation = np.array(resultForward[0])
                resultTranslation = np.array(resultForward[1])

                # Check to make sure each value is roughly equal, up to 2pi
                for resultVal, expectedVal in zip(
                    rotation.flatten(), resultRotation.flatten()
                ):
                    self.assertAlmostEqual(
                        resultVal,
                        expectedVal,
                        delta=epsilon,
                        msg=f"Failed on test {i + 1} of {bot.casename} configuration {configNum + 1} with \nqVals: {str(q)}\n and result: {str(result[0])} \
                        \nRotation matrix expected: \n{str(rotation)}\nGot: \n{str(resultRotation)} \
                        \nTranslation expected: {str(translation)}\nGot: {str(resultTranslation)} \
                        \nhMatrix: \n{str(hMatrix)}\npMatrix: \n{str(pMatrix)}",
                    )
                for resultVal, expectedVal in zip(translation, resultTranslation):
                    self.assertAlmostEqual(
                        resultVal,
                        expectedVal,
                        delta=epsilon,
                        msg=f"Failed on test {i + 1} of {bot.casename} configuration {configNum + 1} with \nqVals: {str(q)}\n and result: {str(result[0])} \
                                            \nRotation matrix expected: \n{str(rotation)}\nGot: \n{str(resultRotation)} \
                                            \nTranslation expected: {str(translation)}\nGot: {str(resultTranslation)} \
                                            \nhMatrix: \n{str(hMatrix)}\npMatrix: \n{str(pMatrix)}",
                    )


def add_general_test(filename):
    # Setup the robot from the filename
    robot_name = (
        filename.replace("IK_", "")
        .replace("2", "two")
        .replace("3", "three")
        .replace("6", "six")
    )

    # Switch on robot name
    robot_instantiater = getattr(Robot, robot_name)

    testcases = []
    # Open ../../test_cases/FILENAME.csv
    with open(
        os.path.join(os.path.dirname(__file__), "..", "test_cases", filename + ".csv")
    ) as f:
        reader = csv.reader(f)
        # Skip the header
        next(reader)
        for row in reader:
            if not row:
                continue
            # Parse the given h and p vals to get a valid robot configuration
            hVals = [float(x) for x in row[0:18]]
            pVals = [float(x) for x in row[18:39]]
            testcases.append(TestCaseGeneral(hVals, pVals))
    bot = TestBot(filename, robot_instantiater, testcases)
    setattr(
        TestGeneralRobots,
        "test_" + filename,
        lambda self: self.check_general_robot(bot),
    )


class TestHardcodedBots(unittest.TestCase):
    def check_hardcoded_robot(self, bot):
        for testCase in bot.testcases:
            rotationMatrix = np.array(
                [
                    testCase.rotationMatrix[i : i + 3]
                    for i in range(0, len(testCase.rotationMatrix), 3)
                ]
            )
            positionVector = testCase.positionVector
            result = bot.robot.get_ik(rotationMatrix, positionVector)
            self.assertFalse(
                len(result) == 0,
                msg=f"Failed on {bot.casename} with rotation matrix: \n{rotationMatrix} and position vector: \n{positionVector}. \nNo result was returned.",
            )


def add_hardcoded_test(filename):
    # Setup the robot from the filename
    robot = getattr(Robot, filename.replace("IRB_6640", "irb6640"))()

    testcases = []
    # Open ../../test_cases/FILENAME.csv
    with open(
        os.path.join(os.path.dirname(__file__), "..", "test_cases", filename + ".csv")
    ) as f:
        reader = csv.reader(f)
        # Skip the header
        next(reader)
        for row in reader:
            if not row:
                continue
            # Parse the given h and p vals to get a valid robot configuration
            rotationMatrix = [float(x) for x in row[0:9]]
            positionVector = [float(x) for x in row[9:12]]
            testcases.append(TestCaseHardCoded(rotationMatrix, positionVector))
    bot = TestBot(filename, robot, testcases)
    setattr(
        TestHardcodedBots,
        "test_" + filename,
        lambda self: self.check_hardcoded_robot(bot),
    )


if __name__ == "__main__":
    # Generate classes for each general test case
    for filename in generalFilenames:
        add_general_test(filename)
    for filename in hardcodedFilenames:
        add_hardcoded_test(filename)
    unittest.main()
