'''
    Ostirion Multiple Moving Averages Demostration
    version 1.0
    Copyright (C) 2021  Ostirion

    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/>.
    
    Contact: www.ostirion.net/contact
'''

from Execution.ImmediateExecutionModel import ImmediateExecutionModel
import numpy as np
import pandas as pd
from CustomUniverseSelectionModel import CustomUniverseSelectionModel as CU


class OptimumSupports(QCAlgorithm):

    def Initialize(self):
        self.test_years = 5
        self.SetStartDate(datetime.today() - timedelta(days=365*5))
        self.SetEndDate(datetime.today())
        self.SetCash(1000000) 
        res = Resolution.Daily
        self.AddUniverseSelection(CU(100))
        self.UniverseSettings.Resolution = res
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        self.AddAlpha(SRAlphaModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())

class SRAlphaModel(AlphaModel):

    def __init__(self):

        self.symbol_data = {}
        self.training_length = 1100
        self.windows = np.linspace(5,100,20).astype(int)
        self.day_counter = 0
        self.training_period = 60
        
    def Update(self, algorithm, data):
        self.day_counter += 1
        insights = []
        if data.HasData == False: return []

        # Find the symbols that have all data:
        common_list = list(set(data.Keys).intersection(self.symbol_data.keys()))

        for s in common_list:
            if not data.get(s): continue
            s_data = self.symbol_data.get(s)
            past_status = s_data.ma_status
            current_status = data.get(s).Close > s_data.ma.Current.Value
            s_data.ma_status = data.get(s).Close > s_data.ma.Current.Value

            #Retrain symbol if needed:
            if self.day_counter%self.training_period == 0:
                algorithm.Debug('Retraining...')
                history = algorithm.History(s, self.training_length, Resolution.Daily)
                best_period = self.FindBestAverage(history, self.windows)
                s_data = SymbolData(s, algorithm, best_period)
                s_data.best_period = best_period
                # Warm-up indicators for added securities
                s_data.WarmUpIndicators(history)

            if algorithm.Securities.get(s).Invested:
                continue
            if current_status == past_status:
                continue
            if current_status:
                # Crossing up: moves down
                d = InsightDirection.Down
            else: d = InsightDirection.Up
            period =  timedelta(days=int(s_data.best_period))

            # Skip erroneous averages.
            if period == -1: 
                continue
            insight = Insight(s, period, InsightType.Price, d, 0.02, 1,"EMA Balance", 1.00)
            insights.append(insight)

        return insights

    def OnSecuritiesChanged(self, algorithm, changes):

        for security in changes.AddedSecurities:

            if security.Symbol not in self.symbol_data:
                history = algorithm.History(security.Symbol, self.training_length, Resolution.Daily)
                # Error in history length:
                if history.empty: continue
                best_period = self.FindBestAverage(history, self.windows)
                if best_period == -1: continue
                self.symbol_data[security.Symbol] = SymbolData(security.Symbol, algorithm, best_period)
                self.symbol_data[security.Symbol].best_period = best_period

                # Warm-up indicators for added securities
                self.symbol_data.get(security.Symbol).WarmUpIndicators(history)

        for security in changes.RemovedSecurities:
            if security.Symbol in self.symbol_data:
                self.symbol_data.pop(security.Symbol)
        return

    def FindBestAverage(self, history, windows):
        prices = history['close'].unstack(level=0)
        if len(prices) < 100:
            #Insufficient history length.
            return -1
        averages = prices.copy()    
        for win in windows:
            averages['MA_' + str(win)] = prices.rolling(window=win).mean()    
        averages.dropna(inplace=True)
        price_col = prices.columns[0]    
        for win in windows:
            averages['d_MA_'+str(win)] = averages[price_col] - averages['MA_' + str(win)]
            shift = averages['d_MA_'+str(win)].shift(1)
            diff = averages['d_MA_'+str(win)]
            averages['c_'+str(win)] = np.sign(shift) != np.sign(diff)     
        cross_columns = [col for col in averages.columns if 'c_' in col]    
        rates = {}    
        for column in cross_columns:
            rates[column] = averages[column].value_counts(normalize=True)[True]    
        results = {}
        for window in windows:
            results[window] = averages[[price_col, 'MA_'+str(window), 'd_MA_'+str(window),
                                        'c_'+str(window)]]    
            averages[str(window)+'_f'] = averages[price_col].shift(-window)        
            averages['ctype_'+str(window)] = averages['d_MA_'+str(window)] > 0        
            averages['FPdir_'+str(window)] = averages[str(window)+'_f'] > averages[price_col]
            averages['H_'+str(window)] = averages['ctype_'+str(window)] != averages['FPdir_'+str(window)]

            try:
                res = averages['H_'+str(window)].value_counts(normalize=True)[True]
                results[window]=res
            except:
                # There are no True values:
                results[window]=-1

        return max(results, key=results.get)   

class SymbolData(object):

    def __init__(self, symbol, algorithm, period):
        self.symbol = symbol
        self.ma = algorithm.SMA(self.symbol, period, Resolution.Daily)
        self.ma_status = True
        self.best_period = False

    def WarmUpIndicators(self, history):
        if self.symbol in history.index:
            for time, row in history.loc[str(self.symbol)].iterrows():
                if not history.empty:
                    self.ma.Update(time, row["close"])

# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# 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 clr import AddReference
AddReference("System")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm.Framework")

from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from itertools import groupby
from math import ceil

class CustomUniverseSelectionModel(FundamentalUniverseSelectionModel):
    '''Defines the QC500 universe as a universe selection model for framework algorithm
    For details: https://github.com/QuantConnect/Lean/pull/1663'''

    def __init__(self, number, filterFineData = True, universeSettings = None, securityInitializer = None):
        '''Initializes a new default instance of the QC500UniverseSelectionModel'''
        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.numberOfSymbolsCoarse = 1000
        self.numberOfSymbolsFine = number
        self.dollarVolumeBySymbol = {}
        self.lastMonth = -1

    def SelectCoarse(self, algorithm, coarse):
        '''Performs coarse selection for the QC500 constituents.
        The stocks must have fundamental data
        The stock must have positive previous-day close price
        The stock must have positive volume on the previous trading day'''
        if algorithm.Time.month == self.lastMonth:
            return Universe.Unchanged

        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData and x.Volume > 0 and x.Price > 0],
                                     key = lambda x: x.DollarVolume, reverse=True)[:self.numberOfSymbolsCoarse]

        self.dollarVolumeBySymbol = {x.Symbol:x.DollarVolume for x in sortedByDollarVolume}

        # If no security has met the QC500 criteria, the universe is unchanged.
        # A new selection will be attempted on the next trading day as self.lastMonth is not updated
        if len(self.dollarVolumeBySymbol) == 0:
            return Universe.Unchanged

        # return the symbol objects our sorted collection
        return list(self.dollarVolumeBySymbol.keys())

    def SelectFine(self, algorithm, fine):
        '''Performs fine selection for the QC500 constituents
        The company's headquarter must in the U.S.
        The stock must be traded on either the NYSE or NASDAQ
        At least half a year since its initial public offering
        The stock's market cap must be greater than 500 million'''

        sortedBySector = sorted([x for x in fine if x.CompanyReference.CountryId == "USA"
                                        and x.CompanyReference.PrimaryExchangeID in ["NYS","NAS"]
                                        and (algorithm.Time - x.SecurityReference.IPODate).days > 180
                                        and x.MarketCap > 5e8],
                               key = lambda x: x.CompanyReference.IndustryTemplateCode)

        count = len(sortedBySector)

        # If no security has met the QC500 criteria, the universe is unchanged.
        # A new selection will be attempted on the next trading day as self.lastMonth is not updated
        if count == 0:
            return Universe.Unchanged

        # Update self.lastMonth after all QC500 criteria checks passed
        self.lastMonth = algorithm.Time.month

        percent = self.numberOfSymbolsFine / count
        sortedByDollarVolume = []

        # select stocks with top dollar volume in every single sector
        for code, g in groupby(sortedBySector, lambda x: x.CompanyReference.IndustryTemplateCode):
            y = sorted(g, key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse = True)
            c = ceil(len(y) * percent)
            sortedByDollarVolume.extend(y[:c])

        sortedByDollarVolume = sorted(sortedByDollarVolume, key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse=True)
        return [x.Symbol for x in sortedByDollarVolume[:self.numberOfSymbolsFine]]