from fileRW import ReadFile
import pandas as pd


class StockInfo:

    stock_ID = ''
    all_data = dict()

    def __init__(self, stock_id):
        '''
        给定股票代码，构造该股票的实例，可获取股票的各项因子数据。
        :param stock_id: 股票代码
        '''
        self.stock_ID = stock_id
        rf = ReadFile()
        self.all_data = rf.read_pickle()

    def get_all_info_of_one_stack(self, start_date='2010-01-04', end_date='2019-12-31'):
        '''
        获取单个股票在给定时间段内的全部信息
        :param start_date: 开始日期
        :param end_date: 结束日期
        :return: DataFrame 一行是一天，一列是一个因子。
        '''

        days = list(self.all_data.keys())
        start_index = 0
        end_index = len(days) - 1

        # 考虑start_date和end_date不是交易日的情况
        for i in range(len(days)):
            if days[i] >= start_date:
                start_index = i
                break
        for i in range(len(days)):
            if days[i] <= end_date:
                end_index = i
        days = days[start_index: end_index + 1]

        data = pd.DataFrame(columns=list(self.all_data["2010-01-04"].columns),
                            index=days)

        for i in list(days):
            # print(self.all_data[i].loc[self.stock_ID])
            data.loc[i] = self.all_data[i].loc[self.stock_ID]
            # print(data.head())

        return data

    def get_close_price(self, start_date='2010-01-04', end_date='2019-12-31'):
        '''
        获取单个股票在给定时间段内的每日收盘价，类似的也可以仿照得到其他信息。
        :param start_date: 开始时间
        :param end_date: 结束时间
        :return: 每日收盘价的Series
        '''
        data = self.get_all_info_of_one_stack(start_date, end_date)
        close_prices = data.loc[:, 'close']

        return close_prices

    def get_open_price(self):
        '''
        如果需要用到该函数，直接仿照 get_close_price() 来写即可。
        :return: 
        '''
        pass


# 调用举例
if __name__ == '__main__':

    si = StockInfo('600150.XSHG')
    data = si.get_all_info_of_one_stack(start_date='2010-01-17', end_date='2011-06-12')
    # data就是一个指定股票的DataFrame，一行是一天，一列是一个因子。

    close_prices = si.get_close_price()
    # close_prices 是一个Series，为日期到收盘价的映射