backend/tax/optimizer.py

"""Tax optimisation strategies for portfolio operations.

This module implements:
- Tax-loss harvesting opportunity identification
- Tax-optimised rebalancing algorithms
- Qualified dividend holding period tracking
- Lot selection optimisation
"""

from datetime import date, timedelta
from decimal import Decimal
from typing import List, Dict, Optional, Tuple
from collections import defaultdict
import logging

from backend.tax.models import (
    TaxLot,
    Transaction,
    TaxLossHarvestingOpportunity,
    TaxOptimizedSale,
    CostBasisMethod,
    TransactionType,
)
from backend.tax.calculator import TaxCalculator

logger = logging.getLogger(__name__)


class TaxOptimizer:
    """Optimise portfolio operations for tax efficiency."""

    def __init__(self, tax_calculator: TaxCalculator):
        """Initialise tax optimiser.

        Args:
            tax_calculator: TaxCalculator instance for tax calculations
        """
        self.tax_calculator = tax_calculator

    def identify_tax_loss_harvesting_opportunities(
        self,
        current_prices: Dict[str, Decimal],
        min_loss_threshold: Decimal = Decimal("100"),
        min_loss_percentage: Decimal = Decimal("0.05"),
    ) -> List[TaxLossHarvestingOpportunity]:
        """Identify tax-loss harvesting opportunities.

        Tax-loss harvesting involves selling securities at a loss to offset capital gains
        and reduce tax liability.

        Args:
            current_prices: Dictionary of ticker to current price
            min_loss_threshold: Minimum dollar loss to consider (default $100)
            min_loss_percentage: Minimum loss percentage to consider (default 5%)

        Returns:
            List of harvesting opportunities
        """
        opportunities = []

        for ticker, lots in self.tax_calculator.tax_lots.items():
            if ticker not in current_prices:
                continue

            current_price = current_prices[ticker]

            # Calculate unrealised loss for each lot
            for lot in lots:
                unrealised_gain_loss = (current_price - lot.cost_basis_per_share) * lot.quantity

                # Only interested in losses
                if unrealised_gain_loss >= 0:
                    continue

                loss_percentage = (
                    (current_price - lot.cost_basis_per_share) / lot.cost_basis_per_share
                )

                # Check if loss meets thresholds
                if (
                    abs(unrealised_gain_loss) < min_loss_threshold
                    or abs(loss_percentage) < min_loss_percentage
                ):
                    continue

                # Estimate tax savings
                tax_rate = self.tax_calculator.get_capital_gains_rate(
                    is_long_term=lot.is_long_term
                )
                estimated_savings = abs(unrealised_gain_loss) * tax_rate

                # Determine wash sale risk (simplified - would need transaction history)
                wash_sale_risk = False  # Would check recent purchases

                # Generate recommendation
                if lot.is_long_term:
                    recommendation = (
                        f"Sell {lot.quantity} shares at ${current_price:.2f} "
                        f"to harvest ${abs(unrealised_gain_loss):,.2f} long-term loss. "
                        f"Estimated tax savings: ${estimated_savings:,.2f}. "
                        f"Consider reinvesting in similar but not identical security to maintain exposure."
                    )
                else:
                    recommendation = (
                        f"Sell {lot.quantity} shares at ${current_price:.2f} "
                        f"to harvest ${abs(unrealised_gain_loss):,.2f} short-term loss. "
                        f"Estimated tax savings: ${estimated_savings:,.2f}. "
                        f"Wait 31 days before repurchasing to avoid wash sale."
                    )

                opportunity = TaxLossHarvestingOpportunity(
                    ticker=ticker,
                    current_price=current_price,
                    unrealised_loss=unrealised_gain_loss,
                    loss_percentage=loss_percentage,
                    quantity=lot.quantity,
                    holding_period_days=lot.holding_period_days,
                    is_long_term=lot.is_long_term,
                    estimated_tax_savings=estimated_savings,
                    wash_sale_risk=wash_sale_risk,
                    recommended_action=recommendation,
                    lot_ids=[lot.lot_id],
                )

                opportunities.append(opportunity)

        # Sort by tax savings (highest first)
        opportunities.sort(key=lambda x: x.estimated_tax_savings, reverse=True)

        return opportunities

    def optimize_lot_selection(
        self,
        ticker: str,
        quantity_to_sell: Decimal,
        current_price: Decimal,
        goal: str = "minimize_tax",
    ) -> TaxOptimizedSale:
        """Optimise which lots to sell for tax efficiency.

        Args:
            ticker: Stock ticker symbol
            quantity_to_sell: Number of shares to sell
            current_price: Current market price
            goal: Optimisation goal - 'minimize_tax', 'harvest_losses', or 'maximize_long_term'

        Returns:
            TaxOptimizedSale with recommended lots
        """
        lots = self.tax_calculator.tax_lots.get(ticker, [])
        if not lots:
            raise ValueError(f"No lots available for {ticker}")

        available_lots = [lot for lot in lots]
        selected_lots = []
        remaining_qty = quantity_to_sell

        if goal == "minimize_tax":
            # Sell lots with highest cost basis first (HIFO) to minimize gains
            available_lots.sort(key=lambda x: x.cost_basis_per_share, reverse=True)
            method = CostBasisMethod.HIFO
            rationale = (
                "Selected lots with highest cost basis (HIFO) to minimise capital gains tax. "
                "This strategy reduces immediate tax liability by selling shares with the least appreciation first."
            )

        elif goal == "harvest_losses":
            # Sell lots at a loss, prioritising largest losses
            loss_lots = [
                lot
                for lot in available_lots
                if (current_price - lot.cost_basis_per_share) < 0
            ]
            loss_lots.sort(
                key=lambda x: (current_price - x.cost_basis_per_share) * x.quantity
            )
            available_lots = loss_lots
            method = CostBasisMethod.SPECIFIC_ID
            rationale = (
                "Selected lots showing losses to harvest tax losses. "
                "Losses can offset capital gains and reduce overall tax liability."
            )

        elif goal == "maximize_long_term":
            # Prioritise long-term lots to get preferential tax rate
            long_term_lots = [lot for lot in available_lots if lot.is_long_term]
            long_term_lots.sort(key=lambda x: x.acquisition_date)
            available_lots = long_term_lots + [
                lot for lot in available_lots if not lot.is_long_term
            ]
            method = CostBasisMethod.SPECIFIC_ID
            rationale = (
                "Prioritised long-term holdings (>365 days) to benefit from lower long-term "
                "capital gains tax rates (0%, 15%, or 20% vs ordinary income rates up to 37%)."
            )

        else:
            raise ValueError(f"Unknown optimisation goal: {goal}")

        # Select lots
        for lot in available_lots:
            if remaining_qty <= 0:
                break

            qty_from_lot = min(remaining_qty, lot.quantity)

            # Create a copy of the lot with adjusted quantity
            selected_lot = TaxLot(
                lot_id=lot.lot_id,
                ticker=lot.ticker,
                acquisition_date=lot.acquisition_date,
                quantity=qty_from_lot,
                original_quantity=lot.original_quantity,
                cost_basis_per_share=lot.cost_basis_per_share,
                total_cost_basis=qty_from_lot * lot.cost_basis_per_share,
            )
            selected_lots.append(selected_lot)
            remaining_qty -= qty_from_lot

        if not selected_lots:
            raise ValueError(f"No suitable lots found for goal: {goal}")

        # Calculate estimated gain/loss and tax
        total_proceeds = quantity_to_sell * current_price
        total_cost_basis = sum(lot.total_cost_basis for lot in selected_lots)
        estimated_gain_loss = total_proceeds - total_cost_basis

        # Determine if majority is long-term
        long_term_qty = sum(
            lot.quantity for lot in selected_lots if lot.is_long_term
        )
        is_long_term = long_term_qty > (quantity_to_sell / 2)

        # Estimate tax
        tax_rate = self.tax_calculator.get_capital_gains_rate(is_long_term=is_long_term)
        estimated_tax = max(Decimal("0"), estimated_gain_loss * tax_rate)

        return TaxOptimizedSale(
            ticker=ticker,
            quantity_to_sell=quantity_to_sell,
            lots_to_sell=selected_lots,
            cost_basis_method=method,
            estimated_gain_loss=estimated_gain_loss,
            estimated_tax=estimated_tax,
            is_long_term=is_long_term,
            rationale=rationale,
        )

    def suggest_rebalancing_strategy(
        self,
        target_allocations: Dict[str, Decimal],
        current_prices: Dict[str, Decimal],
        min_trade_amount: Decimal = Decimal("100"),
    ) -> List[Dict[str, any]]:
        """Suggest tax-optimised rebalancing trades.

        Args:
            target_allocations: Target allocation percentages by ticker
            current_prices: Current market prices by ticker
            min_trade_amount: Minimum trade amount to consider (default $100)

        Returns:
            List of suggested trades with tax implications
        """
        suggestions = []

        # Calculate current allocations
        total_value = Decimal("0")
        current_values = {}

        for ticker, lots in self.tax_calculator.tax_lots.items():
            if ticker not in current_prices:
                continue

            current_price = current_prices[ticker]
            ticker_value = sum(lot.quantity * current_price for lot in lots)
            current_values[ticker] = ticker_value
            total_value += ticker_value

        if total_value == 0:
            return suggestions

        # Calculate required trades
        for ticker, target_pct in target_allocations.items():
            current_value = current_values.get(ticker, Decimal("0"))
            target_value = total_value * target_pct
            difference = target_value - current_value

            if abs(difference) < min_trade_amount:
                continue

            if difference < 0:
                # Need to sell
                quantity_to_sell = abs(difference) / current_prices[ticker]

                # Optimise lot selection to minimise tax
                try:
                    optimised_sale = self.optimize_lot_selection(
                        ticker=ticker,
                        quantity_to_sell=quantity_to_sell,
                        current_price=current_prices[ticker],
                        goal="minimize_tax",
                    )

                    suggestions.append(
                        {
                            "action": "SELL",
                            "ticker": ticker,
                            "quantity": quantity_to_sell,
                            "estimated_proceeds": abs(difference),
                            "estimated_gain_loss": optimised_sale.estimated_gain_loss,
                            "estimated_tax": optimised_sale.estimated_tax,
                            "lots": optimised_sale.lots_to_sell,
                            "rationale": optimised_sale.rationale,
                        }
                    )
                except ValueError as e:
                    logger.warning(f"Could not optimise sale for {ticker}: {e}")

            else:
                # Need to buy
                quantity_to_buy = difference / current_prices[ticker]
                suggestions.append(
                    {
                        "action": "BUY",
                        "ticker": ticker,
                        "quantity": quantity_to_buy,
                        "estimated_cost": difference,
                        "estimated_gain_loss": Decimal("0"),
                        "estimated_tax": Decimal("0"),
                        "rationale": f"Purchase to reach target allocation of {target_pct * 100:.1f}%",
                    }
                )

        return suggestions

    def check_qualified_dividend_holding_period(
        self,
        ticker: str,
        dividend_ex_date: date,
        min_holding_days: int = 60,
    ) -> Tuple[bool, int]:
        """Check if holdings qualify for qualified dividend treatment.

        Qualified dividends are taxed at lower long-term capital gains rates.
        To qualify, stock must be held for >60 days during the 121-day period
        beginning 60 days before the ex-dividend date.

        Args:
            ticker: Stock ticker symbol
            dividend_ex_date: Ex-dividend date
            min_holding_days: Minimum holding days required (default 60)

        Returns:
            Tuple of (qualifies, actual holding days)
        """
        lots = self.tax_calculator.tax_lots.get(ticker, [])
        if not lots:
            return False, 0

        # Calculate 121-day period
        period_start = dividend_ex_date - timedelta(days=60)
        period_end = dividend_ex_date + timedelta(days=60)

        # Check if any lot qualifies
        max_holding_days = 0
        for lot in lots:
            # Calculate days held during the period
            hold_start = max(lot.acquisition_date, period_start)
            hold_end = min(date.today(), period_end)

            if hold_start <= hold_end:
                days_held = (hold_end - hold_start).days
                max_holding_days = max(max_holding_days, days_held)

        qualifies = max_holding_days > min_holding_days

        return qualifies, max_holding_days

    def compare_cost_basis_methods(
        self,
        ticker: str,
        quantity_to_sell: Decimal,
        sale_price: Decimal,
    ) -> Dict[CostBasisMethod, Dict[str, Decimal]]:
        """Compare tax implications across different cost basis methods.

        Args:
            ticker: Stock ticker symbol
            quantity_to_sell: Quantity to sell
            sale_price: Sale price per share

        Returns:
            Dictionary mapping cost basis method to tax implications
        """
        results = {}

        # Save current state
        original_lots = {
            t: [
                TaxLot(
                    lot_id=lot.lot_id,
                    ticker=lot.ticker,
                    acquisition_date=lot.acquisition_date,
                    quantity=lot.quantity,
                    original_quantity=lot.original_quantity,
                    cost_basis_per_share=lot.cost_basis_per_share,
                    total_cost_basis=lot.total_cost_basis,
                )
                for lot in lots
            ]
            for t, lots in self.tax_calculator.tax_lots.items()
        }

        methods = [
            CostBasisMethod.FIFO,
            CostBasisMethod.LIFO,
            CostBasisMethod.HIFO,
            CostBasisMethod.AVERAGE,
        ]

        for method in methods:
            # Restore state
            self.tax_calculator.tax_lots = {
                t: [
                    TaxLot(
                        lot_id=lot.lot_id,
                        ticker=lot.ticker,
                        acquisition_date=lot.acquisition_date,
                        quantity=lot.quantity,
                        original_quantity=lot.original_quantity,
                        cost_basis_per_share=lot.cost_basis_per_share,
                        total_cost_basis=lot.total_cost_basis,
                    )
                    for lot in lots
                ]
                for t, lots in original_lots.items()
            }

            try:
                if method == CostBasisMethod.FIFO:
                    gains = self.tax_calculator.calculate_sale_fifo(
                        ticker, date.today(), quantity_to_sell, sale_price
                    )
                elif method == CostBasisMethod.LIFO:
                    gains = self.tax_calculator.calculate_sale_lifo(
                        ticker, date.today(), quantity_to_sell, sale_price
                    )
                elif method == CostBasisMethod.HIFO:
                    gains = self.tax_calculator.calculate_sale_hifo(
                        ticker, date.today(), quantity_to_sell, sale_price
                    )
                elif method == CostBasisMethod.AVERAGE:
                    gains = self.tax_calculator.calculate_sale_average(
                        ticker, date.today(), quantity_to_sell, sale_price
                    )

                total_gain_loss = sum(g.gain_loss for g in gains)
                tax_liability, _ = self.tax_calculator.calculate_tax_liability(gains)

                results[method] = {
                    "total_gain_loss": total_gain_loss,
                    "tax_liability": tax_liability,
                    "long_term_portion": sum(
                        g.gain_loss for g in gains if g.is_long_term
                    ),
                    "short_term_portion": sum(
                        g.gain_loss for g in gains if not g.is_long_term
                    ),
                }

            except Exception as e:
                logger.error(f"Error calculating with {method}: {e}")
                results[method] = {
                    "error": str(e),
                }

        # Restore original state
        self.tax_calculator.tax_lots = original_lots

        return results

    def generate_tax_optimization_report(
        self,
        current_prices: Dict[str, Decimal],
    ) -> Dict[str, any]:
        """Generate comprehensive tax optimisation report.

        Args:
            current_prices: Current market prices by ticker

        Returns:
            Dictionary with comprehensive tax optimisation recommendations
        """
        report = {
            "harvesting_opportunities": [],
            "qualified_dividends": {},
            "unrealised_positions": {},
            "recommendations": [],
        }

        # Identify harvesting opportunities
        opportunities = self.identify_tax_loss_harvesting_opportunities(current_prices)
        report["harvesting_opportunities"] = [
            {
                "ticker": opp.ticker,
                "unrealised_loss": float(opp.unrealised_loss),
                "estimated_tax_savings": float(opp.estimated_tax_savings),
                "recommendation": opp.recommended_action,
            }
            for opp in opportunities[:5]  # Top 5
        ]

        # Calculate unrealised gains/losses
        total_unrealised_gains = Decimal("0")
        total_unrealised_losses = Decimal("0")

        for ticker, price in current_prices.items():
            gain, qty = self.tax_calculator.get_unrealised_gains(ticker, price)
            if gain > 0:
                total_unrealised_gains += gain
            else:
                total_unrealised_losses += abs(gain)

            report["unrealised_positions"][ticker] = {
                "quantity": float(qty),
                "unrealised_gain_loss": float(gain),
                "current_value": float(qty * price),
            }

        # Generate recommendations
        if opportunities:
            total_potential_savings = sum(opp.estimated_tax_savings for opp in opportunities)
            report["recommendations"].append(
                f"Tax-loss harvesting opportunities identified with potential savings of "
                f"${total_potential_savings:,.2f}. Consider harvesting losses before year-end."
            )

        if total_unrealised_gains > total_unrealised_losses:
            report["recommendations"].append(
                f"Portfolio has net unrealised gains of ${total_unrealised_gains - total_unrealised_losses:,.2f}. "
                f"Consider tax-loss harvesting to offset future gains."
            )

        report["recommendations"].append(
            "IMPORTANT: This analysis is for informational purposes only. "
            "Consult a qualified tax professional before making any investment decisions."
        )

        return report