from pybotics.robot import Robot
from pybotics.predefined_models import ur10

nominal_robot = Robot.from_parameters(ur10())

import pandas as pd

def display_robot_kinematics(robot: Robot):
    df = pd.DataFrame(robot.kinematic_chain.matrix)
    df.columns = ["alpha", "a", "theta", "d"]
    display(df)

display_robot_kinematics(nominal_robot)

"""
Explanation: Robot Calibration
Nominal Robot

A nominal robot model:
Represents what the robot manufacturer intended as a kinematic model
Is mathematically ideal
End of explanation
"""


import numpy as np
from copy import deepcopy

real_robot = deepcopy(nominal_robot)

# let's pretend our real robot has small joint offsets
# in real life, this would be a joint mastering issue (level-1 calibration)
# https://en.wikipedia.org/wiki/Robot_calibration
for link in real_robot.kinematic_chain.links:
    link.theta += np.random.uniform(
        low=np.deg2rad(-0.1),
        high=np.deg2rad(0.1)
    )

display_robot_kinematics(real_robot)

"""
Explanation: Real Robots

Real robots do not conform perfectly to the nominal parameters
Small errors in the robot model can generate large errors in Cartesian position
Sources of errors include, but are not limited to:
Kinematic errors
Mechanical tolerances
Angle offsets


Non-kinematic errors
Joint stiffness
Gravity
Temperature
Friction
End of explanation
"""


joints = []
positions = []
for i in range(1000):
    q = real_robot.random_joints()
    pose = real_robot.fk(q)
    
    joints.append(q)
    positions.append(pose[:-1,-1])

pd.DataFrame(joints).describe()

pd.DataFrame(positions, columns=['x','y','z']).describe()

"""
Explanation: Get Real (aka Measured) Poses

In real life, these poses would be measured using metrology equipment (e.g., laser tracker, CMM)
End of explanation
"""


from sklearn.model_selection import train_test_split
split = train_test_split(joints, positions, test_size=0.3)

train_joints = split[0]
test_joints = split[1]

train_positions = split[2]
test_positions = split[3]

"""
Explanation: Split Calibration and Validation Measures

A portion of the measured configurations and positions should be set aside for validation after calibration (i.e., optimization)
This is to prevent/check the optimized model for overfitting
End of explanation
"""


from pybotics.optimization import compute_absolute_errors

nominal_errors = compute_absolute_errors(
    qs=test_joints,
    positions=test_positions,
    robot=nominal_robot
)

display(pd.Series(nominal_errors).describe())

"""
Explanation: Get Nominal Position Errors

These nominal model is our starting point for calibration
The errors are in millimetres
End of explanation
"""


from pybotics.optimization import OptimizationHandler

# init calibration handler
handler = OptimizationHandler(nominal_robot)

# set handler to solve for theta parameters
kc_mask_matrix = np.zeros_like(nominal_robot.kinematic_chain.matrix, dtype=bool)
kc_mask_matrix[:,2] = True
display(kc_mask_matrix)

handler.kinematic_chain_mask = kc_mask_matrix.ravel()

from scipy.optimize import least_squares
from pybotics.optimization import optimize_accuracy

# run optimization
result = least_squares(
    fun=optimize_accuracy,
    x0=handler.generate_optimization_vector(),
    args=(handler, train_joints, train_positions),
    verbose=2
)  # type: scipy.optimize.OptimizeResult

"""
Explanation: Calibration
End of explanation
"""


calibrated_robot = handler.robot
calibrated_errors = compute_absolute_errors(
    qs=test_joints,
    positions=test_positions,
    robot=calibrated_robot
)

display(pd.Series(calibrated_errors).describe())

import matplotlib.pyplot as plt

%matplotlib inline

plt.xscale("log")
plt.hist(nominal_errors, color="C0", label="Nominal")
plt.hist(calibrated_errors, color="C1", label="Calibrated")

plt.legend()
plt.xlabel("Absolute Error [mm]")
plt.ylabel("Frequency")

"""
Explanation: Results

A calibrated robot model is never perfect in real life
The goal is often to reduce the max error under a desired threshold
End of explanation
"""


