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

//@version=4
strategy(title="VWMA with kNN Machine Learning: MFI/ADX by lastguru", shorttitle="VWMA + kNN: MFI/ADX", overlay=true, default_qty_type=strategy.percent_of_equity, default_qty_value=100)

/////////
// kNN //
/////////

// Define storage arrays for: parameter 1, parameter 2, price, result (up = 1; down = -1)
var knn1 = array.new_float(1, 0)
var knn2 = array.new_float(1, 0)
var knnp = array.new_float(1, 0)
var knnr = array.new_float(1, 0)

// Store the previous trade; buffer the current one until results are in
_knnStore (p1, p2, src) =>
    var prevp1 = 0.0
    var prevp2 = 0.0
    var prevsrc = 0.0
    
    array.push(knn1, prevp1)
    array.push(knn2, prevp2)
    array.push(knnp, prevsrc)
    array.push(knnr, src >= prevsrc ? 1 : -1)
    
    prevp1 := p1
    prevp2 := p2
    prevsrc := src

// Get neighbours by getting k smallest distances from the distance array, and then getting all results with these distances
_knnGet(arr1, arr2, k) =>
    sarr = array.copy(arr1)
    array.sort(sarr)
    ss = array.slice(sarr, 0, min(k, array.size(sarr)))
    m = array.max(ss)
    out = array.new_float(0)
    for i = 0 to array.size(arr1) - 1
        if (array.get(arr1, i) <= m)
            array.push(out, array.get(arr2, i))
    out

// Create a distance array from the two given parameters
_knnDistance(p1, p2) =>
    dist = array.new_float(0)
    n = array.size(knn1) - 1
    for i = 0 to n
        d = sqrt( pow(p1 - array.get(knn1, i), 2) + pow(p2 - array.get(knn2, i), 2) )
        array.push(dist, d)
    dist

// Make a prediction, finding k nearest neighbours
_knn(p1, p2, k) =>
    slice = _knnGet(_knnDistance(p1, p2), array.copy(knnr), k)
    knn = array.sum(slice)

////////////
// Inputs //
////////////

SRC = input(title="Source", type=input.source, defval=open)
FAST = input(title="Fast Length", type=input.integer, defval=13)
SLOW = input(title="Slow Length", type=input.integer, defval=19)
APPLY = input(title = "Apply kNN filter", defval=true)
FILTER = input(title="Filter Length", type=input.integer, defval=13)
SMOOTH = input(title="Filter Smoothing", type=input.integer, defval=6)
// When DIST is 0, KNN default was 23
KNN = input(title="kNN nearest neighbors (k)", type=input.integer, defval=45)
DIST = input(title="kNN minimum difference", type=input.integer, defval=2)
BACKGROUND = input(title = "Draw background", defval=false)

////////
// MA //
////////
fastMA = vwma(SRC, FAST)
slowMA = vwma(SRC, SLOW)

/////////
// DMI //
/////////

// Wilder's Smoothing (Running Moving Average)
_rma(src, length) =>
    out = 0.0
    out := ((length - 1) * nz(out[1]) + src) / length

// DMI (Directional Movement Index)
_dmi (len, smooth) =>
    up = change(high)
    down = -change(low)
    plusDM = na(up) ? na : (up > down and up > 0 ? up : 0)
    minusDM = na(down) ? na : (down > up and down > 0 ? down : 0)
    trur = _rma(tr, len)
    plus = fixnan(100 * _rma(plusDM, len) / trur)
    minus = fixnan(100 * _rma(minusDM, len) / trur)
    sum = plus + minus
    adx = 100 * _rma(abs(plus - minus) / (sum == 0 ? 1 : sum), smooth)
    [plus, minus, adx]

[diplus, diminus, adx] = _dmi(FILTER, SMOOTH)

/////////
// MFI //
/////////

// common RSI function
_rsi(upper, lower) =>
    if lower == 0
        100
    if upper == 0
        0
	100.0 - (100.0 / (1.0 + upper / lower))

mfiUp = sum(volume * (change(ohlc4) <= 0 ? 0 : ohlc4), FILTER)
mfiDown = sum(volume * (change(ohlc4) >= 0 ? 0 : ohlc4), FILTER)
mfi = _rsi(mfiUp, mfiDown)

////////////
// Filter //
////////////

longCondition = crossover(fastMA, slowMA)
shortCondition = crossunder(fastMA, slowMA)

if (longCondition or shortCondition)
    _knnStore(adx, mfi, SRC)
filter = _knn(adx, mfi, KNN)

/////////////
// Actions //
/////////////

bgcolor(BACKGROUND ? filter >= 0 ? color.green : color.red : na)
plot(fastMA, color=color.red)
plot(slowMA, color=color.green)

if (longCondition and (not APPLY or filter >= DIST))
    strategy.entry("Long", strategy.long)
if (shortCondition and (not APPLY or filter <= -DIST))
    strategy.entry("Short", strategy.short)