"""

Owl/tax2026
---

A retirement planner using linear programming optimization.

See companion document for a complete explanation and description
of all variables and parameters.

Module to handle all tax calculations.

Copyright &copy; 2026 - Martin-D. Lacasse

Disclaimers: This code is for educational purposes only and does not constitute financial advice.

"""

import numpy as np
from datetime import date


##############################################################################
# Prepare the data.

taxBracketNames = ["10%", "12/15%", "22/25%", "24/28%", "32/33%", "35%", "37/40%"]

rates_OBBBA = np.array([0.10, 0.12, 0.22, 0.24, 0.32, 0.35, 0.370])
rates_preTCJA = np.array([0.10, 0.15, 0.25, 0.28, 0.33, 0.35, 0.396])

###############################################################################
# Start of section where rates need to be actualized every year.
###############################################################################
# Single [0] and married filing jointly [1].

# These are 2025 current.
taxBrackets_OBBBA = np.array(
    [
        [12400, 50400, 105700, 201775, 256225, 640600, 9999999],
        [24800, 100800, 211400, 403550, 512450, 768700, 9999999],
    ]
)

irmaaBrackets = np.array(
    [
        [0, 106000, 133000, 167000, 200000, 500000],
        [0, 212000, 266000, 334000, 400000, 750000],
    ]
)

# Index [0] stores the standard Medicare part B premium.
# Following values are incremental IRMAA part B monthly fees.
irmaaFees = 12 * np.array([185.00, 74.00, 111.00, 110.90, 111.00, 37.00])

# Make projection for pre-TCJA using 2017 to current year.
# taxBrackets_2017 = np.array(
#    [ [9325, 37950, 91900, 191650, 416700, 418400, 9999999],
#      [18650, 75900, 153100, 233350, 416700, 470700, 9999999],
#    ])
#
# stdDeduction_2017 = [6350, 12700]
#
# For 2025, I used a 30.5% adjustment from 2017, rounded to closest 50.
#
# These are speculated.
taxBrackets_preTCJA = np.array(
    [
        [12150, 49550, 119950, 250200, 544000, 546200, 9999999],  # Single
        [24350, 99100, 199850, 304600, 543950, 614450, 9999999],  # MFJ
    ]
)

# These are 2025 current.
stdDeduction_OBBBA = np.array([16100, 32200])    # Single, MFJ
# These are speculated (adjusted for inflation to 2026). TODO
stdDeduction_preTCJA = np.array([8300, 16600])  # Single, MFJ

# These are current (adjusted for inflation) per individual.
extra65Deduction = np.array([2000, 1600])       # Single, MFJ

# Thresholds setting capital gains brackets 0%, 15%, 20% (adjusted for inflation).
capGainRates = np.array(
    [
        [48350, 533400],
        [96700, 600050],
    ]
)

# Thresholds for net investment income tax (not adjusted for inflation).
niitThreshold = np.array([200000, 250000])
niitRate = 0.038

# Thresholds for 65+ bonus for circumventing tax on social security.
bonusThreshold = np.array([75000, 150000])

###############################################################################
# End of section where rates need to be actualized every year.
###############################################################################


def mediVals(yobs, horizons, gamma_n, Nn, Nq):
    """
    Return tuple (nm, L, C) of year index when Medicare starts and vectors L, and C
    defining end points of constant piecewise linear functions representing IRMAA fees.
    """
    thisyear = date.today().year
    assert Nq == len(irmaaFees), f"Inconsistent value of Nq: {Nq}."
    assert Nq == len(irmaaBrackets[0]), "Inconsistent IRMAA brackets array."
    Ni = len(yobs)
    # What index year will Medicare start? 65 - age.
    nm = 65 - (thisyear - yobs)
    nm = np.min(nm)
    # Has it already started?
    nm = max(0, nm)
    Nmed = Nn - nm

    L = np.zeros((Nmed, Nq-1))
    C = np.zeros((Nmed, Nq))

    # Year starts at offset nm in the plan.
    for nn in range(Nmed):
        imed = 0
        n = nm + nn
        if thisyear + n - yobs[0] >= 65 and n < horizons[0]:
            imed += 1
        if Ni == 2 and thisyear + n - yobs[1] >= 65 and n < horizons[1]:
            imed += 1
        if imed:
            status = 0 if Ni == 1 else 1 if n < horizons[0] and n < horizons[1] else 0
            L[nn] = gamma_n[n] * irmaaBrackets[status][1:]
            C[nn] = imed * gamma_n[n] * irmaaFees
        else:
            raise RuntimeError("mediVals: This should never happen.")

    return nm, L, C


def capitalGainTaxRate(Ni, magi_n, gamma_n, nd, Nn):
    """
    Return an array of decimal rates for capital gains.
    Parameter nd is the index year of first passing of a spouse, if applicable,
    nd == Nn for single individuals.
    """
    status = Ni - 1
    cgRate_n = np.zeros(Nn)

    for n in range(Nn):
        if status and n == nd:
            status -= 1

        if magi_n[n] > gamma_n[n] * capGainRates[status][1]:
            cgRate_n[n] = 0.20
        elif magi_n[n] > gamma_n[n] * capGainRates[status][0]:
            cgRate_n[n] = 0.15

    return cgRate_n


def mediCosts(yobs, horizons, magi, prevmagi, gamma_n, Nn):
    """
    Compute Medicare costs directly.
    """
    thisyear = date.today().year
    Ni = len(yobs)
    costs = np.zeros(Nn)
    for n in range(Nn):
        status = 0 if Ni == 1 else 1 if n < horizons[0] and n < horizons[1] else 0
        for i in range(Ni):
            if thisyear + n - yobs[i] >= 65 and n < horizons[i]:
                # Start with the (inflation-adjusted) basic Medicare part B premium.
                costs[n] += gamma_n[n] * irmaaFees[0]
                if n < 2:
                    mymagi = prevmagi[n]
                else:
                    mymagi = magi[n - 2]
                for q in range(1, 6):
                    if mymagi > gamma_n[n] * irmaaBrackets[status][q]:
                        costs[n] += gamma_n[n] * irmaaFees[q]

    return costs


def taxParams(yobs, i_d, n_d, N_n, gamma_n, MAGI_n, yOBBBA=2099):
    """
    Input is year of birth, index of shortest-lived individual,
    lifespan of shortest-lived individual, total number of years
    in the plan, and the year that preTCJA rates might come back.

    It returns 3 time series:
    1) Standard deductions at year n (sigma_n).
    2) Tax rate in year n (theta_tn)
    3) Delta from top to bottom of tax brackets (Delta_tn)
    This is pure speculation on future values.
    Returned values are not indexed for inflation.
    """
    # Compute the deltas in-place between brackets, starting from the end.
    deltaBrackets_OBBBA = np.array(taxBrackets_OBBBA)
    deltaBrackets_preTCJA = np.array(taxBrackets_preTCJA)
    for t in range(6, 0, -1):
        for i in range(2):
            deltaBrackets_OBBBA[i, t] -= deltaBrackets_OBBBA[i, t - 1]
            deltaBrackets_preTCJA[i, t] -= deltaBrackets_preTCJA[i, t - 1]

    # Prepare the 3 arrays to return - use transpose for easy slicing.
    sigmaBar = np.zeros((N_n))
    Delta = np.zeros((N_n, 7))
    theta = np.zeros((N_n, 7))

    filingStatus = len(yobs) - 1
    souls = list(range(len(yobs)))
    thisyear = date.today().year

    for n in range(N_n):
        # First check if shortest-lived individual is still with us.
        if n == n_d:
            souls.remove(i_d)
            filingStatus -= 1

        if thisyear + n < yOBBBA:
            sigmaBar[n] = stdDeduction_OBBBA[filingStatus] * gamma_n[n]
            Delta[n, :] = deltaBrackets_OBBBA[filingStatus, :]
        else:
            sigmaBar[n] = stdDeduction_preTCJA[filingStatus] * gamma_n[n]
            Delta[n, :] = deltaBrackets_preTCJA[filingStatus, :]

        # Add 65+ additional exemption(s) and "bonus" phasing out.
        for i in souls:
            if thisyear + n - yobs[i] >= 65:
                sigmaBar[n] += extra65Deduction[filingStatus] * gamma_n[n]
                if thisyear + n <= 2028:
                    sigmaBar[n] += 6000 * max(0, 1 - 0.06*max(0, MAGI_n[n] - bonusThreshold[filingStatus]))

        # Fill in future tax rates for year n.
        if thisyear + n < yOBBBA:
            theta[n, :] = rates_OBBBA[:]
        else:
            theta[n, :] = rates_preTCJA[:]

    Delta = Delta.transpose()
    theta = theta.transpose()

    # Return series unadjusted for inflation, except for sigmaBar, in STD order.
    return sigmaBar, theta, Delta


def taxBrackets(N_i, n_d, N_n, yOBBBA=2099):
    """
    Return dictionary containing future tax brackets
    unadjusted for inflation for plotting.
    """
    if not (0 < N_i <= 2):
        raise ValueError(f"Cannot process {N_i} individuals.")

    n_d = min(n_d, N_n)
    status = N_i - 1

    # Number of years left in OBBBA from this year.
    thisyear = date.today().year
    if yOBBBA < thisyear:
        raise ValueError(f"Expiration year {yOBBBA} cannot be in the past.")

    ytc = yOBBBA - thisyear

    data = {}
    for t in range(len(taxBracketNames) - 1):
        array = np.zeros(N_n)
        for n in range(N_n):
            stat = status if n < n_d else 0
            array[n] = taxBrackets_OBBBA[stat][t] if n < ytc else taxBrackets_preTCJA[stat][t]

        data[taxBracketNames[t]] = array

    return data


def rho_in(yobs, N_n):
    """
    Return Required Minimum Distribution fractions for each individual.
    This implementation does not support spouses with more than
    10-year difference.
    It starts at age 73 until it goes to 75 in 2033.
    """
    # Notice that table starts at age 72.
    rmdTable = [
        27.4,
        26.5,
        25.5,
        24.6,
        23.7,
        22.9,
        22.0,
        21.1,
        20.2,
        19.4,
        18.5,
        17.7,
        16.8,
        16.0,
        15.2,
        14.4,
        13.7,
        12.9,
        12.2,
        11.5,
        10.8,
        10.1,
        9.5,
        8.9,
        8.4,
        7.8,
        7.3,
        6.8,
        6.4,
        6.0,
        5.6,
        5.2,
        4.9,
        4.6,
    ]

    N_i = len(yobs)
    if N_i == 2 and abs(yobs[0] - yobs[1]) > 10:
        raise RuntimeError("RMD: Unsupported age difference of more than 10 years.")

    rho = np.zeros((N_i, N_n))
    thisyear = date.today().year
    for i in range(N_i):
        agenow = thisyear - yobs[i]
        # Account for increase of RMD age between 2023 and 2032.
        yrmd = 70 if yobs[i] < 1949 else 72 if 1949 <= yobs[i] <= 1950 else 73 if 1951 <= yobs[i] <= 1959 else 75
        for n in range(N_n):
            yage = agenow + n

            if yage < yrmd:
                pass  # rho[i][n] = 0
            else:
                rho[i][n] = 1.0 / rmdTable[yage - 72]

    return rho
