# -*- coding: utf-8 -*-
#
# Copyright 2017 Ricequant, Inc
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import datetime

import numpy as np
import numpy.lib.recfunctions as nprf

from rqalpha.interface import AbstractEventSource
from rqalpha.events import Event, EVENT
from rqalpha.environment import Environment
from rqalpha.utils import get_account_type
from rqalpha.utils.exception import CustomException, CustomError, patch_user_exc
from rqalpha.utils.datetime_func import convert_int_to_datetime, convert_date_to_date_int, convert_dt_to_int
from rqalpha.const import ACCOUNT_TYPE
from rqalpha.utils.i18n import gettext as _
from rqalpha.model.tick import Tick


ONE_MINUTE = datetime.timedelta(minutes=1)


class SimulationEventSource(AbstractEventSource):
    def __init__(self, env, account_list):
        self._env = env
        self._account_list = account_list
        self._universe_changed = False
        Environment.get_instance().event_bus.add_listener(EVENT.POST_UNIVERSE_CHANGED, self._on_universe_changed)

    def _on_universe_changed(self, event):
        self._universe_changed = True

    def _get_universe(self):
        universe = Environment.get_instance().get_universe()
        if len(universe) == 0 and ACCOUNT_TYPE.STOCK not in self._account_list:
            error = CustomError()
            error.set_msg("Current universe is empty. Please use subscribe function before trade")
            raise patch_user_exc(CustomException(error))
        return universe

    # [BEGIN] minute event helper
    @staticmethod
    def _get_stock_trading_minutes(trading_date):
        trading_minutes = set()
        current_dt = datetime.datetime.combine(trading_date, datetime.time(9, 31))
        am_end_dt = current_dt.replace(hour=11, minute=30)
        pm_start_dt = current_dt.replace(hour=13, minute=1)
        pm_end_dt = current_dt.replace(hour=15, minute=0)
        delta_minute = datetime.timedelta(minutes=1)
        while current_dt <= am_end_dt:
            trading_minutes.add(current_dt)
            current_dt += delta_minute

        current_dt = pm_start_dt
        while current_dt <= pm_end_dt:
            trading_minutes.add(current_dt)
            current_dt += delta_minute
        return trading_minutes

    def _get_future_trading_minutes(self, trading_date):
        trading_minutes = set()
        universe = self._get_universe()
        for order_book_id in universe:
            if get_account_type(order_book_id) == ACCOUNT_TYPE.STOCK:
                continue
            trading_minutes.update(self._env.data_proxy.get_trading_minutes_for(order_book_id, trading_date))
        return set([convert_int_to_datetime(minute) for minute in trading_minutes])

    def _get_trading_minutes(self, trading_date):
        trading_minutes = set()
        for account_type in self._account_list:
            if account_type == ACCOUNT_TYPE.STOCK:
                trading_minutes = trading_minutes.union(self._get_stock_trading_minutes(trading_date))
            elif account_type == ACCOUNT_TYPE.FUTURE:
                trading_minutes = trading_minutes.union(self._get_future_trading_minutes(trading_date))
        return sorted(list(trading_minutes))
    # [END] minute event helper

    def _get_merge_ticks(self, date, last_dt):
        self.today_ticks_cache = {}
        data_proxy = self._env.data_proxy
        self.order_book_id_map = {}

        # get today ticks
        for idx, order_book_id in enumerate(self._get_universe()):
            self.order_book_id_map[idx] = order_book_id
            ticks = data_proxy.get_ticks(order_book_id, date)
            if ticks is None:
                continue
            # append idx columns
            ticks = nprf.append_fields(ticks, 'order_book_id', np.full(len(ticks), idx), usemask=False)
            self.today_ticks_cache[order_book_id] = ticks

        # merge ticks arrays into one
        ticks_list = list(self.today_ticks_cache.values())
        if len(ticks_list) == 0:
            return []
        merge_ticks = np.concatenate(ticks_list)

        merge_ticks.sort(kind='mergesort', order=("date", "time"))

        time_arr = merge_ticks["date"].astype(np.uint64) * 1000000000 + merge_ticks["time"]
        last_idx = time_arr.searchsorted(convert_dt_to_int(last_dt) * 1000)
        merge_ticks = merge_ticks[last_idx:]

        return merge_ticks

    def events(self, start_date, end_date, frequency):
        if frequency == "1d":
            # 根据起始日期和结束日期，获取所有的交易日，然后再循环获取每一个交易日
            for day in self._env.data_proxy.get_trading_dates(start_date, end_date):
                date = day.to_pydatetime()
                dt_before_trading = date.replace(hour=0, minute=0)
                dt_bar = date.replace(hour=15, minute=0)
                dt_after_trading = date.replace(hour=15, minute=30)
                dt_settlement = date.replace(hour=17, minute=0)
                yield Event(EVENT.BEFORE_TRADING, calendar_dt=dt_before_trading, trading_dt=dt_before_trading)
                yield Event(EVENT.BAR, calendar_dt=dt_bar, trading_dt=dt_bar)

                yield Event(EVENT.AFTER_TRADING, calendar_dt=dt_after_trading, trading_dt=dt_after_trading)
                yield Event(EVENT.SETTLEMENT, calendar_dt=dt_settlement, trading_dt=dt_settlement)
        elif frequency == '1m':
            for day in self._env.data_proxy.get_trading_dates(start_date, end_date):
                before_trading_flag = True
                date = day.to_pydatetime()
                last_dt = None
                done = False

                dt_before_day_trading = date.replace(hour=8, minute=30)

                while True:
                    if done:
                        break
                    exit_loop = True
                    trading_minutes = self._get_trading_minutes(date)
                    for calendar_dt in trading_minutes:
                        if last_dt is not None and calendar_dt < last_dt:
                            continue

                        if calendar_dt < dt_before_day_trading:
                            trading_dt = calendar_dt.replace(year=date.year,
                                                             month=date.month,
                                                             day=date.day)
                        else:
                            trading_dt = calendar_dt
                        if before_trading_flag:
                            before_trading_flag = False
                            before_trading_dt = trading_dt - datetime.timedelta(minutes=30)
                            yield Event(EVENT.BEFORE_TRADING, calendar_dt=before_trading_dt,
                                        trading_dt=before_trading_dt)
                        if self._universe_changed:
                            self._universe_changed = False
                            last_dt = calendar_dt
                            exit_loop = False
                            break
                        # yield handle bar
                        yield Event(EVENT.BAR, calendar_dt=calendar_dt, trading_dt=trading_dt)
                    if exit_loop:
                        done = True

                dt = date.replace(hour=15, minute=30)
                yield Event(EVENT.AFTER_TRADING, calendar_dt=dt, trading_dt=dt)

                dt = date.replace(hour=17, minute=0)
                yield Event(EVENT.SETTLEMENT, calendar_dt=dt, trading_dt=dt)
        elif frequency == "tick":
            data_proxy = self._env.data_proxy
            for day in data_proxy.get_trading_dates(start_date, end_date):
                date = day.to_pydatetime()
                last_dt = date
                done = False

                before_trading_dt = date.replace(hour=8, minute=30)
                yield Event(EVENT.BEFORE_TRADING, calendar_dt=before_trading_dt,
                            trading_dt=before_trading_dt)

                while True:
                    if done:
                        break
                    exit_loop = True

                    merge_ticks = self._get_merge_ticks(date, last_dt)

                    for idx in range(len(merge_ticks)):
                        raw_tick = merge_ticks[idx]
                        order_book_id = self.order_book_id_map[raw_tick[-1]]

                        tick = Tick(order_book_id, raw_tick)
                        dt = tick.datetime

                        yield Event(EVENT.TICK, calendar_dt=dt, trading_dt=dt, tick=tick)

                        if self._universe_changed:
                            self._universe_changed = False
                            last_dt = dt
                            exit_loop = False
                            break

                    if exit_loop:
                        done = True

                dt = date.replace(hour=15, minute=30)
                yield Event(EVENT.AFTER_TRADING, calendar_dt=dt, trading_dt=dt)

                dt = date.replace(hour=17, minute=0)
                yield Event(EVENT.SETTLEMENT, calendar_dt=dt, trading_dt=dt)
        else:
            raise NotImplementedError(_("Frequency {} is not support.").format(frequency))
