import numpy as np
from config import DEFAULT_CONFIG
from typing import Dict, Tuple, Any

try:
    from hft_simulator import HFTSimulator, LimitOrderBook
    from hft_strategies import ExecutionBridgeStrategy
    from datetime import datetime
    _HAS_HFT = True
except ImportError:
    _HAS_HFT = False

# ─────────────────────────────────────────────
# MARKET IMPACT MODEL (Almgren-Chriss)
# ─────────────────────────────────────────────
def calibrate_gamma(adv_value: float, daily_volatility: float) -> float:
    """
    Dynamically calibrates the Almgren-Chriss impact coefficient (gamma)
    based on the asset's historical volume and volatility profile.
    Highly volatile, illiquid assets receive a higher gamma (more price impact).
    """
    base_gamma = 0.10
    vol_scalar = max(0.5, min(2.5, daily_volatility / 0.015))  # Normalize around 1.5% daily vol
    liq_scalar = max(0.5, min(3.0, 50_000_000.0 / max(1.0, adv_value))) # Normalize around $50M ADV
    
    return base_gamma * vol_scalar * liq_scalar

def estimate_market_impact(
    trade_value: float, 
    adv_value: float, 
    daily_volatility: float, 
    gamma: float = None,
    max_participation: float = 0.15
) -> float:
    """
    Square-root market impact model based on Almgren-Chriss (2001).
    
    Estimates the price impact (in decimal percentage) of a trade based on its size 
    relative to the Average Daily Volume (ADV) and the asset's inherent volatility.
    
    Formula: ΔP/P = γ * σ * √(Q/V)
    
    Args:
        trade_value: The total dollar size of the intended trade (Q).
        adv_value: The Average Daily Volume in dollars (V).
        daily_volatility: The daily standard deviation of the asset's returns (σ).
        gamma: Scaling constant (typically 0.1 - 0.3 for daily impact models).
        max_participation: The maximum % of ADV we are willing to be. If our trade
                           exceeds this, we assume the rest is unfilled.
                           
    Returns:
        float: The estimated price impact as a decimal (e.g., 0.0032 = 32 bps).
    """
    if adv_value <= 0 or trade_value == 0:
        return 0.0
        
    participation_rate = abs(trade_value) / adv_value
    
    if gamma is None:
        gamma = calibrate_gamma(adv_value, daily_volatility)
        
    # Impact = γ * σ * √(Q / V)
    # We do NOT cap participation for the penalty calculation. Capping it mathematically mutes 
    # the marginal penalty for huge trades, incentivizing the optimizer to trade infinitely.
    impact_decimal = gamma * daily_volatility * np.sqrt(participation_rate)
    
    return float(impact_decimal)


def estimate_liquidity_penalty_rate(
    weight_delta: float, 
    capital: float, 
    adv_value: float, 
    daily_volatility: float, 
    gamma: float = None
) -> float:
    """
    Helper function for structural optimizer integration.
    Calculates the exact fractional rate to penalize the objective function.
    
    Args:
        weight_delta: The proposed change in portfolio weight (Δw).
        capital: Total portfolio capital.
        adv_value: Average Daily Volume of the asset.
        daily_volatility: Asset's daily volatility.
        gamma: Almgren-Chriss impact scaling factor.
        
    Returns:
        float: The penalty coefficient to apply inside a convex solver.
    """
    trade_value = abs(weight_delta * capital)
    if adv_value <= 0 or trade_value == 0:
        return 0.0
        
    if gamma is None:
        gamma = calibrate_gamma(adv_value, daily_volatility)
        
    # Rate = γ * σ * √( (Δw * Cap) / ADV )
    return float(gamma * daily_volatility * np.sqrt(trade_value / adv_value))


# ─────────────────────────────────────────────
# EXECUTION SIMULATOR (Intraday Proxy & HFT)
# ─────────────────────────────────────────────

def simulate_execution_hft(
    target_trades: Dict[str, float], 
    market_data: Dict[str, dict], 
    latency_ms: int = 10, 
    duration_ms: int = 1000
) -> Dict[str, dict]:
    """
    Run execution through HFT simulator for microsecond-level slippage analysis.
    Falls back to standard simulation if HFT module is unavailable.
    """
    if not _HAS_HFT or not target_trades:
        return simulate_execution(target_trades, market_data)

    symbols = list(target_trades.keys())
    start_time = datetime.now()
    duration_sec = duration_ms / 1000.0
    
    sim = HFTSimulator(symbols, start_time, duration_sec, tick_interval_ms=1, latency_ms=latency_ms)
    
    # Initialize books with current market prices
    initial_prices = {sym: market_data.get(sym, {}).get("price", 100.0) for sym in symbols}
    sim.initialize_books(initial_prices)
    
    # Setup execution strategies for each trade
    bridges = {}
    for sym, qty in target_trades.items():
        if qty == 0: continue
        side = 'buy' if qty > 0 else 'sell'
        bridge = ExecutionBridgeStrategy(sym, abs(qty), side, chunks=20)
        bridges[sym] = bridge
        sim.add_strategy(bridge)
        
    # Run the simulation
    results = sim.run()
    
    # Compile the results per asset
    exec_results = {}
    for sym, qty in target_trades.items():
        if qty == 0: continue
        
        bridge = bridges[sym]
        initial_price = initial_prices[sym]
        
        # Determine actual executed quantity and average price
        sym_trades = [t for t in results.trades if t.symbol == sym and (t.buyer_order_id >= 5000000 or t.seller_order_id >= 5000000)]
        
        executed_qty = sum(t.quantity for t in sym_trades)
        if executed_qty > 0:
            avg_price = sum(t.price * t.quantity for t in sym_trades) / executed_qty
            slippage = abs(avg_price - initial_price) / initial_price
        else:
            avg_price = initial_price
            slippage = 0.005 # Default penalty for unexecuted 
            
        exec_results[sym] = {
            "target_qty": qty,
            "executed_qty": executed_qty if qty > 0 else -executed_qty,
            "avg_price": avg_price,
            "slippage_bps": slippage * 10000,
            "fill_rate": executed_qty / abs(qty) if abs(qty) > 0 else 1.0
        }
        
    return exec_results

def simulate_execution(
    target_trades: Dict[str, float], 
    market_data: Dict[str, dict], 
    execution_style: str = 'close',
    max_participation: float = 0.10,
    adv_proxy: float = 50_000_000.0
) -> Tuple[Dict[str, dict], float]:
    """
    Simulates realistic trade execution, accounting for intraday slippage, 
    market impact, and liquidity constraints.
    
    Args:
        target_trades: Dict of {ticker: trade_value_in_dollars}. Positive = Buy, Negative = Sell.
        market_data: Dict containing daily metrics for each ticker:
                     {'open', 'high', 'low', 'close', 'adv', 'volatility'}
        execution_style: 'close', 'twap', 'vwap', or 'aggressive'.
        max_participation: Hard limit on % of ADV we can trade in one day.
        
    Returns:
        Tuple containing:
        - Dict of filled trade details {ticker: {'filled_value', 'execution_price', 'slippage_bps'}}
        - Float of total unfilled dollar value (if liquidity was constrained)
    """
    execution_results = {}
    total_unfilled = 0.0
    
    for ticker, trade_val in target_trades.items():
        if abs(trade_val) < 1e-4:
            continue
            
        data = market_data.get(ticker, {})
        close_px = data.get('close', 1.0)
        open_px = data.get('open', close_px)
        high_px = data.get('high', close_px)
        low_px = data.get('low', close_px)
        adv = data.get('adv', adv_proxy)      # Use parameterized ADV proxy fallback
        vol = data.get('volatility', 0.015)      # Default 1.5% daily vol fallback
        
        # 1. Determine Base Execution Price (Intraday Proxy)
        # Because we operate on daily data, we proxy intraday execution algos
        if execution_style == 'close':
            base_px = close_px
        elif execution_style == 'twap':
            # Time-Weighted Average Price proxy
            base_px = (open_px + high_px + low_px + close_px) / 4.0
        elif execution_style == 'vwap':
            # Volume-Weighted Average Price proxy
            base_px = (open_px + high_px + low_px + close_px) / 4.0
        elif execution_style == 'aggressive':
            # Simulates paying the spread and crossing the book immediately near the open
            base_px = open_px * (1.0 + (vol * 0.5) if trade_val > 0 else 1.0 - (vol * 0.5))
        else:
            base_px = close_px
            
        # 2. Liquidity Constraint (Unfilled Quantity)
        max_trade_size = adv * max_participation
        if abs(trade_val) > max_trade_size:
            filled_val = max_trade_size * np.sign(trade_val)
            unfilled = abs(trade_val) - max_trade_size
            total_unfilled += unfilled
        else:
            filled_val = trade_val
            
        # 3. Market Impact Calculation
        impact_decimal = estimate_market_impact(
            trade_value=filled_val,
            adv_value=adv,
            daily_volatility=vol,
            max_participation=max_participation
        )
        
        # 4. Final Execution Price
        # If buying, impact pushes price UP. If selling, impact pushes price DOWN.
        direction = 1.0 if filled_val > 0 else -1.0
        final_px = base_px * (1.0 + (impact_decimal * direction))
        
        # Calculate absolute slippage in Basis Points relative to the ideal close price
        slippage_bps = abs((final_px - close_px) / close_px) * 10000.0
        
        execution_results[ticker] = {
            'intended_val': trade_val,
            'filled_val': filled_val,
            'base_price': base_px,
            'execution_price': final_px,
            'slippage_bps': slippage_bps,
            'impact_decimal': impact_decimal
        }
        
    return execution_results, total_unfilled

# ─────────────────────────────────────────────
# LIVE EXECUTION ENGINE (Interactive Brokers)
# ─────────────────────────────────────────────
import math
import logging
import asyncio
import pandas as pd
from dataclasses import dataclass
from typing import List, Optional
from config import DEFAULT_ADV

logger = logging.getLogger(__name__)

try:
    from ib_async import IB, Stock, MarketOrder, LimitOrder, Order
    _IB_AVAILABLE = True
except ImportError:
    _IB_AVAILABLE = False
    logger.warning("ib_async not installed. Live execution disabled.")

@dataclass
class GeneratedOrder:
    ticker: str
    action: str  # 'BUY' or 'SELL'
    quantity: int
    algo_strategy: str
    limit_price: float = 0.0

class OrderGenerator:
    """
    Translates mathematical target weights into integer share orders,
    implementing Tiered Execution based on % of Average Daily Volume.
    """
    def __init__(self, adv_data: Dict[str, float]):
        self.adv_data = adv_data

    def _determine_algo(self, trade_dollars: float, adv_dollars: float) -> str:
        """
        Tiered Execution mapping based on % of ADV.
        MOC < 1%, VWAP < 5%, TWAP < 15%, IB_ARRIVAL > 15%
        """
        if adv_dollars <= 0:
            return "VWAP" # Fallback
            
        participation = trade_dollars / adv_dollars
        if participation > 0.50:
            return "IB_ARRIVAL"  # Large orders need arrival-price algo, not passive VWAP
        elif participation < 0.01:
            return "MOC"
        elif participation < 0.05:
            return "VWAP"
        elif participation < 0.15:
            return "TWAP"
        else:
            return "IB_ARRIVAL"

    def generate_orders(
        self, 
        target_weights: pd.Series, 
        current_positions: Dict[str, int], 
        prices: Dict[str, float], 
        capital: float,
        adv_proxy: float = 50_000_000.0
    ) -> List[GeneratedOrder]:
        orders = []
        
        for ticker, tgt_weight in target_weights.items():
            if ticker == 'CASH':
                continue
                
            price = prices.get(ticker)
            if not price or price <= 0:
                continue
                
            target_dollars = tgt_weight * capital
            target_shares = int(math.floor(target_dollars / price))
            current_shares = current_positions.get(ticker, 0)
            
            diff_shares = target_shares - current_shares
            if diff_shares == 0:
                continue
                
            action = "BUY" if diff_shares > 0 else "SELL"
            qty = abs(diff_shares)
            trade_dollars = qty * price
            
            adv = self.adv_data.get(ticker, adv_proxy)
            algo = self._determine_algo(trade_dollars, adv)
            
            orders.append(GeneratedOrder(
                ticker=ticker,
                action=action,
                quantity=qty,
                algo_strategy=algo
            ))
            
        return orders

    def export_orders_to_csv(self, orders: List[GeneratedOrder], filepath: str = None) -> str:
        """
        Exports a list of GeneratedOrders to a CSV file compatible with IBKR BasketTrader.
        """
        import csv
        import datetime
        import os
        
        if not filepath:
            date_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            filepath = f"order_blotter_{date_str}.csv"
            
        with open(filepath, mode='w', newline='') as file:
            writer = csv.writer(file)
            writer.writerow(['Action', 'Quantity', 'Symbol', 'SecType', 'Exchange', 'Currency', 'TimeInForce', 'OrderType', 'AlgoStrategy', 'LmtPrice'])
            
            for o in orders:
                order_type = 'LMT' if o.limit_price > 0 else 'MKT'
                writer.writerow([
                    o.action,
                    o.quantity,
                    o.ticker,
                    'STK',
                    'SMART',
                    'USD',
                    'DAY',
                    order_type,
                    o.algo_strategy,
                    o.limit_price if o.limit_price > 0 else ''
                ])
                
        logger.info(f"Exported {len(orders)} orders to {filepath}")
        return filepath

# ─────────────────────────────────────────────
# REBALANCING SCHEDULER & DRIFT MONITOR
# ─────────────────────────────────────────────
class RebalanceScheduler:
    def __init__(self, frequency: str = 'monthly', drift_threshold: float = 0.05):
        """
        frequency: 'daily', 'weekly', 'monthly', 'quarterly', 'drift_only'
        drift_threshold: absolute weight deviation that triggers an off-cycle rebalance
        """
        self.frequency = frequency
        self.drift_threshold = drift_threshold
        
    def is_rebalance_due(self, current_date: pd.Timestamp, last_rebalance_date: pd.Timestamp, target_w: pd.Series, current_w: pd.Series) -> bool:
        if last_rebalance_date is None:
            return True
            
        # 1. Check Drift Trigger
        if current_w is not None and len(current_w) > 0:
            aligned_target, aligned_current = target_w.align(current_w, fill_value=0.0)
            drift = (aligned_target - aligned_current).abs()
            if drift.max() >= self.drift_threshold:
                logger.info(f"Drift triggered rebalance! Max drift: {drift.max():.2%} >= {self.drift_threshold:.2%}")
                return True
                
        # 2. Check Time Schedule
        if self.frequency == 'drift_only':
            return False
            
        days_since = (current_date - last_rebalance_date).days
        if self.frequency == 'daily' and days_since >= 1:
            return True
        if self.frequency == 'weekly' and days_since >= 7:
            return True
        if self.frequency == 'monthly' and days_since >= 28:
            return True
        if self.frequency == 'quarterly' and days_since >= 90:
            return True
            
        return False

class BrokerBridge:
    """
    Interface to Interactive Brokers via ib_async.
    Handles position reconciliation and human-in-the-loop CLI approval.
    """
    def __init__(self, host: str = '127.0.0.1', port: int = 7497, client_id: int = 1):
        self.host = host
        self.port = port
        self.client_id = client_id
        self.ib = None
        if _IB_AVAILABLE:
            self.ib = IB()
            
    async def connect_async(self):
        if not _IB_AVAILABLE:
            raise RuntimeError("ib_async is not installed. Cannot connect to IBKR.")
        if not self.ib.isConnected():
            await self.ib.connectAsync(self.host, self.port, clientId=self.client_id)
            
    def disconnect(self):
        if self.ib and self.ib.isConnected():
            self.ib.disconnect()

    async def get_positions_async(self) -> Dict[str, int]:
        await self.connect_async()
        # Wait a brief moment to ensure positions are synced from IBKR servers
        await asyncio.sleep(1)
        pos_dict = {}
        for pos in self.ib.positions():
            if pos.contract.secType == 'STK':
                pos_dict[pos.contract.symbol] = int(pos.position)
        return pos_dict

    async def reconcile_positions_async(self, db_positions: Dict[str, int]) -> dict:
        """
        Reconcile internal DB positions against true broker positions.
        Returns a dict of mismatches, empty if perfectly reconciled.
        """
        broker_positions = await self.get_positions_async()
        mismatches = {}
        
        all_tickers = set(db_positions.keys()).union(set(broker_positions.keys()))
        for t in all_tickers:
            db_qty = db_positions.get(t, 0)
            br_qty = broker_positions.get(t, 0)
            if db_qty != br_qty:
                logger.error(f"Reconciliation failure for {t}: DB={db_qty}, Broker={br_qty}")
                mismatches[t] = {'db': db_qty, 'broker': br_qty, 'diff': db_qty - br_qty}
                
        return mismatches

    def build_ib_order(self, gen_order: GeneratedOrder) -> Order:
        """Translates our GeneratedOrder into an IB Order with Algorithmic parameters."""
        contract = Stock(gen_order.ticker, 'SMART', 'USD')
        
        if gen_order.algo_strategy == "MOC":
            order = MarketOrder(gen_order.action, gen_order.quantity)
            order.tif = 'AUC' # Market On Close
        elif gen_order.algo_strategy == "VWAP":
            order = MarketOrder(gen_order.action, gen_order.quantity)
            order.algoStrategy = 'Vwap'
            order.algoParams = []
        elif gen_order.algo_strategy == "TWAP":
            order = MarketOrder(gen_order.action, gen_order.quantity)
            order.algoStrategy = 'Twap'
            order.algoParams = []
        elif gen_order.algo_strategy == "IB_ARRIVAL":
            order = MarketOrder(gen_order.action, gen_order.quantity)
            order.algoStrategy = 'ArrivalPx'
            order.algoParams = []
        else:
            order = MarketOrder(gen_order.action, gen_order.quantity)
            
        return contract, order

    async def execute_orders_async(self, orders: List[GeneratedOrder], require_approval: bool = True):
        await self.connect_async()
        
        if not orders:
            logger.info("No orders to execute.")
            return

        print("\n--- PROPOSED TRADES ---")
        for o in orders:
            print(f" {o.action} {o.quantity} shares of {o.ticker} [Algo: {o.algo_strategy}]")
            
        if require_approval:
            # Wrap the blocking input() in a thread so the async event loop isn't starved
            ans = await asyncio.to_thread(input, "\nApprove execution? (y/n): ")
            if ans.lower().strip() != 'y':
                logger.info("Execution aborted by human.")
                return
                
        for o in orders:
            contract, ib_order = self.build_ib_order(o)
            await self.ib.qualifyContractsAsync(contract)
            trade = self.ib.placeOrder(contract, ib_order)
            logger.info(f"Placed {o.action} order for {o.ticker}")

def execute_emergency_liquidation(state):
    """
    Rapidly constructs sell market orders for all active positions, bypassing 
    standard transition/impact scheduling to protect capital.
    """
    logger.critical("EMERGENCY LIQUIDATION INITIATED!")
    
    positions = state.current_weights if state and hasattr(state, 'current_weights') else {}
    if hasattr(positions, "items"):
        for ticker, weight in positions.items():
            if ticker != 'CASH' and float(weight) > 0:
                logger.critical(f"EMERGENCY SELL MKT: {ticker} (Weight: {float(weight):.2%})")
                
    return True

# ─────────────────────────────────────────────
# CRYPTO ARBITRAGE EXECUTION
# ─────────────────────────────────────────────

def execute_crypto_arbitrage(opportunity: Dict[str, Any], capital: float = 10000.0) -> Dict[str, Any]:
    """
    Simulates or executes a cross-exchange crypto arbitrage trade.
    """
    try:
        from crypto_arb import simulate_arbitrage_trade
        # In a live environment, this would initialize ccxt exchange instances 
        # with API keys and execute market/limit orders simultaneously.
        # For now, we simulate the execution to calculate latency-adjusted net profit.
        
        result = simulate_arbitrage_trade(opportunity, capital, latency_ms=50.0)
        logger.info(f"Crypto Arb Execution: {opportunity['buy_exchange']} -> {opportunity['sell_exchange']} | Status: {result['status']} | Net: ${result['net_profit_usd']:.2f}")
        return result
    except ImportError:
        logger.error("crypto_arb module not found.")
        return {"status": "ERROR", "message": "crypto_arb module not found."}