import pandas as pd
import numpy as np
import pickle
import pre_handle
from fileRW import ReadFile


class Strategy:

    def __init__(self):
        pass

    def renew(self, dict1, list_len, key_list, day_idx):
        # 查找出下标为day_idx前的三天的收益的series，索引为股票名，值为受益，
        # 并且series是按照从高到低排序,然后将股票持仓更新到shares_list
        sub = dict1[key_list[day_idx - 1]]['close'] - dict1[key_list[day_idx - 3]]['open']
        sub_series = pd.Series(sub, index=sub.index)
        sub_series = sub_series.sort_values(ascending=False)

        return pd.Series(sub_series.index[:list_len], index=np.arange(list_len))

    def get_income(self, start, end, dict1, share_series, list_len):
        # 获得一个持仓周期结束的时候涨幅涨了多少
        sub = (dict1[end]['close'] - dict1[start]['open']) / dict1[start]['open']
        sub_series = pd.Series(sub, index=sub.index)
        sum = 0

        for i in range(list_len):
            sum += sub_series[share_series[i]]

        return sum / list_len

    def increase_strategy(self, i, pre_day_idx, key_list, dict1, share_series, position_num, open_cycle, hier_pct):
        '''
        涨幅策略：
        将第四天设置为起始天，设置一个持仓周期比如一个月(交易日一共有22天)，然后卖出再买入
        买入策略是：买入三天中最大涨幅的股票，然后计算直到最后一天的涨幅
        需要的结果是每天的收益，所以结果是一个一个series，索引是日期，值是收益率
        需要设置的变量：当前股票开始持仓的时间，持仓周期（一个月），开始是第几天（第四天）,设置要持仓的股票数量（设为10只，然后每只金钱总量相同），当前持仓的股票是哪些（一个Series，标签为普通序号，值为股票名称）
        需要将每天的持仓情况都存下来，这里其实就是在shares_list里面
        '''

        if i < 2:
            return pre_day_idx, hier_pct, share_series

        elif i == 2:
            sub = dict1[key_list[i]]['close'] - dict1[key_list[0]]['open']
            sub_series = pd.Series(sub, index=sub.index)
            sub_series = sub_series.sort_values(ascending=False)
            share_series = pd.Series(sub_series.index[:position_num], index=np.arange(position_num))
            return pre_day_idx, hier_pct, share_series

        elif i - pre_day_idx == open_cycle or i == len(key_list) - 1:
            # 一个持仓周期到了,需要计算新的收益，然后更新新的持仓，然后更新新的pre_day_idx
            new_income = self.get_income(key_list[pre_day_idx], key_list[i], dict1, share_series,
                                         position_num)  # 获得本周期内的涨幅
            hier_pct = (1 + hier_pct) * (new_income + 1) - 1  # 更新新的总收益率
            share_series = self.renew(dict1, position_num, key_list, i)  # 更新当前持仓
            pre_day_idx = i

        return pre_day_idx, hier_pct, share_series

    def double_average_strategy(self, i, dict1, key_list, shorter_timelen, longer_timelen, position_state):
        """
        双均线策略：
        两条均线，一条是短均线，一天是长均线，这里取5日短均线和10日长均线，对所有的股票并行使用双均线策略，
        我们这里使用所有的股票，持仓的话，为了简化，直接用那三只股票之外的21只股票，每只股票初始的金钱都相同
        每当短均线超过长均线就买入，每当长均线超过短均线就卖出

        然后需要记录每天的一个收益，所以需要记录每天是否持有该股票，如果持有该股票，
        那么今天为止的收益就是(今天的价格减去本周期开始持仓时的价格)除以开始持仓时的价格和开始持仓的收益的累积
        如果没有持有该股票，当天的收益就是昨天的收益
        因为需要计算每只股票的近5天的收盘价的均值以及近10天的收盘价的均值

        返回值需要是包含每天的收益，以及每天的持仓，所以一个是series索引为日期，值为收益率
        另一个是字典键为日期，值为series，索引为股票名称，值为每只股票的累积收益率
        需要设置的变量：
        需要dataframe记录一下每只股票的开仓时间，行索引是股票名称，列索引有开仓时间以及股票是否持有（0代表未持有，1代表持有）
        以及如果处于持有的状态就是开仓时候的收益率，如果没持有就是上次持有的最后一天时候的收益率

        现在的算法处理流程就是每过一天就将一天的收益存入series，索引为日期，值为收益率（也就是每只股票的收益的平均值，
        通过dataframe中的最后一列的收益率以及从开始持仓到今天的收益率的乘积得到）
        """

        shr_series = dict1[key_list[i - shorter_timelen]].copy()['close']
        for j in range(shorter_timelen - 1):
            shr_series += dict1[key_list[i - j - 1]]['close']
        shr_series /= shorter_timelen
        lnger_series = dict1[key_list[i - longer_timelen]].copy()['close']
        for j in range(longer_timelen - 1):
            lnger_series += dict1[key_list[i - 1 - j]]['close']
        lnger_series /= longer_timelen

        return shr_series, lnger_series