from trade import *
from levels import *

class OptionsOvernightContrarian(QCAlgorithm):

    def Initialize(self):
        
        # Settings
        self.SetStartDate(2020, 1, 31)
        self.SetEndDate(2020, 3, 5)
        self.SetCash(1000000)
        self.SetWarmup(timedelta(7))
        self.EnableAutomaticIndicatorWarmUp = True 
        self.orderDate = None
        
        # Apply Robinhood Fees
        self.SetSecurityInitializer(lambda security: security.SetFeeModel(ConstantFeeModel(0)))
        
        # Select Underlying and Configure
        self.ticker = "SPY"
        self.equity = self.AddEquity(self.ticker)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        # Select Target Asset to Trade
        self.option = self.AddOption(self.ticker)
        self.option.SetFilter(self.OptionFilterUniverse)
        
        # Exit before any assignments
        self.Schedule.On(self.DateRules.EveryDay(self.ticker), self.TimeRules.AfterMarketOpen(self.ticker, 1), self.MarketOpen)
        self.Schedule.On(self.DateRules.EveryDay(self.ticker), self.TimeRules.BeforeMarketClose(self.ticker, 5), self.BeforeMarketClose)

        # Trade Tracker
        self.Trade = TradeManagement(self, self.ticker, self.option.Symbol)
        
        # Support Resistance Detection:
        self.SupportResistance = SupportResistance(self, self.ticker)

    def OnData(self, data):
        
        # Ignore dividends, splits etc
        if not self.equity.Exchange.ExchangeOpen or self.IsWarmingUp: 
            return 
        
        # Manage existing positions:
        self.Trade.ManageOpenPositions()
        
        # Consider new ones:
        # if dayRange > 0:
        self.Trade.Create(OrderDirection.Buy)
        # else:
            # self.Trade.Create(OrderDirection.Sell)





















#################################################################################################
# SCRATCH #######################################################################################
#################################################################################################

    def MarketOpen(self): 
        pass
    
    def BeforeMarketClose(self):
        pass

    def OptionFilterUniverse(self, universe):
        # Select puts 2-3 strikes OOM, expiring at least 2 days out; but no more than 7
        return universe.IncludeWeeklys().Strikes(-3, 3).Expiration(2, 7)
        
    def StrikeVisualization(self, contracts):
        strikes = f"{self.Time} - {self.equity.Price} :: [["
        for x in self.contracts:
            strikes = strikes + f"{x.Strike}-{x.Expiry} ]] [["
        self.Log(strikes)
        self.Log(f"{self.Time} - Contract: {self.contracts[0]} | Strike: {self.contracts[0].Strike} | Underlying: {self.equity.Price} | Delta: {self.equity.Price - self.contracts[0].Strike}")

class SupportResistance:
    
    # Get the S&R's for the last 2 weeks.
    def __init__(self, algorithm, ticker):
        self.ticker = ticker
        self.daily = RollingWindow[TradeBar](14);
        self.min    = algorithm.MIN(self.ticker, 390, Resolution.Minute)     # Range of today; breaking out of new highs/lows
        self.max    = algorithm.MAX(self.ticker, 390, Resolution.Minute)
        algorithm.Consolidate(self.ticker, Resolution.Daily, self.SaveDailyBars)
        
    def NextSupport(self):
        pass
    
    def NextResistance(self):
        pass
    
    # Build a 14 day historical rolling window of underlying prices.
    def SaveDailyBars(self, bar):
        self.daily.Add(bar)

import random
import string
import math

class TradeManagement:
    
    def __init__(self, algorithm, ticker, optionSymbol):
        self.Trades     = {}
        self.Algorithm  = algorithm
        self.Ticker = ticker
        self.OptionSymbol = optionSymbol
        
        # Volatility indicators for position sizing
        self.atr    = self.Algorithm.ATR(self.Ticker, 3, MovingAverageType.Simple, Resolution.Daily)
        self.stddev = self.Algorithm.STD(self.Ticker, 6, Resolution.Hour)
    
    # Manage Open Positions
    def ManageOpenPositions(self):
        
        for t in self.OpenTrades():
            
            # Scan for Profit-Loss
            if t.UnrealizedProfit() > t.Target:
                self.Close(t, "Profit")
                
            elif t.UnrealizedProfit() < -t.Target:
                self.Close(t, "Loss")
                
            # Stop Assignment
            if t.ExpiringSoon():
                self.Close(t, "Expiring")
    
    
    # Base target contract count on the number of contracts to hit the profit assuming 2SD move.
    def ContractSizing(self, targetProfit): 
        expectedDollarMovement = 2 * self.stddev.Current.Value * 100
        # At least 1 contract
        contracts = min(5, math.ceil(targetProfit / expectedDollarMovement))
        return int(contracts)
        
    # Base target profit per position on the volatility of the last few days.
    def TargetProfitEstimate(self):
        return round(self.atr.Current.Value * 100)
        
    def OpenTrades(self):
        return [t for t in self.Trades.values() if t.IsOpen()]
    
    # Place a trade in the direction signalled
    def Create(self, direction):
        
        symbol = self.SelectOptionContract(direction)
        if symbol is None:
            return
        
        # If we already hold; skip
        alreadyOpen = [c for c in self.OpenTrades() if c.Symbol == symbol]
        if len(alreadyOpen) > 0:
            return
        
        # If today's the expiry don't trade
        # if (symbol.ID.Date < self.Algorithm.Time):
        #     return
        
        targetProfit = self.TargetProfitEstimate()
        size = self.ContractSizing(targetProfit)
        price = self.Algorithm.Securities[symbol].Price
        
        tag = f"Opened for target of ${targetProfit} at ${price}"
        ticket  = self.Algorithm.MarketOrder(symbol, size, False, tag)
        trade   = Trade(self.Algorithm, self.Algorithm.Securities[symbol], ticket, targetProfit)
        
        self.Trades[trade.Id] = trade
        return trade
        
    def Close(self, trade, reason):
        trade.Close(reason)
        del self.Trades[trade.Id]
        
    # Select OOM Option Contract with Given Bias
    def SelectOptionContract(self, direction):
        
        if direction == OrderDirection.Buy:
            right = OptionRight.Call
        else:
            right = OptionRight.Put
        
        chain = self.Algorithm.CurrentSlice.OptionChains.GetValue(self.OptionSymbol)
        if chain is None: 
            self.Algorithm.Log(f"{self.Algorithm.Time} - No option chains with {self.OptionSymbol}. Missing data for Mar-2nd.")
            return None
        
        # Select contracts expirying tomorrow at least.
        chain = [x for x in chain if (x.Right == right and x.Expiry > self.Algorithm.Time)]
        reverseSort = right == OptionRight.Call         # Reverse sort of a call
        contracts = sorted(sorted(chain, key = lambda x: abs(x.Strike), reverse=reverseSort), key = lambda x: x.Expiry)
        
        # No contracts found
        if len(contracts) == 0:
            return None;
            
        return contracts[0].Symbol


class Trade:
    
    def __init__(self, algorithm, contract, ticket, target): 
        self.Id         = ''.join(random.choice(string.ascii_lowercase) for i in range(10))
        self.Ticket     = ticket 
        self.Contract   = contract
        self.Symbol     = contract.Symbol
        self.Algorithm  = algorithm
        self.Open       = True
        self.Target     = target
        
    def ExpiringSoon(self): 
        expiry = self.Symbol.ID.Date + timedelta(hours=16)
        if ( (expiry - self.Algorithm.Time) < timedelta(minutes=10)):
            self.Algorithm.Log(f"{self.Symbol} Close to Expiry: {expiry} - {self.Algorithm.Time} < 10min" )
            return True
        else:
            return False
        
    def UnrealizedProfitPercent(self):
        return (self.Contract.Holdings.Price - self.Ticket.AverageFillPrice) / self.Contract.Holdings.Price
    
    def UnrealizedProfit(self):
        return 100*(self.Contract.Holdings.Price - self.Ticket.AverageFillPrice) * self.Ticket.Quantity
        
    def IsOpen(self):
        return self.Open
        
    def Close(self, reason):
        self.Open = False
        tag = f"Close | {reason} | {self.UnrealizedProfit()} Net"
        self.Algorithm.Liquidate(self.Symbol, tag)