#
# Authors: MN
#
import numpy as np
import spgrep
import spglib
from yambopy.dbs.latticedb import YamboLatticeDB
from yambopy.tools.citations import citation

class Crystal_Symmetries:
    @citation("A. Togo et al., Spglib, Sci. Technol. Adv. Mater. Meth. 4 (2024)")
    def __init__(self, latdb, magnetic_moments=None, tol=1e-5):
        """
        Class to handle  crystal symmetries
        #
        Given a lattice database, finds all symmetries of a  crystal.
        Works irrespective of whether the SAVE has symmetries or not.
        #
        Parameters:
        -----------
        latdb : object
            Lattice database object with required attributes:
        magnetic_moments : array (natom) or (natom,3) or None, optional
            magnetic_moments of each atom (default: None)
            In collinear case, natom is sufficient, but in non-collinear 
            case provide (mx, my, mz) for each atom.
        tol : float, optional
            Symmetry tolerance (default: 1e-5)
        Members:
        --------
        dataset : dict
            The full symmetry dataset from spglib
        spacegroup_type : dict
            Space group type information
        rotations : ndarray
            Rotation matrices in Cartesian coordinates (n_sym, 3, 3)
        translations : ndarray
            Translation vectors in Cartesian coordinates (n_sym, 3)
        international_symbol : str
            International symbol of the space group
        hall_symbol : str
            Hall symbol of the space group
        wyckoffs : list
            Wyckoff letters for each atom
        pointgroup : str
            Point group symbol
        spacegroup_number : int
            International space group number
        """
        # Get lattice and atomic information
        lattice = latdb.lat
        numbers = np.rint(latdb.atomic_numbers).astype(int)
        lat_inv = np.linalg.inv(lattice)
        positions = latdb.car_atomic_positions @ lat_inv
        positions = positions - np.floor(positions)
        #
        self.magnetic = not (magnetic_moments is None)
        #
        # Get symmetry dataset
        if self.magnetic :
            cell = (lattice, positions, numbers, magnetic_moments)
            print("Currently Magnetic systems not supported.")
            exit()
        else :
            cell = (lattice, positions, numbers)
            self.dataset = spglib.get_symmetry_dataset(cell, symprec=tol)
            self.spacegroup_type = spglib.get_spacegroup_type(self.dataset.hall_number)
        #
        # Convert rotation and fractional translations to Cartesian units
        self.rotations = lattice.T[None,:,:] @ self.dataset.rotations @ lat_inv.T[None,:,:]
        self.translations = self.dataset.translations @ lattice
        #
        self.international_symbol = self.dataset.international
        self.hall_symbol = self.dataset.hall
        self.wyckoffs = self.dataset.wyckoffs
        self.pointgroup = self.dataset.pointgroup
        self.pointgroup_schoenflies = self.spacegroup_type.pointgroup_schoenflies
        self.spacegroup_number = self.dataset.number
    #
    #
    def write_symmetry_info_to_file(self, filename="sym_ops.dat"):
        """Write symmetries in cart units to file."""
        sym_matrices = self.rotations
        trans_vectors = self.translations
        # 
        with open(filename, "w") as f:
            # Write basic info
            f.write("# Basic symmetry information generated by Yambopy using spglib.\n")
            f.write("# International symbol : %s\n" %(str(self.international_symbol)))
            f.write("# Hall symbol          : %s\n" %(str(self.hall_symbol)))
            f.write("# Wyckoffs             : %s\n" %(str(self.wyckoffs)))
            f.write("# Pointgroup           : %s (%s)\n" %(str(self.pointgroup),str(self.pointgroup_schoenflies)))
            f.write("# Spacegroup           : %s\n" %(str(self.spacegroup_number)))
            #
            f.write("# \n")
            f.write("# Coset representatives \n")
            f.write("# \n")
            #
            f.write("# index         symmetry matrix (cart units).         fractional translation (cart units)\n")
            f.write("# -------------------------------------------------------------------------------------\n")

            for i, (S, v) in enumerate(zip(sym_matrices, trans_vectors), start=1):

                # First matrix row printed on same line as index
                row0 = " ".join(f"{S[0, c]:10.6f}" for c in range(3))
                line = f"{i:>{4}d}" + " " * (12 - 4) + row0
                f.write(line + "\n")

                # Second row includes translation vector
                row1 = " ".join(f"{S[1, c]:10.6f}" for c in range(3))
                vec  = " ".join(f"{v[k]:10.6f}" for k in range(3))
                line = " " * 12 + row1
                line += " " * (52 - len(line)) + vec
                f.write(line + "\n")

                # Third row, matrix only
                row2 = " ".join(f"{S[2, c]:10.6f}" for c in range(3))
                line = " " * 12 + row2
                f.write(line + "\n")

                f.write("#-------------------------------------------------------------------------------------\n")


## Test
if __name__ == "__main__":
    import os
    np.set_printoptions(suppress=True)
    latdb = YamboLatticeDB.from_db_file(os.path.join('.', 'SAVE'))
    symm = Crystal_Symmetries(latdb,tol=1e-4)
    print(symm.wyckoffs)
    print(symm.pointgroup_schoenflies)

