import pandas as pd

def normalise(series, equal_ls=True):
    if equal_ls:
        series -= series.mean()
    sum = series.abs().sum()
    return series.apply(lambda x: x/sum)
    

class ValueAlphaModel():
    
    def __init__(self):
        pass
    
    def GenerateAlphaScores(self, algorithm, securities):
        # algorithm.Log(f"Generating alpha scores for {len(securities)} securities...")
        
        fcf_y = pd.DataFrame.from_records(
            [
                {
                    'symbol': security.Symbol,
                    'fcf_y': security.ValuationRatios.CashReturn
                } for security in securities
            ])
        
        fcf_y.set_index('symbol', inplace=True)
        
        fcf_y['alpha_score'] = normalise(fcf_y['fcf_y'], True)
        
        return fcf_y

from universe_selection import FactorUniverseSelectionModel
from alpha_model import ValueAlphaModel
from portfolio_construction import OptimisationPortfolioConstructionModel
from execution import Execution

class TradingBot(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2003, 1, 1)
        self.SetCash(100000)
        
        # Data resolution
        self.UniverseSettings.Resolution = Resolution.Minute
        
        # Universe selection model
        self.securities = []
        self.CustomUniverseSelectionModel = FactorUniverseSelectionModel(self)
        self.AddUniverse(self.CustomUniverseSelectionModel.SelectCoarse, self.CustomUniverseSelectionModel.SelectFine)
        
        # Alpha model
        self.CustomAlphaModel = ValueAlphaModel()
        
        # Portfolio construction model
        self.CustomPortfolioConstructionModel = OptimisationPortfolioConstructionModel()
        
        # Execution model
        self.CustomExecution = Execution()
        
        # Add SPY for trading days data
        self.AddEquity('SPY', Resolution.Daily)
        
        # Schedule rebalancing
        self.Schedule.On(self.DateRules.EveryDay('SPY'), self.TimeRules.At(13, 0), Action(self.RebalancePortfolio))

    def OnData(self, data):
        pass
    
    def RebalancePortfolio(self): 
        alpha_df = self.CustomAlphaModel.GenerateAlphaScores(self, self.securities)
        portfolio = self.CustomPortfolioConstructionModel.GenerateOptimalPortfolio(self, alpha_df)
        self.CustomExecution.ExecutePortfolio(self, portfolio)

import pandas as pd


class OptimisationPortfolioConstructionModel():

    def __init__(self):
        pass

    def GenerateOptimalPortfolio(self, algorithm, alpha_df):
        # algorithm.Log("Generating target portfolio...")
        
        alpha_portfolio = self.CalcAlphaPortfolio(algorithm, alpha_df)
        
        optimal_portfolio = self.Optimise(algorithm, self.AddZeroHoldings(algorithm, alpha_portfolio))
        
        # algorithm.Log(f"Created a portfolio of {len(optimal_portfolio[optimal_portfolio != 0])} securities (liquidating {len(optimal_portfolio[optimal_portfolio == 0])} securities)")
        
        return optimal_portfolio

    def CalcAlphaPortfolio(self, algorithm, alpha_df):
        portfolio = alpha_df['alpha_score']
        portfolio.name = 'weight'
        port_sum = portfolio.abs().sum()
        if port_sum != 1:
            # algorithm.Log(f"Alpha scores don't add up to 1: {port_sum}")
            portfolio /= port_sum
        return portfolio

    def AddZeroHoldings(self, algorithm, portfolio):
        zero_holding_securities = [s.Symbol for s in algorithm.Portfolio.Values if s.Invested and s.Symbol not in portfolio.index]
        for security in zero_holding_securities:
            portfolio.loc[str(security)] = 0
        return portfolio
    
    def Optimise(self, algorithm, portfolio):
        return portfolio

class FactorUniverseSelectionModel():
    
    def __init__(self, algorithm):
        self.algorithm = algorithm
    
    def SelectCoarse(self, coarse):
        # self.algorithm.Log("Generating universe...")
        universe = self.FilterDollarPriceVolume(coarse)
        return [c.Symbol for c in universe]

    def SelectFine(self, fine):
        universe = self.FilterFactor(self.FilterFinancials(fine))
        # self.algorithm.Log(f"Universe consists of {len(universe)} securities")
        self.algorithm.securities = universe
        return [f.Symbol for f in universe]
    
    def FilterDollarPriceVolume(self, coarse):
        filter_dollar_price = [c for c in coarse if c.Price > 1]
        sorted_dollar_volume = sorted([c for c in filter_dollar_price if c.HasFundamentalData], key=lambda c: c.DollarVolume, reverse=True)
        return sorted_dollar_volume[:1000]

    def FilterFinancials(self, fine):
        filter_financials = [f for f in fine if f.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.FinancialServices]
        return filter_financials
    
    def FilterFactor(self, fine):
        filter_factor = sorted(fine, key=lambda f: f.ValuationRatios.CashReturn, reverse=True)
        return filter_factor[:20] + filter_factor[-20:]

class Execution():
    
    def __init__(self):
        pass
    
    def ExecutePortfolio(self, algorithm, portfolio):
        # algorithm.Log(f"Executing portfolio trades...")
        
        liquidate_securities = portfolio[portfolio == 0].index
        holding_port = portfolio[portfolio != 0]
        
        self.LiquidateSecurities(algorithm, liquidate_securities)
        
        self.SetPortfolioHoldings(algorithm, holding_port)
    
    def LiquidateSecurities(self, algorithm, securities):
        # algorithm.Log(f"Liquidating {len(securities)} securities...")
        for security in securities:
            algorithm.Liquidate(security)
        # algorithm.Log(f"Successfully liquidated {len(securities)} securities")
    
    def SetPortfolioHoldings(self, algorithm, portfolio):
        # algorithm.Log(f"Setting portfolio holdings for {len(portfolio)} securities...")
        for security, weight in portfolio.iteritems():
            algorithm.SetHoldings(security, weight)
        # algorithm.Log(f"Successfully set all holdings")