import math

import numpy as np
from numba import njit


def cc(s_map, gt):
    s_map_norm = (s_map - np.mean(s_map)) / np.std(s_map)
    gt_norm = (gt - np.mean(gt)) / np.std(gt)
    a = s_map_norm
    b = gt_norm
    r = (a * b).sum() / math.sqrt((a * a).sum() * (b * b).sum())
    return r


def NSS(saliencyMap, fixationMap, msg=False):
    if not saliencyMap.max() == 0:
        saliencyMap = saliencyMap.astype(float) / saliencyMap.max()
    if not saliencyMap.std(ddof=1) == 0:
        saliencyMap = (saliencyMap - saliencyMap.mean()) / saliencyMap.std(ddof=1)
    score = saliencyMap[fixationMap.astype(bool)].mean()
    return score


def kldiv(s_map, gt):
    s_map = s_map / (np.sum(s_map) * 1.0)
    gt = gt / (np.sum(gt) * 1.0)
    eps = 2.2204e-16
    return np.sum(gt * np.log(eps + gt / (s_map + eps)))


@njit
def euclidean_distance(human_scanpath, simulated_scanpath):
    dist = np.zeros(len(human_scanpath))
    for i in range(len(human_scanpath)):
        P = human_scanpath[i]
        Q = simulated_scanpath[i]
        dist[i] = np.sqrt((P[0] - Q[0]) ** 2 + (P[1] - Q[1]) ** 2)
    return dist.sum()


def _Levenshtein_Dmatrix_initializer(len1, len2):
    Dmatrix = []
    for i in range(len1):
        Dmatrix.append([0] * len2)
    for i in range(len1):
        Dmatrix[i][0] = i
    for j in range(len2):
        Dmatrix[0][j] = j
    return Dmatrix


def _Levenshtein_cost_step(Dmatrix, string_1, string_2, i, j, substitution_cost=1):
    char_1 = string_1[i - 1]
    char_2 = string_2[j - 1]
    # insertion
    insertion = Dmatrix[i - 1][j] + 1
    # deletion
    deletion = Dmatrix[i][j - 1] + 1
    # substitution
    substitution = Dmatrix[i - 1][j - 1] + substitution_cost * (char_1 != char_2)
    # pick the cheapest
    Dmatrix[i][j] = min(insertion, deletion, substitution)


def _Levenshtein(string_1, string_2, substitution_cost=1):
    # get strings lengths and initialize Distances-matrix
    len1 = len(string_1)
    len2 = len(string_2)
    Dmatrix = _Levenshtein_Dmatrix_initializer(len1 + 1, len2 + 1)

    # compute cost for each step in dynamic programming
    for i in range(len1):
        for j in range(len2):
            _Levenshtein_cost_step(
                Dmatrix,
                string_1,
                string_2,
                i + 1,
                j + 1,
                substitution_cost=substitution_cost,
            )

    return Dmatrix[len1][len2]


def _scanpath_to_string(scanpath, height, width, depth, n):
    height_step, width_step, depth_step = height // n, width // n, depth // n
    string = ""
    for i in range(np.shape(scanpath)[0]):
        fixation = scanpath[i].astype(np.int32)
        correspondent_square = (
            (max(0, fixation[0]) // width_step)
            + (max(0, fixation[1]) // height_step) * n
            + (max(0, fixation[2]) // depth_step) * n * n
        )
        string += chr(correspondent_square)
    return string


def string_edit_distance(
    stimulus,  # matrix
    human_scanpath,
    simulated_scanpath,
    n=5,
    substitution_cost=1,
    msg=False,
):
    height, width, depth = np.shape(stimulus)[0:3]
    string_1 = _scanpath_to_string(human_scanpath, height, width, depth, n)
    string_2 = _scanpath_to_string(simulated_scanpath, height, width, depth, n)

    if msg:
        print((string_1, string_2))

    return _Levenshtein(string_1, string_2)
