from AlgorithmImports import *

#lean project-create --language python "Test"
#lean cloud backtest "Test" --push --open
#lean backtest Test

from datetime import timedelta
from decimal import Decimal

class Test(QCAlgorithm):
    def Initialize(self):
        self.SetCash(100000)  # Set Strategy Cash

        self.SetStartDate(2014, 3, 24)
        self.SetEndDate(2014,3, 27)

        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.TestFilter)
        self.AddAlpha(TestAlphaModel(self))

        self.SetPortfolioConstruction(TestPortfolioConstructionModel(portfolioBias = PortfolioBias.Long))

        self.Settings.RebalancePortfolioOnInsightChanges = True
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        self.SetExecution(TestExecutionModel()) 

        #self.SetRiskManagement(CompositeRiskManagementModel(
        #    TrailingStopRiskManagementModel(0.1),
        #    MaximumDrawdownPercentPerSecurity()
        #))


    def TestFilter(self, universe):
        selected = []
        
        universe = [c for c in universe if c.HasFundamentalData] 
        universe = sorted(universe, key=lambda c: c.DollarVolume, reverse=True)[100:103]

        for coarse in universe:
            selected.append(coarse.Symbol)
            self.Log(f"[LOG->ADDED] {self.Time}: {coarse.Symbol}")
            
        return selected

    def OnOrderEvent(self, orderEvent: OrderEvent) -> None:
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        if orderEvent.Status == OrderStatus.Filled:
            self.Log(f"[LOG->ORDER] -> {self.Time}: {order.Type}: {orderEvent}")

class TestAlphaModel(AlphaModel):
    def __init__(self, algo):
        self.algo = algo

    def Update(self, algo, slice):
        insights = []

        for symbol in slice.Keys:
            if slice.ContainsKey(symbol) and slice[symbol] is not None:
                if (algo.Portfolio[symbol].Invested):
                    insights.append(Insight.Price(symbol, timedelta(days=1), InsightDirection.Flat, None, None, None, 0.05))
                    algo.Log(f"[LOG->INSIGHT] Flat -> {algo.Time}: {symbol}")
                else:
                    insights.append(Insight.Price(symbol, timedelta(days=2), InsightDirection.Up, None, None, None, 0.05))
                    algo.Log(f"[LOG->INSIGHT] Up -> {algo.Time}: {symbol}")
        return insights

    def OnSecuritiesChanged(self, algo, changes):
        algo.Log(f"Changes ({algo.Time}):: {changes}")


class TestPortfolioConstructionModel(PortfolioConstructionModel):
    insightCollection = InsightCollection()

    def __init__(self, portfolioBias = PortfolioBias.Long):
        super().__init__()
        self.portfolioBias = portfolioBias
        self.SetRebalancingFunc(Resolution.Daily)


    def CreateTargets(self, algo, insights):
        targets = []

        if (insights.Length > 0):
            self.insightCollection.AddRange(insights)

        percents = self.DetermineTargetPercent(insights)

        for insight in insights:
            percent = float(percents.get(insight))
            algo.Log(f"CreateTargets Add -> {insight.Symbol} {insight.Direction} {percent} ")

            if (insight.Direction == InsightDirection.Up):
                target = PortfolioTarget.Percent(algo, insight.Symbol, percents.get(insight))
                if target != None:
                    targets.append(target)

        expiredInsights = self.insightCollection.RemoveExpiredInsights(algo.UtcTime);

        for insight in expiredInsights:
            algo.Log(f"CreateTargets Expired -> {insight.Symbol} {insight.Direction} ")
            target = PortfolioTarget.Percent(algo, insight.Symbol, 0)
            if target != None and algo.Portfolio[insight.Symbol].Invested:
                targets.append(target)

        return targets

    def DetermineTargetPercent(self, activeInsights):
        result = {}

        count = sum(x.Direction != InsightDirection.Flat and self.RespectPortfolioBias(x) for x in activeInsights)
        percent = 0 if count == 0 else 1.0 / count
        for insight in activeInsights:
            result[insight] = (insight.Direction if self.RespectPortfolioBias(insight) else InsightDirection.Flat) * insight.Weight
        return result

    def RespectPortfolioBias(self, insight):
        return self.portfolioBias == PortfolioBias.LongShort or insight.Direction == self.portfolioBias

class TestExecutionModel(ExecutionModel):
    def __init__(self):
        self.targetsCollection = PortfolioTargetCollection()

    def Execute(self, algo, targets):
        self.targetsCollection.AddRange(targets)
        if not self.targetsCollection.IsEmpty:
            for target in self.targetsCollection.OrderByMarginImpact(algo):
                security = algo.Securities[target.Symbol]
                quantity = OrderSizing.GetUnorderedQuantity(algo, target, security)
                if quantity != 0:
                    aboveMinimumPortfolio = BuyingPowerModelExtensions.AboveMinimumOrderMarginPortfolioPercentage(security.BuyingPowerModel, security, quantity, algo.Portfolio, algo.Settings.MinimumOrderMarginPortfolioPercentage)
                    if aboveMinimumPortfolio:
                        algo.MarketOrder(security, quantity)

            self.targetsCollection.ClearFulfilled(algo)