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

//@version=5

// @description least_squares_regression: Least squares regression algorithm
library("least_squares_regression", overlay=true)

// @function basic_lsr: Basic least squares regression algorithm
// @param series int[] t: time scale value array corresponding to price
// @param series float[] p: price scale value array corresponding to time
// @param series int array_size: the length of regression array
// @returns reg_slop, reg_intercept, reg_level, reg_stdev
export basic_lsr(series int[] t, series float[] p, series int array_size) =>
	p_t = array.new_float(array_size, 0)
	t_2 = array.new_float(array_size, 0)
	
	//init array
	for i = 0 to array_size - 1
	    array.set(p_t, i, array.get(p, i) * array.get(t, i))
	    array.set(t_2, i, array.get(t, i) * array.get(t, i))
	
	//get slop and intercept
	sum_t = array.sum(t)
	sum_p = array.sum(p)
	sum_t_2 = array.sum(t_2)
	sum_p_t = array.sum(p_t)
	a = (array_size * sum_p_t - sum_p * sum_t) / (array_size * sum_t_2 - math.pow(sum_t, 2))
	b = (sum_p - a * sum_t) / array_size
    
    //get regression prediction value
    _p = array.new_float(array_size, 0)
    for i = 0 to array_size - 1
        array.set(_p, i, a * array.get(t, i) + b)
    
    //get regression level 
    float r1 = 0
    float r2 = 0
    p_avg = array.avg(p)
    for i = 0 to array_size - 1
        r1 += math.pow(array.get(_p, i) - p_avg, 2)
        r2 += math.pow(array.get(p, i) - p_avg, 2)
    r = -math.log10(1 - r1 / r2)
    
    //get regression standard devition
    _s = array.new_float(array_size, 0)
    for i = 0 to array_size - 1
        array.set(_s, i, array.get(p, i) - array.get(_p, i))
    s = array.stdev(_s)
	
	//return results
	[a, b, r, s]

// @function top_trend_line_lsr: Trend line fitting based on least square algorithm
// @param series int[] t: time scale value array corresponding to price
// @param series float[] p: price scale value array corresponding to time
// @param series int array_size: the length of regression array
// @param string reg_type: regression type in 'top' and 'bottom'
// @param series int max_iter: maximum fitting iterations
// @param series int min_points: the threshold of regression point numbers
// @returns reg_slop, reg_intercept, reg_level, reg_stdev, reg_point_num
export trend_line_lsr(series int[] t, series float[] p, series int array_size, string reg_type, series int max_iter, series int min_points) =>
	[a, b, r, s] = basic_lsr(t, p, array_size)
	x1 = array.get(t, 0)
	x2 = array.get(t, array_size - 1)
	n = array_size //Function arguments cannot be mutable, what the fuck!

	for i = 0 to max_iter - 1
	    //exit conditions
    	if n < min_points
    	    break
    	   
    	//init new tick and price array
        _t = array.new_int(n, 0)
        _p = array.new_float(n, 0)
        _n = 0
        
        //fine the ground truth values that bigger than predictions
        for j = 0 to n - 1
            if reg_type == 'top'
                if (a * array.get(t, j) + b) < array.get(p, j)
                    array.set(_t, _n, array.get(t, j))
                    array.set(_p, _n, array.get(p, j))
                    _n += 1
            else if reg_type == 'bottom'
                if (a * array.get(t, j) + b) > array.get(p, j)
                    array.set(_t, _n, array.get(t, j))
                    array.set(_p, _n, array.get(p, j))
                    _n += 1
            else
                break
        
        //exit if new array size is less than threshold
        if _n < min_points
            break
        
        //override result if r of new array is bigger than last array
        [_a, _b, _r, _s] = basic_lsr(_t, _p, _n)
        if _r > r
            x1 := array.get(_t, 0)
            x2 := array.get(_t, _n - 1)
            a := _a
            b := _b
            r := _r
            s := _s
            n := _n
    
    //after for loop 
    [x1, x2, a, b, r, s, n]
