'''
    Ostirion Linear Controller BTC Price 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
'''


import numpy as np
import pandas as pd


class CryptoController(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2021, 1, 25)
        self.SetCash(1000000)
        self.SetBrokerageModel(BrokerageName.Bitfinex, AccountType.Margin)

        res = Resolution.Hour
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())
        self.Settings.RebalancePortfolioOnInsightChanges = True
        self.SetExecution(ImmediateExecutionModel())
        self.UniverseSettings.Resolution = res

        # Crypto to trade:
        crypto = ['BTCUSD']

        # K-Gain parameter, 4 as stated in the paper:
        try:
            self.K = float(self.GetParameter("K"))
        except:
            self.Debug('No K-gain parameter detected. Using default')
            self.K = float(4)

        # Create the symbols:
        self.crypto = [self.AddCrypto(ticker, res).Symbol for ticker in crypto]
        self.AddAlpha(CryptoControl(self.crypto, self.K))


class CryptoControl(AlphaModel):


    def __init__(self,
                 crypto,
                 K):

        self.Name = 'Crypto Controller'
        self.crypto = crypto
        self.K = K
        self.total_assets = len(self.crypto)
        self.init = False

        # Dictionaries for data:
        self.changes = {}
        self.loads = {}
        self.gains = {}
        
        # Low pass filter for very low weights, positions below
        # epsilon (0.005%) are set to 0 to avoid insights
        # that would result in no trade:
        self.epsilon = 0.005
        
        # Subsampler cycle of insights, skipping every square root
        # of the gain to allow intermediate rebalancing.
        self.subsampler = int(np.sqrt(int(self.K)))
        self.sample = 0

    def Update(self, algorithm, data):
        
        # Skip if no data in slice and clear insights:
        if not data:
            algorithm.Debug('No data. Skipping period.')
            return []
        
        self.sample += 1
        insights = []

        # Initialize the rate of change indicator,
        # setting the defaults for the first value.
        if not self.init:
            request_data = 120
            
            for symbol in self.crypto:
                h = algorithm.History
                self.loads[symbol] = (self.K*h(symbol, request_data, Resolution.Hour).loc[symbol]['close']).pct_change().sum()
                self.changes[symbol] = algorithm.ROC(symbol, 1)
            initial_load = sum(list(self.loads.values()))
            self.loads = {k:v/initial_load for k,v in self.loads.items()}
            self.init = True

        # Compute the gain g for each crypto symbol:
        for symbol in self.crypto:
            self.gains[symbol] = self.changes[symbol].Current.Value*self.K
            self.loads[symbol] = min(max(self.loads[symbol] + self.gains[symbol], 0), 1)
        
        # Subsample to avoid sync with rebalance and insight periods.
        if self.sample % self.subsampler != 0: return[]

        # Obtain insight time delta and magnitudes.
        t_delta = timedelta(hours=self.subsampler)
        magnitude_values = np.array(list(self.gains.values()))
        magnitude = abs(np.mean(magnitude_values))*100
        

        for symbol in self.crypto:
            w = self.loads[symbol] 
            g = self.gains[symbol]
            if w <= self.epsilon: continue
            if abs(g) <= self.epsilon/10000: continue

            insights.append(Insight(symbol, t_delta,
                                    InsightType.Price,
                                    InsightDirection.Up,
                                    magnitude, int(True),
                                    self.Name, w))
                                    

        return insights

import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import seaborn as sns


def plot_df(df, color='blue', size=(16, 7), legend='Close Price', y_label='Price in USD', title=None, kind='line'):
    plt.style.use('dark_background')
    plt.rcParams["figure.figsize"] = size
    ax = df.plot(kind=kind)
    plt.title(title)
    plt.ylabel(y_label)
    x = 0.1
    y = 0.4
    plt.text(x, y, 'www.ostirion.net', fontsize=15, transform=ax.transAxes)
    plt.legend(ncol=int(len(df.columns) / 2))
    date_form = mdates.DateFormatter("%m-%Y")
    plt.xticks(rotation=45);
    plt.show()
    
def plot_corr_hm(df, title='Title', size=(16, 7), annot = True):
    corr = df.corr()
    plt.style.use('dark_background')
    plt.rcParams["figure.figsize"] = size
    mask = np.triu(np.ones_like(corr, dtype=bool))
    cmap = sns.color_palette("RdBu")
    ax = sns.heatmap(corr, mask=mask, vmax=.3, center=0, cmap=cmap, annot=annot,
                     square=True, linewidths=0, cbar_kws={"shrink": .5}, fmt='g')
    ax.set_title(title)
    plt.setp(ax.get_yticklabels(), rotation=0);
    plt.setp(ax.get_xticklabels(), rotation=90);
    plt.show()

def plot_cm(df, title='Title', size=(16,7)):
    plt.style.use('dark_background')
    plt.rcParams["figure.figsize"] = size
    cmap = sns.color_palette("Blues")
    ax = sns.heatmap(df, cmap=cmap, annot=True, linewidths=0, cbar_kws={"shrink": .5}, fmt='g')
    ax.set_title(title)
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.setp(ax.get_xticklabels(), rotation=0);

def plot_hm(df, title='Title', size=(16, 7), annot = True, x_rot=90):

    plt.style.use('dark_background')
    plt.rcParams["figure.figsize"] = size
    
    cmap = sns.color_palette("RdBu")
    ax = sns.heatmap(df, vmax=.3, center=0, cmap=cmap, annot=annot,
                     square=True, linewidths=0, cbar_kws={"shrink": .5}, fmt='g')
    ax.set_title(title)
    plt.setp(ax.get_yticklabels(), rotation=0);
    plt.setp(ax.get_xticklabels(), rotation=x_rot);
    plt.show()