import warnings;

from __future__ import print_function
from __future__ import division

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
%matplotlib inline
sns.set_context(rc={'figure.figsize': (14, 7) } )
figzize_me = figsize =(14, 7)
# import warnings; 
# warnings.filterwarnings('ignore')

import os
import sys
# 使用insert 0即只使用github，避免交叉使用了pip安装的abupy，导致的版本不一致问题
sys.path.insert(0, os.path.abspath('../'))
import abupy
# 使用沙盒数据，目的是和书中一样的数据环境
abupy.env.enable_example_env_ipython()

from abupy import AbuFactorBuyBreak
from abupy import AbuBenchmark
# buy_factors 60日向上突破，42日向上突破两个因子
buy_factors = [{'xd': 60, 'class': AbuFactorBuyBreak}, 
               {'xd': 42, 'class': AbuFactorBuyBreak}]
benchmark = AbuBenchmark()

from abupy import AbuPickTimeWorker 
from abupy import AbuCapital
from abupy import AbuKLManager

capital = AbuCapital(1000000, benchmark)
kl_pd_manager = AbuKLManager(benchmark, capital)
# 获取TSLA的交易数据
kl_pd = kl_pd_manager.get_pick_time_kl_pd('usTSLA')
abu_worker = AbuPickTimeWorker(capital, kl_pd, benchmark, buy_factors, None)
%time abu_worker.fit()

from abupy import ABuTradeProxy
orders_pd, action_pd, _ = ABuTradeProxy.trade_summary(abu_worker.orders, kl_pd, draw=True)

from abupy import ABuTradeExecute
ABuTradeExecute.apply_action_to_capital(capital, action_pd, kl_pd_manager)
capital.capital_pd.capital_blance.plot()

from abupy import AbuFactorSellBreak
# 使用120天向下突破为卖出信号
sell_factor1 = {'xd': 120, 'class': AbuFactorSellBreak}

from abupy import ABuPickTimeExecute
# 只使用120天向下突破为卖出因子
sell_factors = [sell_factor1]
capital = AbuCapital(1000000, benchmark)
orders_pd, action_pd, _ = ABuPickTimeExecute.do_symbols_with_same_factors(
                          ['usTSLA'], benchmark, buy_factors, sell_factors, capital, show=True)

from abupy import AbuFactorAtrNStop
# 趋势跟踪策略止盈要大于止损设置值，这里0.5，3.0
sell_factor2 = {'stop_loss_n':0.5, 'stop_win_n':3.0, 'class': AbuFactorAtrNStop}
# 两个卖出因子策略并行同时生效
sell_factors = [sell_factor1, sell_factor2]
capital = AbuCapital(1000000, benchmark)
orders_pd, action_pd, _ = ABuPickTimeExecute.do_symbols_with_same_factors(
                          ['usTSLA'], benchmark, buy_factors, sell_factors, capital, show=True)

from abupy import AbuFactorPreAtrNStop
# 暴跌止损卖出因子形成dict
sell_factor3 =  {'class': AbuFactorPreAtrNStop, 'pre_atr_n':1.0}
# 三个卖出因子同时生效，组成sell_factors
sell_factors = [sell_factor1, sell_factor2, sell_factor3]
capital = AbuCapital(1000000, benchmark)
orders_pd, action_pd, _ = ABuPickTimeExecute.do_symbols_with_same_factors(
                          ['usTSLA'], benchmark, buy_factors, sell_factors, capital, show=True)

from abupy import AbuFactorCloseAtrNStop
# 保护止盈卖出因子组成dict
sell_factor4 = {'class': AbuFactorCloseAtrNStop, 'close_atr_n':1.5}
# 四个卖出因子同时并行生效
sell_factors = [sell_factor1, sell_factor2, sell_factor3, sell_factor4]
capital = AbuCapital(1000000, benchmark)
orders_pd, action_pd, _ = ABuPickTimeExecute.do_symbols_with_same_factors(
                          ['usTSLA'], benchmark, buy_factors, sell_factors, capital, show=True)

from abupy import AbuSlippageBuyBase

# 修改g_open_down_rate的值为0.02
g_open_down_rate = 0.02

class AbuSlippageBuyMean2(AbuSlippageBuyBase):
    def fit_price(self):
        if (self.kl_pd_buy.open / self.kl_pd_buy.pre_close) < (
            1 - g_open_down_rate):
            # 开盘下跌K_OPEN_DOWN_RATE以上，单子失效
            print(self.factor_name + 'open down threshold')
            return np.inf
        # 买入价格为当天均价
        self.buy_price = np.mean(
            [self.kl_pd_buy['high'], self.kl_pd_buy['low']])
        return self.buy_price

# 只针对60使用AbuSlippageBuyMean2
buy_factors2 = [{'slippage': AbuSlippageBuyMean2, 'xd': 60,
                 'class': AbuFactorBuyBreak},
                {'xd': 42, 'class': AbuFactorBuyBreak}]

sell_factors = [sell_factor1, sell_factor2, sell_factor3,
                sell_factor4]
capital = AbuCapital(1000000, benchmark)
orders_pd, action_pd, _ = \
    ABuPickTimeExecute.do_symbols_with_same_factors(['usTSLA'],
                                                    benchmark,
                                                    buy_factors2,
                                                    sell_factors,
                                                    capital,
                                                    show=True)

# 我们假定choice_symbols是我们选股模块的结果，
choice_symbols = ['usTSLA', 'usNOAH', 'usSFUN', 'usBIDU', 'usAAPL', 'usGOOG', 'usWUBA', 'usVIPS']
capital = AbuCapital(1000000, benchmark)
orders_pd, action_pd, all_fit_symbols_cnt = ABuPickTimeExecute.do_symbols_with_same_factors(choice_symbols, benchmark, buy_factors, sell_factors, capital, show=False)

orders_pd[:10].filter(['symbol', 'buy_price', 'buy_cnt', 'buy_factor', 'buy_pos', 
                       'sell_date', 'sell_type_extra', 'sell_type', 'profit'])

action_pd[:10]

from abupy import AbuMetricsBase 
metrics = AbuMetricsBase(orders_pd, action_pd, capital, benchmark)
metrics.fit_metrics()
metrics.plot_returns_cmp(only_show_returns=True)

metrics.gains_mean, -metrics.losses_mean

from abupy import AbuKellyPosition
# 42d使用AbuKellyPosition，60d仍然使用默认仓位管理类
buy_factors2 = [{'xd': 60, 'class': AbuFactorBuyBreak},
                {'xd': 42, 'position': {'class': AbuKellyPosition, 'win_rate': metrics.win_rate, 
                                        'gains_mean': metrics.gains_mean, 'losses_mean': -metrics.losses_mean},
                 'class': AbuFactorBuyBreak}]
capital = AbuCapital(1000000, benchmark)
orders_pd, action_pd, all_fit_symbols_cnt = ABuPickTimeExecute.do_symbols_with_same_factors(choice_symbols, benchmark, buy_factors2, sell_factors, capital, show=False)

orders_pd[:10].filter(['symbol', 'buy_cnt', 'buy_factor', 'buy_pos'])

# 选定noah和sfun
target_symbols = ['usSFUN', 'usNOAH']
# 针对sfun只使用42d向上突破作为买入因子
buy_factors_sfun = [{'xd': 42, 'class': AbuFactorBuyBreak}]
# 针对sfun只使用60d向下突破作为卖出因子
sell_factors_sfun = [{'xd': 60, 'class': AbuFactorSellBreak}]

# 针对noah只使用21d向上突破作为买入因子
buy_factors_noah = [{'xd': 21, 'class': AbuFactorBuyBreak}]
# 针对noah只使用42d向下突破作为卖出因子
sell_factors_noah = [{'xd': 42, 'class': AbuFactorSellBreak}]

factor_dict = dict()
# 构建SFUN独立的buy_factors，sell_factors的dict
factor_dict['usSFUN'] = {'buy_factors':buy_factors_sfun, 'sell_factors':sell_factors_sfun}
# 构建NOAH独立的buy_factors，sell_factors的dict
factor_dict['usNOAH'] = {'buy_factors':buy_factors_noah, 'sell_factors':sell_factors_noah}
# 初始化资金
capital = AbuCapital(1000000, benchmark)
# 使用do_symbols_with_diff_factors执行
orders_pd, action_pd, all_fit_symbols = ABuPickTimeExecute.do_symbols_with_diff_factors(
                                        target_symbols, benchmark, factor_dict, capital)

pd.crosstab(orders_pd.buy_factor, orders_pd.symbol)

# 要关闭沙盒数据环境，因为沙盒里就那几个股票的历史数据, 下面要随机做300个股票
abupy.env.disable_example_env_ipython()

# 关闭沙盒后，首先基准要从非沙盒环境换取，否则数据对不齐，无法正常运行
benchmark = AbuBenchmark()
capital = AbuCapital(1000000, benchmark)

from abupy import ABuMarket
# 当传入choice_symbols为None时代表对整个市场的所有股票进行回测
choice_symbols = None
# 顺序获取市场后50支股票
choice_symbols = ABuMarket.all_symbol()[-50:]
# 随机获取50支股票
choice_symbols = ABuMarket.choice_symbols(50)

from abupy import AbuPickTimeMaster
capital = AbuCapital(1000000, benchmark)
orders_pd, action_pd, _ = AbuPickTimeMaster.do_symbols_with_same_factors_process(
                          choice_symbols, benchmark, buy_factors, sell_factors, 
                          capital)

metrics = AbuMetricsBase(orders_pd, action_pd, capital, benchmark)
metrics.fit_metrics()
metrics.plot_returns_cmp(only_show_returns=True)

# 继续使用沙盒数据环境
abupy.env.enable_example_env_ipython()

# 开启沙盒后，基准也要从沙盒环境取，否则数据对不齐，无法正常运行
benchmark = AbuBenchmark()
capital = AbuCapital(1000000, benchmark)
kl_pd_manager = AbuKLManager(benchmark, capital)

from abupy import AbuPickRegressAngMinMax
from abupy import AbuPickStockWorker
# 选股条件threshold_ang_min=0.0, 即要求股票走势为向上上升趋势
stock_pickers = [{'class': AbuPickRegressAngMinMax, 
                'threshold_ang_min':0.0, 'reversed':False}]

# 从这几个股票里进行选股，只是为了演示方便，一般的选股都会是数量比较多的情况比如全市场股票
choice_symbols = ['usNOAH', 'usSFUN', 'usBIDU', 'usAAPL', 'usGOOG', 'usTSLA', 'usWUBA', 'usVIPS']

capital = AbuCapital(1000000, benchmark)
kl_pd_manager = AbuKLManager(benchmark, capital)
stock_pick = AbuPickStockWorker(capital, benchmark, kl_pd_manager, choice_symbols=choice_symbols, stock_pickers=stock_pickers)
stock_pick.fit()
stock_pick.choice_symbols

from abupy import ABuRegUtil
# 从kl_pd_manager缓存中获取选股走势数据，注意get_pick_stock_kl_pd为选股数据，get_pick_time_kl_pd为择时
kl_pd_noah = kl_pd_manager.get_pick_stock_kl_pd('usNOAH')
# 绘制并计算角度
deg = ABuRegUtil.calc_regress_deg(kl_pd_noah.close)
print('noah 选股周期内角度={}'.format(round(deg, 3)))

from abupy import ABuPickStockExecute
stock_pickers = [{'class': AbuPickRegressAngMinMax, 
                 'threshold_ang_min':0.0, 'threshold_ang_max':10.0, 'reversed':False}]

ABuPickStockExecute.do_pick_stock_work(choice_symbols, benchmark, capital, stock_pickers)

kl_pd_sfun = kl_pd_manager.get_pick_stock_kl_pd('usSFUN')
print('sfun 选股周期内角度={}'.format(round(ABuRegUtil.calc_regress_deg(kl_pd_sfun.close), 3)))

# 和上面的代码唯一的区别就是reversed=True
stock_pickers = [{'class': AbuPickRegressAngMinMax, 
                 'threshold_ang_min':0.0, 'threshold_ang_max':10.0, 'reversed':True}]
ABuPickStockExecute.do_pick_stock_work(choice_symbols, benchmark, capital, stock_pickers)

from abupy import AbuPickStockPriceMinMax

stock_pickers = [{'class': AbuPickRegressAngMinMax, 
                 'threshold_ang_min':0.0, 'reversed':False}, 
                 {'class': AbuPickStockPriceMinMax, 'threshold_price_min':50.0,
                 'reversed':False}]
ABuPickStockExecute.do_pick_stock_work(choice_symbols, benchmark, capital, stock_pickers)

# 继续关闭沙盒环境，因为下面要choice_symbols
abupy.env.disable_example_env_ipython()

# 关闭沙盒后，首先基准要从非沙盒环境换取，否则数据对不齐，无法正常运行
benchmark = AbuBenchmark()
capital = AbuCapital(1000000, benchmark)
kl_pd_manager = AbuKLManager(benchmark, capital)

from abupy import ABuMarket, AbuPickStockMaster
# 首先随抽取50支股票
choice_symbols = ABuMarket.choice_symbols(50)

# 股价在15-50之间
stock_pickers = [{'class': AbuPickStockPriceMinMax, 'threshold_price_min':15.0,
                'threshold_price_max':50.0, 'reversed':False}]
%time cs = AbuPickStockMaster.do_pick_stock_with_process(capital, benchmark, stock_pickers, choice_symbols)

cs