import polars as pl
from polars import col


def plan_correct(lf: pl.LazyFrame) -> pl.LazyFrame:
    """ 数据纠正计划 """
    return (
        lf.with_columns(
            col('code').str.slice(2) + '.' + col('code').str.slice(0, 2).str.to_uppercase(),
            col('trade_date').str.to_date(),
            col('time').str.to_time(),
            pl.when((col('bid1')==0) | (col('ask1')==0)) # 修正涨跌停
            .then(col('bid1') + col('ask1'))
            .otherwise((col('bid1') + col('ask1')) / 2)
            .alias('close'),
        )
        .with_columns(
            # 修正无挂单close
            pl.when(col('close')==0)
            .then(
                pl.when(col('now')==0).then(col('pre_close'))# 开盘无挂单修正
                .otherwise(col('now'))
            ).otherwise(col('close')).alias('close'),
        )
    )