from pathlib import Path

import numpy as np
import pandas as pd
from tool_data import get_stock_filted
import matplotlib.pyplot as plt
import seaborn as sns
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

def pic_heatmap(data, market='NYSE'):
    plt.figure(figsize=[4,4])
    ax = plt.gca()
    sns.heatmap(data, vmin=0, vmax=3, cmap=['white', 'blue', 'black', 'red'], cbar=False,
                xticklabels=10,
                yticklabels=10,
                square=True,
                # linecolor='black',
                # linewidths=0.1
                )
    ax.xaxis.set_ticks_position('top')
    thickness = 2
    ax = plt.gca()
    ax.spines['left'].set_color('black')
    ax.spines['left'].set_linewidth(thickness)
    ax.spines['bottom'].set_linewidth(thickness)
    ax.spines['right'].set_linewidth(thickness)
    ax.spines['top'].set_linewidth(thickness)
    # ax.get_xaxis().set_visible(True)
    # ax.get_yaxis().set_visible(True)
    idx2stock, stock2idx = get_stock_filted(market)

    labels = []
    for i in range(50):
        labels.append(idx2stock[i] + '\n(' + str(i) + ')')
    plt.xticks(range(50), labels, rotation=30)
    plt.yticks(range(50), labels, rotation=60)
    id_list = list(range(0, 50, 10))
    id_list.append(49)
    ax.set_xticks(id_list)
    ax.set_yticks(id_list)

    plt.title(market)
    # print(idx2stock)
    # print(stock2idx)
    plt.savefig(Path('./') / ('hotplot_' + market), dpi=300, bbox_inches='tight')
    plt.show()

def pic_(df):
    plt.plot(df)
    plt.show()
    pass


if __name__ == '__main__':
    # # 调用魔法方法 使得每次显示结果时不用调用plt.show()方法
    # sns.set(style='whitegrid', color_codes=True)
    # # 设置默认风格
    # np.random.seed(sum(map(ord, 'categorical')))
    # # 当我们设置相同的seed，每次生成的随机数相同。
    # # 如果不设置seed，则每次会生成不同的随机数
    # # map() 会根据提供的函数对指定序列做映射。
    # # 即返回'distributions'该字符串每个字母的二进制数
    # # 并用sum函数求和，将和作为种子，这里似乎是习惯上的写法。
    #
    # data = np.random.rand(3, 3)
    # h = sns.heatmap(data, cmap='YlGnBu', cbar=False)
    #
    # cb = h.figure.colorbar(h.collections[0])  # 显示colorbar
    # cb.ax.tick_params(labelsize=28)  # 设置colorbar刻度字体大小。
    #
    # # cax = plt.gcf().axes[-1]
    # # cax.tick_params(labelsize=20)

    import numpy as np
    import pylab as plt
    import matplotlib as mpl

    if __name__ == '__main__':
        fig = plt.figure(figsize=(3, 8))
        cmap = mpl.cm.Set1
        ax3 = fig.add_axes([0.3, 0.2, 0.2, 0.5])  # 四个参数分别是左、下、宽、长
        norm = mpl.colors.Normalize(vmin=4, vmax=0)
        bounds = [round(elem, 2) for elem in np.linspace(4, 0, 5)]  #
        cb3 = mpl.colorbar.ColorbarBase(ax3, cmap=cmap,
                                        norm = norm,
               # to use 'extend', you must
               # specify two extra boundaries:
               boundaries = [1.2] + bounds + [2.6],
               extend = 'both',
               ticks = bounds,  # optional
                spacing = 'proportional',
                orientation = 'vertical')
        plt.show()

    plt.show()
