Probability Envelopes (PBE)
// This source code is subject to the terms of the Mozilla Public License 2.0 at https://mozilla.org/MPL/2.0/
// © peacefulLizard50262

//@version=5
indicator("Probability Envelopes", "PBE", precision = 3, overlay = true, max_bars_back = 100)

useLogReturns = input(false, title="Use Log Returns")
multiplier = input.float(2, "Multiplier", 0.05, 10, 0.05)
filter = input.int(6, "Smoothing Length", 1, 1000)
ripple = input.float(0.05, "Ripple", 0.01, 0.99, 0.01)
style  = input.bool(false, "Filter Style")
rounding = input.int(6, "Precision")
max_back = input.int(5000, "Max Bars Back", 0, tooltip = "Adjust this to limit the indicator's look back period. Try to minimize this if you can.")

// Custom cosh function
cosh(float x) =>
    (math.exp(x) + math.exp(-x)) / 2

// Custom acosh function
acosh(float x) =>
    x < 1 ? na : math.log(x + math.sqrt(x * x - 1))

// Custom sinh function
sinh(float x) =>
    (math.exp(x) - math.exp(-x)) / 2

// Custom asinh function
asinh(float x) =>
    math.log(x + math.sqrt(x * x + 1))

// Custom inverse tangent function
atan(float x) =>
    math.pi / 2 - math.atan(1 / x)

// Chebyshev Type I Moving Average
chebyshevI(float src, int len, float ripple) =>
    a = 0.
    b = 0.
    g = 0.
    chebyshev = 0.
    
    a := cosh(1 / len * acosh(1 / (1 - ripple)))
    b := sinh(1 / len * asinh(1 / ripple))
    g := (a - b) / (a + b)
    chebyshev := (1 - g) * src + g * nz(chebyshev[1])
    chebyshev

// Chebyshev Type II Moving Average
chebyshevII(float src, int len, float ripple) =>
    a = 0.
    b = 0.
    g = 0.
    chebyshev = 0.
    
    a := cosh(1 / len * acosh(1 / ripple))
    b := sinh(1 / len * asinh(ripple))
    g := (a - b) / (a + b)
    chebyshev := (1 - g) * src + g * nz(chebyshev[1], src)
    chebyshev

// adjust the ripple value based on your requirement (between 0 and 1)

chebyshev(float src, int length, float ripple, bool style) =>
    not style ?
     chebyshevI(src, length, ripple) :
     chebyshevII(src, length, ripple)

percent(a, b)=>
    ((a-b)/b) * 100

log_return(a, b) =>
    math.log(a / b)

ema(float source = close, float length = 9)=>
    alpha = 2 / (length + 1)
    var float smoothed = na
    smoothed := alpha * source + (1 - alpha) * nz(smoothed[1])

addPercentage(number, percentage) =>
    result = number * (1 + percentage / 100)
    result

addLogReturn(number, log_return) =>
    result = number * math.exp(log_return)
    result

add(a, b) =>
    switch useLogReturns
        true => addLogReturn(a, b)
        false => addPercentage(a, b)

var movement_size_green = array.new<float>(1)
var count_green = array.new<int>(1)
var movement_size_red = array.new<float>(1)
var count_red = array.new<int>(1)
var float expected_mean = 0
var float mean_green = 0
var float std_dev_green = 0
var float mean_red = 0
var float std_dev_red = 0
var float expected_high_price = 0
var float expected_low_price = 0


if bar_index > max_back
    change = useLogReturns ? math.round(log_return(close, close[1]), rounding) : math.round(percent(close, close[1]), rounding)

    if close > open
        if array.includes(movement_size_green, change)
            index = array.indexof(movement_size_green, change)
            array.set(count_green, index, array.get(count_green, index) + 1)
        else
            array.push(movement_size_green, change)
            array.push(count_green, 1)

    if close < open
        if array.includes(movement_size_red, change)
            index = array.indexof(movement_size_red, change)
            array.set(count_red, index, array.get(count_red, index) + 1)
        else
            array.push(movement_size_red, change)
            array.push(count_red, 1)

    probability_green = array.new<float>(array.size(count_green))
    probability_red   = array.new<float>(array.size(count_red))

    count_green_sum = array.sum(count_green)
    count_red_sum = array.sum(count_red)

    for i = 0 to array.size(count_green) - 1
        probability = array.get(count_green, i) / count_green_sum * 100
        array.set(probability_green, i, probability)

    for i = 0 to array.size(count_red) - 1
        probability = array.get(count_red, i) / count_red_sum * 100
        array.set(probability_red, i, probability)

    mean_green := array.avg(movement_size_green) 
    mean_red := array.avg(movement_size_red) 

    squared_diff_green = array.new<float>(array.size(movement_size_green))
    squared_diff_red = array.new<float>(array.size(movement_size_red))

    for i = 0 to array.size(movement_size_green) - 1
        diff = array.get(movement_size_green, i) - mean_green
        array.set(squared_diff_green, i, diff * diff)

    for i = 0 to array.size(movement_size_red) - 1
        diff = array.get(movement_size_red, i) - mean_red
        array.set(squared_diff_red, i, diff * diff)

    variance_green = array.avg(squared_diff_green)
    variance_red = array.avg(squared_diff_red)

    std_dev_green := math.sqrt(variance_green)
    std_dev_red := math.sqrt(variance_red)

    // Calculate the expected green and red movements
    expected_green_movement = 0.0
    expected_red_movement = 0.0
    for i = 0 to array.size(movement_size_green) - 1
        expected_green_movement := nz(expected_green_movement[1]) + array.get(movement_size_green, i) * array.get(probability_green, i) / 100

    for i = 0 to array.size(movement_size_red) - 1
        expected_red_movement := nz(expected_red_movement[1]) + array.get(movement_size_red, i) * array.get(probability_red, i) / 100

    // Calculate the total expected movement (weighted average)
    expected_mean_movement = (count_green_sum * expected_green_movement + count_red_sum * expected_red_movement) / (count_green_sum + count_red_sum)

    // Apply the expected mean movement to the current open price
    expected_mean := chebyshev(close + expected_mean_movement * (useLogReturns ? 1 : 0.01), filter, ripple, style)
plot(expected_mean, "Expected Mean", color=color.orange)

plot(add(expected_mean, mean_green + std_dev_green * multiplier), title="Lower Green Band", color=color.green)
plot(add(expected_mean, mean_red - std_dev_red * multiplier), title="Upper Red Band", color=color.red)
