"""
Tracker utilities
"""
import numpy as np

from scipy.optimize import linear_sum_assignment
from scipy.linalg import solve_triangular

from .utils import boxes_iou


# infinity cost matrix value
INFINITY_COST = 1e5
# inverse chi squared distribution
chi2inv95 = {
    1: 3.8415,
    2: 5.9915,
    3: 7.8147,
    4: 9.4877,
    5: 11.070,
    6: 12.592,
    7: 14.067,
    8: 15.507,
    9: 16.919
}


def keypoints_to_bbox(keypoints, scale=1.0):
    """Convert keypoints to square bounding box

    Args:
        keypoints (np.ndarray): input keypoints
        scale (float, optional): bounding box scale ratio. Defaults to 1.0.

    Returns:
        np.array: constructed bounding box in [x_center, y_center, width, height], where width == height
    """
    # min and max coordinates
    x1, y1 = np.min(keypoints, axis=0)
    x2, y2 = np.max(keypoints, axis=0)
    # convert to [x, y, w, h] format
    w, h = x2 - x1, y2 - y1
    x, y = x1 + w / 2, y1 + h / 2
    # convert to square, based on the biggest rectangle side
    side = w if w > h else h
    # scale bounding box size
    side *= scale
    return np.array([x, y, side, side], dtype=np.float32)


def gating_distance(track, measurements):
    """Gate distance between track and measurements with squared Mahalanobis distance

    Args:
        track (Track): input track
        measurements (_type_): _description_

    Returns:
        np.array: distance between track distribution and measurements
    """
    # project track state and covariance to measurement space
    mean, covariance = track.box_kalman.project()
    # solve Mahalanobis distance between detections and track state distribution
    cholesky_factor = np.linalg.cholesky(covariance[:3, :3])
    d = measurements[:, :3] - mean[:3]
    z = solve_triangular(cholesky_factor, d.T, lower=True, check_finite=False, overwrite_b=True)
    squared_maha = np.sum(z * z, axis=0)
    return squared_maha


def gate_cost_matrix(cost_matrix, tracks, detections, track_indices, detection_indices):
    """Gate cost matrix with squared Mahalanobis distance

    Args:
        cost_matrix (np.ndarray): assignment cost matrix
        tracks (list): list of tracked objects
        detections (list): list of detected objects
        track_indices (list): list of indices of inspected tracks
        detection_indices (list): list of indices of inspected detections

    Returns:
        np.ndarray: gated cost matrix
    """
    # gating threshold from inverse chi squared
    gating_threshold = chi2inv95[3]
    # convert measurements states to array
    measurements = np.asarray([detections[i].state_vector() for i in detection_indices])
    # for every track
    for row, track_idx in enumerate(track_indices):
        track = tracks[track_idx]
        # calculate gating distance and set cost matrix to infinity cost if gating distance is larger then threshold
        distance = gating_distance(track, measurements)
        cost_matrix[row, distance > gating_threshold] = INFINITY_COST
    return cost_matrix


def euclidean_distance(detection_features, track_features):
    """Squared distance between track and detection feature matrices (keypoints)

    Args:
        detection_features (list): list of feature matrices of detected objects
        track_features (list): list of feature matrices of tracked objects

    Returns:
        np.ndarray: distance matrix between tracks and detections
    """
    # convert detection features to tensor of [1, N, M, 2] - N is number of objects and M is number of keypoints
    detection_features = np.asarray(detection_features)[np.newaxis, :, ...]
    # convert track features to tensor of [N, 1, M, 2] - N is number of objects and M is number of keypoints
    track_features = np.asarray(track_features)[:, np.newaxis, ...]
    # if either tracks or detections are empty, return empty matrix
    if len(detection_features) == 0 or len(track_features) == 0:
        return np.zeros((len(detection_features), len(track_features)))
    # distance between features, clipped for numerical stability
    distance = np.sqrt(np.square(detection_features - track_features).sum(axis=(-1, -2)))
    distance = np.clip(distance, 1e-5, float(np.inf))
    return distance


def gate_cost(tracks, detections, track_indices, detection_indices):
    """Gated assignment cost matrix

    Args:
        tracks (list): list of tracked objects
        detections (list): list of detected objects
        track_indices (list): list of indices of inspected tracks
        detection_indices (list): list of indices of inspected detections

    Returns:
        np.ndarray: distance matrix between tracks and detections
    """
    # convert features to tensor
    detection_features = np.array([detections[i].feature for i in detection_indices])
    track_features = np.array([tracks[i].feature for i in track_indices])
    # calculate cost matrix and gate it
    cost_matrix = euclidean_distance(detection_features, track_features)
    # cost_matrix = gate_cost_matrix(cost_matrix, tracks, detections, track_indices, detection_indices)
    return cost_matrix


def cascade_match(max_distance, cascade_depth, tracks, detections, track_indices, detection_indices):
    """Cascade matching algorithm for object assignment, first assign tracked objects that have been updated more recently

    Args:
        max_distance (float): max distance between assigned objects, assignments overstepping this boundary are invalid
        cascade_depth (int): number of cascade levels
        tracks (list): list of tracked objects
        detections (list): list of detected objects
        track_indices (list): list of indices of inspected tracks
        detection_indices (list): list of indices of inspected detections

    Returns:
        list, list, list: lists of matched indices, unmatched track indices and unmatched detection indices
    """
    # init all detections as unmatched
    unmatched_detections = detection_indices
    matches = list()
    # for every cascade level
    for level in range(cascade_depth):
        # no detections left -> terminate matching
        if len(unmatched_detections) == 0:
            break
        # tracks at current cascade level
        track_indices_level = [i for i in track_indices if tracks[i].time_since_update == 1 + level]
        # nothing to match at this level -> next level
        if len(track_indices_level) == 0:
            continue
        # match selected tracks with unmatched detections
        matches_level, _, unmatched_detections = min_cost_match(gate_cost, max_distance, tracks, detections, track_indices_level, unmatched_detections)
        # add matches to list, unmatched detections are updated with assignment
        matches += matches_level
    # unmatched tracks are track indices that are not present in matches
    unmatched_tracks = list(set(track_indices) - set(i for i, _ in matches))
    return matches, unmatched_tracks, unmatched_detections


def iou_cost(tracks, detections, track_indices, detection_indices):
    """Intersection over Union cost matrix

    Args:
        tracks (list): list of tracked objects
        detections (list): list of detected objects
        track_indices (list): list of indices of inspected tracks
        detection_indices (list): list of indices of inspected detections

    Returns:
        np.ndarray: iou matrix between tracks and detections
    """
    # get tracked objects' bounding boxes, transform them to [N, 1, 4] where N is number of tracked objects
    tracks_boxes = np.array([tracks[i].bbox for i in track_indices], dtype=np.float32).reshape(-1, 1, 4)
    # get detected objects' bounding boxes, transform them to [1, N, 4] where N is number of detected objects
    detection_boxes = np.array([detections[i].bbox for i in detection_indices], dtype=np.float32).reshape(1, -1, 4)
    # calculate iou and inverse it so the results is cost
    cost = 1.0 - boxes_iou(tracks_boxes, detection_boxes, center_format=True, scale=1.15)
    # set iou cost of objects that have not been updated in the last time step to infinity
    track_invalid_mask = [tracks[i].time_since_update > 1 for i in track_indices]
    cost[track_invalid_mask, :] = INFINITY_COST
    return cost


def min_cost_match(cost_func, max_distance, tracks, detections, track_indices, detection_indices):
    """Perform cost matrix assignment

    Args:
        cost_func (callable): function constructing cost matrix of passed tracks and detections
        max_distance (float): max distance between assigned objects, assignments overstepping this boundary are invalid
        tracks (list): list of tracked objects
        detections (list): list of detected objects
        track_indices (list): list of indices of inspected tracks
        detection_indices (list): list of indices of inspected detections

    Returns:
        list, list, list: lists of matched indices, unmatched track indices and unmatched detection indices
    """
    # if either detections or tracks are empty, don't perform assignment
    if len(detection_indices) == 0 or len(track_indices) == 0:
        return list(), track_indices, detection_indices
    # calculate cost matrix
    cost_matrix = cost_func(tracks, detections, track_indices, detection_indices)
    # check maximum distance
    cost_matrix[cost_matrix > max_distance] = max_distance + 1e-5
    # perform assignment
    rows, cols = linear_sum_assignment(cost_matrix)
    # convert assignment results to lists of matches, unmatched tracks and unmatched detections
    matches, unmatched_tracks, unmatched_detections = list(), list(), list()
    # unmatched detection are detection indices absent from assignment columns
    for col, detection_idx in enumerate(detection_indices):
        if col not in cols:
            unmatched_detections.append(detection_idx)
    # unmatched tracks are track indices absent from assignment rows
    for row, track_idx in enumerate(track_indices):
        if row not in rows:
            unmatched_tracks.append(track_idx)
    # matches from returned assignment matrix
    for row, col in zip(rows, cols):
        track_idx = track_indices[row]
        detection_idx = detection_indices[col]
        # if maximum distance is crossed -> add objects to unmatched lists
        if cost_matrix[row, col] > max_distance:
            unmatched_tracks.append(track_idx)
            unmatched_detections.append(detection_idx)
        # distance is withing boundary -> add to matches
        else:
            matches.append((track_idx, detection_idx))
    return matches, unmatched_tracks, unmatched_detections
