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

//@version=5
library("LinearRegressionLibrary")

export RepeatedMedian(float y, int n, int lastBar) =>
    //
    // Repeated median regression (robust regression) algorithm:
    // For each point, get slopes of lines connecting it with other data points, 
    // and calculate the median. Then calculate the slope of regression line as 
    // the median of the median slopes. 
    // Finally, calculate the intercept of the regression line.
    //
    // Input:
    // y       :: float, series,  source time series, e.g. close
    // n       :: integet,        number of points in the lookback window
    // lastBar :: integer,        index of the last bar in the selected interval
    // Output:
    // mSlope  :: float,          slope of the regression line
    // mInter  :: float,          intercept of the regression line
    // 
    var float[] S  = array.new_float(0)
    var float[] I  = array.new_float(0)
    float   m      = 0.0
    float   mSlope = 0.0
    float   mInter = 0.0
    bool    isDone = false
    //
    if bar_index[1] <= lastBar and bar_index >= lastBar and n > 1 and not isDone
        isDone := true
        for i = 0 to n - 1
            float[] S1 = array.new_float(0)
            for j = 0 to n - 1
                if i != j
                    array.push(S1, (y[i] - y[j]) / (j - i))
            m  := array.median(S1)
            array.push(S, m)
            // Control the array size 
            if array.size(S) > n
                array.remove(S,0)
        //
        // Calculate intercept
        float   mSlopeLocal = array.median(S)
        for i = 0 to n - 1
            array.push(I, y[n - 1 - i] - mSlopeLocal * i)
            // Control the array size 
            if array.size(I) > n
                array.remove(I,0)
    mSlope := array.median(S)
    mInter := array.median(I)
    //
    [mSlope, mInter]

export TheilSen(float y, int n, int lastBar) =>
    //
    // Theil-Sen estimator (robust regression) algorithm:
    // Get slopes of all lines connecting the data points in the sample. 
    // Then calculate the slope of regression line as the median of the slopes. 
    // Finally, calculate the intercept of the regression line.
    //
    // Input:
    // y       :: float, series,  source time series, e.g. close
    // n       :: integet,        number of points in the lookback window
    // lastBar :: integer,        index of the last bar in the selected interval
    // Output:
    // mSlope  :: float,          slope of the regression line
    // mInter  :: float,          intercept of the regression line
    // 
    var float[] S   = array.new_float(0)
    var float[] I   = array.new_float(0)
    float   tsSlope = 0.0
    float   tsInter = 0.0
    bool    isDone  = false
    //
    if bar_index[1] <= lastBar and bar_index >= lastBar and n > 1 and not isDone
        isDone := true
        for i = 0 to n - 2
            for j = i + 1 to n - 1
                array.push(S, (y[i] - y[j]) / (j - i))
                // Control the array size 
                if array.size(S) > n*(n-1)/2
                    array.remove(S,0)
        //
        // Calculate intercept
        float   mSlopeLocal = array.median(S)
        for i = 0 to n - 1
            array.push(I, y[n - 1 - i] - mSlopeLocal * i)
            // Control the array size 
            if array.size(I) > n
                array.remove(I,0)
    tsSlope := array.median(S)
    tsInter := array.median(I)
    //
    [tsSlope, tsInter]

export OrdinaryLeastSquares(float y, int n, int lastBar) =>
    //
    // Ordinary least square regression (non-robust) algorithm.
    //
    // Input:
    // y        :: float, series,  source time series, e.g. close
    // n        :: integet,        number of points in the lookback window
    // lastBar  :: integer,        index of the last bar in the selected interval
    // Output:
    // olsSlope :: float,          slope of the regression line
    // olsInter :: float,          intercept of the regression line
    // 
    var float   olsSlope = 0.0
    var float   olsInter = 0.0
    bool        isDone   = false
    //
    if bar_index[1] <= lastBar and bar_index >= lastBar and n > 1 and not isDone
        isDone     := true
        float meanX = (n+1)/2
        float meanY = 0.0
        for i  = 0 to n - 1
            meanY  += y[n - 1 - i]/n
        float sum1  = 0.0
        float sum2  = 0.0
        for i  = 0 to n - 1
            sum1   += (i - meanX)*(y[n - 1 - i] - meanY)
            sum2   += (i - meanX)*(i - meanX)
        olsSlope   := sum1/sum2
        olsInter   := meanY - olsSlope*meanX
    //
    [olsSlope, olsInter]

export metricRMSE(float y, int n, int lastBar, float slope, float intercept) =>
    //
    // Model performance metric: root-mean-square error (RMSE) of the regression.
    //
    // Input:
    //    y           :: float,series,   input time series, e.g. close
    //    n           :: integer,        number of points in the interval
    //    lastBar     :: integer,        index of the last bar in the interval
    //    slope       :: float,          model slope
    //    intercept   :: float,          model intercept
    // Output:
    //    rmse        :: float,     root mean square error 
    //
    float rmse = na
    if bar_index[1] <= lastBar and bar_index >= lastBar and n > 1 and na(rmse)
        rmse := 0.0
        for i = 0 to n-1
            rmse := rmse + math.pow(y[i] - intercept - slope*(n - i),2)/n
    math.sqrt(rmse)
    
export metricMAE(float y, int n, int lastBar, float slope, float intercept) =>
    //
    // Model performance metric: mean absolute error (MAE) of the regression.
    //
    // Input:
    //    y           :: float,series,   input time series, e.g. close
    //    n           :: integer,        number of points in the interval
    //    lastBar     :: integer,        index of the last bar in the interval
    //    slope       :: float,          model slope
    //    intercept   :: float,          model intercept
    // Output:
    //    mae         :: float,          mean absolute error 
    //
    var float mae = na
    if bar_index[1] <= lastBar and bar_index >= lastBar and n > 1 and na(mae)
        mae := 0.0
        for i = 0 to n-1
            mae += math.abs(y[i] - intercept - slope*(n - i))/n
    mae

export metricR2(float y, int n, int lastBar, float slope, float intercept) =>
    //
    // Model performance metric: coefficient of determination (R-square) of the regression.
    //
    // Input:
    //    y           :: float,series,   input time series, e.g. close
    //    n           :: integer,        number of points in the interval
    //    lastBar     :: integer,        index of the last bar in the interval
    //    slope       :: float,          model slope
    //    intercept   :: float,          model intercept
    // Output:
    //    Rsq         :: float,          R-square score 
    //
    var float   Rsq = 0.0
    bool        isDone   = false
    //
    if bar_index[1] <= lastBar and bar_index >= lastBar and n > 1 and not isDone
        isDone     := true
        float meanY = 0.0
        for i  = 0 to n - 1
            meanY  += y[i]/n
        float u  = 0.0
        float v  = 0.0
        for i  = 0 to n - 1
            float hatY = intercept + slope*(n - i)
            u += (y[i] - hatY)*(y[i] - hatY)
            v += (y[i] - meanY)*(y[i] - meanY)
        Rsq   := 1 - u/v
    //
    Rsq
