from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Algorithm.Framework")
AddReference("QuantConnect.Common")

from QuantConnect import *
from QuantConnect.Parameters import *
from QuantConnect.Benchmarks import *
from QuantConnect.Brokerages import *
from QuantConnect.Util import *
from QuantConnect.Interfaces import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Selection 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.Indicators import *
from QuantConnect.Data import *
from QuantConnect.Data.Consolidators import *
from QuantConnect.Data.Custom import *
from QuantConnect.Data.Fundamental import *
from QuantConnect.Data.Market import *
from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Notifications import *
from QuantConnect.Orders import *
from QuantConnect.Orders.Fees import *
from QuantConnect.Orders.Fills import *
from QuantConnect.Orders.Slippage import *
from QuantConnect.Scheduling import *
from QuantConnect.Securities import *
from QuantConnect.Securities.Equity import *
from QuantConnect.Securities.Forex import *
from QuantConnect.Securities.Interfaces import *
from datetime import date, datetime, timedelta
from QuantConnect.Python import *
from QuantConnect.Storage import *
QCAlgorithmFramework = QCAlgorithm
QCAlgorithmFrameworkBridge = QCAlgorithm

import calendar


class FirstDaySPY(QCAlgorithmFramework):
    """Buy on close on the last day of the month and sell on the close of the next day"""
    def Initialize(self):
        self.SetStartDate(2020, 1, 28)  # Set Start Date
        self.SetEndDate(2020, 9, 5)  
        self.SetCash(100000)  # Set Strategy Cash

        # Parameters
        self.maximumDuration = 1  # maximum duration to keep trades open

        # Universe
        self.UniverseSettings.Resolution = Resolution.Daily
        # symbols = [ Symbol.Create("SPY", SecurityType.Equity, Market.USA) ]
        self.AddEquity("SPY", Resolution.Daily)
        # self.AddUniverseSelection( ManualUniverseSelectionModel(symbols) )
        # self.AddUniverseSelection(CoarseFundamentalUniverseSelectionModel(self.SelectCoarse))
        self.AddAlpha(FirstDayAlphaModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel()) 
        self.AddRiskManagement(FixedDurationTrades(self.maximumDuration))

        # TODO: verify this, I want to buy on close at the end of the last day and sell on close of the next day
        self.SetExecution(ImmediateExecutionModel())

    def SelectCoarse(self, coarse):
        self.Debug(f"In universe selection time: {self.Time}")
        # TODO: smooth that
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        isETF = [x.Symbol for x in sortedByDollarVolume if not x.HasFundamentalData]
        self.Debug(f"Total coarse stocks: {len(sortedByDollarVolume)} and total nonFundamental: {len(isETF)}")
        return isETF[:10]

    def OnSecuritiesChanged(self, changes):
        self._changes = changes
        self.Debug(f"OnSecuritiesChanged({self.UtcTime}):: {changes}")
        self.Log(f"OnSecuritiesChanged({self.UtcTime}):: {changes}")

    # def OnData(self, data):
    #     """OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
    #         Arguments:
    #             data: Slice object keyed by symbol containing the stock data
    #     """
    #     # if not self.Portfolio.Invested:
    #     #     self.SetHoldings("SPY", 1)
    #     #     self.Debug("Purchased Stock")



class FixedDurationTrades(RiskManagementModel):
    def __init__(self, maximumDuration = 1):
        """Provides an implementation of IRiskManagementModel that limits the maximum number of periods
        that we can stay in a trade"""
        self.maximumDuration = maximumDuration
        self.investedTime = {}

    def ManageRisk(self, algorithm, targets):
        riskAdjustedTargets = list()

        for kvp in algorithm.Securities:
            symbol = kvp.Key
            security = kvp.Value

            # Remove if not invested
            if not security.Invested:
                self.investedTime.pop(symbol, None)
                continue

            # Add newly invested securities
            if symbol not in self.investedTime:
                self.investedTime[symbol] = 0   # Set to 0 periods invested
                continue

            self.investedTime[symbol] += 1
            if self.investedTime[symbol] >= self.maximumDuration:
                # liquidate securities that have reached max duration
                riskAdjustedTargets.append(PortfolioTarget(symbol, 0.))
                algorithm.Debug(f"Setting liquidation target for {symbol}")
        return riskAdjustedTargets


class FirstDayAlphaModel(AlphaModel):
    Name = "FirstDayAlphaModel"

    def Update(self, algorithm, slice):
         # Updates this alpha model with the latest data from the algorithm.
         # This is called each time the algorithm receives data for subscribed securities
         # Generate insights on the securities in the universe.

        insights = []
        DateRules.MonthEnd
        algorithm.Debug(f"Inside Alpha Model at time: {algorithm.Time}")
        
        symbols = list(algorithm.ActiveSecurities.Keys)
        algorithm.Debug(f"Active symbols: {[s.Value for s in symbols]}")

        # insight expiry time and direction
        insightExpiry = Expiry.EndOfDay(algorithm.Time)
        insightDirection = InsightDirection.Up  # insight direction

        # TODO: weight by how many often this worked in the past
        for symbol in symbols:
            # add insights if it's the last tradable day of the month
            if list(algorithm.DateRules.MonthEnd(symbol, 0).GetDates(algorithm.Time, algorithm.Time)):
                algorithm.Debug(f"Last Day of the month, Emit insight for symbol {symbol.Value}")
                insights.append(Insight.Price(symbol, insightExpiry, insightDirection,
                                              None, None, None, None))
        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
         # Handle security changes in from your universe model.
         pass