# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import copy

import pytest  # noqa: F401
import torch

import theseus as th

from .common import (
    check_another_theseus_function_is_copy,
    check_another_theseus_tensor_is_copy,
)


def test_copy_scale_cost_weight():
    scale = th.Variable(torch.tensor(1.0))
    p1 = th.ScaleCostWeight(scale, name="scale_cost_weight")
    for the_copy in [p1.copy(), copy.deepcopy(p1)]:
        check_another_theseus_function_is_copy(p1, the_copy, new_name=f"{p1.name}_copy")
        check_another_theseus_tensor_is_copy(scale, the_copy.scale)


def test_copy_diagonal_cost_weight():
    diagonal = th.Variable(torch.ones(3))
    p1 = th.DiagonalCostWeight(diagonal, name="diagonal_cost_weight")
    for the_copy in [p1.copy(), copy.deepcopy(p1)]:
        check_another_theseus_function_is_copy(p1, the_copy, new_name=f"{p1.name}_copy")
        check_another_theseus_tensor_is_copy(diagonal, the_copy.diagonal)


def test_scale_cost_weight():
    for dim in [1, 2, 10]:
        for batch_size in [1, 2, 10]:
            v1 = th.Vector(tensor=torch.ones(batch_size, dim))
            z = th.Vector(tensor=torch.zeros(batch_size, dim))
            scale = torch.randn(1).item()
            expected_err = torch.ones(batch_size, dim) * scale
            expected_jac = (torch.eye(dim).unsqueeze(0) * scale).expand(
                batch_size, dim, dim
            )

            def _check(cw):
                cf1 = th.Difference(v1, z, cw)
                jacobians, error = cf1.weighted_jacobians_error()
                assert error.allclose(expected_err)
                assert jacobians[0].allclose(expected_jac)

            _check(th.ScaleCostWeight(scale))
            _check(th.ScaleCostWeight(torch.ones(1) * scale))
            _check(th.ScaleCostWeight(torch.ones(batch_size) * scale))
            _check(th.ScaleCostWeight(torch.ones(1, 1) * scale))
            _check(th.ScaleCostWeight(torch.ones(batch_size, 1) * scale))
            _check(th.ScaleCostWeight(th.Variable(torch.ones(1, 1) * scale)))

            batched_scale = torch.randn(batch_size, 1)
            expected_err = batched_scale * torch.ones(batch_size, dim)
            expected_jac = (torch.eye(dim).unsqueeze(0)).expand(
                batch_size, dim, dim
            ) * batched_scale.view(-1, 1, 1)

            cf1 = th.Difference(v1, z, th.ScaleCostWeight(batched_scale))
            jacobians, error = cf1.weighted_jacobians_error()
            assert error.allclose(expected_err)
            assert jacobians[0].allclose(expected_jac)


def test_diagonal_cost_weight():
    for dim in [1, 2, 10]:
        for batch_size in [1, 2, 10]:
            v1 = th.Vector(tensor=torch.ones(batch_size, dim))
            z = th.Vector(tensor=torch.zeros(batch_size, dim))
            diagonal = torch.randn(dim)
            d_matrix = diagonal.diag().unsqueeze(0)
            expected_err = (d_matrix @ torch.ones(batch_size, dim, 1)).view(
                batch_size, dim
            )
            expected_jac = d_matrix @ (torch.eye(dim).unsqueeze(0)).expand(
                batch_size, dim, dim
            )

            def _check(cw):
                cf1 = th.Difference(v1, z, cw)
                jacobians, error = cf1.weighted_jacobians_error()
                assert error.allclose(expected_err)
                assert jacobians[0].allclose(expected_jac)

            _check(th.DiagonalCostWeight(diagonal.tolist()))
            _check(th.DiagonalCostWeight(diagonal))
            _check(th.DiagonalCostWeight(diagonal.view(1, dim)))
            _check(th.DiagonalCostWeight(th.Variable(diagonal)))

            batched_diagonal = torch.randn(batch_size, dim)
            d_matrix = batched_diagonal.diag_embed()
            expected_err = (d_matrix @ torch.ones(batch_size, dim, 1)).view(
                batch_size, dim
            )
            expected_jac = d_matrix @ (torch.eye(dim).unsqueeze(0)).expand(
                batch_size, dim, dim
            )

            cf1 = th.Difference(v1, z, th.DiagonalCostWeight(batched_diagonal))
            jacobians, error = cf1.weighted_jacobians_error()
            assert error.allclose(expected_err)
            assert jacobians[0].allclose(expected_jac)
