import os

import numpy as np
from dtaidistance import dtw
from cobras_ts.cobras_dtw import COBRAS_DTW
from cobras_ts.querier.labelquerier import LabelQuerier
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import DBA as dba
from matplotlib.patches import Ellipse
import datetime as dt
from queue import Queue
import random


ucr_path = '/Users/ary/Desktop/Thesis/UCR_TS_Archive_2015'
dataset = 'ECG200'
#dataset = 'Beef'
budget = 50
alpha = 0.5
window = 1
iteration = 1
cluster_idx = 0



# load the data
data = np.loadtxt(os.path.join(ucr_path,dataset,dataset + '_TEST'), delimiter=',')
data_train = np.loadtxt(os.path.join(ucr_path,dataset,dataset + '_TRAIN'), delimiter=',')
data = np.vstack((data,data_train))
series = data[:,1:][:20]
labels = data[:,0]

start = 0
end = series.shape[1]

print("series cnt = " + str(series.shape[0]))
print("series shape = " + str(series.shape[1]))

# construct the affinity matrix
dt1 = dt.datetime.now()
dists = dtw.distance_matrix(series, window=int(0.01 * window * series.shape[1]))
dt2 = dt.datetime.now()
print("distance matrix use time " + str((dt2 - dt1).seconds))
dists[dists == np.inf] = 0
dists = dists + dists.T - np.diag(np.diag(dists))
affinities = np.exp(-dists * alpha)

# initialise cobras_dtw with the precomputed affinities
clusterer = COBRAS_DTW(affinities, LabelQuerier(labels), budget)

clustering, intermediate_clusterings, runtimes, ml, cl = clusterer.cluster()
dt3 = dt.datetime.now()
print("cluster use time " + str((dt3 - dt2).seconds))

print("there are " + str(len(clustering.clusters)) + " clusters")

x = np.arange(0, series.shape[1])

def initFigure(fig):
    ax1 = fig.add_subplot(231)  # DBA
    ax2 = fig.add_subplot(232, sharex=ax1, sharey=ax1)  # adjusted dots
    ax3 = fig.add_subplot(233, sharex=ax1, sharey=ax1)  # eclipse
    ax4 = fig.add_subplot(234, sharex=ax1)  # dtw_vertical
    ax5 = fig.add_subplot(235, sharex=ax1, sharey=ax4)  # vertical
    ax6 = fig.add_subplot(236, sharex=ax1, sharey=ax4)  # dtw_horizontal

    ax_overall = [ax1, ax2, ax3]
    ax_statistic = [ax4, ax5, ax6]
    axs = ax_overall + ax_statistic

    # axis color
    for ax in axs:
        ax.spines['bottom'].set_color('blue')
        ax.spines['bottom'].set_linewidth(2)
        ax.spines['left'].set_linewidth(2)
        ax.set_xlabel("time")

    for ax in ax_overall:
        ax.spines['left'].set_color('pink')
    for ax in ax_statistic:
        ax.spines['left'].set_color('yellow')

    x = np.arange(0, series.shape[1])

    # y_label
    ax1.set_ylabel("DBA", color="purple")
    ax2.set_ylabel("Instances", color="blue")
    ax3.set_ylabel("DBA with Warping Deviations", color="blue")
    ax6.set_ylabel("Horizontal Warping Deviation", color="green")
    ax4.set_ylabel("Vertical Warping Deviation", color="red")
    ax5.set_ylabel("Standard Deviation", color="gray")
    return ax1, ax2, ax3, ax4, ax5, ax6

def findAllIndicesOfOneCluster(cluster_idx):
    superinstances = clustering.clusters[cluster_idx].super_instances
    indices_of_current_cluster = []
    for superinstance in superinstances:
        indices_of_current_cluster += superinstance.indices
    return indices_of_current_cluster

def findAllSeriesWithIndices(indices):
    target_series = list(map(lambda idx: series[idx], indices))
    return target_series

def getColorIter(size):
    color_arr = cm.rainbow(np.linspace(0, 1, size))
    colors = iter(color_arr)
    return colors


def plotDBA(plt, series_mean, range = None):
    if (range == None):
        plt.plot(x, series_mean, color='purple')
    elif (range[0] + 1 == range[1]):
        plt.scatter(x[range[0]], series_mean[range[0]], color = 'purple')
    else:
        plt.plot(x[range[0]: range[1]], series_mean[range[0]: range[1]], color = 'purple')

def plotEclipseAroundDBA(ax1, series_mean, series_dtw_horiz_var, series_dtw_vertic_var, span_tuple):
    color_type_num = 10
    color_arr = cm.rainbow(np.linspace(0, 1, color_type_num))
    ells = [Ellipse(xy=(i, series_mean[i]),
                    width=series_dtw_horiz_var[i], height=series_dtw_vertic_var[i], color=color_arr[i%color_type_num])
            for i in np.arange(span_tuple[0], span_tuple[1])
            ]
    for e in ells:
        ax1.add_artist(e)

def plotDeviations(plt, values, color, range = None) :
    if (range == None):
        plt.plot(x, values, color=color)
    elif (range[0] + 1 == range[1]):
        plt.scatter(x[range[0]], values[range[0]], color=color)
    else:
        plt.plot(x[range[0]: range[1]], values[range[0]: range[1]], color=color)

    # adjusted_series_weight_mat 往前映射的距离
    # series_mapping_mat         当前映射的index
    # 注意range是[ a, b )
def plotAdjustedSeries(plt, cluster_id ,series_mapping_mat, adjusted_series_weight_mat, range, series):
    plt.set_title(str(len(series)) + " instances of cluster " + str(cluster_id))
    colors = getColorIter(series_mapping_mat.shape[0])
    for series_index in np.arange(0, series_mapping_mat.shape[0]):
        s1, e1 = dba.getStartAndEndMapping(series_mapping_mat, adjusted_series_weight_mat, series_index, range[1] - 1)
        s2, e2 = dba.getStartAndEndMapping(series_mapping_mat, adjusted_series_weight_mat, series_index, range[0])

        x = np.arange(s2, e1)
        y = list(map(lambda  x: series[series_index][x], x))
        if len(x) == 1:
            plt.scatter(x, y, color=next(colors))
        else:
            plt.plot(x, y, color=next(colors))
        expected_range = np.arange(range[0], range[1])
        non_plotted_range = [item for item in expected_range if item not in x]
        for non_plotted_dot in non_plotted_range:
            plt.scatter(non_plotted_dot, series[series_index][non_plotted_dot], color = 'black')


def plotSelectedSpan(span_tuple, cluster_idx, ax1, ax2, ax3, ax4, ax5, ax6):
    cur_series = findAllSeriesWithIndices(findAllIndicesOfOneCluster(cluster_idx))
    dt4 = dt.datetime.now()

    series_mean, series_dtw_horiz_var, series_dtw_vertic_var, series_vertic_var, adjusted_series_mat, series_mapping_mat, adjusted_series_weight_mat, series_dtw_special_vertic_var = dba.performDBA(cur_series, iteration)
    dt5 = dt.datetime.now()
    print("getting DBA use time " + str((dt5 - dt4).seconds))

    plotDBA(ax1, series_mean, span_tuple)
    #plotRawAverage(ax1, cur_series)
    plotAdjustedSeries(ax2,  cluster_idx, series_mapping_mat, adjusted_series_weight_mat, span_tuple, cur_series)
    plotStatisticsCurves(ax4, ax5, ax6, series_dtw_horiz_var, series_dtw_vertic_var, series_vertic_var, series_dtw_special_vertic_var, span_tuple)
    plotDBA(ax3, series_mean, span_tuple)
    plotEclipseAroundDBA(ax3, series_mean, series_dtw_horiz_var, series_dtw_vertic_var, span_tuple)


def plotStatisticsCurves(ax4, ax5, ax6, series_dtw_horiz_var, series_dtw_vertic_var, series_vertic_var, series_dtw_special_vertic_var, span_tuple = None):
    plotDeviations(ax4, series_dtw_vertic_var, 'red', span_tuple)
    plotDeviations(ax5, series_vertic_var, 'grey', span_tuple)
    plotDeviations(ax6, series_dtw_horiz_var, 'green', span_tuple)
    # plotStatisticsCurve(ax6, series_dtw_special_vertic_var, 'purple', span_tuple)


def plott(i):
    fig = plt.figure(i)
    ax1, ax2, ax3, ax4, ax5, ax6 = initFigure(fig)
    plotSelectedSpan((start, end), i-1, ax1, ax2, ax3, ax4, ax5, ax6)


##for plotting the links
def listToDic(ml):
    dict = {}
    for cml in ml:
        if cml[0] in dict:
            dict[cml[0]].append(cml[1])
        else:
            dict[cml[0]] = [cml[1]]
        if cml[1] in dict:
            dict[cml[1]].append(cml[0])
        else:
            dict[cml[1]] = [cml[0]]
    return dict


ml_dict = listToDic(ml)
cl_dict = listToDic(cl)


def getClusterAndRepresentativeIds(clustering, firstId, secondId):
    found_first_sp = False
    found_second_sp = False
    representative_idx_first = 0
    representative_idx_second = 0
    cluster_idx_first = 0
    cluster_idx_second = 0

    for cluster_idx in range(0, len(clustering.clusters)):
        cluster = clustering.clusters[cluster_idx]
        for sp in cluster.super_instances:
            if found_first_sp == True and found_second_sp == True:
                break
            for idx in sp.indices:
                if found_first_sp == False and firstId == idx:
                    cluster_idx_first = cluster_idx
                    representative_idx_first = sp.representative_idx
                    found_first_sp = True
                if found_second_sp == False and secondId == idx:
                    cluster_idx_second = cluster_idx
                    representative_idx_second = sp.representative_idx
                    found_second_sp = True
    # assert (found_first_sp == True and found_second_sp == True)
    return cluster_idx_first, cluster_idx_second, representative_idx_first, representative_idx_second


# def getMustLinks(clustering, first_idx, second_idx):
#
#     # cluster_idx_left, cluster_idx_right = getClusteringIdxs(clustering, first_repre_idx, second_repre_idx)
#     cluster_idx_left, cluster_idx_right, repre_idx_left, repre_idx_right = getClusterAndRepresentativeIds(clustering, first_idx, second_idx)
#     assert (cluster_idx_left == cluster_idx_right)
#     must_links_between_first_and_second = findMustLinksBetweenSuperintances(ml_dict, repre_idx_left, repre_idx_right)
#     links = must_links_between_first_and_second
#     if len(links) >= 2:
#         if links[0] != first_idx:
#             links = [first_idx] + links
#         if links[len(links) - 1] != second_idx:
#             links = links + [second_idx]
#
#     return cluster_idx_left, links

def getMustLinks(first_idx, second_idx, repre_idx_left, repre_idx_right):
    must_links_between_first_and_second = findMustLinksBetweenSuperintances(ml_dict, repre_idx_left, repre_idx_right)
    links = must_links_between_first_and_second
    if len(links) >= 2:
        if links[0] != first_idx:
            links = [first_idx] + links
        if links[len(links) - 1] != second_idx:
            links = links + [second_idx]

    return links

def getCannotLinks(cluster_idx_left, cluster_idx_right, first_repre_idx, second_repre_idx, mock_1, mock_3):
    # cluster_idx_left, cluster_idx_right = getClusteringIdxs(clustering, first_repre_idx, second_repre_idx)
    # assert (cluster_idx_left != cluster_idx_right)
    cannot_link_between_cluster = getCannotLinkBetweenCluster(cluster_idx_left, cluster_idx_right)
    must_links_between_first_and_cluster = findMustLinksBetweenSuperintances(ml_dict, first_repre_idx, cannot_link_between_cluster[0])
    if len(must_links_between_first_and_cluster) > 0 and must_links_between_first_and_cluster[0] != mock_1:
        must_links_between_first_and_cluster = [mock_1] + must_links_between_first_and_cluster

    must_links_between_second_and_cluster = findMustLinksBetweenSuperintances(ml_dict, cannot_link_between_cluster[1], second_repre_idx)
    if len(must_links_between_second_and_cluster) > 0 and must_links_between_second_and_cluster[len(must_links_between_second_and_cluster) - 1] != mock_3:
        must_links_between_second_and_cluster = must_links_between_second_and_cluster + [mock_3]

    return must_links_between_first_and_cluster, must_links_between_second_and_cluster

def getCannotLinkBetweenCluster(cluster_idx_left, cluster_idx_right):
    right_cluster_sps_indices = []
    for sp in clustering.clusters[cluster_idx_right].super_instances:
        right_cluster_sps_indices.append(sp.representative_idx)

    for sp in clustering.clusters[cluster_idx_left].super_instances:
        sp_repre_idx = sp.representative_idx
        if sp_repre_idx in cl_dict:
            for second in cl_dict[sp_repre_idx]:
                if second in right_cluster_sps_indices:
                    return [sp_repre_idx, second]
    return []
def getClusteringIdxs(clustering, first_repre_idx, second_repre_idx):
    found_first_sp = False
    found_second_sp = False
    cluster_idx_first = 0
    cluster_idx_second = 0
    for cluster_idx in range(0, len(clustering.clusters)):
        cluster = clustering.clusters[cluster_idx]
        for sp in cluster.super_instances:
            if found_first_sp == False and sp.representative_idx == first_repre_idx:
                cluster_idx_first = cluster_idx
                found_first_sp = True
            if found_second_sp == False and sp.representative_idx == second_repre_idx:
                cluster_idx_second = cluster_idx
                found_second_sp = True
            if found_first_sp == True and found_second_sp == True:
                break


    # assert (found_first_sp == True and found_second_sp == True)
    return cluster_idx_first, cluster_idx_second

def checkLoop(path):
    dict = {}
    for e in path:
        if (e in dict):
            return False
        else:
            dict[e] = 1
    return True

def findMustLinksBetweenSuperintances(dict, first_repre_idx, second_repre_idx):
    visited_vertex = []
    visited_path = Queue()
    visited_path.put([first_repre_idx])
    visited_vertex.append(first_repre_idx)

    while (not visited_path.empty()):
        current_path = visited_path.get()
        last_of_path = current_path[len(current_path)-1]

        if last_of_path in dict:
            for next_elem in dict[last_of_path]:
                if (not next_elem in visited_vertex):
                    to_be_added_path = current_path + [next_elem]
                    if (next_elem == second_repre_idx):
                        return to_be_added_path
                    else:
                        visited_vertex.append(next_elem)
                        if checkLoop(to_be_added_path):
                            visited_path.put(to_be_added_path)

    return []

print("the number of clusters is "+ str(len(clustering.clusters)))

def getClusterSpsDict():
    cluster_sps_dict = {}
    for i in range(0, len(clustering.clusters)):
        size_of_sps_in_current_cluster = len(clustering.clusters[i].super_instances)
        cluster_sps_dict[i] = size_of_sps_in_current_cluster
        print(str(i) + " : " + str(size_of_sps_in_current_cluster))
    return cluster_sps_dict

cluster_sps_dict = getClusterSpsDict()

def getRandomMLPair():
    ml_pairs = []
    for cluster_idx in cluster_sps_dict:
        first = random.randint(0, cluster_sps_dict[cluster_idx] - 1)
        second = random.randint(0, cluster_sps_dict[cluster_idx] - 1)
        while (first == second):
            first = random.randint(0, cluster_sps_dict[cluster_idx] - 1)
            second = random.randint(0, cluster_sps_dict[cluster_idx] - 1)
        ml_pairs.append([[cluster_idx, first], [cluster_idx, second]])
    return ml_pairs

def getRandomCLPair():
    cl_pairs = []
    cnt = 5
    for i in range(0, len(cluster_sps_dict)):
        for j in range(i+1, len(cluster_sps_dict)):
            if cnt < 0:
                break
            else:
                cl_pairs.append([[i, 1], [j, 1]])

    return cl_pairs


def getYlabel(instanceid, clusterid):
    return "instance id: " + str(instanceid) + "                    \ncluster    id: " + str(clusterid) + "                     "


def plotMustLink(cluster_id, mock_1, mock_2,  repre_idx_left, repre_idx_right):
    must_link = getMustLinks(mock_1, mock_2, repre_idx_left, repre_idx_right)
    if (len(must_link) == 0):
        print("must 1: " + str(mock_1) + " must 2: " + str(mock_2))
    else:
        fig, axs = plt.subplots(len(must_link), sharex=True, sharey=True)

        fig.suptitle('must link')
        for plot_idx in range(0, len(must_link)):
            axs[plot_idx].plot(x, series[must_link[plot_idx]])
            axs[plot_idx].set_ylabel(getYlabel(must_link[plot_idx], cluster_id), color="blue", rotation=0)

        plt.show()

def plotCannotLink(cluster_id_1, cluster_id_2, repre_idx_left, repre_idx_right, mock_1, mock_3):
    cannot_link_1, cannot_link_2 = getCannotLinks(cluster_id_1, cluster_id_2, repre_idx_left, repre_idx_right, mock_1, mock_3)
    if (len(cannot_link_1) == 0 or len(cannot_link_2) == 0) :
        print("can not 1: " + str(mock_1) + "can not 2: " + str(mock_3))
    else:
        fig2, axs2 = plt.subplots(len(cannot_link_1) + len(cannot_link_2), sharex=True, sharey=True)
        fig2.suptitle('can not link')

        for plot_idx in range(0, len(cannot_link_1)):
            axs2[plot_idx].plot(x, series[cannot_link_1[plot_idx]], color = 'pink')
            axs2[plot_idx].set_ylabel(getYlabel(cannot_link_1[plot_idx], cluster_id_1), color = "pink", rotation=0)
            # axs2[plot_idx].tick_params(axis='y', labelrotation=90)
        for plot_idx in range(0, len(cannot_link_2)):
            axs2[len(cannot_link_1) +  plot_idx].plot(series[cannot_link_2[plot_idx]], color = 'purple')
            axs2[len(cannot_link_1) +  plot_idx].set_ylabel(getYlabel(cannot_link_2[plot_idx], cluster_id_2), color="purple", rotation=0)
        plt.show()

def getMock(cluster_idx, super_instance_idx):
    return clustering.clusters[cluster_idx].super_instances[super_instance_idx].representative_idx

def findAllSeriesWithIndices(indices):
    target_series = list(map(lambda idx: series[idx], indices))
    return target_series

def getColorIter(size):
    color_arr = cm.rainbow(np.linspace(0, 1, size))
    colors = iter(color_arr)
    return colors


def plotDBA(plt, series_mean, range = None):
    if (range == None):
        plt.plot(x, series_mean, color='purple')
    elif (range[0] + 1 == range[1]):
        plt.scatter(x[range[0]], series_mean[range[0]], color = 'purple')
    else:
        plt.plot(x[range[0]: range[1]], series_mean[range[0]: range[1]], color = 'purple')

def plotEclipseAroundDBA(ax1, series_mean, series_dtw_horiz_var, series_dtw_vertic_var, span_tuple):
    color_type_num = 10
    color_arr = cm.rainbow(np.linspace(0, 1, color_type_num))
    ells = [Ellipse(xy=(i, series_mean[i]),
                    width=series_dtw_horiz_var[i], height=series_dtw_vertic_var[i], color=color_arr[i%color_type_num])
            for i in np.arange(span_tuple[0], span_tuple[1])
            ]
    for e in ells:
        ax1.add_artist(e)

def plotDeviations(plt, values, color, range = None) :
    if (range == None):
        plt.plot(x, values, color=color)
    elif (range[0] + 1 == range[1]):
        plt.scatter(x[range[0]], values[range[0]], color=color)
    else:
        plt.plot(x[range[0]: range[1]], values[range[0]: range[1]], color=color)

def plotStatic():
    mock_1 = clustering.clusters[0].super_instances[1].representative_idx
    mock_2 = clustering.clusters[0].super_instances[2].representative_idx
    mock_3 = clustering.clusters[1].super_instances[1].representative_idx

    plotMustLink(mock_1, mock_2)
    plotCannotLink(mock_1, mock_3)


def plotRandom():
    mock_ml_pairs = getRandomMLPair()
    mock_cl_paris = getRandomCLPair()

    for mock_ml_pair in mock_ml_pairs:
        m1 = getMock(mock_ml_pair[0][0], mock_ml_pair[0][1])
        m2 = getMock(mock_ml_pair[1][0], mock_ml_pair[1][1])
        plotMustLink(m1, m2)

    for mock_cl_pari in mock_cl_paris:
        m1 = getMock(mock_cl_pari[0][0], mock_cl_pari[0][1])
        m2 = getMock(mock_cl_pari[1][0], mock_cl_pari[1][1])
        plotCannotLink(m1, m2)

def plotMuskAndCannotLinks():
    plotStatic()
    print("before getting random")
    plotRandom()

def plotAsRequired():
    while (True):
        idx1 = int(input("index1"))
        idx2 = int(input("index2"))
        cluster_idx_left, cluster_idx_right, repre_idx_left, repre_idx_right = getClusterAndRepresentativeIds(
            clustering, idx1, idx2)
        if cluster_idx_left == cluster_idx_right:
            plotMustLink(cluster_idx_left, idx1, idx2, repre_idx_left, repre_idx_right)
        else:
            plotCannotLink(cluster_idx_left, cluster_idx_right, repre_idx_left, repre_idx_right, idx1, idx2)

def main():
    for i in range(1, len(clustering.clusters) + 1):
        plott(i)
    plt.show()
    plotAsRequired()
    # plotMuskAndCannotLinks()


if __name__ == '__main__':
    main()

#experimental
def plotDeviationAroundDBA(ax1, series_mean, series_vertic_var, span_tuple = None):
    lower = series_mean - series_vertic_var
    upper = series_mean + series_vertic_var
    real_x = x
    if span_tuple != None:
        real_x = x[span_tuple[0]: span_tuple[1]]
        lower = lower[span_tuple[0]: span_tuple[1]]
        upper = upper[span_tuple[0]: span_tuple[1]]
    ax1.fill_between(real_x, lower, upper, alpha=0.1)

#experimental
def plotRawAverage(plt, cur_series):
    avg = np.divide(np.sum(cur_series, 0), len(cur_series))
    plt.plot(x, avg, color='purple')

#experimental
def plotMultiCurves(plt, indices, title, showLabel=False):
    plt.set_title(title + ' count=' + str(len(indices)))
    colors = getColorIter(len(indices))
    indices.sort()
    for index in indices:
        cur_color = next(colors)
        y = series[index]
        if showLabel == True:
            plt.plot(x, y, color=cur_color, label="index=" + str(index))
        else:
            plt.plot(x, y, color=cur_color)
    if showLabel == True & len(indices) < 15:
        plt.legend(loc = 9)

#deprecated
def plottt(i):
    # fig1 = plt.figure(1)
    exec('fig{} = plt.figure({})'.format(i, i))
    exec('ax{}, ax{},ax{}, ax{},ax{}, ax{} = initFigure(fig{})'.format(str(i)+'1', str(i)+'2', str(i)+'3', str(i)+'4', str(i)+'5', str(i)+'6', str(i)))
    exec('plotSelectedSpan((start, end), {}, ax{}, ax{},ax{}, ax{},ax{}, ax{})'.format(i-1, str(i)+'1', str(i)+'2', str(i)+'3', str(i)+'4', str(i)+'5', str(i)+'6'))

