from datetime import date
from dateutil.relativedelta import relativedelta

import polars as pl
import pandas as pd
from pyxirr import xirr


def cal_ytm(row, today: date):
    """ 计算可转债到期收益率YTM """
    try:
        trade_date = pd.to_datetime(today)
        len_years = len(row['interst_list'])
        # 计算日期列 （时间戳类型）
        dates_return = [
            pd.Timestamp(row['conv_end_date'] + relativedelta(years=- len_years + i + 1)) 
            for i in range(len_years)
        ]
        # 计算金额列
        value_return = row['interst_list'].copy()
        value_return[-1] = row['maturity_put_price']
        dates_return_pre = [i for i in dates_return if i <= trade_date]
        dates_return_cal = [trade_date] + dates_return[len(dates_return_pre):]
        value_return_cal = [-row['close']] + value_return[len(dates_return_pre):]
        ytm = xirr(dates_return_cal, value_return_cal)
        return ytm
    except Exception as e:
        return None


def plan_ytm(lf: pl.LazyFrame, today: date) -> pl.LazyFrame:
    """ 计算可转债到期收益率YTM """
    return (
        lf.with_columns(
            pl.struct(['conv_end_date', 'interst_list', 'maturity_put_price', 'close'])
            .map_elements(lambda row: cal_ytm(row, today), return_dtype=pl.Float64).alias('ytm'),
        )
    )
