#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Aug 31 22:48:21 2021

@author: apple
"""

import numpy as np 
import pandas as pd
from HRP import seriation

import fastcluster
from scipy.cluster.hierarchy import fcluster

from gap_statistic import OptimalK

from backtest import df_to_matrix

#HERC
def intersection(list1, list2): 
    intersec = [set(list1) & set(list2)]
    return intersec


def compute_allocation(covar, clusters,Z,dimensions):
    numClusters = len(clusters)
    aWeights = np.array([1.] * len(covar))
    
    cWeights = np.array([1.] * numClusters)
    cVar = np.array([0.] * numClusters)
    
    for i, cluster in clusters.items():
        cluster_covar = covar[cluster, :][:, cluster]
        inv_diag = 1 / np.diag(cluster_covar)
        aWeights[cluster] = inv_diag / np.sum(inv_diag)
        
    for i, cluster in clusters.items():
        weights = aWeights[cluster]
        cVar[i - 1] = np.dot(
            weights, np.dot(covar[cluster, :][:, cluster], weights))
        
    for m in range(numClusters - 1):
        left = int(Z[dimensions - 2 - m, 0])
        lc = seriation(Z, dimensions, left)
        
        right = int(Z[dimensions - 2 - m, 1])
        rc = seriation(Z, dimensions, right)


        id_lc = []
        id_rc = []
        
        for i, cluster in clusters.items():
            if sorted(intersection(lc, cluster)) == sorted(cluster):
                id_lc.append(i)
            if sorted(intersection(rc, cluster)) == sorted(cluster):
                id_rc.append(i)


        id_lc = np.array(id_lc) - 1
        id_rc = np.array(id_rc) - 1

        alpha = 0
        lcVar = np.sum(cVar[id_lc])
        rcVar = np.sum(cVar[id_rc])
        alpha = lcVar / (lcVar + rcVar)

        cWeights[id_lc] = cWeights[
            id_lc] * alpha
        cWeights[id_rc] = cWeights[
            id_rc] * (1 - alpha)

    for i, cluster in clusters.items():
        aWeights[cluster] = aWeights[cluster] * cWeights[
            i - 1]
        
    return aWeights

#Dataframe of returns
def HERC(mat_ret):
    #Need to first calculate the optimal number of clusters
    #The mat_ret that goes into this must be a np array of returns
    # correl_mat = mat_ret.corr(method='pearson')
    column_dic = {k:v for v, k in enumerate(mat_ret.columns)}
    
    correl_mat = df_to_matrix(mat_ret.corr(method='pearson'))
    dist = 1 - correl_mat
    dim = len(dist)
    
    tri_a, tri_b = np.triu_indices(dim, k = 1)
    
    
    
    Z = fastcluster.linkage(dist[tri_a, tri_b], method='ward')
    
    optimalK = OptimalK(parallel_backend = 'rust')
    n_clusters = optimalK(mat_ret.values, cluster_array = np.arange(1,len(mat_ret)))
   
    nb_clusters = n_clusters
    clustering_inds = fcluster(Z, nb_clusters, criterion='maxclust')
    clusters = {i: [] for i in range(min(clustering_inds),max(clustering_inds) + 1)}
    for i, v in enumerate(clustering_inds):
        clusters[v].append(i)
    
    HERC_w = compute_allocation(correl_mat, clusters, Z, dim)
    HERC_w = pd.Series(HERC_w)
    my_inverted_dict = dict(map(reversed, column_dic.items()))
    
    HERC_w = HERC_w.rename(index = my_inverted_dict)
    
    return HERC_w
    
[((30, 15, 30, 43), 'numpy.array', 'np.array', ({(30, 24, 30, 42): '[1.0] * numClusters'}, {}), '([1.0] * numClusters)', True, 'import numpy as np\n'), ((31, 11, 31, 39), 'numpy.array', 'np.array', ({(31, 20, 31, 38): '[0.0] * numClusters'}, {}), '([0.0] * numClusters)', True, 'import numpy as np\n'), ((91, 19, 91, 46), 'numpy.triu_indices', 'np.triu_indices', (), '', True, 'import numpy as np\n'), ((95, 8, 95, 62), 'fastcluster.linkage', 'fastcluster.linkage', (), '', False, 'import fastcluster\n'), ((97, 15, 97, 50), 'gap_statistic.OptimalK', 'OptimalK', (), '', False, 'from gap_statistic import OptimalK\n'), ((101, 22, 101, 68), 'scipy.cluster.hierarchy.fcluster', 'fcluster', (), '', False, 'from scipy.cluster.hierarchy import fcluster\n'), ((107, 13, 107, 30), 'pandas.Series', 'pd.Series', ({(107, 23, 107, 29): 'HERC_w'}, {}), '(HERC_w)', True, 'import pandas as pd\n'), ((45, 13, 45, 43), 'HRP.seriation', 'seriation', ({(45, 23, 45, 24): 'Z', (45, 26, 45, 36): 'dimensions', (45, 38, 45, 42): 'left'}, {}), '(Z, dimensions, left)', False, 'from HRP import seriation\n'), ((48, 13, 48, 44), 'HRP.seriation', 'seriation', ({(48, 23, 48, 24): 'Z', (48, 26, 48, 36): 'dimensions', (48, 38, 48, 43): 'right'}, {}), '(Z, dimensions, right)', False, 'from HRP import seriation\n'), ((65, 16, 65, 35), 'numpy.sum', 'np.sum', ({(65, 23, 65, 34): 'cVar[id_lc]'}, {}), '(cVar[id_lc])', True, 'import numpy as np\n'), ((66, 16, 66, 35), 'numpy.sum', 'np.sum', ({(66, 23, 66, 34): 'cVar[id_rc]'}, {}), '(cVar[id_rc])', True, 'import numpy as np\n'), ((35, 23, 35, 45), 'numpy.diag', 'np.diag', ({(35, 31, 35, 44): 'cluster_covar'}, {}), '(cluster_covar)', True, 'import numpy as np\n'), ((36, 39, 36, 55), 'numpy.sum', 'np.sum', ({(36, 46, 36, 54): 'inv_diag'}, {}), '(inv_diag)', True, 'import numpy as np\n'), ((41, 21, 41, 67), 'numpy.dot', 'np.dot', ({(41, 28, 41, 57): 'covar[(cluster), :][:, (cluster)]', (41, 59, 41, 66): 'weights'}, {}), '(covar[(cluster), :][:, (cluster)], weights)', True, 'import numpy as np\n'), ((61, 16, 61, 31), 'numpy.array', 'np.array', ({(61, 25, 61, 30): 'id_lc'}, {}), '(id_lc)', True, 'import numpy as np\n'), ((62, 16, 62, 31), 'numpy.array', 'np.array', ({(62, 25, 62, 30): 'id_rc'}, {}), '(id_rc)', True, 'import numpy as np\n')]