#region imports
from AlgorithmImports import *
#endregion
from datetime import timedelta
import numpy as np
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Brokerages import *
from QuantConnect.Data import *
from QuantConnect.Data.Market import *
from QuantConnect.Orders import *
from QuantConnect.Securities import *
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Algorithm.Framework.Portfolio import *
from QuantConnect.Algorithm.Framework.Execution import *
from QuantConnect.Algorithm.Framework.Risk import *
from QuantConnect.Algorithm.Framework.Selection import *

class MeanReversionAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 1, 1)
        self.SetEndDate(2022, 1, 31)
        self.SetCash(100000)
        self.SetBrokerageModel(BrokerageName.QuantConnectBrokerage, AccountType.Margin)
        self.UniverseSettings.Resolution = Resolution.Minute
        self.SetAlpha(MeanReversionAlphaModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.01))

        self.Log("Algorithm initialized.")

        # Add symbols to the universe
        # self.symbol = self.AddEquity("SPY", Resolution.Minute)
        self.symbolDataBySymbol = {}
        # self.symbol = self.AddEquity("SPY", Market.USA).Symbol
        # self.symbol = self.AddEquity("SPY").Symbol
        # self.symbol = self.AddFutureContract(Futures.Grains.SoybeanOil).Symbol
        self.symbol = self.AddFuture("ZL", Resolution.Minute)

    def OnData(self, data):
        self.Log("Received new data.")

        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if not data.Bars.ContainsKey(symbol) or not symbolData.IsReady:
                continue
            price = data.Bars[symbol].Close
            symbolData.Update(price)
            if symbolData.IsOversold():
                insights.append(Insight.Price(symbol, timedelta(minutes=1), InsightDirection.Up))
            elif symbolData.IsOverbought():
                insights.append(Insight.Price(symbol, timedelta(minutes=1), InsightDirection.Down))
        self.EmitInsights(insights)

    def OnSecuritiesChanged(self, changes):
        for security in changes.RemovedSecurities:
            symbolData = self.symbolDataBySymbol.pop(security.Symbol, None)
            if symbolData is not None:
                symbolData.Dispose()
        addedSymbols = [x.Symbol for x in changes.AddedSecurities if x.Symbol not in self.symbolDataBySymbol]
        history = self.History(addedSymbols, 20, Resolution.Minute)
        for symbol in addedSymbols:
            symbolData = SymbolData(symbol, 20, history)
            self.symbolDataBySymbol[symbol] = symbolData

    def OnOrderEvent(self, orderEvent):
        self.Log(f"Order event: {orderEvent}")

    def OnEndOfAlgorithm(self):
        self.Log("Algorithm ended.")

class MeanReversionAlphaModel(AlphaModel):
    def __init__(self, lookback=20):
        self.lookback = lookback
        self.symbolDataBySymbol = {}

    def Update(self, algorithm, data):
        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if not data.Bars.ContainsKey(symbol) or not symbolData.IsReady:
                continue
            price = data.Bars[symbol].Close
            symbolData.Update(price)
            if symbolData.IsOversold():
                insights.append(Insight.Price(symbol, timedelta(minutes=1), InsightDirection.Up))
            elif symbolData.IsOverbought():
                insights.append(Insight.Price(symbol, timedelta(minutes=1), InsightDirection.Down))
        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        for security in changes.RemovedSecurities:
            symbolData = self.symbolDataBySymbol.pop(security.Symbol, None)
            if symbolData is not None:
                symbolData.Dispose()
        addedSymbols = [x.Symbol for x in changes.AddedSecurities if x.Symbol not in self.symbolDataBySymbol]
        history = algorithm.History(addedSymbols, self.lookback, Resolution.Minute)
        for symbol in addedSymbols:
            symbolData = SymbolData(symbol, self.lookback, history)
            self.symbolDataBySymbol[symbol] = symbolData

class SymbolData:
    def __init__(self, symbol, lookback, history):
        self.symbol = symbol
        self.lookback = lookback
        self.history = history
        self.prices = []
        self.IsReady = False

    def Update(self, price):
        self.prices.append(price)
        if len(self.prices) > self.lookback:
            self.prices.pop(0)
            self.IsReady = True

    def IsOversold(self):
        if not self.IsReady:
            return False
        return self.prices[-1] < np.mean(self.prices) - 2 * np.std(self.prices)

    def IsOverbought(self):
        if not self.IsReady:
            return False
        return self.prices[-1] > np.mean(self.prices) + 2 * np.std(self.prices)

    def Dispose(self):
        pass