from pathlib import Path

from A_Global import path_stock
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
import numpy as np
plt.rcParams['xtick.direction'] = 'in'  # in; out; inout
plt.rcParams['ytick.direction'] = 'in'
col_dict = {
    'r1d': 'CLOSE',
    'r5d': 'avg_5',
    'r20d': 'avg_20',
    'r60d': 'avg_60'
}

industry_dict = {
    'FB.O'  : 'Interactive Media',
    'GOOG.O': 'Interactive Media',
    'AEP.O' : 'Electric Utilities',
    'LNT.O' : 'Electric Utilities',
    'APA.O' : 'Oil & Gas Exploration',
    'FANG.O': 'Oil & Gas Exploration',
    'NVDA.O': 'Semiconductors',
    'AAPL.O': 'Technology Hardware',
    'NXPI.O': 'Semiconductors',
    'XEL.O' : 'Multi-Utilities',
    'EBAY.O': 'Internet',
    'EXC.O' : 'Multi-Utilities'

}
font1 = {
    # 'family': 'Nimbus Roman',
    'weight': 'bold',
    'style': 'normal',
    'size': 12,
}

pos = [231, 232, 233, 234, 235, 236]
def pic_close_2stocks(market, col, s1, s2, title, start_date, end_date, i):
    ax = plt.subplot(pos[i])

    path_dir = path_stock / ('stocks_' + market)
    df1 = pd.read_csv(path_dir / s1, index_col = 0)
    df2 = pd.read_csv(path_dir / s2, index_col = 0)
    close1 = df1.loc[start_date: end_date, [col]]
    close2 = df2.loc[start_date: end_date, [col]]
    # idx = np.arange(len(close1.index))
    year = '2018'
    idx = [d[5:7] + '/' + d[8:] for d in close1.index]

    close1 = MinMaxScaler().fit_transform(close1)
    close2 = MinMaxScaler().fit_transform(close2)
    close1 = close1.reshape(251)
    close2 = close2.reshape(251)

    line1, = plt.plot(idx, close1, label=s1 + ':' + industry_dict[s1], linewidth=2, linestyle='--')
    line2, = plt.plot(idx, close2, label=s2 + ':' + industry_dict[s2], color='brown', linewidth=2, linestyle='-')
    lines = [line1, line2]

    # plt.legend(handles=lines, labels=[s1, s2], loc='best')
    # ax.legend(loc=2, bbox_to_anchor=(1.05,1.0),borderaxespad = 0.)
    plt.legend(bbox_to_anchor=(0, 1.02, 1, 0.2), loc="lower left",
               mode="expand", borderaxespad=0, ncol=1, prop=font1)
    if i == 4:
        plt.xlabel('Trading Days in [' + idx[0] + '/' + year + ',' + idx[-1] + '/' + year + ']', font1)

    if i in [0, 3]: plt.ylabel('Price', font1)
    xticks = list(range(0, len(idx), 70))  # len(idx)//10
    xticks.append(len(idx) - 1)
    ax.set_xticks(xticks)
    labels = ax.get_xticklabels() + ax.get_yticklabels()
    [label.set_fontweight('bold') for label in labels]
    plt.tick_params(labelsize=12, width='2')
    thickness = 2
    ax.spines['left'].set_linewidth(thickness)
    ax.spines['bottom'].set_linewidth(thickness)
    ax.spines['right'].set_linewidth(thickness)
    ax.spines['top'].set_linewidth(thickness)
    if i in [0,1,2]:
        ax.axes.xaxis.set_ticklabels([])

    if i in [1,2,4,5]:
        ax.axes.yaxis.set_ticklabels([])
    if i in [0,1,2]:
        plt.title(title, x=1/2, y=-0.1, fontsize=14, fontfamily='Times New Roman')
    else:
        plt.title(title, x=1/2, y=-0.23, fontsize=14, fontfamily='Times New Roman')
    # plt.show()

    # ax2.set_xlabel('Trading Days in [' + idx[0] + '/' + year + ',' + idx[-1] + '/' + year + ']', font1)
    # ax2.set_ylabel("Price", font1)
'''
画实验一最后的不同关系股票股价变化图

'''
if __name__ == '__main__':
    stock_pair_arr = [
        ('r1d', ('FB.O', 'GOOG.O'), 'SD'),
        ('r5d', ('AEP.O', 'LNT.O'), 'SW'),
        ('r20d', ('APA.O','FANG.O'),'SM'),
        ('r60d', ('AAPL.O', 'NVDA.O'), 'SS'),
        ('r20d', ('NXPI.O', 'XEL.O'),'RM'),
        ('r60d', ('EBAY.O', 'EXC.O'),'RS'),
    ]

    plt.figure(figsize=(16, 9))
    plt.subplots_adjust(hspace=0.3, wspace=0.1)
    for i in range(len(stock_pair_arr)):
        e = stock_pair_arr[i]
        pic_close_2stocks('NASDAQ', col_dict[e[0]], e[1][0], e[1][1], e[2], '2018-01-02', '2018-12-31', i)
    plt.savefig(Path('./') / 'SG-Relations', dpi=300, bbox_inches='tight')
    plt.show()
