from QuantConnect.Data.UniverseSelection import *
import pandas as pd
import numpy as np
from scipy import stats

import decimal as d
import random

class BasicTemplateAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.Debug("%%%%%%%%%%%%% Initialize")
        self.my_universe = []
        
        # Set the initial cash, start, and end dates for backtesting
        self.mycash = 30000
        self.mycash = 5000
        self.cash_per_trade = 5000
        self.SetCash(self.mycash)
        self.SetStartDate(2011,1,4)
        self.SetEndDate(2011,1,5)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin) # Or override account type
        
        # Subscribe to data for securities returned in my universe 
        self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.MinimumTimeInUniverse = 0 
        self.AddUniverse(self.universe_filter_course, self.universe_filter_fine)
        
        # Schedule the 'trade' method to run
        # Add SPY to our data subscription so we can schedule from it (use defaults)

        self.Schedule.On(self.DateRules.EveryDay(),
                         self.TimeRules.At(9, 31),
                         Action(self.setup_trade))
        self.Schedule.On(self.DateRules.EveryDay(),
                         self.TimeRules.At(9, 32),
                         Action(self.my_stocks))
                         
                         
        self.Schedule.On(self.DateRules.EveryDay(),
                         self.TimeRules.At(15, 50),
                         Action(self.liquidate_portfolio))
        
        self.Schedule.On(self.DateRules.EveryDay(),
                         self.TimeRules.At(11, 00),
                         Action(self.cancel_orders))
                         
        overlayPlot = Chart("Overlay Plot")
        ZscorePlot = Chart("ZScore")
        self.AddChart(overlayPlot)
        self.AddChart(ZscorePlot)
        self.contractsAdded = set()
        
        #self.Schedule.On(
        #    self.DateRules.EveryDay('SPY'), 
        #    self.TimeRules.AfterMarketOpen('SPY', 1),       
        #    Action(self.my_trade))
    
    def universe_filter_fine(self, fine_data):
        columns = ['SecurityType', 'ShareClassStatus', 'ExchangeId', 'IsDepositaryReceipt', 'IsPrimaryShare']
        
        data_df = pd.DataFrame.from_records(
                [[getattr(s.SecurityReference, name) for name in columns] for s in fine_data],
                #[[s.ValuationRatios.PERatio, s.ValuationRatios.SalesYield] for s in fine_data],
                index=[s.Symbol for s in fine_data],
                columns=columns,)
        data_df = data_df.query("(SecurityType == 'ST00000001') and (IsDepositaryReceipt == False) and (IsPrimaryShare)")
        
        #data_df = data_df.sort_values(data_df.columns[0], ascending = True)
        data_df.sort_index(inplace=True)
        #data_df = data_df[:10]
        #self.Debug(data_df)
        
        #self.Debug("data is: {}".format(data_df.EVToEBITDA.values))
        self.universe = data_df.index.tolist()
        self.my_universe = data_df.index.tolist()
        
        #self.Debug("fine size is: {}".format(len(self.universe)))
        return self.universe

    def universe_filter_course(self, coarse_data):
        # With data_df already initialized globally, this should now just
        #   populate securities for values matching query instead.
        columns = ['Price', 'DollarVolume', 'HasFundamentalData', 'Volume']
        
        data_df = pd.DataFrame.from_records(
                [[getattr(s, name) for name in columns] for s in coarse_data],
                index=[s.Symbol for s in coarse_data],
                columns=columns,
                coerce_float=True)
                
        data_df = data_df.sort_values(by='Volume', ascending=False)

        data_df = data_df.query("HasFundamentalData " \
            "and (DollarVolume > @data_df.DollarVolume.quantile(.99))" \
            "and (DollarVolume < @data_df.DollarVolume.quantile(1.00))")
        '''
        data_df = data_df.query("HasFundamentalData ")
        data_df = data_df.sort_values('DollarVolume')
        data_df = data_df[data_df.DollarVolume > 1000000]
        data_df = data_df[50:100]
        '''
        self.StocksDailyData =data_df
        if len(data_df) == 0:
            return
        self.my_universe = data_df.index.tolist()
        
        return data_df.index.tolist()

    def setup_trade(self):
        # See how many securities we have in our Securities list
        '''
        for universe in self.UniverseManager.Keys:
            self.Debug("universe name: {}".
                format(universe))
        '''    
        self.Debug("%%%%%%%%%%%%% setup_trade {}".format(self.Time))
        self.sym_buylist = []
        self.otm_calls = set()
        self.buyList = []
        self.stock_list = []
        self.contract = []
        
        self.addedStocks = []
        
        #for universe in self.UniverseManager.Values:
            #self.Debug("universe count: {}".
            #    format(universe.Members.Count))
        if len(self.my_universe) == 0:
            #self.Debug('Empty universe')
            return
        for stock in self.my_universe:
            stockID = self.AddEquity(stock.Value, Resolution.Minute)
            stockID.SetDataNormalizationMode(DataNormalizationMode.Raw)
            
            self.addedStocks.append(stockID)
    def my_stocks(self):
        self.Debug("%%%%%%%%%%%%% my_stocks {}".format(self.Time))
        
        for stockID in self.addedStocks:
            if self.Securities[stockID.Symbol].Price != 0:
                self.stock_list.append(stockID.Symbol)
            self.Debug(stockID.Symbol)
            self.Debug(self.Securities[stockID.Symbol].Price)
            
    def liquidate_portfolio(self):
        
        #self.Debug('Checking Holdings 2222222222222...')
        for kvp in self.Portfolio:
            holdings = kvp.Value
            symbol = holdings.Symbol
            
            #if (holdings.Invested):
                #self.Debug("Symbol: {} -> Price: {}".format(symbol, holdings.AveragePrice))
                
                
                
        for symbol in self.Portfolio.Keys:
            self.SetHoldings(symbol, 0)

      
    def cancel_orders(self):
        #self.Debug('=====================================================================')
        openOrders = self.Transactions.GetOpenOrders()
        if len(openOrders)> 0:
            for x in openOrders:
                 self.Transactions.CancelOrder(x.Id)
################https://www.quantconnect.com/docs/algorithm-reference/overview

def returnNotMatches(a, b):
    return [[x for x in a if x not in b], [x for x in b if x not in a]]
    
def to_dataframe(data_c, properties, labels):
    # Helper function to make a dataframe from the coarse object. 
    # Then we can use all the handy dataframe methods.
    # Set the 'coerce_float' parameter to get meaningful datatypes
    # otherwise all the fields will be un-useful generic objects.

    data = [[getattr(stock, name) for name in properties] for stock in data_c]
    symbols = [stock.Symbol for stock in data_c ]
    
    data_df = pd.DataFrame.from_records(
            data, 
            index=symbols, 
            columns=labels, 
            coerce_float=True)
            
    return data_df