retiretui_engine 0.2.0

//! Federal tax formulas. The rule shapes live here; the yearly values come
//! from [`crate::params`].

use std::collections::BTreeMap;
use std::ops::RangeInclusive;

use crate::params::{BenefitParams, Bracket, StateParams, TaxParams};
use crate::plan::{AccountKind, Dollars, FilingStatus, Person, PlanDate};

/// The age Medicare coverage (and IRMAA exposure) begins.
pub const MEDICARE_AGE: u8 = 65;

/// The earliest age a retirement benefit can be claimed, and whose year's
/// bend points compute it.
pub const EARLIEST_CLAIM_AGE: i16 = 62;

/// The age whose year's average wage a worker's earnings are indexed to.
const INDEXING_AGE: i16 = 60;

/// The age a career filled at one salary is taken to begin, as SSA's own
/// illustrative earners do.
pub const FIRST_WORKING_AGE: i16 = 22;

/// The age past which delaying a claim earns nothing more.
pub const LATEST_CREDIT_AGE: i16 = 70;

/// Months in a year, which the benefit's monthly figures are counted in.
pub const MONTHS_PER_YEAR: i32 = 12;

/// [`EARLIEST_CLAIM_AGE`] in months.
pub const EARLIEST_CLAIM_MONTHS: i32 = EARLIEST_CLAIM_AGE as i32 * MONTHS_PER_YEAR;
/// [`LATEST_CREDIT_AGE`] in months.
pub const LATEST_CREDIT_MONTHS: i32 = LATEST_CREDIT_AGE as i32 * MONTHS_PER_YEAR;
/// The highest indexed years averaged into the AIME, and their months.
const COMPUTATION_YEARS: usize = 35;
const COMPUTATION_MONTHS: f64 = (COMPUTATION_YEARS as i32 * MONTHS_PER_YEAR) as f64;
/// The PIA formula's replacement rates, per bend-point band.
const PIA_RATES: [f64; 3] = [0.90, 0.32, 0.15];
/// Early claims lose 5/9 of 1% a month for 36 months, 5/12 of 1% beyond;
/// delayed claims earn 2/3 of 1% a month for those born 1943 and later.
const EARLY_FIRST_MONTHS: i32 = 36;
const EARLY_FIRST_RATE: f64 = 5.0 / 900.0;
const EARLY_BEYOND_RATE: f64 = 5.0 / 1200.0;
const DELAYED_RATE: f64 = 2.0 / 300.0;
/// Full retirement age rises two months a birth year across 1938-42 and
/// 1955-59, from 65 to 66 and from 66 to 67.
const FRA_65_LAST_BIRTH_YEAR: i16 = 1937;
const FRA_66_FIRST_BIRTH_YEAR: i16 = 1943;
const FRA_66_LAST_BIRTH_YEAR: i16 = 1954;
const FRA_67_FIRST_BIRTH_YEAR: i16 = 1960;
const FRA_STEP_MONTHS: i32 = 2;

/// Years between the MAGI that is measured and the premiums it prices.
pub const IRMAA_LOOKBACK_YEARS: i16 = 2;

/// Whether a person is Medicare-covered in `year`.
#[must_use]
pub fn is_medicare_covered(person: &Person, year: i16) -> bool {
    person.age_in_year(year) >= i16::from(MEDICARE_AGE)
}

/// Birth year from which SECURE 2.0 moves the RMD start to 75.
const RMD_AGE_75_BIRTH_YEAR: i16 = 1960;
const RMD_START_AGE_EARLY: u8 = 73;
const RMD_START_AGE_LATE: u8 = 75;

/// Tax on ordinary taxable income (after deductions) via the bracket walk.
#[must_use]
pub fn ordinary_tax(params: &TaxParams, status: FilingStatus, taxable: Dollars) -> Dollars {
    walk_brackets(params.brackets.for_status(status), taxable)
}

/// A state's income tax: ordinary income and gains alike, with the taxable
/// share of Social Security where the state taxes it, less its deduction.
#[must_use]
pub fn state_tax(
    state: &StateParams,
    status: FilingStatus,
    income: Dollars,
    taxable_social_security: Dollars,
) -> Dollars {
    let benefits = if state.taxes_social_security {
        taxable_social_security
    } else {
        0
    };
    let taxable = income + benefits - state.deduction.get(status);
    walk_brackets(state.brackets.for_status(status), taxable)
}

fn walk_brackets(brackets: &[Bracket], taxable: Dollars) -> Dollars {
    let mut tax = 0.0;
    for (i, bracket) in brackets.iter().enumerate() {
        if taxable <= bracket.over {
            break;
        }
        let top = brackets
            .get(i + 1)
            .map_or(taxable, |next| taxable.min(next.over));
        tax += bracket.rate * (top - bracket.over) as f64;
    }
    tax.round() as Dollars
}

/// Tax on long-term gains stacked on top of ordinary taxable income.
#[must_use]
pub fn ltcg_tax(
    params: &TaxParams,
    status: FilingStatus,
    ordinary_taxable: Dollars,
    gains: Dollars,
) -> Dollars {
    if gains <= 0 {
        return 0;
    }
    let stack_bottom = ordinary_taxable.max(0);
    let zero_top = params.ltcg.zero_until.get(status);
    let fifteen_top = params.ltcg.fifteen_until.get(status);
    let in_zero = (zero_top - stack_bottom).clamp(0, gains);
    let above_zero = gains - in_zero;
    let in_middle = (fifteen_top - stack_bottom.max(zero_top)).clamp(0, above_zero);
    let in_top = above_zero - in_middle;
    let tax = params.ltcg.middle_rate * in_middle as f64 + params.ltcg.top_rate * in_top as f64;
    tax.round() as Dollars
}

/// The taxable portion of a Social Security benefit, from provisional income
/// (`other_income` plus half the benefit) against the statutory thresholds.
#[must_use]
pub fn taxable_social_security(
    params: &TaxParams,
    status: FilingStatus,
    other_income: Dollars,
    benefit: Dollars,
) -> Dollars {
    if benefit <= 0 {
        return 0;
    }
    let provisional = other_income + benefit / 2;
    let base = params.social_security.provisional_base.get(status);
    let upper = params.social_security.provisional_upper.get(status);
    if provisional <= base {
        return 0;
    }
    if provisional <= upper {
        return ((provisional - base) / 2).min(benefit / 2);
    }
    let from_lower_band = ((upper - base) / 2).min(benefit / 2);
    let above_upper = (0.85 * (provisional - upper) as f64).round() as Dollars;
    (above_upper + from_lower_band).min((0.85 * benefit as f64).round() as Dollars)
}

/// Full retirement age in months, from the birth year.
#[must_use]
pub fn full_retirement_months(birth_year: i16) -> i32 {
    let steps = |from: i16| FRA_STEP_MONTHS * i32::from(birth_year - from);
    if birth_year <= FRA_65_LAST_BIRTH_YEAR {
        65 * MONTHS_PER_YEAR
    } else if birth_year < FRA_66_FIRST_BIRTH_YEAR {
        65 * MONTHS_PER_YEAR + steps(FRA_65_LAST_BIRTH_YEAR)
    } else if birth_year <= FRA_66_LAST_BIRTH_YEAR {
        66 * MONTHS_PER_YEAR
    } else if birth_year < FRA_67_FIRST_BIRTH_YEAR {
        66 * MONTHS_PER_YEAR + steps(FRA_66_LAST_BIRTH_YEAR)
    } else {
        67 * MONTHS_PER_YEAR
    }
}

/// The annual retirement benefit for a claim at `age_months` of age: each
/// year's covered earnings capped at its own wage base and indexed to the
/// average wage of the year the worker turns 60, the highest 35 years
/// averaged monthly, the PIA through the bend points of the year they turn
/// 62 truncated to the dime, carried through each of `colas` in turn and
/// truncated to the dime again after each, then reduced or credited month
/// by month against full retirement age and truncated to the dollar.
/// Claims past 70 earn 70's credit. Without `colas` the benefit is in the
/// dollars of the year the worker turns 62.
#[must_use]
pub fn social_security_benefit(
    params: &BenefitParams,
    birth_year: i16,
    age_months: i32,
    earnings: &BTreeMap<i16, Dollars>,
    colas: &[f64],
) -> Dollars {
    let pia = (colas.iter()).fold(
        primary_insurance_amount(params, birth_year, earnings),
        |pia, cola| to_the_dime(pia * (1.0 + cola)),
    );
    let monthly = (pia * claim_factor(birth_year, age_months) * 100.0).round() / 100.0;
    monthly.floor() as Dollars * MONTHS_PER_YEAR as Dollars
}

fn to_the_dime(amount: f64) -> f64 {
    ((amount * 100.0).round() / 10.0).floor() / 10.0
}

/// A career at one real wage, as SSA's Quick Calculator fills a record
/// from current earnings: `salary`, in the dollars of `paid_in`, scaled to
/// each of `years` by that year's average wage over `paid_in`'s, so every
/// year indexes back to the salary; a year the index does not carry is
/// left out.
#[must_use]
pub fn earnings_at_wage(
    params: &BenefitParams,
    salary: Dollars,
    paid_in: i16,
    years: RangeInclusive<i16>,
) -> BTreeMap<i16, Dollars> {
    years
        .filter_map(|year| {
            let scaled = salary as f64 * params.wage_ratio(year, paid_in)?;
            Some((year, scaled.round() as Dollars))
        })
        .collect()
}

/// A career before `start_year` at one `salary` in that year's dollars,
/// through [`earnings_at_wage`] from the year the worker turned
/// [`FIRST_WORKING_AGE`].
#[must_use]
pub fn career_before(
    params: &BenefitParams,
    birth_year: i16,
    salary: Dollars,
    start_year: i16,
) -> BTreeMap<i16, Dollars> {
    let from = birth_year + FIRST_WORKING_AGE;
    earnings_at_wage(params, salary, start_year, from..=start_year - 1)
}

/// The age in months a claim at `age_months` is paid at for the rest of
/// the calendar year it is made in, whose January the worker was
/// `january_months` old in: credits past full retirement age earned that
/// year are paid from the next January, except for a claim at 70 or later,
/// credited at once.
#[must_use]
pub fn claim_year_age(birth_year: i16, age_months: i32, january_months: i32) -> i32 {
    let full = full_retirement_months(birth_year);
    if age_months > full && age_months < LATEST_CREDIT_MONTHS {
        january_months.clamp(full, age_months)
    } else {
        age_months
    }
}

/// A calendar month as one count across years: `year * 12 + month - 1`.
#[must_use]
pub fn month_index(year: i16, month: i8) -> i32 {
    i32::from(year) * MONTHS_PER_YEAR + i32::from(month) - 1
}

/// The month someone born on `birth` attains `age`: an age is attained the
/// day before the birthday, so a month early for a birthday on the 1st.
#[must_use]
pub fn attained_month(birth: PlanDate, age: u8) -> i32 {
    month_index(birth.year() + i16::from(age), birth.0.month()) - i32::from(birth.0.day() == 1)
}

/// Someone born on `birth`'s age in `month`, in whole months.
#[must_use]
pub fn age_months(birth: PlanDate, month: i32) -> i32 {
    month - attained_month(birth, 0)
}

/// The first month a retirement benefit can be paid for: 62 must be held
/// throughout it, so the month 62 is attained only when it is attained on
/// that month's first day.
#[must_use]
pub fn first_claim_month(birth: PlanDate) -> i32 {
    let attained = attained_month(birth, EARLIEST_CLAIM_AGE as u8);
    attained + i32::from(birth.0.day() != 2)
}

/// The share of `year` a benefit first paid for `first_paid` is paid for:
/// the months from it, out of twelve; all of a later year and none of an
/// earlier one.
#[must_use]
pub fn claim_year_share(first_paid: i32, year: i16) -> f64 {
    let months = (month_index(year + 1, 1) - first_paid).clamp(0, MONTHS_PER_YEAR);
    f64::from(months) / f64::from(MONTHS_PER_YEAR)
}

fn primary_insurance_amount(
    params: &BenefitParams,
    birth_year: i16,
    earnings: &BTreeMap<i16, Dollars>,
) -> f64 {
    let index_year = birth_year + INDEXING_AGE;
    let to_index_year = |year| {
        if year >= index_year {
            1.0
        } else {
            params.wage_ratio(index_year, year).unwrap_or(1.0)
        }
    };
    let mut indexed: Vec<f64> = earnings
        .iter()
        .map(|(&year, &amount)| amount.min(params.wage_base(year)) as f64 * to_index_year(year))
        .collect();
    indexed.sort_by(|a, b| b.total_cmp(a));
    let total: f64 = indexed.iter().take(COMPUTATION_YEARS).sum();
    let aime = (total / COMPUTATION_MONTHS).floor();
    let eligibility_year = birth_year + EARLIEST_CLAIM_AGE;
    let [first, second] = params
        .bend_points(eligibility_year)
        .map(|point| point as f64);
    let pia = PIA_RATES[0] * aime.min(first)
        + PIA_RATES[1] * (aime.min(second) - first).max(0.0)
        + PIA_RATES[2] * (aime - second).max(0.0);
    to_the_dime(pia)
}

fn claim_factor(birth_year: i16, age_months: i32) -> f64 {
    let full = full_retirement_months(birth_year);
    let claim = age_months.min(LATEST_CREDIT_MONTHS);
    if claim >= full {
        return 1.0 + f64::from(claim - full) * DELAYED_RATE;
    }
    let early = full - claim;
    let first = early.min(EARLY_FIRST_MONTHS);
    1.0 - f64::from(first) * EARLY_FIRST_RATE - f64::from(early - first) * EARLY_BEYOND_RATE
}

/// The annual IRMAA surcharge per covered person for a lookback MAGI.
/// Tiers are cliffs: one dollar over a threshold buys the whole tier.
#[must_use]
pub fn irmaa_surcharge(
    params: &TaxParams,
    status: FilingStatus,
    magi: Dollars,
    part_d: bool,
) -> Dollars {
    params
        .irmaa
        .iter()
        .rev()
        .find(|tier| magi > tier.magi_over.get(status))
        .map_or(0, |tier| tier.part_b + if part_d { tier.part_d } else { 0 })
}

/// The MAGI at which IRMAA tier `tier` (zero-indexed) begins - the ceiling
/// that keeps a household below it; `None` past the table.
#[must_use]
pub fn irmaa_threshold(params: &TaxParams, status: FilingStatus, tier: u8) -> Option<Dollars> {
    params
        .irmaa
        .get(usize::from(tier))
        .map(|entry| entry.magi_over.get(status))
}

/// The required minimum distribution for the year a person reaches `age`,
/// from the prior year-end `balance`. Zero below the table's first age; ages
/// past the table's last row keep its final divisor.
#[must_use]
pub fn rmd(params: &TaxParams, age: u8, balance: Dollars) -> Dollars {
    if balance <= 0 {
        return 0;
    }
    let divisors = &params.rmd.divisors;
    let Some(first) = divisors.first() else {
        return 0;
    };
    if age < first.age {
        return 0;
    }
    let divisor = divisors
        .iter()
        .rev()
        .find(|row| row.age <= age)
        .map_or(first.divisor, |row| row.divisor);
    (balance as f64 / divisor).round() as Dollars
}

/// The age a person's RMDs begin under SECURE 2.0, from their birth year.
#[must_use]
pub fn rmd_start_age(birth_year: i16) -> u8 {
    if birth_year >= RMD_AGE_75_BIRTH_YEAR {
        RMD_START_AGE_LATE
    } else {
        RMD_START_AGE_EARLY
    }
}

/// Whether early withdrawals from this wrapper escape the penalty
/// (governmental 457(b) distributions are penalty-free at any age).
#[must_use]
pub fn is_penalty_exempt(kind: AccountKind) -> bool {
    matches!(kind, AccountKind::K457b)
}

/// The accounts of one person that share one employee limit for the year.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
pub enum LimitPool {
    /// 401(k) and 403(b), one elective limit between them.
    EmployerPlan,
    /// A 457(b), limited on its own.
    Deferred457,
    /// A SIMPLE IRA, limited on its own.
    Simple,
    /// Traditional and Roth IRAs together.
    Ira,
    /// An HSA, whose employer contributions count too.
    Hsa,
}

/// The pool the kind's employee contributions count against, or `None`
/// where the law limits nothing.
#[must_use]
pub fn limit_pool(kind: AccountKind) -> Option<LimitPool> {
    match kind {
        AccountKind::K401k | AccountKind::K403b => Some(LimitPool::EmployerPlan),
        AccountKind::K457b => Some(LimitPool::Deferred457),
        AccountKind::SimpleIra => Some(LimitPool::Simple),
        AccountKind::Ira => Some(LimitPool::Ira),
        AccountKind::Hsa => Some(LimitPool::Hsa),
        AccountKind::K414k | AccountKind::SepIra | AccountKind::Brokerage | AccountKind::Cash => {
            None
        }
    }
}

/// The employee contribution limit for an account kind at an age (the age
/// reached during the year), or `None` when no employee limit applies.
/// HSA family coverage is assumed for a married household.
#[must_use]
pub fn employee_limit(
    params: &TaxParams,
    status: FilingStatus,
    kind: AccountKind,
    age: u8,
) -> Option<Dollars> {
    let limits = &params.limits;
    let plan_catch_up = |base_50: Dollars, base_60: Dollars| match age {
        60..=63 => base_60,
        50.. => base_50,
        _ => 0,
    };
    let limit = match limit_pool(kind)? {
        LimitPool::EmployerPlan | LimitPool::Deferred457 => {
            limits.employer_plan
                + plan_catch_up(
                    limits.employer_plan_catch_up_50,
                    limits.employer_plan_catch_up_60,
                )
        }
        LimitPool::Simple => {
            limits.simple + plan_catch_up(limits.simple_catch_up_50, limits.simple_catch_up_60)
        }
        LimitPool::Ira => limits.ira + if age >= 50 { limits.ira_catch_up_50 } else { 0 },
        LimitPool::Hsa => {
            let coverage = match status {
                FilingStatus::Single => limits.hsa_self,
                FilingStatus::MarriedJoint => limits.hsa_family,
            };
            coverage + if age >= 55 { limits.hsa_catch_up_55 } else { 0 }
        }
    };
    Some(limit)
}