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'

plt.rcParams['font.sans-serif']=['arial']
# plt.rcParams.update({'font.size': 13})

'''画Introducion的关系图'''
if __name__ == '__main__':
    market      = 'NASDAQ'
    start_date  = '2018-01-02'
    end_date    = '2018-12-31'
    plt.figure(figsize=(6,6.5))

    ax1 = plt.subplot(211)
    s1 = 'FB.O'
    s2 = 'GOOG.O'

    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, ['CLOSE']]
    close2 = df2.loc[start_date: end_date, ['CLOSE']]
    idx = np.arange(len(close1.index))
    year = '2018'
    idx = [d[5:7] + '/' + d[8:] for d in close1.index]
    # print(idx)
    # # print(close1)
    # exit()

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

    line1, = plt.plot(idx, close1, linewidth=2, linestyle='--')
    line2, = plt.plot(idx, close2, color='brown', linewidth=2, linestyle='-')
    lines = [line1, line2]
    thickness = 2
    ax1.spines['left'].set_linewidth(thickness)
    ax1.spines['bottom'].set_linewidth(thickness)
    ax1.spines['right'].set_linewidth(thickness)
    ax1.spines['top'].set_linewidth(thickness)

    # ax = plt.gca()
    # ax.spines["top"].set_color(None)
    # ax.spines["right"].set_color(None)

    xticks = list(range(0, len(idx), 50)) # len(idx)//10
    # xticks.append(len(idx)-1)
    ax1.set_xticks(xticks)

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

    font2 = {
            # 'family': 'Nimbus Roman',
             'weight': 'normal',
             'style': 'normal',
             'size': 13,
             }

    ax1.set_xlabel('Trading Days in [' + idx[0] + '/' + year + ',' + idx[-1] + '/' + year + ']', font1)
    ax1.set_ylabel("Price", font1)

    plt.tick_params(labelsize=12, width='2')

    plt.legend(handles=lines, labels=[s1, s2], loc='best', prop=font1)
    # plt.xlabel('Trading Days', loc='right') # , fontsize=10
    # plt.ylabel('Price')
    labels = ax1.get_xticklabels() + ax1.get_yticklabels()
    # [label.set_fontname('Nimbus Roman') for label in labels]
    [label.set_fontweight('bold') for label in labels]

    plt.title('(a) Industry relation', x=1/2, y=-0.3, fontsize=14, fontfamily='Times New Roman')

    ax2 = plt.subplot(212)
    s1 = 'NXPI.O'
    s2 = 'XEL.O'

    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, ['CLOSE']]
    close2 = df2.loc[start_date: end_date, ['CLOSE']]
    # idx = np.arange(len(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, linewidth=2, linestyle='--')
    line2, = plt.plot(idx, close2, color='brown', linewidth=2, linestyle='-')
    lines = [line1, line2]
    ax2.spines['left'].set_linewidth(thickness)
    ax2.spines['bottom'].set_linewidth(thickness)
    ax2.spines['right'].set_linewidth(thickness)
    ax2.spines['top'].set_linewidth(thickness)

    ax2.set_xlabel('Trading Days in [' + idx[0] + '/' + year + ',' + idx[-1] + '/' + year + ']', font1)
    ax2.set_ylabel("Price", font1)

    xticks = list(range(0, len(idx), 50)) # len(idx)//10
    # xticks.append(len(idx)-1)
    ax2.set_xticks(xticks)

    plt.tick_params(labelsize=12, width='2')

    plt.legend(handles=lines, labels=[s1, s2], loc=[0.55, 0.7], prop=font1)

    # plt.xlabel('Trading Days', loc='right') # , fontsize=10
    # plt.ylabel('Price') # , fontsize=10
    labels = ax2.get_xticklabels() + ax2.get_yticklabels()
    # [label.set_fontname('Nimbus Roman') for label in labels]
    [label.set_fontweight('bold') for label in labels]
    plt.title('(b) Non-industry relation', x=1/2, y=-0.3, fontsize=14, fontfamily='Times New Roman')

    plt.tight_layout()
    plt.savefig(Path('./') / 'SG-Relations-Intro', dpi=300, bbox_inches='tight')
    plt.show()
