import math
import numpy as np
from scipy.optimize import fmin
N = 6
x0 = 1

# Ranges:
# x in [0, ..., N)
# u in [0, ..., N - 1)
# v in [0, ..., N)


def get_state_values(u: np.ndarray) -> np.ndarray:
    x = np.zeros((N))
    x[0] = x0
    for i in range(1, N):
        x[i] = 2 * x[i - 1] + u[i - 1]
    return x


def constrained_f(u: np.ndarray, v: np.ndarray, t: float) -> float:
    x = get_state_values(u)
    f = 0
    for i in range(N - 1):
        f += math.pow(x[i], 2) + 2 * math.pow(u[i], 2)
    cost = 0
    # Constraint imposed on control value:
    for i in range(N - 1):
        cost += (u[i] - v[i, 0]) * max(0, u[i] - v[i, 0])
        cost += (-u[i] - v[i, 1]) * max(0, -u[i] - v[i, 1])
    # Constraint imposed on the final state value:
    cost += math.pow((x[N-1] - v[N - 1, 0]), 2)
    return f + (t/2) * cost


def main():
    alpha = 2
    beta = 2
    epsilon = 0.1
    t = 0.1
    c = 0.1
    x = np.zeros((N))
    u = np.zeros((N - 1))
    v = np.zeros((N, 2))
    lower = 1
    upper = 3
    last = 20
    for i in range(N - 1):
        v[i, 0] = upper
        v[i, 1] = lower
    v[N - 1][0] = last

    while True:
        f_old = constrained_f(u, v, t)
        u = fmin(constrained_f, u, args=(v, t), disp=False)
        f_new = constrained_f(u, v, t)

        r = np.zeros((N, 2))
        for i in range(N - 1):
            r[i, 0] = u[i] - upper
            r[i, 1] = -u[i] - lower
        r[N - 1, 0] = x[N - 1] - last

        gamma = np.zeros((N, 2))
        for i in range(N - 1):
            gamma[i, 0] = v[i, 0] + r[i, 0] - upper
            gamma[i, 1] = v[i, 1] + r[i, 1] - lower
        gamma[N - 1, 0] = x[N - 1] + r[N - 1, 0] - last
        gamma_sum = np.sum(gamma)
        if gamma_sum <= epsilon:
            if np.abs(r[N - 1, 0]) <= np.abs(gamma[N - 1, 0]) 
                    and np.abs(f_old - f_new) <= epsilon:
                break
        else:
            if gamma_sum < c:
                for i in range(N - 1):
                    v[i, 0] = upper - r[i, 0]
                    v[i, 1] = lower - r[i, 1]
                v[N - 1, 0] = last - r[N - 1, 0]
                c = alpha * gamma_sum
            else:
                t = beta * t
                for i in range(N - 1):
                    v[i, 0] = upper - r[i, 0] / beta
                    v[i, 1] = lower - r[i, 1] / beta
                v[N - 1, 0] = last - r[N - 1, 0] / beta

    print(u)
    print(get_state_values(u))
    print(f_new)


if __name__ == '__main__':
    main()