"""
Useful References:
-QC.Securities Namespace link
https://www.quantconnect.com/lean/documentation/topic26198.html
-QC SecurityExchange Class Members
https://www.quantconnect.com/lean/documentation/topic26905.html
-QC SecurityExchangeHours Class Members
https://www.quantconnect.com/lean/documentation/topic26922.html

-Data Consolidation Example Algo
https://github.com/QuantConnect/Lean/blob/master/Algorithm.Python/DataConsolidationAlgorithm.py
-Link to list of built-in indicators for QC (from Github)
https://github.com/QuantConnect/Lean/tree/master/Indicators
"""
################################################################################
# BACKTEST DETAILS
CASH = 1e6
START_DATE = '05-01-2002' #'05-01-1971' # MM-DD-YYYY format
END_DATE   = None #'12-31-1971' # MM-DD-YYYY format (or None for to date)
TIMEZONE = 'US/Central' # e.g. 'US/Eastern', 'US/Central', 'US/Pacific'

# Lookback period in months for fundamental indices trend scores
TREND_PERIOD = 12
# Turn on/off using the simplified version basing the desired country weights on 
#  blended scores only and not the rankings against the other countries
# When using this logic, it is essentially looking at absolute momentum rather than cross-sectional momentum
SIMPLIFIED_WEIGHTS = False

# Set the currency exchange to use
CURRENCY_EXCHANGE = 'Oanda' # either 'Oanda' or 'FXCM'

# Set the benchmark - must be an equity/etf
BENCHMARK = "SPY"

# Set your quandl authorization code
QUANDL_AUTH_CODE = 'ZYmz4yBbyvxrejZ44hKo'

# DEFINE COUNTRIES TO TRACK
COUNTRIES = [
    'Australia',
    'Canada',
    # 'France',
    # 'Germany',
    # 'Italy',
    'Europe',
    'Japan',
    'UK',
    # 'USA'
    ]
    
# DEFINE THE WEIGHTS FOR PORTFOLIO CONSTRUCTION
WEIGHT_EA = 0.5 # economic activity weight
WEIGHT_IN = 0.5 # inflation weight
    
# DEFINE CURRENCIES TO TRADE
TRADE_CURRENCIES = True # turn on/off using this data
CURRENCIES = {}
CURRENCIES['Australia'] = 'AUDUSD' # QC oanda data starts 2002-05-05, fxcm starts 2007-03-30
CURRENCIES['Canada']    = 'USDCAD' # QC oanda data starts 2002-05-06, fxcm starts 2007-03-30
CURRENCIES['Europe']    = 'EURUSD' # QC oanda data starts 2002-05-05, fxcm starts 2007-03-30
CURRENCIES['Japan']     = 'USDJPY' # QC oanda data starts 2002-05-05, fxcm starts 2007-03-30
CURRENCIES['UK']        = 'GBPUSD' # QC oanda data starts 2002-05-05, fxcm starts 2007-03-30

# Instead trade based on FRED exchange rates
TRADE_FRED_RATES = False # turn on/off using this data
FRED = {}
# Enter FRED rate codes below
FRED['Australia'] = 'DEXUSAL'
FRED['Canada']    = 'DEXCAUS'
FRED['Europe']    = 'DEXUSEU'
FRED['Japan']     = 'DEXJPUS'
FRED['UK']        = 'DEXUSUK'

# Enter CSV download links below
# The countries need to be ALL CAPS for this
FRED_LINKS = {}
FRED_LINKS['AUSTRALIA'] = 'https://docs.google.com/spreadsheets/d/e/2PACX-1vSaLzwcbVI_2PIQwha5IaIPk8oCymJq5fZtAw-VS20mOty-iEY53MANo1GliCos8TaiLr48RoD_WwQ5/pub?gid=367664124&single=true&output=csv'
FRED_LINKS['CANADA']    = 'https://docs.google.com/spreadsheets/d/e/2PACX-1vSaLzwcbVI_2PIQwha5IaIPk8oCymJq5fZtAw-VS20mOty-iEY53MANo1GliCos8TaiLr48RoD_WwQ5/pub?gid=0&single=true&output=csv'
FRED_LINKS['EUROPE']    = 'https://docs.google.com/spreadsheets/d/e/2PACX-1vSaLzwcbVI_2PIQwha5IaIPk8oCymJq5fZtAw-VS20mOty-iEY53MANo1GliCos8TaiLr48RoD_WwQ5/pub?gid=748699328&single=true&output=csv'
FRED_LINKS['JAPAN']     = 'https://docs.google.com/spreadsheets/d/e/2PACX-1vSaLzwcbVI_2PIQwha5IaIPk8oCymJq5fZtAw-VS20mOty-iEY53MANo1GliCos8TaiLr48RoD_WwQ5/pub?gid=1378346043&single=true&output=csv'
FRED_LINKS['UK']        = 'https://docs.google.com/spreadsheets/d/e/2PACX-1vSaLzwcbVI_2PIQwha5IaIPk8oCymJq5fZtAw-VS20mOty-iEY53MANo1GliCos8TaiLr48RoD_WwQ5/pub?gid=1282896252&single=true&output=csv'

# QUANDL DATASET CODES
# Country codes for datasets
CODE = {}
CODE['Australia'] = 'AUS'
CODE['Canada']  = 'CAN'
CODE['Europe']  = 'OECDE'
# CODE['France']  = 'FRA'
# CODE['Germany'] = 'DEU'
# CODE['Italy']   = 'ITA'
CODE['Japan']   = 'JPN'
CODE['UK']      = 'GBR'
# CODE['USA']     = 'USA'

# Monthly Industrial Production
IP = {}
for country in COUNTRIES:
    IP[country] = 'OECD/KEI_PRINTO01_' + CODE[country] + '_ST_M'

# # REF: https://www.quandl.com/data/OECD/KEI_PRINTO01_CAN_ST_M-Industrial-production-s-a-Canada-Level-ratio-or-index-Monthly
# IP['Canada']  = 'OECD/KEI_PRINTO01_CAN_ST_M'
# # REF: https://www.quandl.com/data/OECD/KEI_PRINTO01_FRA_ST_M-Industrial-production-s-a-France-Level-ratio-or-index-Monthly
# IP['France']  = 'OECD/KEI_PRINTO01_FRA_ST_M'
# # REF: https://www.quandl.com/data/OECD/KEI_PRINTO01_DEU_ST_M-Industrial-production-s-a-Germany-Level-ratio-or-index-Monthly
# IP['Germany'] = 'OECD/KEI_PRINTO01_DEU_ST_M'
# # REF: https://www.quandl.com/data/OECD/KEI_PRINTO01_ITA_ST_M-Industrial-production-s-a-Italy-Level-ratio-or-index-Monthly
# IP['Italy']   = 'OECD/KEI_PRINTO01_ITA_ST_M'
# # REF: https://www.quandl.com/data/OECD/KEI_PRINTO01_JPN_ST_M-Industrial-production-s-a-Japan-Level-ratio-or-index-Monthly
# IP['Japan']   = 'OECD/KEI_PRINTO01_JPN_ST_M'
# # REF: https://www.quandl.com/data/OECD/KEI_PRINTO01_GBR_ST_M-Industrial-production-s-a-United-Kingdom-Level-ratio-or-index-Monthly
# IP['UK']      = 'OECD/KEI_PRINTO01_GBR_ST_M'
# # REF: https://www.quandl.com/data/OECD/KEI_PRINTO01_USA_ST_M-Industrial-production-s-a-United-States-Level-ratio-or-index-Monthly
# IP['USA']     = 'OECD/KEI_PRINTO01_USA_ST_M'

# Monthly Retail Sales
RS = {}
for country in COUNTRIES:
    RS[country] = 'OECD/KEI_SLRTTO01_' + CODE[country] + '_ST_M'

# # REF: https://www.quandl.com/data/OECD/KEI_SLRTTO01_CAN_ST_M-Retail-trade-Volume-s-a-Canada-Level-ratio-or-index-Monthly
# RS['Canada'] = 'OECD/KEI_SLRTTO01_CAN_ST_M'
# # REF: https://www.quandl.com/data/OECD/KEI_SLRTTO01_USA_ST_M-Retail-trade-Volume-s-a-United-States-Level-ratio-or-index-Monthly
# RS['USA']    = 'OECD/KEI_SLRTTO01_USA_ST_M'

# Monthly Unemployment
UE = {}
for country in COUNTRIES:
    UE[country] = 'OECD/STLABOUR_' + CODE[country] + '_LFHUTTTT_ST_M'

# # REF: https://www.quandl.com/data/OECD/STLABOUR_CAN_LFHUTTTT_ST_M-Canada-Unemployment-Monthly-Total-All-Persons-Level-Rate-Or-Quantity-Series
# UE['Canada'] = 'OECD/STLABOUR_CAN_LFHUTTTT_ST_M'
# # REF: https://www.quandl.com/data/OECD/STLABOUR_USA_LFHUTTTT_ST_M-United-States-Harmonised-Unemployment-Monthly-Total-All-Persons-Level-Rate-Or-Quantity-Series
# UE['USA']    = 'OECD/STLABOUR_USA_LFHUTTTT_ST_M'

# Monthly Consumer Prices
CP = {}
for country in COUNTRIES:
    CP[country] = 'OECD/KEI_CPALTT01_' + CODE[country] + '_ST_M'

# # REF: https://www.quandl.com/data/OECD/KEI_CPALTT01_CAN_ST_M-Consumer-prices-all-items-Canada-Level-ratio-or-index-Monthly
# CP['Canada'] = 'OECD/KEI_CPALTT01_CAN_ST_M'
# # REF: https://www.quandl.com/data/OECD/KEI_CPALTT01_USA_ST_M-Consumer-prices-all-items-United-States-Level-ratio-or-index-Monthly
# CP['USA']    = 'OECD/KEI_CPALTT01_USA_ST_M'

# Monthly Producer Prices
PP = {}
for country in COUNTRIES:
    PP[country] = 'OECD/KEI_PIEAMP01_' + CODE[country] + '_ST_M'

# # REF: https://www.quandl.com/data/OECD/KEI_PIEAMP01_CAN_ST_M-Producer-prices-Manufacturing-Canada-Level-ratio-or-index-Monthly
# PP['Canada'] = 'OECD/KEI_PIEAMP01_CAN_ST_M'
# # REF: https://www.quandl.com/data/OECD/KEI_PIEAMP01_USA_ST_M-Producer-prices-Manufacturing-United-States-Level-ratio-or-index-Monthly
# PP['USA']    = 'OECD/KEI_PIEAMP01_USA_ST_M'


# Set the number of warmup days
CALENDAR_WARMUP_DAYS = 50

# Turn on/off specific logs
PRINT_WARMUP = False # print logs during warmup for debugging
# PRINT_BARS   = False # print desired OHLC bars for debugging
PRINT_SIGNALS = True # print desired trading signals

################################################################################ 
############################ END OF ALL USER INPUTS ############################
################################################################################

# VALIDATE USER INPUTS - DO NOT CHANGE BELOW

#-------------------------------------------------------------------------------
import datetime as DT
# Verify start date
try:
    START_DATE_DT = DT.datetime.strptime(START_DATE, '%m-%d-%Y')
except:
    raise ValueError("Invalid START_DATE format ({}). Must be in MM-DD-YYYY "
        "format.".format(START_DATE))
        
# Verify end date
try:
    if END_DATE:
        END_DATE_DT = DT.datetime.strptime(END_DATE, '%m-%d-%Y')
except:
    raise ValueError("Invalid END_DATE format ({}). Must be in MM-DD-YYYY "
        "format or set to None to run to date.".format(END_DATE))
        
# Verify CURRENCY_EXCHANGE
if CURRENCY_EXCHANGE not in ['Oanda', 'FXCM']:
    raise ValueError("Invalid CURRENCY_EXCHANGE ({}). Must be 'Oanda' or "
        "'FXCM'.".format(CURRENCY_EXCHANGE))
        

###############################################################################
# Standard library imports
import datetime as DT
# from datetime import timedelta
# from datetime import date
from dateutil.parser import parse
import decimal
import math
# import numpy as np
# import pandas as pd
import pytz

# QuantConnect specific imports
import QuantConnect as qc
from QuantConnect.Python import PythonQuandl

# Import from files
from notes_and_inputs import *

################################################################################
class CustomTradingStrategy(QCAlgorithm):
    def Initialize(self):
        """Initialize algorithm."""
        # Set backtest details
        self.SetStartDate(
            START_DATE_DT.year, START_DATE_DT.month, START_DATE_DT.day)
        if END_DATE:
            self.SetEndDate(
                END_DATE_DT.year, END_DATE_DT.month, END_DATE_DT.day)
        self.SetCash(CASH)
        self.SetTimeZone(TIMEZONE)
        
        # Set your personal token necessary for restricted dataset
        Quandl.SetAuthCode(QUANDL_AUTH_CODE)
        
        # Set up dictionaries to hold all Quandl custom datasets
        self.industrial_production = {}
        self.retail_sales = {}
        self.unemployment = {}
        self.consumer_prices = {}
        self.producer_prices = {}
        self.symbols = {}
        
        # Loop through all desired countries to track
        for country in COUNTRIES:
            # Get all industrial production datasets
            self.industrial_production[country] = self.AddData(
                QuandlCustomColumns, 
                IP[country], 
                Resolution.Daily, 
                TimeZones.NewYork
                ).Symbol
            # Get all retail sales datasets
            self.retail_sales[country] = self.AddData(
                QuandlCustomColumns, 
                RS[country], 
                Resolution.Daily, 
                TimeZones.NewYork
                ).Symbol
            # Get all unemployment datasets
            self.unemployment[country] = self.AddData(
                QuandlCustomColumns, 
                UE[country], 
                Resolution.Daily, 
                TimeZones.NewYork
                ).Symbol
            # Get all consumer prices datasets
            self.consumer_prices[country] = self.AddData(
                QuandlCustomColumns, 
                CP[country], 
                Resolution.Daily, 
                TimeZones.NewYork
                ).Symbol
            # Get all producer prices datasets
            self.producer_prices[country] = self.AddData(
                QuandlCustomColumns, 
                PP[country], 
                Resolution.Daily, 
                TimeZones.NewYork
                ).Symbol
                
            # Add data to trade
            if TRADE_CURRENCIES:
                # Add built-in currency data from the desired exchange
                if CURRENCY_EXCHANGE == 'Oanda': # either 'Oanda' or 'FXCM'
                    exchange = Market.Oanda
                else:
                    exchange = Market.FXCM
                self.symbols[country] = \
                    self.AddForex(CURRENCIES[country], Resolution.Hour, exchange).Symbol
            elif TRADE_FRED_RATES:
                # Add FRED rate data
                self.symbols[country] =\
                    self.AddData(MyCustomData, country, Resolution.Daily).Symbol
                
        # Set benchmark for scheduling purposes
        self.bm = self.AddEquity(BENCHMARK, Resolution.Hour).Symbol
        
        # Rebalance the portfolio at the end of every month at market close
        self.Schedule.On(
            self.DateRules.MonthEnd(self.bm),
            self.TimeRules.BeforeMarketClose(self.bm, 0), 
            self.Rebalance)
            
        # Initialize required variables
        self.portfolio_targets = []
        
        try:
            self.TREND_PERIOD = int(self.GetParameter("TREND_PERIOD"))
        except:
            self.TREND_PERIOD = TREND_PERIOD
            
        # Warm up the indicators prior to the start date
        self.SetWarmUp(timedelta(CALENDAR_WARMUP_DAYS))
        
#-------------------------------------------------------------------------------
    def OnData(self, data):
        """Built-in event handler for new data."""
        # Check for new portfolio targets
        if len(self.portfolio_targets) > 0:
            # Update portfolio
            self.SetHoldings(self.portfolio_targets)
            # Empty portfolio targets list
            self.portfolio_targets = []
            
#-------------------------------------------------------------------------------
    def Rebalance(self):
        """Rebalance the portfolio."""   
        # First calculate the economic momentum scores and get the desired
        #  country weights
        self.CalculateScores()
        
        # Loop through each country and get the portfolio targets for each
        self.portfolio_targets = []
        for country in COUNTRIES:
            # Get the desired country weight
            if country in self.weights:
                target_weight = self.weights[country]
            else:
                target_weight = 0
            
            # Get the instrument to trade for the country
            symbol = self.symbols[country]
            
# # Verify that data is available for the symbol
# if not self.Securities[symbol].HasData:
#     continue
            
            # Check if USD is the base or quote currency
            # 'USD' is the base currency if at beginning of symbol
            usd_base = False
            if TRADE_CURRENCIES and 'USD' == str(symbol)[:3]:
                usd_base = True
            elif TRADE_FRED_RATES:
                # Get the FRED rate code
                fred_code = str(FRED[country])[3:] # remove 'DEX' at beginning
                if 'US' == fred_code[:2]:
                    usd_base = True
            if usd_base:
                # Want to sell to go long the country currency
                # So inverse the target weight
                target_weight = -target_weight
            # else: # 'USD' is the quote currency
            
            # Set the target allocation for the currency
            self.portfolio_targets.append(
                PortfolioTarget(symbol, target_weight)
                )
                
#-------------------------------------------------------------------------------
    def CalculateScores(self):
        """Calculate the economic momentum scores for all countries.""" 
        # Create dictionaries to store country specific index scores
        ea_index = {} # economic activity index
        in_index = {} # inflation index
        weight = {} # overall desired weight for each country
        
        # Loop through each country
        for country in COUNTRIES:
            # Get history for all 5 of the fundamental factors
            length = (TREND_PERIOD+1+6)*31 # get an extra 5-6 months of data
            hist_ip = self.History([self.industrial_production[country]], length)
            hist_rs = self.History([self.retail_sales[country]], length)
            hist_ue = self.History([self.unemployment[country]], length)
            hist_cp = self.History([self.consumer_prices[country]], length)
            hist_pp = self.History([self.producer_prices[country]], length)
            
# we can test different lookback periods in QC's optimization feature, but it takes some extra compute power (discuss with Jon later - need to just get it working for now)
            # Get the economic activity index
            if min([len(hist_ip), len(hist_rs), len(hist_ue)]) > TREND_PERIOD+1:
                # Get the industrial production log growth rate
                ip = math.log(hist_ip.iloc[-1].value) \
                    -math.log(hist_ip.iloc[-TREND_PERIOD-1].value)
                # Get the retail sales log growth rate
                rs = math.log(hist_rs.iloc[-1].value) \
                    -math.log(hist_rs.iloc[-TREND_PERIOD-1].value)
                # Get the unemployment log growth rate
                ue = math.log(hist_ue.iloc[-1].value) \
                    -math.log(hist_ue.iloc[-TREND_PERIOD-1].value)
                # Calculate the economic activity index
                # equal weighted average of factors above
                ea_index[country] = (ip+rs+ue)/3.0
            else:
                ea_index[country] = 0
                
            # Get the inflation index
            if min([len(hist_cp), len(hist_pp)]) > TREND_PERIOD+1:
                # Get the consumer prices log growth rate
                cp = math.log(hist_cp.iloc[-1].value) \
                    -math.log(hist_cp.iloc[-TREND_PERIOD-1].value)
                pp = math.log(hist_pp.iloc[-1].value) \
                    -math.log(hist_pp.iloc[-TREND_PERIOD-1].value)
                # Calculate the inflation index
                # equal weighted average of factors above
                in_index[country] = (cp+pp)/2.0
            else:
                in_index[country] = 0
                
        # Check if SIMPLIFIED_WEIGHTS is used
        if SIMPLIFIED_WEIGHTS:
            # Loop through each country
            for country in COUNTRIES:
                # Get the overall desired weight based on raw blended score
                blended_score = WEIGHT_EA*ea_index[country] + WEIGHT_IN*in_index[country]
                if blended_score > 0:
                    weight[country] = 1
                elif blended_score < 0:
                    weight[country] = -1
                else:
                    weight[country] = 0
        else:   
            # Now get rankings for each index
            # reverse=True for those we want to sort in descending order
            sorted_ea = sorted(ea_index.items(), key=lambda x: x[1], reverse=True)
            sorted_in = sorted(in_index.items(), key=lambda x: x[1], reverse=True)
            
            ea_ranks = list(range(1,len(sorted_ea)+1))
            ea_avg_rank = sum(ea_ranks)/len(ea_ranks)
            in_ranks = list(range(1,len(sorted_in)+1))
            in_avg_rank = sum(in_ranks)/len(in_ranks)
            
            # Loop through each country
            for country in COUNTRIES:
# paper mentions a 3month lag on using this data for the portfolio
                # Get the economic activity score
                if country in ea_index:
                    rank_ea = [x[0] for x in sorted_ea].index(country)+1
                    inverse_rank_ea = ea_ranks[-rank_ea]
# this formula doesn't make sense to me - discuss with Jon
                    score_ea = inverse_rank_ea - ea_avg_rank # sum([x[0] for x in sorted_ea])/len(sorted_ea)
                else:
                    score_ea = 0
                # Get the inflation score
                if country in in_index:
                    rank_in = [x[0] for x in sorted_in].index(country)+1
                    inverse_rank_in = in_ranks[-rank_in]
# this formula doesn't make sense to me - discuss with Jon
                    score_in = inverse_rank_in - in_avg_rank # sum([x[0] for x in sorted_in])/len(sorted_in)
                else:
                    score_in = 0
                # Get the overall desired weight
                weight[country] = WEIGHT_EA*score_ea + WEIGHT_IN*score_in
                
        # NOTE: weight adds up to 0 / longs + shorts
        # Make sure absolute values of long and short weights add to 1
        if sum([abs(x[1]) for x in weight.items()]) != 1:
            # Rescale so absolute value of long weights and short weights add to 1
            weights_sum = sum([abs(x[1]) for x in weight.items()])
            if weights_sum > 0:
                factor = 1.0/weights_sum
            else:
                factor = 0
            for country in COUNTRIES:
                weight[country] = factor*weight[country]
            
        # Save weights and return
        self.weights = weight
        
        return
    
#-------------------------------------------------------------------------------
    def CustomSecurityInitializer(self, security):
        """
        Define models to be used for securities as they are added to the 
        algorithm's universe.
        """
        # Define the data normalization mode
        security.SetDataNormalizationMode(DataNormalizationMode.Adjusted)
        
        # Define the fee model to use for the security
        # security.SetFeeModel()
        # Define the slippage model to use for the security
        # security.SetSlippageModel()
        # Define the fill model to use for the security
        # security.SetFillModel()
        # Define the buying power model to use for the security
        # security.SetBuyingPowerModel()
        
################################################################################
class QuandlCustomColumns(PythonQuandl):
    """
    Quandl often doesn't use close columns so need to tell LEAN which is the 
    "value" column.
    REF: https://www.quantconnect.com/forum/discussion/11566/kei-based-strategy/p1
    """
    def __init__(self):
        # Define ValueColumnName: cannot be None, Empty or non-existant column name
        self.ValueColumnName = "Value"
        
################################################################################
class MyCustomData(PythonData):
    """
    Custom Data Class
    REF: https://www.quantconnect.com/forum/discussion/4079/python-best-practise-for-using-consolidator-on-custom-data/p1
    """
    def GetSource(self, config, date, isLiveMode):
        # Get file specific to the asset symbol
        country = config.Symbol.Value
        # Get the custom url link for the FRED data for the desired country
        file = FRED_LINKS[country]
        return SubscriptionDataSource(
            file, SubscriptionTransportMedium.RemoteFile)

    def Reader(self, config, line, date, isLiveMode):
        # Create new object
        asset = MyCustomData()
        asset.Symbol = config.Symbol
        
        # If first character is not a digit, return
        if not (line.strip() and line[0].isdigit()):
            return None
            
        # try:
        # Example File Format:
        # Date                      Price  
        # 2017-01-02 09:02:00+01:00 64.88
        
        # Comma separated file, so split data row by comma
        data = line.split(',')
        
        # If price is invalid, return None
        try:
            value = float(data[1])
        except:
            return None
        # Return None if value is 0
        if value == 0:
            return None
            
        # Set time of data at close 16:00 / 57600s = 16hr*60min/hr*60s/min
        # Set time of data at close 16:01 / 57660s = 16hr*60min/hr*60s/min
        # asset.Time = parse(data[0]) #+ timedelta(seconds=57660)
        asset.Time = DT.datetime.strptime(data[0], "%m/%d/%Y")
        
        # Set the value used for filling positions
        asset.Value = value #decimal.Decimal(data[1])
        asset["Close"] = value
        
        return asset
            
        # except ValueError:
        #     # Do nothing
        #     return None