ma_function
// This source code is subject to the terms of the Mozilla Public License 2.0 at https://mozilla.org/MPL/2.0/
// © boitoki

//@version=5

// @description: This library is a package of various MAs.
// @sample:
//   import boitoki/ma_function/{version} as ma
//
//   ma.calc('SMA', close, 10)
//  
library("ma_function")

export sma (float x, int y) =>
    sum = 0.0
    for i = 0 to y - 1
        sum := sum + x[i] / y
    sum

export ema (float src, int length) =>
    alpha = 2 / (length + 1)
    sum = 0.0
    sum := na(sum[1]) ? src : alpha * src + (1 - alpha) * nz(sum[1])

export kama (float xPrice, int Length)=>
    xvnoise = math.abs(xPrice - xPrice[1])
    nfastend = 0.666
    nslowend = 0.0645
    nsignal = math.abs(xPrice - xPrice[Length])
    nnoise = math.sum(xvnoise, Length)
    nefratio = (nnoise != 0 ? nsignal / nnoise : 0)
    nsmooth = math.pow(nefratio * (nfastend - nslowend) + nslowend, 2) 
    nAMA = 0.0
    nAMA := nz(nAMA[1]) + nsmooth * (xPrice - nz(nAMA[1]))

export zlema (float src, int length) =>
    zxLag = length/2==math.round(length/2) ? length/2 : (length - 1) / 2
    zxEMAData = (src + (src - src[zxLag]))
    ZLEMA = ema(zxEMAData, length)

export mav (float source, int length) => 
    ((ta.sma(source, length)[1]*(length-1)) + source) / length

export dema (float src, int len) =>
    e = ema(src, len)
    2 * e - ema(e, len)

export tema (float src, int len) =>
    e = ema(src, len)
    3 * (e - ema(e, len)) + ema(ema(e, len), len)

export hma (float src, int len) =>
    ta.wma(2 * ta.wma(src, len / 2) - ta.wma(src, len), math.round(math.sqrt(len)))

export  tma (float source, int length) =>
    ta.sma(ta.sma(source, math.ceil(length / 2)), math.floor(length / 2) + 1)

export ama (float source, int length, int fast, int slow) =>
    fastAlpha = 2 / (fast + 1)
    slowAlpha = 2 / (slow + 1)
    
    hh = ta.highest(length + 1)
    ll = ta.lowest(length + 1)
    mltp = hh - ll != 0 ? math.abs(2 * source - ll - hh) / (hh - ll) : 0
    ssc = mltp * (fastAlpha - slowAlpha) + slowAlpha
    
    ama = 0.0
    ama := nz(ama[1]) + math.pow(ssc, 2) * (source - nz(ama[1]))

    ama

export rma (float src, int length) =>
	alpha = 1 / length
	sum = 0.0
	sum := na(sum[1]) ? sma(src, length) : alpha * src + (1 - alpha) * nz(sum[1])

export bwma (float src, int length, int beta, int alpha) =>
    var b = array.new_float(0)
    if barstate.isfirst
        for i = 0 to length-1
            x = i / (length-1)
            w = math.pow(x, alpha-1) * math.pow(1-x, beta-1)
            array.push(b,w)

    den = array.sum(b)
    sum = 0.
    for i = 0 to length-1
        sum := sum + src[i]*array.get(b,i)
    filt = sum/den

// @function This function provides the MA specified in the argument.
// @param name MA's name
// @param source Source value
// @param length Look back periods
// @returns MA value
export ma_function (simple string name, float source, int length) =>
	switch name
	    'ALMA'  => ta.alma(source, length,0.85, 6)
	    'AMA'   => ama(source, length, 2, 30)
        'BWMA'  => bwma(source, length, 3, 3)
	    'DEMA'  => dema(source, length)
	    'EMA'   => ema(source, length)
	    'HMA'   => hma(source, length)
	    'KAMA'  => kama(source, length)
	    'LSMA'  => ta.linreg(source, length, 0)
	    'MAV'   => mav(source, length)
	    'RMA'   => rma(source, length)
	    'SMA'   => ta.sma(source, length)
	    'SWMA'  => ta.swma(source)
	    'TEMA'  => tema(source, length)
	    'TMA'   => tma(source, length)
	    'VWMA'  => ta.vwma(source, length)
	    'WMA'   => ta.wma(source, length)
	    'ZLEMA' => zlema(source, length)

// This function is a syntax sugar of `ma_function`
export calc (simple string name, float source, int length) =>
    ma_function(str.upper(name), source, length)

i_source    = input.source(close, 'Source')
i_type      = input.string('SMA', 'Type', options=['ALMA', 'AMA', 'BWMA', 'DEMA', 'EMA', 'HMA', 'KAMA', 'LSMA', 'MAV', 'RMA', 'SMA', 'SWMA', 'TEMA', 'TMA', 'VWMA', 'WMA', 'ZLEMA'])
i_len       = input.int(10, 'Length')
i_color     = input.color(color.blue, 'Color')

plot(calc(i_type, i_source, i_len), 'MA', color=i_color)
