pension_engine.py

"""
PensionSight — Core Financial Calculation Engine
All pure math, no external dependencies.
"""

import math
from dataclasses import dataclass, field
from typing import Optional


# ─────────────────────────────────────────────
# NPS SCHEME CONSTANTS (2025 PFRDA norms)
# ─────────────────────────────────────────────

NPS_SCHEMES = {
    "LC25":  {"name": "Life Cycle 25 – Low",       "max_equity": 0.25, "equity_exit_age": 55},
    "LC50":  {"name": "Life Cycle 50 – Moderate",  "max_equity": 0.50, "equity_exit_age": 55},
    "LC75":  {"name": "Life Cycle 75 – High",      "max_equity": 0.75, "equity_exit_age": 55},
    "AGGRESSIVE": {"name": "Life Cycle Aggressive","max_equity": 0.35, "equity_exit_age": 55},
    "ACTIVE": {"name": "Active Choice",            "max_equity": 0.75, "equity_exit_age": None},
}

# Historical NPS fund CAGR approximations (2004-2025 verified ranges)
SCENARIO_RETURNS = {
    "conservative": {
        "equity": 0.10,
        "corporate_bond": 0.07,
        "govt_bond": 0.06,
        "label": "Conservative (low market)"
    },
    "realistic": {
        "equity": 0.13,
        "corporate_bond": 0.08,
        "govt_bond": 0.07,
        "label": "Realistic (average market)"
    },
    "optimistic": {
        "equity": 0.16,
        "corporate_bond": 0.095,
        "govt_bond": 0.08,
        "label": "Optimistic (strong market)"
    }
}

ANNUITY_RATE_DEFAULT = 0.06        # 6% annuity rate (ASP average 2025)
ANNUITY_MIN_PERCENT  = 0.40        # 40% minimum annuity purchase (PFRDA 2025 norm)
INFLATION_RATE       = 0.06        # 6% CPI (RBI average)
MIN_RETIREMENT_AGE   = 60
MAX_DEFERRAL_AGE     = 75
TAX_EXEMPT_LUMPSUM   = 0.60        # 60% lumpsum is tax-free under NPS


# ─────────────────────────────────────────────
# DATA MODELS
# ─────────────────────────────────────────────

@dataclass
class SubscriberProfile:
    current_age: int
    retirement_age: int
    monthly_contribution: float
    existing_corpus: float = 0.0
    annual_contribution_increase: float = 0.0   # % per year step-up
    sector: str = "non_government"               # government / non_government / vatsalya
    scheme: str = "LC50"
    annuity_percent: float = 0.40
    annuity_rate: float = ANNUITY_RATE_DEFAULT
    deferral_age: Optional[int] = None           # defer exit beyond 60
    employer_contribution: float = 0.0           # monthly employer top-up
    desired_monthly_pension: Optional[float] = None

    # Gig worker fields
    is_gig_worker: bool = False
    monthly_incomes: list = field(default_factory=list)   # list of monthly income values

    @property
    def investment_years(self) -> int:
        end_age = self.deferral_age if self.deferral_age else self.retirement_age
        return max(1, end_age - self.current_age)

    @property
    def total_monthly_contribution(self) -> float:
        return self.monthly_contribution + self.employer_contribution


@dataclass
class ScenarioResult:
    scenario: str
    label: str
    blended_return: float
    projected_corpus: float
    real_corpus: float                # inflation-adjusted
    lumpsum_withdrawal: float
    annuity_corpus: float
    monthly_pension: float
    real_monthly_pension: float
    total_contributions: float
    wealth_gained: float
    investment_years: int


@dataclass
class ReverseplanResult:
    desired_monthly_pension: float
    required_corpus: float
    required_monthly_sip: float
    scenario: str
    investment_years: int
    current_age: int
    retirement_age: int


@dataclass
class NudgeResult:
    nudge_type: str
    message: str
    impact_rupees: float
    current_value: float
    improved_value: float


# ─────────────────────────────────────────────
# CORE CALCULATION ENGINE
# ─────────────────────────────────────────────

class PensionCalculationEngine:

    def _blended_return(self, scheme_key: str, scenario: str) -> float:
        """Calculate blended return based on scheme allocation and scenario."""
        scheme = NPS_SCHEMES.get(scheme_key, NPS_SCHEMES["LC50"])
        s = SCENARIO_RETURNS[scenario]

        eq   = scheme["max_equity"]
        gb   = (1 - eq) * 0.6    # 60% of debt in govt bonds
        cb   = (1 - eq) * 0.4    # 40% of debt in corporate bonds

        return round(eq * s["equity"] + cb * s["corporate_bond"] + gb * s["govt_bond"], 4)

    def _future_value_step_up_sip(
        self,
        monthly_sip: float,
        annual_return: float,
        years: int,
        annual_step_up_pct: float = 0.0,
        existing_corpus: float = 0.0
    ) -> float:
        """
        Future value of a step-up SIP (SIP increases by annual_step_up_pct every year).
        Also compounds any existing corpus.
        """
        monthly_rate = annual_return / 12
        total_months = years * 12

        # Compound existing corpus
        fv_existing = existing_corpus * ((1 + monthly_rate) ** total_months)

        # Step-up SIP: iterate year by year
        fv_sip = 0.0
        for year in range(years):
            sip_this_year = monthly_sip * ((1 + annual_step_up_pct / 100) ** year)
            months_remaining = (years - year) * 12
            # FV of 12 equal monthly payments with months_remaining left to compound
            for m in range(12):
                months_to_end = months_remaining - m
                fv_sip += sip_this_year * ((1 + monthly_rate) ** months_to_end)

        return fv_existing + fv_sip

    def _gig_worker_avg_monthly(self, monthly_incomes: list) -> float:
        """Average monthly income for gig workers from a list of monthly values."""
        if not monthly_incomes:
            return 0.0
        return sum(monthly_incomes) / len(monthly_incomes)

    def _monthly_pension_from_corpus(self, annuity_corpus: float, annuity_rate: float) -> float:
        """Monthly pension from annuity corpus at given annual annuity rate."""
        return (annuity_corpus * annuity_rate) / 12

    def _real_value(self, nominal: float, years: int) -> float:
        """Deflate nominal value to today's purchasing power."""
        return nominal / ((1 + INFLATION_RATE) ** years)

    # ── Main projection ──────────────────────────────────────────────────────

    def project_all_scenarios(self, profile: SubscriberProfile) -> dict:
        """Run 3-scenario projection for a subscriber profile."""
        results = {}

        # Effective monthly contribution (handle gig worker)
        if profile.is_gig_worker and profile.monthly_incomes:
            avg_income = self._gig_worker_avg_monthly(profile.monthly_incomes)
            # Use contribution as % of average income if provided, else use monthly_contribution
            effective_monthly = profile.monthly_contribution if profile.monthly_contribution > 0 \
                                 else avg_income * 0.10   # default 10% if not set
        else:
            effective_monthly = profile.total_monthly_contribution

        years = profile.investment_years
        total_contributions = effective_monthly * 12 * years  # rough estimate

        for scenario in ["conservative", "realistic", "optimistic"]:
            r = self._blended_return(profile.scheme, scenario)
            corpus = self._future_value_step_up_sip(
                monthly_sip=effective_monthly,
                annual_return=r,
                years=years,
                annual_step_up_pct=profile.annual_contribution_increase,
                existing_corpus=profile.existing_corpus
            )

            annuity_corpus     = corpus * max(profile.annuity_percent, ANNUITY_MIN_PERCENT)
            lumpsum            = corpus * (1 - max(profile.annuity_percent, ANNUITY_MIN_PERCENT))
            monthly_pension    = self._monthly_pension_from_corpus(annuity_corpus, profile.annuity_rate)
            real_corpus        = self._real_value(corpus, years)
            real_monthly_pen   = self._real_value(monthly_pension, years)

            results[scenario] = ScenarioResult(
                scenario=scenario,
                label=SCENARIO_RETURNS[scenario]["label"],
                blended_return=round(r * 100, 2),
                projected_corpus=round(corpus),
                real_corpus=round(real_corpus),
                lumpsum_withdrawal=round(lumpsum),
                annuity_corpus=round(annuity_corpus),
                monthly_pension=round(monthly_pension),
                real_monthly_pension=round(real_monthly_pen),
                total_contributions=round(total_contributions),
                wealth_gained=round(corpus - total_contributions),
                investment_years=years
            )

        return results

    # ── Reverse Planner ──────────────────────────────────────────────────────

    def reverse_plan(
        self,
        desired_monthly_pension: float,
        current_age: int,
        retirement_age: int,
        scheme: str = "LC50",
        scenario: str = "realistic",
        annuity_percent: float = 0.40,
        annuity_rate: float = ANNUITY_RATE_DEFAULT,
        existing_corpus: float = 0.0,
        annual_step_up: float = 0.0
    ) -> ReverseplanResult:
        """
        Given a desired monthly pension, calculate required monthly SIP today.
        Works backwards: pension → annuity corpus → total corpus → SIP.
        """
        years = max(1, retirement_age - current_age)
        r     = self._blended_return(scheme, scenario)
        monthly_rate = r / 12
        total_months = years * 12

        # Step 1: Annuity corpus needed for desired pension
        annual_pension     = desired_monthly_pension * 12
        required_annuity_c = annual_pension / annuity_rate

        # Step 2: Total corpus needed (annuity is annuity_percent of corpus)
        effective_annuity_pct = max(annuity_percent, ANNUITY_MIN_PERCENT)
        required_corpus = required_annuity_c / effective_annuity_pct

        # Step 3: Subtract compounded existing corpus
        fv_existing       = existing_corpus * ((1 + monthly_rate) ** total_months)
        corpus_from_sip   = max(0, required_corpus - fv_existing)

        # Step 4: Required monthly SIP (standard FV of annuity formula)
        if corpus_from_sip == 0:
            required_sip = 0.0
        elif annual_step_up == 0:
            # Simple SIP formula: FV = SIP × [((1+r)^n - 1)/r] × (1+r)
            fv_factor    = ((1 + monthly_rate) ** total_months - 1) / monthly_rate * (1 + monthly_rate)
            required_sip = corpus_from_sip / fv_factor
        else:
            # Step-up SIP: binary search (harder to invert analytically)
            required_sip = self._binary_search_sip(
                target_corpus=corpus_from_sip,
                annual_return=r,
                years=years,
                annual_step_up=annual_step_up
            )

        return ReverseplanResult(
            desired_monthly_pension=desired_monthly_pension,
            required_corpus=round(required_corpus),
            required_monthly_sip=round(required_sip),
            scenario=scenario,
            investment_years=years,
            current_age=current_age,
            retirement_age=retirement_age
        )

    def _binary_search_sip(
        self,
        target_corpus: float,
        annual_return: float,
        years: int,
        annual_step_up: float,
        tolerance: float = 100
    ) -> float:
        """Binary search for required SIP when step-up is involved."""
        lo, hi = 100.0, target_corpus
        for _ in range(60):
            mid = (lo + hi) / 2
            fv  = self._future_value_step_up_sip(mid, annual_return, years, annual_step_up)
            if abs(fv - target_corpus) < tolerance:
                return mid
            if fv < target_corpus:
                lo = mid
            else:
                hi = mid
        return (lo + hi) / 2

    # ── AI Nudge Engine ───────────────────────────────────────────────────────

    def generate_nudges(self, profile: SubscriberProfile) -> list:
        """
        Rule-based nudge engine: compares current vs improved scenarios.
        Returns list of NudgeResult objects sorted by impact.
        """
        nudges = []
        base_results  = self.project_all_scenarios(profile)
        base_corpus   = base_results["realistic"].projected_corpus
        base_pension  = base_results["realistic"].monthly_pension
        years         = profile.investment_years

        # Nudge 1: Increase SIP by ₹500
        if profile.monthly_contribution < 50000:
            boost = 500
            improved = self._sim_corpus_change(profile, sip_delta=boost)
            gain = improved - base_corpus
            if gain > 0:
                nudges.append(NudgeResult(
                    nudge_type="sip_increase",
                    message=f"Increasing your monthly SIP by ₹{boost:,} adds ₹{gain:,.0f} to your corpus — that's ₹{round(gain * ANNUITY_MIN_PERCENT * ANNUITY_RATE_DEFAULT / 12):,}/month extra pension.",
                    impact_rupees=gain,
                    current_value=base_corpus,
                    improved_value=improved
                ))

        # Nudge 2: Add annual step-up of 5%
        if profile.annual_contribution_increase < 5:
            improved = self._sim_corpus_change(profile, step_up_delta=5)
            gain = improved - base_corpus
            if gain > 0:
                nudges.append(NudgeResult(
                    nudge_type="step_up",
                    message=f"Adding a 5% annual step-up to your SIP (raise contributions by 5% each year) grows your corpus by ₹{gain:,.0f} — ₹{round(gain * ANNUITY_MIN_PERCENT * ANNUITY_RATE_DEFAULT / 12):,}/month more pension.",
                    impact_rupees=gain,
                    current_value=base_corpus,
                    improved_value=improved
                ))

        # Nudge 3: Start 2 years earlier
        if profile.current_age > 25:
            improved = self._sim_corpus_change(profile, age_delta=-2)
            gain = improved - base_corpus
            if gain > 0:
                nudges.append(NudgeResult(
                    nudge_type="early_start",
                    message=f"If you had started NPS 2 years earlier, your corpus would be ₹{gain:,.0f} higher. Starting early is the single most powerful retirement lever.",
                    impact_rupees=gain,
                    current_value=base_corpus,
                    improved_value=improved
                ))

        # Nudge 4: Retire 2 years later (defer)
        if profile.retirement_age <= 60:
            improved = self._sim_corpus_change(profile, retire_later=2)
            gain = improved - base_corpus
            if gain > 0:
                nudges.append(NudgeResult(
                    nudge_type="defer_retirement",
                    message=f"Deferring retirement by 2 years to age {profile.retirement_age + 2} adds ₹{gain:,.0f} to your corpus — ₹{round(gain * ANNUITY_MIN_PERCENT * ANNUITY_RATE_DEFAULT / 12):,}/month more pension.",
                    impact_rupees=gain,
                    current_value=base_corpus,
                    improved_value=improved
                ))

        # Nudge 5: Switch to higher growth scheme
        if profile.scheme in ["LC25", "LC50"] and profile.current_age < 45:
            improved = self._sim_corpus_change(profile, scheme_upgrade="LC75")
            gain = improved - base_corpus
            if gain > 0:
                nudges.append(NudgeResult(
                    nudge_type="scheme_upgrade",
                    message=f"Switching from {NPS_SCHEMES[profile.scheme]['name']} to LC75 (higher equity at your young age) could add ₹{gain:,.0f} to your corpus.",
                    impact_rupees=gain,
                    current_value=base_corpus,
                    improved_value=improved
                ))

        nudges.sort(key=lambda x: x.impact_rupees, reverse=True)
        return nudges

    def _sim_corpus_change(
        self, profile, sip_delta=0, step_up_delta=0,
        age_delta=0, retire_later=0, scheme_upgrade=None
    ) -> float:
        """Simulate corpus with a single parameter change."""
        import copy
        p = copy.copy(profile)
        p.monthly_contribution        += sip_delta
        p.annual_contribution_increase = max(0, p.annual_contribution_increase + step_up_delta)
        p.current_age                 = max(18, p.current_age + age_delta)
        p.retirement_age              += retire_later
        if scheme_upgrade:
            p.scheme = scheme_upgrade
        return self.project_all_scenarios(p)["realistic"].projected_corpus