import numpy as np
import sys
import ast
import argparse




''' PROBLEM 1 :
Linear solution (just one loop through the array)'''
def problem1 (arr: np.ndarray, n: int):
solutions = np.zeros(n)
first_index = get_first_index(arr)
if first_index == -1:
returnsolutions - 1
else:
solutions[0:first_index] = -1
idx_at = first_index
while idx_at < n-1:
idx_at += 1
if arr[idx_at] == 1:
solutions[idx_at] = 0
first_index = idx_at
else:
solutions[idx_at] = (idx_at - first_index)
returnsolutions


def TestProblem1():

arr1 = np.array([0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0])
first_ = get_first_index(arr1)
solutions = problem1(arr1, len(arr1))
explicit_solution = np.array([-1, -1, -1, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 0, 1, 2, 3, 4, 5, 0, 1,2 ,3 ,4, 0, 1, 2])
assert np.linalg.norm(solutions - explicit_solution) < 1.0E-7


''' PROBLEM 2
IDEA : represent the kernel as a matrix, use masking
to set as many values to be 0,
then the formula of the convolution is simply
the product of the kernel with arr'''
def problem2(arr: np.ndarray, n: int, w: int, h1: int, h2: int):
EXP_CUTOFF = 6

# column vector
i_axis = np.arange(n).reshape(n, 1)
# row vector
j_axis = np.arange(n).reshape(1, n)
# diff matrix
diff_index_matrix = i_axis - j_axis

# apply mask (conditions on sum) + exponential kernel cutoff
tmp_ker_matrix = diff_index_matrix * (diff_index_matrix >=0 ) * (diff_index_matrix < w) * (np.abs(diff_index_matrix / h1) <= EXP_CUTOFF)

# kernel matrix
ker_matrix = np.exp(-tmp_ker_matrix / h1) * np.cos(tmp_ker_matrix / h2)
ker_matrix = ker_matrix * (diff_index_matrix >=0 ) * (diff_index_matrix < w) * (np.abs(diff_index_matrix / h1) <= EXP_CUTOFF)

returnnp.dot(ker_matrix, arr)


def explicit_problem2(arr: np.ndarray, n: int, w: int, h1: int, h2: int):
b = arr * 0
for i in range(0, n):
b_i = 0
for j in range(0, n):
if (j <= i) and (j >= 0) and (i - w < j):
diff_i_j = i - j
b_i += arr[j] * np.exp(-diff_i_j / h1) * np.cos(diff_i_j / h2)
b[i] = b_i
returnb


def TestProblem2():
arr_test = np.array([1, -1, 2, -1, 2, 5, -0.5, -2.5, 1, 0.3, -7, 1, 2, 3, 4, -7, -0.2, 0.3, 0.4, 5.0])

n_test = len(arr_test)
h1_test = (n_test // 2) + 2
h2_test = n_test - 2
w_test = (n_test // 10) + 2

solution_explicit = explicit_problem2(arr_test, n_test, w_test, h1_test, h2_test)
solution_vectorised = problem2(arr_test, n_test, w_test, h1_test, h2_test)
relative_diff = np.linalg.norm(solution_vectorised - solution_explicit) / np.linalg.norm(solution_explicit)
assert relative_diff < 1.0E-5


''' PROBLEM 3
IDEA : Induction : X(i+1) := X(i) + A(i+1) * SMOOTHING(i+1) +
Sum_{ [S(i), S(i+1)[ } A(k) * SMOOTHING(k)

So if build the array of Deltas
Delta(i) = Sum_{ [S(i), S(i+1)[ } A(k) * SMOOTHING(k)

I have X(i+1) := X(i) + A(i+1) * SMOOTHING(i+1) + Delta(i)'''
def problem3(arr: np.ndarray, n: int, s_vec: np.ndarray):
EXP_CUTOFF = 6
solutions = np.zeros(n)
delta_arr = np.zeros(n)
for i in range(0, n-1):
s_i = s_vec[i]
s_i_incr = s_vec[i+1]
min_s = min(s_i, s_i_incr)
max_s = max(s_i, s_i_incr)
if min_s < max_s:
for k in range(min_s, max_s):
exp_arg = 50 * k / n
if exp_arg <= EXP_CUTOFF:
delta_arr[i] += (arr[k] * np.exp(-exp_arg))

solutions[0] = arr[0]
for i in range(1, n):
s_i = s_vec[i]
s_i_decr = s_vec[i-1]
exp_arg = 50 * i / n
new_term = arr[i] * (np.exp(-exp_arg) if exp_arg <= EXP_CUTOFF else 0)
if s_i == i:
solutions[i] = new_term
elif s_i == s_i_decr:
solutions[i] = solutions[i-1] + new_term
elif s_i < s_i_decr:
solutions[i] = solutions[i-1] + new_term + delta_arr[i-1]
else:
solutions[i] = solutions[i-1] + new_term - delta_arr[i-1]

returnsolutions


def explicit_problem3(arr: np.ndarray, n: int, s_vec: np.ndarray):
solutions = np.zeros(n)
solutions[0] = arr[0]
for i in range(1, n):
for k in range(s_vec[i], i+1):
solutions[i] += (arr[k] * np.exp(-(50 * k / n)))
returnsolutions


def TestProblem3():
arr_test = np.array([1, -1, 2, -1, 2, 5, -0.5, -2.5, 1, 0.3, -7, 1, 2, 3, 4, -7, -0.2, 0.3, 0.4, 5.0])
s_vec_test = np.array([0, 1, 1, 2, 1, 1, 3, 5, 7, 7, 7, 5, 6, 10, 5, 11, 6, 0, 5, 18])
n_test = len(arr_test)
assert n_test == len(s_vec_test)
explicit_solution = explicit_problem3(arr_test, n_test, s_vec_test)
fast_solution = problem3(arr_test, n_test, s_vec_test)
relative_diff = np.linalg.norm(explicit_solution - fast_solution) / np.linalg.norm(explicit_solution)
assert relative_diff < 1.0E-3


def main():
parser = argparse.ArgumentParser(
description="Call one of problems : P1, P2, P3 with the associated list(s) of numbers"
)

# Define arguments
parser.add_argument(
"--problem", "-p",
type=str,
required=True,
choices=["P1", "P2", "P3"],
help="Which problem to run"
)
parser.add_argument(
"--primary_arr", "-a",
type=str, # "[3, 12, 17]"
required=True,
help="List of primary array (space-separated)"
)
parser.add_argument(
"--other_params", "-o",
type=str, # "[3, 12, 17]"
required=False,
help="List of params for problem 2 : w, h1, h2 (space-separated)"
)
parser.add_argument(
"--second_arr", "-s",
type=str, # "[0, 1, 0, 1, 2, 0, 5, 1, 2, 6]"
required=False,
help="List of for kernel arr of S arr (problem 3) (space-separated)"
)

args = parser.parse_args()

problem = args.problem
arr = ast.literal_eval(args.primary_arr)
arr_np = np.array(arr)
n = len(arr_np)

if problem == "P1":
print("Launch problem 1")
returnproblem1(arr_np, n)
elif problem == "P2":
print("Launch problem 2")
params = ast.literal_eval(args.other_params)
w = int(params[0])
h1 = int(params[1])
h2 = int(params[2])
returnproblem2(arr_np, n, w, h1, h2)
elif problem == "P3":
print("Launch problem 3")
s_vec = ast.literal_eval(args.second_arr)
s_vec_np = np.array(s_vec)
returnproblem3(arr_np, n, s_vec_np)


if __name__ == "__main__":
TestProblem1()
print(f" Test Problem 1 passed")
TestProblem2()
print(f" Test Problem 2 passed")
TestProblem3()
print(f" Test Problem 3 passed")

print(" Problem 1 ex command line :")
cmd1 = ''' python.exe C:\problems.py --problem P3 --primary_arr "[0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0]" '''
print(cmd1)

print(" Problem 2 ex command line :" )
cmd2 = ''' python.exe C:\problems.py --problem P3 --primary_arr "[1, -1, 2, -1, 2, 5, -0.5, -2.5, 1, 0.3, -7, 1, 2, 3, 4, -7, -0.2, 0.3, 0.4, 5.0]" --other_params "[3, 12, 17]" '''
print(cmd2)

print(" Problem 3 ex command line :" )
cmd3 = ''' python.exe C:\problems.py --problem P3 --primary_arr "[1, -1, 2, -1, 2, 5, -0.5, -2.5, 1, 0.3, -7, 1, 2, 3, 4, -7, -0.2, 0.3, 0.4, 5.0]" --second_arr "[0, 1, 1, 2, 1, 1, 3, 5, 7, 7, 7, 5, 6, 10, 5, 11, 6, 0, 5, 18]"'''
print(cmd3)
main()