# Just because QC requires a main.py file to build correctly.
# Launch example_usage.ipynb and inspect notebook_graphing_modules.py!

class non(QCAlgorithm):

    def Initialize(self):
        # Set Start Date so that backtest has 5+ years of data
        self.SetStartDate(2015, 1, 6)

        # No need to set End Date as the final submission will be tested
        # up until the review date

        # Set $1m Strategy Cash to trade significant AUM
        self.SetCash(1000000)

        # Add a relevant benchmark, with the default being SPY
        self.AddEquity('SPY')
        self.SetBenchmark('SPY')

        # Use the Alpha Streams Brokerage Model, developed in conjunction with
        # funds to model their actual fees, costs, etc.
        # Please do not add any additional reality modelling, such as Slippage, Fees, Buying Power, etc.
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())




    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        # if not self.Portfolio.Invested:
        #   self.SetHoldings("SPY", 1)

'''
Created January 2020 by Ian Rolin Larson | www.ianrlarson.com
----released for unlimited modification and distribution----

Usage: from notebook_graphing_modules import g

    Pass g a string ticker, eg. "GOOG", and it will output a bokeh plot in notebook
using its default analysis module, included in this file.
Pass it a list of tickers and it will output a column of full-sized plots--so you can check your
trading logic on multiple symbols simultaneously.

    g takes 2 important arguments, the symbol or list of symbols and the name of a function--without parentheses.
(There are 2 optional keyword arguments, 'bars' and resolution. Defaults to 360 and 'd'.)
This analysis function is what you're supposed to create. Each dataframe passes this function
(open,high,low,close) prices. You can use objects or global variables to store previous values for more
complex analysis. You return either a start price (whichever price)
and an end price (whichever extreme you want to draw to) or (None,None). These prices can be anywhere on the chart, but
but I recommend drawing from the top or bottom of a candlestick to the high or low respectively.
This will highlight a wick--the highlight is rectangular and as wide as a candle body
By default the chart candles are black and white for clarity and the drawing
is in red to stand out. Multiple colors could be used by modifying the functions herein to add another
column to the dataframe (hence the analyze module would have 3 outputs) so you can select different colors
to indicate various types of data.

    d takes the same inputs g does but outputs the dataframes instead of a complete unit
    plot takes a dataframe or list of dfs with start, end, date columns and outputs plot(s)
These functions exist so you have direct access to all the important functionality in a modularly
testable way.

    Most of the dependencies at the beginning are necessary so that we can initiazlize the quantbook
and get data as in module 1 from quantconnect so we can generate dataframes. They came from the 
research environment documentation...we need to populate the namespace to use research environment
commands. Also, this makes it so we can just import the function of interest, primarily g, but still
get data. It keeps things under wraps so you can concentrate on your analysis functions.

analyze -- takes OHLC values and outputs (start, end) drawing prices or (None,None)
mod1 -- takes the symbol(s) and generates dataframes
mod2 -- takes the dataframe(s), adds a date column, and uses a function--analyze by default--to add start and end columns
d -- puts all the above together in an easy-to-use module
plot -- takes the modified dataframe(s) from the "d" pipeline and generates plot(s) from them
g -- packages d and plot for a truly seamless 1-call experience
'''

#import jupyter dependencies so this file will build and not run into 'not found' errors
from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Jupyter")
AddReference("QuantConnect.Indicators")
from System import *
from QuantConnect import *
from QuantConnect.Data.Market import TradeBar, QuoteBar
from QuantConnect.Jupyter import *
from QuantConnect.Indicators import *
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import pandas as pd

#initialize quantbook so don't have to do it in the notebook
qb = QuantBook()

##----------CUSTOM MODULE AND GRAPHING CODE FOLLOWS:--------------


##EXAMPLE, DEFAULT WICK ANALYZE MODULE--DOES NOT DO ANYTHING USEFUL
#global variables needed for your analyze function:
count = -1

#analyze each candle in a dataframe
def analyze(o,h,l,c):
    '''
    input: open, high, low, close
    output: start, end -- these will either be price values, or there will be nothing returned
    '''
    global count #declare global variables to keep track of values to persist from one candle to the next
    count += 1 #count starts at -1 to compensate for advanced position
    
    #test implimentation logic -- does nothing useful, for illustrative purposes only
    if count % 5 == 0:
        e=h
        if c>o:
            s = c
        else:
            s = o
        return (s, e)
    else:
        return (None, None)
    
##module 1 as above:

def mod1(symbol,bars,resolution):
    '''Requires qb to be set up already, subscribes to new asset data
    input: string or list of strings, where these are tickers; optional int "bars="; optional "resolution="'d','h','m'"
    Output: Pandas dataframe or list of dataframes of the OHLC series(s) requested'''
    
    def action(i):
        sym = qb.AddEquity(i)
        return qb.History([i], bars, res[resolution])
    
    res = {
        'd' : Resolution.Daily,
        'h' : Resolution.Hour,
        'm' : Resolution.Minute
    }
    
    if type(symbol) == str:
        return action(symbol)
    else:
        assert type(symbol) == list, 'Symbol(s)--tickers--need to be passed in either as a single string or list of strings. \
        If you get an error from the quant connect modules it\'s likely your symbol(s) are invalid.'
        return [action(e)for e in symbol]

##module 2 as above:
def mod2(data,func):
    
    def action(df):
        '''takes in dataframe, then feeds OHLC to ANALYSIS module "analyze", then returns dataframe(s)'''
        # iterating over multiple columns
        #call analyze in a loop, pass it row values, update new columns
        #doing list comprehension instead as from here: https://stackoverflow.com/questions/16476924/how-to-iterate-over-rows-in-a-dataframe-in-pandas/55557758#55557758
        #result = [f(row[0], ..., row[n]) for row in df[['col1', ...,'coln']].values]
        # add date column so can graph it better
        df['date']=[pd.to_datetime(e[1]) for e in df.index]
        #apply analysis function to OHLC values
        result = [func(row[3],row[1],row[2],row[0]) for row in df[['close','high','low','open']].values]
        #add columns to dataframe with NA's
        df['start'] = [e[0] for e in result]
        df['end'] = [e[1] for e in result]
        
        return df
    
    if type(data) == list:
        return [action(e)for e in data]
    else:
        return action(data)
        
#package them all together for easy usage in the notebook
def d(symbol,func=analyze, bars=360, resolution='d'):
    '''Takes all the modules contained in this file as input, assembles the graphing pipeline, shows the plot'''
    return mod2(mod1(symbol,bars,resolution),func)
    
##MULTI PLOTTING MODULE
'''input list of dfs or df with date, start, end columns specifying wicks to highlight, output plot or multiple plots'''
def plot(df):
    from bokeh.plotting import figure, show
    from bokeh.io import output_notebook
    from bokeh.layouts import column
    
    w = 12*60*60*1000 # half day in ms
    TOOLS = "pan,wheel_zoom,box_zoom,reset,save"
    
    def pl(df):
        inc = df.close > df.open
        dec = df.open > df.close
        #Plot width of 963 was the maximum fit in my QC jupyter notebook without having to scroll
        p = figure(x_axis_type="datetime", tools=TOOLS, plot_width=963, title = (df.index[0][0]).split()[0], toolbar_location = "above", active_scroll = 'wheel_zoom')
        p.xaxis.major_label_orientation = .7 #use this instead of Pi/4...why have an additional import?
        p.grid.grid_line_alpha=0.3
        #candles
        p.segment(df.date, df.high, df.date, df.low, color="black")
        p.vbar(df.date[inc], w, df.open[inc], df.close[inc], fill_color='white', line_color="black")
        p.vbar(df.date[dec], w, df.open[dec], df.close[dec], fill_color='black', line_color="black")
        #highlight wicks
        #p.circle(df.date, df.target, color="red", fill_alpha=0.2, size=df.radius)
        p.vbar(df.date, w, df.start, df.end, fill_color="red", line_color="red", fill_alpha=0.2)
        return p
    
    if type(df) == list:
        p = [pl(i) for i in df]
    else:
        p = pl(df)
     
    output_notebook()
    show(column(p))
    
#put it ALL together...including the graphing!!   
def g(symbol,func=analyze, bars=360, resolution='d'):
    plot(d(symbol,func,bars,resolution))