from pathlib import Path

import numpy
from matplotlib import pyplot
from matplotlib.figure import figaspect, SubplotParams
from mpl_toolkits.mplot3d import Axes3D

from transformations import calculate_rotation_matrix, centroid, convert_rotation_matrix_to_angles

expected_points_file_path = Path('expected_points.txt')
actual_points_file_path = Path('actual_points.txt')

plot_3d: Axes3D = pyplot.figure(subplotpars=SubplotParams(0, 0, 1, 1, 0, 0),
                                figsize=(10, 8)).add_subplot(projection='3d')


def read_points_from_file(points_file_path: Path) -> numpy.ndarray:
    """
    Чтение пути из файла

    :param points_file_path: путь к файлу
    :return: Путь, вычитанный из файла

    """
    with points_file_path.open('r') as points_file:
        return numpy.array([[float(coordinate) for coordinate in point_line.strip().split(' ')]
                            for point_line in points_file])


def rsmd(first_path: numpy.ndarray, second_path: numpy.ndarray) -> float:
    """
    Расчёт среднеквадратичное отклонение между путями

    :param first_path: первый путь
    :param second_path: второй путь
    :return: СКО

    """
    diff = first_path - second_path
    return numpy.sqrt((diff * diff).sum() / first_path.shape[0])


if __name__ == '__main__':
    expected_points = read_points_from_file(expected_points_file_path)
    actual_points = read_points_from_file(actual_points_file_path)

    plot_3d.plot(expected_points[:, 0], expected_points[:, 1], expected_points[:, 2],
                 's-k', markersize=12, linewidth=3,
                 label='Expected Points')
    plot_3d.plot(actual_points[:, 0], actual_points[:, 1], actual_points[:, 2],
                 's-g', markersize=12, linewidth=3,
                 label='Actual Points')

    for calculation_points_number in range(len(actual_points), 1, -1):
        print(f"{15 * '='} Calc on {calculation_points_number:2} points {15 * '='}")
        rotation_matrix = calculate_rotation_matrix(actual_points,
                                                    expected_points,
                                                    calculation_points_number=calculation_points_number)
        angles = convert_rotation_matrix_to_angles(rotation_matrix)

        actual_points_rotated = numpy.dot(actual_points, rotation_matrix)

        actual_points_for_offset = actual_points_rotated[:calculation_points_number]
        expected_points_for_offset = expected_points[:calculation_points_number]

        offset = centroid(actual_points_for_offset) - centroid(expected_points_for_offset)
        actual_points_converted = actual_points_rotated - offset

        print(f"    {angles = }")
        print(f"    offset = {tuple(offset)}")
        print(f"    error = {rsmd(expected_points, actual_points_converted)}")

        plot_3d.plot(actual_points_converted[:, 0], actual_points_converted[:, 1], actual_points_converted[:, 2],
                     'o-', markersize=9, linewidth=3,
                     label=f'Transformed actual points (on {calculation_points_number} points)')

    plot_3d.grid(True)
    plot_3d.tick_params(labelsize=15)

    pyplot.legend()
    pyplot.show()