"""
Diversification limit manager for MoERegression.

Prevents over-diversification by limiting number of simultaneous trades.
This reduces correlation risk and improves portfolio quality.
"""

import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import Dict, List, Optional


class DiversificationManager:
    """Manage diversification limits for trading positions."""

    def __init__(self, config: Dict = None):
        """Initialize diversification parameters."""
        self.config = config or self.get_default_config()
        self.open_positions: Dict[int, Dict] = {}  # {trade_id: position_data}
        self.entry_time = {}

    def get_default_config(self) -> Dict:
        """Get default configuration."""
        return {
            # Maximum number of simultaneous trades
            'max_open_positions': 3,

            # Maximum trades per ticker
            'max_trades_per_ticker': 2,

            # Maximum trades per direction (LONG/SHORT)
            'max_trades_per_direction': 2,

            # Maximum trades by market type (MOEX/Forex/Crypto)
            'max_trades_per_market': 2,

            # Minimum time between positions
            'min_position_age_hours': 12,

            # Cooldown period after a losing trade
            'cooldown_hours_after_loss': 4,

            # Risk limit per position (% of portfolio)
            'max_position_risk_pct': 10.0
        }

    def can_open_position(self, signal: Dict, current_positions: List[Dict],
                         portfolio_value: float = 100000) -> Dict:
        """
        Check if a new position can be opened.

        Returns dict with:
            - can_open: bool
            - reason: str (if cannot open)
            - score: float (0-1, how open is allowed)
        """
        results = {
            'can_open': True,
            'reason': None,
            'score': 1.0
        }

        # 1. Limit max open positions
        if len(current_positions) >= self.config['max_open_positions']:
            results['can_open'] = False
            results['reason'] = f'Max positions reached: {len(current_positions)}/{self.config["max_open_positions"]}'
            results['score'] = 0.0
            return results

        # 2. Limit per ticker
        tickers_with_positions = [p['ticker'] for p in current_positions]
        if signal['ticker'] in tickers_with_positions:
            count = tickers_with_positions.count(signal['ticker'])
            if count >= self.config['max_trades_per_ticker']:
                results['can_open'] = False
                results['reason'] = f'Too many positions in {signal["ticker"]}: {count}/{self.config["max_trades_per_ticker"]}'
                results['score'] = 0.5
                return results

        # 3. Limit per direction
        directions = [p['direction'] for p in current_positions]
        if signal['direction'] in directions:
            count = directions.count(signal['direction'])
            if count >= self.config['max_trades_per_direction']:
                results['can_open'] = False
                results['reason'] = f'Too many {signal["direction"]} positions: {count}/{self.config["max_trades_per_direction"]}'
                results['score'] = 0.5
                return results

        # 4. Check position age (prevent opening multiple positions too close)
        # Build list of timestamps, ignoring missing or None values
        timestamps = []
        for p in current_positions:
            if p.get('entry_time') and isinstance(p['entry_time'], (int, float)):
                timestamps.append(p['entry_time'])
        latest_position_time = max(timestamps, default=0)

        if latest_position_time > 0:
            time_diff = datetime.now() - datetime.fromtimestamp(latest_position_time)

            # Minimum age between positions
            min_age = timedelta(hours=self.config['min_position_age_hours'])
            if time_diff < min_age:
                results['can_open'] = False
                results['reason'] = f'Position too recent: {time_diff.total_seconds()/3600:.1f}h < {self.config["min_position_age_hours"]}h'
                results['score'] = 0.3
                return results

        # 5. Check cooldown after loss
        recent_losers = []
        for p in current_positions:
            if not p.get('entry_time') or not isinstance(p['entry_time'], (int, float)):
                continue
            if (datetime.now() - datetime.fromtimestamp(p['entry_time'])).total_seconds() < 3600 and p.get('pnl', 0) < 0:
                recent_losers.append(p)

        if recent_losers:
            cooldown_hours = self.config['cooldown_hours_after_loss']
            results['can_open'] = False
            results['reason'] = f'Cooldown active: {len(recent_losers)} losing positions in last hour'
            results['score'] = 0.2
            return results

        # 6. Calculate position risk
        if portfolio_value > 0:
            position_risk_pct = signal.get('atr_entry', 0) / signal.get('entry_price', 1) * 100

            if position_risk_pct > self.config['max_position_risk_pct']:
                results['can_open'] = False
                results['reason'] = f'Position too risky: {position_risk_pct:.1f}% > {self.config["max_position_risk_pct"]}%'
                results['score'] = 0.4
                return results

        # All checks passed
        results['can_open'] = True
        results['reason'] = 'Position approved'
        results['score'] = 1.0
        return results

    def open_position(self, position_id: int, signal: Dict, entry_time: float):
        """Register a new open position."""
        self.open_positions[position_id] = {
            'ticker': signal['ticker'],
            'direction': signal['direction'],
            'entry_price': signal['entry_price'],
            'atr_entry': signal['atr_entry'],
            'entry_time': entry_time,
            'pnl': 0
        }
        self.entry_time[position_id] = entry_time

    def update_position(self, position_id: int, current_price: float, atr: float):
        """Update position with current price and ATR."""
        if position_id in self.open_positions:
            pos = self.open_positions[position_id]
            pos['current_price'] = current_price
            pos['current_atr'] = atr

    def close_position(self, position_id: int, exit_price: float, pnl: float, exit_time: float):
        """Close a position and remove it from tracking."""
        if position_id in self.open_positions:
            self.open_positions[position_id]['exit_price'] = exit_price
            self.open_positions[position_id]['exit_time'] = exit_time
            self.open_positions[position_id]['pnl'] = pnl

            # Remove from active positions
            del self.open_positions[position_id]
            del self.entry_time[position_id]

    def get_open_positions(self) -> List[Dict]:
        """Get list of currently open positions."""
        return list(self.open_positions.values())

    def get_position_by_id(self, position_id: int) -> Optional[Dict]:
        """Get position by ID."""
        return self.open_positions.get(position_id)

    def get_position_by_ticker(self, ticker: str) -> Optional[Dict]:
        """Get all positions for a ticker."""
        return [p for p in self.open_positions.values() if p['ticker'] == ticker]

    def get_portfolio_risk(self, portfolio_value: float = 100000) -> float:
        """Calculate total portfolio risk (sum of ATR % of prices)."""
        total_risk = 0.0
        for pos in self.open_positions.values():
            atr_pct = pos['current_atr'] / pos['entry_price']
            total_risk += atr_pct

        return total_risk

    def get_unrealized_pnl(self, current_prices: Dict[int, float]) -> float:
        """Calculate unrealized PnL for all open positions."""
        total_pnl = 0.0

        for pos_id, pos in self.open_positions.items():
            current_price = current_prices.get(pos_id, pos['entry_price'])
            pnl = (current_price - pos['entry_price']) * pos.get('volume', 1000) / 1000
            total_pnl += pnl

        return total_pnl

    def calculate_diversification_score(self) -> float:
        """Calculate portfolio diversification score (0-1)."""
        if not self.open_positions:
            return 1.0

        # Score is based on:
        # 1. Number of tickers (more tickers = more diversified)
        tickers_count = len(set(p['ticker'] for p in self.open_positions.values()))
        max_tickers = self.config['max_open_positions']
        ticker_score = min(tickers_count / max_tickers, 1.0)

        # 2. Balance between directions
        directions = [p['direction'] for p in self.open_positions.values()]
        long_count = directions.count('LONG')
        short_count = directions.count('SHORT')
        direction_score = min(long_count, short_count) / max(long_count, short_count) if (long_count + short_count) > 0 else 1.0

        # Weighted average
        return 0.6 * ticker_score + 0.4 * direction_score

    def calculate_diversification_report(self) -> str:
        """Calculate and return diversification report string."""
        return create_diversification_report(
            current_positions=self.get_open_positions(),
            open_count=len(self.get_open_positions()),
            max_positions=self.config['max_open_positions']
        )


def create_diversification_report(current_positions: List[Dict], open_count: int,
                                   max_positions: int = 3) -> str:
    """Create human-readable diversification report."""
    if open_count == 0:
        return "No open positions"

    # Categorize by ticker
    ticker_stats = {}
    for pos in current_positions:
        ticker = pos['ticker']
        if ticker not in ticker_stats:
            ticker_stats[ticker] = {'count': 0, 'directions': set()}
        ticker_stats[ticker]['count'] += 1
        ticker_stats[ticker]['directions'].add(pos['direction'])

    report = f"\n📊 Diversification Report ({open_count}/{max_positions} positions):\n"
    report += "-" * 60 + "\n"

    for ticker, stats in ticker_stats.items():
        directions_str = ', '.join(stats['directions'])
        report += f"  {ticker:<10} {stats['count']} pos | {directions_str}\n"

    return report

    def check_diversification(self, ticker: str, direction: str) -> tuple[bool, str, float]:
        """
        Check if we can open a new position respecting all diversification rules.
        
        Returns:
            (can_open: bool, reason: str, score: float)
        """
        if self.diversification_disabled:
            return True, 'diversification_disabled', 1.0

        # Rule 1: Max open positions
        open_positions = self.get_open_positions()
        if len(open_positions) >= self.config['max_open_positions']:
            return False, f'max_positions_reached ({len(open_positions)}/{self.config["max_open_positions"]})', 0.0

        # Rule 2: Max trades per ticker
        ticker_positions = [p for p in open_positions if p['ticker'] == ticker]
        if len(ticker_positions) >= self.config['max_trades_per_ticker']:
            return False, f'max_per_ticker_reached ({ticker}: {len(ticker_positions)}/{self.config["max_trades_per_ticker"]})', 0.0

        # Rule 3: Max trades per direction
        direction_positions = [p for p in open_positions if p['direction'] == direction]
        if len(direction_positions) >= self.config['max_trades_per_direction']:
            return False, f'max_per_direction_reached ({direction}: {len(direction_positions)}/{self.config["max_trades_per_direction"]})', 0.0

        # Rule 4: Min position age (already closed positions don't count)
        # This rule applies to NEW positions, not closed ones
        # So we don't need to check it for opening new positions

        # Rule 5: Cooldown after loss
        # Check last closed position for this ticker
        ticker_closed = self._get_last_closed_by_ticker(ticker)
        if ticker_closed and ticker_closed['close_reason'] == 'SL':
            now_ts = int(time.time())
            close_ts = int(ticker_closed.get('exit_time', 0))
            # fallback: если exit_time битый, используем entry_time
            if close_ts < int(ticker_closed.get('entry_time', 0)):
                logger.warning(f"[{ticker}] exit_time={close_ts} < entry_time={ticker_closed.get('entry_time', 0)}, "
                              f"использую entry_time как fallback для cooldown")
                close_ts = int(ticker_closed.get('entry_time', 0))
            delta = now_ts - close_ts
            if delta < self.config['cooldown_hours_after_loss'] * 3600:
                hours_left = int((self.config['cooldown_hours_after_loss'] * 3600 - delta) / 3600)
                return False, f'cooldown_after_loss ({hours_left}h remaining)', 0.0

        # All checks passed - calculate diversification score
        # Score is based on:
        # 1. Number of tickers (more tickers = more diversified)
        tickers_count = len(set(p['ticker'] for p in open_positions))
        max_tickers = self.config['max_open_positions']
        ticker_score = min(tickers_count / max_tickers, 1.0)

        # 2. Balance between directions
        directions = [p['direction'] for p in open_positions]
        long_count = directions.count('LONG')
        short_count = directions.count('SHORT')
        direction_score = min(long_count, short_count) / max(long_count, short_count) if (long_count + short_count) > 0 else 1.0

        # Weighted average
        score = 0.6 * ticker_score + 0.4 * direction_score

        return True, 'ok', score