from System.Collections.Generic import List
from QuantConnect.Data.UniverseSelection import *
import operator
from math import ceil,floor
import numpy as np
import pandas as pd
from scipy import stats

class AltmanZScoreAlgorithm(QCAlgorithm):
    def Initialize(self):

        self.SetStartDate(2016,01,01)  #Set Start Date
        #self.SetEndDate(2017,05,02)    #Set End Date
        self.SetCash(100000)            #Set Strategy Cash
        self.flag1 = 1
        self.flag2 = 0
        self.flag3 = 0
    
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.UniverseSettings.Resolution = Resolution.Daily        
        self.UniverseSettings.Leverage   = 1
        
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.AddEquity("SPY")
        self.__numberOfSymbols = 100
        self.__numberOfSymbolsFine = 10

        self._changes = SecurityChanges.None
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY"), Action(self.Rebalancing))
        self.SetWarmUp(200)
        
        self.splotName = 'Strategy Info'
        sPlot = Chart(self.splotName)
        sPlot.AddSeries(Series('Leverage',  SeriesType.Line, 0))
        self.AddChart(sPlot)


    def CoarseSelectionFunction(self, coarse):
        if self.flag1:
            CoarseWithFundamental = [x for x in coarse if (x.HasFundamentalData) and (float(x.Price) > 5)]
            sortedByDollarVolume = sorted(CoarseWithFundamental, key=lambda x: x.DollarVolume, reverse=True) 
            top = sortedByDollarVolume[:self.__numberOfSymbols]
            
            self.prices = {}
            for x in top:
                self.prices[x.Symbol] = x.Price
                
            return [i.Symbol for i in top]
        else:
            return []


    def FineSelectionFunction(self, fine):
        if self.flag1:
            self.flag1 = 0
            self.flag2 = 1
            filtered_fine = [x for x in fine if x.FinancialStatements.BalanceSheet.TotalAssets.TwelveMonths and
                             x.FinancialStatements.BalanceSheet.TotalLiabilities.TwelveMonths and
                             x.FinancialStatements.BalanceSheet.WorkingCapital.TwelveMonths and
                             x.FinancialStatements.BalanceSheet.RetainedEarnings.TwelveMonths and
                             x.FinancialStatements.IncomeStatement.EBIT.TwelveMonths and
                             x.FinancialStatements.IncomeStatement.TotalRevenue.TwelveMonths and
                             x.EarningReports.BasicAverageShares.TwelveMonths]
                             
            sortedByfactor1 = [x for x in filtered_fine if ZScore(x.FinancialStatements.BalanceSheet.TotalAssets.TwelveMonths,
                             x.FinancialStatements.BalanceSheet.TotalLiabilities.TwelveMonths,
                             x.FinancialStatements.BalanceSheet.WorkingCapital.TwelveMonths,
                             x.FinancialStatements.BalanceSheet.RetainedEarnings.TwelveMonths,
                             x.FinancialStatements.IncomeStatement.EBIT.TwelveMonths,
                             x.FinancialStatements.IncomeStatement.TotalRevenue.TwelveMonths,
                             x.EarningReports.BasicAverageShares.TwelveMonths,self.prices[x.Symbol]).ObjectiveScore() > 1.81]

            filtered_finer = [x for x in sortedByfactor1 if x.ValuationRatios.EVToEBITDA]
    
            sortedByfactor2 = sorted(filtered_finer, key=lambda x: x.ValuationRatios.EVToEBITDA, reverse=False)

            topFine = sortedByfactor2[:self.__numberOfSymbolsFine]

            self.flag3 = self.flag3 + 1            
            
            return [i.Symbol for i in topFine]

        else:
            return []


    def OnData(self, data):
        if self.flag3 > 0:
            if self.flag2 == 1:
                self.flag2 = 0
                for security in self._changes.RemovedSecurities:
                    if security.Invested:
                        self.Liquidate(security.Symbol)
                 
                for security in self._changes.AddedSecurities:
                    self.SetHoldings(security.Symbol, 1./float(self.__numberOfSymbolsFine))    

                self._changes = SecurityChanges.None;

    # this event fires whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
        self._changes = changes
    
    def Rebalancing(self):
        self.flag1 += 1
        self.account_leverage = self.Portfolio.TotalAbsoluteHoldingsCost / self.Portfolio.TotalPortfolioValue
        self.Plot(self.splotName,'Leverage', float(self.account_leverage))

        
    
class ZScore(object):
    def __init__(self, totalassets, totalliabilities, workingcapital, retainedearnings, ebit, totalrevenue, shares, price):
        self.totalassets = float(totalassets)
        self.totalliabilities = float(totalliabilities)
        self.workingcapital = float(workingcapital)
        self.retainedearnings = float(retainedearnings)
        self.ebit = float(ebit)
        self.totalrevenue = float(totalrevenue)
        self.shares = float(shares)
        self.price = float(price)
    
    def ObjectiveScore(self):
        X1 = 1.2 * (self.workingcapital / self.totalassets)
        X2 = 1.4 * (self.retainedearnings / self.totalassets)
        X3 = 3.3 * (self.ebit / self.totalassets)
        X4 = 0.6 * ((self.shares * self.price) / self.totalliabilities)
        X5 = 1.0 * (self.totalrevenue / self.totalassets)
        return X1 + X2 + X3 + X4 + X5