ATR Bands (ATRBANDS)

// Licensed under the Apache License, Version 2.0
// © mihakralj
//@version=6
indicator("ATR Bands (ATRBANDS)", "ATRBANDS", overlay=true)

//@function Calculates ATR Bands using ATR for width
//@param source Source series for the center line
//@param length Period for ATR and MA calculations
//@param multiplier ATR multiplier for band width
//@returns tuple with [middle, upper, lower] band values
//@optimized Uses RMA with warmup compensator, O(1) complexity per bar
atrbands(series float source, simple int length, simple float multiplier) =>
    if length <= 0 or multiplier <= 0.0
        runtime.error("Length and multiplier must be greater than 0")
    var float prevClose = close
    float tr1 = high - low
    float tr2 = math.abs(high - prevClose)
    float tr3 = math.abs(low - prevClose)
    float trueRange = math.max(tr1, tr2, tr3)
    prevClose := close
    var int p = math.max(1, length)
    var int head = 0
    var int count = 0
    var array<float> bufferSource = array.new_float(p, na)
    var array<float> bufferTR = array.new_float(p, na)
    var float sumSource = 0.0
    var float sumTR = 0.0
    float oldestSource = array.get(bufferSource, head)
    float oldestTR = array.get(bufferTR, head)
    if not na(oldestSource)
        sumSource -= oldestSource
        sumTR -= oldestTR
        count -= 1
    float currentSource = nz(source)
    float currentTR = nz(trueRange)
    sumSource += currentSource
    sumTR += currentTR
    count += 1
    array.set(bufferSource, head, currentSource)
    array.set(bufferTR, head, currentTR)
    head := (head + 1) % p
    var float EPSILON = 1e-10
    var float raw_rma = 0.0
    var float e = 1.0
    float atrValue = na
    if not na(trueRange)
        float alpha = 1.0 / float(length)
        raw_rma := (raw_rma * (length - 1) + trueRange) / length
        e := (1 - alpha) * e
        atrValue := e > EPSILON ? raw_rma / (1.0 - e) : raw_rma
    float middleBand = nz(sumSource / count, source)
    float width = nz(atrValue * multiplier)
    [middleBand, middleBand + width, middleBand - width]

// ---------- Main loop ----------

// Inputs
i_source = input.source(close, "Source")
i_length = input.int(20, "Length", minval=1)
i_mult = input.float(2.0, "ATR Multiplier", minval=0.001)

// Calculation
[middle, upper, lower] = atrbands(i_source, i_length, i_mult)

// Plot
plot(middle, "Middle", color=color.yellow, linewidth=2)
p1 = plot(upper, "Upper", color=color.yellow, linewidth=2)
p2 = plot(lower, "Lower", color=color.yellow, linewidth=2)
fill(p1, p2, color=color.new(color.blue, 90), title="Band Fill")