import numpy as np
from scipy.optimize import least_squares

# Эталонные значения параметров
K1_REFERENCE = 0.577
K2_REFERENCE = 0.422

# Диапазоны из задания
K1_BOUNDS = (1e-5, 1e2)
K2_BOUNDS = (1e-2, 1.0)

T_BEGIN = 0.0
T_END = 10.0
H = 0.1


def analytical_solution(t, k1, k2):
    """
    Аналитическое решение задачи Коши.
    """
    t = np.asarray(t, dtype=float)

    x1 = np.exp(-k1 * t)

    if np.isclose(k1, k2):
        x2 = k1 * t * np.exp(-k1 * t)
    else:
        x2 = (
            k1 / (k2 - k1)
            * (np.exp(-k1 * t) - np.exp(-k2 * t))
        )

    x3 = 1.0 - x1 - x2

    return np.column_stack((x1, x2, x3))


def derivative(state, k1, k2):
    """
    Правая часть системы дифференциальных уравнений.
    """
    x1, x2, x3 = state

    return np.array([
        -k1 * x1,
        k1 * x1 - k2 * x2,
        k2 * x2
    ])


def rk4_step(state, h, k1, k2):
    """
    Один шаг метода Рунге--Кутты четвёртого порядка.
    """
    q1 = derivative(state, k1, k2)
    q2 = derivative(state + h * q1 / 2.0, k1, k2)
    q3 = derivative(state + h * q2 / 2.0, k1, k2)
    q4 = derivative(state + h * q3, k1, k2)

    return state + h / 6.0 * (q1 + 2.0 * q2 + 2.0 * q3 + q4)


def rk4_solution(t, k1, k2):
    """
    Решение задачи Коши методом RK4.
    """
    solution = np.zeros((len(t), 3), dtype=float)
    solution[0] = [1.0, 0.0, 0.0]

    for index in range(1, len(t)):
        h = t[index] - t[index - 1]

        solution[index] = rk4_step(
            solution[index - 1],
            h,
            k1,
            k2
        )

    return solution


# Моменты времени из задания:
# tj = 0.1j, j = 0, ..., 100
time = np.linspace(T_BEGIN, T_END, 101)

# Эталонные данные
reference_solution = analytical_solution(time, K1_REFERENCE, K2_REFERENCE)


def residuals(parameters):
    """
    Вектор невязок между моделью и эталонными данными.
    """
    k1, k2 = parameters

    numerical_solution = rk4_solution(time, k1, k2)

    # Получается вектор из 101 * 3 = 303 невязок
    return (numerical_solution - reference_solution).ravel()


# Начальное приближение
initial_parameters = np.array([0.3, 0.7])

result = least_squares(
    residuals,
    initial_parameters,
    bounds=(
        [K1_BOUNDS[0], K2_BOUNDS[0]],
        [K1_BOUNDS[1], K2_BOUNDS[1]]
    )
)

# result.x - solution
found_k1, found_k2 = result.x

absolute_error_k1 = abs(found_k1 - K1_REFERENCE)
absolute_error_k2 = abs(found_k2 - K2_REFERENCE)

relative_error_k1 = absolute_error_k1 / K1_REFERENCE * 100.0
relative_error_k2 = absolute_error_k2 / K2_REFERENCE * 100.0
criterion_value = np.sum(result.fun ** 2)

print("Results")
print(f"Found k1:{found_k1:.4f}")
print(f"Theoretical k1:{K1_REFERENCE:.4f}")
print(f"Absolute error of k1:{absolute_error_k1:.4e}")
print(f"Relative error of k1:{relative_error_k1:.4e} %")
print()

print(f"Found k2:{found_k2:.4f}")
print(f"Theoretical k2:{K2_REFERENCE:.4f}")
print(f"Absolute error of k2:{absolute_error_k2:.4e}")
print(f"Relative error of k2:{relative_error_k2:.4e} %")
print()