def InitCharts(algorithm):
    performance_plot = Chart('Performance Breakdown')
    performance_plot.AddSeries(Series('Total Fees', SeriesType.Line, 0))
    performance_plot.AddSeries(Series('Total Gross Profit', SeriesType.Line, 0))
    algorithm.AddChart(performance_plot)

    exposure_plot = Chart('Exposure/Leverage')
    exposure_plot.AddSeries(Series('Gross', SeriesType.Line, 0))
    exposure_plot.AddSeries(Series('Net', SeriesType.Line, 0))
    algorithm.AddChart(exposure_plot)

    country_exposure_plot = Chart('Country Exposure')
    for etf, country in algorithm.etf_country.items():
        country_exposure_plot.AddSeries(Series(country, SeriesType.Line, 0))


def PlotPerformanceChart(algorithm):
    algorithm.Plot('Performance Breakdown', 'Total Fees', algorithm.Portfolio.TotalFees)
    algorithm.Plot('Performance Breakdown', 'Total Gross Profit', algorithm.Portfolio.TotalProfit)


def PlotExposureChart(algorithm):
    long_val = 0
    short_val = 0
    for security, v in algorithm.Portfolio.items():
        if v.Invested:
            val = v.AbsoluteHoldingsValue
            if v.IsLong:
                long_val += val
            elif v.IsShort:
                short_val += val

    total_equity = algorithm.Portfolio.TotalPortfolioValue
    gross = (long_val + short_val) / total_equity
    net = (long_val - short_val) / total_equity
    algorithm.Plot('Exposure/Leverage', 'Gross', gross)
    algorithm.Plot('Exposure/Leverage', 'Net', net)


def PlotCountryExposureChart(algorithm):
    for etf, country in algorithm.etf_country.items():
        exposure = algorithm.Securities[etf].Holdings.HoldingsValue / algorithm.Portfolio.TotalHoldingsValue
        algorithm.Plot('Country Exposure', country, exposure)

import pandas as pd
import numpy as np
import cvxpy as cv


class OptimisationPortfolioConstructionModel:

    def __init__(self, turnover, max_wt, longshort, mkt_neutral):
        self.turnover = turnover
        self.max_wt = max_wt
        self.longshort = longshort
        self.mkt_neutral = mkt_neutral

    def GenerateOptimalPortfolio(self, algorithm, alpha_df):
        # algorithm.Log("Generating target portfolio...")

        alphas = alpha_df['alpha_score']

        optimal_portfolio = self.Optimise(algorithm, self.AddZeroHoldings(algorithm, alphas))

        # algorithm.Log(f"Created a portfolio of {len(optimal_portfolio[optimal_portfolio != 0])} securities")

        return optimal_portfolio

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

    def Optimise(self, algorithm, alphas):
        invested_securities = [security for security in algorithm.Portfolio.Values if security.Invested]
        if len(invested_securities) == 0:
            algorithm.Log('Initial portfolio rebalance')
            self.initial_rebalance = True
            turnover = 1
            initial_portfolio = pd.DataFrame(columns=['symbol', 'weight', 'alpha']).set_index('symbol')
        else:
            self.initial_rebalance = False
            turnover = self.turnover
            initial_portfolio = pd.DataFrame.from_records(
                [
                    {
                        'symbol': str(security.Symbol),
                        'weight': security.HoldingsValue / algorithm.Portfolio.TotalHoldingsValue,
                        'alpha': alphas.loc[security] if security in alphas.index else 0,
                    } for security in invested_securities
                ]).set_index('symbol')
        for security, alpha in alphas.iteritems():
            if security not in initial_portfolio.index:
                initial_portfolio.loc[security, 'weight'] = 0
                initial_portfolio.loc[security, 'alpha'] = alpha

        for i in range(int(turnover * 100), 101, 1):
            to = i / 100
            optimiser = Optimiser(initial_portfolio, turnover=to, max_wt=self.max_wt, mkt_neutral=self.mkt_neutral)
            optimal_portfolio, optimisation_status = optimiser.optimise()
            if optimisation_status != 'optimal':
                algorithm.Log(f'Optimisation with {to} turnover not feasible: {optimisation_status}')
            else:
                break
        return optimal_portfolio


class Optimiser:

    def __init__(self, initial_portfolio, turnover, max_wt, longshort=True, mkt_neutral=True):
        self.symbols = np.array(initial_portfolio.index)
        self.init_wt = np.array(initial_portfolio['weight'])
        self.opt_wt = cv.Variable(self.init_wt.shape)
        self.alpha = np.array(initial_portfolio['alpha'])
        self.longshort = longshort
        self.mkt_neutral = mkt_neutral
        self.turnover = turnover
        self.max_wt = max_wt
        if self.longshort:
            self.min_wt = -self.max_wt
            self.net_exposure = 0
            self.gross_exposure = 1
        else:
            self.min_wt = 0
            self.net_exposure = 1
            self.gross_exposure = 1

    def optimise(self):
        constraints = self.get_constraints()
        optimisation = cv.Problem(cv.Maximize(cv.sum(self.opt_wt * self.alpha)), constraints)
        optimisation.solve()
        status = optimisation.status
        if status == 'optimal':
            optimal_portfolio = pd.Series(
                np.round(optimisation.solution.primal_vars[list(optimisation.solution.primal_vars.keys())[0]], 3),
                index=self.symbols)
        else:
            optimal_portfolio = pd.Series(np.round(self.init_wt, 3), index=self.symbols)
        return optimal_portfolio, status

    def get_constraints(self):
        min_wt = self.opt_wt >= self.min_wt
        max_wt = self.opt_wt <= self.max_wt
        turnover = cv.sum(cv.abs(self.opt_wt - self.init_wt)) <= self.turnover * 2
        net_exposure = cv.sum(self.opt_wt) == self.net_exposure
        gross_exposure = cv.sum(cv.abs(self.opt_wt)) <= self.gross_exposure
        if self.mkt_neutral:
            return [min_wt, max_wt, turnover, net_exposure, gross_exposure]
        else:
            return [min_wt, max_wt, turnover, gross_exposure]

import pandas as pd

from io import StringIO
from datetime import timedelta

from portfolio_construction import OptimisationPortfolioConstructionModel
from execution import Execution
from charting import InitCharts, PlotPerformanceChart, PlotExposureChart, PlotCountryExposureChart


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


class StockifySentiment(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 1)  # Set Start Date
        self.SetEndDate(2020, 5, 20)
        self.SetCash(100000)  # Set Strategy Cash

        # Weighting style - normalise or alpha_max (alpha maximisation w/ optimisation)
        self.weighting_style = 'normalise'

        # Market neutral
        self.mkt_neutral = True

        # Audio feature to use
        self.audio_feature = 'valence'

        # Get data
        self.data, self.etf_list, self.etf_country = self.DataSetup()

        # Add ETFs
        for etf in self.etf_list:
            self.AddEquity(etf, Resolution.Minute)

        # Portfolio construction model
        self.CustomPortfolioConstructionModel = OptimisationPortfolioConstructionModel(turnover=1, max_wt=0.2,
                                                                                       longshort=True,
                                                                                       mkt_neutral=self.mkt_neutral)

        # Execution model
        self.CustomExecution = Execution(liq_tol=0.005)

        # Schedule rebalancing
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Wednesday), self.TimeRules.BeforeMarketClose('IVV', 210),
                         Action(self.RebalancePortfolio))

        # Init charting
        InitCharts(self)

        # Schedule charting
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Wednesday), self.TimeRules.BeforeMarketClose('IVV', 5),
                         Action(self.PlotCharts))

    def OnData(self, data):
        pass

    def RebalancePortfolio(self):
        df = self.data.loc[self.Time - timedelta(7):self.Time].reset_index().set_index('symbol')
        
        if self.weighting_style == 'normalise':
            portfolio = normalise(df['alpha_score'], equal_ls=self.mkt_neutral)
        elif self.weighting_style == 'alpha_max':
            df = df[['alpha_score']]
            portfolio = self.CustomPortfolioConstructionModel.GenerateOptimalPortfolio(self, df)
        else:
            raise Exception('Invalid weighting style')
        
        self.CustomExecution.ExecutePortfolio(self, portfolio)

    def PlotCharts(self):
        PlotPerformanceChart(self)
        PlotExposureChart(self)
        PlotCountryExposureChart(self)

    def DataSetup(self):
        df = pd.read_csv(StringIO(
            self.Download('https://raw.githubusercontent.com/Ollie-Hooper/StockifySentiment/master/data/scores.csv')))
        data = df[['date', 'country', f's_{self.audio_feature}']].copy()
        data['date'] = pd.to_datetime(data['date'])
        data.rename(columns={f's_{self.audio_feature}': 'alpha_score'}, inplace=True)
        etf_df = pd.read_csv(StringIO(
            self.Download('https://raw.githubusercontent.com/Ollie-Hooper/StockifySentiment/master/data/etf.csv')))
        data = pd.merge(data, etf_df)
        data = data.sort_values('date')
        data.set_index(['date', 'symbol'], inplace=True)
        data = data[['alpha_score']]
        return data, etf_df['symbol'].to_list(), etf_df.set_index('symbol')['country_name'].to_dict()

class Execution():

    def __init__(self, liq_tol):
        self.liq_tol = liq_tol

    def ExecutePortfolio(self, algorithm, portfolio):
        # algorithm.Log(f"Executing portfolio trades...")

        liquidate_securities = portfolio[abs(portfolio) < self.liq_tol].index
        holding_port = portfolio[abs(portfolio) >= self.liq_tol]

        self.LiquidateSecurities(algorithm, liquidate_securities)

        self.SetPortfolioHoldings(algorithm, holding_port)

    def LiquidateSecurities(self, algorithm, securities):
        liquid_count = 0
        for security in securities:
            if algorithm.Securities[security].Invested:
                algorithm.Liquidate(security)
                liquid_count += 1
        # algorithm.Log(f"Successfully liquidated {liquid_count} 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")