Support & Resistance AI (K means/median) [ThinkLogicAI]
// This source code is subject to the terms of the Mozilla Public License 2.0 at https://mozilla.org/MPL/2.0/
// © ThinkLogicAI


//@version=5
indicator("K clustering", overlay = true, max_lines_count = 10, max_labels_count = 251)

//inputs 
k = input.int(3, "Number of clusters", minval = 3, maxval = 5, step = 1)
cent_meth = input.string("K means", title="Cluster Method", options=["K means", "K median"]) 
train_n = input.int(252,"Bars back from last bar to train on", minval = 20, maxval = 1000, step = 1)
plot_band = input.bool(true, "Show SD bands")

//variables needed
x = close
max_bars_back(x, 5000)
var clust = array.new<float>(k, 0)
var sd_clust = array.new<float>(k, 0)
var n_clust = array.new<float>(k,0)
var conv = array.new<int>(0)
var bar_lab = array.new<label>(0)
var table tab = table.new(position.top_right,2, k+1)
var line ts = na
var line k1 = na
var line k2 = na
var line k3 = na
var line k4 = na
var line k5 = na
var line k6 = na
var line sh_1 = na
var line sl_1 = na
var line sh_2 = na
var line sl_2 = na
var line sh_3 = na
var line sl_3 = na
var line sh_4 = na
var line sl_4 = na
var line sh_5 = na
var line sl_5 = na
var line sh_6 = na
var line sl_6 = na



//fuctions used
euc_dist(p,q) => 
    //distance metric
    //p is the close vector and q is the cluster cecnter
    math.sqrt(math.pow(p-q,2))

subtract_array(a,b)=>
    //subtract arrays
    out = array.new<float>(array.size(a),na)
    for i = 0 to array.size(a) -1
        array.set(out,i, array.get(a,i) - array.get(b,i))
    array.sum(out)

center_method(arry)=>
    //calculate either the mean or median based on user input
    float out = 0.00
    if cent_meth == "K means"//k means clustering
        out := array.avg(arry)
    else
        out := array.median((arry))//k meedian clustering
    out

sd_method(arry, mu, ddof)=>
    //my will either be the mean or the median based on user input
    float sum_dif = 0.00
    if cent_meth == "k means"
        sum_dif := array.stdev(arry)
    else
        if arry.size() > 0
            for i = 0 to arry.size() - 1
                sum_dif += math.pow(array.get(arry, i) - mu,2)
            sum_dif := math.sqrt(sum_dif/(array.size(arry)-ddof))
    sum_dif

    

vector_params(k, a1, a2, a3, a4, a5, a6)=>
    //calculate mean, sd, and array size
    mu = array.new<float>(0)
    sd = array.new<float>(0)
    N = array.new<float>(0)
    mean_ar1 = center_method(a1)
    mean_ar2 = center_method(a2)
    mean_ar3 = center_method(a3)
    mean_ar4 = center_method(a4)
    mean_ar5 = center_method(a5)
    mean_ar6 = center_method(a6)
    sd_ar1 = sd_method(a1, mean_ar1, 0)
    sd_ar2 = sd_method(a2, mean_ar2, 0)
    sd_ar3 = sd_method(a3, mean_ar3, 0)
    sd_ar4 = sd_method(a4, mean_ar4, 0)
    sd_ar5 = sd_method(a5, mean_ar5, 0)
    sd_ar6 = sd_method(a6, mean_ar6, 0)
    n1 = array.size(a1)
    n2 = array.size(a2)
    n3 = array.size(a3)
    n4 = array.size(a4)
    n5 = array.size(a5)
    n6 = array.size(a6)
    array.push(mu, mean_ar1)
    array.push(mu, mean_ar2)
    array.push(mu, mean_ar3)
    array.push(mu, mean_ar4)
    array.push(mu, mean_ar5)
    array.push(mu, mean_ar6)
    array.push(sd, sd_ar1)
    array.push(sd, sd_ar2)
    array.push(sd, sd_ar3)
    array.push(sd, sd_ar4)
    array.push(sd, sd_ar5)
    array.push(sd, sd_ar6)
    array.push(N, n1)
    array.push(N, n2)
    array.push(N, n3)
    array.push(N, n4)
    array.push(N, n5)
    array.push(N, n6)
    [array.slice(mu,0,k), array.slice(sd, 0, k), array.slice(N, 0, k)]


// k means/median clustering
if barstate.islast
    err = 100.00
    for i = 0 to array.size(clust) - 1
        array.set(clust, i, i == 0 ? x[math.floor(train_n*.02)] : x[math.floor(train_n*((i/k)))])
    
    //run while clusters keep changing 
    while err > .01
        c1 = array.new<float>(0)
        c2 = array.new<float>(0)
        c3 = array.new<float>(0)
        c4 = array.new<float>(0)
        c5 = array.new<float>(0)
        c6 = array.new<float>(0)
        hold = array.new<int>(0)
        for i = 1 to train_n
            int idx = na
            dist = 1000000.00
            for j = 0 to k-1
                d_temp = euc_dist(x[i], array.get(clust,j))
                if d_temp < dist
                    dist := d_temp
                    idx := j
            if idx == 0
                array.push(c1, x[i])
                array.push(hold, 0)
            else if idx == 1
                array.push(c2, x[i])
                array.push(hold, 1)
            else if idx == 2
                array.push(c3, x[i])
                array.push(hold, 2)
            else if idx == 3
                array.push(c4, x[i])
                array.push(hold, 3)
            else if idx == 4
                array.push(c5, x[i])
                array.push(hold, 4)
            else if idx == 5
                array.push(c6, x[i])
                array.push(hold, 5)
        //after looping though data
        [mu, sd, n] = vector_params(k, c1, c2, c3, c4, c5, c6)
        err := math.pow(subtract_array(clust, mu),2)
        clust := array.copy(mu)
        sd_clust := array.copy(sd)
        array.push(conv,1)
        conv := array.copy(hold)
        n_clust := array.copy(n)
        
    //populate table
    table.cell(tab, 0, 0, "Cluster",  bgcolor = color.new(color.gray, 50), text_color = color.white)
    table.cell(tab, 1, 0, "Density", bgcolor = color.new(color.gray, 50), text_color = color.white)

    //visuals, table update, and stats all based on user defined number of clusters
    if k == 3
        n1 = n_clust.get(0)
        n2 = n_clust.get(1)      
        n3 = n_clust.get(2)
        denom = n1 + n2 + n3
        table.cell(tab, 1, 1, str.tostring(100*math.round(n1 / denom,3)) + "%", bgcolor = color.new(color.green,50), text_color = color.white)
        table.cell(tab, 1, 2, str.tostring(100*math.round(n2 / denom,3)) + "%", bgcolor = color.new(color.red, 50), text_color = color.white)
        table.cell(tab, 1, 3, str.tostring(100*math.round(n3 / denom,3)) + "%", bgcolor = color.new(color.blue,50), text_color = color.white)
        table.cell(tab, 0, 1, "1", bgcolor = color.gray, text_color = color.white)
        table.cell(tab, 0, 2, "2", bgcolor = color.gray, text_color = color.white)
        table.cell(tab, 0, 3, "3", bgcolor = color.gray, text_color = color.white)
        if n1 > 0
            k1 := line.new(bar_index - 1, array.get(clust,0), bar_index, array.get(clust,0),xloc = xloc.bar_time,extend = extend.both, color = color.green, width = 2)
            if plot_band
                sh_1 := line.new(bar_index - 1 , array.get(clust,0) + array.get(sd_clust, 0), bar_index, array.get(clust,0) + array.get(sd_clust, 0),xloc = xloc.bar_time,extend = extend.both, color = color.white, width = 1)
                sl_1 := line.new(bar_index - 1, array.get(clust,0) - array.get(sd_clust, 0), bar_index, array.get(clust,0) - array.get(sd_clust, 0),xloc = xloc.bar_time,extend = extend.both, color = color.white, width = 1)    
                linefill.new(sh_1, sl_1, color.new(color.green, 90))
        if n2 > 0
            k2 := line.new(bar_index - 1, array.get(clust,1), bar_index, array.get(clust,1),xloc = xloc.bar_time,extend = extend.both, color = color.red, width = 2)
            if plot_band
                sh_2 := line.new(bar_index - 1, array.get(clust,1) + array.get(sd_clust, 1), bar_index, array.get(clust,1) + array.get(sd_clust, 1),xloc = xloc.bar_time,extend = extend.both, color = color.white, width = 1)
                sl_2 := line.new(bar_index - 1, array.get(clust,1) - array.get(sd_clust, 1), bar_index, array.get(clust,1) - array.get(sd_clust, 1),xloc = xloc.bar_time,extend = extend.both, color = color.white, width = 1)
                linefill.new(sh_2, sl_2, color.new(color.red, 90))
        if n3 > 0    
            k3 := line.new(bar_index - 1, array.get(clust,2), bar_index, array.get(clust,2),xloc = xloc.bar_time,extend = extend.both, color = color.blue, width = 2)
            if plot_band
                sh_3 := line.new(bar_index - 1, array.get(clust,2) + array.get(sd_clust, 2), bar_index, array.get(clust,2) + array.get(sd_clust, 2),xloc = xloc.bar_time,extend = extend.both, color = color.white, width = 1)
                sl_3 := line.new(bar_index - 1, array.get(clust,2) - array.get(sd_clust, 2), bar_index, array.get(clust,2) - array.get(sd_clust, 2),xloc = xloc.bar_time,extend = extend.both, color = color.white, width = 1)
                linefill.new(sh_3, sl_3, color.new(color.blue, 90))
    else if k == 4
        n1 = n_clust.get(0)
        n2 = n_clust.get(1)      
        n3 = n_clust.get(2)
        n4 = n_clust.get(3)
        denom = n1 + n2 + n3 + n4
        table.cell(tab, 1, 1, str.tostring(100*math.round(n1 / denom,3)) + "%", bgcolor = color.new(color.green,50), text_color = color.white)
        table.cell(tab, 1, 2, str.tostring(100*math.round(n2 / denom,3)) + "%", bgcolor = color.new(color.red,50), text_color = color.white)
        table.cell(tab, 1, 3, str.tostring(100*math.round(n3 / denom,3)) + "%", bgcolor = color.new(color.blue,50), text_color = color.white)
        table.cell(tab, 1, 4, str.tostring(100*math.round(n4 / denom,3)) + "%", bgcolor = color.new(color.orange,50), text_color = color.white)
        table.cell(tab, 0, 1, "1", bgcolor = color.gray, text_color = color.white)
        table.cell(tab, 0, 2, "2", bgcolor = color.gray, text_color = color.white)
        table.cell(tab, 0, 3, "3", bgcolor = color.gray, text_color = color.white)
        table.cell(tab, 0, 4, "4", bgcolor = color.gray, text_color = color.white)
        if n1 > 0
            k1 := line.new(bar_index - 1, array.get(clust,0), bar_index, array.get(clust,0),extend = extend.both, color = color.green, width = 2)
            if plot_band
                sh_1 := line.new(bar_index - 1, array.get(clust,0) + array.get(sd_clust, 0), bar_index, array.get(clust,0) + array.get(sd_clust, 0),extend = extend.both, color = color.white, width = 1)
                sl_1 := line.new(bar_index - 1, array.get(clust,0) - array.get(sd_clust, 0), bar_index, array.get(clust,0) - array.get(sd_clust, 0),extend = extend.both, color = color.white, width = 1)    
                linefill.new(sh_1, sl_1, color.new(color.green, 90))
        if n2 > 0
            k2 := line.new(bar_index - 1, array.get(clust,1), bar_index, array.get(clust,1),extend = extend.both, color = color.red, width = 2)
            if plot_band
                sh_2 := line.new(bar_index - 1, array.get(clust,1) + array.get(sd_clust, 1), bar_index, array.get(clust,1) + array.get(sd_clust, 1),extend = extend.both, color = color.white, width = 1)
                sl_2 := line.new(bar_index - 1, array.get(clust,1) - array.get(sd_clust, 1), bar_index, array.get(clust,1) - array.get(sd_clust, 1),extend = extend.both, color = color.white, width = 1)
                linefill.new(sh_2, sl_2, color.new(color.red, 90))
        if n3 > 0    
            k3 := line.new(bar_index - 1, array.get(clust,2), bar_index, array.get(clust,2),extend = extend.both, color = color.blue, width = 2)
            if plot_band
                sh_3 := line.new(bar_index - 1, array.get(clust,2) + array.get(sd_clust, 2), bar_index, array.get(clust,2) + array.get(sd_clust, 2),extend = extend.both, color = color.white, width = 1)
                sl_3 := line.new(bar_index - 1, array.get(clust,2) - array.get(sd_clust, 2), bar_index, array.get(clust,2) - array.get(sd_clust, 2),extend = extend.both, color = color.white, width = 1)
                linefill.new(sh_3, sl_3, color.new(color.blue, 90))
        if n4 > 0
            k4 := line.new(bar_index - 1, array.get(clust,3), bar_index, array.get(clust,3),extend = extend.both, color = color.orange, width = 2)
            if plot_band
                sh_4 := line.new(bar_index - 1, array.get(clust,3) + array.get(sd_clust, 3), bar_index, array.get(clust,3) + array.get(sd_clust, 3),extend = extend.both, color = color.white, width = 1)
                sl_4 := line.new(bar_index - 1, array.get(clust,3) - array.get(sd_clust, 3), bar_index, array.get(clust,3) - array.get(sd_clust, 3),extend = extend.both, color = color.white, width = 1)
                linefill.new(sh_4, sl_4, color.new(color.orange, 90))

    else if k == 5
        n1 = n_clust.get(0)
        n2 = n_clust.get(1)      
        n3 = n_clust.get(2)
        n4 = n_clust.get(3)
        n5 = n_clust.get(4)
        denom = n1 + n2 + n3 + n4 + n5
        table.cell(tab, 1, 1, str.tostring(100*math.round(n1 / denom,3)) + "%", bgcolor = color.new(color.green,50), text_color = color.white)
        table.cell(tab, 1, 2, str.tostring(100*math.round(n2 / denom,3)) + "%", bgcolor = color.new(color.red,50), text_color = color.white)
        table.cell(tab, 1, 3, str.tostring(100*math.round(n3 / denom,3)) + "%", bgcolor = color.new(color.blue,50), text_color = color.white)
        table.cell(tab, 1, 4, str.tostring(100*math.round(n4 / denom,3)) + "%", bgcolor = color.new(color.orange,50), text_color = color.white)
        table.cell(tab, 1, 5, str.tostring(100*math.round(n5 / denom,3)) + "%", bgcolor = color.new(color.yellow,50), text_color = color.white)
        table.cell(tab, 0, 1, "1", bgcolor = color.gray, text_color = color.white)
        table.cell(tab, 0, 2, "2", bgcolor = color.gray, text_color = color.white)
        table.cell(tab, 0, 3, "3", bgcolor = color.gray, text_color = color.white)
        table.cell(tab, 0, 4, "4", bgcolor = color.gray, text_color = color.white)
        table.cell(tab, 0, 5, "5", bgcolor = color.gray, text_color = color.white)
        if n1 > 0
            k1 := line.new(bar_index - 1, array.get(clust,0), bar_index, array.get(clust,0),extend = extend.both, color = color.green, width = 2)
            if plot_band
                sh_1 := line.new(bar_index - 1, array.get(clust,0) + array.get(sd_clust, 0), bar_index, array.get(clust,0) + array.get(sd_clust, 0),extend = extend.both, color = color.white, width = 1)
                sl_1 := line.new(bar_index - 1, array.get(clust,0) - array.get(sd_clust, 0), bar_index, array.get(clust,0) - array.get(sd_clust, 0),extend = extend.both, color = color.white, width = 1)    
                linefill.new(sh_1, sl_1, color.new(color.green, 90))
        if n2 > 0
            k2 := line.new(bar_index - 1, array.get(clust,1), bar_index, array.get(clust,1),extend = extend.both, color = color.red, width = 2)
            if plot_band
                sh_2 := line.new(bar_index - 1, array.get(clust,1) + array.get(sd_clust, 1), bar_index, array.get(clust,1) + array.get(sd_clust, 1),extend = extend.both, color = color.white, width = 1)
                sl_2 := line.new(bar_index - 1, array.get(clust,1) - array.get(sd_clust, 1), bar_index, array.get(clust,1) - array.get(sd_clust, 1),extend = extend.both, color = color.white, width = 1)
                linefill.new(sh_2, sl_2, color.new(color.red, 90))
        if n3 > 0    
            k3 := line.new(bar_index - 1, array.get(clust,2), bar_index, array.get(clust,2),extend = extend.both, color = color.blue, width = 2)
            if plot_band
                sh_3 := line.new(bar_index - 1, array.get(clust,2) + array.get(sd_clust, 2), bar_index, array.get(clust,2) + array.get(sd_clust, 2),extend = extend.both, color = color.white, width = 1)
                sl_3 := line.new(bar_index - 1, array.get(clust,2) - array.get(sd_clust, 2), bar_index, array.get(clust,2) - array.get(sd_clust, 2),extend = extend.both, color = color.white, width = 1)
                linefill.new(sh_3, sl_3, color.new(color.blue, 90))
        if n4 > 0
            k4 := line.new(bar_index - 1, array.get(clust,3), bar_index, array.get(clust,3),extend = extend.both, color = color.orange, width = 2)
            if plot_band
                sh_4 := line.new(bar_index - 1, array.get(clust,3) + array.get(sd_clust, 3), bar_index, array.get(clust,3) + array.get(sd_clust, 3),extend = extend.both, color = color.white, width = 1)
                sl_4 := line.new(bar_index - 1, array.get(clust,3) - array.get(sd_clust, 3), bar_index, array.get(clust,3) - array.get(sd_clust, 3),extend = extend.both, color = color.white, width = 1)
                linefill.new(sh_4, sl_4, color.new(color.orange, 90))
            
        if n5 > 0
            k5 := line.new(bar_index - 1, array.get(clust,4), bar_index, array.get(clust,4),extend = extend.both, color = color.yellow, width = 2)
            if plot_band    
                sh_5 := line.new(bar_index - 1, array.get(clust,4) + array.get(sd_clust, 4), bar_index, array.get(clust,4) + array.get(sd_clust, 4),extend = extend.both, color = color.white, width = 1)
                sl_5 := line.new(bar_index - 1, array.get(clust,4) - array.get(sd_clust, 4), bar_index, array.get(clust,4) - array.get(sd_clust, 4),extend = extend.both, color = color.white, width = 1)    
                linefill.new(sh_5, sl_5, color.new(color.yellow, 90))




