#region imports
from AlgorithmImports import *
#endregion

def InitCharts(self):
    
    self.sectordict = {'Technology' : 311, 'BasicMaterials' : 101, 'ConsumerCyclical' : 102, 'FinancialServices' : 103, 'RealEstate' : 104, 'ConsumerDefensive' : 205, 'Healthcare' : 206, 'Utilities' : 207, 'CommunicationServices' : 308, 'Energy' : 309, 'Industrials' : 310}
    
    sector_plot = Chart('Sector Performance')
    for sector in self.sectordict.keys():
        sector_plot.AddSeries(Series(sector, SeriesType.Line, 0)) # Option for stacked area instead is below
        #sector_plot.AddSeries(Series(sector, SeriesType.StackedArea, 0))
    self.AddChart(sector_plot)

    sectorweight_plot = Chart('Sector Weightings')
    for sector in self.sectordict.keys():
        sectorweight_plot.AddSeries(Series(sector, SeriesType.Line, 0, "%"))
    self.AddChart(sectorweight_plot)

def PlotSectorChart(self):

    for sector in list(self.sectordict.keys()):
        valueinsector = []
        for security, v in self.Portfolio.items():   
            if v.Invested:
                if self.dictionary[security.Value].AssetClassification.MorningstarSectorCode == self.sectordict[sector]:
                    valueinsector.append(v.NetProfit + v.UnrealizedProfit) #total net and unrealized profit of invested security   
            if v.LastTradeProfit != 0 and v.Invested == False: 
                if self.dictionary[security.Value].AssetClassification.MorningstarSectorCode == self.sectordict[sector]:
                    valueinsector.append(v.NetProfit) #total net profit of security that was once invested but is no longer invested
                    
        Sectorprofitvalue = sum([f for f in valueinsector])
        if Sectorprofitvalue != 0: #net profit is -1 for first month because fees. Graph looks better starting from the origin
            self.Plot('Sector Performance', f'{sector}', Sectorprofitvalue)
        else:
            pass

def PlotSectorWeightingsChart(self):

    for sector in list(self.sectordict.keys()):
            weightsinsector=[]
            for security, v in self.Portfolio.items():
                if v.Invested:
                    if self.dictionary[security.Value].AssetClassification.MorningstarSectorCode == self.sectordict[sector]:
                        weightsinsector.append(100*v.AbsoluteHoldingsValue/self.Portfolio.TotalHoldingsValue)
            Totalweightingofsector = sum([abs(f) for f in weightsinsector])
            self.Plot('Sector Weightings', f'{sector}', Totalweightingofsector)
    

        

#region imports
from AlgorithmImports import *
#endregion
from universe_selection import FactorUniverseSelectionModel
from charting import InitCharts, PlotSectorChart, PlotSectorWeightingsChart

class TradingBot(QCAlgorithm):

    def Initialize(self): 

        self.SetStartDate(2020, 9, 1)
        self.SetEndDate(2023, 9, 1)
        self.SetCash(100000)
        
        self.SetWarmUp(14, Resolution.Daily)
        
        # Data resolution
        self.UniverseSettings.Resolution = Resolution.Minute

        # Universe selection model
        self.securities = [] 
        self.dictionary = {} # Dictionary to hold fundamental data of past and present holdings. This is used rather than pulling from Securities dictionary as fine fundamental data is not always available for all stocks on everyday so bugs are avoided this way  
        self.CustomUniverseSelectionModel = FactorUniverseSelectionModel(self) 
        self.AddUniverse(self.CustomUniverseSelectionModel.SelectCoarse, self.CustomUniverseSelectionModel.SelectFine)
        
        # Add SPY for trading days data
        self.AddEquity('SPY', Resolution.Daily)
        self.SetBenchmark("SPY")

        # Schedule rebalancing
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Friday), self.TimeRules.At(12,15,0), Action(self.Execute))
        
        #Initiate charting
        InitCharts(self)

        #Schedule charting once per week
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Friday), self.TimeRules.BeforeMarketClose("SPY", 0), Action(self.PlotCharts))
                
    def OnData(self, data):
        self.slice = data

    def PlotCharts(self):
        if self.IsWarmingUp:
            pass 
        else:
            PlotSectorChart(self)
            PlotSectorWeightingsChart(self)
        
    def Execute(self):
        if self.IsWarmingUp:
            pass 
        else:
            self.Liquidate()
            self.dictionary.update({f.Symbol.Value : f for f in self.securities})
            portfoliodataframe = pd.DataFrame.from_records( 
                [
                    {
                        'symbol': security.Symbol,
                        'weight': 1/len(self.securities) 
                    } for security in self.securities
                ]).set_index('symbol')
            self.Log(portfoliodataframe)
            self.Log(f"Setting portfolio holdings for {len(portfoliodataframe.index)} securities...")
            for security, weight in portfoliodataframe.iterrows():
                if self.slice.Bars.ContainsKey(security) and self.slice[security] is not None:
                    self.SetHoldings(security, weight)
            self.Log(f'BREAK - Successfully set all holdings')

  
        
        

#region imports
from AlgorithmImports import *
#endregion

import numpy as np

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(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 5 < c.Price] 
        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 FilterFactor(self, fine):
        stocksinsector = [f for f in fine if f.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.FinancialServices]
        sortedbycashreturn = sorted(stocksinsector, key=lambda f: f.ValuationRatios.CashReturn, reverse=True)
        universe = sortedbycashreturn[:10] 
        return universe