# Copyright 2023-2026 The PULGON Project Developers
#
# 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.

import argparse
import itertools
import logging
import warnings
from fractions import Fraction
from typing import List, Optional, Tuple, Union

import ase
import numpy as np
from ase import Atoms
from ase.io import read
from pymatgen.core.operations import SymmOp
from pymatgen.util.coord import find_in_coord_list

from pulgon_tools.utils import (
    angle_between_points,
    decimal_places,
    refine_cell,
)

warnings.filterwarnings(
    "ignore",
    category=RuntimeWarning,
    message="invalid value encountered in divide",
)


class CyclicGroupAnalyzer:
    """Analyze the generalized translational group
    (cyclic group).

    The general outline of the algorithm is as follows:

    1. Find primitive cell and get the pure translation
       distance in z direction.
    2. Get the potential monomers and the corresponding
       translational vector.
    3. Detect the possible cyclic group:
       - If rotational symmetry exists, it's a screw-axis
         group TQ(f) generated by (CQ|f).
       - If mirror symmetry exists, it's a glide plane
         group T'(f) generated by (sigmaV|f).

    """

    def __init__(
        self,
        atom: ase.atoms.Atoms,
        tolerance: float = 0.001,
    ) -> None:
        """
        Args:
            atom: Line group structure to determine
                the generalized translational group
            tolerance: distance tolerance for symmetry detection

        """
        logging.debug(
            "--------------------start detecting generalized translational group"
        )
        alongOZ = self._check_if_along_OZ(atom)
        if not alongOZ:
            logging.error(
                "Error while detect cyclic group: The axis direction is not OZ"
            )
        else:
            self._symprec = tolerance
            self._layer_symprec = 0.05

            round_symprec = decimal_places(tolerance)
            self._round_symprec = round_symprec
            self._zaxis = np.array([0, 0, 1])

            self._atom = self._find_axis_center_of_nanotube(atom)
            self._primitive = self._find_primitive()
            self.supercell_mutiple = int(
                len(self._atom.numbers) / len(self._primitive.numbers)
            )
            self._pure_trans = self._primitive.cell[2, 2]
            self._analyze()

    def _analyze(self) -> None:
        """Detect all possible monomers and determine their cyclic groups.

        Sets self.cyclic_group, self.monomers, and self._sym_operations.
        """
        monomer, potential_trans = self._potential_translation()

        logging.debug("There are %d monomer candidates" % len(monomer))
        (
            self.cyclic_group,
            self.monomers,
            self._sym_operations,
        ) = self._get_translations(monomer, potential_trans)

    def _find_axis_center_of_nanotube(
        self, atom: ase.atoms.Atoms
    ) -> ase.atoms.Atoms:
        """Shift the center of mass to scaled position (x,y) = (0.5, 0.5).

        Args:
            atom: initial structure

        Returns: centralized structure

        """
        n_st = atom.copy()
        center = self._get_center_of_mass_periodic(atom)

        pos = (
            np.remainder(atom.get_scaled_positions() - center + 0.5, [1, 1, 1])
            @ atom.cell
        )
        atoms = Atoms(
            cell=n_st.cell,
            numbers=n_st.numbers,
            positions=pos,
        )
        return atoms

    def _center_of_nanotube(self, atom: ase.atoms.Atoms) -> ase.atoms.Atoms:
        """Shift the center of mass to scaled position (x,y) = (0, 0).

        Args:
            atom: initial structure

        Returns: centralized structure

        """
        n_st = atom.copy()
        center = self._get_center_of_mass_periodic(atom)
        pos = (
            np.remainder(atom.get_scaled_positions() - center, [1, 1, 1])
            @ atom.cell
        )
        atoms = Atoms(
            cell=n_st.cell,
            numbers=n_st.numbers,
            positions=pos,
        )
        return atoms

    def _get_center_of_mass_periodic(
        self, atom: ase.atoms.Atoms
    ) -> np.ndarray:
        """Compute the center of mass in scaled
        coordinates with periodic wrapping.

        Uses circular mean for x and y; standard mean for z.

        Args:
            atom: ASE Atoms object.

        Returns:
            np.ndarray of scaled center-of-mass coordinates [cx, cy, cz].
        """
        L = np.array([1, 1, 1])
        x = atom.get_scaled_positions()
        theta = 2.0 * np.pi * x / L
        mass = atom.get_masses()
        mtheta = (
            np.arctan2(
                (-np.sin(theta) * np.expand_dims(mass, axis=1)).sum(axis=0)
                / len(theta),
                (-np.cos(theta) * np.expand_dims(mass, axis=1)).sum(axis=0)
                / len(theta),
            )
            + np.pi
        )
        center = L * mtheta / 2.0 / np.pi

        tmp = atom.get_center_of_mass(scaled=True)
        center[2] = tmp[2]
        return center

    def _get_translations(
        self, monomer_atoms: ase.atoms.Atoms, potential_tans: list
    ) -> tuple:
        """

        Args:
            monomer_atoms: possible monomers
            potential_tans: translational distances in z direction of monomers

        Returns: cyclic groups and monomers

        """
        cyclic_group, mono, sym_op = [], [], []
        invariant_op = SymmOp(
            [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0]]
        )
        for ii, monomer in enumerate(monomer_atoms):
            logging.debug("---Start detecting Item %d monomer" % (ii + 1))
            tran = potential_tans[ii]
            ind = int(np.round(1 / tran))

            if ind - 1 / tran > self._symprec:
                logging.error("Selecting wrong translational vector")
                continue

            if len(monomer) == len(self._primitive):
                logging.debug(
                    "The monomer is self._primitive, so this cyclic group is T"
                )
                cyclic_group.append("T")
                mono.append(monomer)
                sym_op.append(invariant_op)
            else:
                # detect rotation
                logging.debug(
                    "The scaled translational distance is %s " % (1 / ind)
                )
                logging.debug("Start detecting rotation")

                rotation, Q, tmp_sym_operations = self._detect_rotation(
                    monomer, tran * self._pure_trans, ind
                )
                if rotation:
                    logging.debug(
                        "Append rotational cyclic group, Q is 360/degree=%s"
                        % Q
                    )
                    cyclic_group.append(
                        "(C%s|T%s(%s))"
                        % (Q, ind, np.round(tran * self._pure_trans, 3))
                    )
                    mono.append(monomer)
                    tmp_sym_operations.insert(0, invariant_op)
                    sym_op.append(tmp_sym_operations)
                else:
                    logging.debug("None of the candidate rotations works")

                if (
                    ind == 2 and abs(tran - 0.5) < self._layer_symprec
                ):  # only 2 layer in primitive cell
                    # detect mirror
                    logging.debug(
                        "The scaled translational distance is 1/2,"
                        " start detecting mirror symmetry"
                    )

                    mirror, tmp_sym_operations = self._detect_mirror(
                        monomer, self._pure_trans / 2
                    )
                    if mirror:
                        logging.debug(
                            "Mirror plane exist, append to cyclic group"
                        )
                        cyclic_group.append(
                            "T'(%s)" % np.round(self._pure_trans / 2, 3)
                        )
                        mono.append(monomer)

                        sym_op.append([invariant_op, tmp_sym_operations])

                    else:
                        logging.debug("Mirror does not exist")
        return cyclic_group, mono, sym_op

    def _detect_possible_helical_angle(
        self, ind: int, monomer: ase.atoms.Atoms
    ) -> np.ndarray:
        """Compute candidate helical rotation angles between monomer layers.

        Args:
            ind: number of monomer layers in the primitive cell.
            monomer: ASE Atoms of the base monomer.

        Returns:
            np.ndarray of unique candidate angles in degrees.
        """
        scaled = monomer.get_scaled_positions()
        z = scaled[:, 2]

        min_z = np.min(z)
        min_mask = np.abs(z - min_z) < self._symprec
        points_start = (scaled[min_mask] @ self._primitive.cell)[:, :2]

        # find out the points on upper layer
        target_z = (min_z + 1 / ind) % 1.0

        prim_z = self._primitive.get_scaled_positions()[:, 2]
        itp = np.where(np.abs(prim_z - target_z) < self._symprec)[0]

        if len(itp) == 0:
            raise RuntimeError(f"No atoms found at target z={target_z:.4f}")

        point_end = (
            self._primitive.get_scaled_positions() @ self._primitive.cell
        )[itp][:, :2]
        center = ([0.5, 0.5, 0] @ self._primitive.cell)[:2]

        # calculate the angle
        arr = np.array(
            [
                angle_between_points(ps, center, pe)
                for ps in points_start
                for pe in point_end
            ]
        )

        angles = np.unique(np.round(arr, self._round_symprec))
        angles = angles[np.abs(angles) > self._symprec]

        return angles

    def _detect_rotation(
        self,
        monomer: ase.atoms.Atoms,
        tran: np.float64,
        ind: int,
    ) -> Tuple[bool, Union[Fraction, int], Optional[list]]:
        """Detect helical/rotational symmetry between
        monomer layers in a nanotube.

        For each candidate rotation angle, verifies that
        rotating the base monomer by `angle * layer` and
        translating by `tran * layer` maps onto existing
        atoms in the primitive cell. Tests both positive
        and negative rotation directions independently.

        Args:
            monomer: Base monomer (layer 0) extracted
                from the primitive cell.
            tran: Translational displacement (Cartesian,
                along z) per monomer layer.
            ind: Number of monomer layers contained in
                the primitive cell.

        Returns:
            A tuple of (found, Q, sym_ops) where:
                - found (bool): True if a valid rotational
                  symmetry is detected.
                - Q (Fraction | int): Rotational index,
                  defined as 360 / angle. Returns 1 if no
                  symmetry is found.
                - sym_ops (list[SymmOp] | None): Symmetry
                  operations (one per layer, for the valid
                  rotation direction). None if not found.

        Notes:
            - When ind == 1, no inter-layer rotation can
              be defined; returns False.
            - Positive and negative rotation directions
              are tested separately to avoid mixing
              inconsistent symmetry operations.
            - Q is computed via Fraction to avoid
              floating-point precision issues.
        """
        if ind == 1:
            return False, 1, None

        coords = self._primitive.get_scaled_positions() @ self._primitive.cell
        center = [0.5, 0.5, 0] @ self._primitive.cell
        pot_angle = self._detect_possible_helical_angle(ind, monomer)
        logging.debug("Candidate rotational degrees are: %s" % str(pot_angle))

        for test_ind in pot_angle:
            for sign, label in [(1, "pos"), (-1, "neg")]:
                sym_ops = []
                all_pass = True

                for layer in range(1, ind):
                    angle = sign * test_ind * layer
                    op = SymmOp.from_origin_axis_angle(
                        origin=center, axis=self._zaxis, angle=angle
                    )
                    op = SymmOp.from_rotation_and_translation(
                        op.rotation_matrix,
                        op.translation_vector + [0, 0, tran * layer],
                    )

                    layer_pass = all(
                        len(
                            tmp := find_in_coord_list(
                                coords,
                                op.operate(site.position),
                                self._symprec,
                            )
                        )
                        == 1
                        and self._primitive.numbers[tmp[0]] == site.number
                        for site in monomer
                    )

                    if not layer_pass:
                        all_pass = False
                        break
                    sym_ops.append(op)
                if all_pass:
                    Q = (
                        Fraction(360, int(round(test_ind))).limit_denominator()
                        if float(test_ind).is_integer()
                        else Fraction(360 / test_ind).limit_denominator(1000)
                    )
                    return True, Q, sym_ops

        return False, 1, None

    def _detect_mirror(
        self,
        monomer: ase.atoms.Atoms,
        tran: np.float64,
    ) -> Tuple[bool, Optional[SymmOp]]:
        """

        Args:
            monomer: monomer candidates
            tran: the translational distance of monomer candidates

        Returns: judge if the mirror symmetry exist

        """

        coords = self._primitive.positions
        diff_st_ind = np.array(
            [
                find_in_coord_list(coords, coord, self._symprec)
                # for coord in sites_pos
                for coord in monomer.positions
            ]
        )

        if diff_st_ind.ndim > 1:
            diff_st_ind = diff_st_ind.T[0]
        # diff_st: diff_st + monomer = primitive cell
        diff_st = self._primitive[
            np.setdiff1d(range(len(coords)), diff_st_ind)
        ]
        for itp1, itp2 in itertools.combinations_with_replacement(
            range(len(monomer)), 2
        ):
            s1, s2 = monomer[itp1], diff_st[itp2]

            if (
                s1.number == s2.number
                and (s1.position[2] + tran - s2.position[2]) < self._symprec
            ):
                normal = s1.position - s2.position
                normal[2] = 0
                op = SymmOp.reflection(
                    normal, origin=([0.5, 0.5, 0.5] @ self._primitive.cell)
                )
                op = SymmOp.from_rotation_and_translation(
                    op.rotation_matrix,
                    op.translation_vector + np.array([0, 0, tran]),
                )
                itp = []
                for site in monomer:
                    coord = op.operate(site.position)
                    tmp = find_in_coord_list(
                        diff_st.positions, coord, self._symprec
                    )
                    itp.append(
                        len(tmp) == 1
                        and diff_st.numbers[tmp[0]] == site.number
                    )
                if np.array(itp).all():
                    return True, op
        return False, None

    def _get_monomer_ind(
        self, z_round: np.ndarray, z_uniq: np.ndarray
    ) -> Tuple[List[np.ndarray], List[np.ndarray]]:
        monomer_ind = [np.where(z_round == val)[0] for val in z_uniq]

        monomer_ind_sum = []
        accumulated = np.array([], dtype=np.int32)
        for ind in monomer_ind:
            accumulated = np.sort(np.append(accumulated, ind))
            monomer_ind_sum.append(accumulated.copy())
        return monomer_ind, monomer_ind_sum

    def _potential_translation(self) -> tuple:
        """Generate the potential monomer and the scaled
        translational distance in z axis.

        Returns: possible monomers and translational distances

        """
        z = self._primitive.get_scaled_positions()[:, 2]
        z_round = np.round(z, self._round_symprec)

        z_uniq, counts = np.unique(z_round, return_counts=True)
        monomer_ind, monomer_ind_sum = self._get_monomer_ind(z_round, z_uniq)

        z_uniq_high_accurate = np.array(
            [z[tmp_ind].mean().item() for tmp_ind in monomer_ind]
        )

        potential_trans = np.append(
            (z_uniq_high_accurate - z_uniq_high_accurate[0])[1:], 1
        )

        translation, monomer = [], []
        for ii in range(len(z_uniq)):
            monomer_num = counts[: ii + 1].sum()
            # check the atomic number and layer number
            # of potential monomer; check the translational
            # distance corresponds to the layer numbers
            if (
                len(self._primitive) % monomer_num
                == 0  # atom number is multiple of monomer
                and len(z_uniq) % (ii + 1) == 0  #
                and abs(len(z_uniq) / (ii + 1) - 1 / potential_trans[ii])
                < self._layer_symprec
            ):
                if len(self._primitive) == monomer_num:
                    # if the monomer is the whole structure
                    monomer.append(self._primitive)
                    translation.append(1)
                else:
                    monomer.append(self._primitive[monomer_ind_sum[ii]])
                    translation.append(potential_trans[ii])
        return monomer, translation

    def _find_primitive(self) -> ase.atoms.Atoms:
        """Find the primitive cell of a 1D line-group
        structure along z-axis."""
        scaled_pos = self._atom.get_scaled_positions()
        x_y = scaled_pos[:, :2]
        z = scaled_pos[:, 2]

        # Find candidate matching atoms in xy-plane
        indices = find_in_coord_list(x_y[1:], x_y[0], atol=self._symprec)
        indices = indices + 1 if indices.size > 0 else np.array([], dtype=int)

        if len(indices) == 0:
            logging.debug("It's already a primitive cell.")
            return self._atom

        potential_z = (z[indices] - z[0]) % 1
        trans_z = []

        # Check which translation preserves the cell
        for dz in potential_z:
            shifted_pos = scaled_pos + np.array([0, 0, dz])
            pos1, num1 = refine_cell(
                shifted_pos, self._atom.numbers, symprec=self._round_symprec
            )
            pos2, num2 = refine_cell(
                scaled_pos, self._atom.numbers, symprec=self._round_symprec
            )
            if np.allclose(pos1, pos2, atol=self._symprec) and np.array_equal(
                num1, num2
            ):
                trans_z.append(dz)

        if not trans_z:
            logging.debug("No valid translation found. Already primitive.")
            return self._atom

        # Build new primitive cell
        pure_z = min(trans_z)
        atom_indices = np.where(z < pure_z - self._symprec)[0]
        cell = np.zeros((3, 3))
        cell[0] = self._atom.cell[0]
        cell[1] = self._atom.cell[1]
        cell[2] = [0, 0, pure_z * self._atom.cell[2, 2]]

        atom = Atoms(
            numbers=self._atom.numbers[atom_indices],
            positions=self._atom.positions[atom_indices],
            cell=cell,
        )
        logging.debug(f"Primitive cell found with z-translation {pure_z}.")
        return atom

    def _check_if_along_OZ(self, atom: ase.atoms.Atoms) -> bool:
        """Check whether the cell's periodic axis is aligned along z.

        Args:
            atom: ASE Atoms object.

        Returns:
            bool: True if the third cell vector is along z only.
        """
        if (
            np.isclose(atom.cell[2, :2], [0, 0]).all()
            and np.isclose(atom.cell[:2, 2], [0, 0]).all()
        ):
            return True
        else:
            return False

    def get_cyclic_group(self) -> tuple:
        """Return the detected cyclic group symbols and corresponding monomers.

        Returns:
            tuple of (cyclic_group, monomers).
        """
        return self.cyclic_group, self.monomers

    def get_cyclic_group_and_op(self) -> tuple:
        """Return cyclic group symbols, monomers, and symmetry operations.

        Returns:
            tuple of (cyclic_group, monomers, sym_operations).
        """
        return self.cyclic_group, self.monomers, self._sym_operations

    def get_generators(self) -> np.ndarray:
        """Return the generator affine matrix of the cyclic group.

        Returns:
            np.ndarray: 4x4 affine matrix of the generator operation.
        """
        if len(self._sym_operations) == 1:
            return self._sym_operations[0].affine_matrix
        else:
            op = SymmOp.from_rotation_and_translation(
                rotation_matrix=self._sym_operations[0][1].rotation_matrix,
                translation_vec=[
                    0,
                    0,
                    (
                        self._sym_operations[0][1].translation_vector
                        @ np.linalg.inv(self._atom.cell)
                        % 1
                    )[2],
                ],
            )
            return op.affine_matrix


def main():
    """CLI entry point for detecting the generalized translational group."""
    parser = argparse.ArgumentParser(
        description="Try to detect the generalized translational group of a line group structure"
    )
    parser.add_argument(
        "-p",
        "--POSCAR",
        help="path to the file from which coordinates will be read",
    )
    parser.add_argument(
        "-t",
        "--tolerance",
        default=1e-3,
        type=float,
        help="Tolerance for atomic positions",
    )
    parser.add_argument(
        "-o",
        "--enable_log",
        action="store_true",
        help="Enable the output of detection process",
    )
    args = parser.parse_args()

    flag_log = args.enable_log
    if flag_log:
        logging.basicConfig(
            level=logging.DEBUG,
            format="%(asctime)s - %(levelname)s - %(message)s",
        )

    st_name = args.POSCAR
    st = read(st_name)

    cyclic = CyclicGroupAnalyzer(st, tolerance=args.tolerance)
    cy, monomers = cyclic.get_cyclic_group()
    print("The generalized translational group symbol:")
    for ii, cg in enumerate(cy):
        print("Monomer %s:" % monomers[ii].symbols, cg)


if __name__ == "__main__":
    main()
