# -*- coding: UTF-8 -*-

import os
import time
import tushare as ts
import pandas as pd
from datetime import datetime

from tma import DATA_PATH

# A股交易日历
# --------------------------------------------------------------------

FILE_CALENDAR = os.path.join(DATA_PATH, 'calendar.csv')
if os.path.exists(FILE_CALENDAR) and \
        time.time() - os.path.getmtime(FILE_CALENDAR) < 3600 * 24:
    trade_calendar = pd.read_csv(FILE_CALENDAR, encoding='utf-8')
else:
    trade_calendar = ts.trade_cal()  # tushare提供的交易日历
    trade_calendar.to_csv(FILE_CALENDAR, index=False, encoding='utf-8')


def is_trade_day(date):
    """判断date日期是不是交易日

    :param date: str or datetime.date, 如 2018-03-15
    :return: Bool
    """
    trade_day = trade_calendar[trade_calendar["isOpen"] == 1]
    trade_day_list = list(trade_day['calendarDate'])
    if isinstance(date, datetime):
        date = str(date.date())
    if date in trade_day_list:
        return True
    else:
        return False


def is_in_trade_time():
    """判断当前是否是交易时间"""
    today = str(datetime.now().date())
    if not is_trade_day(today):  # 如果今天不是交易日，返回False
        return False
    t1 = today + " " + "09:30:00"
    t1 = datetime.strptime(t1, "%Y-%m-%d %H:%M:%S")
    t2 = today + " " + "11:30:00"
    t2 = datetime.strptime(t2, "%Y-%m-%d %H:%M:%S")
    t3 = today + " " + "13:00:00"
    t3 = datetime.strptime(t3, "%Y-%m-%d %H:%M:%S")
    t4 = today + " " + "15:00:00"
    t4 = datetime.strptime(t4, "%Y-%m-%d %H:%M:%S")
    if t1 <= datetime.now() <= t2 or t3 <= datetime.now() <= t4:
        return True
    else:
        return False


def get_recent_trade_days(date, n=10):
    """返回任一日期date前后的n个交易日日期

    :param n: int
        返回交易日的数量n。
        如果 n > 0，则返回date日期未来的n个交易日日期；
        如果 n < 0，则返回date日期过去的n个交易日日期。
    :param date: str  默认值 today
    :return: list
        n个交易日日期列表
    """
    cal_date = list(trade_calendar['calendarDate'])
    date_index = cal_date.index(date)
    if n >= 0:
        recent_days = cal_date[date_index:date_index + n + 20]
    else:
        recent_days = cal_date[date_index + n - 20:date_index + 1]
        recent_days = recent_days[::-1]

    res = []
    for d in recent_days:
        if is_trade_day(d):
            res.append(d)
        if len(res) == abs(n):
            break
    return res


class Calendar:
    """A股交易日历"""

    def __init__(self):
        self.calendar = trade_calendar

    @staticmethod
    def is_trade_time():
        return is_in_trade_time()

    @staticmethod
    def is_trade_day(date):
        return is_trade_day(date)

    @staticmethod
    def recent_trade_days(n, date=None):
        if date is None:
            date = datetime.now().date().__str__()
        return get_recent_trade_days(date=date, n=n)


# --------------------------------------------------------------------

def debug_print(msg, level="INFO"):
    from tma import DEBUG, logger
    LOGS = {
        "INFO": logger.info,
        "DEBUG": logger.debug,
        "EXCEPTION": logger.exception
    }
    if DEBUG:
        LOGS[level](msg)
