paye_calculator.py

"""
Comprehensive PAYE Calculator for Nigeria Tax Act 2026.

Features:
- Full deduction calculations (pension, NHF, rent relief)
- Progressive tax band computation
- Minimum tax rule application
- Validation and confidence scoring
- WhatsApp and Web formatted outputs
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Any
from datetime import date
import json

from tax_config import (
    get_regime, get_active_regime, TaxRegimeConfig, TaxBand,
    NTA_2026_CONFIG, format_bands
)


@dataclass
class DeductionBreakdown:
    """Breakdown of all deductions applied."""
    pension_contribution: float = 0.0
    nhf_contribution: float = 0.0
    nhis_contribution: float = 0.0
    life_insurance: float = 0.0
    rent_relief: float = 0.0
    cra_amount: float = 0.0  # For PITA regime
    other_deductions: float = 0.0
    
    @property
    def total(self) -> float:
        return (
            self.pension_contribution +
            self.nhf_contribution +
            self.nhis_contribution +
            self.life_insurance +
            self.rent_relief +
            self.cra_amount +
            self.other_deductions
        )


@dataclass
class BandCalculation:
    """Details of tax calculated in a single band."""
    band_lower: float
    band_upper: float
    rate: float
    taxable_in_band: float
    tax_amount: float
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "range": f"N{self.band_lower:,.0f} - N{self.band_upper:,.0f}",
            "rate": f"{self.rate * 100:.0f}%",
            "taxable": self.taxable_in_band,
            "tax": self.tax_amount
        }


@dataclass
class ValidationResult:
    """Validation result for a calculation."""
    is_valid: bool
    errors: List[str] = field(default_factory=list)
    warnings: List[str] = field(default_factory=list)
    confidence: float = 1.0


@dataclass
class PAYECalculation:
    """Complete PAYE calculation result."""
    # Income
    gross_annual_income: float
    gross_monthly_income: float
    
    # Deductions
    deductions: DeductionBreakdown
    
    # Taxable income
    taxable_income: float
    
    # Tax computation
    band_calculations: List[BandCalculation]
    computed_tax: float
    minimum_tax: float
    final_tax: float
    
    # Rates
    effective_rate: float
    marginal_rate: float
    
    # Net pay
    annual_net_pay: float
    monthly_net_pay: float
    monthly_tax: float
    
    # Metadata
    regime: str
    calculation_date: date
    legal_citations: List[str]
    validation: ValidationResult
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return {
            "gross_income": {
                "annual": self.gross_annual_income,
                "monthly": self.gross_monthly_income
            },
            "deductions": {
                "pension": self.deductions.pension_contribution,
                "nhf": self.deductions.nhf_contribution,
                "nhis": self.deductions.nhis_contribution,
                "rent_relief": self.deductions.rent_relief,
                "cra": self.deductions.cra_amount,
                "other": self.deductions.other_deductions,
                "total": self.deductions.total
            },
            "taxable_income": self.taxable_income,
            "tax": {
                "computed": self.computed_tax,
                "minimum": self.minimum_tax,
                "final": self.final_tax,
                "monthly": self.monthly_tax
            },
            "rates": {
                "effective_percent": self.effective_rate,
                "marginal_percent": self.marginal_rate
            },
            "net_pay": {
                "annual": self.annual_net_pay,
                "monthly": self.monthly_net_pay
            },
            "band_breakdown": [b.to_dict() for b in self.band_calculations],
            "metadata": {
                "regime": self.regime,
                "calculation_date": self.calculation_date.isoformat(),
                "legal_citations": self.legal_citations,
                "confidence": self.validation.confidence,
                "warnings": self.validation.warnings
            }
        }


class PAYECalculator:
    """
    Comprehensive PAYE Calculator for Nigerian tax.
    
    Supports NTA 2026 and PITA 2025 regimes.
    """
    
    def __init__(self, regime_code: str = None):
        """
        Initialize calculator with a tax regime.
        
        Args:
            regime_code: Tax regime code (default: NTA_2026)
        """
        self.regime = get_regime(regime_code)
    
    def calculate(
        self,
        gross_income: float,
        period: str = "annual",
        pension_contribution: float = None,
        nhf_contribution: float = None,
        nhis_contribution: float = None,
        life_insurance: float = 0.0,
        annual_rent_paid: float = 0.0,
        other_deductions: float = 0.0,
        apply_minimum_tax: bool = True,
    ) -> PAYECalculation:
        """
        Calculate PAYE tax with full deductions.
        
        Args:
            gross_income: Gross income amount
            period: 'annual' or 'monthly'
            pension_contribution: Employee pension (default: 8% of gross)
            nhf_contribution: NHF contribution (default: 2.5% of gross)
            nhis_contribution: NHIS contribution (default: None - not mandatory)
            life_insurance: Life insurance premium paid
            annual_rent_paid: Rent paid (for NTA 2026 rent relief)
            other_deductions: Other allowable deductions
            apply_minimum_tax: Whether to apply minimum tax rule
            
        Returns:
            PAYECalculation with complete breakdown
        """
        # Normalize to annual
        if period.lower() == "monthly":
            gross_annual = gross_income * 12
        else:
            gross_annual = gross_income
        
        gross_monthly = gross_annual / 12
        
        # Calculate deductions
        deductions = self._calculate_deductions(
            gross_annual=gross_annual,
            pension_contribution=pension_contribution,
            nhf_contribution=nhf_contribution,
            nhis_contribution=nhis_contribution,
            life_insurance=life_insurance,
            annual_rent_paid=annual_rent_paid,
            other_deductions=other_deductions
        )
        
        # Calculate taxable income
        taxable_income = max(0, gross_annual - deductions.total)
        
        # Apply progressive bands
        band_calcs, computed_tax, marginal_rate = self._apply_bands(taxable_income)
        
        # Minimum tax - only applies if regime has it (NTA 2026 does NOT)
        minimum_tax = gross_annual * self.regime.minimum_tax_rate
        
        # NTA 2026 has no minimum tax rule - just use computed
        final_tax = computed_tax
        
        # Check minimum wage exemption
        annual_min_wage = self.regime.minimum_wage_monthly * 12
        if gross_annual <= annual_min_wage:
            final_tax = 0.0
        
        # Calculate rates
        effective_rate = (final_tax / gross_annual * 100) if gross_annual > 0 else 0.0
        
        # Net pay
        annual_net = gross_annual - final_tax - deductions.pension_contribution - deductions.nhf_contribution
        monthly_net = annual_net / 12
        monthly_tax = final_tax / 12
        
        # Validation
        validation = self._validate(
            gross_annual=gross_annual,
            taxable_income=taxable_income,
            final_tax=final_tax,
            deductions=deductions
        )
        
        # Legal citation - single authority, no per-line citations
        citations = [self.regime.authority]
        
        return PAYECalculation(
            gross_annual_income=gross_annual,
            gross_monthly_income=gross_monthly,
            deductions=deductions,
            taxable_income=taxable_income,
            band_calculations=band_calcs,
            computed_tax=computed_tax,
            minimum_tax=minimum_tax,
            final_tax=final_tax,
            effective_rate=effective_rate,
            marginal_rate=marginal_rate,
            annual_net_pay=annual_net,
            monthly_net_pay=monthly_net,
            monthly_tax=monthly_tax,
            regime=self.regime.name,
            calculation_date=date.today(),
            legal_citations=citations,
            validation=validation
        )
    
    def _calculate_deductions(
        self,
        gross_annual: float,
        pension_contribution: float,
        nhf_contribution: float,
        nhis_contribution: float,
        life_insurance: float,
        annual_rent_paid: float,
        other_deductions: float
    ) -> DeductionBreakdown:
        """Calculate all deductions."""
        
        # Pension: default to 8% employee contribution
        if pension_contribution is None:
            pension = gross_annual * self.regime.pension_rate
        else:
            pension = pension_contribution
        
        # NHF: default to 2.5% (mandatory for employers with 5+ staff)
        if nhf_contribution is None:
            nhf = gross_annual * self.regime.nhf_rate
        else:
            nhf = nhf_contribution
        
        # NHIS: not mandatory, only if enrolled
        nhis = nhis_contribution or 0.0
        
        # CRA (for PITA regime)
        cra = 0.0
        if self.regime.cra_enabled:
            cra_base = max(
                self.regime.cra_fixed_amount,
                gross_annual * self.regime.cra_percent_of_gross
            )
            cra = cra_base + (gross_annual * self.regime.cra_additional_percent)
        
        # Rent relief (for NTA 2026)
        rent_relief = 0.0
        if self.regime.rent_relief_enabled and annual_rent_paid > 0:
            rent_relief = min(
                self.regime.rent_relief_cap,
                annual_rent_paid * self.regime.rent_relief_percent
            )
        
        return DeductionBreakdown(
            pension_contribution=pension,
            nhf_contribution=nhf,
            nhis_contribution=nhis,
            life_insurance=life_insurance,
            rent_relief=rent_relief,
            cra_amount=cra,
            other_deductions=other_deductions
        )
    
    def _apply_bands(
        self,
        taxable_income: float
    ) -> Tuple[List[BandCalculation], float, float]:
        """Apply progressive tax bands."""
        band_calcs: List[BandCalculation] = []
        total_tax = 0.0
        remaining = taxable_income
        marginal_rate = 0.0
        
        for band in self.regime.bands:
            if remaining <= 0:
                break
            
            band_width = band.upper - band.lower
            taxable_in_band = min(remaining, band_width)
            
            if taxable_in_band > 0:
                tax_in_band = taxable_in_band * band.rate
                total_tax += tax_in_band
                marginal_rate = band.rate * 100
                
                band_calcs.append(BandCalculation(
                    band_lower=band.lower,
                    band_upper=min(band.upper, band.lower + taxable_in_band),
                    rate=band.rate,
                    taxable_in_band=taxable_in_band,
                    tax_amount=tax_in_band
                ))
                
                remaining -= taxable_in_band
        
        return band_calcs, total_tax, marginal_rate
    
    def _validate(
        self,
        gross_annual: float,
        taxable_income: float,
        final_tax: float,
        deductions: DeductionBreakdown
    ) -> ValidationResult:
        """Validate calculation for sanity."""
        errors = []
        warnings = []
        
        # Tax should never exceed income
        if final_tax > gross_annual:
            errors.append("CRITICAL: Tax exceeds gross income")
        
        # Effective rate should be reasonable
        effective_rate = (final_tax / gross_annual * 100) if gross_annual > 0 else 0
        if effective_rate > 30:
            warnings.append(f"High effective rate: {effective_rate:.1f}%")
        
        # Taxable income should not be negative
        if taxable_income < 0:
            errors.append("Taxable income is negative")
        
        # Deductions should not exceed gross
        if deductions.total > gross_annual:
            errors.append("Total deductions exceed gross income")
        
        # Pension sanity check (should be ~8%)
        expected_pension = gross_annual * 0.08
        if abs(deductions.pension_contribution - expected_pension) > expected_pension * 0.5:
            warnings.append("Pension contribution differs from standard 8%")
        
        # Calculate confidence
        confidence = 1.0
        confidence -= len(errors) * 0.3
        confidence -= len(warnings) * 0.1
        confidence = max(0.0, min(1.0, confidence))
        
        return ValidationResult(
            is_valid=len(errors) == 0,
            errors=errors,
            warnings=warnings,
            confidence=confidence
        )
    
    # ========== OUTPUT FORMATTERS ==========
    
    def format_whatsapp(self, calc: PAYECalculation) -> str:
        """Format for WhatsApp (concise, no emojis)."""
        lines = []
        
        # Header
        lines.append("*TAX CALCULATION SUMMARY*")
        lines.append("")
        
        # Key figures
        lines.append(f"Gross Income: N{calc.gross_monthly_income:,.0f}/month")
        lines.append(f"Tax Payable: N{calc.monthly_tax:,.0f}/month")
        lines.append(f"Take-Home: N{calc.monthly_net_pay:,.0f}/month")
        lines.append(f"Effective Rate: {calc.effective_rate:.1f}%")
        lines.append("")
        
        # Deductions summary
        lines.append("*Deductions Applied:*")
        if calc.deductions.pension_contribution > 0:
            lines.append(f"- Pension (8%): N{calc.deductions.pension_contribution:,.0f}")
        if calc.deductions.nhf_contribution > 0:
            lines.append(f"- NHF (2.5%): N{calc.deductions.nhf_contribution:,.0f}")
        if calc.deductions.rent_relief > 0:
            lines.append(f"- Rent Relief: N{calc.deductions.rent_relief:,.0f}")
        lines.append(f"Total Deductions: N{calc.deductions.total:,.0f}")
        lines.append("")
        
        # Tax breakdown
        lines.append("*Tax Breakdown:*")
        for band in calc.band_calculations:
            if band.tax_amount > 0:
                lines.append(
                    f"- {band.rate*100:.0f}% on N{band.taxable_in_band:,.0f} = N{band.tax_amount:,.0f}"
                )
            else:
                lines.append(f"- First N{band.taxable_in_band:,.0f}: TAX FREE")
        lines.append("")
        lines.append("_Powered by Kaanta_")
        
        return "\n".join(lines)
    
    def format_web(self, calc: PAYECalculation) -> Dict[str, Any]:
        """Format for Web (structured JSON for rendering)."""
        return {
            "summary": {
                "headline": f"You pay N{calc.monthly_tax:,.0f} monthly tax on N{calc.gross_monthly_income:,.0f} income",
                "effective_rate": f"{calc.effective_rate:.1f}%",
                "take_home": calc.monthly_net_pay,
            },
            "income": {
                "gross_monthly": calc.gross_monthly_income,
                "gross_annual": calc.gross_annual_income,
                "net_monthly": calc.monthly_net_pay,
                "net_annual": calc.annual_net_pay,
            },
            "deductions": {
                "items": [
                    {"name": "Pension (8%)", "amount": calc.deductions.pension_contribution},
                    {"name": "NHF (2.5%)", "amount": calc.deductions.nhf_contribution},
                    {"name": "Rent Relief", "amount": calc.deductions.rent_relief},
                ],
                "total": calc.deductions.total,
            },
            "tax": {
                "taxable_income": calc.taxable_income,
                "computed": calc.computed_tax,
                "minimum": calc.minimum_tax,
                "final": calc.final_tax,
                "monthly": calc.monthly_tax,
                "bands": [
                    {
                        "range": f"N{b.band_lower:,.0f} - N{b.band_upper:,.0f}",
                        "rate": b.rate * 100,
                        "amount": b.taxable_in_band,
                        "tax": b.tax_amount,
                    }
                    for b in calc.band_calculations
                ],
            },
            "rates": {
                "effective": calc.effective_rate,
                "marginal": calc.marginal_rate,
            },
            "legal": {
                "regime": calc.regime,
                "citations": calc.legal_citations,
                "date": calc.calculation_date.isoformat(),
            },
            "validation": {
                "confidence": calc.validation.confidence,
                "warnings": calc.validation.warnings,
                "is_valid": calc.validation.is_valid,
            }
        }
    
    def format_detailed(self, calc: PAYECalculation) -> str:
        """Format detailed breakdown for reports."""
        lines = []
        
        lines.append("=" * 60)
        lines.append("PERSONAL INCOME TAX CALCULATION")
        lines.append(f"Regime: {calc.regime}")
        lines.append(f"Date: {calc.calculation_date.isoformat()}")
        lines.append("=" * 60)
        lines.append("")
        
        # Income
        lines.append("INCOME")
        lines.append("-" * 40)
        lines.append(f"Gross Annual Income:     N{calc.gross_annual_income:>15,.2f}")
        lines.append(f"Gross Monthly Income:    N{calc.gross_monthly_income:>15,.2f}")
        lines.append("")
        
        # Deductions
        lines.append("DEDUCTIONS")
        lines.append("-" * 40)
        if calc.deductions.pension_contribution > 0:
            lines.append(f"Pension Contribution:    N{calc.deductions.pension_contribution:>15,.2f}")
        if calc.deductions.nhf_contribution > 0:
            lines.append(f"NHF Contribution:        N{calc.deductions.nhf_contribution:>15,.2f}")
        if calc.deductions.nhis_contribution > 0:
            lines.append(f"NHIS Contribution:       N{calc.deductions.nhis_contribution:>15,.2f}")
        if calc.deductions.rent_relief > 0:
            lines.append(f"Rent Relief:             N{calc.deductions.rent_relief:>15,.2f}")
        if calc.deductions.cra_amount > 0:
            lines.append(f"CRA:                     N{calc.deductions.cra_amount:>15,.2f}")
        lines.append("-" * 40)
        lines.append(f"TOTAL DEDUCTIONS:        N{calc.deductions.total:>15,.2f}")
        lines.append("")
        
        # Taxable income
        lines.append("TAXABLE INCOME")
        lines.append("-" * 40)
        lines.append(f"Gross Income:            N{calc.gross_annual_income:>15,.2f}")
        lines.append(f"Less: Total Deductions:  N{calc.deductions.total:>15,.2f}")
        lines.append("-" * 40)
        lines.append(f"TAXABLE INCOME:          N{calc.taxable_income:>15,.2f}")
        lines.append("")
        
        # Tax computation
        lines.append("TAX COMPUTATION (Progressive Bands)")
        lines.append("-" * 40)
        for band in calc.band_calculations:
            rate_str = f"{band.rate*100:.0f}%"
            if band.rate == 0:
                lines.append(f"N{band.band_lower:>12,.0f} - N{band.band_upper:>12,.0f}  TAX FREE")
            else:
                lines.append(
                    f"N{band.band_lower:>12,.0f} - N{band.band_upper:>12,.0f}  "
                    f"{rate_str:>5} x N{band.taxable_in_band:>12,.0f} = N{band.tax_amount:>12,.2f}"
                )
        lines.append("-" * 40)
        lines.append(f"FINAL TAX PAYABLE:       N{calc.final_tax:>15,.2f}")
        lines.append("")
        
        # Summary
        lines.append("SUMMARY")
        lines.append("-" * 40)
        lines.append(f"Monthly Tax:             N{calc.monthly_tax:>15,.2f}")
        lines.append(f"Monthly Take-Home:       N{calc.monthly_net_pay:>15,.2f}")
        lines.append(f"Effective Tax Rate:      {calc.effective_rate:>14.2f}%")
        lines.append(f"Marginal Tax Rate:       {calc.marginal_rate:>14.0f}%")
        lines.append(\"\")
        # Validation
        if calc.validation.warnings:
            lines.append("NOTES")
            lines.append("-" * 40)
            for warning in calc.validation.warnings:
                lines.append(f"* {warning}")
            lines.append("")
        
        lines.append("=" * 60)
        lines.append("Calculated by Kaanta AI")
        lines.append("=" * 60)
        
        return "\n".join(lines)


# Convenience function
def calculate_paye(
    income: float,
    period: str = "monthly",
    rent_paid: float = 0,
    regime: str = "NTA_2026"
) -> PAYECalculation:
    """
    Quick PAYE calculation.
    
    Args:
        income: Income amount
        period: 'monthly' or 'annual'
        rent_paid: Annual rent paid (for rent relief)
        regime: Tax regime code
        
    Returns:
        PAYECalculation
    """
    calc = PAYECalculator(regime)
    return calc.calculate(
        gross_income=income,
        period=period,
        annual_rent_paid=rent_paid
    )


if __name__ == "__main__":
    # Test the calculator
    print("Testing PAYE Calculator\n")
    
    calc = PAYECalculator("NTA_2026")
    
    # Test case 1: N500,000/month
    result = calc.calculate(gross_income=500_000, period="monthly")
    print(calc.format_detailed(result))
    
    print("\n" + "=" * 60 + "\n")
    print("WhatsApp Format:")
    print(calc.format_whatsapp(result))