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

//@version=5

// @description The most used array utility functions
library("racille_arrayutils")

// Notes
// A -> set of arrays
// R -> set of real numbers
// 0 -> null set
// a -> argument set

// -------------------------------------------------------------------------------------------------
// { f : 0 -> a
// -------------------------------------------------------------------------------------------------

// @function returns sin function as a parameter to calculate the function_array()
export func_sin() => "sin"

// @function returns cos function as a parameter to calculate the function_array()
export func_cos() => "cos"

// @function returns tan function as a parameter to calculate the function_array()
export func_tan() => "tan"

// @function returns cot function as a parameter to calculate the function_array()
export func_cot() => "cot"

// @function returns asin function as a parameter to calculate the function_array()
export func_asin() => "asin"

// @function returns acos function as a parameter to calculate the function_array()
export func_acos() => "acos"

// @function returns atan function as a parameter to calculate the function_array()
export func_atan() => "atan"

// @function returns acot function as a parameter to calculate the function_array()
export func_acot() => "acot"

// @function returns sqrt function as a parameter to calculate the function_array()
export func_sqrt() => "pow|0.5"

// @function returns pow function as a parameter to calculate the function_array()
// @param x - power
export func_pow(float x) => "pow|" + str.tostring(x)

// -------------------------------------------------------------------------------------------------
// } f : 0 -> a
// -------------------------------------------------------------------------------------------------

// -------------------------------------------------------------------------------------------------
// { f : R -> a
// -------------------------------------------------------------------------------------------------

// @function returns exp function as a parameter to calculate the function_array()
// @param x - base
export func_exp(float x = math.e) => "exp|" + str.tostring(x)

// @function returns log function as a parameter to calculate the function_array()
// @param x - base
export func_log(float x = math.e) => "log|" + str.tostring(x)

// -------------------------------------------------------------------------------------------------
// } f : R -> a
// -------------------------------------------------------------------------------------------------


// -------------------------------------------------------------------------------------------------
// { f : A, a -> A
// -------------------------------------------------------------------------------------------------

// validate that string is numeric
__number_validator(string s) =>
	ret = true
	for i = 0 to str.length(s) - 1 by 1
		if not str.contains("0123456789.", str.substring(s, i, i))
			ret := false
	
	ret

// validate that input string is an appropriate function parameter func_*
__func_validator(string func) =>
	if func == "sin" or func == "cos" or func == "tan" or func == "cot" or func == "asin" or func == "acos" or func == "atan" or func == "acot"
		true
	else 
		if str.contains(func, '|')
			t_key = array.get(str.split(func, '|'), 0)
			t_val = array.get(str.split(func, '|'), 1)
			(t_key == "log" or t_key == "pow" or t_key == "exp") and __number_validator(t_val)
		else
			false

// utility that classificates func as simple (without additional parameter) or complex
__func_clasificator(string func) =>
	func == "sin" or func == "cos" or func == "tan" or func == "cot" or func == "asin" or func == "acos" or func == "atan" or func == "acot" ? "simple" : "complex"

__simple_func_replace(float arg, string func) =>
	switch func
		"sin" => math.sin(arg)
		"cos" => math.cos(arg)
		"tan" => math.tan(arg)
		"cot" => 1/math.tan(arg)
		"asin" => math.asin(arg)
		"acos" => math.acos(arg)
		"atan" => math.atan(arg)
		"acot" => math.pi/2 - math.atan(arg)

__complex_func_replace(float arg, string func, float parameter) =>
	switch func
		"exp" => math.pow(parameter, arg)
		"pow" => math.pow(arg, parameter)
		"log" => math.log(arg) / math.log(parameter)

// utility for func_replace_array
__func_replace_util(float arg, string func) =>
	switch __func_clasificator(func)
		"simple" => __simple_func_replace(arg, func)
		"complex" => __complex_func_replace(arg, array.get(str.split(func, '|'), 0), str.tonumber(array.get(str.split(func, '|'), 1)))


// @function replace each element of array with func(element) and returns a new array
// @param arr  - array
// @param func - function to replace. Must be one of func_*()
// @returns new array, where each element is func(element)
export apply_func(float[] arr, string func) =>
	if not __func_validator(func)
		runtime.error("Invalid function parameter passed as an func_replace_array argument")

	res = array.new_float()
    if array.size(arr)
        for i = 0 to array.size(arr) - 1 by 1
			array.push(res, __func_replace_util(array.get(arr, i), func))

    res

// -------------------------------------------------------------------------------------------------
// { f : A, R -> A
// -------------------------------------------------------------------------------------------------
// @function multiplies each element of array by multiplier and returns a new array
// @param arr - array
// @param x   - multiplier
// @returns new array, where each element is multiplied on x
export multiply_array(float[] arr, float x) =>
	res = array.new_float()
    if array.size(arr)
        for i = 0 to array.size(arr) - 1 by 1
            array.push(res, array.get(arr, i) * x)

    res


// @function multiplies each element of array by multiplier and returns a new array
// @param arr - array
// @param n   - multiplier
// @returns new array, where each element is multiplied on n
export multiply_array(float[] arr, int n) =>
	multiply_array(arr, float(n))


// @function multiplies each element of array by multiplier and returns a new array
// @param arr - array
// @param n   - multiplier
// @returns new array, where each element is multiplied on n
export multiply_array(int[] arr, int n) =>
	res = array.new_int()
    if array.size(arr)
        for i = 0 to array.size(arr) - 1 by 1
            array.push(res, array.get(arr, i) * n)

    res


// @function divides each element of array by divider and returns a new array
// @param arr - array
// @param x   - divider
// @returns new array, where each element is divided by x
export divide_array(float[] arr, float x) =>
	res = array.new_float()
    if array.size(arr) and x != 0
        for i = 0 to array.size(arr) - 1 by 1
            array.push(res, array.get(arr, i) / x)

    res


// @function divides each element of array by divider and returns a new array
// @param arr - array
// @param n   - divider
// @returns new array, where each element is divided by n
export divide_array(float[] arr, int n) =>
	multiply_array(arr, float(n))


// @function adds increment to each element of array and returns a new array
// @param arr - array
// @param x   - increment
// @returns new array, where each element is increased by x
export increase_array(float[] arr, float x) =>
	res = array.new_float()
    if array.size(arr)
        for i = 0 to array.size(arr) - 1 by 1
            array.push(res, array.get(arr, i) + x)

    res


// @function adds increment to each element of array and returns a new array
// @param arr - array
// @param n   - increment
// @returns new array, where each element is increased by n
export increase_array(float[] arr, int n) =>
	increase_array(arr, float(n))


// @function adds increment to each element of array and returns a new array
// @param arr - array
// @param n   - increment
// @returns new array, where each element is multiplied on x
export increase_array(int[] arr, int n) =>
	res = array.new_int()
    if array.size(arr)
        for i = 0 to array.size(arr) - 1
            array.push(res, array.get(arr, i) + n)

    res


// @function substracts decrement to each element of array and returns a new array
// @param arr - array
// @param x   - decrement
// @returns new array, where each element is decreased by x
export decrease_array(float[] arr, float x) =>
	increase_array(arr, -x)


// @function substracts decrement to each element of array and returns a new array
// @param arr - array
// @param n   - decrement
// @returns new array, where each element is decreased by n
export decrease_array(float[] arr, int n) =>
	increase_array(arr, -n)


// @function substracts decrement to each element of array and returns a new array
// @param arr - array
// @param n   - decrement
// @returns new array, where each element is decreased by x
export decrease_array(int[] arr, int n) =>
	increase_array(arr, -n)


// @function changes each elements sign of array and returns a new array
// @param arr - array
// @returns new array, where each element is of different sign
export negate_array(float[] arr) =>
	multiply_array(arr, -1)


// -------------------------------------------------------------------------------------------------
// } f : A, R -> A
// -------------------------------------------------------------------------------------------------


// -------------------------------------------------------------------------------------------------
// { f : A, A -> A
// -------------------------------------------------------------------------------------------------


// @function calculates elementwise sum of two arrays
export array_sum(float[] arr1, float[] arr2) =>
    res = array.new_float()
    if array.size(arr1) and array.size(arr2)
        for i = 0 to math.min(array.size(arr1), array.size(arr2)) - 1
            array.push(res, array.get(arr1, i) + array.get(arr2, i))

    res


// @function calculates elementwise difference of two arrays
export array_diff(float[] arr1, float[] arr2) =>
    array_sum(arr1, negate_array(arr2))


// @function calculates elementwise product of two arrays
export array_product(float[] arr1, float[] arr2) =>
    res = array.new_float()
    if array.size(arr1) and array.size(arr2)
        for i = 0 to math.min(array.size(arr1), array.size(arr2)) - 1
            array.push(res, array.get(arr1, i) * array.get(arr2, i))

    res


// @function calculates elementwise division of two arrays
export array_division(float[] arr1, float[] arr2) =>
    array_product(arr1, apply_func(arr2, func_pow(-1)))


// -------------------------------------------------------------------------------------------------
// } f : A, A -> A
// -------------------------------------------------------------------------------------------------

// -------------------------------------------------------------------------------------------------
// { f : A -> R
// -------------------------------------------------------------------------------------------------

// @function calculates the product of array elements
export array_element_prod(float[] arr) =>
	ret = 0.0
	if array.size(arr)
		ret := 1
        for i = 0 to array.size(arr) - 1
            ret := ret * array.get(arr, i)
	
	ret


__naive_geomean(float[] arr) =>
	if array.size(arr)
		math.pow(array_element_prod(arr), 1/array.size(arr))
	else
		0.0

__log_geomean(float[] arr) =>
	if array.size(arr)
		math.exp(array.avg(apply_func(arr, func_log())))
	else
		0.0

// @function calculates the geometric mean of array elements
// @param arr - array
// @param method - calculation method. Can be chosen from ["auto", "naive", "log"]
// @returns geometric mean of array elements. 0 if array is empty
export array_geomean(float[] arr, string method = "auto") =>
	if method != "auto" and method != "naive" and method != "log"
		runtime.error("Invalid usage of array_geomean(). Unexpected calculation method")
	sz_ = array.size(arr)
	if sz_ < 10 or method == "naive"
		__naive_geomean(arr)
	else
		__log_geomean(arr)

// -------------------------------------------------------------------------------------------------
// } f : A -> R
// -------------------------------------------------------------------------------------------------

// Testing

PASSED = "✅"
FAILED = "❌"

test_result_tostr(string name, bool pass) => name + " " + (pass ? PASSED : FAILED) + "\n"

test_AR_to_A() =>
	string result = "f: A, R -> A suite:\n"
	float[] i_float = array.from(4.0, 5.5, 2.9)

	float[] o_float_plus2p5 = increase_array(i_float, 2.5)
	float[] o_float_minus3 = decrease_array(i_float, 3)
	float[] o_float_mult2 = multiply_array(i_float, 2)
	float[] o_float_div2p2  = divide_array(i_float, 2.2)

	bool validate_plus = true
	bool validate_minus = true
	bool validate_mult = true
	bool validate_div = true

	for i = 0 to array.size(i_float) - 1
		validate_plus  := validate_plus and (array.get(i_float, i) + 2.5 == array.get(o_float_plus2p5, i))
		validate_minus := validate_minus and (array.get(i_float, i) - 3 ==  array.get(o_float_minus3, i))
		validate_mult := validate_mult and (array.get(i_float, i) * 2 ==  array.get(o_float_mult2, i))
		validate_div := validate_div and (array.get(i_float, i) / 2.2 == array.get(o_float_div2p2, i))

	result := result + test_result_tostr("Addition", validate_plus)
	result := result + test_result_tostr("Subtraction", validate_minus)
	result := result + test_result_tostr("Multiplication", validate_mult)
	result := result + test_result_tostr("Division", validate_div)

	result

test_AA_to_A() =>
	string result = "f: A, A -> A suite:\n"
	float[] i1 = array.from(2.5, 3.5, 5.5)
	float[] i2 = array.from(0.5, 1.75, 11.0)

	float[] o_plus = array_sum(i1, i2)
	float[] o_minus = array_diff(i1, i2)
	float[] o_mult = array_product(i1, i2)
	float[] o_div  = array_division(i1, i2)

	bool validate_plus = true
	bool validate_minus = true
	bool validate_mult = true
	bool validate_div = true

	for i = 0 to array.size(i1) - 1
		validate_plus := validate_plus and (array.get(i1, i) + array.get(i2, i) == array.get(o_plus, i))
		validate_minus := validate_minus and (array.get(i1, i) - array.get(i2, i) == array.get(o_minus, i))
		validate_mult := validate_mult and (array.get(i1, i) * array.get(i2, i) == array.get(o_mult, i))
		validate_div := validate_div and (array.get(i1, i) / array.get(i2, i) == array.get(o_div, i))


	result := result + test_result_tostr("Addition", validate_plus)
	result := result + test_result_tostr("Subtraction", validate_minus)
	result := result + test_result_tostr("Multiplication", validate_mult)
	result := result + test_result_tostr("Division", validate_div)

	result

test_A_to_R() =>
	string result = "f: A -> R suite:\n"

	float[] i_pos = array.from(2.5, 4.0, 3.3)
	float[] i_neg = array.from(2.5, -4.0, 3.3)
	float expected_prod = 2.5 * 4.0 * 3.3
	float expected_geomean = math.pow(expected_prod, 1./3)

	result := result + test_result_tostr("Product", expected_prod == array_element_prod(i_pos))
	result := result + test_result_tostr("Naive Geomean, pos x", expected_geomean == array_geomean(i_pos, "naive"))
	result := result + test_result_tostr("Log Geomean, pos x", expected_geomean == array_geomean(i_pos, "log"))
	result := result + test_result_tostr("Naive Geomean, neg x", -expected_geomean == array_geomean(i_neg, "naive"))
	result := result + test_result_tostr("Log Geomean, neg x", -expected_geomean == array_geomean(i_neg, "log"))

	result

test_Aa_to_A_common() =>
	string result = "f: A, a -> A, common suite:\n"

	float[] i_pos = array.from(0.25, 0.81, 1.3)
	float[] i_neg = array.from(0.25, -0.81, 1.3)

	float[] o_sqrt_pos = apply_func(i_pos, func_sqrt())
	float[] o_pow3_neg = apply_func(i_neg, func_pow(3))
	float[] o_pow2p81_pos = apply_func(i_pos, func_pow(2.81))
	float[] o_exp_pos = apply_func(i_pos, func_exp())
	float[] o_exp2_neg = apply_func(i_neg, func_pow(2))
	float[] o_ln_pos = apply_func(i_pos, func_log())
	float[] o_log4_pos = apply_func(i_pos, func_log(4))
	
	validate_sqrt = true
	validate_powneg = true
	validate_powpos = true
	validate_exppos = true
	validate_expneg = true
	validate_ln = true
	validate_log4 = true

	for i = 0 to array.size(i_pos) - 1
		validate_sqrt := validate_sqrt and (math.sqrt(array.get(i_pos, i)) == array.get(o_sqrt_pos, i))
		validate_powneg := validate_powneg and (math.pow(array.get(i_neg, i), 3) == array.get(o_pow3_neg, i))
		validate_powpos := validate_powpos and (math.pow(array.get(i_pos, i), 2.81) == array.get(o_pow2p81_pos, i))
		validate_exppos := validate_exppos and (math.exp(array.get(i_pos, i)) == array.get(o_exp_pos, i))
		validate_expneg := validate_expneg and (math.pow(2, array.get(i_neg, i)) == array.get(o_exp2_neg, i))
		validate_ln := validate_ln and (math.log(array.get(i_pos, i)) == array.get(o_ln_pos, i))
		validate_log4 := validate_log4 and (math.log(array.get(i_pos, i))/math.log(4) == array.get(o_log4_pos, i))

	result := result + test_result_tostr("Sqrt", validate_sqrt)
	result := result + test_result_tostr("x^2, any x", validate_powneg)
	result := result + test_result_tostr("x^2.81, pos x", validate_powpos)
	result := result + test_result_tostr("Exp, pos x", validate_exppos)
	result := result + test_result_tostr("Exp, neg x", validate_expneg)
	result := result + test_result_tostr("Ln", validate_ln)
	result := result + test_result_tostr("Log 4", validate_log4)

	result

test_Aa_to_A_trig() =>
	string result = "f: A, a -> A, trigonometric suite:\n"

	float[] i_pos = array.from(0.25, 0.81, math.pi/4)

	float[] o_sin = apply_func(i_pos, func_sin())
	float[] o_cos = apply_func(i_pos, func_cos())
	float[] o_tan = apply_func(i_pos, func_tan())
	float[] o_cot = apply_func(i_pos, func_cot())
	float[] o_asin = apply_func(i_pos, func_asin())
	float[] o_acos = apply_func(i_pos, func_acos())
	float[] o_atan = apply_func(i_pos, func_atan())
	float[] o_acot = apply_func(i_pos, func_acot())

	validate_sin = true
	validate_cos = true
	validate_tan = true
	validate_cot = true
	validate_asin = true
	validate_acos = true
	validate_atan = true
	validate_acot = true


	for i = 0 to array.size(i_pos) - 1
		validate_sin := validate_sin and (math.sin(array.get(i_pos, i)) == array.get(o_sin, i))
		validate_cos := validate_cos and (math.cos(array.get(i_pos, i)) == array.get(o_cos, i))
		validate_tan := validate_tan and (math.tan(array.get(i_pos, i)) == array.get(o_tan, i))
		validate_cot := validate_cot and (1/math.tan(array.get(i_pos, i)) == array.get(o_cot, i))
		validate_asin := validate_asin and (math.asin(array.get(i_pos, i)) == array.get(o_asin, i))
		validate_acos := validate_acos and (math.acos(array.get(i_pos, i)) == array.get(o_acos, i))
		validate_atan := validate_atan and (math.atan(array.get(i_pos, i)) == array.get(o_atan, i))
		validate_acot := validate_acot and (math.pi/2 - math.atan(array.get(i_pos, i)) == array.get(o_acot, i))

	result := result + test_result_tostr("Sin", validate_sin)
	result := result + test_result_tostr("Cos", validate_cos)
	result := result + test_result_tostr("Tan", validate_tan)
	result := result + test_result_tostr("Cot", validate_cot)
	result := result + test_result_tostr("Asin", validate_asin)
	result := result + test_result_tostr("Acos", validate_acos)
	result := result + test_result_tostr("Atan", validate_atan)
	result := result + test_result_tostr("Acot", validate_acot)

	result

run_all_tests() =>
	result = "Racille array utils unit tests:\n\n"
	result := result + test_A_to_R() + "\n"
	result := result + test_AA_to_A() + "\n"
	result := result + test_AR_to_A() + "\n"
	result := result + test_Aa_to_A_common() + "\n"
	result := result + test_Aa_to_A_trig()

	result

if barstate.islast
	label.new(bar_index, close, run_all_tests(), textalign = text.align_left)
