# -*- coding: UTF8 -*-
"""
Provides metrics for the evaluation of SLAM algorithms.
author: Michael Grupp

This file is part of evo (github.com/MichaelGrupp/evo).

evo is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

evo is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with evo.  If not, see <http://www.gnu.org/licenses/>.
"""

import abc
import itertools
import logging
import math
import sys
import typing
from enum import Enum, unique

import numpy as np

from evo import EvoException
from evo.core import filters, trajectory
from evo.core.result import Result
from evo.core import lie_algebra as lie
from evo.core.units import Unit, ANGLE_UNITS, LENGTH_UNITS, METER_SCALE_FACTORS

if sys.version_info[0] >= 3 and sys.version_info[1] >= 4:
    ABC = abc.ABC
else:
    ABC = abc.ABCMeta("ABC", (), {})

logger = logging.getLogger(__name__)

PathPair = tuple[trajectory.PosePath3D, trajectory.PosePath3D]


class MetricsException(EvoException):
    pass


@unique
class StatisticsType(Enum):
    rmse = "rmse"
    mean = "mean"
    median = "median"
    std = "std"
    min = "min"
    max = "max"
    sse = "sse"


@unique
class PoseRelation(Enum):
    full_transformation = "full transformation"
    translation_part = "translation part"
    rotation_part = "rotation part"
    rotation_angle_rad = "rotation angle in radians"
    rotation_angle_deg = "rotation angle in degrees"
    point_distance = "point distance"
    point_distance_error_ratio = "point distance error ratio"


class Metric(ABC):
    @abc.abstractmethod
    def process_data(self, data):
        return

    @abc.abstractmethod
    def get_statistic(self, statistics_type):
        return

    @abc.abstractmethod
    def get_all_statistics(self):
        return

    @abc.abstractmethod
    def get_result(self):
        return


class PE(Metric):
    """
    Abstract base class of pose error metrics.
    """

    def __init__(self):
        self.unit: Unit = Unit.none
        self.error: np.ndarray = np.array([])

    def __str__(self) -> str:
        return "PE metric base class"

    @abc.abstractmethod
    def process_data(self, data):
        return

    def change_unit(self, new_unit: Unit) -> None:
        if self.unit is new_unit:
            return

        if self.unit in (Unit.none, Unit.frames, Unit.percent, Unit.seconds):
            raise MetricsException(f"{self.unit} does not support conversions")

        bad_combinations = list(itertools.product(ANGLE_UNITS, LENGTH_UNITS))
        if any(
            combi in bad_combinations
            for combi in itertools.permutations((self.unit, new_unit))
        ):
            raise MetricsException(f"cannot convert {self.unit} to {new_unit}")

        if len(self.error) == 0:
            raise MetricsException(
                "error array is empty - "
                "please process data before changing the unit"
            )

        if self.unit in LENGTH_UNITS and new_unit in LENGTH_UNITS:
            # Convert first to meters, then to the final length unit.
            self.error *= (
                METER_SCALE_FACTORS[self.unit] / METER_SCALE_FACTORS[new_unit]
            )
        elif self.unit is Unit.radians and new_unit is Unit.degrees:
            self.error = np.rad2deg(self.error)
        elif self.unit is Unit.degrees and new_unit is Unit.radians:
            self.error = np.deg2rad(self.error)
        else:
            raise MetricsException(
                f"unknown unit combination {(self.unit, new_unit)}"
            )
        self.unit = new_unit

    def get_statistic(self, statistics_type: StatisticsType) -> float:
        if statistics_type == StatisticsType.rmse:
            squared_errors = np.power(self.error, 2)
            return math.sqrt(np.mean(squared_errors))
        elif statistics_type == StatisticsType.sse:
            squared_errors = np.power(self.error, 2)
            return np.sum(squared_errors)
        elif statistics_type == StatisticsType.mean:
            return float(np.mean(self.error))
        elif statistics_type == StatisticsType.median:
            return np.median(self.error)
        elif statistics_type == StatisticsType.max:
            return np.max(self.error)
        elif statistics_type == StatisticsType.min:
            return np.min(self.error)
        elif statistics_type == StatisticsType.std:
            return float(np.std(self.error))
        else:
            raise MetricsException("unsupported statistics_type")

    def get_all_statistics(self) -> dict[str, float]:
        """
        :return: a dictionary {StatisticsType.value : float}
        """
        statistics = {}
        for s in StatisticsType:
            try:
                statistics[s.value] = self.get_statistic(s)
            except MetricsException as e:
                if "unsupported statistics_type" not in str(e):
                    raise
        return statistics

    def get_result(
        self, ref_name: str = "reference", est_name: str = "estimate"
    ) -> Result:
        """
        Wrap the result in Result object.
        :param ref_name: optional, label of the reference data
        :param est_name: optional, label of the estimated data
        :return:
        """
        result = Result()
        metric_name = self.__class__.__name__
        result.add_info(
            {
                "title": str(self),
                "ref_name": ref_name,
                "est_name": est_name,
                "label": f"{metric_name} ({self.unit.value})",
            }
        )
        result.add_stats(self.get_all_statistics())
        if hasattr(self, "error"):
            result.add_np_array("error_array", self.error)
        return result


class RPE(PE):
    """
    RPE: relative pose error
    metric for investigating the odometry drift
    """

    def __init__(
        self,
        pose_relation: PoseRelation = PoseRelation.translation_part,
        delta: float = 1.0,
        delta_unit: Unit = Unit.frames,
        rel_delta_tol: float = 0.1,
        all_pairs: bool = False,
        pairs_from_reference: bool = False,
    ):
        if delta < 0:
            raise MetricsException("delta must be a positive number")
        if (
            delta_unit == Unit.frames
            and not isinstance(delta, int)
            and not delta.is_integer()
        ):
            raise MetricsException(
                f"delta must be integer for delta unit {delta_unit}"
            )
        self.delta = int(delta) if delta_unit == Unit.frames else delta
        self.delta_unit = delta_unit
        self.rel_delta_tol = rel_delta_tol
        self.pose_relation = pose_relation
        self.all_pairs = all_pairs
        self.pairs_from_reference = pairs_from_reference
        self.E: list[np.ndarray] = []
        self.error = np.array([])
        self.delta_ids: list[int] = []
        if pose_relation in (
            PoseRelation.translation_part,
            PoseRelation.point_distance,
        ):
            self.unit = Unit.meters
        elif pose_relation == PoseRelation.point_distance_error_ratio:
            self.unit = Unit.percent
        elif pose_relation == PoseRelation.rotation_angle_deg:
            self.unit = Unit.degrees
        elif pose_relation == PoseRelation.rotation_angle_rad:
            self.unit = Unit.radians
        else:
            # dimension-less
            self.unit = Unit.none

    def __str__(self) -> str:
        title = (
            f"RPE w.r.t. {self.pose_relation.value} ({self.unit.value})"
            f"\nfor delta = {self.delta} ({self.delta_unit.value})"
        )
        if self.all_pairs:
            title += " using all pairs"
        else:
            title += " using consecutive pairs"
        return title

    @staticmethod
    def rpe_base(
        Q_i: np.ndarray,
        Q_i_delta: np.ndarray,
        P_i: np.ndarray,
        P_i_delta: np.ndarray,
    ) -> np.ndarray:
        """
        Computes the relative SE(3) error pose for a single pose pair
        following the notation of the TUM RGB-D paper.
        :param Q_i: reference SE(3) pose at i
        :param Q_i_delta: reference SE(3) pose at i+delta
        :param P_i: estimated SE(3) pose at i
        :param P_i_delta: estimated SE(3) pose at i+delta
        :return: the RPE matrix E_i in SE(3)
        """
        Q_rel = lie.relative_se3(Q_i, Q_i_delta)
        P_rel = lie.relative_se3(P_i, P_i_delta)
        E_i = lie.relative_se3(Q_rel, P_rel)
        return E_i

    def process_data(self, data: PathPair) -> None:
        """
        Calculates the RPE on a batch of SE(3) poses from trajectories.
        :param data: tuple (traj_ref, traj_est) with:
        traj_ref: reference evo.trajectory.PosePath or derived
        traj_est: estimated evo.trajectory.PosePath or derived
        """
        if len(data) != 2:
            raise MetricsException(
                "please provide data tuple as: (traj_ref, traj_est)"
            )
        traj_ref, traj_est = data
        if traj_ref.num_poses != traj_est.num_poses:
            raise MetricsException(
                "trajectories must have same number of poses"
            )

        id_pairs = id_pairs_from_delta(
            (
                traj_ref.poses_se3
                if self.pairs_from_reference
                else traj_est.poses_se3
            ),
            self.delta,
            self.delta_unit,
            self.rel_delta_tol,
            all_pairs=self.all_pairs,
        )

        # Store flat id list e.g. for plotting.
        self.delta_ids = [j for i, j in id_pairs]

        if self.pose_relation in (
            PoseRelation.point_distance,
            PoseRelation.point_distance_error_ratio,
        ):
            # Only compares the magnitude of the point distance instead of
            # doing the full vector comparison of 'translation_part'.
            # Can be directly calculated on positions instead of full poses.
            ref_distances = np.array(
                [
                    np.linalg.norm(
                        traj_ref.positions_xyz[i] - traj_ref.positions_xyz[j]
                    )
                    for i, j in id_pairs
                ]
            )
            est_distances = np.array(
                [
                    np.linalg.norm(
                        traj_est.positions_xyz[i] - traj_est.positions_xyz[j]
                    )
                    for i, j in id_pairs
                ]
            )
            self.error = np.abs(ref_distances - est_distances)
            if self.pose_relation == PoseRelation.point_distance_error_ratio:
                nonzero = ref_distances.nonzero()[0]
                if nonzero.size != ref_distances.size:
                    logger.warning(
                        f"Ignoring {ref_distances.size - nonzero.size} zero "
                        "divisions in ratio calculations."
                    )
                    self.delta_ids = [self.delta_ids[i] for i in nonzero]
                self.error = (
                    np.divide(self.error[nonzero], ref_distances[nonzero])
                    * 100
                )
        else:
            # All other pose relations require the full pose error.
            self.E = [
                self.rpe_base(
                    traj_ref.poses_se3[i],
                    traj_ref.poses_se3[j],
                    traj_est.poses_se3[i],
                    traj_est.poses_se3[j],
                )
                for i, j in id_pairs
            ]

        logger.debug(
            f"Compared {len(self.E)} relative pose pairs, "
            f"delta = {self.delta} ({self.delta_unit.value}) with "
            f"{'all pairs.' if self.all_pairs else 'consecutive pairs.'}"
        )

        logger.debug(
            f"Calculating RPE for {self.pose_relation.value} pose relation..."
        )

        if self.pose_relation in (
            PoseRelation.point_distance,
            PoseRelation.point_distance_error_ratio,
        ):
            # Already computed, see above.
            pass
        elif self.pose_relation == PoseRelation.translation_part:
            self.error = np.array(
                [np.linalg.norm(E_i[:3, 3]) for E_i in self.E]
            )
        elif self.pose_relation == PoseRelation.rotation_part:
            # ideal: rot(E_i) = 3x3 identity
            self.error = np.array(
                [
                    np.linalg.norm(lie.so3_from_se3(E_i) - np.eye(3))
                    for E_i in self.E
                ]
            )
        elif self.pose_relation == PoseRelation.full_transformation:
            # ideal: E_i = 4x4 identity
            self.error = np.array(
                [np.linalg.norm(E_i - np.eye(4)) for E_i in self.E]
            )
        elif self.pose_relation == PoseRelation.rotation_angle_rad:
            self.error = np.array(
                [abs(lie.so3_log_angle(E_i[:3, :3])) for E_i in self.E]
            )
        elif self.pose_relation == PoseRelation.rotation_angle_deg:
            self.error = np.array(
                [abs(lie.so3_log_angle(E_i[:3, :3], True)) for E_i in self.E]
            )
        else:
            raise MetricsException(
                "unsupported pose_relation: ", self.pose_relation
            )


class APE(PE):
    """
    APE: absolute pose error
    metric for investigating the global consistency of a SLAM trajectory
    """

    def __init__(
        self, pose_relation: PoseRelation = PoseRelation.translation_part
    ):
        self.pose_relation = pose_relation
        self.E: list[np.ndarray] = []
        self.error = np.array([])
        if pose_relation in (
            PoseRelation.translation_part,
            PoseRelation.point_distance,
        ):
            self.unit = Unit.meters
        elif pose_relation == PoseRelation.rotation_angle_deg:
            self.unit = Unit.degrees
        elif pose_relation == PoseRelation.rotation_angle_rad:
            self.unit = Unit.radians
        else:
            self.unit = Unit.none  # dimension-less

    def __str__(self) -> str:
        title = "APE w.r.t. "
        title += (
            str(self.pose_relation.value)
            + " "
            + ("(" + self.unit.value + ")" if self.unit else "")
        )
        return title

    @staticmethod
    def ape_base(x_t: np.ndarray, x_t_star: np.ndarray) -> np.ndarray:
        """
        Computes the absolute error pose for a single SE(3) pose pair
        following the notation of the Kümmerle paper.
        :param x_t: estimated absolute pose at t
        :param x_t_star: reference absolute pose at t
        .:return: the delta pose
        """
        return lie.relative_se3(x_t, x_t_star)

    def process_data(self, data: PathPair) -> None:
        """
        Calculates the APE on a batch of SE(3) poses from trajectories.
        :param data: tuple (traj_ref, traj_est) with:
        traj_ref: reference evo.trajectory.PosePath or derived
        traj_est: estimated evo.trajectory.PosePath or derived
        """
        if len(data) != 2:
            raise MetricsException(
                "please provide data tuple as: (traj_ref, traj_est)"
            )
        traj_ref, traj_est = data
        if traj_ref.num_poses != traj_est.num_poses:
            raise MetricsException(
                "trajectories must have same number of poses"
            )

        if self.pose_relation in (
            PoseRelation.translation_part,
            PoseRelation.point_distance,
        ):
            # Translation part of APE is equivalent to distance between poses,
            # we don't require full SE(3) matrices for faster computation.
            self.E = traj_est.positions_xyz - traj_ref.positions_xyz
        else:
            self.E = [
                self.ape_base(x_t, x_t_star)
                for x_t, x_t_star in zip(
                    traj_est.poses_se3, traj_ref.poses_se3
                )
            ]
        logger.debug(f"Compared {len(self.E)} absolute pose pairs.")
        logger.debug(
            f"Calculating APE for {self.pose_relation.value} pose relation..."
        )

        if self.pose_relation in (
            PoseRelation.translation_part,
            PoseRelation.point_distance,
        ):
            # E is an array of position vectors only in this case
            self.error = np.array([np.linalg.norm(E_i) for E_i in self.E])
        elif self.pose_relation == PoseRelation.rotation_part:
            self.error = np.array(
                [
                    np.linalg.norm(lie.so3_from_se3(E_i) - np.eye(3))
                    for E_i in self.E
                ]
            )
        elif self.pose_relation == PoseRelation.full_transformation:
            self.error = np.array(
                [np.linalg.norm(E_i - np.eye(4)) for E_i in self.E]
            )
        elif self.pose_relation == PoseRelation.rotation_angle_rad:
            self.error = np.array(
                [abs(lie.so3_log_angle(E_i[:3, :3])) for E_i in self.E]
            )
        elif self.pose_relation == PoseRelation.rotation_angle_deg:
            self.error = np.array(
                [abs(lie.so3_log_angle(E_i[:3, :3], True)) for E_i in self.E]
            )
        else:
            raise MetricsException("unsupported pose_relation")


def id_pairs_from_delta(
    poses: typing.Sequence[np.ndarray],
    delta: float,
    delta_unit: Unit,
    rel_tol: float = 0.1,
    all_pairs: bool = False,
) -> filters.IdPairs:
    """
    high-level function - get index tuples of pairs with distance==delta
    from a pose list
    :param poses: list of SE(3) poses
    :param delta: the interval step for indices
    :param delta_unit: unit of delta (metrics.Unit enum member)
    :param rel_tol: relative tolerance to accept or reject deltas
    :param all_pairs: use all pairs instead of consecutive pairs
    :return: list of index tuples (pairs)
    """
    if delta_unit == Unit.frames:
        id_pairs = filters.filter_pairs_by_index(poses, int(delta), all_pairs)
    elif delta_unit == Unit.meters:
        id_pairs = filters.filter_pairs_by_path(
            poses, delta, delta * rel_tol, all_pairs
        )
    elif delta_unit in {Unit.degrees, Unit.radians}:
        use_degrees = delta_unit == Unit.degrees
        id_pairs = filters.filter_pairs_by_angle(
            poses, delta, delta * rel_tol, use_degrees, all_pairs
        )
    else:
        raise filters.FilterException(f"unsupported delta unit: {delta_unit}")

    if len(id_pairs) == 0:
        raise filters.FilterException(
            f"delta = {delta} ({delta_unit.value}) produced an empty index list - try lower values "
            f"or a less strict tolerance"
        )

    logger.debug(
        f"Found {len(id_pairs)} pairs with delta {delta} ({delta_unit.value}) "
        f"among {len(poses)} poses "
        + ("using consecutive pairs." if not all_pairs else "using all pairs.")
    )

    return id_pairs