#
# Original File:
# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect
# Corporation.
#
# Changes:
# The universe selection model is extended to take parameters as
# optional arguments.
# Ostirion SLU Copyright 2021
# Madrid, Spain
# Hector Barrio - hbarrio@ostirion.net.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at:
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from itertools import groupby
from math import ceil
from clr import AddReference
import numpy as np
from typing import List, Set, Tuple, Dict
AddReference("System")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm.Framework")


class FlexibleUniverseSelectionModel(FundamentalUniverseSelectionModel):

    '''
    Class representing a parametrically selected securities universe.
    Attributes:
        n_coarse (int): Number of securities in the coarse selection.
        n_fine (int): Number of securities in fine selection.
        age (int): Minimum time since IPO.
        recent (int): Maximum time from IPO.
        vol_lim (float): Minimum daily volume of each security.
        min_price (float): Minimum price of each security.
        max_price (float): Maximum price of each security.
        period (str): "Month" or "Day". Recalculate the universe every period.
        m_cap_lim (float): Minimum market cap of security to be considered.
        markets (list[str]): Markets in which the security trades.
        c_id (str): Code of the country of origin of securities.
        from_top (bool): Take the top (True) or bottom (False) volume securities.
        restrict_country (bool): Restrict the country of origin and market for securities.
        verbose (bool): False for silent, True for announcing size and components.
    '''

    def __init__(self: None,
                 n_coarse: int=1000,
                 n_fine: int=500,
                 age: int=1250,
                 recent: int=-1,
                 vol_lim: int=0,
                 min_price: int=0,
                 max_price: float=np.Inf,
                 period: str='Month',
                 m_cap_lim: float=5e8,
                 markets: List[str]=["NYS", "NAS"],
                 c_id: str='USA',
                 from_top: bool=True,
                 restrict_country: bool=True,
                 verbose: bool=False,
                 filterFineData: bool=True,
                 universeSettings: UniverseSettings=None,
                 securityInitializer: SecurityInitializer=None) -> None:

        super().__init__(filterFineData, universeSettings, securityInitializer)

        # Parameter settings:
        self.n_symbols_coarse = n_coarse
        self.n_symbols_fine = n_fine
        self.age = age
        self.recent = recent
        self.vol_lim = vol_lim
        self.min_price = min_price
        self.max_price = max_price
        self.period = period
        self.m_cap_lim = m_cap_lim
        self.markets = markets
        self.c_id = c_id
        self.reverse = from_top
        self.restrict_country = restrict_country
        self.verbose = verbose

        self.usd_vol = {}
        self.last_month = -1

    def SelectCoarse(self,
                     algorithm: QCAlgorithm,
                     coarse: CoarseFundamental) -> FineFundamental:

        '''
        Coarse unviverse selection method.
        Args:
            algorithm (QCAlgorithm): Current algorithm instance.
            coarse (CoarseFundamental): QC Coarse universe object.
        Returns:
            fine (FineFundamental): QC fine universe object.
        '''

        if self.period == 'Month':
            if algorithm.Time.month == self.last_month:
                return Universe.Unchanged
        elif self.period != 'Day':
            algoithm.Log('Period not valid.. Choose "Day" or "Month". Defaulting to "Month".')

        c = coarse
        usd_vol = sorted([x for x in c if
                          x.HasFundamentalData and
                          x.Volume > self.vol_lim and
                          self.max_price > x.Price > self.min_price],
                         key=lambda x: x.DollarVolume,
                         reverse=self.reverse)[:self.n_symbols_coarse]

        self.usd_vol = {x.Symbol: x.DollarVolume for x in usd_vol}

        if len(self.usd_vol) == 0:
            return Universe.Unchanged

        return list(self.usd_vol.keys())

    def SelectFine(self,
                   algorithm: QCAlgorithm,
                   fine: FineFundamental) -> FineFundamental:

        '''
        Coarse unviverse selection method.
        Args:
            algorithm (QCAlgorithm): Current algorithm instance.
            fine (FineFundamental): QC fine universe object.
        Returns:
            new_universe (FineFundamental): QC fine universe object.
        '''

        f = fine
        a = algorithm
        sort_sector = sorted([x for x in f if
                              x.MarketCap > self.m_cap_lim],
                             key=lambda x: x.CompanyReference.IndustryTemplateCode)

        count = len(sort_sector)

        if count == 0:
            return Universe.Unchanged

        if self.recent != -1:
            sort_sector = [x for x in sort_sector if
                           (a.Time -
                            x.SecurityReference.IPODate).days < self.recent]
        else:
            sort_sector = [x for x in sort_sector if
                           (a.Time -
                            x.SecurityReference.IPODate).days > self.age]

        if self.restrict_country:
            sort_sector = [x for x in sort_sector if
                           x.CompanyReference.CountryId == self.c_id and
                           x.CompanyReference.PrimaryExchangeID in self.markets]

        self.last_month = a.Time.month

        percent = self.n_symbols_fine / count
        sort_usd_vol = []

        for c, g in groupby(sort_sector,
                            lambda x: x.CompanyReference.IndustryTemplateCode):
            y = sorted(g, key=lambda x: self.usd_vol[x.Symbol],
                       reverse=self.reverse)
            c = ceil(len(y) * percent)
            sort_usd_vol.extend(y[:c])

        sort_usd_vol = sorted(sort_usd_vol,
                              key=lambda x: self.usd_vol[x.Symbol],
                              reverse=self.reverse)
        new_universe = [x.Symbol for x in sort_usd_vol[:self.n_symbols_fine]]

        if self.verbose:
            for s in new_universe:
                algorithm.Log('Adding: '+str(s.Symbol))
            algorithm.Log('Universe members: ' + str(len(new_universe)))

        return new_universe

'''
 *******************************************************
 Copyright (C) 2021 Ostirion SLU
 Madrid, Spain
 Hector Barrio <hbarrio@ostirion.net>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.

 *******************************************************
'''
import numpy as np
from FlexibleUniverseSelectionModel import FlexibleUniverseSelectionModel as fusm


class CorrAtTop(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 5, 1)
        self.SetEndDate(datetime.today())
        self.SetCash(1000000)

        self.SetBrokerageModel(AlphaStreamsBrokerageModel())

        res = Resolution.Daily
        self.market = self.AddEquity('SPY', res).Symbol
        
        # Add TIPS:
        self.tip = self.AddEquity('TIP', res).Symbol

        # Get Parameters:
        # Universe size that conforms the top:
        try:
            self.n_stocks = int(self.GetParameter("n_stocks"))
        except:
            self.Log('Defaulting stocks parameter.')
            self.n_stocks = 15
    
        # Period to compute correlation
        try:
            self.corr_period = int(self.GetParameter("corr_period"))
        except:
            self.Log('Defaulting period parameter.')
            self.corr_period = 60
    
        # Minimum correlation for "sell" signal:
        try:
            self.min_corr = float(self.GetParameter("min_corr"))
        except:
            self.Log('Defaulting minimum correlation parameter.')
            self.min_corr = 0.2
    
        # Risk control parameters:
        try:
            self.rc = float(self.GetParameter("risk_factor"))
        except:
            self.Log('Defaulting risk parameter.')
            self.rc = 0.03

        self.UniverseSettings.Resolution = res
        universe = fusm(n_fine=self.n_stocks)
        self.AddUniverseSelection(universe)

        self.SetRiskManagement(TrailingStopRiskManagementModel(self.rc))

        self.AddAlpha(CorrAtTopAlphaModel(self.market,
                                          self.tip,
                                          self.corr_period,
                                          self.min_corr,
                                          self.rc))

        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())

        self.SetExecution(ImmediateExecutionModel())


class CorrAtTopAlphaModel(AlphaModel):
    """
    """

    def __init__(self, market, tip, corr_period, min_corr, risk):

        self.symbol_data = {}
        self.market = market
        self.tip = tip

        # These are taken from parameter grid:
        self.period = corr_period
        self.min_corr = min_corr
        self.mom_vol_period = 22
        self.mom_limit = 0
        self.vol_limit = 0.01

        # Normal, non-parameter variables:
        self.Name = 'Correlation at Top'
        self.fut_ret = risk  # Future returns are not calculated.
        self.counter = False
        self.refresh = 2

    def Update(self, algorithm, data):

        insights = []
        if not data:
            return []

        symbols = data.keys()

        if not self.counter or self.counter % self.refresh == 0:
            if self.market in symbols:
                # Compute Market Momentum and Volatilty:
                market_price = algorithm.History(self.market,
                                      self.mom_vol_period,
                                      Resolution.Daily).unstack(level=0)['close']
                self.vol, self.mom = self.compute_vol_mom(market_price)
                symbols.remove(self.market)
                symbols.remove(self.tip)

            price = algorithm.History(symbols,
                                      self.period,
                                      Resolution.Daily).unstack(level=0)['close']
            self.corr = price.corr().mean().mean()
            
            # Debugs and plots:
            '''
            algorithm.Debug(str(len(symbols)))
            algorithm.Debug(str(self.vol))
            algorithm.Debug(str(self.mom))
            algorithm.Plot("corr", "Correlation", self.corr)
            algorithm.Plot("mom", "Momentum", self.mom)
            algorithm.Plot("vol", "Volatility", self.vol)
            '''

        # Inelegant counter, to be replaced by
        # timer.
        self.counter += 1
        p = timedelta(days=self.refresh)
        condition = self.corr < self.min_corr or \
                    self.mom < self.mom_limit or \
                    self.vol > self.vol_limit

        if condition:
            algorithm.Debug('Correlation: '+str(self.corr))
            direction = InsightDirection.Up
            algorithm.Debug('Low Correlation, dropping positions.')
            insights.append(Insight(self.tip, p, InsightType.Price,
                                direction, self.fut_ret, 1, self.Name, 1))
        else:
            direction = InsightDirection.Up

        insights.append(Insight(self.market, p, InsightType.Price,
                                direction, self.fut_ret, 1, self.Name, 1))

        return insights

    def compute_vol_mom(self, prices):
        rets = prices.pct_change().dropna()
        vol = np.std(rets)[-1]
        mom = np.mean(rets)[-1]
        return vol, mom