from scipy import optimize
from utils import get_monthly_sip_nav_df


def calculate_xnpv(rate, cashflows):
    chron_order = sorted(cashflows, key=lambda x: x[0])
    t0 = chron_order[0][0]
    return sum([cf/(1+rate)**((t-t0).days/365.0) for (t,cf) in chron_order])

def calculate_xirr(cashflows, guess=0.1):
    return optimize.newton(lambda r: calculate_xnpv(r,cashflows), guess)

def get_investment_xirr(investment_df, start_date, end_date, SIP_date, lumpsum_amount, sip_amount):
    # Get the monthly NAVs
    if(sip_amount == 0):
        sip_amount = 1000

    monthly_nav_df = get_monthly_sip_nav_df(investment_df, start_date, end_date, SIP_date)
    
    # Initialize lists to store cash flows and their corresponding dates
    cash_flows = []
    dates = []
    
    # Add the lumpsum investment at the start
    cash_flows.append(-lumpsum_amount)
    dates.append(start_date)
    
    # Calculate initial units from lumpsum investment
    initial_units = lumpsum_amount / monthly_nav_df['nav'].iloc[0]
    total_units = initial_units
    
    # Iterate over each row and record the SIP investments
    for _, row in monthly_nav_df.iterrows():
        cash_flows.append(-sip_amount)  # SIP investment is negative cash flow
        dates.append(row['date'])
        total_units += sip_amount / row['nav']
    
    # Add the final value as a positive cash flow
    final_value = total_units * monthly_nav_df['nav'].iloc[-1]
    cash_flows.append(final_value)
    dates.append(monthly_nav_df['date'].iloc[-1])

    portfolio_XIRR = calculate_xirr(list(zip(dates, cash_flows)))
    
    return portfolio_XIRR * 100

def get_investment_sip_absolute_returns(investment_df, sip_amount, lumpsum_amount, stepup, start_date, end_date, SIP_Date):
    # start_date = pd.Timestamp(start_date)
    # end_date = pd.Timestamp(end_date)

    scheme_df_monthly = get_monthly_sip_nav_df(investment_df, start_date, end_date, SIP_Date)

    total_investment = lumpsum_amount
    current_sip_amount = sip_amount

    # do calculation for upfront investment
    units_bought = lumpsum_amount / scheme_df_monthly.iloc[0]['nav']
    units_accumulated = units_bought
    previous_year = start_date.year

    for _, row in scheme_df_monthly.iloc[:-1].iterrows():
        # Check if a year has passed and increase SIP amount accordingly
        if row['date'].year > previous_year:
            current_sip_amount += current_sip_amount * (stepup / 100)
            previous_year = row['date'].year

        units_bought = current_sip_amount / row['nav']
        units_accumulated += units_bought
        total_investment += current_sip_amount

    final_value = units_accumulated * scheme_df_monthly.iloc[-1]['nav']
    total_return = (final_value - total_investment) / total_investment * 100

    return total_return, final_value, total_investment