# src/models/risk.py
import numpy as np
from typing import Dict, Optional
from config import USE_KELLY, USE_DYNAMIC_SIZING, KELLY_FRACTION, MAX_RISK_PER_TRADE


class KellySizing:
    """Kelly Criterion для расчёта размера позиции."""
    def __init__(self, kelly_fraction: float = KELLY_FRACTION, max_risk: float = MAX_RISK_PER_TRADE):
        self.kelly_fraction = kelly_fraction
        self.max_risk = max_risk

    def calculate(self, win_rate: float, avg_win: float, avg_loss: float) -> float:
        """Возвращает долю капитала для ставки."""
        if avg_loss == 0 or win_rate <= 0 or win_rate >= 1:
            return 0.0
        
        win_loss_ratio = avg_win / avg_loss
        kelly = (win_rate * win_loss_ratio - (1 - win_rate)) / win_loss_ratio
        
        kelly = max(0, kelly)
        sized_kelly = kelly * self.kelly_fraction
        return min(sized_kelly, self.max_risk)


class DynamicPositionSizer:
    """Динамический расчёт размера позиции на основе уверенности модели."""
    def __init__(self, base_risk: float = MAX_RISK_PER_TRADE, confidence_weights: Optional[Dict] = None):
        self.base_risk = base_risk
        self.confidence_weights = confidence_weights or {
            'HIGH': 1.0,
            'MEDIUM': 0.7,
            'LOW': 0.3
        }

    def calculate(self, prob: float, confidence: str, ev: float, rr_ratio: float) -> float:
        """Рассчитывает размер позиции на основе вероятности, уверенности и EV."""
        conf_mult = self.confidence_weights.get(confidence, 0.5)
        
        if ev <= 0:
            return 0.0
        
        base_size = self.base_risk * conf_mult
        
        prob_edge = max(0, prob - 0.5) * 2
        
        dynamic_size = base_size * prob_edge * (1 + ev / rr_ratio)
        
        return min(dynamic_size, self.base_risk)


class RiskManager:
    """Комбинированное управление риском."""
    def __init__(self):
        self.kelly = KellySizing() if USE_KELLY else None
        self.sizer = DynamicPositionSizer() if USE_DYNAMIC_SIZING else None

    def compute_position_size(self, direction: str, prob: float, confidence: str, 
                             ev: float, rr_ratio: float, 
                             historical_winrate: Optional[float] = None,
                             avg_win: Optional[float] = None,
                             avg_loss: Optional[float] = None) -> Dict:
        """Возвращает рекомендуемый размер позиции и метаинформацию."""
        result = {
            'kelly_fraction': 0.0,
            'dynamic_fraction': 0.0,
            'final_fraction': 0.0,
            'recommended_risk': 'NONE'
        }

        if self.kelly and historical_winrate is not None and avg_win is not None and avg_loss is not None:
            result['kelly_fraction'] = self.kelly.calculate(historical_winrate, avg_win, avg_loss)

        if self.sizer:
            result['dynamic_fraction'] = self.sizer.calculate(prob, confidence, ev, rr_ratio)

        if result['kelly_fraction'] > 0 and result['dynamic_fraction'] > 0:
            result['final_fraction'] = min(result['kelly_fraction'], result['dynamic_fraction'])
        elif result['kelly_fraction'] > 0:
            result['final_fraction'] = result['kelly_fraction']
        elif result['dynamic_fraction'] > 0:
            result['final_fraction'] = result['dynamic_fraction']

        if result['final_fraction'] >= MAX_RISK_PER_TRADE * 0.8:
            result['recommended_risk'] = 'FULL'
        elif result['final_fraction'] >= MAX_RISK_PER_TRADE * 0.4:
            result['recommended_risk'] = 'HALF'
        elif result['final_fraction'] > 0:
            result['recommended_risk'] = 'MINI'
        else:
            result['recommended_risk'] = 'NONE'

        return result