import numpy as np

from common.consts import AU
from sim_scenes.solar_system.halley_comet_lib import HalleyCometParams
from sim_lab.halley_comet_research_calc import calc_simulator


def target_function(x, y, z):
    params = HalleyCometParams(
        start_time='1982-09-24 00:00:00',
        # init_velocity=[-2.836 + (x / 1000), 4.705 + (y / 1000), 8.85 + (z / 1000)],
        init_velocity=[x, y, z],
        init_position=[0, -5 * AU, -10 * AU]
    )
    return calc_simulator(params)


def loss_function(a, b, c):
    # 0.586AU, 35.1AU, 13817天
    target_a = 0.586
    target_b = 35.1
    target_c = 13817.0

    loss_a = (a - target_a) ** 2
    loss_b = (b - target_b) ** 2
    loss_c = (c - target_c) ** 2

    total_loss = loss_a + loss_b + loss_c

    return total_loss


def compute_gradients(x, y, z, epsilon=1e-5):
    # 数值方法计算梯度
    grad_x = (loss_function(*target_function(x + epsilon, y, z)) - loss_function(*target_function(x, y, z))) / epsilon
    grad_y = (loss_function(*target_function(x, y + epsilon, z)) - loss_function(*target_function(x, y, z))) / epsilon
    grad_z = (loss_function(*target_function(x, y, z + epsilon)) - loss_function(*target_function(x, y, z))) / epsilon

    return grad_x, grad_y, grad_z


def adaptive_learning_rate(iteration):
    # 自适应学习率调整
    return 0.01 / np.sqrt(iteration + 1)


def gradient_descent(init_velocity, iterations):
    x, y, z = init_velocity

    for i in range(1, iterations + 1):  # 从1开始，避免除以0
        a, b, c = target_function(x, y, z)
        loss = loss_function(a, b, c)

        # 计算梯度
        grad_x, grad_y, grad_z = compute_gradients(x, y, z)

        # 自适应学习率调整
        learning_rate = adaptive_learning_rate(i)

        # 更新参数
        x = x - learning_rate * grad_x
        y = y - learning_rate * grad_y
        z = z - learning_rate * grad_z

        # 打印当前迭代的结果
        if i % 100 == 0:
            print(f"Iteration {i}: Loss = {loss}, a = {a}, b = {b}, c = {c}")

    return x, y, z


# 设置初始值
init_velocity = [-2.816, 4.725, 8.87]  # 请替换为实际的初始值
iterations = 1000

# 运行梯度下降算法
final_x, final_y, final_z = gradient_descent(init_velocity, iterations)

# 打印最终结果
final_a, final_b, final_c = target_function(final_x, final_y, final_z)
print(f"Final Result: a = {final_a}, b = {final_b}, c = {final_c}")