from sangreal_wind.sangreal_calendar import Monthly, step_trade_dt
from sangreal_wind.utils.datetime_handle import dt_handle
from sangreal_wind.utils.engines 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 universe not in ('A', 'HK'):
        raise ValueError('universe only support A or HK!')
    begin_dt, end_dt = dt_handle(begin_dt), dt_handle(end_dt)
    if universe == 'A':
        table = getattr(WIND_DB, 'AShareEODPrices'.upper())
    else:
        table = getattr(WIND_DB, 'HKSHAREEODPRICES'.upper())
    query = WIND_DB.query(table.S_INFO_WINDCODE, table.TRADE_DT,
                          table.S_DQ_ADJCLOSE)
    if sid is not None:
        if isinstance(sid, str):
            query = query.filter(table.S_INFO_WINDCODE == sid)
        else:
            query = query.filter(table.S_INFO_WINDCODE.in_(sid))

    if trade_dt is not None:
        begin_dt = end_dt = dt_handle(trade_dt)
    df = query.filter(
        table.TRADE_DT >= step_trade_dt(begin_dt, -1), table.TRADE_DT <= end_dt).order_by(
            table.TRADE_DT).to_df()
    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)
    table = getattr(WIND_DB, 'ASHAREMONTHLYYIELD'.upper())
    query = WIND_DB.query(table.S_INFO_WINDCODE, table.TRADE_DT,
                          table.S_MQ_PCTCHANGE)
    if sid is not None:
        if isinstance(sid, str):
            query = query.filter(table.S_INFO_WINDCODE == sid)
        else:
            query = query.filter(table.S_INFO_WINDCODE.in_(sid))

    if trade_dt is not None:
        trade_dt = MONTH.prev(trade_dt)
        df = query.filter(table.TRADE_DT == trade_dt).order_by(
            table.TRADE_DT).to_df()
    else:
        df = query.filter(
            table.TRADE_DT >= begin_dt, table.TRADE_DT <= end_dt).order_by(
                table.TRADE_DT).to_df()
    df.columns = ['sid', 'trade_dt', 'close']
    df.close = df.close / 100.0
    df = df.pivot(values='close', index='trade_dt', columns='sid')
    df.dropna(how='all', inplace=True)
    return df.T


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)
