# -*- coding: utf-8 -*-
"""
Consolidated tax calculation engine for the Social Security Tax Torpedo app.

Merges the best of Python_version.py (OOP structure, all 4 filing statuses)
and Tax_calc_breakdown.ipynb (evolved functions, auto-computed base_amount_cap).

All functions are importable and self-contained -- no top-level execution.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Tuple, Dict, Optional
import numpy as np
import pandas as pd

Bracket = Tuple[float, float, float]  # (lower_inclusive, upper_exclusive, rate)


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class SSBThresholds:
    """
    IRS provisional-income thresholds for Social Security taxation.
    These differ by filing status.

    Typical values:
      - Single / HOH: t1=25,000  t2=34,000
      - MFJ:          t1=32,000  t2=44,000
      - MFS (lived with spouse): effectively 0 / 0 (special rule)
    """
    t1: float
    t2: float


@dataclass(frozen=True)
class TaxConfig:
    """
    Complete tax configuration for one filing status.
    ssb_base_amount_cap is auto-computed as 0.5 * (t2 - t1).
    """
    name: str
    brackets: List[Bracket]
    standard_deduction: float
    ssb_thresholds: SSBThresholds
    ssb_base_amount_cap: float = field(init=False)

    def __post_init__(self):
        object.__setattr__(
            self,
            "ssb_base_amount_cap",
            0.5 * (self.ssb_thresholds.t2 - self.ssb_thresholds.t1),
        )


# ---------------------------------------------------------------------------
# 2016 tax bracket data (all 4 filing statuses)
# Rates: 10%, 15%, 25%, 27.75%, 33%, 35%, 39.6%
# Note: 2016 also had a personal exemption of $4,050 per person,
#       which is NOT modeled here. Only the standard deduction is used.
# ---------------------------------------------------------------------------

TAX_BRACKETS_SGL = [
    (0, 9275, 0.10),
    (9275, 37650, 0.15),
    (37650, 91150, 0.25),
    (91150, 190150, 0.2775),
    (190150, 413350, 0.33),
    (413350, 415050, 0.35),
    (415050, float("inf"), 0.396),
]

TAX_BRACKETS_MFJ = [
    (0, 18550, 0.10),
    (18550, 75300, 0.15),
    (75300, 151900, 0.25),
    (151900, 231450, 0.2775),
    (231450, 413350, 0.33),
    (413350, 466950, 0.35),
    (466950, float("inf"), 0.396),
]

TAX_BRACKETS_HOH = [
    (0, 13250, 0.10),
    (13250, 50400, 0.15),
    (50400, 130150, 0.25),
    (130150, 210800, 0.2775),
    (210800, 413350, 0.33),
    (413350, 441000, 0.35),
    (441000, float("inf"), 0.396),
]

TAX_BRACKETS_MFS = [
    (0, 9275, 0.10),
    (9275, 37650, 0.15),
    (37650, 75950, 0.25),
    (75950, 115725, 0.2775),
    (115725, 206675, 0.33),
    (206675, 233475, 0.35),
    (233475, float("inf"), 0.396),
]

# Standard deductions (2016) + personal exemptions folded in
# Personal exemption: $4,050 per person (1 for SGL/HOH/MFS, 2 for MFJ)
STD_DEDUCTION_SGL = 6300 + 4050          # $10,350
STD_DEDUCTION_MFJ = 12600 + 4050 + 4050  # $20,700
STD_DEDUCTION_HOH = 9300 + 4050          # $13,350
STD_DEDUCTION_MFS = 6300 + 4050          # $10,350

# Filing-status configurations
CFG_SGL = TaxConfig(
    name="Single",
    brackets=TAX_BRACKETS_SGL,
    standard_deduction=STD_DEDUCTION_SGL,
    ssb_thresholds=SSBThresholds(t1=25000, t2=34000),
)

CFG_MFJ = TaxConfig(
    name="Married Filing Jointly",
    brackets=TAX_BRACKETS_MFJ,
    standard_deduction=STD_DEDUCTION_MFJ,
    ssb_thresholds=SSBThresholds(t1=32000, t2=44000),
)

CFG_HOH = TaxConfig(
    name="Head of Household",
    brackets=TAX_BRACKETS_HOH,
    standard_deduction=STD_DEDUCTION_HOH,
    ssb_thresholds=SSBThresholds(t1=25000, t2=34000),
)

CFG_MFS = TaxConfig(
    name="Married Filing Separately",
    brackets=TAX_BRACKETS_MFS,
    standard_deduction=STD_DEDUCTION_MFS,
    ssb_thresholds=SSBThresholds(t1=25000, t2=34000),
)

CONFIGS: Dict[str, TaxConfig] = {
    "SGL": CFG_SGL,
    "MFJ": CFG_MFJ,
    "HOH": CFG_HOH,
    "MFS": CFG_MFS,
}


# ---------------------------------------------------------------------------
# Core tax calculation functions
# ---------------------------------------------------------------------------

def ssb_tax(other_income: float, ssb: float, cfg: TaxConfig) -> float:
    """
    Taxable Social Security benefits (IRS worksheet-style).

    PI = other_income + 0.5 * ssb

    Tier 1 (PI <= t1):  taxable = 0
    Tier 2 (t1 < PI <= t2):  taxable = min(0.5*(PI-t1), 0.5*ssb)
    Tier 3 (PI > t2):  taxable = min(0.85*ssb, base_amount + 0.85*(PI-t2))
    """
    t1, t2 = cfg.ssb_thresholds.t1, cfg.ssb_thresholds.t2
    pi = other_income + 0.5 * ssb

    if pi <= t1:
        return 0.0
    if pi <= t2:
        return min(0.5 * (pi - t1), 0.5 * ssb)

    base_amount = min(0.5 * (t2 - t1), 0.5 * ssb)
    candidate = base_amount + 0.85 * (pi - t2)
    return min(0.85 * ssb, candidate)


def bracket_tax(taxable_income: float, cfg: TaxConfig) -> float:
    """Progressive tax from brackets."""
    if taxable_income <= 0:
        return 0.0
    tax = 0.0
    for lower, upper, rate in cfg.brackets:
        if taxable_income <= lower:
            break
        amt = min(taxable_income, upper) - lower
        if amt > 0:
            tax += amt * rate
        if taxable_income <= upper:
            break
    return float(tax)


def compute_baseline_tax(other_income: float, cfg: TaxConfig) -> float:
    """Tax ignoring Social Security (baseline reference)."""
    taxable_income = max(0.0, other_income - cfg.standard_deduction)
    return bracket_tax(taxable_income, cfg)


def tax_with_ssb(other_income: float, ssb: float, cfg: TaxConfig) -> float:
    """Total federal tax including SSB taxation."""
    taxable_ssb = ssb_tax(other_income, ssb, cfg)
    taxable_income = max(0.0, other_income + taxable_ssb - cfg.standard_deduction)
    return bracket_tax(taxable_income, cfg)


def tax_with_ssb_detail(other_income: float, ssb: float, cfg: TaxConfig) -> Dict[str, float]:
    """
    Full tax calculation returning all intermediate values.
    """
    taxable_ssb_val = ssb_tax(other_income, ssb, cfg)
    pi = other_income + 0.5 * ssb
    agi = other_income + taxable_ssb_val
    taxable_income = max(0.0, agi - cfg.standard_deduction)
    tax = bracket_tax(taxable_income, cfg)

    return {
        "other_income": float(other_income),
        "ssb": float(ssb),
        "provisional_income": float(pi),
        "taxable_ssb": float(taxable_ssb_val),
        "pct_ssb_taxable": float(taxable_ssb_val / ssb * 100) if ssb > 0 else 0.0,
        "agi": float(agi),
        "standard_deduction": float(cfg.standard_deduction),
        "taxable_income": float(taxable_income),
        "tax": float(tax),
        "effective_rate": float(tax / other_income * 100) if other_income > 0 else 0.0,
    }


def bracket_marginal_rate(other_income: float, cfg: TaxConfig) -> float:
    """Marginal bracket rate (step function, ignoring SSB effects)."""
    ti = float(max(0.0, other_income - cfg.standard_deduction))
    for lower, upper, rate in cfg.brackets:
        if ti >= lower and ti < upper:
            return float(rate)
    return float(cfg.brackets[-1][2])


def total_marginal_rate(
    other_income: float, ssb: float, cfg: TaxConfig, delta: float = 100.0
) -> float:
    """
    Total marginal tax rate via finite difference (includes SSB torpedo effect).
    """
    t1 = tax_with_ssb(other_income, ssb, cfg)
    t2 = tax_with_ssb(other_income + delta, ssb, cfg)
    return (t2 - t1) / delta


def find_torpedo_bounds(
    cfg: TaxConfig, ssb: float, x_max: float = 200000
) -> Tuple[Optional[float], Optional[float]]:
    """
    Find the zero-point (where tax first > 0) and confluence point
    (where taxable SSB reaches 85% of SSB).

    The zero point is the highest Other Income where the marginal rate
    is still 0% (i.e. tax(OI) == 0 AND tax(OI + delta) == 0), so that
    it plots cleanly inside the No-Tax Zone.

    Returns (zero_point, confluence_point). Either may be None.
    """
    delta = 100.0  # must match the delta used for marginal-rate computation

    # --- Coarse scan to bracket the transition points ---
    xs = np.linspace(0, x_max, 5000)
    ssb_vals = np.array([ssb_tax(x, ssb, cfg) for x in xs])
    tax_vals = np.array([tax_with_ssb(x, ssb, cfg) for x in xs], dtype=float)

    zero_point = None
    confluence_point = None

    # Find approximate zero point (first x where tax > 0)
    approx_zp_lo = None
    approx_zp_hi = None
    prev_x = 0.0
    for x, tv in zip(xs, tax_vals):
        if tv > 0 and approx_zp_lo is None:
            approx_zp_lo = prev_x
            approx_zp_hi = x
        prev_x = x

    # Binary search for the precise boundary where tax transitions from 0 to >0
    if approx_zp_lo is not None:
        lo, hi = approx_zp_lo, approx_zp_hi
        for _ in range(50):  # ~15 decimal digits of precision
            mid = (lo + hi) / 2
            if tax_with_ssb(mid, ssb, cfg) > 0:
                hi = mid
            else:
                lo = mid
        # lo is the highest income where tax == 0
        # Step back by delta so that marginal rate (tax(x+delta) - tax(x))/delta == 0
        zero_point = max(0.0, lo - delta)

    # Find confluence point
    for x, sv in zip(xs, ssb_vals):
        if zero_point is not None and sv >= 0.85 * ssb:
            confluence_point = float(x)
            break

    return zero_point, confluence_point


def classify_zone(
    other_income: float, ssb: float, cfg: TaxConfig,
    zero_point: Optional[float] = None, confluence_point: Optional[float] = None,
) -> str:
    """
    Classify the user's income into one of three zones:
      - 'No-Tax Zone'   (green)  -- tax == 0
      - 'High-Tax Zone'  (red)   -- in the torpedo
      - 'Same-Old Zone'  (blue)  -- past the torpedo
    """
    my_tax = tax_with_ssb(other_income, ssb, cfg)

    if my_tax <= 0:
        return "No-Tax Zone"

    if confluence_point is not None and other_income >= confluence_point:
        return "Same-Old Zone"

    return "High-Tax Zone"


def bracket_breakdown(taxable_income: float, cfg: TaxConfig) -> pd.DataFrame:
    """DataFrame showing tax in each bracket."""
    rows = []
    ti = float(max(0.0, taxable_income))
    for lower, upper, rate in cfg.brackets:
        if ti <= lower:
            amt = 0.0
        else:
            amt = max(0.0, min(ti, upper) - lower)
        rows.append({
            "lower": lower,
            "upper": upper if np.isfinite(upper) else np.inf,
            "rate": rate,
            "taxed_amount": amt,
            "tax_in_bracket": amt * rate,
        })
        if ti <= upper:
            break

    df = pd.DataFrame(rows)
    df["tax_in_bracket"] = df["tax_in_bracket"].astype(float)
    df["cum_tax"] = df["tax_in_bracket"].cumsum()
    return df