"""
TradeSense AI — Institutional Capacity & Market Impact Engine
Implements:
1. Almgren-Chriss non-linear square-root market impact modeling.
2. Capacity sensitivity curves across Base (gamma=0.10), Adverse (gamma=0.25), and Severe (gamma=0.50) scenarios.
3. AUM scaling tiers ($100K -> $100M) and Breakeven Capacity identification.
"""

from typing import List, Dict, Any, Optional, Tuple
from datetime import date, datetime, timezone
import numpy as np
import pandas as pd
from pydantic import BaseModel, Field


class AUMTierResult(BaseModel):
    """Execution and return performance at a specific AUM scaling tier."""
    aum_usd: float
    aum_label: str
    order_size_usd: float
    avg_daily_volume_usd: float
    participation_rate_pct: float
    fixed_slippage_bps: float
    market_impact_bps: float
    total_friction_drag_bps: float
    net_return_pct: float
    net_sharpe: float
    is_institutionally_viable: bool  # True if Net Sharpe >= 1.0 and Participation <= 5%


class CapacityScenarioResult(BaseModel):
    """Capacity analysis under a specific market impact parameter scenario."""
    scenario_name: str
    gamma_impact_coeff: float
    tiers: List[AUMTierResult]
    breakeven_capacity_usd: float
    max_viable_aum_usd: float


class CapacityAnalysisReport(BaseModel):
    """Consolidated institutional capacity report."""
    base_scenario: CapacityScenarioResult
    adverse_scenario: CapacityScenarioResult
    severe_scenario: CapacityScenarioResult
    recommended_deployment_capacity_usd: float
    capacity_classification: str


class CapacityEngine:
    """
    Simulates non-linear market impact, participation rates, and AUM scalability.
    """

    AUM_TIERS = [
        (100_000, "$100K"),
        (1_000_000, "$1M"),
        (5_000_000, "$5M"),
        (10_000_000, "$10M"),
        (25_000_000, "$25M"),
        (50_000_000, "$50M"),
        (100_000_000, "$100M"),
    ]

    @classmethod
    def evaluate_aum_capacity(
        cls,
        base_annual_return_pct: float = 75.4,
        base_annual_volatility_pct: float = 16.0,
        annual_turnover: float = 8.5,
        avg_asset_adv_usd: float = 10_000_000_000.0,  # $10B ADV for mega-caps
        gamma_impact: float = 0.10,
        scenario_name: str = "Base_Impact (gamma=0.10)",
    ) -> CapacityScenarioResult:
        """
        Calculates square-root Almgren-Chriss market impact across AUM tiers.
        """
        tiers: List[AUMTierResult] = []
        max_viable = 0.0
        breakeven = 100_000_000.0

        for aum, label in cls.AUM_TIERS:
            # Rebalance size per asset (assuming 4 active tranche legs)
            order_usd = (aum / 4.0)
            part_rate = (order_usd / (avg_asset_adv_usd + 1e-9)) * 100.0

            # Almgren-Chriss impact in bps = gamma * vol * sqrt(Order / ADV) * 10,000
            impact_bps = gamma_impact * (base_annual_volatility_pct / 100.0) * np.sqrt(order_usd / avg_asset_adv_usd) * 10000.0
            fixed_slippage_bps = 5.0

            total_friction_bps = fixed_slippage_bps + impact_bps
            # Annual friction drag = (Total bps / 10,000) * Annual Turnover * 100%
            drag_pct = (total_friction_bps / 10000.0) * annual_turnover * 100.0

            net_ret = base_annual_return_pct - drag_pct
            net_sr = net_ret / (base_annual_volatility_pct + 1e-9)

            is_viable = (net_sr >= 1.0) and (part_rate <= 5.0)
            if is_viable:
                max_viable = aum
            elif breakeven == 100_000_000.0 and net_sr < 1.0:
                breakeven = aum

            tiers.append(AUMTierResult(
                aum_usd=float(aum),
                aum_label=label,
                order_size_usd=round(order_usd, 2),
                avg_daily_volume_usd=avg_asset_adv_usd,
                participation_rate_pct=round(part_rate, 4),
                fixed_slippage_bps=fixed_slippage_bps,
                market_impact_bps=round(float(impact_bps), 2),
                total_friction_drag_bps=round(float(total_friction_bps), 2),
                net_return_pct=round(float(net_ret), 2),
                net_sharpe=round(float(net_sr), 2),
                is_institutionally_viable=is_viable,
            ))

        return CapacityScenarioResult(
            scenario_name=scenario_name,
            gamma_impact_coeff=gamma_impact,
            tiers=tiers,
            breakeven_capacity_usd=breakeven,
            max_viable_aum_usd=max_viable,
        )

    @classmethod
    def run_capacity_stress_suite(
        cls,
        base_annual_return_pct: float = 75.4,
        base_annual_volatility_pct: float = 16.0,
        annual_turnover: float = 8.5,
    ) -> CapacityAnalysisReport:
        """
        Runs full sensitivity across Base (0.10), Adverse (0.25), and Severe (0.50).
        """
        base = cls.evaluate_aum_capacity(
            base_annual_return_pct, base_annual_volatility_pct, annual_turnover,
            gamma_impact=0.10, scenario_name="Base_Market_Impact (gamma=0.10)"
        )
        adverse = cls.evaluate_aum_capacity(
            base_annual_return_pct, base_annual_volatility_pct, annual_turnover,
            gamma_impact=0.25, scenario_name="Adverse_Market_Impact (gamma=0.25)"
        )
        severe = cls.evaluate_aum_capacity(
            base_annual_return_pct, base_annual_volatility_pct, annual_turnover,
            gamma_impact=0.50, scenario_name="Severe_Market_Impact (gamma=0.50)"
        )

        rec_cap = adverse.max_viable_aum_usd

        if rec_cap >= 25_000_000.0:
            classification = "HIGH_INSTITUTIONAL_CAPACITY (>= $25M)"
        elif rec_cap >= 5_000_000.0:
            classification = "MODERATE_CAPACITY ($5M - $25M)"
        else:
            classification = "NICHE_CONSTRAINED_CAPACITY (< $5M)"

        return CapacityAnalysisReport(
            base_scenario=base,
            adverse_scenario=adverse,
            severe_scenario=severe,
            recommended_deployment_capacity_usd=rec_cap,
            capacity_classification=classification,
        )