# Copyright 2021 The NetKet Authors - All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
2D planar symmetry groups and transformations.

This module provides functions to construct point groups and symmetry operations
for 2D systems, including rotations, reflections, and glide operations.
It includes common 2D point groups like cyclic (C) and dihedral (D) groups,
as well as specific groups for rectangles and squares.
"""

import numpy as np

from netket.utils.types import Array
from netket.utils.moduletools import export, hide_unexported

from netket.utils.group._point_group import PGSymmetry, PointGroup
from netket.utils.group._semigroup import Identity


hide_unexported(__name__)

__all__ = [
    "rotation_group",
    "dihedral",
]


@export
def rotation(angle: float) -> PGSymmetry:
    """
    Returns a 2D rotation by `angle` degrees.

    Args:
       angle: the angle of rotation in degrees.
    """
    angle = np.radians(angle)
    return PGSymmetry(
        np.asarray([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]])
    )


@export
def C(n: int) -> PointGroup:
    """
    Returns the Cₙ `PointGroup` of 2D rotational symmetries.

    Arguments:
        n: the index of the rotation group (the smallest rotation angle is 360°/n)

    Returns:
        a `PointGroup` implementing Cₙ
    """
    return PointGroup(
        [Identity()] + [rotation(360 / n * i) for i in range(1, n)], ndim=2
    )


rotation_group = C


@export
def reflection(angle: float) -> PGSymmetry:
    """
    Returns a 2D reflection across an axis at angle `angle` to the +x direction.

    Args:
        angle: the angle between the +x axis and the reflection axis.
    """
    axis = np.radians(angle) * 2  # the mirror matrix is written in terms of 2φ
    return PGSymmetry(
        np.asarray([[np.cos(axis), np.sin(axis)], [np.sin(axis), -np.cos(axis)]])
    )


@export
def reflection_group(angle: float) -> PointGroup:
    """
    Returns the Z₂ `PointGroup` containing the identity and a reflection across an
    axis at angle `angle` to the +x.

    Args:
        angle: the angle between the +x axis and the reflection axis.
    """
    return PointGroup([Identity(), reflection(angle)], ndim=2)


@export
def glide(trans: Array, origin: Array = (0, 0)) -> PGSymmetry:
    """
    Returns a 2D glide composed of translation by `trans` and reflection across
    its direction.

    Arguments:
        trans: translation vector
        origin: a point on the glide axis, defaults to the origin
    """
    axis = np.arctan2(trans[1], trans[0]) * 2
    W = np.asarray([[np.cos(axis), np.sin(axis)], [np.sin(axis), -np.cos(axis)]])
    w = np.asarray(trans) + (np.eye(2) - W) @ np.asarray(origin)
    return PGSymmetry(W, w)


@export
def glide_group(trans: Array, origin: Array = (0, 0)) -> PGSymmetry:
    """
    Returns the Z_2 `PointGroup`containing the identity and a  2D glide composed
    of translation by `trans` and reflection across its direction.

    Arguments:
        trans: translation vector
        origin: a point on the glide axis, defaults to the origin
    The output is only a valid `PointGroup` after supplying a `unit_cell`
    consistent with the glide axis; otherwise, operations like `product_table`
    will fail.
    """
    return PointGroup([Identity(), glide(trans, origin)], ndim=2)


@export
def D(n: int, axis: float = 0) -> PointGroup:
    """
    Returns the 2D dihedral `PointGroup` :math:`D_n` generated by a 360°/n rotation
    and a reflection.

    Arguments:
        n: index of the dihedral group
        axis: optional, the angle of one reflection axis with the +x direction
            (default: 0)

    Returns:
        a `PointGroup` object implementing :math:`D_n`
    """
    return reflection_group(axis) @ C(n)


dihedral = D


@export
def rectangle() -> PointGroup:
    """
    The symmetry group of a rectangle aligned with the Cartesian axes (Vierergruppe).
    """
    return D(2)


@export
def square() -> PointGroup:
    """The symmetry group of a square aligned with the Cartesian axes."""
    return D(4)
