KNN ATR Dual Range Predictions [SS]
// This source code is subject to the terms of the Mozilla Public License 2.0 at https://mozilla.org/MPL/2.0/
//  /$$$$$$   /$$                                                         /$$                                            
// /$$__  $$ | $$                                                        | $$                                            
//| $$  \__//$$$$$$    /$$$$$$  /$$    /$$ /$$$$$$   /$$$$$$   /$$$$$$$ /$$$$$$    /$$$$$$  /$$    /$$ /$$$$$$   /$$$$$$$
//|  $$$$$$|_  $$_/   /$$__  $$|  $$  /$$//$$__  $$ /$$__  $$ /$$_____/|_  $$_/   /$$__  $$|  $$  /$$//$$__  $$ /$$_____/
// \____  $$ | $$    | $$$$$$$$ \  $$/$$/| $$$$$$$$| $$  \__/|  $$$$$$   | $$    | $$$$$$$$ \  $$/$$/| $$$$$$$$|  $$$$$$ 
// /$$  \ $$ | $$ /$$| $$_____/  \  $$$/ | $$_____/| $$       \____  $$  | $$ /$$| $$_____/  \  $$$/ | $$_____/ \____  $$
//|  $$$$$$/ |  $$$$/|  $$$$$$$   \  $/  |  $$$$$$$| $$       /$$$$$$$/  |  $$$$/|  $$$$$$$   \  $/  |  $$$$$$$ /$$$$$$$/
// \______/   \___/   \_______/    \_/    \_______/|__/      |_______/    \___/   \_______/    \_/    \_______/|_______/ 

// © Steversteves
//@version=5
indicator("KNN ATR Dual Range Predictions [SS]", overlay=true, max_bars_back = 5000)

// Groups // 
g1 = "Settings"
g2 = "Primary Targets" 
g3 = "Secondary Targets"
// Settings 
train = input.int(150, "Training Time", group = g1) 
reg_type = input.string("Last Instance", "KNN Regression Type", ["Last Instance", "Cluster"], group = g1)
clust = input.int(3, "Clusters", group = g1) 
tol = input.float(0.5, "Tolerance", group = g1) 
stats = input.bool(true, "Show Regression Statistics", group = g1)

// Primary Targets 
prim_tf = input.bool(true, "Show Primary Timeframe Targets", group = g2)
fill = input.bool(true, "Fill Primary Targets", group = g2) 

// Secondary Targets 
alt_tf = input.bool(true, "Show Secondary Timeframe Targets", group = g3)
timeframe = input.timeframe("D", "Secondary Timeframe", group = g3)

// Regression 

knn_regression(float y, float x, int len, string last_instance_or_cluster_or_avg, int clusters, float tolerance) => 
    int index_of = 0
    float result = 0.0
    float error = 0.0 
    int inst = 0 
    float y_result = 0.0 
    float cor = 0.0 
    float r2 = 0.0 
    y_in = y[1] 
    ed_y = math.abs(math.sqrt(math.pow(y[1] - x[1],2)))
    ed_x = math.abs(math.sqrt(math.pow(y - x, 2)))
    ed_array = array.new<float>() 
    y_array = array.new<float>() 
    cluster_array = array.new<float>() 
    cluster_avg = array.new<float>() 
    for i = 0 to len 
        array.push(ed_array, ed_y[i]) 
        array.push(y_array, y_in[i])
    // Find KNN 
    if last_instance_or_cluster_or_avg == "Last Instance"
        if array.includes(ed_array, ed_x) 
            index_of := array.indexof(ed_array, ed_x) 
            result := array.get(y_array, index_of) 
        else if result[1] < 0 
            result := result[1] 
        else if result[2] < 0 
            result := result[2] 
        else if result[3] < 0 
            result := result[3] 
        else 
            na 
    else if last_instance_or_cluster_or_avg == "Avg" 
        for i = 0 to array.size(ed_array) - 1 
            if array.get(ed_array, i) == ed_x
                y_result += array.get(y_array, i)
                inst += 1 
                
        result := y_result / inst 
    else if last_instance_or_cluster_or_avg == "Cluster" 
        for i = 0 to array.size(ed_array) - 1 
            if array.get(ed_array, i) >= ed_x - tolerance and array.get(ed_array, i) <= ed_x + tolerance 
                array.push(cluster_array, array.get(y_array, i))
        if array.size(cluster_array) - 1 > clusters 
            for i = 0 to clusters 
                array.push(cluster_avg, array.get(cluster_array, i))
            result := array.avg(cluster_avg) 


    // Error 
    bool above_result = y > result 
    bool below_result = y < result 
    above_arr = array.new<float>() 
    below_arr = array.new<float>() 
    above_dif = y - result 
    below_dif = result - y 
    for i = 0 to len 
        if above_result[i]
            array.push(above_arr, above_dif[i])
        else if below_result[i] 
            array.push(below_arr, below_dif[i])
        else 
            na 
    avg1 = array.avg(above_arr) 
    avg2 = array.avg(below_arr) 
    error := math.avg(avg1, avg2) 
    // Correlation 
    result_array = array.new<float>() 
    for i = 0 to array.size(y_array) - 1 
        array.push(result_array, result[i])
    cov = array.covariance(result_array, y_array) 
    y_sd = array.stdev(y_array) 
    r_sd = array.stdev(result_array) 
    cor := cov / (y_sd * r_sd) 
    // R2 
    r2 := math.pow(cor, 2) 
    [result, error, cor, r2]

// Calcs 
[hitoop_tf, optolo_tf, op_tf, ti] = request.security(syminfo.tickerid, timeframe, [high[1] - open[1], open[1] - low[1], open, time], lookahead = barmerge.lookahead_on)

hitoop = (high[1] - open[1]) 
optolo = (open[1] - low[1]) 

[h_r, h_er, h_cor, h_r2] = knn_regression(hitoop, hitoop[1], train, str.tostring(reg_type), clust, tol) 
[l_r, l_er, l_cor, l_r2] = knn_regression(optolo, optolo[1], train, str.tostring(reg_type), clust, tol) 

[h_tf_r, h_tf_er, h_tf_cor, h_tf_r2] = knn_regression(hitoop_tf, hitoop_tf[1], train, str.tostring(reg_type), clust, tol) 
[l_tf_r, l_tf_er, l_tf_cor, l_tf_r2] = knn_regression(optolo_tf, optolo_tf[1], train, str.tostring(reg_type), clust, tol) 

// Colours 
color black = color.rgb(0, 0, 0) 
color white = color.white 
color gray = color.gray 
color greenfill = color.new(color.lime, 90) 
color redfill = color.new(color.red, 90) 

// Plots 
a = plot(prim_tf ? ta.sma(open + h_r, 14) : na, color=greenfill, linewidth=2) 
b = plot(prim_tf ? ta.sma(open + h_r + h_er, 14) : na, color=greenfill, linewidth=2) 
c = plot(prim_tf ? ta.sma(open - l_r, 14) : na, color=redfill, linewidth=2) 
d = plot(prim_tf ? ta.sma(open - l_r - l_er, 14) : na, color=redfill, linewidth=2) 

fill(a, b, color = fill ? greenfill : na)
fill(c, d, color = fill ? redfill : na) 

// Secondary TF Plots 
var line hitp_1 = na 
var line hitp_2 = na 
var line lotp_1 = na 
var line lotp_2 = na 


if barstate.isconfirmed and alt_tf 
    line.delete(hitp_1), line.delete(hitp_2), line.delete(lotp_1), line.delete(lotp_2)
    hitp_1 := line.new(bar_index[10], op_tf + h_tf_r, bar_index, op_tf + h_tf_r, extend = extend.right, color = color.lime, width=3)
    hitp_2 := line.new(bar_index[10], op_tf + h_tf_r + h_tf_er, bar_index, op_tf + h_tf_r + h_tf_er, extend = extend.right, color = color.lime, width=3)
    lotp_1 := line.new(bar_index[10], op_tf - l_tf_r, bar_index, op_tf - l_tf_r, extend = extend.right, color = color.red, width=3)
    lotp_2 := line.new(bar_index[10], op_tf - l_tf_r - l_tf_er, bar_index, op_tf - l_tf_r - l_tf_er, extend = extend.right, color = color.red, width=3)


// Stats 

if stats
    var table stats_tbl = table.new(position.middle_right, 3, 6, bgcolor = black, frame_color = color.aqua, frame_width = 3) 
    table.cell(stats_tbl, 1, 1, text = "Statistics Table", bgcolor = black, text_color = white)    
    table.cell(stats_tbl, 1, 2, text = "Primary Timeframe Data", bgcolor = gray, text_color = white)  
    table.cell(stats_tbl, 1, 3, text = "High Correlation: " + str.tostring(math.round(h_cor,2)) + "\n High R2: " + str.tostring(math.round(h_r2,2)) + "\n Low Correlation: " + str.tostring(math.round(l_cor,2)) + "\n Low R2: " + str.tostring(math.round(l_r2,2)), bgcolor = black, text_color = white)
    table.cell(stats_tbl, 1, 4, text = "Secondary Timeframe Data", bgcolor = gray, text_color = white)   
    table.cell(stats_tbl, 1, 5, text = "High Correlation: " + str.tostring(math.round(h_tf_cor,2)) + "\n High R2: " + str.tostring(math.round(h_tf_r2,2)) + "\n Low Correlation: " + str.tostring(math.round(l_tf_cor,2)) + "\n Low R2: " + str.tostring(math.round(l_tf_r2,2)), bgcolor = black, text_color = white)
