from functools import lru_cache
from sqlalchemy import func

from sangreal_wind.utils import dt_handle
from sangreal_wind.utils.industry_base import (
    IndustryAPI, ZXIndustryAPI, WindHKIndustryAPI, HKCITICSIndustryAPI,
    _zx_industry_api, _wind_hk_industry_api, _hk_citics_industry_api
)


class DynamicIndustry:

    def __init__(self, ind=None):
        self.ind = ind

    def preview(self, trade_dt, adjust=False):
        # adjust {bool} -- [由于中信变更行业分类，是否调整兼容之前的代码] (default: {True})
        all_stk = get_industry(trade_dt=trade_dt, level=1, sid=None, adjust=adjust)
        if self.ind is not None:
            return set(all_stk[all_stk["ind"] == self.ind].index)
        return set(all_stk.index)


def get_industry(trade_dt, sid=None, level=1, adjust=False):
    """[get industry of stock 中信行业]

    Arguments:
        trade_dt {[str or datetime]} -- [trade_dt]

    Keyword Arguments:
        sid {[str or iterable]} -- [sids of stocks] (default: {None})
        level {int} -- [level of zx industry] (default: {1})
        adjust {bool} -- [由于中信变更行业分类，是否调整兼容之前的代码] (default: {True})

    Returns:
        [pd.DataFrame] -- [sid: ind]
    """
    trade_dt = dt_handle(trade_dt)
    return _zx_industry_api.get_industry_at_date(trade_dt, sid, level, adjust)


@lru_cache()
def get_industry_all(level=1, adjust=False):
    """[summary]
    adjust {bool} -- [由于中信变更行业分类，是否调整兼容之前的代码] (default: {True})
    """
    return _zx_industry_api.get_all_industries_data(level, adjust)

@lru_cache()
def get_industry_all_windhk(level=1):
    """[summary]
    adjust {bool} -- [由于中信变更行业分类，是否调整兼容之前的代码] (default: {True})
    """
    return _wind_hk_industry_api.get_all_industries_data(level)


def get_industry_windhk(trade_dt, sid=None, level=1):
    """[get industry of stock 中信行业]

    Arguments:
        trade_dt {[str or datetime]} -- [trade_dt]

    Keyword Arguments:
        sid {[str or iterable]} -- [sids of stocks] (default: {None})
        level {int} -- [level of zx industry] (default: {1})
        adjust {bool} -- [由于中信变更行业分类，是否调整兼容之前的代码] (default: {True})

    Returns:
        [pd.DataFrame] -- [sid: ind]
    """
    trade_dt = dt_handle(trade_dt)
    df = _wind_hk_industry_api.get_all_industries_data(level)
    df = _wind_hk_industry_api.filter_by_stocks(df, sid)
    df = _wind_hk_industry_api.filter_by_trade_date(df, trade_dt)
    return df.set_index("sid")[["ind"]].copy()

@lru_cache()
def get_industry_all_hk():
    """[summary]"""
    return _hk_citics_industry_api.get_all_industries_data()


def get_industry_hk(trade_dt, sid=None):
    """[get industry of stock 中信行业]"""
    return _hk_citics_industry_api.get_industry_at_date(trade_dt, sid)


def get_industry_sp(
    trade_dt, sid=None, split=["银行", "非银行金融", "综合金融"], adjust=False
):
    """[将split中部分中信一级行业转换为相应的二级行业]

    Arguments:
        trade_dt {[str]} -- [description]

    Keyword Arguments:
        sid      {[str or iterable]} -- [sids of stocks] (default: {None})
        split      {list} -- [industry which convert level1 to level2] (default: {['银行', '非银行金融']})
        adjust {bool} -- [由于中信变更行业分类，是否调整兼容之前的代码] (default: {True})

    Returns:
        [pd.DataFrame] -- [sid: ind]
    """
    trade_dt = dt_handle(trade_dt)
    df = get_industry_all(level=1, adjust=adjust)
    if sid is not None:
        sid = {sid} if isinstance(sid, str) else set(sid)
        df = df[df["sid"].isin(sid)]

    df = df.loc[
        (df["entry_dt"] <= trade_dt)
        & ((df["out_dt"] >= trade_dt) | (df["out_dt"].isnull()))
    ].copy()

    split_sid = df[df["ind"].isin(split)]
    normal_sid = df[~(df["ind"].isin(split))]
    df1 = get_industry_all(level=2, adjust=adjust)
    df1 = df1[df1["sid"].isin(split_sid.sid)]

    df1 = df1.loc[
        (df1["entry_dt"] <= trade_dt)
        & ((df1["out_dt"] >= trade_dt) | (df1["out_dt"].isnull()))
    ].copy()
    # 将一级和二级合并
    df = normal_sid.append(df1, ignore_index=True)

    return df.set_index("sid")[["ind"]]


if __name__ == "__main__":
    print(len(DynamicIndustry().preview("20180101")))
