module fundit::rbsaCalculator

use fundit::sqlUtilities;
use fundit::performanceDataPuller;
use fundit::operationDataPuller;

/*
 *   RBSA 计算
 *   @param: ret <TABLE>: historical return (double) vector which contains the same number of return as index
 *           index_ret <TABLE>: historical index return table  which each row is an index
 *           is_long <BOOL>: true - long-only, false - long-short
 *   @return: table
 * 
 *   Create  20240703  模仿python代码在Dolphin中实现，具体计算逻辑完全不懂                          Joey
 *                     原代码见： http://gogs.fundit.cn/FundIt/FinanceCalcPython/src/dev36/pf_scical/v1/calc_rbsa_use_osqp.py
 *                     Python官方示例见：https://osqp.org/docs/examples/least-squares.html
 *                     Dolphin官方示例见：https://docs.dolphindb.cn/zh/funcs/o/osqp.html
 *
 */
defg cal_rbsa(ret, index_ret, is_long) {

    // 窗口长度
    m = ret.size()
    // 指数个数
    n = index_ret.cols()
    
    P0 = matrix(float, n, m+n)
    P1 = concatMatrix([matrix(float, m, n), eye(m)])
    P = concatMatrix([P0, P1], false)
    q = array(float, m+n, (m+n)*10, 0)
    
    A0 = concatMatrix( [matrix(index_ret), -eye(m)])
    A1 = concatMatrix( [matrix(take(1, n)).transpose(), matrix(float, 1, m)])
    A2 = concatMatrix( [eye(n), matrix(float, n, m)])
    A = concatMatrix( [A0, A1, A2], false)
    
    // join 1 是为了限制所有权重加总为100%
    // 下限
    lb =(ret join 1) join array(float, n, n*10, iif(is_long == true, 0, -2))
    // 上限
    ub=(ret join 1) join array(float, n, n*10, iif(is_long == true, 1, 2))
    
    res = osqp( q, P, A, lb, ub)
    
    return res
 }


/*
 * 滚动 rbsa
 * @param ret <TABLE>: return table, at least with "effective_date" and "ret" as columns
 * @param index_ret <TABLE>: index return table, with "effective_date" and all index ids as columns
 * @param is_long <BOOL>: boolean. true means weightings could be negative values
 * @param window <INT>: number of return in a window
 * @param step <INT>: rolling step
 * 
 * TODO: use rolling()
 * 
 * @return <TABLE> with "effective_date", "index_id" and "weights" columns
 */
def cal_rolling_rbsa(ret, index_ret, is_long, window, step) {

	// 找到所有指数全有数据的最早日期
    v_start_date = EXEC effective_date.max() AS start_date 
                   FROM (SELECT entity_id, effective_date.min() AS effective_date FROM index_ret WHERE ret IS NOT NULL GROUP BY entity_id);

	m_index_ret = SELECT ret FROM index_ret WHERE effective_date >= v_start_date PIVOT BY effective_date, entity_id;
	
    t = SELECT * FROM ej(ret, m_index_ret, 'effective_date') ORDER BY ret.effective_date;
    t.nullFill!(0)

    // not sure why this doesn't work
    // rolling(cal_rbsa{,,is_long}, (t.ret, t.slice(, ret.cols():).matrix()), window, step)

    // 指数个数
    n = m_index_ret.cols() - 1

    // 计算起始位置
    i = (t.size() - window) % step

    // 运行rbsa计算次数
    cnt = (t.size() - i - window) / step + 1;
    
    tb = table(max(cnt,1):0, ["effective_date", "price_date", "index_id", "weights", "alpha", "r2", "adj_r2"], [STRING, DATE, STRING, DOUBLE, DOUBLE, DOUBLE, DOUBLE]);

    if(t.size() >= max(window, step) && cnt > 0) {
    	
        do {

			alpha = 0;
        	r2 = 0;
        	adj_r2 = 0;

            v_ret = t.ret[i:(i+window)];
            
            t_index_ret = t.slice( i:(i+window), ret.cols(): );
        
            // 传入window个收益
            res = cal_rbsa(v_ret, t_index_ret, is_long);

            if(res[0] == 'solved') {

				m_predict_ret = t_index_ret.matrix() ** res[1][0:n];
				
				alpha = v_ret.mean() - m_predict_ret.mean();
				SSR = sum2(m_predict_ret - v_ret.mean());
				SST = sum2(v_ret - v_ret.mean());
				if(SST == 0) {
					// 当SST=0, 先计算SSE再计算SST
					SSE = sum2(v_ret - m_predict_ret);
					SST = SSE + SSR;
				}
				
				if(SST != 0) {
					r2 = SSR/SST;
					adj_r2 = 1 - (1 - r2) * (window - 1) / (window - n - 1);
				}
				
				for(j in 1..n) {
	    	        tb.tableInsert(t.effective_date[i+window-1], t.price_date[i+window-1], m_index_ret.colNames()[j], res[1][j-1].round(4), alpha, r2, adj_r2);
				}
            }

            // 往前推进step个收益
            i = i + step
            
            cnt -= 1
    
        } while( cnt > 0)
        
    }

    return tb

}

/*
 *  计算单基金或组合的RBSA
 * 
 *  @param entity_type <STRING>: 目标基金/组合的类型
 *  @param entity_id <STRING>: 目标基金/组合的ID
 *  @param index_ids <VECTOR>: 基准指数IDs
 *  @param freq <STRING>: m, w, d
 *  @param start_day <DATE>
 *  @param end_day <DATE>
 *  @param is_long <BOOL>: 是否只考虑纯多头
 *  @param window <INT>: 窗口（必须多于基准指数个数）
 *  @param step <INT>: 步长
 *  
 *  @return <TABLE>: entity_id, effective_date, price_date, index_id, weights, alternative_id, level, alpha, r2, adj_r2
 *  
 *  TODO: 数字与界面和数据库都对不上
 *  
 *  Example: cal_single_entity_RBSA('MF', 'MF00003PW1', ['IN00000008', 'IN00000077', 'IN0000007G', 'IN0000009M'], 'w', 1900.01.01, 2024.11.15, true, 24, 24);
 *           cal_single_entity_RBSA('PF', 166002, ['FA00000VML', 'FA00000VMM', 'FA00000VMN', 'FA00000VMO', 'IN0000007G'], 'w', 2020.01.01, 2024.11.08, true, 24, 24);
 *           cal_single_entity_RBSA('MF', 'MF000200KQ', ['IN00000008', 'IN00000077', 'IN0000007G', 'IN0000009M'], 'w', 1900.01.01, 2024.11.16, true, 24, 24);
 */
def cal_single_entity_RBSA(entity_type, entity_id, index_ids, freq='w', start_day=1900.01.01, end_day=2099.12.31, is_long=true, window=24, step=24) {
// entity_type='MF'
// entity_id= 'MF00003PW1'
// index_ids=['IN00000008', 'IN00000077', 'IN0000007G', 'IN0000009M']
// freq='w'
// start_day=2001.01.19
// end_day=2024.11.16
// is_long=true
// window=48
// step=13

    tb_result = table(100:0, ["entity_id", "effective_date", "index_id", "weights", "alpha", "r2", "adj_r2"], 
                             [iif(entity_type=='PF', INT, STRING), STRING, STRING, DOUBLE, DOUBLE, DOUBLE, DOUBLE]);

    v_entity = array(iif(entity_type=='PF', INT, STRING));
    v_entity.append!(entity_id);
    
	entity_ret = get_entity_return(entity_type, v_entity, freq, start_day, end_day, true);

	// 数据长度不够，按照顺序依次分别用母基金(4), 指数(3)的数据来代替 (level=2, 基金经理的数据在Java里好像没用？)
	level = 1
	alternative_id = NULL;
	if(entity_ret.isVoid() || entity_ret.size() < window) {
		if(entity_type IN ['MF', 'HF']) {

			fund_info = get_fund_info(v_entity);
			p_fund_id = fund_info.p_fund_id[0];
			primary_benchmark_id = fund_info.benchmark_id[0];
			if(p_fund_id != NULL) {
				entity_ret = get_entity_return(entity_type, v_entity.replace(entity_id, p_fund_id) , freq, start_day, end_day, true);
				alternative_id = p_fund_id;
				level = 4;
			} else if(primary_benchmark_id != NULL) {
				entity_ret = get_entity_return(entity_type, v_entity.replace(entity_id, primary_benchmark_id) , freq, start_day, end_day, true);
				alternative_id = primary_benchmark_id;
				level = 3;
			} else {
				return tb_result;
			}
		} else if(entity_type == 'PF'){

			portfolio_info = get_portfolio_info(v_entity);
			primary_benchmark_id = portfolio_info.benchmark_id[0];
			if(primary_benchmark_id != NULL) {
				entity_ret = get_entity_return(entity_type, v_entity.replace(entity_id, primary_benchmark_id) , freq, start_day, end_day, true);
				alternative_id  = primary_benchmark_id;
				level = 3;
			} else
				return tb_result;
		}
	}

    // 因为用来做基准指数的可能是指数、因子、基金等等任何时间序列数据，所以不用填 entity_type
    index_ret = get_entity_return(NULL, index_ids, freq, start_day, end_day, true);

	if(index_ret.isVoid() || index_ret.size() == 0) return tb_result;

	tb_result = SELECT entity_id, effective_date, price_date, index_id, weights, alternative_id, level, alpha, r2, adj_r2
	            FROM cal_rolling_rbsa(entity_ret, index_ret, is_long, window, step);

	return tb_result;
	
}


/*
 *  计算单基金或组合的RBSA
 * 
 *  @param entity_type <STRING>: MF, HF, PF
 *  @param entity_ret <TABLE>: [COLUMNS] entity_id, price_date, ret
 *  @param index_ret <TABLE>: [COLUMNS] entity_id, price_date, ret
 *  @param freq <STRING>: m, w, d
 *  @param start_day <DATE>
 *  @param end_day <DATE>
 *  @param is_long <BOOL>: 是否只考虑纯多头
 *  @param window <INT>: 窗口（必须多于基准指数个数）
 *  @param step <INT>: 步长
 *  
 *  @return <TABLE>: entity_id, effective_date, price_date, index_id, weights, alternative_id, level, alpha, r2, adj_r2
 *  
 *  TODO: 数字与界面和数据库都对不上
 *  
 */
def cal_entity_RBSA(entity_type, entity_ret, index_ret, freq='w', start_day=1900.01.01, end_day=2099.12.31, is_long=true, window=24, step=24) {
// entity_type='MF'
// freq='w'
// start_day=2024.01.05
// end_day=today()
// is_long=true
// window=48
// step=13

    tb_result = table(100:0, ["entity_id", "effective_date", "index_id", "weights", "alpha", "r2", "adj_r2"], 
                             [iif(entity_type=='PF', INT, STRING), STRING, STRING, DOUBLE, DOUBLE, DOUBLE, DOUBLE]);

	if(index_ret.isVoid() || index_ret.size() == 0) return tb_result;

	tb_entity_ret = entity_ret;

	// 数据长度不够，按照顺序依次分别用母基金(4), 指数(3)的数据来代替 (level=2, 基金经理的数据在Java里好像没用？)
	level = 1
	alternative_id = NULL;
	if(tb_entity_ret.isVoid() || tb_entity_ret.size() < window) {
		if(entity_type IN ['MF', 'HF']) {

			fund_info = get_fund_info(tb_entity_ret.entity_id);
			p_fund_id = fund_info.p_fund_id;
			primary_benchmark_id = fund_info.benchmark_id;
			if(p_fund_id != NULL) {
				tb_entity_ret = SELECT entity_id, price_date, ret FROM get_entity_return(entity_type, p_fund_id , freq, start_day, end_day, true);
				alternative_id = p_fund_id[0];
				level = 4;
			} else if(primary_benchmark_id != NULL) {
				tb_entity_ret = SELECT entity_id, price_date, ret FROM get_entity_return(entity_type, primary_benchmark_id, freq, start_day, end_day, true);
				alternative_id = primary_benchmark_id[0];
				level = 3;
			} else {
				return tb_result;
			}
		} 
	}

	tb_entity_ret.addColumn('effective_date', STRING);
	tb_index_ret = index_ret;
	tb_index_ret.addColumn('effective_date', STRING);

	if(freq == 'm') {

		UPDATE tb_entity_ret SET effective_date = price_date.temporalFormat('yyyy-MM');
		UPDATE tb_index_ret SET effective_date = price_date.temporalFormat('yyyy-MM');
		
	} else if(freq == 'w') {

		UPDATE tb_entity_ret SET effective_date = get_year_week(price_date);
		UPDATE tb_index_ret SET effective_date = get_year_week(price_date);

	} else {

		UPDATE tb_entity_ret SET effective_date = price_date$STRING;
		UPDATE tb_index_ret SET effective_date = price_date$STRING;
	}

	tb_result = SELECT entity_ret.entity_id[0] AS entity_id, effective_date, price_date, index_id, weights, alternative_id, level, alpha, r2, adj_r2
	            FROM cal_rolling_rbsa(tb_entity_ret, tb_index_ret, is_long, window, step);

	return tb_result;
	
}