calculate_taxes = True
if calculate_taxes:
    from taxes import CalculateTaxes
    
class UncoupledMultidimensionalAutosequencers(QCAlgorithm):

    def Initialize(self):
        yr=int(self.GetParameter('yr'))
        self.SetStartDate(2015, 2, 1)  
        # self.SetEndDate(2020, 3, 31)
        self.SetCash(1000000) 
        
        tickers = [('TMF', 1/3), ('VTV', 1/2), ('TQQQ', 1/6), ('UPRO', 1/6)]#, ('SDS', 0.25)]#, ('TLT', 0.5)]#, ('TLT', 0)]
        # tickers = [('SPXU', -1.5/3)]#, ('UPRO', 0.5/2), ('TMF', 0.5/3)]
        benchmark = 'QQQ'
        s = self.AddEquity(benchmark, Resolution.Minute).Symbol
        self.SetBenchmark(benchmark)
       
        self.my_universe = {}
        
        for t, w in tickers:
            s = self.AddEquity(t, Resolution.Minute).Symbol
            self.my_universe[s] = (s, w, 0, 0, 0)
            
        trade_schedule = [150]#, 90, 150, 210, 270, 360]
        
        for sc in trade_schedule:
            self.Schedule.On(self.DateRules.EveryDay(benchmark), self.TimeRules.AfterMarketOpen(benchmark, sc), self.Trade)
        self.rebalance = self.Time
        
        if calculate_taxes:
            self.calculateTaxes = CalculateTaxes()
            self.calculateTaxes.MyInitialize(self)
            
        self.Schedule.On(self.DateRules.MonthEnd(), self.TimeRules.BeforeMarketClose(benchmark, 30), self.EOM)
    
    
    def OnData(self, slice):
        if calculate_taxes:
            self.calculateTaxes.OnData(slice)
            
    def OnOrderEvent(self, orderEvent):
        if calculate_taxes:
            self.calculateTaxes.OnOrderEvent(orderEvent)
            
    def EOM(self):
        if calculate_taxes:
            self.calculateTaxes.CalculateUnrealizedTaxableGains()
            
    def OnEndOfAlgorithm(self):
        if calculate_taxes:
            self.calculateTaxes.CalculateUnrealizedTaxableGains()
            
    def Trade(self):
        
        if  self.rebalance > self.Time:
            return
        
        ttl=0
        for s, ow, lastPrice,  lpmin, lpmx  in self.my_universe.values():
            currentPrice = self.Securities[s].Price
            # w=ow
            # pct = 1- lastPrice/currentPrice
            # if currentPrice < lastPrice:
                
            # elif currentPrice  > lastPrice:
            #     pct = 1- lpmin/currentPrice
            # if pct_min >-0.025 :
            #     pct=pct_min
            # elif pct_mx >0.025:
            #     pct = pct_mx
            # else:
            # pct =(pct_min,-pct_mx)
            # if abs(pct) >0.025  :
            #     w=ow*1+pct*60
            #     w=max(0,min(1.5 -ttl, w))
            #     self.my_universe[s] = (s, w, currentPrice, currentPrice, currentPrice  )
            #     # self.SetHoldings(s, w)
            # elif currentPrice < lpmin :
            #     self.my_universe[s] = (s, w, currentPrice, currentPrice, lpmx )
            #     # self.SetHoldings(s, w)
            # elif currentPrice > lpmx :
            #     self.my_universe[s] = (s, w, currentPrice, lpmin, currentPrice )
            self.SetHoldings(s, ow)
                
                
            # ttl+=w
             
        self.rebalance = self.Time + timedelta(days=30)

import pandas as pd
from scipy import stats
from System import *
import math
import numpy as np  
import traceback
import sys
import statistics 
from datetime import datetime
from statistics import mode
from operator import itemgetter 
class CalculateTaxes:
    def MyInitialize(self, algo, longTermFederalTaxRate=0.2, netInvestmentTaxRate=0.038, \
            shortTermFederalTaxRate = 0.37, longTermStateTaxRate=0.0575, shortTermStateTaxRate =0.0575):
        self.a = algo
        self.longTermFederalTaxRate = longTermFederalTaxRate
        self.netInvestmentTaxRate = netInvestmentTaxRate
        self.longTermStateTaxRate = longTermStateTaxRate
        self.shortTermFederalTaxRate = shortTermFederalTaxRate
        self.shortTermStateTaxRate = shortTermStateTaxRate
        self.totalLongTermTaxRate = self.longTermFederalTaxRate + self.netInvestmentTaxRate + self.longTermStateTaxRate
        self.totalShortTermTaxRate = self.shortTermFederalTaxRate + self.shortTermStateTaxRate
        self.transactionDictLifo = {}
        
        self.totalLongTermTaxesOwedOnLifoBasis = 0
        self.totalShortTermTaxesOwedOnLifoBasis = 0
        
        self.unrealizedTaxableGains = 0
        self.realizedTaxableGains = 0
        self.slice = None
        
        
    def OnData(self, slice):
        self.slice = slice
        # self.ProcessDividends()
        # self.ProcessSplits()
        
    def ProcessDividends(self):
        if self.slice is None:
            return
        if len(self.slice.Dividends) ==0 :
            return
        invested_stocks = [(x.Key, x.Value.Quantity) for x in self.Portfolio if x.Value.Invested and  x.Key.SecurityType == SecurityType.Equity and x.Value.Quantity > 0 ]
        
        for s, q in invested_stocks:
            if self.slice.Dividends.ContainsKey(s):
                amount_received = self.slice.Dividends[s].Distribution * q
                #assume all of it is qualified
                tax = amount_received * self.totalLongTermTaxRate
                self.totalLongTermTaxesOwedOnLifoBasis += tax
                
                self.realizedTaxableGains += amount_received
                
    
    def ProcessSplits(self):
        if self.slice is None:
            return
        if len(self.slice.Splits) ==0 :
            return
        invested_stocks = [(x.Key, x.Value.Quantity) for x in self.Portfolio if x.Value.Invested and  x.Key.SecurityType == SecurityType.Equity and x.Value.Quantity != 0 ]
        
        for s, q in invested_stocks:
            t = s.Value
            if not self.slice.Splits.ContainsKey(t):
                continue
            spySplit = self.slice.Splits[t]
            if spySplit.Type == 0:
                self.myLog(f'{t} stock will split next trading day', 25)
            if spySplit.Type == 1:
                self.myLog(f"{t} Split type: {spySplit.Type}, Split factor: {spySplit.SplitFactor}, Reference price: {spySplit.ReferencePrice}", 25)
                
                transactionList = self.transactionDictLifo.get(s)
                if transactionList is None:
                    continue
                
                updatedTransactionList = []
                
                for s3, t, q, p in transactionList:
                    q = q / spySplit.SplitFactor
                    p = p * spySplit.SplitFactor
                    transaction = (s3, t, q, p)
                    updatedTransactionList.append(transaction)
                self.transactionDictLifo[s] = updatedTransactionList
    
    
    
    def OnOrderEvent(self, orderEvent):
        # order = self.Transactions.GetOrderById(orderEvent.OrderId)
        if orderEvent.Status != OrderStatus.Filled and orderEvent.Status != OrderStatus.PartiallyFilled : 
            return
        
        s = orderEvent.Symbol
        self.AddOrderToTransactions(orderEvent)
        
        self.CalculateTaxesOnTheLatestTransaction(s)
         
    
    def update_display(self):    
        self.a.SetRuntimeStatistic("Unrealized Gain (LIFO)", "${:,.0f}".format(self.unrealizedTaxableGains))
        # self.SetRuntimeStatistic("Realized Gain(LIFO)", "${:,.0f}".format(self.realizedTaxableGains))
        
        self.a.SetRuntimeStatistic("Taxes Owed", "${:,.0f}".format(self.totalLongTermTaxesOwedOnLifoBasis + self.totalShortTermTaxesOwedOnLifoBasis))
        
        # unrealizedgainseries.AddPoint(self.Time, self.unrealizedTaxableGains)
        # taxesowedseries.AddPoint(self.Time, (self.totalLongTermTaxesOwedOnLifoBasis) + (self.totalShortTermTaxesOwedOnLifoBasis))
        
        # self.Plot('Strategy Equity', "Unrealized Gain (LIFO)", self.unrealizedTaxableGains)
        # self.Plot('Strategy Equity', "Realized Gain(LIFO)", self.realizedTaxableGains)
        
        # self.Plot('Strategy Equity', "Taxes Owed", (self.totalLongTermTaxesOwedOnLifoBasis) + (self.totalShortTermTaxesOwedOnLifoBasis))
        pass
        
    def AddOrderToTransactions(self, orderEvent):
        
        q = orderEvent.FillQuantity
        p =  orderEvent.FillPrice
        s = orderEvent.Symbol
        
        if s.SecurityType == SecurityType.Option:
            q = q*100
            
        latestTransaction = (s, self.a.Time, q, p)
        
        transactionList = self.transactionDictLifo.get(s)
        if transactionList is None:
            transactionList  = [latestTransaction]
        else:
            transactionList.append(latestTransaction)
            
        self.transactionDictLifo[s] = transactionList
    
    def CalculateUnrealizedTaxableGains(self):
        
        self.unrealizedTaxableGains = 0
        for s, tls in self.transactionDictLifo.items():
            
            latest = self.a.Securities[s].Price
            for s, t, q, p in tls:
                profit = (latest - p) *q
                self.unrealizedTaxableGains += profit
        self.update_display() 
        
    def CalculateTaxesOnTheLatestTransaction(self, s):
        transactionList = self.transactionDictLifo.get(s)
        
        n = len(transactionList) - 1
        if n < 1:
            return 
        updatedTransactionList =[]
        latestTransaction = transactionList[-1]
        remainingQtyToBeAdressed=latestTransaction[2]
        
        totalLongTermTaxesOwedOnTheTransaction, totalShortTermTaxesOwedOnTheTransaction = 0, 0 
        for i in range(n):
            index = n - i-1
            prevTransaction = transactionList[index]
            
            if remainingQtyToBeAdressed != 0:
                longTermTaxesOwed, shortTermTaxesOwed, qtyAddressedByTheTransaction, remainingQtyToBeAdressed = self.CalculateTaxesOnThePair(latestTransaction, prevTransaction)
                latestTransaction = (latestTransaction[0], latestTransaction[1], remainingQtyToBeAdressed, latestTransaction[3])
                prevTransaction = (prevTransaction[0], prevTransaction[1],  prevTransaction[2] - qtyAddressedByTheTransaction, prevTransaction[3])
                totalLongTermTaxesOwedOnTheTransaction += longTermTaxesOwed
                totalShortTermTaxesOwedOnTheTransaction += shortTermTaxesOwed
            
        
            if prevTransaction[2] != 0:
                updatedTransactionList =[ prevTransaction ]+ updatedTransactionList
        
        if remainingQtyToBeAdressed != 0:
            updatedTransactionList.append(latestTransaction)
            
        self.transactionDictLifo[s] = updatedTransactionList
        self.totalLongTermTaxesOwedOnLifoBasis += totalLongTermTaxesOwedOnTheTransaction
        self.totalShortTermTaxesOwedOnLifoBasis += totalShortTermTaxesOwedOnTheTransaction
        
            
    def CalculateTaxesOnThePair(self, latestTransaction, prevTransaction):
        
        latestTransactionQuantity = latestTransaction[2]
        
         
        latestFillPrice = latestTransaction[3]
        prevFillPrice = prevTransaction[3]
        
        prevTransactionQuantity = prevTransaction[2]
        longTermTaxesOwed, shortTermTaxesOwed, qtyAddressedByTheTransaction, remainingQtyToBeAdressed = 0, 0, 0, latestTransactionQuantity
        
        if prevTransactionQuantity * latestTransactionQuantity < 0:

            abs_qtyAddressedByTheTransaction = min(abs(prevTransactionQuantity), abs(latestTransactionQuantity))
            
            qtyAddressedByTheTransaction = -abs_qtyAddressedByTheTransaction* latestTransactionQuantity/abs(latestTransactionQuantity)
            profit = (latestFillPrice - prevFillPrice)* qtyAddressedByTheTransaction
            
            latestDate = latestTransaction[1]
            pevDate = prevTransaction[1]
            
            numOfYears = (latestDate-pevDate).days/365
            
            if numOfYears >= 1:
                longTermTaxesOwed = profit * self.totalLongTermTaxRate
            else:
                shortTermTaxesOwed = profit * self.totalShortTermTaxRate
        
            remainingQtyToBeAdressed = latestTransactionQuantity + qtyAddressedByTheTransaction
            self.realizedTaxableGains += profit 
            
        return longTermTaxesOwed, shortTermTaxesOwed, qtyAddressedByTheTransaction, remainingQtyToBeAdressed