from alpha import LongShortCBOPAlpha
from universe import BestWorstCBOPLargeCap

class CashBasedOperatingProfitability(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 8, 1)
        self.SetEndDate(2015, 10, 15)
        self.SetCash(1000000)
        
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverseSelection(BestWorstCBOPLargeCap())
        
        self.AddAlpha(LongShortCBOPAlpha())
        
        self.Settings.RebalancePortfolioOnInsightChanges = False
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel(lambda time: None))
        
        self.SetExecution(ImmediateExecutionModel())
        
        self.SetBrokerageModel(BrokerageName.AlphaStreams)

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class BestWorstCBOPLargeCap(FundamentalUniverseSelectionModel):
    month = -1
    
    def __init__(self, decile_denominator=10):
        self.decile_denominator = decile_denominator
        super().__init__(True)


    def SelectCoarse(self, algorithm, coarse):
        if self.month == algorithm.Time.month:
            return Universe.Unchanged
        self.month = algorithm.Time.month
        
        return [x.Symbol for x in coarse if x.HasFundamentalData and x.Price > 5]
        
    
    def SelectFine(self, algorithm, fine):
        fine = [x for x in fine if (x.EarningReports.BasicAverageShares.ThreeMonths != 0 and \
                                    x.EarningReports.BasicEPS.TwelveMonths != 0 and \
                                    x.ValuationRatios.PERatio != 0 and \
                                    x.FinancialStatements.IncomeStatement.TotalRevenue.TwelveMonths != 0 and \
                                    x.FinancialStatements.IncomeStatement.SellingAndMarketingExpense.TwelveMonths != 0 and \
                                    x.FinancialStatements.IncomeStatement.CostOfRevenue.TwelveMonths != 0 and \
                                    x.FinancialStatements.IncomeStatement.GeneralAndAdministrativeExpense.TwelveMonths != 0)]
        
        if len(fine) == 0: 
            return []
        
        market_cap_by_symbol = {}
        cbop_by_symbol = {}
        
        # Market cap and cbop calc.
        for stock in fine:
            symbol = stock.Symbol
            
            market_cap_by_symbol[symbol] = float(stock.CompanyProfile.MarketCap)
            
            income_statement = stock.FinancialStatements.IncomeStatement
            cbop_by_symbol[symbol] = income_statement.TotalRevenue.TwelveMonths - \
                                     income_statement.SellingAndMarketingExpense.TwelveMonths - \
                                     income_statement.CostOfRevenue.TwelveMonths - \
                                     income_statement.GeneralAndAdministrativeExpense.TwelveMonths
                                     
        # Sorting by market cap and cbop
        top_market_cap = sorted(market_cap_by_symbol.items(), key = lambda x:x[1], reverse=True)[:int(len(market_cap_by_symbol)/2)]
        
        sorted_by_cbop = sorted(top_market_cap, key = lambda x:cbop_by_symbol[x[0]], reverse=True)
        decile = int(len(sorted_by_cbop)/self.decile_denominator)
        longs = [symbol for symbol, _ in sorted_by_cbop[:decile]]
        shorts = [symbol for symbol, _ in sorted_by_cbop[-decile:]]

        return longs + shorts

from dateutil.relativedelta import relativedelta

class LongShortCBOPAlpha(AlphaModel):
    symbol_data_by_symbol = {}
    month = -1

    def Update(self, algorithm, data):
        if data.Time.month == self.month:
            return []
        self.month = data.Time.month
        
        # Filter for available data
        has_data = {symbol: symbol_data for symbol, symbol_data in self.symbol_data_by_symbol.items() if data.ContainsKey(symbol)}
        
        # Split into long/short from CBOP
        sorted_by_cbop = sorted(has_data.items(), key=lambda x: x[1].CBOP, reverse=True)
        half = int(len(sorted_by_cbop) / 2)
        longs = [ symbol for symbol, _ in sorted_by_cbop[:half] ]
        shorts = [ symbol for symbol, _ in sorted_by_cbop[-half:] ]
        
        # Gather market cap info
        market_cap_by_symbol = {symbol: symbol_data.MarketCap for symbol, symbol_data in has_data.items()}
        total_market_cap = sum(market_cap_by_symbol.values())
        
        # Generate insights
        insights = []
        hold_duration = Expiry.EndOfMonth(data.Time)
        for symbol in longs:
            weight = market_cap_by_symbol[symbol] / total_market_cap
            insight = Insight.Price(symbol, hold_duration, InsightDirection.Up, None, None, None, weight)
            insights.append(insight)
        for symbol in shorts:
            weight = market_cap_by_symbol[symbol] / total_market_cap
            insight = Insight.Price(symbol, hold_duration, InsightDirection.Down, None, None, None, weight)
            insights.append(insight)
        
        return insights
        
    def OnSecuritiesChanged(self, algorithm, changes):
        for security in changes.AddedSecurities:
            self.symbol_data_by_symbol[security.Symbol] = SymbolData(security)
        
        for security in changes.RemovedSecurities:
            self.symbol_data_by_symbol.pop(security.Symbol, None)
            
class SymbolData:
    def __init__(self, security):
        self.security = security
        
    @property
    def MarketCap(self):
        return self.security.Fundamentals.CompanyProfile.MarketCap
    
    @property
    def CBOP(self):
        income_statement = self.security.Fundamentals.FinancialStatements.IncomeStatement
        return income_statement.TotalRevenue.TwelveMonths - \
                income_statement.SellingAndMarketingExpense.TwelveMonths - \
                income_statement.CostOfRevenue.TwelveMonths - \
                income_statement.GeneralAndAdministrativeExpense.TwelveMonths