from datetime import datetime

import polars as pl
from optionprice import Option


def cal_option_value(row):
    """ 计算期权价值 """
    try:
        option = Option(
            european=False,
            kind='call',
            s0=row['close_stk'],  # 现价
            k=row['conv_price'],  # 转股价格
            sigma=row['volatility_stk'],  # 年化波动率
            r=row['free_risk'] / 100,  # 无无风险利率
            # dv=row.dv_ratio # 股息率
            start=datetime.now().strftime("%Y-%m-%d"),
            end=row['conv_end_date'].strftime("%Y-%m-%d")
        )
        option_value = 100 / row['conv_price'] * option.getPrice()
        return option_value
    except Exception:
        return None


def plan_option(lf: pl.LazyFrame) -> pl.LazyFrame:
    """ 计划期权数据 """
    return (
        lf
        .with_columns(
            pl.struct(['close_stk', 'conv_price', 'volatility_stk', 'free_risk', 'conv_end_date'])
            .map_elements(cal_option_value, return_dtype=pl.Float64)
            .alias('option_value'),
        )
        .with_columns(
            (pl.col('pure_value') + pl.col('option_value')).alias('theory_value'),
        )
        .with_columns(
            (-((pl.col('theory_value') / pl.col('close')) - 1)).alias('theory_bias'),
        )
    )
