import numpy as np
from matplotlib import pyplot as plt
from scipy.cluster.hierarchy import linkage, dendrogram
from scipy.spatial.distance import pdist, squareform


def hierarchical_clustering(df, vmin=None, vmax=None, figsize=(8,8), top_height=2, left_width=2,
                            xmaxticks=None, ymaxticks=None, metric='euclidean', cmap=None):
    '''based on heatmap function from
    http://nbviewer.ipython.org/github/herrfz/dataanalysis/
    blob/master/week3/svd_pca.ipynb
    Generates a heatmap from the input matrix.
    '''

    df = df.copy()

    # Subplot sizes
    total_width, total_height = figsize 
    
    main_h = 1-(top_height/total_height)
    main_w = 1-(left_width/total_width)
    
    gap_x = (0.1/total_width)
    gap_y = (0.1/total_height) 
    
    left_h = main_h
    left_w = 1-main_w

    top_h = 1-main_h
    top_w = main_w
    

    ydim, xdim = df.shape

    if xmaxticks is None:
        xmaxticks = int( 5*main_w*total_width )
    if ymaxticks is None:
        ymaxticks = int( 5*main_h*total_height )
    
    dm = df.fillna(0).values
    D1 = squareform(pdist(dm, metric=metric))
    D2 = squareform(pdist(dm.T, metric=metric))
    fig = plt.figure(figsize=figsize)
    fig.set_tight_layout(False)

    
    # add left dendrogram
    ax1 = fig.add_axes([0, 0, left_w-gap_x, left_h], frameon=False)
    Y = linkage(D1, method='complete')
    Z1 = dendrogram(Y, orientation='left', color_threshold=0, above_threshold_color='k')
    ax1.set_xticks([])
    ax1.set_yticks([])
    # add top dendrogram
    ax2 = fig.add_axes([left_w, main_h+gap_y, top_w, top_h-gap_y], frameon=False)
    Y = linkage(D2, method='complete')
    Z2 = dendrogram(Y, color_threshold=0, above_threshold_color='k')
    ax2.set_xticks([])
    ax2.set_yticks([])
    # add matrix plot
    axmatrix = fig.add_axes([left_w, 0, main_w, main_h])
    idx1 = Z1['leaves']
    idx2 = Z2['leaves']
    D = dm[idx1, :]
    D = D[:, idx2]

    if cmap is None: cmap='coolwarm'
    fig = axmatrix.matshow(D[::-1], aspect='auto', cmap=cmap,
                           vmin=vmin, vmax=vmax)

    axmatrix.set_xticks([])
    axmatrix.set_yticks([])
    
    ax = plt.gca()
    ax.yaxis.tick_right()
    ax.xaxis.tick_bottom()

    clustered = df.iloc[Z1['leaves'][::-1], Z2['leaves']]

    ndx_y = np.linspace(0,len(clustered.index)-1, ymaxticks)
    ndx_x = np.linspace(0,len(clustered.columns)-1, xmaxticks)
    ndx_y = [int(i) for i in ndx_y]
    ndx_x = [int(i) for i in ndx_x]

    _ = plt.yticks(ndx_y, clustered.iloc[ndx_y].index)
    _ = plt.xticks(ndx_x, clustered.columns[ndx_x], rotation=90)

    return clustered, fig, Z1['leaves'][::-1], Z2['leaves']