ZigZag Chart with Supertrend
// This source code is subject to the terms of the Mozilla Public License 2.0 at https://mozilla.org/MPL/2.0/
// © LonesomeTheBlue

//@version=4
study("ZigZag Chart with SuperTrend",  max_bars_back = 500, max_lines_count = 500, explicit_plot_zorder = true)
prd = input(defval = 4, title="ZigZag Period", minval = 2, maxval = 50, group = "Setup")
showmast = input(defval = "Supertrend", title = "Show => ", options = ["Moving Average", "Supertrend"], group = "Setup", inline = "mova")
malen = input(defval = 20, title=" MA Length", minval = 2, maxval = 200, group = "Setup", tooltip = "Used for Moving Average")
atrlen = input(defval = 10, title=" ATR Length", minval = 2, maxval = 50, group = "Setup", tooltip = "Used for Supertrend")
atrmult = input(defval = 2., title=" ATR Multiplier", minval = 0.1, maxval = 10., step = 0.1, group = "Setup", tooltip = "Used for Supertrend")
bodycolup = input(defval = color.lime, title = "Body Color", inline = "bcol", group = "Colors")
bodycoldown = input(defval = color.red, title = "", inline = "bcol", group = "Colors")
bodycoltop = input(defval = #b5b5b8, title = "Wick Color", inline = "wcol", group = "Colors")
bodycolbottom = input(defval = #b5b5b8, title = "", inline = "wcol", group = "Colors")
chbarcolor = input(defval = true, title = "Change Bar Color by Supertrend", tooltip = "Supertrend should be activated using the option above", group = "Colors")

// get/keep highest/lowest closing prices
float ph = highestbars(close, prd) == 0 ? close : na
float pl = lowestbars(close, prd) == 0 ? close : na
var dir = 0
dir := iff(ph and na(pl), 1, iff(pl and na(ph), -1, dir))
var max_array_size = 320
var zigzag = array.new_float(0)

add_to_zigzag(value, bindex)=>
    array.unshift(zigzag, bindex)
    array.unshift(zigzag, value)
    if array.size(zigzag) > max_array_size
        array.pop(zigzag)
        array.pop(zigzag)
    
update_zigzag(value, bindex)=>
    if array.size(zigzag) == 0
        add_to_zigzag(value, bindex)
    else
        if (dir == 1 and value > array.get(zigzag, 0)) or (dir == -1 and value < array.get(zigzag, 0))
            array.set(zigzag, 0, value)
            array.set(zigzag, 1, bindex)
        0.

bool dirchanged = (dir != dir[1])
if ph or pl
    if dirchanged
        add_to_zigzag(dir == 1 ? ph : pl, bar_index)
    else
        update_zigzag(dir == 1 ? ph : pl, bar_index)

// definition for candles, moving averages and supertrend
var zzcandles = array.new_line(0)
var zztopwicks = array.new_line(0)
var zzbottomwicks = array.new_line(0)
var maarray = array.new_line(0)
var float zz_high = na
var float zz_low = na
var zigzagtrueRange = array.new_float(0)
var zigzagatr = array.new_float(0)
varip float st_up = na
varip float st_dn = na
varip float st_trendup = na
varip float st_trenddn = na
float st_trendup_1 = st_trendup
float st_trenddn_1 = st_trenddn
varip int st_trend = 1
var zzstdir = array.new_float(0)
var zzsupertrend = array.new_float(0)

rma_sum(length)=>
    array.size(zigzagtrueRange) >= length ? array.avg(array.slice(zigzagtrueRange, 0, length)) : na
    
// calculate rma for truerange of zigzag candles 
zigzag_rma(length)=>
	float alpha = 1. / length
	float ret = array.size(zigzagatr) == 0 ? rma_sum(length) : 
	                                         na(array.get(zigzagatr, 0)) ? rma_sum(length) : 
	                                         alpha * array.get(zigzagtrueRange, 0) + (1 - alpha) * array.get(zigzagatr, 0)

// shift left the candles
shift_left_candles()=>
    // 3 "for" loops, not to get loop error
    for x = 0 to (array.size(zzcandles) > 0 ? array.size(zzcandles) - 1 : na)
        line.set_x1(array.get(zzcandles, x), bar_index - x)
        line.set_x2(array.get(zzcandles, x), bar_index - x)
    for x = 0 to (array.size(zztopwicks) > 0 ? array.size(zztopwicks) - 1 : na)
        line.set_x1(array.get(zztopwicks, x), bar_index - x)
        line.set_x2(array.get(zztopwicks, x), bar_index - x)
    for x = 0 to (array.size(zzbottomwicks) > 0 ? array.size(zzbottomwicks) - 1 : na)
        line.set_x1(array.get(zzbottomwicks, x), bar_index - x)
        line.set_x2(array.get(zzbottomwicks, x), bar_index - x)

// calculate ATR for zigzag candles and keep them in an array
zigzag_atr(length, insert)=>
    float hi = array.size(zztopwicks) > 0 ? line.get_y1(array.get(zztopwicks, 0)) : na
    float hi_1 = array.size(zztopwicks) > 1 ? line.get_y1(array.get(zztopwicks, 1)) : na
    float cl = array.size(zzcandles) > 0 ? line.get_y2(array.get(zzcandles, 0)) : na
    float cl_1 = array.size(zzcandles) > 1 ? line.get_y2(array.get(zzcandles, 1)) : na
    float lo = array.size(zzbottomwicks) > 0 ? line.get_y1(array.get(zzbottomwicks, 0)) : na
    
    if insert
        if array.size(zigzagtrueRange) > 100
            array.pop(zigzagtrueRange)
        if array.size(zigzagatr) > 100
            array.pop(zigzagatr)
        array.unshift(zigzagtrueRange, na(hi_1) ? hi - lo : max(max(hi - lo, abs(hi - cl[1])), abs(lo - cl_1)))
        array.unshift(zigzagatr, zigzag_rma(length))
    else
        array.set(zigzagtrueRange, 0, na(hi_1) ? hi - lo : max(max(hi - lo, abs(hi - cl[1])), abs(lo - cl_1)))
        array.set(zigzagatr, 0, zigzag_rma(length))

// calculate/show zigzag candles, calculate zigza atr, zigzag supertrend and keep them in the arrays
if array.size(zigzag) >= 4
    zz_high := max(zz_high, high)
    zz_low := min(zz_low, low)
    if dirchanged
        // draw/keep candle body
        array.unshift(zzcandles, 
                      line.new(x1 = bar_index, y1 = array.get(zigzag, 2), x2 = bar_index, y2 = array.get(zigzag, 0), 
                               color = array.get(zigzag, 0) >= array.get(zigzag, 2) ? bodycolup : bodycoldown, 
                               width = 5))
        
        // draw/keep top wicks
        zz_high := max(high, array.get(zigzag, 0), array.get(zigzag, 2))
        zz_low := min(low, array.get(zigzag, 0), array.get(zigzag, 2))
        array.unshift(zztopwicks,
                      line.new(x1 = bar_index, y1 = zz_high, x2 = bar_index, y2 = max(array.get(zigzag, 2), array.get(zigzag, 0)), 
                               color = bodycoltop))
        
        // draw/keep bottom wicks
        array.unshift(zzbottomwicks,
                      line.new(x1 = bar_index, y1 = zz_low, x2 = bar_index, y2 = min(array.get(zigzag, 2), array.get(zigzag, 0)), 
                               color = bodycolbottom))
        
        // remove old array elements
        if array.size(zzcandles) > 125
            line.delete(array.pop(zzcandles))
            line.delete(array.pop(zztopwicks))
            line.delete(array.pop(zzbottomwicks))
        
        // insert new atr value to the array
        zigzag_atr(atrlen, true)
    else
        // set close/high/low of the candle
        line.set_y2(array.get(zzcandles, 0), array.get(zigzag, 0))
        line.set_y1(array.get(zztopwicks, 0), zz_high)
        line.set_y2(array.get(zztopwicks, 0), max(array.get(zigzag, 2), array.get(zigzag, 0)))
        line.set_y1(array.get(zzbottomwicks, 0), zz_low)
        line.set_y2(array.get(zzbottomwicks, 0), min(array.get(zigzag, 2), array.get(zigzag, 0)))
        
        // update atr value in the array
        zigzag_atr(atrlen, false)
    
    // shift left the candles
    shift_left_candles()
    
    // calculate zigzag supertrend
    float cl = array.size(zzcandles) > 0 ? line.get_y2(array.get(zzcandles, 0)) : na
    float cl_1 = array.size(zzcandles) > 1 ? line.get_y2(array.get(zzcandles, 1)) : na
    float hi = array.size(zztopwicks) > 0 ? line.get_y1(array.get(zztopwicks, 0)) : na
    float lo = array.size(zzbottomwicks) > 0 ? line.get_y1(array.get(zzbottomwicks, 0)) : na
    hl2_ = (hi + lo) / 2
    st_up := hl2_ - (atrmult * array.get(zigzagatr, 0))
    st_dn := hl2_ + (atrmult * array.get(zigzagatr, 0))
    
    st_trendup := cl_1 > st_trendup_1 ? max(st_up, st_trendup_1) : st_up
    st_trenddn := cl_1 < st_trenddn_1 ? min(st_dn, st_trenddn_1) : st_dn
    
    st_trend := cl > st_trenddn_1 ? 1 : cl < st_trendup_1 ? -1 : nz(st_trend, 1)
    if dirchanged
        if array.size(zzsupertrend) > 124
            array.pop(zzsupertrend)
            array.pop(zzstdir)
        array.unshift(zzsupertrend, st_trend == 1 ? st_trendup : st_trenddn)
        array.unshift(zzstdir, st_trend)
    else
        array.set(zzsupertrend, 0, st_trend == 1 ? st_trendup : st_trenddn)
        array.set(zzstdir, 0, st_trend)

// show zigzag moving average if enabled
if showmast == "Moving Average" and array.size(zigzag) > malen * 2 + 2
    for x = 0 to (array.size(maarray) > 0 ? array.size(maarray) - 1 : na)
        line.delete(array.pop(maarray))
    
    float oldpoint = na
    for x = 0 to (array.size(zigzag) > 0 ? array.size(zigzag) - 1 : na) by 2
        if x > array.size(zigzag) - malen * 2 - 1 or x >= 250
            break
        float zzma = 0.0
        for y = x to x + malen * 2 - 1 by 2
            zzma += array.get(zigzag, y) / malen
        
        if not na(oldpoint)
            array.unshift(maarray, 
                          line.new(x1 = bar_index - round(x / 2) + 1, y1 = oldpoint, x2 = bar_index - round(x / 2), y2 = zzma, 
                          color = color.blue, 
                          width = 2))
        oldpoint := zzma

// show zigzag supertrend if enabled
if showmast == "Supertrend" and array.size(zigzag) > atrlen * 2 + 2
    var zzstlines = array.new_line(0)
    var zzstlabels = array.new_label(0)
    for x = 0 to (array.size(zzstlines) > 0 ? array.size(zzstlines) - 1 : na)
        line.delete(array.pop(zzstlines))
    for x = 0 to (array.size(zzstlabels) > 0 ? array.size(zzstlabels) - 1 : na)
        label.delete(array.pop(zzstlabels))
    
    // 2 "for" loops, not to get loop error
    for x = 0 to (array.size(zzsupertrend) > 1 ? array.size(zzsupertrend) - 2 : na)
        if array.get(zzstdir, x) == array.get(zzstdir, x + 1)
            array.unshift(zzstlines, 
                          line.new(x1 = bar_index - x, y1 = array.get(zzsupertrend, x), x2 = bar_index - x - 1, y2 = array.get(zzsupertrend, x + 1), 
                                   color = array.get(zzstdir, x) == 1 ? bodycolup : bodycoldown, 
                                   width = 2))
                                   
    for x = 0 to (array.size(zzsupertrend) > 1 ? array.size(zzsupertrend) - 2 : na)    
        if array.get(zzstdir, x) != array.get(zzstdir, x + 1)
            array.unshift(zzstlabels, 
                          label.new( x = bar_index - x, y = array.get(zzsupertrend, x), 
                                     color = array.get(zzstdir, x) == 1 ? bodycolup : bodycoldown,
                                     style = array.get(zzstdir, x) == 1 ? label.style_triangleup : label.style_triangledown, 
                                     size = size.tiny))

// changr bar color by zigzag supertrend if enabled
barcolor(chbarcolor and showmast == "Supertrend" ? (st_trend == 1 ? bodycolup : bodycoldown) : na)
 