from Execution.ImmediateExecutionModel import ImmediateExecutionModel

from datetime import time

from collections import defaultdict

import numpy as np

from io import StringIO

import pandas as pd

import talib

class MicroShort(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)  # Set Start Date
        self.SetEndDate(2020, 12, 31)  # Set Start Date
        # self.SetEndDate(2019, 1, 1)  # Set Start Date
        
        #self.SetStartDate(2014, 1, 1)  # Set Start Date
        #self.SetEndDate(2014, 10, 1)  # Set Start Date
        
        self.SetCash(3000)  # Set Strategy Cash
        self.SetExecution(ImmediateExecutionModel())
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        
        self.UniverseSettings.Resolution = Resolution.Minute
        
        self.UniverseSettings.ExtendedMarketHours = True
        
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
       # self.SetSecurityInitializer(self.CustomSecurityInitializer)
        
        self.AddEquity("SPY")

        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(12, 0), self.CoverShorts)
        # self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(9, 35), self.CancelOrders)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(9, 20), self.FireOrders)
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday),self.TimeRules.At(9, 35),self.RebalanceVTI)
        
        self.openOrders = []
        self.stopOrders = []
        
        self.numberOfGappers = 20
        
        self.numberOfStocks = 5
        
        self.numberOfSymbolsCoarse = 100
    
    def OnData(self, data):
        pass
    
    def CustomSecurityInitializer(self, security):
        security.SetSlippageModel(ConstantSlippageModel(0.000))
        
    def RebalanceVTI(self):
        self.Log('Rebalancing SPY')
        self.SetHoldings('SPY', 0.4)

    def FireOrders(self):
        
        self.pumpedStocks = {}
        NightChange = dict()
        
        # Gappers = dict()
        
        VolumeGappers = dict()
        
        self.openOrders = []
        self.stopOrders = []
        
        # Find Gappers
        for security in self.ActiveSecurities.Values:
            
            if 'SPY' in str(security.Symbol):
                continue
            
            if security.HasData:
                    closehist = self.History(security.Symbol, 2, Resolution.Daily)
                            
                    if str(security.Symbol) not in closehist.index or closehist.empty:
                        #algorithm.Log(f'{ticker} not found in history data frame.')
                        continue
            
                    closebars = np.asarray(closehist['close'])
                    
                    if self.LiveMode:
                        closeyest = closebars[1]
                        # self.Log('On {} comparing with the close of {} from {}'.format(str(self.Time),closeyest,np.asarray(closehist.index.values)[1]))
                    else:
                        closeyest = closebars[0]
                        # self.Log('On {} comparing with the close of {} from {}'.format(str(self.Time),closeyest,np.asarray(closehist.index.values)[0]))
                        
                    todayopen = security.Open
                    
                    sigmahist = self.History(security.Symbol, 100, Resolution.Daily)
                    sigmahist_pct = sigmahist.pct_change()
                    
                    try:
                         sigmahist_close = np.asarray(sigmahist['close'])
                         sigmahist_pct_close = np.asarray(sigmahist_pct['close'])
                    except KeyError:
                         continue
                     
                    ema = talib.EMA(sigmahist_close,timeperiod=20)[-1]
                    gap = todayopen/ closeyest - 1
                    
                    if todayopen < ema:
                        # self.Debug('Above EMA')
                        continue
                    
                    sigmahist_pct_close = sigmahist_pct_close - np.nanmean(sigmahist_pct_close)
                    sigma = np.nanstd(sigmahist_pct_close)
                    
                    
                    if gap/sigma < 2:
                        continue
                    
                    NightChange[security.Symbol] = gap
        
        Gappers = dict(sorted(NightChange.items(), key = lambda kv: (-round(kv[1], 6), kv[0]))[:self.numberOfGappers])
        
        self.Log('Scanning for Gappers ...')
        
        # Rank after Pre-Market Volume
        for s,g in Gappers.items():
            
            # self.Debug('Found {} up {:.2f}%'.format(s,g*100))
            
            premarketh = self.History(s, 300, Resolution.Minute)
            
            if premarketh.empty:
                # self.Debug('PM empty')
                continue
            
            # totalsum = premarketh.sum()
            
            # if totalsum['volume'] > 1e5:
            #     VolumeGappers[s] = totalsum['volume'] 
            
            try:
                dollarvol = (premarketh.close*premarketh.volume).sum()
            except AttributeError:
                print('Attribute Error')
                continue
            
            # if dollarvol > 1e4:
            VolumeGappers[s] = dollarvol
        
        StocksToTrade = dict(sorted(VolumeGappers.items(), key = lambda kv: (-round(kv[1], 6), kv[0]))[0:self.numberOfStocks])
    
        self.Log('Scanning for Pre-Market-Volume...')
        
        for key,value in StocksToTrade.items():

            # self.Debug('SELL')
            
            sec = self.ActiveSecurities[key]
            
            #Due to Margin requirements, take only half
            available = self.Portfolio.TotalPortfolioValue*0.5
            
            #Short Sell out of the gates
            shares = round(available/len(StocksToTrade)/sec.Price)
        
            self.Log('Shorting {} shares of {} at {:.2f} = {:.2f} with {:.2e} pre-market DollarVolume'.format(shares,str(key),sec.Price,shares*sec.Price,float(value)))
                              
            oshort = self.MarketOrder(key,-shares, asynchronous = True)
            # oshort = self.MarketOrder(key, -shares, asynchronous = True)
            self.openOrders.append(oshort)

    def OnOrderEvent(self, fill):
        # Short Order Went Through, Issue a Stop Loss
        if (fill.Status == OrderStatus.Filled or fill.Status == OrderStatus.PartiallyFilled) and fill.FillQuantity < 0 and not 'SPY' in str(fill.Symbol):
            stopo = self.StopMarketOrder(fill.Symbol, abs(fill.FillQuantity), round(fill.FillPrice*1.2,2))
            self.Log('Setting Stop Loss for {} with Stop {}'.format(fill.Symbol,round(fill.FillPrice*1.2,2)))
            self.stopOrders.append(stopo)
        
        if (fill.Status == OrderStatus.Canceled):
            self.Log('Order cancelled')
    
    # def CancelOrders(self):
    #     for o in self.openOrders:
    #         if o.Status == OrderStatus.PartiallyFilled or o.Status == OrderStatus.Submitted:
    #             o.Cancel('Short Order for {} could not be filled completely till 9:32, cancelling'.format(o.Symbol))
    #             self.openOrders.remove(o)
    
    def CoverShorts(self):
        invested = [ x.Symbol for x in self.Portfolio.Values if x.Invested ]
        
        for s in invested:
            if not str(s) == 'SPY':
                self.Liquidate(s)
            
    def CoarseSelectionFunction(self, coarse):
        
        filtered = [x for x in coarse if x.HasFundamentalData
                                      and  5000000 > x.DollarVolume > 1000000
                                      and  5 > x.Price > 0.1]

        # sort the stocks by dollar volume and take the top 500
        top = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)[:self.numberOfSymbolsCoarse]
        
        self.Log('Selecting Universe coarsely ...')

        self.symbols = [ i.Symbol for i in top ]
        
        return self.symbols
    
    def FineSelectionFunction(self, fine):
        # Tries to avoid Short Squeezes
        
        self.Log('Selecting Universe finely ...')
        
        filtered = [x for x in fine if x.EarningReports.BasicAverageShares.ThreeMonths > 100000]
        return [f.Symbol for f in filtered]