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 f(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)


from scipy.optimize import minimize
import numpy as np



# 定义损失函数
def loss_function(x, target):
    # 定义损失函数，这里使用简单的平方差，并调整权重
    weights = np.array([10, 1, 0.1])  # 调整权重，使近日点和远日点的偏差对最终结果的影响更大
    return np.sum(weights * (x - target) ** 2)


# 定义目标函数
def target_function(v):
    vx, vy, vz = v
    result = f(vx, vy, vz)
    # 计算损失函数，这里假设目标是 [0.586, 35.1, 13817]
    target = np.array([0.586, 35.1, 13817])
    loss = loss_function(result, target)
    return loss


# 设置初始速度
initial_velocity = [-2.816, 4.725, 8.87]  # 请替换为实际的初始值
# 设置初始速度时添加随机扰动
initial_velocity = [-2.816 + np.random.uniform(-1, 1), 4.725 + np.random.uniform(-1, 1), 8.87 + np.random.uniform(-1, 1)]


# 运行优化算法
result = minimize(target_function, initial_velocity, method='trust-constr', options={'maxiter': 10000, 'disp': True},
                  bounds=[(-100, 100), (-100, 100), (-100, 100)])

# 打印最终结果
final_velocity = result.x
print(f"Optimal Velocity: {final_velocity}")

#
# # # 定义目标函数
# # def f(vx, vy, vz):
# #     # 这里需要根据你的实际情况定义哈雷彗星轨道的计算
# #     # 返回近日点距离（单位：AU）、远日点距离（单位：AU）、以及哈雷彗星近日点到远日点花了多长时间（单位：天）
# #     # 这里使用一个简单的示例，你需要根据实际情况修改
# #     return np.array([0.586, 35.1, 13817])
#
# # 定义损失函数
# def loss_function(x, target):
#     # 定义损失函数，这里使用简单的平方差
#     return (x - target) ** 2
#
#
# # 定义目标函数
# def target_function(v):
#     vx, vy, vz = v
#     result = f(vx, vy, vz)
#     # 计算损失函数，这里假设目标是 [0.586, 35.1, 13817]
#     target = np.array([0.586, 35.1, 13817])
#     loss = np.sum([loss_function(result[i], target[i]) for i in range(len(result))])
#     return loss
#
#
# # 设置初始速度
# initial_velocity = [-2.816, 4.725, 8.87]  # 请替换为实际的初始值
#
# # 运行优化算法
# # result = minimize(target_function, initial_velocity, method='BFGS',
# #                   options={'maxiter': 1000, 'disp': True})  # 使用BFGS优化算法，你可以尝试其他算法
#
# # 运行优化算法
# # result = minimize(target_function, initial_velocity, method='L-BFGS-B', options={'maxiter': 10000, 'disp': True})
#
# # 运行优化算法
# result = minimize(target_function, initial_velocity, method='trust-constr', options={'maxiter': 10000, 'disp': True},
#                   bounds=[(-100, 100), (-100, 100), (-100, 100)])
#
# # 打印最终结果
# final_velocity = result.x
# print(f"Optimal Velocity: {final_velocity}")