from AlgorithmImports import *
from datetime import timedelta
import numpy as np
from scipy.stats import norm

class EuropeanOptionsArbitrageAlgorithm(QCAlgorithm):
    
    def initialize(self):
        self.set_start_date(2012, 1, 1)
        self.set_cash(1000000)
        
        # Evropské indexy
        self._symbols = ["SPX", "NDX", "RUT"]
        self._underlying = {}
        self._option_chains = {}
        
        # Přidání indexů a jejich opcí
        for symbol_str in self._symbols:
            index = self.add_index(symbol_str, Resolution.MINUTE)
            self._underlying[symbol_str] = index.symbol
            
            # Přidání opcí - evropské opce
            option = self.add_index_option(index.symbol, Resolution.MINUTE)
            option.set_filter(lambda u: u.include_weeklys().strikes(-1, 1).expiration(0, 60))
            self._option_chains[symbol_str] = option.symbol
        
        # Black-Scholes kalkulátor
        self._bs_calculator = BlackScholesCalculator()
        
        # Sledování otevřených pozic
        self._open_positions = {}
        
        # Naplánování kontroly před zavřením trhu (15:45 EST)
        self.schedule.on(
            self.date_rules.every_day(),
            self.time_rules.at(15, 45),
            self._check_arbitrage_opportunities
        )
        
        # Nastavení pro seed prices
        self.settings.seed_initial_prices = True
        
        self.debug("Algoritmus inicializován pro arbitráž evropských opcí")
    
    def _check_arbitrage_opportunities(self):
        """Kontrola arbitrážních příležitostí před zavřením trhu"""
        
        # Kontrola cash reserve (minimálně 80%)
        cash_ratio = self.portfolio.cash / self.portfolio.total_portfolio_value
        if cash_ratio < 0.80:
            self.debug(f"Cash reserve ({cash_ratio:.1%}) pod 80%, žádné nové pozice")
            return
        
        for symbol_str in self._symbols:
            # Kontrola zda už máme otevřenou pozici na tomto tickeru
            if self._has_open_position_for_ticker(symbol_str):
                continue
            
            underlying_symbol = self._underlying[symbol_str]
            
            # Získání aktuální ceny underlying
            if not self.securities.contains_key(underlying_symbol):
                continue
                
            underlying_price = self.securities[underlying_symbol].price
            if underlying_price == 0:
                continue
            
            # Získání option chain
            option_symbol = self._option_chains[symbol_str]
            chain = self.current_slice.option_chains.get(option_symbol)
            
            if chain is None or len(chain) == 0:
                continue
            
            # Odhad implicitní volatility z ATM opcí
            iv = self._estimate_implied_volatility(chain, underlying_price)
            if iv == 0:
                continue
            
            # Risk-free rate (aktuální)
            risk_free_rate = 0.05  # 5% aproximace
            
            # Kontrola každé opce v chainu
            for contract in chain:
                # Kontrola likvidity
                if contract.bid_price == 0 or contract.ask_price == 0:
                    continue
                
                # Výpočet času do expirace
                time_to_expiry = (contract.expiry - self.time).days / 365.0
                if time_to_expiry <= 0:
                    continue
                
                # Tržní cena (bid pro nákup, použijeme ask protože kupujeme)
                market_price = contract.ask_price
                
                # Výpočet Black-Scholes fair value
                if contract.right == OptionRight.CALL:
                    bs_price = self._bs_calculator.call_option(
                        underlying_price,
                        contract.strike,
                        time_to_expiry,
                        risk_free_rate,
                        iv
                    )
                else:  # PUT
                    bs_price = self._bs_calculator.put_option(
                        underlying_price,
                        contract.strike,
                        time_to_expiry,
                        risk_free_rate,
                        iv
                    )
                
                # Kontrola arbitrážní příležitosti
                # Kupujeme když je tržní cena nižší než BS fair value
                discount = (bs_price - market_price) / bs_price if bs_price > 0 else 0
                
                if discount > 0.10:  # Minimálně 10% discount pro arbitráž
                    self._execute_arbitrage(contract, market_price, bs_price, underlying_price, symbol_str)
    
    def _has_open_position_for_ticker(self, ticker):
        """Kontrola zda existuje otevřená pozice pro daný ticker"""
        for symbol, position in self._open_positions.items():
            if position.get('ticker') == ticker:
                return True
        return False
    
    def _estimate_implied_volatility(self, chain, underlying_price):
        """Odhad průměrné IV z ATM opcí"""
        atm_ivs = []
        
        for contract in chain:
            # Pouze ATM opce (±5% od underlying)
            if abs(contract.strike - underlying_price) / underlying_price < 0.05:
                if hasattr(contract, 'implied_volatility') and contract.implied_volatility > 0:
                    atm_ivs.append(contract.implied_volatility)
        
        if len(atm_ivs) > 0:
            return np.mean(atm_ivs)
        
        return 0.25  # Default 25% pokud není k dispozici
    
    def _execute_arbitrage(self, contract, market_price, bs_price, underlying_price, ticker):
        """Provedení arbitrážního obchodu"""
        
        # Kontrola zda už nemáme pozici na tomto tickeru
        if self._has_open_position_for_ticker(ticker):
            return
        
        # Kontrola zda už nemáme tuto konkrétní pozici
        if contract.symbol in self._open_positions:
            return
        
        # Maximálně 5% portfolia na jeden trade
        max_position_value = self.portfolio.total_portfolio_value * 0.01
        
        # Počet kontraktů (každý kontrakt = 100 akcií)
        contract_value = market_price * 100
        if contract_value == 0:
            return
            
        quantity = int(max_position_value / contract_value)
        
        if quantity < 1:
            return
        
        # Market order nákup
        ticket = self.market_order(contract.symbol, quantity)
        
        if ticket.status == OrderStatus.FILLED:
            self._open_positions[contract.symbol] = {
                'ticker': ticker,
                'quantity': quantity,
                'entry_price': market_price,
                'bs_price': bs_price,
                'expiry': contract.expiry,
                'underlying': underlying_price,
                'strike': contract.strike,
                'right': contract.right
            }
            
            self.debug(f"ARBITRÁŽ {ticker}: {contract.symbol} | Množství: {quantity} | "
                      f"Tržní cena: ${market_price:.2f} | BS cena: ${bs_price:.2f} | "
                      f"Discount: {((bs_price - market_price) / bs_price * 100):.1f}%")
    
    def on_data(self, data):
        """Event handler pro příchozí data"""
        # Kontrola expirovaných opcí
        expired_symbols = []
        for symbol, position in self._open_positions.items():
            if self.time >= position['expiry']:
                expired_symbols.append(symbol)
        
        # Odstranění expirovaných pozic ze sledování
        for symbol in expired_symbols:
            if symbol in self._open_positions:
                position = self._open_positions[symbol]
                ticker = position.get('ticker', 'N/A')
                self.debug(f"EXPIROVÁNO {ticker}: {symbol} | Drženo do splatnosti")
                del self._open_positions[symbol]


class BlackScholesCalculator:
    """Black-Scholes model pro oceňování evropských opcí"""
    
    def call_option(self, S, K, T, r, sigma):
        """
        Výpočet ceny evropské call opce
        S: Aktuální cena underlying
        K: Strike price
        T: Čas do expirace (roky)
        r: Risk-free rate
        sigma: Volatilita (IV)
        """
        if T <= 0 or sigma <= 0:
            return max(S - K, 0)  # Intrinsic value
        
        d1 = (np.log(S / K) + (r + 0.5 * sigma ** 2) * T) / (sigma * np.sqrt(T))
        d2 = d1 - sigma * np.sqrt(T)
        
        call_price = S * norm.cdf(d1) - K * np.exp(-r * T) * norm.cdf(d2)
        return max(call_price, 0)
    
    def put_option(self, S, K, T, r, sigma):
        """
        Výpočet ceny evropské put opce
        S: Aktuální cena underlying
        K: Strike price
        T: Čas do expirace (roky)
        r: Risk-free rate
        sigma: Volatilita (IV)
        """
        if T <= 0 or sigma <= 0:
            return max(K - S, 0)  # Intrinsic value
        
        d1 = (np.log(S / K) + (r + 0.5 * sigma ** 2) * T) / (sigma * np.sqrt(T))
        d2 = d1 - sigma * np.sqrt(T)
        
        put_price = K * np.exp(-r * T) * norm.cdf(-d2) - S * norm.cdf(-d1)
        return max(put_price, 0)
    
    def delta(self, S, K, T, r, sigma, option_type='call'):
        """Výpočet delta (citlivost na změnu ceny underlying)"""
        if T <= 0 or sigma <= 0:
            return 0
        
        d1 = (np.log(S / K) + (r + 0.5 * sigma ** 2) * T) / (sigma * np.sqrt(T))
        
        if option_type.lower() == 'call':
            return norm.cdf(d1)
        else:  # put
            return norm.cdf(d1) - 1
    
    def gamma(self, S, K, T, r, sigma):
        """Výpočet gamma (citlivost delta na změnu ceny)"""
        if T <= 0 or sigma <= 0:
            return 0
        
        d1 = (np.log(S / K) + (r + 0.5 * sigma ** 2) * T) / (sigma * np.sqrt(T))
        return norm.pdf(d1) / (S * sigma * np.sqrt(T))
    
    def vega(self, S, K, T, r, sigma):
        """Výpočet vega (citlivost na změnu volatility)"""
        if T <= 0 or sigma <= 0:
            return 0
        
        d1 = (np.log(S / K) + (r + 0.5 * sigma ** 2) * T) / (sigma * np.sqrt(T))
        return S * norm.pdf(d1) * np.sqrt(T) / 100
    
    def theta(self, S, K, T, r, sigma, option_type='call'):
        """Výpočet theta (time decay)"""
        if T <= 0 or sigma <= 0:
            return 0
        
        d1 = (np.log(S / K) + (r + 0.5 * sigma ** 2) * T) / (sigma * np.sqrt(T))
        d2 = d1 - sigma * np.sqrt(T)
        
        if option_type.lower() == 'call':
            theta = (-S * norm.pdf(d1) * sigma / (2 * np.sqrt(T)) 
                    - r * K * np.exp(-r * T) * norm.cdf(d2)) / 365
        else:  # put
            theta = (-S * norm.pdf(d1) * sigma / (2 * np.sqrt(T)) 
                    + r * K * np.exp(-r * T) * norm.cdf(-d2)) / 365
        
        return theta