from datetime import date
from dateutil.relativedelta import relativedelta

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

from .realtime_derived_plan import calculate_cb_future_cashflows


def cal_ytm(row, today: date):
    """ 计算可转债到期收益率YTM """
    try:
        cashflows = calculate_cb_future_cashflows(row, today)
        if cashflows:
            dates_return_cal = [pd.Timestamp(today)] + [
                pd.Timestamp(cashflow_date) for cashflow_date, _ in cashflows
            ]
            value_return_cal = [-row['close']] + [cash for _, cash in cashflows]
        else:
            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:
        return None


def plan_ytm(lf: pl.LazyFrame, today: date) -> pl.LazyFrame:
    """ 计算可转债到期收益率YTM """
    if "pay_interest_day_list" not in set(lf.collect_schema().names()):
        lf = lf.with_columns(
            pl.lit(None, dtype=pl.List(pl.Date)).alias("pay_interest_day_list")
        )
    return (
        lf.with_columns(
            pl.struct([
                'conv_end_date',
                'pay_interest_day_list',
                'interst_list',
                'maturity_put_price',
                'close',
                'call_event_status_window',
                'call_event_reg_date_window',
                'call_event_price_window',
            ])
            .map_elements(lambda row: cal_ytm(row, today), return_dtype=pl.Float64).alias('ytm'),
        )
    )
