# 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/>.


from jaxtyping import Float
import torch

from torchlensmaker.types import (
    BatchTensor,
    ScalarTensor,
    Batch2DTensor,
    Batch3DTensor,
    MaskTensor,
    Tf,
)

from torchlensmaker.kinematics.homogeneous_geometry import (
    hom_scale_2d,
    hom_translate_2d,
    kinematic_chain_extend_2d,
    hom_scale_3d,
    hom_translate_3d,
    kinematic_chain_extend_3d,
)


def lens_diameter_domain_2d(
    points: Batch2DTensor, diameter: ScalarTensor
) -> MaskTensor:
    return torch.abs(points[..., 1]) <= diameter / 2


def lens_diameter_domain_3d(
    points: Batch3DTensor, diameter: ScalarTensor
) -> MaskTensor:
    _, y, z = points.unbind(-1)
    r2 = y**2 + z**2
    return r2 <= (diameter / 2) ** 2


def lens_diameter_implicit_domain_2d(
    F: BatchTensor,
    points: Batch2DTensor,
    diameter: ScalarTensor,
    tol: float,
) -> MaskTensor:
    in_diameter = torch.abs(points[..., 1]) <= diameter / 2
    in_F = torch.abs(F) <= tol
    return torch.logical_and(in_diameter, in_F)


def lens_diameter_implicit_domain_3d(
    F: BatchTensor,
    points: Batch3DTensor,
    diameter: ScalarTensor,
    tol: float,
) -> MaskTensor:
    _, y, z = points.unbind(-1)
    r2 = y**2 + z**2
    in_diameter = r2 <= (diameter / 2) ** 2
    in_F = torch.abs(F) <= tol
    return torch.logical_and(in_diameter, in_F)


def anchor_transforms_2d(
    anchor0: ScalarTensor,
    anchor1: ScalarTensor,
    scale: ScalarTensor,
    base: Tf,
) -> tuple[Tf, Tf]:
    """
    Compute transforms required to position a surface that has anchors and a scale.

    They are:
    - the "surface transform": the tf that applies to the surface itself to
      position it in global frame
    - the "next transform": the tf that applies to the next element in a
      sequential system where the surface is the latest element of the kinematic
      chain
    """
    # First anchor transform
    t0_x = -scale * anchor0
    t0_y = torch.zeros_like(t0_x)
    tf0 = hom_translate_2d(torch.stack((t0_x, t0_y), dim=-1))

    # Second anchor transform
    t1_x = scale * anchor1
    t1_y = torch.zeros_like(t0_y)
    tf1 = hom_translate_2d(torch.stack((t1_x, t1_y), dim=-1))

    # Compose with the existing kinematic chain
    tfscale = hom_scale_2d(scale)
    tf_surface = kinematic_chain_extend_2d(base, [tf0, tfscale])
    tf_next = kinematic_chain_extend_2d(base, [tf0, tf1])

    return tf_surface, tf_next


def anchor_transforms_3d(
    anchor0: ScalarTensor,
    anchor1: ScalarTensor,
    scale: ScalarTensor,
    base: Tf,
) -> tuple[Tf, Tf]:
    """
    Compute transforms required to position a surface that has anchors and a scale.

    They are:
    - the "surface transform": the tf that applies to the surface itself to
      position it in global frame
    - the "next transform": the tf that applies to the next element in a
      sequential system where the surface is the latest element of the kinematic
      chain
    """
    # First anchor transform
    t0_x = -scale * anchor0
    t0_y = torch.zeros_like(t0_x)
    t0_z = torch.zeros_like(t0_x)
    tf0 = hom_translate_3d(torch.stack((t0_x, t0_y, t0_z), dim=-1))

    # Second anchor transform
    t1_x = scale * anchor1
    t1_y = torch.zeros_like(t1_x)
    t1_z = torch.zeros_like(t1_x)
    tf1 = hom_translate_3d(torch.stack((t1_x, t1_y, t1_z), dim=-1))

    # Compose with the existing kinematic chain
    tfscale = hom_scale_3d(scale)
    tf_surface = kinematic_chain_extend_3d(base, [tf0, tfscale])
    tf_next = kinematic_chain_extend_3d(base, [tf0, tf1])

    return tf_surface, tf_next
