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

//   __    __            __       __  __                  __       __                        __      __    __              __      _______             __    __                                          __ 
//  /  |  /  |          /  |  _  /  |/  |                /  \     /  |                      /  |    /  \  /  |            /  |    /       \           /  \  /  |                                        /  |
//  $$ |  $$ |  ______  $$ | / \ $$ |$$ |____    ______  $$  \   /$$ | __    __   _______  _$$ |_   $$  \ $$ |  ______   _$$ |_   $$$$$$$  |  ______  $$  \ $$ |  ______   _____  ____    ______    ____$$ |
//  $$ |__$$ | /      \ $$ |/$  \$$ |$$      \  /      \ $$$  \ /$$$ |/  |  /  | /       |/ $$   |  $$$  \$$ | /      \ / $$   |  $$ |__$$ | /      \ $$$  \$$ | /      \ /     \/    \  /      \  /    $$ |
//  $$    $$ |/$$$$$$  |$$ /$$$  $$ |$$$$$$$  |/$$$$$$  |$$$$  /$$$$ |$$ |  $$ |/$$$$$$$/ $$$$$$/   $$$$  $$ |/$$$$$$  |$$$$$$/   $$    $$< /$$$$$$  |$$$$  $$ | $$$$$$  |$$$$$$ $$$$  |/$$$$$$  |/$$$$$$$ |
//  $$$$$$$$ |$$    $$ |$$ $$/$$ $$ |$$ |  $$ |$$ |  $$ |$$ $$ $$/$$ |$$ |  $$ |$$      \   $$ | __ $$ $$ $$ |$$ |  $$ |  $$ | __ $$$$$$$  |$$    $$ |$$ $$ $$ | /    $$ |$$ | $$ | $$ |$$    $$ |$$ |  $$ |
//  $$ |  $$ |$$$$$$$$/ $$$$/  $$$$ |$$ |  $$ |$$ \__$$ |$$ |$$$/ $$ |$$ \__$$ | $$$$$$  |  $$ |/  |$$ |$$$$ |$$ \__$$ |  $$ |/  |$$ |__$$ |$$$$$$$$/ $$ |$$$$ |/$$$$$$$ |$$ | $$ | $$ |$$$$$$$$/ $$ \__$$ |
//  $$ |  $$ |$$       |$$$/    $$$ |$$ |  $$ |$$    $$/ $$ | $/  $$ |$$    $$/ /     $$/   $$  $$/ $$ | $$$ |$$    $$/   $$  $$/ $$    $$/ $$       |$$ | $$$ |$$    $$ |$$ | $$ | $$ |$$       |$$    $$ |
//  $$/   $$/  $$$$$$$/ $$/      $$/ $$/   $$/  $$$$$$/  $$/      $$/  $$$$$$/  $$$$$$$/     $$$$/  $$/   $$/  $$$$$$/     $$$$/  $$$$$$$/   $$$$$$$/ $$/   $$/  $$$$$$$/ $$/  $$/  $$/  $$$$$$$/  $$$$$$$/ 
//                                                                                                                                                                                                          
//                                                                                                                                                                                                          
//
//@version=5

// @description Implementation of functions related to data correlation calculations. Formulas have been transformed in such a way that we avoid running loops and instead make use of time series to gradually build the data we need to perform calculation. This allows the calculations to run on unbound series, and/or higher number of samples
library("DataCorrelation")
import HeWhoMustNotBeNamed/BinaryInsertionSort/1 as bis

// @function Calculates chatterjee correlation between two series. Formula is - ξnₓᵧ = 1 - (3 * ∑ |rᵢ₊₁ - rᵢ|)/ (n²-1)
// @param x First series for which correlation need to be calculated
// @param y Second series for which correlation need to be calculated
// @param sampleSize number of samples to be considered for calculattion of correlation. Default is 20000
// @param plotSize How many historical values need to be plotted on chart.
// @returns float correlation - Chatterjee correlation value if falls within plotSize, else returns na
export chatterjeeCorrelation(float x, float y, simple int sampleSize = 20000, simple int plotSize = 500) =>
	float correlation = na

	if(bar_index > last_bar_index-sampleSize)
		[xArray, xSortedArray, xSortIndices] = bis.get_sorted_arrays(x)
		[yArray, ySortedArray, ySortIndices] = bis.get_sorted_arrays(y)

		if(bar_index >= last_bar_index-plotSize)
            coefficient = 0.0
            n = array.size(xArray)
			index = array.get(xSortIndices, 0)
			for i=1 to n > 1? n -1 : na
				indexNext = array.get(xSortIndices, i)
				coefficient := coefficient + math.abs(array.get(ySortIndices, indexNext)-array.get(ySortIndices, index))
				index := indexNext

			correlation := 1 - 3*coefficient/(math.pow(n,2)-1)
	correlation

// @function Calculates spearman correlation between two series. Formula is - ρ = 1 - (6∑dᵢ²/n(n²-1))
// @param x First series for which correlation need to be calculated
// @param y Second series for which correlation need to be calculated
// @param sampleSize number of samples to be considered for calculattion of correlation. Default is 20000
// @param plotSize How many historical values need to be plotted on chart.
// @returns float correlation - Spearman correlation value if falls within plotSize, else returns na
export spearmanCorrelation(float x, float y, simple int sampleSize = 20000, simple int plotSize = 500) =>
	float correlation = na

	if(bar_index > last_bar_index-sampleSize)
		[xArray, xSortedArray, xSortIndices] = bis.get_sorted_arrays(x)
		[yArray, ySortedArray, ySortIndices] = bis.get_sorted_arrays(y)
		if(bar_index >= last_bar_index-plotSize)
			n = array.size(xArray)
			dSquareSum = 0.0
            for i=0 to n > 0? n -1 : na
                d = array.get(xSortIndices, i) - array.get(ySortIndices, i)
                dSquareSum += d*d
            correlation := 1 - (6*dSquareSum/(math.pow(n, 3)-n))
    correlation


// @function Calculates covariance between two series of unbound length. Formula is Covₓᵧ = ∑ ((xᵢ-x̄)(yᵢ-ȳ)) / (n-1) for sample and Covₓᵧ = ∑ ((xᵢ-x̄)(yᵢ-ȳ)) / n for population
// @param x First series for which covariance need to be calculated
// @param y Second series for which covariance need to be calculated
// @param include boolean flag used for selectively including sample
// @param biased boolean flag representing population covariance instead of sample covariance
// @returns float covariance - covariance of selective samples of two series x, y
export covariance(float x, float y, bool include=true, simple bool biased = true)=>
    var float xySum = 0.0
    var float xSum = 0.0
    var float ySum = 0.0
    var n=0
    var float covariance = na
    if(include)
        xy = x*y
        n+=1
        xSum += x
        ySum += y
        xySum += xy
        xMean = xSum/n
        yMean = ySum/n

        covariance := (xySum + (n*xMean*yMean) - (yMean*xSum) - (xMean*ySum))/(n-(biased?0:1))
    covariance

// @function Calculates Standard Deviation of a series. Formula is σ = √( ∑(xᵢ-x̄)² / n ) for sample and σ = √( ∑(xᵢ-x̄)² / (n-1) ) for population
// @param x Series for which Standard Deviation need to be calculated
// @param include boolean flag used for selectively including sample
// @param biased boolean flag representing population covariance instead of sample covariance
// @returns float stddev - standard deviation of selective samples of series x
export stddev(float x,  bool include=true, simple bool biased = true)=>
    var float xSquareSum = 0.0
    var float xSum = 0.0
    var n=0
    var float stddev = na
    if(include)
        n+=1
        xSquareSum += x*x
        xSum += x
        xMean = xSum/n
        stddev := math.sqrt( (xSquareSum + n*xMean*xMean - 2*xMean*xSum)/(n-(biased?0:1)) )
    stddev

// @function Calculates pearson correlation between two series of unbound length. Formula is r = Covₓᵧ / σₓσᵧ
// @param x First series for which correlation need to be calculated
// @param y Second series for which correlation need to be calculated
// @param include boolean flag used for selectively including sample
// @returns float correlation - correlation between selective samples of two series x, y
export correlation(float x, float y, bool include=true)=>
    covariance = covariance(x,y, include)
    sdX = stddev(x, include)
    sdY = stddev(y, include)
    covariance/(sdX * sdY)