""" 
Covered puts 
"""

from datetime import timedelta


class AdaptableYellowGreenJaguar(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 1)
        self.SetEndDate(2020, 1, 1)
        self.SetCash(100000)
        self.ticker = "SPY"
        self.equity = self.AddEquity(self.ticker, Resolution.Minute)
        self.option = self.AddOption(self.ticker, Resolution.Minute)
        self.symbol = self.option.Symbol
        self.option.SetFilter(-100, 0, timedelta(0), timedelta(30))
        self.SetBenchmark(self.equity.Symbol)
        # TODO: Schedule once a month
        
    def OnData(self, data):
        if not self.Portfolio[self.ticker].Invested:
            self.SetHoldings(self.equity.Symbol, 0.99)
        
        option_invested = [x.Key for x in self.Portfolio if x.Value.Invested and 
                           x.Value.Type==SecurityType.Option]
        if len(option_invested) < 1:
            self.TradeOptions(data) 
            
    def TradeOptions(self, data):
        for i in data.OptionChains:
            if i.Key != self.symbol: continue
            chain = i.Value
            call = [x for x in chain if x.Right == OptionRight.Put] 
            # sorted the contracts according to their expiration dates and choose the ATM options
            contracts = sorted(sorted(call, key = lambda x: abs(chain.Underlying.Price - x.Strike)), 
                                            key = lambda x: x.Expiry, reverse=True)
            [self.print_option(c) for c in contracts]
            if len(contracts) == 0: return    
            put = contracts[0].Symbol
            self.SetHoldings(self.equity.Symbol, 0.01)
    
    def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))
                
    
    def print_option(self, contract):
        self.Debug(f"{contract.Symbol.Value}, strike: {contract.Strike}, Delta: {contract.Greeks.Delta}, Vega: {contract.Greeks.Vega}, Last Price: {contract.LastPrice}, Underlying Price: {contract.UnderlyingLastPrice}")

class StaticBarbell(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2005, 1, 1)  # Set Start Date
        self.SetCash(10000)  # Set Strategy Cash
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
        self.stock = self.AddEquity("SPXL", Resolution.Daily)  # SPXL, FIHD
        self.safe = self.AddEquity("IEF", Resolution.Daily)  # IEF
        self.insurance = self.AddEquity("VIXY", Resolution.Daily)
        self.insu_ratio = 0.01
        self.stock_ratio = 0.20
        self.safe_ratio = 1.0 - self.insu_ratio - self.stock_ratio
        self.Schedule.On(self.DateRules.WeekStart("SPXL"),
                         self.TimeRules.At(10, 0, 0),
                         self.RebalancePortfolio)
    
    def RebalancePortfolio(self):
        self.SetPosition(self.stock, self.stock_ratio)
        self.SetPosition(self.safe, self.safe_ratio)
        self.SetPosition(self.insurance, self.insu_ratio)
            
    def SetPosition(self, asset, share):
        if self.IsMarketOpen(asset.Symbol) and asset.Price > 0:
            self.SetHoldings(asset.Symbol, share)

"""
ML Version of Black Swan Model
@version: 0.1
"""

import clr
clr.AddReference("System")
clr.AddReference("QuantConnect.Algorithm")
clr.AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework.Execution import *

import pandas as pd
pd.set_option('mode.use_inf_as_na', True)
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import MinMaxScaler
from sklearn.neural_network import MLPRegressor

import features as ft


STEPS = [("scaler", MinMaxScaler()),
         ("model", MLPRegressor(n_iter_no_change=1,
                                early_stopping=True))]


class MLCryptoAlgo(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2010, 1, 1)
        self.SetCash(100000)
        self.Settings.FreePortfolioValuePercentage = 0.05

        self.resolution = Resolution.Daily
        self.tickers = ["SPXL", "IEF", "VIXY"]
        self.SetBenchmark("SPY")
        [self.AddEquity(t, self.resolution) for t in self.tickers]
        self.lookback = 252
        self.train_days = timedelta(365 * 5)
        self.test_days = timedelta(365 * 2)
        self.test_start = None

        self.model = Pipeline(steps=STEPS)
        self.Train(self.DateRules.WeekStart(),
                   self.TimeRules.At(0, 0),
                   self.train_model)
        self.Schedule.On(self.DateRules.EveryDay("SPXL"),
                         self.TimeRules.AfterMarketOpen("SPXL", 5),
                         self.trade)

    def train_model(self):
        self.test_start = self.Time-self.test_days
        train_start = self.test_start-self.train_days
        x_train, y_train = self.get_data(train_start, self.test_start)
        self.model.fit(x_train, y_train)
        x_test, y_test = self.get_data(self.test_start, self.Time)
        score = self.model.score(x_test, y_test)
        self.Debug(f"Model score {score:.2f}")

    def trade(self):
        if self.test_start is not None:
            self.Transactions.CancelOpenOrders()
            x_pred = self.get_data(self.test_start, self.Time, include_y=False)
            x_pred = x_pred.sort_index().groupby("symbol").last()
            symbols = x_pred.index.get_level_values("symbol")
            y_pred = pd.Series(self.model.predict(x_pred), index=symbols)
            positions = y_pred.clip(lower=0)
            positions = (positions/sum(positions)).fillna(0)
            self.Debug(f"{self.Time} - {positions}")
            [self.SetHoldings(sym, pos) for sym, pos in positions.iteritems()]

    def get_data(self, start, end, include_y=True):
        tickers = list(self.ActiveSecurities.Keys)
        history = self.History(tickers, start, end, self.resolution)
        history.eval("volatility = high - low", inplace=True)
        features = ft.sequence(history[["close", "volume", "volatility"]],
                               lookback=self.lookback).dropna()
        target = history["close"].groupby("symbol").pct_change(1).shift(-1)
        if include_y:
            target = target.reindex_like(features).dropna()
            return features.loc[target.index], target
        else:
            return features

class StaticBarbell(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2005, 1, 1)  # Set Start Date
        self.SetCash(10000)  # Set Strategy Cash
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
        self.stock = self.AddEquity("SPXL", Resolution.Daily)  # SPXL, FIHD
        self.safe = self.AddEquity("IEF", Resolution.Daily)  # IEF
        self.insurance = self.AddEquity("VIXY", Resolution.Daily)
        self.insu_ratio = 0.05
        self.stock_ratio = 0.9
        self.safe_ratio = 1.0 - self.insu_ratio - self.stock_ratio
        self.Schedule.On(self.DateRules.WeekStart("SPXL"),
                         self.TimeRules.At(10, 0, 0),
                         self.RebalancePortfolio)
    
    def RebalancePortfolio(self):
        self.SetPosition(self.stock, self.stock_ratio)
        self.SetPosition(self.safe, self.safe_ratio)
        self.SetPosition(self.insurance, self.insu_ratio)
            
    def SetPosition(self, asset, share):
        if self.IsMarketOpen(asset.Symbol) and asset.Price > 0:
            self.SetHoldings(asset.Symbol, share)

class StaticBarbell(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2005, 1, 1)  # Set Start Date
        self.SetCash(10000)  # Set Strategy Cash
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
        self.positions = {"SPXL": 0.2, "IEF": 0.8} # VIXY?
        [self.AddEquity(ticker) for ticker in self.positions]
        self.Schedule.On(self.DateRules.MonthStart("SPXL"),
                         self.TimeRules.At(10, 0, 0),
                         self.RebalancePortfolio)
    
    def RebalancePortfolio(self):
        [self.SetHoldings(t, p) for t, p in self.positions.items()]