from __future__ import annotations

import option
import spread

from AlgorithmImports import *


class FocusedApricotScorpion(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)
        self.SetCash(300000)
        self.SetSecurityInitializer(self.SecurityInitializer)
        
        self.spx = self.AddIndex("SPX").Symbol
        self.spxw = self.AddIndexOption(self.spx, "SPXW").Symbol

        self.Schedule.On(
            self.DateRules.EveryDay(self.spx),
            self.TimeRules.AfterMarketOpen(self.spx, 30),
            self._after_open_of_market,
        )

    def SecurityInitializer(self, security):
        security.SetMarketPrice(self.GetLastKnownPrice(security))

        if security.Type in [SecurityType.Equity, SecurityType.Index]:
            periods = 5
            security.VolatilityModel = StandardDeviationOfReturnsVolatilityModel(periods)
            trade_bars = self.History[TradeBar](security.Symbol, periods, Resolution.Daily)
            for trade_bar in trade_bars:
                security.VolatilityModel.Update(security, trade_bar)
                
        # if security.Type is SecurityType.Option:
        #     security.PriceModel = OptionPriceModels.BjerksundStensland()

    def OnData(self, data: Slice):
        if not self.Portfolio.Invested:
            return

    def _after_open_of_market(self):
        option_symbols = self.OptionChainProvider.GetOptionContractList(self.spxw, self.Time)
        option_chains = option.sort_listed_options(option_symbols)
        today = self.Time.date()
        
        # Check if today is an option expiration day
        if today not in option_chains:
            return

        # Get the straddle price
        chain = option_chains[today]
        spot_price = self.Securities[self.spx].Close
        volatility = self.Securities[self.spx].VolatilityModel.Volatility
        atm_options = option.get_atm_options(spot_price, chain)
        leg_1 = self.AddOptionContract(atm_options[0].symbol)
        leg_2 = self.AddOptionContract(atm_options[1].symbol)
        straddle = spread.get_straddle_price(spot_price, volatility, leg_1, leg_2)

        # Filter trades based on factors
        cond_0 = straddle.percent_price > 1
        cond_1 = straddle.percent_price < 1

        cond_2 = straddle.relative_percent_price > 4
        cond_3 = straddle.relative_percent_price < 5
        # if cond_1 and cond_2:
        #     return
        if cond_3:
            return

        # Send straddle order
        trade_risk = 10_000
        try:
            quantity = round(trade_risk / straddle.straddle_premium)
            legs = []
            legs.append(Leg.Create(leg_1.Symbol, 1))
            legs.append(Leg.Create(leg_2.Symbol, 1))
            self.ComboMarketOrder(legs, quantity, tag=str(straddle.relative_percent_price))
        except ZeroDivisionError:
            print("Missing straddle price.")

from __future__ import annotations
from collections import defaultdict
import datetime as dt

from AlgorithmImports import *


def sort_listed_options(symbols: list) -> dict:
    """Group all active options for symbol by expiration date."""
    chains = defaultdict(list)
    for s in symbols:
        option = OptionSymbol(s).compute()
        chains[option.expiration].append(option)
    return chains

def get_atm_options(spot_price: float, options: list) -> list:
    sorted_options = sorted(options, key=lambda x: x.distance_from_atm(spot_price))
    return sorted_options[:2]


class OptionSymbol:
    def __init__(self, symbol):
        self.symbol = symbol

    def compute(self) -> self:
        self.expiration = self._expiration()
        self.strike = self._strike()
        self.put_call = self._put_call()
        return self

    def distance_from_atm(self, price) -> float:
        return abs(self.strike - price)

    def _expiration(self) -> dt.date:
        return dt.datetime.strptime(self.symbol.Value[6:12], "%y%m%d").date()
    
    def _strike(self) -> float:
        return float(self.symbol.Value[-8:]) / 1000

    def _put_call(self) -> str:
        return self.symbol.Value[12:13]

from __future__ import annotations
from dataclasses import dataclass

from AlgorithmImports import *


def get_straddle_price(spot_price: float, volatility: float, leg_1, leg_2) -> StraddlePrice:
    straddle_price = _mid_price(leg_1) + _mid_price(leg_2)
    return StraddlePrice(straddle_price, spot_price, volatility)

def _mid_price(contract) -> float:
    return (contract.AskPrice + contract.BidPrice) / 2


@dataclass
class StraddlePrice:
    """Keep record of straddle price and percent price."""
    straddle_price: float
    stock_price: float
    volatility: float

    @property
    def straddle_premium(self) -> float:
        return self.straddle_price * 100

    @property
    def percent_price(self) -> float:
        """Straddle price as percentage of stock price."""
        return round(self._percent_price() * 100, 3)

    @property
    def relative_percent_price(self) -> float:
        """Straddle percent price divided by historical volatility."""
        return round(self._percent_price() / self.volatility * 100, 3)

    def _percent_price(self) -> float:
        return self.straddle_price / self.stock_price