
import datetime
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_factors 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):
            # 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}')
            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}')
            return
        res_df = pd.DataFrame(s_list)
        if res_df.empty:
            return
        return res_df[fund_id]
