from sangreal_wind.sangreal_calendar import Monthly, step_trade_dt
from sangreal_wind.utils.datetime_handle import dt_handle
from sangreal_wind.utils.return_base import (
    ReturnDataAPI, DailyReturnAPI, MonthlyReturnAPI,
    _daily_return_api, _monthly_return_api
)
from sangreal_wind.utils.config_manager import WIND_DB

# 月度类
MONTH = Monthly(-1)


def get_daily_ret(
    sid=None,
    trade_dt=None,
    begin_dt="20030101",
    end_dt="20990101",
    universe="A",
):
    """[get daily_ret of stocks,]

    Keyword Arguments:
        sid {[sid or iterable]} -- [stock windcode] (default: {None})
        begin_dt {str or datetime} -- [begin_dt] (default: {'20030101'})
        end_dt {str or datetime} -- [end_dt] (default: {'20990101'})
        trade_dt {[str or datetime]} -- [trade_dt] (default: {None})

    Returns:
        ret {pd.DataFrame} -- [sid: trade_dt]
    """
    if trade_dt is not None:
        begin_dt = end_dt = dt_handle(trade_dt)

    begin_dt, end_dt = dt_handle(begin_dt), dt_handle(end_dt)

    # 🎯 使用统一的API实例，消除代码重复
    df = _daily_return_api.get_return_data(sid, begin_dt, end_dt, universe)
    df.columns = ["sid", "trade_dt", "pct_change"]
    df = df.pivot(values="pct_change", index="trade_dt", columns="sid")

    # # 防止出现0的情况，强制缺失na
    df = df.pct_change(fill_method=None)
    df.dropna(how="all", inplace=True)
    return df.T


def get_monthly_ret(
    sid=None,
    trade_dt=None,
    begin_dt="20030101",
    end_dt="20990101",
):
    """[get monthly_ret of stocks,]

    Keyword Arguments:
        sid {[sid or iterable]} -- [stock windcode] (default: {None})
        begin_dt {str or datetime} -- [begin_dt] (default: {'20030101'})
        end_dt {str or datetime} -- [end_dt] (default: {'20990101'})
        trade_dt {[str or datetime]} -- [trade_dt] (default: {None})

    Returns:
        ret {pd.DataFrame} -- [sid: trade_dt]
    """
    begin_dt, end_dt = dt_handle(begin_dt), dt_handle(end_dt)

    # 🎯 使用统一的API实例，消除代码重复
    if trade_dt is not None:
        trade_dt = MONTH.prev(trade_dt)
        df = _monthly_return_api.get_return_data(sid, begin_dt, end_dt, trade_dt)
        # 单个日期查询返回原始格式
        df.columns = ["sid", "trade_dt", "close"]
        df.close = df.close / 100.0
        return df
    else:
        # 日期范围查询返回透视格式
        return _monthly_return_api.get_return_data(sid, begin_dt, end_dt)


if __name__ == "__main__":
    # df = get_daily_ret(begin_dt='20181101')
    # print(df.head())
    df = get_daily_ret(begin_dt="20180101", end_dt="20181223")
    print(df)
