import copy
import os
import warnings
from copy import deepcopy
from functools import partial
from typing import List, Optional, Tuple

from scipy.optimize import minimize

from src.postprocessing.extract_poses import load_robot_poses, load_computed_poses
import matplotlib.pyplot as plt
import numpy as np
import open3d as o3d


def rotation_matrix_x(theta):
    return np.array(
        [[1, 0.0, 0.0],
         [0.0, np.cos(theta), -np.sin(theta)],
         [0.0, np.sin(theta), np.cos(theta)]]
    )


def rotation_matrix_y(theta):
    return np.array(
        [[np.cos(theta), 0.0, np.sin(theta)],
         [0.0, 1.0, 0.0],
         [-np.sin(theta), 0.0, np.cos(theta)]]
    )


def rotation_matrix_z(theta):
    return np.array(
        [[np.cos(theta), -np.sin(theta), 0.0],
         [np.sin(theta), np.cos(theta), 0.0],
         [0.0, 0.0, 1.0]]
    )


class Pose:
    def __init__(self, id, pos_vec, rot_mat):
        self.id = id
        self.pos_vec = pos_vec
        self.rot_mat = rot_mat


def extract_robot_camera_poses(real_poses):
    poses = []
    for id, rm in real_poses.items():
        rot_mat = rm[:3, :3]
        pos_vec = np.array(rm[:3, 3]).reshape(3, 1)
        poses.append(Pose(id, pos_vec, rot_mat))
    print("Collected a total of {} robot camera poses".format(len(poses)))

    return poses


def extract_inferred_camera_poses(computed_poses):
    poses = []
    for id, data in computed_poses.items():
        rot_mat = np.array(data["rotation"]).reshape(3, 3).astype("float")
        pos_vec = np.array(data["center"]).reshape(3, 1).astype("float")
        poses.append(Pose(id, pos_vec, rot_mat))

    print("Collected a total of {} computed poses".format(len(poses)))
    return poses


def plot_camera_positions_matplotlib(ax, poses, color, marker):
    for pose in poses:
        ax.scatter(pose.pos_vec[0], pose.pos_vec[1], pose.pos_vec[2], color=color, marker=marker, alpha=0.6)
    plt.axis("auto")


def plot_camera_positions_open3d(ax, poses, color):
    for pose in poses:
        ax.scatter(pose.pos_vec[0], pose.pos_vec[1], pose.pos_vec[2], color=color)
    plt.axis("auto")


def plot_camera_poses(ax, poses, color):
    for pose in poses:
        ax.quiver(pose.pos_vec[0], pose.pos_vec[1], pose.pos_vec[2], pose.rot_mat[:, 0],
                  pose.rot_mat[:, 1], pose.rot_mat[:, 2], length=0.1,
                  normalize=True, color=color)
    plt.axis("auto")


def optimal_translation(r_reals, r_computeds, r_estimate):
    l = []

    # print(r_estimate)
    for i, (r_real, r_computed) in enumerate(zip(r_reals, r_computeds)):
        l.append(np.linalg.norm(r_real - r_estimate))
    return np.sum(l)


def get_3d_rotation_matrix(alpha, beta, gamma):
    return rotation_matrix_x(alpha) @ rotation_matrix_y(beta) @ rotation_matrix_z(gamma)


def get_scaling_matrix(a, b, c):
    return np.array([
        [a, 0.0, 0.0],
        [0.0, b, 0.0],
        [0.0, 0.0, c], ]
    )


def optimal_rotation(r_reals, r_computeds, rotation_angles):
    l = []
    alpha, beta, gamma = rotation_angles[0], rotation_angles[1], rotation_angles[2]
    S = get_3d_rotation_matrix(alpha, beta, gamma)
    # print(S)
    for i, (r_real, r_computed) in enumerate(zip(r_reals, r_computeds)):
        l.append(np.linalg.norm(r_real - np.matmul(S, r_computed)))
    return np.sum(l)


def optimal_scaling(r_reals, r_computeds, scaling_factors):
    l = []
    a, b, c = scaling_factors[0], scaling_factors[1], scaling_factors[2]
    scaling = get_scaling_matrix(a, b, c)
    # print(scaling)
    for i, (r_real, r_computed) in enumerate(zip(r_reals, r_computeds)):
        l.append(np.linalg.norm(r_real - np.matmul(scaling, r_computed)))
    return np.sum(l)


def order_poses_by_id(poses):
    return sorted(poses, key=lambda x: x.id)


def optimize_rotation(real_poses, computed_poses):
    x0_Omega = np.random.uniform(-2 * np.pi, +2 * np.pi, [1, 3])
    ordered_real_poses = order_poses_by_id(real_poses)
    ordered_computed_poses = order_poses_by_id(computed_poses)
    print("Robot camera poses ordered by filename: {}".format([pose.id for pose in ordered_real_poses]))
    print("Computed poses ordered by filename: {}".format([pose.id for pose in ordered_computed_poses]))

    # print(compute_geometric_center(ordered_real_poses))
    # print(compute_geometric_center(ordered_computed_poses))

    r_reals = [pose.pos_vec for pose in ordered_real_poses]
    r_computeds = [pose.pos_vec for pose in ordered_computed_poses]

    partial_func_optimal_rotation = partial(optimal_rotation, r_reals, r_computeds)
    res2 = minimize(partial_func_optimal_rotation, x0_Omega, method='Nelder-Mead', tol=1e-14,
                    options={"maxiter": 200000, "disp": True})
    alpha = res2.x[0]
    beta = res2.x[1]
    gamma = res2.x[2]
    Omega = get_3d_rotation_matrix(alpha, beta, gamma)
    # print(alpha, beta, gamma)

    return Omega


def optimize_scaling(real_poses, computed_poses):
    x0_scaling = np.random.uniform(-5, 5, [1, 3])
    ordered_real_poses = order_poses_by_id(real_poses)
    ordered_computed_poses = order_poses_by_id(computed_poses)

    # print(compute_geometric_center(ordered_real_poses))
    # print(compute_geometric_center(ordered_computed_poses))

    r_reals = [pose.pos_vec for pose in ordered_real_poses]
    r_computeds = [pose.pos_vec for pose in ordered_computed_poses]

    partial_func_optimal_scaling = partial(optimal_scaling, r_reals, r_computeds)
    res2 = minimize(partial_func_optimal_scaling, x0_scaling, method='Nelder-Mead', tol=1e-11,
                    options={"maxiter": 200000, "disp": True})
    a = res2.x[0]
    b = res2.x[1]
    c = res2.x[2]

    # print(a, b, c)
    return get_scaling_matrix(a, b, c)


def transform_poses(poses, scaling: Optional[np.ndarray] = None,
                    rotation: Optional[np.ndarray] = None,
                    translation: Optional[np.ndarray] = None):
    transformed_poses = deepcopy(poses)

    if translation is not None:
        for pose in transformed_poses:
            pose.pos_vec = pose.pos_vec + translation

    if rotation is not None:
        for pose in transformed_poses:
            pose.pos_vec = rotation @ pose.pos_vec

    if scaling is not None:
        for pose in transformed_poses:
            pose.pos_vec = scaling @ pose.pos_vec

    return transformed_poses


def compute_centroid(vertices: np.ndarray):
    xs = np.mean(vertices[:, 0])
    ys = np.mean(vertices[:, 1])
    zs = np.mean(vertices[:, 2])
    return np.array([xs, ys, zs]).reshape(3, 1)


def compute_geometric_center(poses: List[Pose]):
    vertices = np.zeros([len(poses), 3])
    for index, pose in enumerate(poses):
        vertices[index, :] = np.transpose(pose.pos_vec)
    centroid = compute_centroid(vertices)
    return centroid


def translate_poses(computed_poses, r_estimated):
    c = deepcopy(computed_poses)
    for pose in c:
        pose.pos_vec -= r_estimated

    return c


def check_omega(Omega: np.ndarray):
    # check if determinant of matrix = 1.0 (rotation matrix should satisfy this condition)
    return np.round(np.linalg.det(Omega), 2) == 1.0


def center_poses(poses):
    centroid = compute_geometric_center(poses)
    centered_poses = translate_poses(poses, centroid)
    return centered_poses, centroid


class Transformation:
    def __init__(self, translation: np.array, rotation: np.array, scaling: np.array,
                 computed_cameras_centroid: np.array):
        # order of transformation should be
        self.translation = translation
        self.rotation = rotation
        self.check_rotation()
        self.scaling = scaling
        self.computed_cameras_centroid = computed_cameras_centroid

    def check_rotation(self):
        if not check_omega(self.rotation):
            raise UserWarning("Determinant of matrix is not 1.0, not a valid rotation matrix...")

    def serialize(self, path_to_file: str):
        with open(path_to_file, "w") as outfile:
            outfile.write("Translation vector:\n{}\n\n".format(str(np.round(self.translation, 4))))
            outfile.write("Rotation matrix:\n{}\n\n".format(str(np.round(self.rotation, 4))))
            outfile.write("Scaling matrix:\n{}\n\n".format(str(np.round(self.scaling, 4))))
            outfile.write("Centroid of computed camera positions:\n{}\n\n".format(
                str(np.round(self.computed_cameras_centroid, 4))))


def pipeline(real_poses, computed_poses):
    if len(real_poses) != len(computed_poses):
        warnings.warn(
            "Number of computed camera poses is not equal to actual camera poses. Optimization may not converge...")

    if len(real_poses) < 20:
        warnings.warn(
            "Number of taken photos is too small (less than 20), which may compromise Meshroom's ability to estimate camera "
            "positions correctly, thus compromising the estimation of the transformation matrix...")

    _, real_centroid = center_poses(real_poses)
    _, computed_centroid = center_poses(computed_poses)

    # center all poses
    centered_real_poses = transform_poses(poses=real_poses, translation=-real_centroid)
    poses = transform_poses(poses=computed_poses, translation=-computed_centroid)
    print("Centroid of Centered computed poses: {}".format(center_poses(centered_real_poses)[1]))
    print("Centroid of Centered real poses: {}".format(center_poses(poses)[1]))

    # find optimal orientation
    rotation = optimize_rotation(centered_real_poses, poses)
    poses = transform_poses(poses=poses, rotation=rotation)

    # find optimal scaling
    scaling = optimize_scaling(centered_real_poses, poses)
    poses = transform_poses(poses=poses, scaling=scaling)

    # translate to centroid of real camera positions
    poses = transform_poses(poses=poses, translation=real_centroid)

    # for t_pose, dt_pose in zip(rotated_poses, doubly_transformed_poses):
    #     print("before scaling: {}".format(t_pose.pos_vec))
    #     print("after scaling: {}".format(dt_pose.pos_vec))
    # TODO: add some sanity checks

    return poses, Transformation(translation=real_centroid, rotation=rotation, scaling=scaling,
                                 computed_cameras_centroid=computed_centroid)


def transform_mesh(mesh: o3d.geometry.TriangleMesh,
                   transformation: Transformation) -> o3d.geometry.TriangleMesh:
    # bring model to centroid of computed cameras array
    mesh_centered = copy.deepcopy(mesh).translate(-transformation.computed_cameras_centroid)

    # rotate
    mesh_rotated = mesh_centered.rotate(transformation.rotation, center=(0, 0, 0))

    # scale
    mesh_scaled = mesh_rotated.scale(np.mean(np.diag(transformation.scaling)), center=(0, 0, 0))

    # translate back
    mesh_translated = copy.deepcopy(mesh_scaled).translate(transformation.translation)

    return mesh_translated


def convert_poses_to_open3d(poses: List[Pose]):
    points = []
    for pose in poses:
        points.append(pose.pos_vec.transpose().tolist()[0])

    return o3d.utility.Vector3dVector(points)


def convert_poses_to_o3d_point_cloud(poses: List[Pose],
                                     color: Optional[Tuple[int, int, int]]) -> o3d.geometry.PointCloud:
    real_points = convert_poses_to_open3d(poses)
    point_cloud = o3d.geometry.PointCloud(
        points=real_points,
    )

    point_cloud.colors = o3d.utility.Vector3dVector([color for i in range(len(real_points))])

    return point_cloud


def transform_model_to_world_coordinates(path_to_poses_dir: str, path_to_cameras_sfm: str, path_to_computed_mesh: str,
                                         path_to_transformed_mesh_dir: str, show_plot=False, exclude_poses=None):
    if not exclude_poses:
        exclude_poses = []

    # optimization and transformation
    print("Extracting robot camera poses from provided .npy files...")
    real_poses = extract_robot_camera_poses(load_robot_poses(
        path_to_poses_dir))
    print("Extracting computed camera poses from Meshroom .sfm file...")
    computed_poses = extract_inferred_camera_poses(load_computed_poses(
        path_to_cameras_sfm))

    # exclude erroneously compurted poses (manual input)
    real_poses = list(filter(lambda x: x.id not in exclude_poses, real_poses))
    computed_poses = list(filter(lambda x: x.id not in exclude_poses, computed_poses))

    print("Running optimization pipeline to obtain a transformation estimate...")
    transformed_poses, transformation = pipeline(real_poses, computed_poses)
    print("Loading computed .obj...")
    mesh = o3d.io.read_triangle_mesh(path_to_computed_mesh)
    print("Applying transformation to computed .obj...")
    transformed_mesh = transform_mesh(mesh, transformation)

    if show_plot:
        # open3d plots
        coord_frame = o3d.geometry.TriangleMesh.create_coordinate_frame()
        real_points = convert_poses_to_open3d(real_poses)
        real_cameras_pcl = o3d.geometry.PointCloud(points=real_points, )
        real_cameras_pcl.colors = o3d.utility.Vector3dVector([[1, 0, 0] for i in range(len(real_points))])
        real_cameras_pcl = convert_poses_to_o3d_point_cloud(poses=real_poses, color=(1, 0, 0))
        computed_cameras_pcl = convert_poses_to_o3d_point_cloud(poses=computed_poses, color=(0, 1, 0))
        transformed_cameras_pcl = convert_poses_to_o3d_point_cloud(poses=transformed_poses, color=(0, 0, 1))
        o3d.visualization.draw_geometries(
            [coord_frame, mesh, transformed_mesh, real_cameras_pcl, computed_cameras_pcl, transformed_cameras_pcl])

        # matplotlib plots
        fig = plt.figure()
        ax = fig.gca(projection='3d')
        plot_func = plot_camera_positions_matplotlib
        plot_func(ax, real_poses, color='r', marker="o")
        plot_func(ax, computed_poses, color='g', marker="o")
        plot_func(ax, transformed_poses, color='b', marker="*")
        plt.show()

    print("Writing output to new .obj...")

    o3d.io.write_triangle_mesh(os.path.join(path_to_transformed_mesh_dir, "transformed_mesh.obj"),
                               transformed_mesh, print_progress=True)
    transformation.serialize(os.path.join(path_to_transformed_mesh_dir, "transformation_parameters.txt"))


if __name__ == '__main__':
    transform_model_to_world_coordinates(
        path_to_poses_dir="/mnt/storage/roboweldar/simulation_test_1/new_dataset_reconstruction",
        path_to_cameras_sfm="/mnt/storage/roboweldar/simulation_test_1/MeshroomCache/StructureFromMotion/578f830addc4cce0c6fccaca48fd23068ffff1a3/cameras.sfm",
        path_to_computed_mesh="/mnt/storage/roboweldar/simulation_test_1/MeshroomCache/Texturing/4e42bb83b74f64fe0fbec93a87a58355fa145bd6/texturedMesh.obj",
        path_to_transformed_mesh_dir="/mnt/storage/roboweldar/simulation_test_1/MeshroomCache/Texturing/4e42bb83b74f64fe0fbec93a87a58355fa145bd6/transformed_mesh",
        show_plot=True,
        exclude_poses=None)
