'''
添加列 均价 5日 20日 60日
添加列 均价涨跌幅
'''

import os
from pathlib import Path
import pandas as pd
from pandas import Series

from A_Global import path_stock, path_stock_tmp, path_dat, raw_cols, CLOSE, OPEN



def avg_n(df, n):
    res = Series([0 for _ in range(len(df.index))], index=df.index, name=CLOSE)
    for i_ in range(n):
        res = res + df[CLOSE].shift(i_)
    return round(res / n, 2)

def call_add_column(market):
    path = path_stock / ('stocks_' + market.lower())

    files = os.listdir(path)
    i = 0
    for file in files:
        df = pd.read_csv(path / file, index_col=0)

        df['avg_5']     = avg_n(df, 5)
        df['avg_20']    = avg_n(df, 20)
        df['avg_60']    = avg_n(df, 60)

        df['r1d']       = round((df[CLOSE] - df[CLOSE].shift(1)) / df[CLOSE].shift(1), 3)
        df['r5d']       = round((df['avg_5'] - df['avg_5'].shift(1)) / df['avg_5'].shift(1), 3)
        df['r20d']      = round((df['avg_20'] - df['avg_20'].shift(1)) / df['avg_20'].shift(1), 3)
        df['r60d']      = round((df['avg_60'] - df['avg_60'].shift(1)) / df['avg_60'].shift(1), 3)

        df['label']     = round((df[OPEN].shift(-1) / df[OPEN]), 3)

        df.to_csv(path / file)

        i += 1
        print('ADD COL OVER', i, file)
    return i

def call_del_column(market):
    path = path_stock / ('stocks_' + market.lower())

    files = os.listdir(path)
    i = 0
    for file in files:
        df = pd.read_csv(path / file, index_col=0)
        df = df[raw_cols]

        df.to_csv(path / file)

        i += 1
        print('DEL COL OVER', i, file)
    return i

def call_add_label(idx):
    from A_Global import path_index
    path_idx = path_index / idx
    df = pd.read_csv(path_idx, index_col=0)
    df['label'] = round((df[CLOSE].shift(-1) / df[CLOSE]), 3)
    print(df)
    df.to_csv(path_idx)

if __name__ == '__main__':
    market = 'HKEX'

    ''' for stock '''
    n_stock = call_add_column(market)
    print('add columns OVER -', n_stock, 'stocks in', market)

    ''' roll back'''
    # call_del_column('')

    ''' for index '''
    # call_add_label('HSI.HI')