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

//@version=5

// @description TODO:
library(title='FunctionPolynomialRegression', overlay=true)

// @function Method to return a polynomial regression channel using (X,Y) sample points.
// @param sample_x float array, sample data X points.
// @param sample_y float array, sample data Y points.
// @returns tuple with:
//      _predictions:   Array with adjusted Y values.
//      _max_dev:       Max deviation from the mean.
//      _min_dev:       Min deviation from the mean.
//      _stdev/_sizeX:  Average deviation from the mean.
export polyreg(float[] sample_x, float[] sample_y) => //{
    int _sizeY = array.size(id=sample_y)
    int _sizeX = array.size(id=sample_x)
    //
    switch
        (_sizeX < 1)        => runtime.error('FunctionPolynomialRegression -> polyreg(): "sample_x" has the wrong size.')
        (_sizeX != _sizeY)  => runtime.error('FunctionPolynomialRegression -> polyreg(): "sample_x" and "sample_y" size does not match.')
    //
    float _meanX = array.sum(id=sample_x) / _sizeX
    float _meanY = array.sum(id=sample_y) / _sizeX
    float _meanXY = 0.0
    float _meanY2 = 0.0
    float _meanX2 = 0.0
    float _meanX3 = 0.0
    float _meanX4 = 0.0
    float _meanX2Y = 0.0
    //
    for _i = 0 to _sizeY - 1 by 1
        float _Xi = array.get(id=sample_x, index=_i)
        float _Yi = array.get(id=sample_y, index=_i)
        _meanXY := _meanXY + _Xi * _Yi
        _meanY2 := _meanY2 + math.pow(_Yi, 2)
        _meanX2 := _meanX2 + math.pow(_Xi, 2)
        _meanX3 := _meanX3 + math.pow(_Xi, 3)
        _meanX4 := _meanX4 + math.pow(_Xi, 4)
        _meanX2Y := _meanX2Y + math.pow(_Xi, 2) * _Yi
        _meanX2Y
    _meanXY := _meanXY / _sizeX
    _meanY2 := _meanY2 / _sizeX
    _meanX2 := _meanX2 / _sizeX
    _meanX3 := _meanX3 / _sizeX
    _meanX4 := _meanX4 / _sizeX
    _meanX2Y := _meanX2Y / _sizeX
    _meanX2Y
    //-----------|covs
    float _sXX = _meanX2 - _meanX * _meanX
    float _sXY = _meanXY - _meanX * _meanY
    float _sXX2 = _meanX3 - _meanX * _meanX2
    float _sX2X2 = _meanX4 - _meanX2 * _meanX2
    float _sX2Y = _meanX2Y - _meanX2 * _meanY
    //-----------|
    float _b = (_sXY * _sX2X2 - _sX2Y * _sXX2) / (_sXX * _sX2X2 - _sXX2 * _sXX2)
    float _c = (_sX2Y * _sXX - _sXY * _sXX2) / (_sXX * _sX2X2 - _sXX2 * _sXX2)
    float _a = _meanY - _b * _meanX - _c * _meanX2
    //-----------|
    float[] _predictions = array.new_float(size=0, initial_value=0.0)
    float _max_dev = 0.0
    float _min_dev = 0.0
    float _stdev = 0.0
    for _i = 0 to _sizeX - 1 by 1
        float _Xi = array.get(id=sample_x, index=_i)
        float _vector = _a + _b * _Xi + _c * _Xi * _Xi
        array.push(id=_predictions, value=_vector)
        //
        float _Yi = array.get(id=sample_y, index=_i)
        float _diff = _Yi - _vector
        if _diff > _max_dev
            _max_dev := _diff
            _max_dev
        if _diff < _min_dev
            _min_dev := _diff
            _min_dev
        _stdev := _stdev + math.abs(_diff)
        _stdev
    [_predictions, _max_dev, _min_dev, _stdev / _sizeX]
//{ usage:
//{ remarks:
//      language D: https://rosettacode.org/wiki/Polynomial_regression
//}}}

// @function Method for drawing the Polynomial Regression into chart.
// @param sample_x float array, sample point X value.
// @param sample_y float array, sample point Y value.
// @param extend string, default=extend.none, extend lines.
// @param mid_color color, default=color.blue, middle line color.
// @param mid_style string, default=line.style_solid, middle line style.
// @param mid_width int, default=2, middle line width.
// @param std_color color, default=color.aqua, standard deviation line color.
// @param std_style string, default=line.style_dashed, standard deviation line style.
// @param std_width int, default=1, standard deviation line width.
// @param max_color color, default=color.purple, max range line color.
// @param max_style string, default=line.style_dotted, max line style.
// @param max_width int, default=1, max line width.
// @returns line array.
export draw (
     float[] sample_x, float[] sample_y, int n_segments=10, string extend=extend.none, 
     color mid_color=color.blue, string mid_style=line.style_solid, int mid_width=2,
     color std_color=color.aqua, string std_style=line.style_dashed, int std_width=1,
     color max_color=color.purple, string max_style=line.style_dotted, int max_width=1
     ) => //{
    //
    [P, Pmax, Pmin, Pstdev] = polyreg(sample_x, sample_y)
    string _xloc = xloc.bar_index
    //
    var line[] _lines_mid = array.new_line(n_segments)
    var line[] _lines_max = array.new_line(n_segments)  
    var line[] _lines_min = array.new_line(n_segments)
    var line[] _lines_upper = array.new_line(n_segments)
    var line[] _lines_lower = array.new_line(n_segments)
    //
    for _i=0 to (n_segments - 1)
        string _extend = extend.none
        if _i == 0
            switch (extend)
                (extend.left) => _extend := extend.left
                (extend.both) => _extend := extend.left
                => _extend := extend.none
        if _i == (n_segments - 1)
            switch (extend)
                (extend.right) => _extend := extend.right
                (extend.both) => _extend := extend.right
                => _extend := extend.none
        //
        array.set(_lines_mid    , _i, line.new(bar_index, open, bar_index, open, _xloc, _extend, mid_color, mid_style, mid_width))
        array.set(_lines_max    , _i, line.new(bar_index, open, bar_index, open, _xloc, _extend, max_color, max_style, max_width))
        array.set(_lines_min    , _i, line.new(bar_index, open, bar_index, open, _xloc, _extend, max_color, max_style, max_width))
        array.set(_lines_upper  , _i, line.new(bar_index, open, bar_index, open, _xloc, _extend, std_color, std_style, std_width))
        array.set(_lines_lower  , _i, line.new(bar_index, open, bar_index, open, _xloc, _extend, std_color, std_style, std_width))
    //
    int pr_size = array.size(id=P)
    int pr_step = math.max(pr_size / n_segments, 1)
    for _i = 0 to pr_size - pr_step - 1 by pr_step
        int _next_step_index = _i + pr_step
        int _line = _i / pr_step
        //
        line.set_xy1(id=array.get(_lines_mid  , _line), x=int(array.get(id=sample_x, index=_i)), y=array.get(id=P, index=_i))
        line.set_xy2(id=array.get(_lines_mid  , _line), x=int(array.get(id=sample_x, index=_i + pr_step)), y=array.get(id=P, index=_i + pr_step))
        line.set_xy1(id=array.get(_lines_max  , _line), x=int(array.get(id=sample_x, index=_i)), y=array.get(id=P, index=_i) + Pmax)
        line.set_xy2(id=array.get(_lines_max  , _line), x=int(array.get(id=sample_x, index=_i + pr_step)), y=array.get(id=P, index=_i + pr_step) + Pmax)
        line.set_xy1(id=array.get(_lines_min  , _line), x=int(array.get(id=sample_x, index=_i)), y=array.get(id=P, index=_i) + Pmin)
        line.set_xy2(id=array.get(_lines_min  , _line), x=int(array.get(id=sample_x, index=_i + pr_step)), y=array.get(id=P, index=_i + pr_step) + Pmin)
        line.set_xy1(id=array.get(_lines_upper, _line), x=int(array.get(id=sample_x, index=_i)), y=array.get(id=P, index=_i) + Pstdev)
        line.set_xy2(id=array.get(_lines_upper, _line), x=int(array.get(id=sample_x, index=_i + pr_step)), y=array.get(id=P, index=_i + pr_step) + Pstdev)
        line.set_xy1(id=array.get(_lines_lower, _line), x=int(array.get(id=sample_x, index=_i)), y=array.get(id=P, index=_i) - Pstdev)
        line.set_xy2(id=array.get(_lines_lower, _line), x=int(array.get(id=sample_x, index=_i + pr_step)), y=array.get(id=P, index=_i + pr_step) - Pstdev)
    //
    [_lines_mid, _lines_max, _lines_min, _lines_upper, _lines_lower]
//{ usage:
//{ remarks:
//}}}

int length = input(10)
var float[] prices = array.new_float(size=length, initial_value=open)
var int[] indices = array.new_int(size=length, initial_value=0)
int fractal_size = input(2)
if ta.pivothigh(fractal_size, fractal_size)
    e = array.pop(id=prices)
    i = array.pop(id=indices)
    array.insert(id=prices, index=0, value=high[fractal_size])
    array.insert(id=indices, index=0, value=bar_index[fractal_size])
if ta.pivotlow(fractal_size, fractal_size)
    e = array.pop(id=prices)
    i = array.pop(id=indices)
    array.insert(id=prices, index=0, value=low[fractal_size])
    array.insert(id=indices, index=0, value=bar_index[fractal_size])

draw(indices, prices)