kNN
// This source code is subject to the terms of the Mozilla Public License 2.0 at https://mozilla.org/MPL/2.0/
// © lastguru, 2022. All right reserved

//@version=5
// @description Collection of experimental kNN functions
library("kNN", true)

// @function Store the previous trade; buffer the current one until results are in. Results are binary: up/down
// @param knn knn matrix
// @param p1 feature 1 value
// @param p2 feature 2 value
// @param src current price
// @param maxrows limit the matrix size to this number of rows (0 of no limit)
// @returns modified knn matrix
export knnStore(matrix<float> knn, float p1, float p2, float src, int maxrows = 0) =>
    var prevp1 = 0.0
    var prevp2 = 0.0
    var prevsrc = 0.0
    
    if (maxrows > 0 and matrix.rows(knn) > maxrows)
        matrix.remove_row(knn, 0)
    
    // Add to storage matrix: parameter 1, parameter 2, price, result (up = 1; down = -1)
    matrix.add_row(knn, matrix.rows(knn), array.from(prevp1, prevp2, prevsrc, src >= prevsrc ? 1 : -1))

    prevp1 := p1
    prevp2 := p2
    prevsrc := src
    
    knn

// @function Store the previous trade; buffer the current one until results are in. Results are in percents
// @param knn knn matrix
// @param p1 feature 1 value
// @param p2 feature 2 value
// @param src current price
// @param maxrows limit the matrix size to this number of rows (0 of no limit)
// @returns modified knn matrix
export knnStorePercent(matrix<float> knn, float p1, float p2, float src, int maxrows = 0) =>
    var prevp1 = 0.0
    var prevp2 = 0.0
    var prevsrc = 0.0
    
    if (maxrows > 0 and matrix.rows(knn) > maxrows)
        matrix.remove_row(knn, 0)
    
    // Add to storage matrix: parameter 1, parameter 2, price, result percentage
    matrix.add_row(knn, matrix.rows(knn), array.from(prevp1, prevp2, prevsrc, 100 * (src - prevsrc) / prevsrc))

    prevp1 := p1
    prevp2 := p2
    prevsrc := src
    
    knn

// @function Get neighbours by getting k results with the smallest distances
// @param distance distance array
// @param result result array
// @returns array slice of k results
export knnGet(float[] distance, float[] result, int k) =>
    m = matrix.new<float>(0, 0, 0)
    matrix.add_col(m, 0, distance)
    matrix.add_col(m, 1, result)
    matrix.sort(m, 0)
    
    r = matrix.col(m, 1)
    out = array.slice(r, 0, math.min(k, array.size(r)-1))

// @function Create a distance array from the two given parameters
// @param knn knn matrix
// @param p1 feature 1 value
// @param p2 feature 2 value
// @returns distance array
export knnDistance(matrix<float> knn, float p1, float p2) =>
    dist = array.new_float(0)
    n = matrix.rows(knn) - 1
    for i = 0 to n
        d = math.sqrt( math.pow(p1 - matrix.get(knn, i, 0), 2) + math.pow(p2 - matrix.get(knn, i, 1), 2) )
        array.push(dist, d)
    dist

// @function Make a prediction, finding k nearest neighbours and summing them up
// @param knn knn matrix
// @param p1 feature 1 value
// @param p2 feature 2 value
// @param k sum k nearest neighbors
// @returns sum of k nearest neighbors
export knnSum(matrix<float> knn, float p1, float p2, int k) =>
    slice = (matrix.rows(knn) == 0) ? array.from(0.0) : knnGet(knnDistance(knn, p1, p2), matrix.col(knn, 3), k)
    out = array.sum(slice)

// @function execute kNN filter
// @param kNN filter type
// @param k number k
// @param skew kNN minimum difference
// @param depth kNN matrix size limit
// @param price series
// @param long long condition
// @param short short condition
// @param store store the supplied features (if false, only checks the results without storage)
// @param feature1 feature 1 value
// @param feature2 feature 2 value
// @returns [longOK, shortOK] filter output
export doKNN(string type, int k = 5, int skew = 0, int depth, float src, bool long, bool short, bool store, float feature1, float feature2) =>
    var knnM = matrix.new<float>(0, 0, 0)
    longOK = true
    shortOK = true
    
    switch type
        "Boolean" =>
            if (long or short)
                if (store)
                    knnM := knnStore(knnM, feature1, feature2, src, depth)
                
                filter = knnSum(knnM, feature1, feature2, k)
                longOK := filter >= skew
                shortOK := filter <= -skew
        "Percent" =>
            if (long or short)
                if (store)
                    knnM := knnStorePercent(knnM, feature1, feature2, src, depth)
                
                filter = knnSum(knnM, feature1, feature2, k)
                longOK := filter >= skew
                shortOK := filter <= -skew
    
    [longOK, shortOK]