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

# %% auto 0
__all__ = ['do_on_dfs', 'do_on_dfs_class', 'concat_sqls', '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

# %% ../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"]
    if latest_version!=(now_version:=get_config(cfg_name='settings.ini')['version']):
        print(f'您当前使用的是{now_version}，最新版本为{latest_version}，建议您更新')
except Exception:
    try:
        response = requests.get("https://pypi.org/pypi/factor-reader/json", timeout=2)
        latest_version = response.json()["info"]["version"]
        if latest_version!=(now_version:=__version__):
            logger.warning(f'您当前使用的是{now_version}，最新版本为{latest_version}，建议您更新至最新版～')
    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, Union[list, 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, Union[list, tuple]):
            dfs = [
                func(
                    self,
                    i,
                    *[_list_value(i, num) for i in args],
                    **_dict_value(kwargs, num)
                )
                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, Union[list, tuple]) and "date" in list(dfs[0].columns):
            dfs = reduce(lambda x, y: pd.merge(x, y, on=["date", "code"]), dfs)
        return dfs

    return wrapper

# %% ../nbs/factor_reader_core.ipynb 6
class FactorReader(object):
    def __init__(self, token: str) -> None:
        """输入token以连接因子数据库"""
        logger.info("正在初始化并连接因子数据库，请稍等，过程大约需要1分钟……")
        self.__ds_object = deeplake.load("hub://chenzongwei/factor", token=token)
        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)}
        
    @staticmethod


    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("已经是最新的了，无需更新")

    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
    @do_on_dfs_class
    def read_factor(
        self,
        fac_key: str,
        trade_date: int = None,
        start_date: int = None,
        end_date: int = None,
        sql_like: bool = False,
    ) -> pd.DataFrame:
        """通过表名，读取因子数据

        Parameters
        ----------
        fac_key : str
            表的名称或因子的名称
        trade_date : int, optional
            读取单日因子值，形如20230113，指定此参数时，start_date和end_date两个参数将失效, by default None
        start_date : int, optional
            读取因子值的起始日期，形如20130101, by default None
        end_date : int, optional
            读取因子值的终止日期，形如20221231, 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"

        if fac_key in self.keys:
            print(f"正在读取{self.keys_names[fac_key]}的数据")
        elif len(names_related := [i for i in self.names if fac_key in i]) > 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:
            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:
                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 = [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
