import math
import pandas as pd
from io import StringIO

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

class CalibratedTransdimensionalProcessor(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 6, 1)
        self.SetCash(100000)
        
        self.symbols = self.get_symbols()
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverseSelection(ManualUniverseSelectionModel(self.symbols))
        
        self.AddAlpha(TrumpBetaDiversificationModel())
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        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):
        self.thresh = 1/4

    def Update(self, algorithm, slice):
        insights = []
        
        ### Explain how to use custom data properly. 
        trump_beta = pd.Series({k: v.Value for k, v in slice.Get[TrumpBeta]().items()})
        
        if len(trump_beta) < 250:
            return insights
        
        # slice["IBM"].
        # @alexcatarino
        
        low_trump_beta = abs(trump_beta).sort_values()[:math.floor(self.thresh*len(trump_beta))]
        
        for security in low_trump_beta.keys():
        #for security in trump_beta.keys(): # This is for the benchmark
            insight = Insight(security, timedelta(7), InsightType.Price, InsightDirection.Up)
            insights.append(insight)
        
        return insights
    
    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            algorithm.AddData(TrumpBeta, added.Symbol)
        for removed in changes.RemovedSecurities:
            algorithm.RemoveSecurity(removed.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))}
            return None
        
        trump_beta = TrumpBeta()
        trump_beta.Symbol = config.Symbol
        ## Explain this -- 
        trump_beta.EndTime = pd.to_datetime(data[self.columns["Date"]]) + timedelta(days=1)
        
        value = data[self.columns[config.Symbol.Value]]
        if not value: return None
        trump_beta.Value = float(value)

        return trump_beta