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

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

# %% ../nbs/factor_reader_core.ipynb 2
import pandas as pd
from loguru import logger
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
import psycopg2 as pg
from tenacity import retry, stop_after_attempt

# %% ../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_nosqls(func):
    @wraps(func)
    def wrapper(self,fac_key,*args,**kwargs):
        dfs=func(self,fac_key,*args,**kwargs)
        def double_index(res,key):
            res=res.reset_index().assign(name=self._deal_key_name(key)[1])
            res=res.set_index(['name','date'])
            return res
        if (isinstance(dfs, list) or isinstance(dfs, tuple)) and ("date" not in dfs[0].columns):
            dfs=[double_index(i,self._deal_key_name(key)[1]) for i,key in zip(dfs,fac_key)]
            dfs=pd.concat(dfs)
        elif (isinstance(dfs, list) or isinstance(dfs, tuple)) and ("date" in dfs[0].columns):
            dfs = reduce(
                lambda x, y: pd.merge(x, y, on=["date", "code"], how="outer"), dfs
            )
        elif "date" not in dfs.columns:
            dfs=double_index(dfs,fac_key)
        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],
                            "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():

    def __init__(
        self,
        token: str,
        user: str = "admin",
        password: str = "quest",
        host: str = "43.143.223.158",
        port: str = "8812",
        database: str = "qdb",
    ) -> None:
        """通过postgre的psycopg2驱动连接questdb数据库

        Parameters
        ----------
        token : str
            读取数据的凭证
        user : str, optional
            用户名, by default "admin"
        password : str, optional
            密码, by default "quest"
        host : str, optional
            地址, by default "43.143.223.158"
        port : str, optional
            端口, by default "8812"
        database : str, optional
            数据库, by default "qdb"
        """
        self.token = token
        self.user = user
        self.password = password
        self.host = host
        self.port = port
        self.database = database
        infos=self.show_all_factors_information()
        self.keys=list(infos.数据键名)
        self.names=list(infos.因子名称)
        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 __connect(self):
        conn = pg.connect(
            user=self.user,
            password=self.password,
            host=self.host,
            port=self.port,
            database=self.database,
        )
        return conn
    
    @retry(stop=stop_after_attempt(10))
    def __get_data(
        self, sql_order: str
    ) -> pd.DataFrame:
        """以sql命令的方式，从数据库中读取数据

        Parameters
        ----------
        sql_order : str
            sql命令

        Returns
        -------
        pd.DataFrame
            读取的结果
        """
        conn = self.__connect()
        cursor = conn.cursor()
        cursor.execute(sql_order)
        df_data = cursor.fetchall()
        columns = [i[0] for i in cursor.description]
        df = pd.DataFrame(df_data, columns=columns)
        return df
    
    def _deal_key_name(self,input_key):
        names_related = [i for i in self.names if input_key in i]
        if input_key in self.keys:
            fac_key=input_key
            fac_name=self.keys_names[fac_key]
        elif len(names_related) > 0:
            fac_name = names_related[0]
            fac_key = self.names_keys[fac_name]
        else:
            raise ValueError(
                "输入的fac_key参数有误，请输入因子名称或因子键名，可通过`factor_reader.show_all_factors_information`函数来查看可用的因子名称和键名"
            )
        return fac_key,fac_name
    
    def __update_factor(self,table_name:str,df:pd.DataFrame):
        tables=self.__get_data("show tables").table.tolist()
        if table_name in tables:
            logger.info(f'{table_name}已经存在了，即将更新')
            old_end=self.__get_data(f"select max(date) from {table_name}").iloc[0,0]
            new=df[df.index>old_end]
            new=new.stack().reset_index()
            new.columns=['date','code','fac']
        else:
            logger.info(f'{table_name}第一次上传')
            new=df.stack().reset_index()
            new.columns=['date','code','fac']
        self.__write_via_df(new,table_name)
        
    def __write_via_df(
        self,
        df: pd.DataFrame,
        table_name: str,
        symbols= None,
        tuple_col= None,
    ) -> None:
        """通过questdb的python库直接将dataframe写入quested数据库

        Parameters
        ----------
        df : pd.DataFrame
            要写入的dataframe
        table_name : str
            questdb中该表的表名
        symbols : Union[str, bool, list[int], list[str]], optional
            为symbols的那些列的名称, by default None
        tuple_col : Union[str, list[str]], optional
            数据类型为tuple或list的列的名字, by default None
        """
        if tuple_col is None:
            ...
        elif isinstance(tuple_col, str):
            df[tuple_col] = df[tuple_col].apply(str)
        else:
            for t in tuple_col:
                df[t] = df[t].apply(str)
        if symbols is not None:
            with qdbing.Sender(self.host, 9009) as sender:
                sender.dataframe(df, table_name=table_name, symbols=symbols)
        else:
            with qdbing.Sender(self.host, 9009) as sender:
                sender.dataframe(df, table_name=table_name)
    
    @back_for_better
    def show_all_factors_information(self) -> pd.DataFrame:
        df=self.__get_data("select * from factors_info")
        if 'timestamp' in df.columns:
            df=df.drop(columns=['timestamp'])
        return df
    
    @concat_nosqls
    @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 date_ok(x):
            if isinstance(x,int):
                date=pd.Timestamp(str(x))
            elif isinstance(x,float):
                date=pd.Timestamp(str(int(x)))
            else:
                date=pd.Timestamp(x)
            date=datetime.datetime.strftime(date,'%Y-%m-%d')
            return date
        
        '''名称部分处理'''
        fac_key,fac_name=self._deal_key_name(fac_key)
        print(f"正在读取{fac_name}的数据")
        
        '''日期部分处理'''
        if trade_date is not None:
            trade_date=date_ok(trade_date)
            sql_order=f"select * from {fac_key} where date='{trade_date}'"
        elif (start_date is not None) and (end_date is not None):
            start_date=date_ok(start_date)
            end_date=date_ok(end_date)
            sql_order=f"select * from {fac_key} where date>='{start_date}' and date<='{end_date}'"
        elif start_date is not None:
            start_date=date_ok(start_date)
            sql_order=f"select * from {fac_key} where date>='{start_date}'"
        elif end_date is not None:
            end_date=date_ok(end_date)
            sql_order=f"select * from {fac_key} where date<='{end_date}'"
        else:
            sql_order=f"select * from {fac_key}"
        
        res=self.__get_data(sql_order)
        if 'timestamp' in res.columns:
            res=res.drop(columns=['timestamp'])
        if not sql_like:
            res=res.pivot(index='date',columns='code',values='fac')
        else:
            res=res.rename(columns={'fac':self.keys_names[fac_key]})

        return res
