# This file is part of Torch Lens Maker
# Copyright (C) 2024-present Victor Poughon
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.


import pytest
import torch
import torch.nn as nn

from torchlensmaker.core.base_module import (
    MultiForwardModule,
    multiforward,
)


def test_multi_forward() -> None:
    "Test multiforward and multiforward module"

    class ReversibleTranslation2D(MultiForwardModule):
        def __init__(self):
            super().__init__()
            self.translation = nn.Parameter(torch.zeros(2))

        @multiforward
        def forward_kinematic(self, point):
            return point + self.translation

        @multiforward
        def reverse_kinematic(self, translated):
            return translated - self.translation

    model = ReversibleTranslation2D()
    data = torch.tensor([[1.0, 2.0], [3.0, 4.0]])  # Batch of 2 points

    target_forward = torch.tensor([[11.0, 7.0], [13.0, 9.0]])
    target_reverse = torch.tensor([[-9.0, -3.0], [-7.0, -1.0]])

    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

    for i in range(2):
        print("> ", i)

        translated = model.forward_kinematic(data)
        reverse = model.reverse_kinematic(data)

        loss_fwd = nn.MSELoss()(translated, target_forward)
        loss_rev = nn.MSELoss()(reverse, target_reverse)

        # Compute loss dependent on both forward and reverse
        total_loss = loss_fwd + loss_rev

        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()
        print(model.translation)
