
import datetime
import traceback
from typing import Tuple, List, Optional

import numpy as np
import pandas as pd
import statsmodels.api as sm
from pandas.tseries.offsets import DateOffset

from .derived import DerivedFactor
from .utils import FundHoldStockFactor


class StockFactorApi:

    _DECOMP_RET_DATE_OFFSET = DateOffset(years=1)
    _WINDOW_LEN = 60

    @staticmethod
    def get_stock_factor_ret(factor_names: Tuple[str], universe: str) -> Optional[pd.DataFrame]:
        rets: List[pd.Series] = []
        for one in factor_names:
            try:
                factor = DerivedFactor._registry[one]
            except KeyError as e:
                print(f'[stock_factor_api] invalid stock factor name {one} (err_msg){e} (universe){universe}')
                continue
            f_ret = factor.get_ret(universe=universe)
            if f_ret is not None:
                rets.append(f_ret)
            else:
                print(f'[stock_factor_api] retrieve return of {one} failed (universe){universe}')
        if rets:
            return pd.concat(rets, axis=1)

    @staticmethod
    def decomp_ret(fund_ret: pd.Series, factor_names: Tuple[str], universe: str) -> Optional[Tuple[pd.Series, pd.DataFrame]]:

        def _lambda_rolling_reg(x, whole, fund_name):
            if pd.to_datetime(x.index.array[-1]).isoweekday() != 5:
                return 0
            # x是基金每日收益率, 对相应日期的因子收益做回归(加前置的常数项)(使用前WINDOW_LEN日)
            results = sm.OLS(x[:-1], sm.add_constant(whole.drop(columns=fund_name).loc[x.index[:-1], :])).fit()
            # 回归结果中包含常数项系数，这里只取另外几个indicator的系数(即上边的因子收益)
            # T+1日各因子收益的贡献 = T+1日因子收益 * 前WINDOW_LEN日回归系数
            contribution = whole.drop(columns=fund_name).loc[x.index[-1], :] * results.params[1:]
            # T+1日特异收益率 = T+1日基金收益 - T+1日各因子收益的贡献和
            spec_ret = x.array[-1] - sum(contribution)

            # 记录T+1日各因子收益的贡献, 特异收益率, 基金收益率等信息
            fund_attribution_list.append(contribution.append(pd.Series({'spec': spec_ret, 'datetime': x.index.array[-1]})))
            return 0

        ret_rg = StockFactorApi.get_stock_factor_ret(factor_names, universe)
        if ret_rg is None:
            print(f'[stock_factor_api] retrieve return of {factor_names} failed')
            return
        try:
            fund_ret = fund_ret.dropna()
            # 是否只使用部分数据计算因子暴露
            # cal_start_date = (fund_ret.index.array[-1] - StockFactorApi._DECOMP_RET_DATE_OFFSET).date()
            # if fund_ret.index.array[0] > cal_start_date:
            #     print(f'[stock_factor_api] not enough data of fund {fund_ret.name} (start_date){fund_ret.index.array[0]}')
            #     return
            # filtered_fund_ret = fund_ret[fund_ret.index >= cal_start_date]
            # if filtered_fund_ret.empty:
            #     print(f'[stock_factor_api] no data of fund {fund_ret.name} to decomp (start_date){fund_ret.index.array[0]} (end_date){fund_ret.index.array[-1]}')
            #     return
            # rg = filtered_fund_ret.to_frame().join(ret_rg).ffill().dropna()
            # 计算因子收益曲线(包含特异收益率)
            rg = fund_ret.to_frame().join(ret_rg).ffill().dropna()
            fund_attribution_list = []
            # 以下为手写的rolling, 但实际和直接用rolling的时间差不多
            # for i in range(rg.shape[0]):
            #     end_loc = i + StockFactorApi._WINDOW_LEN
            #     if end_loc + 1 >= rg.shape[0]:
            #         break
            #     data_in_window = rg.iloc[i:end_loc, :]
            #     # x是基金每日收益率, 对相应日期的因子收益做回归(加前置的常数项)(使用前WINDOW_LEN日)
            #     results = sm.OLS(data_in_window[fund_ret.name], sm.add_constant(data_in_window.drop(columns=fund_ret.name))).fit()
            #     # 回归结果中包含常数项系数，这里只取另外几个indicator的系数(即上边的因子收益)
            #     # T+1日各因子收益的贡献 = T+1日因子收益 * 前WINDOW_LEN日回归系数
            #     contribution = rg.iloc[end_loc, :].drop(fund_ret.name) * results.params[1:]
            #     # T+1日特异收益率 = T+1日基金收益 - T+1日各因子收益的贡献和
            #     spec_ret = rg.iloc[end_loc, :].at[fund_ret.name] - sum(contribution)

            #     # 记录T+1日各因子收益的贡献, 特异收益率, 基金收益率等信息
            #     fund_attribution_list.append(contribution.append(pd.Series({'spec': spec_ret, 'datetime': rg.index.array[end_loc]})))
            rg[fund_ret.name].rolling(window=StockFactorApi._WINDOW_LEN+1).apply(_lambda_rolling_reg, kwargs={'whole': rg, 'fund_name': fund_ret.name})

            # 计算基金全时长的因子暴露
            Y = rg.iloc[:, 0]
            X = sm.add_constant(rg.iloc[:, 1:])
            model1 = sm.OLS(Y, X)
            resu1 = model1.fit()
            return pd.Series(resu1.params[1:], index=ret_rg.columns, name=fund_ret.name), pd.DataFrame(fund_attribution_list).set_index('datetime').cumsum()
        except Exception as e:
            print(f'[stock_factor_api] abnormal error when decomp return (err_msg){e}')
            traceback.print_exc()
            return

    @staticmethod
    def fac_pos(fund_id: str, factor_names: Tuple[str], report_date: datetime.date) -> Optional[pd.DataFrame]:
        df = FundHoldStockFactor().get()
        if df is None:
            print(f'[fac_pos] failed to get data of fund hold stock')
            return
        df = df[(df.fund_id == fund_id) & (df.datetime == report_date)]
        df = df[df.notna().all(axis=1)]
        if df.empty:
            print(f'[fac_pos] empty df, can not calc (fund_id){fund_id} (report_date){report_date}')
            return

        s_list: List[pd.Series] = []
        try:
            for one in factor_names:
                try:
                    factor = DerivedFactor._registry[one]
                except KeyError as e:
                    print(f'[fac_pos] invalid stock factor name {one} (err_msg){e}')
                    continue
                f_val = factor.get()
                if f_val is None:
                    print(f'[fac_pos] retrieve value of {one} failed')
                    continue
                anjou = []
                weight = []
                for i in range(1, 11):
                    lis = df['rank'+str(i)+'_stock_code'].values
                    mult1 = pd.DataFrame(f_val[lis].iloc[-1])
                    mult1.columns = ['rank'+str(i)+'_fac']
                    mult2 = pd.DataFrame(df['rank'+str(i)+'_stockweight']).set_index(df['rank'+str(i)+'_stock_code'])
                    mult2.columns = ['rank'+str(i)+'_fac']
                    w = mult2.copy()
                    w.index = df.fund_id.to_list()
                    weight.append(w)
                    louis = (mult1 * mult2).replace({np.nan: 0})
                    louis.index = df.fund_id.to_list()
                    anjou.append(louis)
                res = (pd.concat(anjou, axis=1).sum(axis=1) / pd.concat(weight, axis=1).sum(axis=1)).rename(one)
                s_list.append(res)
        except Exception as e:
            print(f'[fac_pos] abnormal error when calc fac pos (err_msg){e}')
            traceback.print_exc()
            return
        res_df = pd.DataFrame(s_list)
        if res_df.empty:
            return
        return res_df[fund_id]
