from datetime import date

import polars as pl

from ldnode.loader.files import CB_HIST_DATA, STK_HIST_DATA


def get_hist_data(lf: pl.LazyFrame, latest_date: date) -> tuple[pl.DataFrame, pl.DataFrame]:
    hist_data = pl.read_parquet(CB_HIST_DATA.get_local_path())
    stk_hist_data = pl.read_parquet(STK_HIST_DATA.get_local_path())
    # 实时数据日期若比历史数据更新，则合并，否则直接用历史数据
    hist_data = hist_data.with_columns(pl.col('trade_date').dt.date().alias('trade_date'))
    stk_hist_data = stk_hist_data.with_columns(pl.col('trade_date').dt.date().alias('trade_date'))
    if latest_date > hist_data['trade_date'].max():
        hist_data = pl.concat([
            hist_data.drop('__index_level_0__'), 
            (
                lf.select([
                    'code', 'trade_date', 
                    'pre_close', 'open', 'high', 'low', 'close', 'vol', 
                    'amount', 'turnover',
                ])
            ).collect()
        ])
    else:
        hist_data = hist_data
    if latest_date > stk_hist_data['trade_date'].max():
        stk_hist_data = pl.concat([
            stk_hist_data.drop('__index_level_0__'), 
            (
                lf.select([
                    'code_stk', 'trade_date', 
                    'pre_close_stk', 'open_stk', 'high_stk', 'low_stk', 'close_stk', 
                    'vol_stk', 'amount_stk',
                ])
                .unique()
                .rename({
                    'code_stk': 'code', 
                    'pre_close_stk': 'pre_close', 'open_stk': 'open', 'high_stk': 'high', 'low_stk': 'low', 'close_stk': 'close', 
                    'vol_stk': 'vol', 'amount_stk': 'amount',
                })
            ).collect()
        ])
    else:
        stk_hist_data = stk_hist_data
    return hist_data, stk_hist_data


def plan_pct_chg(lf: pl.LazyFrame) -> pl.LazyFrame:
    return (
        lf.with_columns(
            (pl.col('close') / pl.col('pre_close') - 1).alias('pct_chg'),
            (pl.col('high') / pl.col('pre_close') - 1).alias('high_pct_chg'),
        )
    )
    
def gen_rolling_sum_expr(column_name: str, period: int) -> pl.Expr:
    return (
        pl.col(column_name)
        .rolling_sum(
            window_size=period,
            min_samples=1
        )
        .over("code")  # 按 code 分组计算
        .alias(f"{column_name}_{period}")
    )


def gen_rolling_pct_chg(period: int, suffix: str = '') -> pl.Expr:
    return (
        (pl.col("pct_chg") + 1)
        .log()  # 对数值进行求和避免浮点溢出
        .rolling_sum(
            window_size=period, 
            min_samples=1
        )
        .over("code")
        .exp()  # 指数还原
        .sub(1)  # 最终结果减1
        .alias(f"pct_chg_{period}{suffix}")
    )


def gen_rolling_mean(column_name: str, period: int) -> pl.Expr:
    return (
        pl.col(column_name)
        .rolling_mean(
            window_size=period,
            min_samples=1
        )
        .over("code")
        .alias(f"{column_name}_ma_{period}")
    )


def plan_moving_factors(lf: pl.LazyFrame, latest_date: date, period: int = 5) -> pl.LazyFrame:    
    rolling_mean_exprs = [gen_rolling_mean(column, period) for column in ['close']]
    rolling_sum_exprs = [gen_rolling_sum_expr(column, period) for column in ['vol', 'amount', 'turnover']]
    rolling_pct_chg_exprs = [gen_rolling_pct_chg(period)]
    
    return (
        lf
        .sort('trade_date', descending=True) # 按日期降序排列
        .group_by("code", maintain_order=True) # 按 code 分组，保持顺序
        .head(period)
        .sort(["code", "trade_date"])  # 重新按 code 和日期升序排列（为滚动窗口准备）
        .with_columns(
            *rolling_mean_exprs,
            *rolling_sum_exprs,
            *rolling_pct_chg_exprs,
        )
        .filter(pl.col("trade_date") == latest_date)
        .select(["code", "close_ma_5", "vol_5", "amount_5", "turnover_5", "pct_chg_5"])
    )


def plan_stk_moving_factors(lf: pl.LazyFrame, latest_date: date, period: int = 5) -> pl.LazyFrame:
    rolling_pct_chg_exprs = [gen_rolling_pct_chg(period, '_stk')]
    
    return (
        lf
        .sort('trade_date', descending=True) # 按日期降序排列
        .group_by("code", maintain_order=True) # 按 code 分组，保持顺序
        .head(period)
        .sort(["code", "trade_date"])  # 重新按 code 和日期升序排列（为滚动窗口准备）
        .with_columns(
            *rolling_pct_chg_exprs,
        )
        .filter(pl.col("trade_date") == latest_date)
        .select(["code", "pct_chg_5_stk"])
    )
    

def plan_calc_volatility(lf: pl.LazyFrame, latest_date: date, period: int = 242, suffix: str = '') -> pl.LazyFrame:
    """ 计算波动率计划 """
    return (
        lf
        .sort(["code", "trade_date"], descending=[False, True])
        .group_by("code", maintain_order=True)
        .agg(pl.all().head(period))  # 已按时间倒序，取头部即最新数据
        .explode(pl.all().exclude("code"))
        .sort(["code", "trade_date"])  # 滚动窗口需要时间升序
        .with_columns(
            (
                pl.col("pct_chg")
                .rolling_std(
                    window_size=period,
                    min_samples=1,
                    ddof=0  # 保持与 Pandas 的 std(ddof=0) 一致
                )
                .over("code") * pl.lit(period).sqrt()
            ).alias(f"volatility{suffix}")
        )
        .filter(pl.col("trade_date") == latest_date)
        .select(["code", f"volatility{suffix}"])
    )
    

def plan_pulse_factor(lf: pl.LazyFrame, latest_date: date, period: int = 242, high_pct_chg: float = 0.045) -> pl.LazyFrame:
    """ 计划脉冲因子 """
    return (
        lf
        .sort(["code", "trade_date"], descending=[False, True])
        .group_by("code")
        .head(period)
        .sort(["code", "trade_date"])
        .with_columns(
            # 生成布尔值并转为整型
            (pl.col("high_pct_chg") >= high_pct_chg).cast(pl.Int8).alias("condition"),
        )
        .with_columns(
            # 分组滚动求和
            pl.col("condition")
            .rolling_sum(window_size=242, min_samples=1)
            .over("code")
        )
        # 筛选最新日期并整理格式
        .filter(pl.col("trade_date") == latest_date)
        .select([
            "code",
            pl.col("condition").alias(f"pulse_{high_pct_chg}_{period}"),
        ])
    )


def plan_history(lf: pl.LazyFrame) -> pl.LazyFrame:
    """ 计划可转债历史数据 """
    latest_date: date = lf.select('trade_date').max().collect().item()
    
    hist_data, stk_hist_data = get_hist_data(lf, latest_date)
    cb_hist_quotation =  hist_data.lazy().pipe(plan_pct_chg)
    stk_hist_quotation = stk_hist_data.lazy().pipe(plan_pct_chg)
    return (
        lf
        .join(
            plan_moving_factors(cb_hist_quotation, latest_date), 
            on='code', 
            how='outer_coalesce',
        )
        .join(
            plan_stk_moving_factors(stk_hist_quotation, latest_date), 
            left_on='code_stk', 
            right_on='code', 
            how='outer_coalesce',
        )
        .with_columns(
            (pl.col('close') / pl.col('close_ma_5') - 1).alias('bias_5'),
            (pl.col('pct_chg_5') - pl.col('pct_chg_5_stk')).alias('alpha_pct_chg_5'),
        )
        .join(
            plan_calc_volatility(cb_hist_quotation, latest_date), 
            on='code', 
            how='outer_coalesce',
        )
        .join(
            plan_calc_volatility(stk_hist_quotation, latest_date, suffix='_stk'), 
            left_on='code_stk', 
            right_on='code', 
            how='outer_coalesce',
        )
        .join(
            plan_pulse_factor(cb_hist_quotation, latest_date), 
            on='code', 
            how='outer_coalesce',
        )
    )
