# -*- coding: utf-8 -*-
"""
因子数据类

提供因子数据的读取和查询功能
"""

import os
import time
import datetime

import numpy as np
import pandas as pd

from .quote_base import QuoteBase, assert_quote_init
from .utils import formatter_list, formatter_st


class QuoteFactor(QuoteBase):
    """因子数据类

    提供因子数据的读取和查询功能，包括：
    - 因子列表查询
    - 因子时序数据
    - 因子快照数据

    """

    @assert_quote_init
    def get_factors(self, category=None) -> pd.DataFrame:
        """获取所有因子列表

        参数:
            category: 因子类别筛选；类型 str/list；默认值 None 获取所有因子
                      可选值: 'basics', 'emotion', 'growth', 'momentum', 'pershare',
                              'quality', 'risk', 'style', 'technical'

        返回值:
            DataFrame 对象：
                category：因子类别
                factor：因子名称
                intro：因子简介

        """
        category = formatter_list(category)

        filename = 'quote-factor-map.csv'
        df = self._fetch_data(filename)
        if df is None:
            raise Exception('fetch data from %s failed.' % filename)
        df = df.copy()
        if category:
            df = df[df['category'].isin(category)]
        df.reset_index(drop=True, inplace=True)
        return df

    @assert_quote_init
    def get_factor_data(self, security, factors=None, start_date=None, end_date=None, count=None) -> pd.DataFrame:
        """获取股票因子时序数据

        参数:
            security: 一支股票代码或者一个股票代码的 list
            factors: 因子名称；类型 str/list；默认值 None 获取所有因子
            start_date: 开始日期，与 count 二选一；类型 str/struct_time/datetime.date/datetime.datetime
            end_date: 结束日期；类型 str/struct_time/datetime.date/datetime.datetime；默认值最新行情日
            count: 结果集行数，与 start_date 二选一

        返回值:
            DataFrame 对象，包含因子数据：
                datetime：日期时间
                code：证券代码
                <factor>：各因子列

        """
        # 参数格式标准化
        security = formatter_list(security)
        factors = formatter_list(factors)
        start_date = formatter_st(start_date) if start_date else None
        end_date = formatter_st(end_date, self.get_quote_date())
        if not start_date and not count:
            raise ValueError('start_date or count must to be set.')
        elif start_date and count:
            raise ValueError('start_date and count cannot be set both.')

        dt_start = np.datetime64(
            datetime.datetime.fromtimestamp(time.mktime(start_date))) if start_date else None
        dt_end = np.datetime64(datetime.datetime.fromtimestamp(time.mktime(end_date)))

        result = pd.DataFrame()
        for c in security:
            p = c.find('.')
            if p < 0:
                continue
            sec_code = c[0:p]
            exc_code = c[p + 1:]
            # 转换交易所代码: XSHG -> SH, XSHE -> SZ（支持两种格式）
            if exc_code == 'XSHG':
                exc_code = 'SH'
            elif exc_code == 'XSHE':
                exc_code = 'SZ'
            filename = os.path.join('factors', exc_code, exc_code + '_FACTORS_' + sec_code + '.csv')
            df = self._fetch_data(filename, date_columns='datetime')
            if df is None:
                continue
            df = df.copy()
            # 日期筛选
            if dt_start and dt_end:
                df = df[(df['datetime'] >= dt_start) & (df['datetime'] <= dt_end)]
            elif dt_end:
                df = df[df['datetime'] <= dt_end]
            if count and count > 0:
                df = df.iloc[-count:]
            # 插入证券代码（统一使用 SH/SZ 后缀）
            exch_suffix = 'SH' if exc_code == 'SH' or exc_code == 'XSHG' else 'SZ'
            df.insert(loc=0, column='code', value=sec_code + '.' + exch_suffix)
            # 因子列筛选
            if factors:
                cols_retain = ['datetime', 'code'] + factors
                cols_all = list(df.columns)
                cols_drop = [x for x in cols_all if x not in cols_retain]
                df.drop(columns=cols_drop, inplace=True)
            result = pd.concat([result, df], sort=False, ignore_index=True)

        result.reset_index(drop=True, inplace=True)
        return result

    @assert_quote_init
    def get_factor_static(self, date=None, factors=None, exch=None) -> pd.DataFrame:
        """获取全市场因子快照数据

        参数:
            date: 行情日期；类型 str/struct_time/datetime.date/datetime.datetime；默认值最新行情日
            factors: 因子名称；类型 str/list；默认值 None 获取所有因子
            exch: 交易所；默认值 ['SH', 'SZ']

        返回值:
            DataFrame 对象，包含因子数据：
                datetime：日期时间
                code：证券代码（完整格式，如 600000.SH）
                <factor>：各因子列

        """
        factors = formatter_list(factors)
        exch = formatter_list(exch, ['SH', 'SZ'])
        curr_date = formatter_st(date, self.get_quote_date())

        result = pd.DataFrame()
        for x in exch:
            filename = os.path.join('static_factors', x, time.strftime('%Y', curr_date),
                                    x + '_STATIC_FACTORS_' + time.strftime('%Y%m%d', curr_date) + '.csv')
            df = self._fetch_data(filename, date_columns='datetime')
            if df is None:
                continue
            df = df.copy()
            # 证券代码补全交易所后缀
            df['code'] = df['code'] + '.' + x
            # 因子列筛选
            if factors:
                cols_retain = ['datetime', 'code'] + factors
                cols_all = list(df.columns)
                cols_drop = [x for x in cols_all if x not in cols_retain]
                df.drop(columns=cols_drop, inplace=True)
            result = pd.concat([result, df], sort=False, ignore_index=True)
        return result