using QuantConnect.Indicators;
using QuantConnect.PythonPredictions;

namespace QuantConnect.Algorithm.CSharp
{
    /// <summary>
    /// Basic algorithm showing how to implement Python code in a C# algorithm
    /// </summary>
    public class PythonInCSharpeExampleAlgorithm : QCAlgorithm
    {
    	PythonMethods python = new PythonMethods();
    	public Resolution resolution = Resolution.Hour;
    	
        /// <summary>
        /// Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.
        /// </summary>
        public override void Initialize()
        {
        	// Initialize the Algorithm
			SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage);
            SetStartDate(2015, 01, 01);
            
            SetCash(1000000);
            SetWarmUp(180);
            SetBenchmark("SPY");
            var spy = AddEquity("SPY", resolution).Symbol;
            AddEquity("AAPL", resolution);
            AddEquity("GM", resolution);
            AddEquity("TSLA", resolution);
			
			//Schedule the HMM model event to trigger. We will update the market state for each symbol every morning
			//and we will refit our model at the beginning of each month
        	Schedule.On(DateRules.EveryDay("SPY"), TimeRules.At(5,0), () => 
	        {
	        	var bars = History<TradeBar>(Securities.Keys, 100, Resolution.Daily);
	        	foreach (Symbol sym in Securities.Keys)
	        	{
	        		// Pass the history and the symbol we are looking to update
	        		python.UpdateMarketState(sym, bars);
	        	}
	        	
	        	// Print the current overall market state
	        	if (python.GetSymbolState(spy) == 1)
	        	{
	        		Debug("Current State of the market = " + python.GetSymbolState(spy));
	        	}
	        	else if (python.GetSymbolState(spy) == 0)
	        	{
	        		Debug("Current State of the market = " + python.GetSymbolState(spy));
	        	}
	        });
	        Schedule.On(DateRules.MonthStart("SPY"), TimeRules.At(19,0), () => 
	        {
        		var bars = History<TradeBar>(Securities.Keys, 900, Resolution.Daily);
	        	foreach (Symbol sym in Securities.Keys)
	        	{
	        		// Pass the extended history to refit the selected symbol models
	        		python.RefitModels(sym, bars);
	        	}
        	});
        }

        /// <summary>
        /// OnData event.
        /// </summary>
        /// <param name="data">Slice object keyed by symbol containing the stock data</param>
        public void OnData(Slice data)
        {
        	// Check each security to see if the HMM state of the symbol matches the market. If it does, enter a position.
        	// Exit the position if the market state changes.
            foreach(var i in Securities.Keys)
        	{
    			if (data.ContainsKey(i))
    			{
					if (Portfolio[i].Quantity > 0 && python.GetSymbolState(Portfolio["SPY"].Symbol) == 1)
					{
						Liquidate(i);
					}
					else if (Portfolio[i].Quantity < 0 && python.GetSymbolState(Portfolio["SPY"].Symbol) == 0)
					{
						Liquidate(i);
					}
					else if (Portfolio[i].Quantity == 0)
					{
	    				// Check to make sure we have a state for the given symbol or market, otherwise do nothing.
	    				if (python.GetSymbolState(i) == -1 || python.GetSymbolState(Portfolio["SPY"].Symbol) == -1)
	    					return;
						if (python.GetSymbolState(i) == python.GetSymbolState(Portfolio["SPY"].Symbol))
						{
							if (python.GetSymbolState(i) == 1)
								MarketOrder(i, -1);
							else if (python.GetSymbolState(i) == 0)
								MarketOrder(i, 1);
						}
					}
				}
			}
    	}
    }
}

using Python.Runtime;

namespace QuantConnect.PythonPredictions
{
	// Python Methods class that allows for calls of the HMM code. The original code can be found here:
	// https://www.quantconnect.com/forum/discussion/6878/from-research-to-production-hidden-markov-models/p1
	class PythonMethods
	{
		// Dictionary containing the current state of the symbol. Will result in a 1 or 0 value
		private Dictionary<Symbol, double> _symbolState = new Dictionary<Symbol, double>();
		
		// A Dictionary of the models the HMM function creates. Used to create a prediction and will be
		// periodically updated
		private Dictionary<Symbol, PyObject> _models = new Dictionary<Symbol, PyObject>();
		
		// Update the market state of a selected symbol
		public void UpdateMarketState(Symbol symbol, IEnumerable<DataDictionary<TradeBar>> bars)
		{ 
			// Convert the history TradeBars to a list that allows for easy conversion into a DataFrame
			// object in python
        	var history = new List<TradeBar>();
        	foreach (var bar in bars)
        	{
        		if (!bar.ContainsKey(symbol))
        			return;
        		history.Add(bar[symbol]);
        	}
        	// Check to see if we already have a model for the given symbol. If not we will have to create one
	    	if (_models.ContainsKey(symbol))
	    	{
	    		var model = _models[symbol];
	    		// This is where we implement our python code. This creates the script that we can then call
	    		using (Py.GIL())
	            {
	                dynamic GetState = PythonEngine.ModuleFromString("GetStateModule",
	                    @"
from clr import AddReference
AddReference(""QuantConnect.Common"")
from QuantConnect import *
from QuantConnect.Data.Custom import *
import numpy as np
import pandas as pd
from statsmodels.tsa.stattools import adfuller
from hmmlearn.hmm import GaussianHMM

def to_dict(TradeBar):
	return{
		'close': TradeBar.Close,
		'EndTime': TradeBar.EndTime,
		'Symbol': TradeBar.Symbol,
	}

def TestStationartiy(returns):
    return adfuller(returns)[1] < 0.05

def PredictState(df, model):
	price = np.array(df.iloc[0].close).reshape((1,1))
	return model.predict(price)[0]

def GetState(history, model):
	state = -1
	df = pd.DataFrame([to_dict(s) for s in history], columns = ['close', 'Symbol', 'EndTime'])
	df.set_index('EndTime')
	#returns = np.array(df.close.pct_change().dropna())
	#returns = np.array(returns).reshape((len(returns),1))
	returns = df.unstack(level = 1).close.transpose().pct_change().dropna()
	stationarity = TestStationartiy(returns)
	if stationarity:
		state = PredictState(df, model)
	        
	return state
").GetAttr("GetState");
	
	                // Run the GetState method we created above
	                double state = GetState(history, model);
	                
	                // Add or update the state of the selected symbol
	                if (!_symbolState.ContainsKey(symbol))
	                	_symbolState.Add(symbol, state);
	                else if (_symbolState[symbol] != state && state != -1)
	                	_symbolState[symbol] = state;
	            }
	    	}
	    	// Create a new model
	    	else
	    	{
	    		// Implementing python code to create the HMM model
	    		using (Py.GIL())
	            {
	                dynamic CreateHMM = PythonEngine.ModuleFromString("CreateHMMModule",
	                    @"
from clr import AddReference
AddReference(""QuantConnect.Common"")
from QuantConnect import *
from QuantConnect.Data.Custom import *
import numpy as np
import pandas as pd
from hmmlearn.hmm import GaussianHMM

def to_dict(TradeBar):
	return{
		'close': TradeBar.Close,
		'EndTime': TradeBar.EndTime,
		'Symbol': TradeBar.Symbol,
	}
	
def CreateHMM(history, symbol):
	df = pd.DataFrame([to_dict(s) for s in history], columns = ['close', 'Symbol', 'EndTime'])
	df.set_index('EndTime')
	returns = np.array(df.close.pct_change().dropna())
	returns = np.array(returns).reshape((len(returns),1))
	_model = GaussianHMM(n_components=2, covariance_type=""full"", n_iter=10000).fit(returns)
	return _model

").GetAttr("CreateHMM");
	                
	                // Call the Create HMM model script we created above and add to the dictionary
	                var _model = CreateHMM(history, symbol);
	                _models[symbol] = _model;
	            }
	    	}
		}
		
		
		// Refit the HMM of a selected symbol
		public void RefitModels(Symbol symbol, IEnumerable<DataDictionary<TradeBar>> bars)
		{
			// Convert the history TradeBars to a list that allows for easy conversion into a DataFrame
			// object in python
        	var history = new List<TradeBar>();
        	foreach (var bar in bars)
        	{
        		if (!bar.ContainsKey(symbol))
        			return;
        		history.Add(bar[symbol]);
        	}
        	// Check to see if we already have a model for the given symbol. If not we will have to create one
	    	if (_models.ContainsKey(symbol))
	    	{
	    		// This is where we implement our python code. This creates the script that we can then call
	    		using (Py.GIL())
	            {
	                dynamic RefitModel = PythonEngine.ModuleFromString("RefitModelModule",
	                    @"
from clr import AddReference
AddReference(""QuantConnect.Common"")
from QuantConnect import *
from QuantConnect.Data.Custom import *
import numpy as np
import pandas as pd
from hmmlearn.hmm import GaussianHMM

def to_dict(TradeBar):
	return{
		'close': TradeBar.Close,
		'EndTime': TradeBar.EndTime,
		'Symbol': TradeBar.Symbol,
	}


def RefitModel(history, symbol, model):
	df = pd.DataFrame([to_dict(s) for s in history], columns = ['close', 'Symbol', 'EndTime'])
	df.set_index('EndTime')
	returns = np.array(df.close.pct_change().dropna())
	returns = np.array(returns).reshape((len(returns),1))
	return model.fit(returns)
").GetAttr("RefitModel");
	
	                // Call the Refit HMM model script we created above and update the dictionary
	                var model = _models[symbol];
	                var _model = RefitModel(history, symbol, model);
	                _models[symbol] = _model;
	            }		
	    	}
	    	// Create a new model
	    	else
	    	{
	    		// Implementing python code to create the HMM model
	    		using (Py.GIL())
	            {
	                dynamic CreateHMM = PythonEngine.ModuleFromString("CreateHMMModule",
	                    @"
from clr import AddReference
AddReference(""QuantConnect.Common"")
from QuantConnect import *
from QuantConnect.Data.Custom import *
import numpy as np
import pandas as pd
from hmmlearn.hmm import GaussianHMM

def to_dict(TradeBar):
	return{
		'close': TradeBar.Close,
		'EndTime': TradeBar.EndTime,
		'Symbol': TradeBar.Symbol,
	}
	
def CreateHMM(history, symbol):
	df = pd.DataFrame([to_dict(s) for s in history], columns = ['close', 'Symbol', 'EndTime'])
	df.set_index('EndTime')
	returns = np.array(df.close.pct_change().dropna())
	returns = np.array(returns).reshape((len(returns),1))
	_model = GaussianHMM(n_components=2, covariance_type=""full"", n_iter=10000).fit(returns)
	return _model

").GetAttr("CreateHMM");
	                
	                // Call the Create HMM model script we created above and add to the dictionary
	                var _model = CreateHMM(history, symbol);
	                _models[symbol] = _model;
	            }
	    	}
		}
		
		// Retreive the symbol state of a selected symbol
		public double GetSymbolState(Symbol symbol)
		{
			// Check to see if we have a state for the symbol and return that value. Otherwise return -1.
			if (_symbolState.ContainsKey(symbol))
				return (_symbolState[symbol]);
			else
				return -1;
		}
	}
}