import numpy as np
import pandas as pd
from QuantConnect.Data import *
from QuantConnect.Data.Custom.Estimize import *


class QuantumTachyonGearbox(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 1)  # Set Start Date
        self.SetCash(1000000)  # Set Strategy Cash
        
        self.SetBenchmark("SPY")

        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        
        self.UniverseSettings.Resolution = Resolution.Minute
        self.SetUniverseSelection(CoarseFundamentalUniverseSelectionModel(self.CoarseSelectionFunction))
        
        self.symbolDataBySymbol = {}
        self.collection = InsightCollection()
        self.all_insights = InsightCollection()

    def CoarseSelectionFunction(self, coarse):
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
    
        return [ x.Symbol for x in sortedByDollarVolume if x.HasFundamentalData ][75:100]


    def CheckExpiredInsights(self):
        expiredInsights = self.all_insights.RemoveExpiredInsights(self.UtcTime)
        self.EmitInsights([Insight.Price(x.Symbol, timedelta(1), InsightDirection.Flat) for x in expiredInsights])
        orders = [self.Liquidate(x.Symbol) for x in expiredInsights]

    def OnData(self, data):
        
        self.CheckExpiredInsights()
 
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if data.ContainsKey(symbolData.release):
                bar = data[symbolData.release]
                if bar is None:
                    continue
                consensus = bar.ConsensusWeightedEpsEstimate
                actual = bar.Eps

                if (actual is None) or (consensus is None):
                    continue
                self.Log(f'{self.Time} :: Release data for {symbolData.Symbol.Value} :: EPS Actual {actual} / EMS Consensus {consensus}')
                insight = Insight.Price(symbol, timedelta(2) if actual-consensus < 0 else timedelta(1), InsightDirection.Down if actual-consensus < 0 else InsightDirection.Up)
                self.collection.Add(insight)
                self.all_insights.Add(insight)

        to_return = [x for x in self.collection]
        self.collection.Clear()
        self.EmitInsights(to_return)
        invested = max(len(to_return) + len([x for x in self.Portfolio.Keys if self.Portfolio[x].Invested]),1)
        orders = [self.SetHoldings(insight.Symbol, 1/(2*invested) if insight.Direction == InsightDirection.Up else -1/(2*invested)) for insight in to_return]

    
    def OnSecuritiesChanged(self, changes):
        symbols = [x.Symbol for x in changes.AddedSecurities]
        for symbol in symbols:
            symbolData = SymbolData(symbol, self)
            symbolData.Initialize(symbol, self)
            self.symbolDataBySymbol[symbol] = symbolData
        
        symbols = [x.Symbol for x in changes.RemovedSecurities]
        for symbol in symbols:
            if symbol in self.symbolDataBySymbol.keys():
                self.symbolDataBySymbol.pop(symbol)


        
class SymbolData:
    
    def __init__(self, symbol, algorithm):
        self.Symbol = symbol
        self.release = None

    def Initialize(self, symbol, algorithm):
        self.release = algorithm.AddData(EstimizeRelease, symbol).Symbol