# 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 copy
import itertools
import logging
from decimal import Decimal
from typing import Any, List, Optional, Tuple, Union

import ase
import cvxpy as cp
import numpy as np
import scipy as sp
import scipy.interpolate
import scipy.sparse as ss
import scipy.spatial.distance
from ase import Atoms
from pymatgen.core.operations import SymmOp
from pymatgen.util.coord import find_in_coord_list

from pulgon_tools.Irreps_tables import line_group_sympy
from pulgon_tools.Irreps_tables_withparities import (
    line_group_sympy_withparities,
)


def e() -> np.ndarray:
    """
    Returns: identity matrix
    """
    mat = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
    return mat


def Cn(n: Union[int, float]) -> np.ndarray:
    """
    Args:
        n: rotate 2*pi/n

    Returns: rotation matrix
    """
    mat = np.array(
        [
            [np.cos(2 * np.pi / n), -np.sin(2 * np.pi / n), 0],
            [np.sin(2 * np.pi / n), np.cos(2 * np.pi / n), 0],
            [0, 0, 1],
        ]
    )
    return mat


def sigmaV() -> np.ndarray:
    """

    Returns: mirror symmetric matrix

    """
    mat = np.array([[1, 0, 0], [0, -1, 0], [0, 0, 1]])
    return mat


def sigmaH() -> np.ndarray:
    """

    Returns: mirror symmetric matrix about x-y plane

    """
    mat = np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]])
    return mat


def U() -> np.ndarray:
    """

    Returns: A symmetric matrix about the x-axis

    """
    mat = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
    return mat


def U_d(fid: Union[float, int]) -> np.ndarray:
    """

    Args:
        fid: the angle between symmetry axis d and
            axis x, d located in the x-y plane

    Returns: A symmetric matrix about the d-axis

    """
    mat = np.array(
        [
            [np.cos(2 * fid), np.sin(2 * fid), 0],
            [np.sin(2 * fid), -np.cos(2 * fid), 0],
            [0, 0, -1],
        ]
    )
    return mat


def S2n(n: Union[int, float]) -> np.ndarray:
    """
    Args:
        n: dihedral group, rotate 2*pi/n

    Returns: rotation and mirror matrix

    """
    mat = np.array(
        [
            [np.cos(np.pi / n), np.sin(np.pi / n), 0],
            [-np.sin(np.pi / n), np.cos(np.pi / n), 0],
            [0, 0, -1],
        ]
    )
    return mat


def sortrows(a: np.ndarray) -> np.ndarray:
    """
    :param a:
    :return: Compare each row in ascending order
    """
    return a[np.lexsort(np.rot90(a))]


def refine_cell(
    scale_pos: np.ndarray, numbers: np.ndarray, symprec: int = 4
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Refine fractional coordinates into [0,1) and remove duplicate atoms.

    Parameters
    ----------
    scale_pos : (N,3) or (3,) ndarray
        Fractional coordinates.
    numbers : (N,) ndarray
        Atomic species indices.
    symprec : int
        Decimal precision used to identify duplicates.

    Returns
    -------
    pos : ndarray
        Refined fractional coordinates.
    numbers : ndarray
        Corresponding atomic species.
    """

    scale_pos = np.asarray(scale_pos)

    # map coordinates to [0,1)
    scale_pos = np.mod(scale_pos, 1.0)

    # single atom case
    if scale_pos.ndim == 1:
        pos = np.round(scale_pos, symprec)
        return pos, numbers

    # round for numerical stability
    scale_pos = np.round(scale_pos, symprec)

    # remove duplicates while preserving order
    _, index = np.unique(scale_pos, axis=0, return_index=True)
    index = np.sort(index)

    pos = scale_pos[index]
    numbers = numbers[index]

    return pos, numbers


def frac_range(
    start: float,
    end: float,
    left: bool = True,
    right: bool = True,
    symprec: float = 0.01,
) -> list:
    """return the integer within the specified range

    Args:
        start: left boundary
        end: right boundary
        left: False means delete the left boundary
            element if it is an integer
        right: False means delete the right boundary
            element if it is an integer
        symprec: system precision

    Returns:

    """
    close = list(
        range(
            np.ceil(start).astype(np.int32), np.floor(end).astype(np.int32) + 1
        )
    )
    if left == False:
        if close[0] - start < symprec:
            close.pop(0)  # delete the left boundary
    if right == False:
        if close[-1] - end < symprec:
            close.pop()  # delete the right boundary
    return close


def decimal_places(x: float) -> int:
    """Return the number of decimal places of x
    using exact Decimal arithmetic."""
    d = Decimal(str(x))
    return abs(d.as_tuple().exponent)


def get_center_of_mass_periodic(atom: 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 find_axis_center_of_nanotube(atom: ase.atoms.Atoms) -> ase.atoms.Atoms:
    """remove the center of structure to (x,y):(0.5,0.5)
    Args:
        atom: initial structure

    Returns: centralized structure

    """
    n_st = atom.copy()
    center = get_center_of_mass_periodic(atom)

    pos = (
        # np.remainder(atom.get_scaled_positions() - center + 0.5, [1, 1, 1])
        np.remainder(
            atom.get_scaled_positions()
            # - [center[0],center[1],0] + [0.5, 0.5, 0],
            - center + [0.5, 0.5, 0],
            [1, 1, 1],
        )
        @ atom.cell
    )

    atoms = Atoms(
        cell=n_st.cell,
        numbers=n_st.numbers,
        positions=pos,
    )
    return atoms


def get_perms(
    atoms: Atoms,
    cyclic_group_ops: List[SymmOp],
    point_group_ops: List[SymmOp],
    symprec: float = 1e-2,
) -> Tuple[np.ndarray, List[SymmOp]]:
    """get the permutation table from symmetry operations

    Args:
        atoms:
        cyclic_group_ops:
        point_group_ops:
        symprec:

    Returns: permutation table
             rotation matrix (RM) = RM from point group @ RM from cyclic group
    """
    combs = list(itertools.product(point_group_ops, cyclic_group_ops))
    coords_car = atoms.positions
    coords_scaled = atoms.get_scaled_positions()
    coords_car_center = (
        atoms.get_scaled_positions() - [0.5, 0.5, 0.5]
    ) @ atoms.cell

    perms, rotation_matrix, translation_vector = [], [], []
    for ii, op in enumerate(combs):
        tmp_perm = np.ones((1, len(atoms.numbers)))[0]
        op1, op2 = op
        rotation_matrix.append(op1.rotation_matrix @ op2.rotation_matrix)
        translation_vector.append(
            op1.translation_vector + op2.translation_vector
        )

        for jj, site in enumerate(atoms):
            pos = (site.scaled_position - [0.5, 0.5, 0.5]) @ atoms.cell

            tmp = op1.operate(pos)
            idx1 = find_in_coord_list(coords_car_center, tmp, symprec)

            tmp1 = op2.operate(coords_car[idx1.item()])
            tmp1 = np.remainder(tmp1 @ np.linalg.inv(atoms.cell), [1, 1, 1])
            idx2 = find_in_coord_list(coords_scaled, tmp1, symprec)

            if idx2.size == 0:
                logging.ERROR("tolerance exceed while calculate perms")
            tmp_perm[jj] = idx2.item()
        perms.append(tmp_perm)
    perms_table, itp = np.unique(perms, axis=0, return_index=True)
    perms_table = perms_table.astype(np.int32)

    rotation_matrix = np.array(rotation_matrix)[itp]
    translation_vector = np.array(translation_vector)[itp]
    sym_operations = [
        SymmOp.from_rotation_and_translation(
            rotation_matrix[ii], translation_vector[ii]
        )
        for ii in range(len(itp))
    ]
    return perms_table, sym_operations


def get_perms_from_ops(
    atoms: Atoms,
    ops_sym: List[SymmOp],
    symprec: float = 1e-2,
    round: int = 4,
) -> np.ndarray:
    """get the permutation table from symmetry operations

    Args:
        atoms:
        symprec:

    Returns: permutation table
    """
    invcell = np.linalg.inv(atoms.cell)
    coords_center = atoms.positions - atoms.get_center_of_mass()
    coords_scaled_center = coords_center @ invcell
    coords_scaled_center[coords_center @ invcell >= 0.5] -= 1
    coords_scaled_center[coords_center @ invcell <= -0.5] += 1
    coords_center = coords_scaled_center @ atoms.cell

    perms = []
    for ii, op in enumerate(ops_sym):
        perms.append([])
        operated = op.operate_multi(coords_center)
        operated_scaled = operated @ invcell
        # much safer than remainder - no floating point issues
        operated_scaled[operated @ invcell >= 0.5] -= 1
        operated_scaled[operated @ invcell <= -0.5] += 1
        operated = operated_scaled @ atoms.cell
        for aid in range(len(coords_center)):
            equivalents = np.where(
                np.linalg.norm(operated[aid] - coords_center, axis=1) < symprec
            )[0]
            if len(equivalents) == 1:
                for eq in equivalents:
                    perms[-1].append(eq)
            else:
                print(ii, aid)
                raise ValueError
    perms_table = np.array(perms).astype(np.int32)
    return perms_table


def get_matrices(
    atoms: Atoms,
    ops_sym: List[SymmOp],
    symprec: float = 1e-5,
) -> List[np.ndarray]:
    """Build 3N x 3N representation matrices from
    symmetry operations and permutations.

    Args:
        atoms: ASE Atoms object.
        ops_sym: list of SymmOp symmetry operations.
        symprec: tolerance for permutation detection.

    Returns:
        list of np.ndarray: representation matrices of shape (3N, 3N).
    """
    perms_table = get_perms_from_ops(atoms, ops_sym, symprec=symprec)

    natoms = len(atoms.numbers)
    matrices = []
    for ii, perm in enumerate(perms_table):
        matrix = np.zeros((3 * natoms, 3 * natoms))
        for jj in range(natoms):
            idx = perm[jj]
            matrix[3 * idx : 3 * (idx + 1), 3 * jj : 3 * (jj + 1)] = ops_sym[
                ii
            ].rotation_matrix.copy()
        matrices.append(matrix)
    return matrices


def get_matrices_withPhase(
    atoms: Atoms,
    ops_sym: List[SymmOp],
    qpoint: float,
    symprec: float = 1e-3,
) -> List[np.ndarray]:
    """Build complex 3N x 3N representation matrices
    with Bloch phase factors.

    Args:
        atoms: ASE Atoms object.
        ops_sym: list of SymmOp symmetry operations.
        qpoint: q-point value along z for the phase factor exp(i*q*f).
        symprec: tolerance for permutation detection.

    Returns:
        list of complex np.ndarray of shape (3N, 3N).
    """
    perms_table = get_perms_from_ops(atoms, ops_sym, symprec=symprec)

    natoms = len(atoms.numbers)
    matrices = []
    for ii, perm in enumerate(perms_table):
        matrix = np.zeros((3 * natoms, 3 * natoms)).astype(np.complex128)
        phasefacter = ops_sym[ii].translation_vector[2]

        for jj in range(natoms):
            idx = perm[jj]
            matrix[3 * idx : 3 * (idx + 1), 3 * jj : 3 * (jj + 1)] = ops_sym[
                ii
            ].rotation_matrix.copy() * np.exp(1j * qpoint * phasefacter)
        matrices.append(matrix)
    return matrices


def affine_matrix_op(
    af1: np.ndarray, af2: np.ndarray, symprec: float = 1e-8
) -> np.ndarray:
    """Definition of group multiplication

    Args:
        af1: group element 1
        af2: group element 2

    Returns:

    """
    af1 = np.round(af1 / symprec).astype(int) * symprec
    af2 = np.round(af2 / symprec).astype(int) * symprec

    ro = af2[:3, :3] @ af1[:3, :3]
    tran = np.remainder(af2[:3, 3] + af2[:3, :3] @ af1[:3, 3], [1, 1, 1])

    af = np.eye(4)
    af[:3, :3] = ro
    af[:3, 3] = tran

    return af


def dimino_affine_matrix_and_character(
    generators: np.ndarray,
    character: np.ndarray,
    symec: float = 0.001,
) -> Tuple[np.ndarray, np.ndarray]:
    """

    Args:
        generators: the generators of point group
        symec: system precision

    Returns: all the group elements and correspond character

    """
    e_in = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])

    if character[0].ndim == 0:
        L_chara = np.array([np.complex128(1)])
        C_chara = np.array([np.complex128(1)])
    else:
        L_chara = np.array([np.complex128(1) * np.eye(character[0].shape[0])])
        C_chara = np.array([np.complex128(1) * np.eye(character[0].shape[0])])

    G = generators
    g, g1 = G[0].copy(), G[0].copy()
    g_chara, g1_chara = character[0].copy(), character[0].copy()
    L = np.array([e_in])
    while not ((g - e_in) < symec).all():
        L = np.vstack((L, [g]))
        L_chara = np.vstack((L_chara, [g_chara]))

        g = affine_matrix_op(g, g1)
        g_chara = np.dot(g_chara, g1_chara)

    for ii in range(len(G)):
        C = np.array([e_in])
        L1 = L.copy()
        L1_chara = L_chara.copy()

        more = True
        while more:
            more = False
            for jj, g in enumerate(list(C)):
                g_chara = C_chara[jj]

                for kk, ss in enumerate(G[: ii + 1]):
                    ss_chara = character[kk]
                    sg = affine_matrix_op(ss, g)
                    sg_chara = np.dot(ss_chara, g_chara)

                    itp = (abs((sg - L).sum(axis=1).sum(axis=1)) < 0.001).any()
                    if not itp:
                        if C.ndim == 3:
                            C = np.vstack((C, [sg]))
                            C_chara = np.vstack((C_chara, [sg_chara]))
                        else:
                            C = np.array((C, sg))
                            C_chara = np.array((C_chara, sg_chara))

                        if L.ndim == 3:
                            L = np.vstack(
                                (
                                    L,
                                    np.array(
                                        [affine_matrix_op(sg, t) for t in L1]
                                    ),
                                )
                            )
                            tmp = np.array(
                                [
                                    np.dot(sg_chara, t_chara)
                                    for t_chara in L1_chara
                                ]
                            )
                            L_chara = np.vstack((L_chara, tmp))
                        else:
                            L = np.array(
                                L,
                                np.array(
                                    [affine_matrix_op(sg, t) for t in L1]
                                ),
                            )
                        more = True
    L_chara_trace = []
    for lc in L_chara:
        if lc.ndim == 1:
            L_chara_trace.append(lc)
        else:
            L_chara_trace.append(np.trace(lc))
    return L, np.array(L_chara_trace)


def brute_force_generate_group(generators: np.ndarray, symec: float = 0.01):
    """Generate all group elements by brute-force
    multiplication of generators.

    Args:
        generators: array of 4x4 affine matrices as group generators.
        symec: tolerance for comparing matrix elements.

    Returns:
        np.ndarray: array of all group element affine matrices.
    """
    e_in = np.eye(4)
    G = generators
    L = np.array([e_in])
    while True:
        numL_old = len(L)
        for g in L:
            for h in G:
                gh = affine_matrix_op(g, h)
                judge = (abs(L - gh) < symec).all(axis=1).all(axis=1).any()
                if not judge:
                    L = np.concatenate((L, [gh]), axis=0)
        numL_new = len(L)
        if numL_old == numL_new:
            break
    return L


def brute_force_generate_group_subsequent(
    generators: np.ndarray, symec: float = 0.01
) -> Tuple[np.ndarray, List[List[int]]]:
    """generate all the group elements by brute force algorithm

    Args:
        generators: The most basic elements used to generate complete groups
        symec: tolerance

    Returns:
        L: all group elements
        L_seq: the multiplication order of the generators
    """

    e_in = np.eye(4)
    G = generators
    L = np.array([e_in])
    L_seq = [[0]]
    while True:
        numL_old = len(L)
        for ii, g in enumerate(L):
            tmp_seq1 = L_seq[ii]
            for jj, h in enumerate(G):
                gh = affine_matrix_op(g, h, symprec=1e-5)
                tmp_seq2 = tmp_seq1 + [jj + 1]
                judge = (abs(L - gh) < symec).all(axis=1).all(axis=1).any()
                if not judge:
                    L = np.concatenate((L, [gh]), axis=0)
                    L_seq.append(tmp_seq2)
        numL_new = len(L)
        if numL_old == numL_new:
            break
    return L, L_seq


def dimino(generators: np.ndarray, symec: int = 4) -> np.ndarray:
    """
    Generate all group elements from a set of generators
    using Dimino's algorithm.

    Dimino's algorithm performs coset enumeration: starting
    from the cyclic group generated by the first generator,
    it iteratively adds new cosets for each subsequent
    generator until the group is closed.

    Reference:
        Dimino, A. (1971). An algorithm for generating
        all group elements from a set of generators.
        See also: computational group theory texts.

    Args:
        generators: Array of shape (n_gen, 3, 3) containing
            the generator matrices of the point group. Each
            matrix is a 3x3 real orthogonal matrix
            representing a symmetry operation.
        symec: Number of decimal places for rounding when
            comparing matrices (to handle floating-point
            near-equalities). Default is 4.

    Returns:
        np.ndarray of shape (|G|, 3, 3): all group elements (symmetry
        operation matrices), including the identity element E.

    Example:
        >>> from pulgon_tools.utils import Cn, sigmaV, e
        >>> gens = np.array([Cn(6), sigmaV()])
        >>> all_ops = dimino(gens, symec=4)
        >>> all_ops.shape  # C6v has 12 elements
        (12, 3, 3)
    """
    G = generators

    # --- Step 1: Build the cyclic subgroup generated by G[0] ---
    # Start from E, repeatedly multiply by G[0] until we return to E.
    g = copy.deepcopy(G[0])
    g1 = copy.deepcopy(G[0])
    L = np.array(
        [e()]
    )  # L holds all known group elements; initialise with identity

    while not (np.round(g, symec) == e()).all():
        L = np.vstack((L, [g]))
        g = np.dot(g, g1)

    L = np.round(L, symec)

    # --- Step 2: Extend L by coset enumeration over each generator ---
    # For each generator G[i], try to find elements s*g
    # (s in G[0..i], g in coset C) that are not yet in L,
    # and if so, add the entire left-coset of that element.
    for ii in range(len(G)):
        C = np.array(
            [e()]
        )  # Current coset representatives; start with identity
        L1 = copy.deepcopy(
            L
        )  # Snapshot of L at the start of this generator's pass
        more = True

        while more:
            more = False
            for g in C:  # iterate over current coset representatives
                for s in G[: ii + 1]:  # multiply by all generators seen so far
                    sg = np.round(np.dot(s, g), symec)

                    # Check whether sg is already in L
                    already_in_L = (
                        (np.round(sg, symec) == L)
                        .all(axis=1)
                        .all(axis=1)
                        .any()
                    )

                    if not already_in_L:
                        # sg is a new coset representative
                        # Add it to C
                        C = np.vstack((C, [sg]))

                        # Add the entire left-coset  sg * L1  to L
                        new_coset = np.array([np.dot(sg, t) for t in L1])
                        L = np.vstack((L, new_coset))

                        more = True  # L grew, so we need another pass

    return L


def dimino_affine_matrix(
    generators: np.ndarray, symec: float = 0.01
) -> np.ndarray:
    """

    Args:
        generators: the generators of point group
        symec: system precision

    Returns: all the group elements and correspond character

    """
    e_in = np.eye(4)

    G = generators
    g, g1 = G[0].copy(), G[0].copy()
    L = np.array([e_in])
    while not ((g - e_in) < symec).all():
        L = np.vstack((L, [g]))
        g = affine_matrix_op(g, g1)
    for ii in range(len(G)):
        C = np.array([e_in])
        L1 = L.copy()
        more = True
        while more:
            more = False
            for g in list(C):
                for ss in G[: ii + 1]:
                    sg = affine_matrix_op(ss, g)
                    itp = (abs((sg - L).sum(axis=1).sum(axis=1)) < symec).any()
                    if not itp:
                        if C.ndim == 3:
                            C = np.vstack((C, [sg]))
                        if L.ndim == 3:
                            L = np.vstack(
                                (
                                    L,
                                    np.array(
                                        [affine_matrix_op(sg, t) for t in L1]
                                    ),
                                )
                            )
                        more = True
    L = np.unique(L, axis=0)
    return L


def dimino_affine_matrix_and_subsequent(
    generators: np.ndarray, symec: float = 0.001
) -> Tuple[np.ndarray, List[List[int]]]:
    """

    Args:
        generators: the generators of point group
        symec: system precision

    Returns: all the group elements and correspond character

    """
    e_in = np.eye(4)

    G = generators
    g, g1 = G[0].copy(), G[0].copy()
    g_subs, g1_subs = [1], [1]
    L = np.array([e_in])
    L_subs = [[0]]

    while not ((g - e_in) < symec).all():
        L = np.vstack((L, [g]))
        L_subs.append(g_subs.copy())

        g = affine_matrix_op(g, g1)
        g_subs = g_subs + g1_subs
    for ii in range(len(G)):
        C = np.array([e_in])
        C_subs = [[0]]

        L1 = L.copy()
        L1_subs = L_subs.copy()
        more = True
        while more:
            more = False
            for jj, g in enumerate(list(C)):
                g_subs = C_subs[jj]
                for kk, ss in enumerate(G[: ii + 1]):
                    ss_subs = [kk + 1]

                    sg = affine_matrix_op(ss, g)
                    sg_subs = ss_subs + g_subs

                    itp = (abs((sg - L).sum(axis=1).sum(axis=1)) < 0.001).any()
                    if not itp:
                        if C.ndim == 3:
                            C = np.vstack((C, [sg]))
                            C_subs.append(sg_subs)
                        else:
                            C = np.array((C, sg))
                            C_subs = C_subs.append(sg_subs)
                        if L.ndim == 3:
                            L = np.vstack(
                                (
                                    L,
                                    np.array(
                                        [affine_matrix_op(sg, t) for t in L1]
                                    ),
                                )
                            )
                            tmp = [sg_subs + t_subs for t_subs in L1_subs]
                            L_subs = L_subs + tmp

                        else:
                            L = np.array(
                                L,
                                np.array(
                                    [affine_matrix_op(sg, t) for t in L1]
                                ),
                            )
                            tmp = [sg_subs + t_subs for t_subs in L1_subs]
                            L_subs = L_subs + tmp
                        more = True
    return L, L_subs


def get_character(
    DictParams: dict, symprec: float = 1e-8
) -> Tuple[list, list, list]:
    """Compute symbolic representation matrices for the line group.

    Args:
        DictParams: dict with keys 'qpoints', 'nrot', 'order', 'family', 'a'.
        symprec: tolerance for symbolic evaluation.

    Returns:
        tuple of (characters, paras_values, paras_symbols).
    """
    characters, paras_values, paras_symbols = line_group_sympy(
        DictParams, symprec
    )
    return characters, paras_values, paras_symbols


def get_character_withparities(
    DictParams: dict, symprec: float = 1e-8
) -> Tuple[List[Any], List[Any], List[Any]]:
    """Compute symbolic representation matrices with parity operations.

    Args:
        DictParams: dict with keys 'qpoints', 'nrot', 'order', 'family', 'a'.
        symprec: tolerance for symbolic evaluation.

    Returns:
        tuple of (characters, paras_values, paras_symbols).
    """
    characters, paras_values, paras_symbols = line_group_sympy_withparities(
        DictParams, symprec
    )
    return characters, paras_values, paras_symbols


def get_character_num(
    DictParams: dict, symprec: float = 1e-8
) -> Tuple[np.ndarray, list, list]:
    """Compute numerical character table from representation matrices.

    Takes the trace of each representation matrix to obtain characters.

    Args:
        DictParams: dict with keys 'qpoints', 'nrot', 'order', 'family', 'a'.
        symprec: tolerance for symbolic evaluation.

    Returns:
        tuple of (characters, paras_values, paras_symbols) where
        characters is a numpy array of traces.
    """
    representation_mat, paras_values, paras_symbols = line_group_sympy(
        DictParams, symprec
    )

    characters = []
    for ii, rep_mat in enumerate(representation_mat):  # loop IR
        if rep_mat.ndim == 1:
            characters.append(rep_mat)
        else:
            characters.append(np.trace(rep_mat, axis1=1, axis2=2))
    characters = np.array(characters)
    return characters, paras_values, paras_symbols


def get_character_num_withparities(
    DictParams: dict, symprec: float = 1e-8
) -> Tuple[np.ndarray, list, list]:
    """Compute numerical character table with parity operations.

    Takes the trace of each representation matrix (including parities)
    to obtain characters.

    Args:
        DictParams: dict with keys 'qpoints', 'nrot', 'order', 'family', 'a'.
        symprec: tolerance for symbolic evaluation.

    Returns:
        tuple of (characters, paras_values, paras_symbols) where
        characters is a numpy array of traces.
    """
    (
        representation_mat,
        paras_values,
        paras_symbols,
    ) = line_group_sympy_withparities(DictParams, symprec)
    characters = []
    for ii, rep_mat in enumerate(representation_mat):  # loop IR
        if rep_mat.ndim == 1:
            characters.append(rep_mat)
        else:
            characters.append(np.trace(rep_mat, axis1=1, axis2=2))
    characters = np.array(characters)
    return characters, paras_values, paras_symbols


def get_sym_constrains_matrices_M(
    ops: np.ndarray,
    permutations: np.ndarray,
    diminsion: int = 3,
) -> ss.spmatrix:
    """M K = 0

    :param ops:
    :param permutations:
    :param diminsion: spatial dimension (default 3)
    :return:

    """

    if permutations.ndim == 2:
        natom = len(permutations[0])
    elif permutations.ndim == 1:
        natom = len(permutations)
        permutations = np.array([permutations])
    else:
        logging.ERROR("error for permutations' ndim")

    if len(ops.shape) == 2:
        ops = np.array([ops])

    size1 = diminsion**2
    I = np.eye(size1)
    M = []

    idx1 = np.repeat(np.arange(natom), natom)
    idx2 = np.tile(np.arange(natom), natom)

    tmp1 = (idx1 * natom + idx2) * size1
    tmp2 = (idx1 * natom + idx2 + 1) * size1
    tmp3 = np.linspace(tmp1, tmp2, size1 + 1).astype(np.int64)[:-1, :].T
    for ii, op in enumerate(ops):
        print("now run in %s operation" % ii)
        perm = permutations[ii]
        C = np.einsum(
            "ij,kl->ikjl",
            op[:diminsion, :diminsion],
            op[:diminsion, :diminsion],
        ).reshape(size1, size1)
        x = ss.csc_matrix((size1 * natom**2, size1 * natom**2))
        if (perm == np.arange(natom)).all():
            M.append(x)
            continue

        pidx1 = perm[idx1]
        pidx2 = perm[idx2]

        ptmp1 = (pidx1 * natom + pidx2) * size1
        ptmp2 = (pidx1 * natom + pidx2 + 1) * size1
        ptmp3 = np.linspace(ptmp1, ptmp2, size1 + 1).astype(np.int64)[:-1, :].T

        itp1 = np.repeat(tmp3, size1, axis=1)
        itp2 = np.tile(tmp3, (1, size1))
        pitp2 = np.tile(ptmp3, (1, size1))
        xl = x.tolil()
        xl[itp1, itp2] = C.flatten()
        xl[itp1, pitp2] = -I.flatten()  #
        M.append(xl)

    M = scipy.sparse.vstack((M))
    return M


def _calc_dists(
    atoms: Atoms, tolerance: float = 1e-3
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Return the distances between atoms in the supercell, their
    degeneracies and the associated displacements along OZ.
    """
    MIN_DELTA = -1
    MAX_DELTA = 1
    positions = atoms.positions
    cell = atoms.cell
    n_satoms = positions.shape[0]
    d2s = np.empty((MAX_DELTA - MIN_DELTA + 1, n_satoms, n_satoms))

    # TODO: This could not be enough and eventually we
    # should do a proper check for the shortest distance.
    for j, j_c in enumerate(range(MIN_DELTA, MAX_DELTA + 1)):
        shifted_positions = positions + (j_c * cell[2, :])[np.newaxis, :]
        d2s[j, :, :] = sp.spatial.distance.cdist(
            positions, shifted_positions, "sqeuclidean"
        )
        # d2s[j, :, :] = sp.spatial.distance.cdist(
        #     positions, shifted_positions)
    d2min = d2s.min(axis=0)
    dmin = np.sqrt(d2min)
    degenerate = np.abs(d2s - d2min) < tolerance
    nequi = degenerate.sum(axis=0, dtype=int)

    maxequi = nequi.max()
    shifts = np.empty((n_satoms, n_satoms, maxequi))
    sorting = np.argsort(np.logical_not(degenerate), axis=0)
    shifts = np.transpose(sorting[:maxequi, :, :], (1, 2, 0)).astype(np.intc)
    shifts = np.asarray(range(MIN_DELTA, MAX_DELTA + 1))[shifts]
    return (dmin, nequi, shifts)


def get_continum_constrains_matrices_M_for_conpact_fc(
    phonon,
) -> ss.coo_array:
    """Build constraint matrix enforcing acoustic and rotational sum rules.

    Constructs a sparse matrix M such that M @ K = 0 encodes translational
    (acoustic), Born-Huang, and Huang invariance conditions on compact
    force constants.

    Args:
        phonon: Phonopy object with force_constants and supercell set.

    Returns:
        scipy sparse COO matrix encoding all continuum constraints.
    """
    IFC = phonon.force_constants.copy()
    scell = phonon.supercell

    symbols = scell.symbols
    cell = scell.cell
    positions = (
        (scell.scaled_positions + np.asarray([0.5, 0.5, 0.0])[np.newaxis, :])
        % 1.0
    ) @ cell
    ase_atoms = ase.Atoms(symbols, positions, cell=cell, pbc=True)
    dists, degeneracy, shifts = _calc_dists(ase_atoms)
    n_satoms = len(symbols)

    # %%
    average_delta = np.zeros((n_satoms, n_satoms, 3))
    for i in range(n_satoms):
        for j in range(n_satoms):
            n_elements = degeneracy[i, j]
            for i_d in range(n_elements):
                average_delta[i, j, :] += (
                    positions[j, :]
                    - positions[i, :]
                    + shifts[i, j, i_d] * cell[2, :]
                )
            average_delta[i, j, :] /= n_elements

    average_pos = scell.scaled_positions @ cell

    # %%
    average_products = np.zeros((n_satoms, n_satoms, 3, 3))
    for i in range(n_satoms):
        for j in range(n_satoms):
            n_elements = degeneracy[i, j]
            for i_d in range(n_elements):
                delta = (
                    positions[j, :]
                    - positions[i, :]
                    + shifts[i, j, i_d] * cell[2, :]
                )
                average_products[i, j, :, :] += np.outer(delta, delta)
            average_products[i, j, :, :] /= n_elements

    M1 = ss.coo_array(([]))

    # Append the acoustic sum rules originating from translations.
    n_atoms, n_satoms = IFC.shape[:2]
    n_rows = M1.shape[0]
    rows = M1.row.tolist()
    cols = M1.col.tolist()
    data = M1.data.tolist()

    # translational sum rules
    for i in range(n_atoms):
        for alpha in range(3):
            for beta in range(3):
                for j in range(n_satoms):
                    rows.append(n_rows)
                    cols.append(
                        np.ravel_multi_index((i, j, alpha, beta), IFC.shape)
                    )
                    data.append(1.0)
                n_rows += 1

    # The same but for rotations (Born-Huang).
    positions = phonon.supercell.positions
    for i in range(n_atoms):
        for alpha in range(3):
            for beta in range(3):
                for gamma in range(3):
                    for j in range(n_satoms):
                        r_ij = average_delta[phonon.primitive.p2s_map[i], j]
                        rows.append(n_rows)
                        cols.append(
                            np.ravel_multi_index(
                                (i, j, alpha, beta), IFC.shape
                            )
                        )
                        data.append(r_ij[gamma])
                        rows.append(n_rows)
                        cols.append(
                            np.ravel_multi_index(
                                (i, j, alpha, gamma), IFC.shape
                            )
                        )
                        data.append(-r_ij[beta])
                    n_rows += 1

    # And fthe Huang invariances, also for rotation.
    for alpha in range(3):
        for beta in range(3):
            for gamma in range(3):
                for delta in range(3):
                    for i in range(n_atoms):
                        for j in range(n_satoms):
                            products = average_products[
                                phonon.primitive.p2s_map[i], j
                            ]
                            rows.append(n_rows)
                            cols.append(
                                np.ravel_multi_index(
                                    (i, j, alpha, beta), IFC.shape
                                )
                            )
                            data.append(products[gamma, delta])
                            rows.append(n_rows)
                            cols.append(
                                np.ravel_multi_index(
                                    (i, j, gamma, delta), IFC.shape
                                )
                            )
                            data.append(-products[alpha, beta])
                    n_rows += 1

    # Make sure the IFC matrix is symmetric.
    for i in range(n_atoms):
        for j in range(n_atoms):
            for alpha in range(3):
                for beta in range(3):
                    rows.append(n_rows)
                    cols.append(
                        np.ravel_multi_index(
                            (i, phonon.primitive.p2s_map[j], alpha, beta),
                            IFC.shape,
                        )
                    )

                    data.append(1.0)
                    rows.append(n_rows)
                    cols.append(
                        np.ravel_multi_index(
                            (j, phonon.primitive.p2s_map[i], beta, alpha),
                            IFC.shape,
                        )
                    )
                    data.append(-1.0)
                    n_rows += 1

    # Rebuild the sparse matrix.
    M1 = ss.coo_array((data, (rows, cols)), shape=(n_rows, IFC.size))
    return M1


def get_IFCSYM_from_cvxpy_M(M: ss.spmatrix, IFC: np.ndarray) -> np.ndarray:
    """Symmetrize force constants by solving a constrained optimization.

    Minimizes ||x - IFC||^2 subject to M @ x = 0 using CVXPY.

    Args:
        M: sparse constraint matrix encoding symmetry/sum rules.
        IFC: original force constant array.

    Returns:
        np.ndarray: symmetrized force constants with the same shape as IFC.
    """
    flat_IFCs = IFC.ravel()
    x = cp.Variable(IFC.size)
    cost = cp.sum_squares(x - flat_IFCs)
    prob = cp.Problem(cp.Minimize(cost), [M @ x == 0])
    prob.solve()
    IFC_sym = x.value.reshape(IFC.shape)
    return IFC_sym


def get_symbols_from_ops(
    ops_sym: List[np.ndarray],
) -> List[str]:
    """
    Get the symbols of point group operations from the rotation matrix.

    Args:
        ops_sym: A list of rotation matrix of symmetry operations.

    Returns:
        symbols: A list of point group symbols.

    Notes:
        The point group operations are as follows:
        E: Identity
        sigmaH: Reflection in the horizontal mirror plane.
        sigmaV: Reflection in the vertical mirror plane.
        U: Reflection in the plane perpendicular to the axis.
        Cn: Rotation by 2 * pi / n about the axis.
        Sn: Rotation by 2 * pi / n about the axis
            followed by reflection in the plane
            perpendicular to the axis.
    """
    symbols = []
    for op in ops_sym:
        op_ro = op[:3, :3]

        val1, _ = np.linalg.eig(op_ro)
        itp1 = np.isclose(val1, 1, atol=1e-4)
        itp2 = np.isclose(val1, -1, atol=1e-4)

        if np.logical_or(itp1, itp2).all():  # test if it is a reflection h,v,U
            if itp1.sum() == 3:
                symbols.append("E")
            elif itp1.sum() == 2 and itp2.sum() == 1:
                tmp = np.where(itp2)[0].item()
                if tmp == 2:
                    symbols.append("sigmaH")
                else:
                    symbols.append("sigmaV")
            elif itp1.sum() == 1 and itp2.sum() == 2 and itp2[2]:
                symbols.append("U")
            else:
                logging.ERROR("The Eigenvalues are very strange")
        else:  # test if it is a rotation Cn
            tmp = 1 / (-np.log(val1) / 2 / np.pi * 1j)[:2]

            if np.isclose(
                tmp[0].real, abs(tmp[1].real), atol=1e-4
            ) and np.isclose(val1[2].real, 1, atol=1e-4):
                num = int(np.round(tmp[0].real))
                symbols.append("C%d" % num)
            elif np.isclose(
                tmp[0].real, abs(tmp[1].real), atol=1e-4
            ) and np.isclose(val1[2].real, -1, atol=1e-4):
                num = int(np.round(tmp[0].real))
                symbols.append("S%d" % num)
            else:
                logging.ERROR("The Eigenvalues are very strange")
    return symbols


def divide_irreps(
    vec: np.ndarray,
    adapted: np.ndarray,
    dimensions: List[int],
) -> np.ndarray:
    """
    Project vectors into IR space according to the adapted basis.

    Parameters
    ----------
    vec : array of shape (n,m), where n is the number
        of vectors and m is the dimension
        The vector to be divided
    adapted : array of shape (m,m). The basis vectors are arranged in columns.
        The adapted basis
    dimensions : list of int. Sum(dimensions) = m.
        The dimension of each irrep

    Returns
    -------
    means : array of shape (n,k), where k is the number of irreps
        The projected length of each irrep
    """
    tmp1 = vec @ adapted.conj()
    start = 0
    means, vectors = [], []
    for im, dim in enumerate(dimensions):
        end = start + dim
        if vec.ndim == 1:
            means.append((np.abs(tmp1[start:end]) ** 2).sum())
        else:
            means.append((np.abs(tmp1[:, start:end]) ** 2).sum(axis=1))
        start = copy.copy(end)
    means = np.array(means)
    if means.ndim > 1:
        means = means.T
    return np.array(means)


def angle_between_points(
    A: Union[list, np.ndarray],
    B: Union[list, np.ndarray],
    C: Union[list, np.ndarray],
) -> float:
    """
    Calculate angle between three points. Point b is the vertex

    Parameters
    ----------
    A : The first point coordinates
    B : The second point coordinates
    C : The third point coordinates

    Returns
    -------
    angle : float
        The angle between points A, B and C in degrees.
    """
    BA = np.array(A) - np.array(B)
    BC = np.array(C) - np.array(B)

    cosine_angle = np.dot(BA, BC) / (np.linalg.norm(BA) * np.linalg.norm(BC))
    cosine_angle = np.clip(cosine_angle, -1.0, 1.0)
    angle_rad = np.arccos(cosine_angle)
    return np.degrees(angle_rad)
