from AlgorithmImports import *


class MultiAssetManager(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2002, 1, 1)
        self.SetEndDate(2022, 1, 1)
        self.SetCash(100000)
        self.resolution = int(self.GetParameter("Resolution"))
        self.equities = self.GetParameter("Equities").split(",")
        self.forexPairs = self.GetParameter("ForexPairs").split(",")
        self.cryptoPairs = self.GetParameter("CryptoPairs").split(",")
        for equity in self.equities:
            if len(equity) < 3:
                continue
            self.AddEquity(equity, self.resolution)
        for forex in self.forexPairs:
            if len(forex) < 3:
                continue
            self.AddForex(forex, self.resolution)
        for crypto in self.cryptoPairs:
            if len(crypto) < 3:
                continue
            self.AddCrypto(crypto, self.resolution)

        allocationRatioVariable = self.GetParameter(
            "AllocationRatio").split(",")
        self.allocationRatio = {}
        allocationSum = 0
        for allocation in allocationRatioVariable:
            [pair, ratio] = allocation.split(":")
            self.allocationRatio[pair] = float(ratio)
            allocationSum += float(ratio)
        if allocationSum > 1:
            raise Exception('Allocation Sum is more than 1')
        else:
            self.Log("Not allocated%:"+str(1-allocationSum))
        self.preparePlotPrices()

        schedulePair = list(self.allocationRatio.keys())[0]
        self.Schedule.On(self.DateRules.MonthStart(),  
                 self.TimeRules.AfterMarketOpen(schedulePair, 10), 
                 self.rebalanceEveryMonth)

    def OnData(self, data):
        self.plotPrices(data)
        self.plotAssets()

    def rebalanceEveryMonth(self):
      self.rebalanceAllocations()

    def preparePlotPrices(self):
        pricesChart = Chart('Prices')
        for pair in self.allocationRatio.keys():
            pricesChart.AddSeries(
                Series(pair, SeriesType.Line, "$"))
        self.AddChart(pricesChart)

    def plotPrices(self, data):
        for pair in self.allocationRatio.keys():
            if pair not in data:
                continue
            if data[pair]:
              price = data[pair].Value
              self.Plot("Prices", pair, price)
    
    def preparePlotAssets(self):
        assetsChart = Chart('Assets')
        for pair in self.allocationRatio.keys():
            assetsChart.AddSeries(
                Series(pair, SeriesType.Line, "$"))
        self.AddChart(assetsChart)

    def plotAssets(self):
        for pair in self.allocationRatio.keys():
            if pair not in self.Portfolio:
                continue
            self.Plot("Assets", pair, self.Portfolio[pair].Quantity * self.Portfolio[pair].Price)


    def rebalanceAllocations(self):
      portfolioTargers = []
      for pair in self.allocationRatio.keys():
        portfolioTargers.append(PortfolioTarget(pair,self.allocationRatio[pair]))
      self.SetHoldings(portfolioTargers)

    # def rebalanceAllocations(self,data):
    #   portfolioTargers = []
    #   for pair in self.allocationRatio.keys():
    #     portfolioTargers.append(PortfolioTarget(pair,self.allocationRatio[pair]))
        
    #   self.SetHoldings([PortfolioTarget("SPY", 0.8), PortfolioTarget("IBM", 0.2)])
    #   allocations = []
    #   portfolioValue = self.Portfolio.TotalPortfolioValue
    #   for pair in self.allocationRatio.keys():
    #         if not self.Portfolio[pair].Invested:
    #             continue
    #         allocations.append([pair,portfolioValue / (self.Portfolio[pair].Price * self.Portfolio[pair].Quantity)])
      
    #   allocations.sort(key=self.allocationSortKey,reverse=True)
    #   self.SetHoldings()


    # def allocationSortKey(element):
    #       return element[1]