import itertools
from typing import Callable, Mapping, Union

import numba as nb
import numpy as np
from pymoo.algorithms.nsga2 import NSGA2
from pymoo.factory import get_crossover, get_mutation, get_sampling
from pymoo.model.problem import Problem
from pymoo.optimize import minimize

MetricFunction = Callable[[np.ndarray], float]
"""Type hint describing a function which accepts an ndarray and calculates some metric about it."""


class ArrayMetric:
    @staticmethod
    @nb.njit(cache=True)
    def inversions(arr: np.ndarray) -> float:
        """
        Calculates the _total_ number of 'inversions' within an array,
        then normalizes the value based on the array length.

        Parameters
        ----------
        arr

        Returns
        -------

        """
        l = arr.shape[0]
        if l < 2:
            return 0
        inv = 0
        for i in range(l):
            for j in range(i + 1, l):
                if arr[i] > arr[j]:
                    inv += 1
        total_inv = l * (l - 1) / 2

        return 1 - (inv / total_inv)

    @staticmethod
    def adjacent_unsorted(arr: np.ndarray) -> float:
        diffs = np.diff(arr)
        asc_sorted_count = diffs[diffs > 0].sum()
        desc_sorted_count = diffs[diffs < 0].sum()

        # We want this metric to be symmetric, so it doesn't matter
        # whether the array is sorted in ascending or descending order.
        # To achieve that, take the larger of the number of ASC and DESC
        # sorted counts.
        sorted_count = max(asc_sorted_count, desc_sorted_count)

        # Normalize over the number of unique elements in the array
        # to account for multiple occurrences of a given element; this is because
        # if we sorted the array, any single instance of an element occurring
        # multiple times could be a range of indices, all of which would be correct.
        # TODO: Need to somehow normalize this over the whole [0.0, 1.0] range,
        #       currently it'd only (it seems) cover from [0.5, 1.0]
        return sorted_count / len(np.unique(arr))

    @staticmethod
    def mean_displacement(arr: np.ndarray) -> float:
        """
        Calculates the mean 'displacement' of the array elements from the position they'd
        be in if the array were to be sorted.

        Notes
        -----
        This function accounts for multiple occurrences of an element, and is also symmetric
        -- reversing an array should result in the same score (to make it so that it's irrelevant
        whether the array is sorted by ascending or descending values).
        """
        # We're dealing with integer key arrays (integers from 0..N), if we called 'unique' on the input
        # we should get the natural numbers up to N back. This allows us to better handle multiple occurrences
        # of a given key -- instead of saying the key could occur at indices X..Y (where Y-X is the number of
        # occurrences of that specific key), we just treat those indices X..Y as the same value. And by the logic
        # above, that means that elements 0..N would occur at indices 0..N, so each element represents the index
        # it'd fall at in the sorted array.
        # Argsort the array, then subtract the value of each element (in the argsort result) from it's index;
        # this gives the "displacement" of each element from it's index if the array was fully sorted.
        # then compute the mean over the differences and return that.
        pass

    @staticmethod
    @nb.njit(cache=True)
    def lnds(arr: np.ndarray) -> float:
        len = arr.shape[0]
        L = np.zeros(len)
        L[0] = 1
        for i in range(1, len):
            for j in range(i):
                if arr[j] <= arr[i]:
                    L[i] = max(L[i], L[j] + 1)

        return np.max(L) / len

    @staticmethod
    @nb.njit(cache=True)
    def pw_diff_sum(arr: np.ndarray) -> float:
        l = arr.shape[0]
        if l < 2:
            return 0
        n = np.min(arr)
        x = np.max(arr)
        diff_sums = np.array([np.sum(np.abs(arr - a)) for a in arr])
        max_sum = (x - n) * l * (l - 1) / 2
        return np.sum(diff_sums) / max_sum

    @staticmethod
    @nb.njit(cache=True)
    def mean_seq_diff(arr: np.ndarray) -> float:
        diffs = np.diff(arr)
        m = np.mean(np.abs(diffs))
        x = np.max(diffs)
        return 1 if (m == 0 or x == 0) else 1 - m / x

    @staticmethod
    @nb.njit(cache=True)
    def mean_abs_dev(arr: np.ndarray) -> float:
        l = arr.shape[0]
        if l < 2:
            return 0
        m = np.mean(arr)
        n = np.min(arr)
        x = np.max(arr)
        abs_devs = np.abs(arr - m)
        dev = np.mean(abs_devs)
        max_dev = ((m - n) + (x - m)) / 2
        return 0 if max_dev == 0 else dev / max_dev

    @staticmethod
    @nb.njit(cache=True)
    def var_ratio(arr: np.ndarray) -> float:
        l = arr.shape[0]
        if l < 2:
            return 0
        return np.var(arr) * 4 / np.square(np.max(arr) - np.min(arr))

    @staticmethod
    @nb.njit(cache=True)
    def davies_bouldin(arr: np.ndarray) -> float:
        """
        Calculates the Davies-Bouldin index of a key array. This provides a measurement of "clustering".

        Parameters
        ----------
        arr : np.ndarray
            The key array to calculate the Davies-Bouldin index over.

        Returns
        -------
        float
            The calculated Davies-Bouldin index of `arr`.

        Notes
        -----
        * https://en.wikipedia.org/wiki/Cluster_analysis#Internal_evaluation
        * https://en.wikipedia.org/wiki/Davies%E2%80%93Bouldin_index
        * https://scikit-learn.org/stable/modules/generated/sklearn.metrics.davies_bouldin_score.html#sklearn.metrics.davies_bouldin_score

        TODO: This could be implemented more efficiently if the input were a riptable Grouping object
        instead of a raw key array.

        See Also
        --------
        ArrayMetric.dunn
        ArrayMetric.silhouette_coeff
        """
        if len(arr) == 0:
            return np.nan

        # Get the max value in the array -- this tells us how many keys there are.
        key_count = arr.max() + 1

        # Arrays to hold the sum of the indices where each key is found, and the number of occurrences of that key.
        key_index_sums = np.zeros(key_count, dtype=np.uint64)
        key_counts = np.zeros(key_count, dtype=np.uint64)

        for i in range(len(arr)):
            key = arr[i]
            key_index_sums[key] += i
            key_counts[key] += 1

        # Compute the centroid (mean index) for each key (representing a cluster).
        centroids = key_index_sums.astype(np.float64) / key_counts.astype(np.float64)
        del key_index_sums

        # For each key/cluster, determine the average (mean) distance of the points (in our case,
        # indices of the occurrences of the key) to the centroid for that key.
        # We actually just compute the total distance from each point to the centroid (per key); since we
        # know how many times each key occurs we can easily compute the average later when needed, and this
        # avoids an extra array allocation / unnecessary memory usage.
        total_dist_to_centroid = np.zeros(key_count, dtype=np.float64)
        for i in range(len(arr)):
            key = arr[i]
            # abs() is used here since we're computing a Euclidian distance (L2 norm) over 1D.
            total_dist_to_centroid[key] += abs(centroids[key] - i)

        # The summation value (added to for each key).
        summation = 0.0

        # Iterate over pairs of keys (direction doesn't matter, so this is a "triangular" nested loop).
        for i in range(key_count - 1):
            at_least_one_pair_seen = False
            # The max value of the intra-vs-intra-cluster distance metric for all j values given this i value.
            max_ij_distance_ratio = np.nan

            sigma_i = total_dist_to_centroid[i] / key_counts[i]

            for j in range(i + 1, key_count):
                sigma_j = total_dist_to_centroid[j] / key_counts[j]

                # Distance over 1D (array indices) between centroids for key/cluster i and j.
                # abs() is used to calculate Euclidian distance (L2 norm) over 1D.
                ij_centroid_dist = abs(centroids[i] - centroids[j])

                # Calculate the intra-vs-inter cluster metric for (i, j).
                ij_distance_ratio = (sigma_i + sigma_j) / ij_centroid_dist

                # If this is the first trip through the loop, set the max distance ratio.
                # Otherwise, if the ratio for this (i, j) pair is larger than the max value
                # we've seen for all other j's (given the current value of i), update the max value.
                if not at_least_one_pair_seen:
                    max_ij_distance_ratio = ij_distance_ratio
                    at_least_one_pair_seen = True
                elif ij_distance_ratio > max_ij_distance_ratio:
                    max_ij_distance_ratio = ij_distance_ratio

            if at_least_one_pair_seen:
                summation += max_ij_distance_ratio

        # Return the average of the summed max-ratios.
        return summation / key_count

    @staticmethod
    @nb.njit(cache=True)
    def dunn(arr: np.ndarray) -> float:
        """
        Calculates the Dunn index of a key array. This provides a measurement of "clustering".

        Parameters
        ----------
        arr : np.ndarray
            The key array to calculate the Dunn index for.

        Returns
        -------
        float
            The calculated Dunn index of `arr`.

        Notes
        -----
        * https://en.wikipedia.org/wiki/Cluster_analysis#Internal_evaluation
        * https://en.wikipedia.org/wiki/Dunn_index

        TODO: This could be implemented more efficiently if the input were a riptable Grouping object
        instead of a raw key array.

        See Also
        --------
        ArrayMetric.davies_bouldin
        ArrayMetric.silhouette_coeff
        """
        if len(arr) == 0:
            return np.nan

        # Get the max value in the array -- this tells us how many keys there are.
        key_count = arr.max() + 1

        # Arrays to hold the sum of the indices where each key is found, and the number of occurrences of that key.
        key_index_sums = np.zeros(key_count, dtype=np.uint64)
        key_counts = np.zeros(key_count, dtype=np.uint64)

        for i in range(len(arr)):
            key = arr[i]
            key_index_sums[key] += i
            key_counts[key] += 1

        # Compute the minimum inter-cluster distance, where the distance is defined
        # as the Euclidian distance between the centroids of the clusters.
        min_intercluster_dist = 1e50
        for i in range(key_count - 1):
            centroid_i = np.float64(key_index_sums[i]) / key_counts[i]
            for j in range(i + 1, key_count):
                centroid_j = np.float64(key_index_sums[j]) / key_counts[j]

                # Euclidian distance over 1D between the two centroids.
                ij_centroid_dist = abs(centroid_i - centroid_j)

                # If the distance between the centroids for keys/clusters i and j
                # is smaller than the current min, update the min.
                if ij_centroid_dist < min_intercluster_dist:
                    min_intercluster_dist = ij_centroid_dist

        del key_index_sums

        # Compute the maximum intra-cluster distance.
        # There are various ways to define intra-cluster distance, and some may provide "better" results than others.
        # E.g. mean of pairwise distances between all points in a cluster.
        # See the article (linked in docstring) for details.
        # Here we use the maximum distance between any two points in the cluster;
        # we're working in 1D, so the extremes for each key/cluster are the first and
        # last occurrences of each key within the input array.
        first_key_index = np.full(key_count, len(arr), dtype=np.int64)
        last_key_index = np.zeros(key_count, dtype=np.int64)

        for i in range(len(arr)):
            key = arr[i]
            if first_key_index[key] > i:
                # This is the first occurrence of this key, so save the array index.
                first_key_index[key] = i

            # Always update the last (greatest) index at which we saw a given key.
            last_key_index[key] = i

        max_intracluster_dist = 0
        for i in range(key_count):
            intracluster_dist = last_key_index[i] - first_key_index[i]

            # Update the max intracluster distance if this cluster's distance is larger
            # than any we've seen so far.
            if intracluster_dist > max_intracluster_dist:
                max_intracluster_dist = intracluster_dist

        # The Dunn index is defined as the ratio between the minimal inter-cluster distance
        # and the maximal intra-cluster distance.
        return min_intercluster_dist / max_intracluster_dist

    @staticmethod
    def silhouette_coeff(arr: np.ndarray) -> float:
        """
        Calculates the silhouette coefficient of a key array. This provides a measurement of "clustering".

        Parameters
        ----------
        arr : np.ndarray
            The key array to calculate the silhouette coefficient for.

        Returns
        -------
        float
            The calculated silhouette coefficient of `arr`.

        Notes
        -----
        * https://en.wikipedia.org/wiki/Cluster_analysis#Internal_evaluation
        * https://en.wikipedia.org/wiki/Silhouette_(clustering)
        * https://scikit-learn.org/stable/modules/generated/sklearn.metrics.silhouette_score.html#sklearn.metrics.silhouette_score

        TODO: This could be implemented more efficiently if the input were a riptable Grouping object
        instead of a raw key array.

        See Also
        --------
        ArrayMetric.davies_bouldin
        ArrayMetric.dunn
        """
        # TODO: Instead of implementing this ourselves, can just call the scikit-learn implementation (link in docstring).
        pass


def calculate_metrics(arr: np.ndarray, metrics: Mapping[str, MetricFunction]) -> Mapping[str, float]:
    """
    Given an array and a table of metric/scoring functions, calculates the metrics for the array.
    """
    return {k: f(arr) for k, f in metrics}


_metrics: Mapping[str, MetricFunction] = {
    "inversions": ArrayMetric.inversions,
    "adjacent_unsorted": ArrayMetric.adjacent_unsorted,
    "lnds": ArrayMetric.lnds,
    "pw_diff_sum": ArrayMetric.pw_diff_sum,
    "mean_seq_diff": ArrayMetric.mean_seq_diff,
    "mean_abs_dev": ArrayMetric.mean_abs_dev,
    "var_ratio": ArrayMetric.var_ratio,
    "davies_bouldin": ArrayMetric.davies_bouldin,
    "dunn_index": ArrayMetric.dunn,
    # "silhouette_coeff": ArrayMetric.silhouette_coeff
}


def score_array(arr: np.ndarray) -> Mapping[str, float]:
    """Scores an array based on a pre-defined set of metric functions."""
    return calculate_metrics(arr, _metrics)


def batch_index_gen_error(
    lo: Union[int, np.int32, np.int64],
    hi: Union[int, np.int32, np.int64],
    size: Union[int, np.int32, np.int64],
    s_metric: MetricFunction,
    c_metric: MetricFunction,
    d_metric: MetricFunction,
) -> None:
    ss = np.linspace(0, 1, 6)
    cs = np.linspace(0, 1, 6)
    ds = np.linspace(0, 1, 6)
    mets = itertools.product(ss, cs, ds)
    index_count = 0
    s_error = 0
    c_error = 0
    d_error = 0
    for (
        s,
        c,
        d,
    ) in mets:
        indexes = gen_index_MOO(lo, hi, size, s_metric, c_metric, d_metric, sort=s, clustered=c, dispersed=d)
        index_count += indexes.shape[0]
        s_error_curr = 0
        c_error_curr = 0
        d_error_curr = 0
        for index in indexes:
            s_exp = s
            c_exp = c
            d_exp = d

            s_obs = s_metric(index)
            c_obs = c_metric(index)
            d_obs = d_metric(index)

            s_error_curr += abs(s_exp - s_obs)
            c_error_curr += abs(c_exp - c_obs)
            d_error_curr += abs(d_exp - d_obs)

            s_error += s_error_curr
            c_error += c_error_curr
            d_error += d_error_curr

        print(
            f"batch errors -  s: {s_error_curr/indexes.shape[0]}    c: {c_error_curr/indexes.shape[0]}   d: {d_error_curr/indexes.shape[0]}    # indexes: {indexes.shape[0]}"
        )

    s_error /= index_count
    c_error /= index_count
    d_error /= index_count

    print(f"s_metric avg error: {s_error}    c_metric avg error: {c_error}     d_metric avg error: {d_error}")


class SCDArrayGen(Problem):
    def __init__(self, lo, hi, index_size, s_metric, c_metric, d_metric, s, c, d):
        self.lo = lo
        self.hi = hi
        self.index_size = index_size
        self.s_metric = s_metric
        self.c_metric = c_metric
        self.d_metric = d_metric
        self.s = s
        self.c = c
        self.d = d

        xl = np.full(index_size, self.lo)
        xu = np.full(index_size, self.hi)

        super().__init__(n_var=self.index_size, n_constr=2, xl=xl, xu=xu, type_var=int)

    def s_error(self, arr):
        return abs(self.s_metric(arr) - self.s)

    def c_error(self, arr):
        return abs(self.c_metric(arr) - self.c)

    def d_error(self, arr):
        return abs(self.d_metric(arr) - self.d)

    def _evaluate(self, x, out, *args, **kwargs):

        # objectives
        se = np.apply_along_axis(self.s_error, 1, x)
        ce = np.apply_along_axis(self.c_error, 1, x)
        de = np.apply_along_axis(self.d_error, 1, x)

        # constraints
        g1 = self.lo - x
        g2 = x - self.hi

        out["F"] = np.column_stack([se, ce, de])
        out["G"] = np.column_stack([g1, g2])


# Evolutionary, Black-Box, Multi-Objective Optimization Algorithm
def gen_index_MOO(lo, hi, size, s_metric, c_metric, d_metric, dist=None, sort=0.0, clustered=0.0, dispersed=0.0):
    max_error = 0.25
    prob = SCDArrayGen(lo, hi, size, s_metric, c_metric, d_metric, sort, clustered, dispersed)
    alg = NSGA2(
        pop_size=100,
        sampling=get_sampling("int_random"),
        crossover=get_crossover("int_two_point"),
        mutation=get_mutation("int_pm"),
    )
    res = minimize(prob, alg, seed=3)

    # optionally screen out solutions that still deviated from the desired values by too much
    # mask = res.F < max_error
    # masksum = np.sum(mask, axis=1) == 3
    return res.X


if __name__ == "__main__":
    # test
    # print(gen_index_MOO(0, 150, 500, ArrayMetric.inversions, ArrayMetric.pw_diff_sum, ArrayMetric.mean_abs_dev, sort=0.5, clustered=0.2, dispersed=0.8))
    lo = 0
    hi = 10000
    size = 100
    s_metric = ArrayMetric.inversions
    c_metric = ArrayMetric.pw_diff_sum
    d_metric = ArrayMetric.mean_abs_dev
    batch_index_gen_error(lo, hi, size, s_metric, c_metric, d_metric)
