#encoding=utf-8
import qlib
import pandas as pd
import pickle
import xgboost as xgb
import numpy as np
import re
from qlib.constant import REG_US
from qlib.utils import exists_qlib_data, init_instance_by_config
from qlib.workflow import R
from qlib.workflow.record_temp import SignalRecord, PortAnaRecord
from qlib.utils import flatten_dict
from qlib.data import LocalExpressionProvider
from qlib.data.ops import Operators, OpsList
from qlib.data.base import Feature
from pyecharts import options as opts
from pyecharts.charts import Kline, Line, Grid
from my_data_handler import MyAlphaHandler

# model_file = r'.\mlruns\1\d6536b056ba84a74be6b33971f443cf6\artifacts\trained_model'
model_file = r'.\mlruns\1\148ef1cd7acd48deac3eadc339ad3008\artifacts\trained_model'
with open(model_file, 'rb') as fi:
    model = pickle.load(fi)

exprs, columns = MyAlphaHandler.get_custom_config()


raw_data = pd.read_csv('../stock_data/TSLA.csv', parse_dates=['time'])
raw_data['data_time'] = raw_data['time'].dt.strftime("%Y-%m-%d %H:%M:00")
raw_data.set_index('time', inplace=True)
raw_data["vwap"] = np.nan
raw_data.sort_index(inplace=True)
# print(raw_data)
class MyFeature(Feature):
    def _load_internal(self, instrument, start_index, end_index, freq):
        print("load", self._name, instrument, start_index, end_index, freq)
        return raw_data.loc[start_index:end_index][self._name]
Operators.register(OpsList + [MyFeature])

def my_parse_field(field):
    if not isinstance(field, str):
        field = str(field)
    for pattern, new in [(r"\$(\w+)", rf'MyFeature("\1")'), (r"(\w+\s*)\(", r"Operators.\1(")]:  # Features  # Operators
        field = re.sub(pattern, new, field)
    return field

obj = dict()
for field in exprs:
    expression = eval(my_parse_field(field))
    series = expression.load('TSLA', "2022-01-02", "2022-02-28", "1min")
    series = series.astype(np.float32)
    obj[field] = series
data = pd.DataFrame(obj)
data.columns = columns

view_time_start = '2022-02-11'
view_time_end = '2022-02-12'

pre_data = raw_data.loc[view_time_start:view_time_end].copy()
pred=model.model.predict(xgb.DMatrix(data.loc[view_time_start:view_time_end]))
pre_data['pred_score'] = pred
records = pre_data.to_dict("records")

cash = 50000
position = {}
hold_thresh = 5
score_thresh = 0.001
x_axises, y_axises, mark_points, money = [], [], [], []
for record in records:
    x_axises.append(record['data_time'])
    y_axises.append([
        record['open'], record['close'], record['low'], record['high']
    ])
    if 'hold_cnt' in position:
        position['hold_cnt'] += 1
    if position and (record['open'] >= position['close'] * 1.01 or record['open'] < position['close'] * 0.995 or record['pred_score'] < -score_thresh or position['hold_cnt'] >= hold_thresh):
        cash += position['amount'] * record['open']
        position = {}
        #print("sell")
        mark_points.append(opts.MarkPointItem(
            coord=[record['data_time'], record['high']],
            symbol='triangle', symbol_size=7, 
            itemstyle_opts=opts.ItemStyleOpts(color="green")
        ))
    elif record['pred_score'] > score_thresh and not position:
        position = dict(record)
        position['amount'] = int(cash / position['open'])
        cash -= position['amount'] * position['open']
        # buy
        #print("buy")
        position['hold_cnt'] = 0
        mark_points.append(opts.MarkPointItem(
            coord=[record['data_time'], record['high']],
            symbol='arrow', symbol_size=7,
            itemstyle_opts=opts.ItemStyleOpts(color="yellow")
        ))
    cur_money = cash
    if position:
        cur_money += position['amount'] * record['close']
    money.append(cur_money)
if position:
    cash += position['amount'] * records[-1]['close']
print("cash:", cash)

kline_graph = (
    Kline()
    .add_xaxis(x_axises)
    .add_yaxis(
        "kline",
        y_axises,        
        markpoint_opts=opts.MarkPointOpts(
            data=mark_points
        ),
    )
    .set_global_opts(
        xaxis_opts=opts.AxisOpts(is_scale=True),
        yaxis_opts=opts.AxisOpts(
            is_scale=True,
            splitarea_opts=opts.SplitAreaOpts(
                is_show=True, areastyle_opts=opts.AreaStyleOpts(opacity=1)
            ),
        ),
        title_opts=opts.TitleOpts(title="%s_%s" % (view_time_start, view_time_end)),
        datazoom_opts=[opts.DataZoomOpts(type_="inside", xaxis_index=[0, 1],)],
    )
)

kline_line = (
    Line()
    .add_xaxis(xaxis_data=x_axises)
    .add_yaxis(
        series_name="cur_money",
        y_axis=money,
        is_smooth=True,
        linestyle_opts=opts.LineStyleOpts(opacity=0.5),
        label_opts=opts.LabelOpts(is_show=False),
        markline_opts=opts.MarkLineOpts(
            data=[opts.MarkLineItem(y=50000)]
        ),
    )
    .set_global_opts(
        xaxis_opts=opts.AxisOpts(
            type_="category",
            grid_index=2,
            axislabel_opts=opts.LabelOpts(is_show=False),            
        ),
        yaxis_opts=opts.AxisOpts(
            min_='dataMin'
        )
    )
)

grid_chart = Grid(init_opts=opts.InitOpts(width='2000px', height='900px'))
grid_chart.add(
    kline_graph,
    grid_opts=opts.GridOpts(pos_left="3%", pos_right="10%", height="50%"),
)
grid_chart.add(
    kline_line,
    grid_opts=opts.GridOpts(
        pos_left="3%", pos_right="10%", pos_top="60%", height="30%"
    ),
)
grid_chart.render("kline_markline.html")
[((25, 17, 25, 51), 'my_data_handler.MyAlphaHandler.get_custom_config', 'MyAlphaHandler.get_custom_config', ({}, {}), '()', False, 'from my_data_handler import MyAlphaHandler\n'), ((28, 11, 28, 70), 'pandas.read_csv', 'pd.read_csv', (), '', True, 'import pandas as pd\n'), ((38, 0, 38, 41), 'qlib.data.ops.Operators.register', 'Operators.register', ({(38, 19, 38, 40): '(OpsList + [MyFeature])'}, {}), '(OpsList + [MyFeature])', False, 'from qlib.data.ops import Operators, OpsList\n'), ((53, 7, 53, 24), 'pandas.DataFrame', 'pd.DataFrame', ({(53, 20, 53, 23): 'obj'}, {}), '(obj)', True, 'import pandas as pd\n'), ((23, 12, 23, 27), 'pickle.load', 'pickle.load', ({(23, 24, 23, 26): 'fi'}, {}), '(fi)', False, 'import pickle\n'), ((60, 25, 60, 77), 'xgboost.DMatrix', 'xgb.DMatrix', ({(60, 37, 60, 76): 'data.loc[view_time_start:view_time_end]'}, {}), '(data.loc[view_time_start:view_time_end])', True, 'import xgboost as xgb\n'), ((44, 16, 44, 43), 're.sub', 're.sub', ({(44, 23, 44, 30): 'pattern', (44, 32, 44, 35): 'new', (44, 37, 44, 42): 'field'}, {}), '(pattern, new, field)', False, 'import re\n'), ((116, 19, 116, 47), 'pyecharts.options.AxisOpts', 'opts.AxisOpts', (), '', True, 'from pyecharts import options as opts\n'), ((123, 19, 123, 83), 'pyecharts.options.TitleOpts', 'opts.TitleOpts', (), '', True, 'from pyecharts import options as opts\n'), ((147, 19, 149, 9), 'pyecharts.options.AxisOpts', 'opts.AxisOpts', (), '', True, 'from pyecharts import options as opts\n'), ((153, 28, 153, 73), 'pyecharts.options.InitOpts', 'opts.InitOpts', (), '', True, 'from pyecharts import options as opts\n'), ((156, 14, 156, 73), 'pyecharts.options.GridOpts', 'opts.GridOpts', (), '', True, 'from pyecharts import options as opts\n'), ((160, 14, 162, 5), 'pyecharts.options.GridOpts', 'opts.GridOpts', (), '', True, 'from pyecharts import options as opts\n'), ((124, 23, 124, 77), 'pyecharts.options.DataZoomOpts', 'opts.DataZoomOpts', (), '', True, 'from pyecharts import options as opts\n'), ((111, 23, 113, 9), 'pyecharts.options.MarkPointOpts', 'opts.MarkPointOpts', (), '', True, 'from pyecharts import options as opts\n'), ((135, 23, 135, 54), 'pyecharts.options.LineStyleOpts', 'opts.LineStyleOpts', (), '', True, 'from pyecharts import options as opts\n'), ((136, 19, 136, 48), 'pyecharts.options.LabelOpts', 'opts.LabelOpts', (), '', True, 'from pyecharts import options as opts\n'), ((145, 27, 145, 56), 'pyecharts.options.LabelOpts', 'opts.LabelOpts', (), '', True, 'from pyecharts import options as opts\n'), ((83, 27, 83, 60), 'pyecharts.options.ItemStyleOpts', 'opts.ItemStyleOpts', (), '', True, 'from pyecharts import options as opts\n'), ((95, 27, 95, 61), 'pyecharts.options.ItemStyleOpts', 'opts.ItemStyleOpts', (), '', True, 'from pyecharts import options as opts\n'), ((106, 4, 106, 11), 'pyecharts.charts.Kline', 'Kline', ({}, {}), '()', False, 'from pyecharts.charts import Kline, Line, Grid\n'), ((120, 45, 120, 74), 'pyecharts.options.AreaStyleOpts', 'opts.AreaStyleOpts', (), '', True, 'from pyecharts import options as opts\n'), ((129, 4, 129, 10), 'pyecharts.charts.Line', 'Line', ({}, {}), '()', False, 'from pyecharts.charts import Kline, Line, Grid\n'), ((138, 18, 138, 44), 'pyecharts.options.MarkLineItem', 'opts.MarkLineItem', (), '', True, 'from pyecharts import options as opts\n')]