# -*- coding: utf-8 -*-
# coding: utf-8
import numpy as np
import scipy as sc
from sklearn import mixture
import collections
from GrASP import Params_GrASP


def Spectral_clustering (event_id, Stadis_array, Station_loc_data_XY, Probability, Threshold, Max_eigen, Std_thres, Normalized_mode = 'on'):
    """ 
    Extract tremor station groups

    Parameters:
    ----------
    event_id: ID for Target datetime
    Stadis_array: Permanent adjacency matrix
    Station_loc_data_XY: Station locations in XY coordinates (unit: meter)
    Probability: Tremor probability
    Threshold: Threshold for tremor probability
    Max_eigen: CCR threshold for the second smallest eigenvalue
    Std_thres: Threshold for Graph std
    Normalized_mode: Use normalized laplacian. default: 'on'

    Returns:
    ----------
    pred_output: Cluster ID 
    pred_keys_output: Key for Cluster ID
    Adjacency_ID_output: Station ID 
    output_type: In case 1, one cluster was extracted; in case 2, two or more clusters were extracted.
    """


    ##### Permanent Adjacency #####
    Adjacency_matrix_All = Stadis_array.copy()
    ##### Temporal weight #####
    Probability_weight = Probability[event_id,:]
    Probability_weight = Probability_weight[:,np.newaxis]
    Weight = np.dot(Probability_weight, Probability_weight.T)
    ##### Temporal Adjacency #####
    Adjacency_matrix_All = Adjacency_matrix_All * Weight 
    Adjacency_ID = np.array([], dtype = 'int')
    for id in range(np.shape(Adjacency_matrix_All)[0]):
        if np.sum(Adjacency_matrix_All[id,:]) > 0:
            Adjacency_ID = np.append(Adjacency_ID, id)

    if len(Adjacency_ID) > 2:
        ##### Adjacency matrix #####
        Adjacency_matrix = []
        for ad_id in Adjacency_ID:
            Adjacency_matrix.append(Adjacency_matrix_All[ad_id, Adjacency_ID])
        Adjacency_matrix = np.array(Adjacency_matrix)
        ##### Degree matrix #####
        Degree_matrix = np.identity(np.shape(Adjacency_matrix)[0])
        Degree_matrix_inv = np.identity(np.shape(Adjacency_matrix)[0])
        for diagonal in range(np.shape(Adjacency_matrix)[0]):
            Degree_matrix[diagonal, diagonal] = np.sum(Adjacency_matrix[:, diagonal])
            Degree_matrix_inv[diagonal, diagonal] = (np.sum(Adjacency_matrix[:, diagonal]))**(-0.5)
        ##### Laplacian matrix #####
        Laplacian_matrix = Degree_matrix - Adjacency_matrix
        ##### Normalized Laplacian #####
        Normalized_Laplacian = np.dot(Degree_matrix_inv, Laplacian_matrix)
        Normalized_Laplacian = np.dot(Normalized_Laplacian, Degree_matrix_inv)
        if Normalized_mode == 'on':
            Laplacian_matrix = Normalized_Laplacian

        ##### Eigenvalue decomposition #####
        rank = np.linalg.matrix_rank(Laplacian_matrix)
        eigval, eigvec = sc.linalg.eigh(Laplacian_matrix)
        eigval = np.round(eigval, 7)
        eigvec = np.round(eigvec, 7)
        zero_number = len(eigval[eigval == 0])

        ##### One cluster #####
        if zero_number == 1:
            pred_output = []
            pred_keys_output = []
            Adjacency_ID_output = []
            output_type = 1
            Station_locx = Station_loc_data_XY[Adjacency_ID, 1]
            Station_locy = Station_loc_data_XY[Adjacency_ID, 2]
            Graph_std = 0.001*np.array([np.std(Station_locx, ddof = 1), np.std(Station_locy, ddof = 1)])
            Graph_std_max = np.max(Graph_std)
            try:
                np.std(Station_locy, ddof=1)
            except RuntimeWarning:
                print(Station_locy)
                quit()
            ##### Cumulative contributiion rate #####
            Accumulation = np.array([])
            for con in range(len(eigval)):
                eigval_sample = eigval[0 : con + 1]
                Accumulation = np.append(Accumulation, np.sum(eigval_sample)/np.sum(eigval))
            Accumulation = Accumulation * 100
            ##### Eigenvectors #####
            if len(Accumulation[Accumulation < Max_eigen]) > 1:
                eigvec_mixture = eigvec[:, eigval <= eigval[1]]
            else:
                eigvec_mixture = eigvec[:, eigval <= eigval[0]]
            
            ##### Graph std threshold #####
            if Graph_std_max <= Std_thres:
                if np.shape(eigvec_mixture)[1] == 1:
                    ##### Gaussian mixture model #####
                    GMM = mixture.GaussianMixture(n_components = np.shape(eigvec_mixture)[1], covariance_type = 'full', init_params = 'k-means++')
                    GMM.fit(eigvec_mixture)
                    pred = GMM.predict(eigvec_mixture)
                    pred_count = collections.Counter(pred)

                    pred_output = pred
                    pred_keys_output = pred_count.keys()
                    Adjacency_ID_output = Adjacency_ID
                    output_type = 1

                elif np.shape(eigvec_mixture)[1] > 1:
                    ##### Gaussian mixture model #####
                    GMM = mixture.GaussianMixture(n_components = np.shape(eigvec_mixture)[1], covariance_type = 'full', init_params = 'k-means++')
                    GMM.fit(eigvec_mixture)
                    pred = GMM.predict(eigvec_mixture)
                    pred_count = collections.Counter(pred)

                    pred_output = pred
                    pred_keys_output = pred_count.keys()
                    Adjacency_ID_output = Adjacency_ID
                    output_type = 1

        ##### Two or more clusters #####
        elif zero_number > 1:
            eigvec_mixture = eigvec[:, eigval <= eigval[zero_number - 1]]
            ##### Gaussian mixture model #####
            GMM = mixture.GaussianMixture(n_components = zero_number, covariance_type = 'full', init_params='k-means++')
            GMM.fit(eigvec_mixture)
            pred = GMM.predict(eigvec_mixture)
            pred_count = collections.Counter(pred)
            pred_output = []
            pred_keys_output = []
            Adjacency_ID_output = []
            output_type = 2

            ##### Each cluster #####
            for multi_id in pred_count.keys():
                Adjacency_ID_multi = Adjacency_ID[pred == multi_id]
                Laplacian_matrix_multi = []
                for sub_id in range(len(Adjacency_ID)):
                    if pred[sub_id] == multi_id:
                        Laplacian_matrix_multi.append(Laplacian_matrix[sub_id, pred == multi_id])
                Laplacian_matrix_multi = np.array(Laplacian_matrix_multi)

                ##### Eigenvalue decomposition #####
                rank = np.linalg.matrix_rank(Laplacian_matrix_multi)
                eigval_multi, eigvec_multi = sc.linalg.eigh(Laplacian_matrix_multi)
                eigval_multi = np.round(eigval_multi, 7) ##### 固有値: 小数点 7 桁以下で丸め込み
                eigvec_multi = np.round(eigvec_multi, 7) ##### 固有ベクトル: 小数点 7 桁以下で丸め込み
                
                Station_locx = Station_loc_data_XY[Adjacency_ID_multi, 1]
                Station_locy = Station_loc_data_XY[Adjacency_ID_multi, 2]
                Graph_std = 0.001*np.array([np.std(Station_locx, ddof=1), np.std(Station_locy, ddof=1)])
                Graph_std_max = np.max(Graph_std)

                ##### Cumulative contributiion rate #####
                Accumulation_multi = np.array([])
                for con_multi in range(len(eigval_multi)):
                    eigval_sample_multi = eigval_multi[0: con_multi+1]
                    Accumulation_multi = np.append(Accumulation_multi, np.sum(eigval_sample_multi)/np.sum(eigval_multi))
                Accumulation_multi = Accumulation_multi * 100
                ##### Eigenvectors #####
                if len(Accumulation_multi[Accumulation_multi < Max_eigen]) > 1:
                    eigvec_mixture_multi = eigvec_multi[:, eigval_multi <= eigval_multi[1]]
                else:
                    eigvec_mixture_multi = eigvec_multi[:, eigval_multi <= eigval_multi[0]]

                ##### Graph std threshold #####
                if Graph_std_max <= Std_thres:
                    if np.shape(eigvec_mixture_multi)[1] == 1:
                        ##### Gaussian mixture model #####
                        GMM = mixture.GaussianMixture(n_components = np.shape(eigvec_mixture_multi)[1], covariance_type = 'full', init_params = 'k-means++')
                        GMM.fit(eigvec_mixture_multi)
                        pred_multi = GMM.predict(eigvec_mixture_multi)
                        pred_count_multi = collections.Counter(pred_multi)

                        pred_output.append(pred_multi)
                        pred_keys_output.append(pred_count_multi.keys())
                        Adjacency_ID_output.append(Adjacency_ID_multi)
                        output_type = 2

                    elif np.shape(eigvec_mixture)[1] > 1:
                        
                        GMM = mixture.GaussianMixture(n_components = np.shape(eigvec_mixture_multi)[1], covariance_type = 'full', init_params = 'k-means++')
                        GMM.fit(eigvec_mixture_multi)
                        pred_multi = GMM.predict(eigvec_mixture_multi)
                        pred_count_multi = collections.Counter(pred_multi)

                        pred_output.append(pred_multi)
                        pred_keys_output.append(pred_count_multi.keys())
                        Adjacency_ID_output.append(Adjacency_ID_multi)
                        output_type = 2
    else: 
        pred_output = []
        pred_keys_output = []
        Adjacency_ID_output =[]
        output_type = 0

    return pred_output, pred_keys_output, Adjacency_ID_output, output_type


def Apply_Tremor (Data, Stadis_array, Station_loc_data_XY):
    """ 
    Apply GrASP for Tremor

    Parameters:
    ----------
    Data: Tremor detection data
    Stadis_array: Permanent adjacency matrix
    Station_loc_data_XY: Station locations in XY coordinates (unit: meter)

    Return:
    ----------
    Save_Candidate: Station association results
    """

    Save_Candidate = []
    YMDHm = Data[:, :5]
    Probability = Data[:, 5:np.shape(Data)[1]]
    Probability[Probability < Params_GrASP.Threshold] = 0

    for event in range(0, np.shape(Probability)[0]):

        ##### Spectral clustering #####
        pred_output, pred_keys_output, Adjacency_ID_output, output_type = Spectral_clustering(event, Stadis_array, Station_loc_data_XY, Probability, Params_GrASP.Threshold, Params_GrASP.Max_eigen, Params_GrASP.Std_thres, Normalized_mode = 'on')

        if output_type == 1:
            pred = pred_output
            pred_keys = pred_keys_output
            Adjacency_ID_list = Adjacency_ID_output
            for class_id in pred_keys:
                Adjacency_ID_class = Adjacency_ID_list[pred == class_id]
                if len(Adjacency_ID_class) >= Params_GrASP.Min_sta:
                    Save_Date = YMDHm[event, :5]
                    Save_Probability = np.zeros(np.shape(Stadis_array)[0])
                    Save_Probability[Adjacency_ID_class] = Probability[event, Adjacency_ID_class]
                    Save_List = np.append(Save_Date, Save_Probability)
                    Save_Candidate.append(list(Save_List))
            
        elif output_type == 2:
            EC_dis_list = []
            Class_id_list =[]
            for out in range(len(pred_output)):
                pred = pred_output[out]
                pred_keys = pred_keys_output[out]
                Adjacency_ID_list = Adjacency_ID_output[out]
                for class_id in pred_keys:
                    Adjacency_ID_class = Adjacency_ID_list[pred == class_id]
                    if len(Adjacency_ID_class) >= Params_GrASP.Min_sta:
                        Save_Date = YMDHm[event, :5]
                        Save_Probability = np.zeros(np.shape(Stadis_array)[0])
                        Save_Probability[Adjacency_ID_class] = Probability[event, Adjacency_ID_class]
                        Save_List = np.append(Save_Date, Save_Probability)
                        Save_Candidate.append(list(Save_List))

    Save_Candidate = np.array(Save_Candidate)
    
    return Save_Candidate

