from io import StringIO

import pandas as pd

data_url = "https://www.dropbox.com/s/skbmxw8yhv6rump/trump_beta.csv?dl=1"

class CalibratedTransdimensionalProcessor(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 6, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.symbols = self.get_symbols()
        
        self.AddUniverseSelection(
           ManualUniverseSelectionModel(self.symbols)
        )
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.AddAlpha(TrumpBetaDiversificationModel(self))
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(rebalance=Resolution.Daily, portfolioBias=PortfolioBias.LongShort))
        self.Settings.RebalancePortfolioOnInsightChanges = False
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        
        self.SetExecution(ImmediateExecutionModel())
    
    
    def get_symbols(self):
        constituents = pd.read_csv(StringIO(self.Download(data_url)), index_col="Date").columns
        return [Symbol.Create(s, SecurityType.Equity, Market.USA) for s in constituents]


class TrumpBetaDiversificationModel:
    
    def __init__(self, algorithm):
        self.thresh = 0.2
        self.df = pd.read_csv(StringIO(algorithm.Download(data_url)), index_col="Date")
        self.df.index = pd.to_datetime(self.df.index)

    def Update(self, algorithm, slice):
        insights = []
        
        self.df.index = pd.to_datetime(self.df.index)
        
        if algorithm.Time not in self.df.index: return []
        
        trump_beta = self.df.loc[algorithm.Time]
        
        low_trump_beta = [security for security, beta in trump_beta.iteritems() if abs(beta) < self.thresh and security in [str(s) for s in slice.keys()]]
        
        for security in trump_beta.keys():
            insight = Insight(security, timedelta(1), InsightType.Price, InsightDirection.Up)
            insights.append(insight)
        
        return insights
    
    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            algorithm.AddData(TrumpBeta, added.Symbol)


class TrumpBeta(PythonData):
    
    def __init__(self):
        self.columns = {}

    def GetSource(self, config, date, isLive):
        return SubscriptionDataSource(data_url, SubscriptionTransportMedium.RemoteFile);


    def Reader(self, config, line, date, isLive):
        data = line.split(',')
        if not (line.strip() and line[0].isdigit()):
            self.columns = {data[i]: i for i in range(0, len(data) - 1)}
            return None
        
        ticker = str(config.Symbol).split('.')[0]
        
        if ticker not in self.columns: return None
        
        trump_beta = TrumpBeta()
        trump_beta.Symbol = config.Symbol
        trump_beta.Time = pd.to_datetime(data[self.columns["Date"]]) + timedelta(days=10) # Make sure we only get this data AFTER trading day - don't want forward bias.
        
        value = data[self.columns[ticker]]
        
        if not value: return None
        
        trump_beta.Value = float(value)

        return trump_beta