import numpy as np
import pyceres
import pycolmap

def optimize_scale(sim3, scale_prior):
    prob = pyceres.Problem()
    prob.add_residual_block(
        pyceres.factors.NormalPrior(scale_prior, np.eye(1)),
        pyceres.TrivialLoss(),
        [sim3.scale],
    )
    prob.set_manifold(
        sim3.scale, pycolmap.manifold.PositiveExponentialManifold(1)
    )
    options = pyceres.SolverOptions()
    summary = pyceres.SolverSummary()
    pyceres.solve(options, prob, summary)


sim3d = pycolmap.Sim3d()
optimize_scale(sim3d, np.array([5.0]))
np.testing.assert_allclose(sim3d.scale, np.array([5.0]), atol=1e-5)
optimize_scale(sim3d, np.array([-1.0]))
np.testing.assert_allclose(sim3d.scale, np.array([0.0]), atol=1e-5)

import numpy as np
import pyceres
import pycolmap
import copy

sim3 = pycolmap.Sim3d()
target_sim3 = pycolmap.Sim3d(
    2.0, np.array([0.0, 0.0, 0.0, 1.0]), np.array([2.0, 3.0, 4.0])
)
ref_points = [np.array([k, k, k]).astype(float) for k in range(10)]

prob = pyceres.Problem()
points = []
for i, ref_point in enumerate(ref_points):
    if i < 3:
        point = ref_point + np.ones(3) * 0.5  # add noise
    else:
        point = ref_point
    points.append(point)
    target_point = target_sim3 * ref_point
    cost = pycolmap.cost_functions.Point3dAlignmentCost(target_point, np.eye(3))
    prob.add_residual_block(
        cost,
        pyceres.TrivialLoss(),
        [points[i], sim3.rotation.quat, sim3.translation, sim3.scale],
    )
    if i >= 3:
        prob.set_parameter_block_constant(ref_points[i])
prob.set_manifold(sim3.rotation.quat, pyceres.QuaternionManifold())
prob.set_manifold(sim3.scale, pycolmap.manifold.PositiveExponentialManifold(1))
options = pyceres.SolverOptions()
summary = pyceres.SolverSummary()
pyceres.solve(options, prob, summary)

np.testing.assert_allclose(points[0], ref_points[0], atol=1e-5)
np.testing.assert_allclose(points[1], ref_points[1], atol=1e-5)
np.testing.assert_allclose(points[2], ref_points[2], atol=1e-5)
np.testing.assert_allclose(sim3.scale, target_sim3.scale, atol=1e-5)
np.testing.assert_allclose(
    sim3.rotation.quat, target_sim3.rotation.quat, atol=1e-5
)
np.testing.assert_allclose(sim3.translation, target_sim3.translation, atol=1e-5)