try:
    import numpy as np
except ImportError as e:
    import sys
    sys.stderr.write(f"An error occurred while importing numpy: {str(e)}\n")
    del sys

try:
    import re
except ImportError as e:
    import sys
    sys.stderr.write(f"An error occurred while importing re: {str(e)}\n")
    del sys

@staticmethod
def read_file(file_location:str=None, strip=True):

    if file_location is None:
        raise ValueError("(!) File location is not set.")
    
    try:
        with open(file_location, 'r') as file:
            if strip:
                for line in file:
                    yield line.strip()  # strip() elimina espacios y saltos de línea al principio y al final
            else:
                for line in file:
                    yield line  # strip() elimina espacios y saltos de línea al principio y al final
    except FileNotFoundError:
        raise FileNotFoundError(f"(!) File not found at {file_location}")
    except IOError:
        raise IOError(f"(!) Error reading file at {file_location}")
    except Exception as e:
        print(f"Error inesperado: {e}")

@staticmethod
def is_number(s):
    if s is None:
        return False
    try:
        float(s)
        return True
    except (ValueError, TypeError):
        return False

class POSCAR:
    def __init__(self, file_location:str=None, name:str=None, **kwargs):
        self._atomCountByType = None

    @property
    def atomCountByType(self):
        return self._atomCountByType

    @atomCountByType.setter
    def atomCountByType(self, value):
        """Allow external or internal code to overwrite the cached counts."""
        self._atomCountByType = np.asarray(value, dtype=int)

    def group_elements_and_positions(self, atomLabelsList:list=None, atomPositions:list=None, regroup=False):
        atomLabelsList = atomLabelsList if atomLabelsList is not None else getattr(self, 'atomLabelsList', None)
        atomPositions = atomPositions if atomPositions is not None else getattr(self, 'atomPositions', None)
        
        if atomLabelsList is None:
            return False

        if regroup:
            element_indices = {}
            for i, label in enumerate(atomLabelsList):
                if label not in element_indices:
                    element_indices[label] = []
                element_indices[label].append(i)

            atomLabelsList_new = []
            atomPositions_new = []
            atomPositions_frac_new = []
            atomicConstraints_new = []
            
            uniqueAtomLabels_new = list(element_indices.keys())  # preserve order
            
            has_positions = atomPositions is not None
            has_frac = hasattr(self, 'atomPositions_fractional') and self.atomPositions_fractional is not None
            has_constraints = hasattr(self, 'atomicConstraints') and self.atomicConstraints is not None

            for label in uniqueAtomLabels_new:
                indices = element_indices[label]
                atomLabelsList_new.extend([label] * len(indices))
                
                if has_positions:
                    atomPositions_new.extend([atomPositions[idx] for idx in indices])
                if has_frac:
                    atomPositions_frac_new.extend([self.atomPositions_fractional[idx] for idx in indices])
                if has_constraints:
                    atomicConstraints_new.extend([self.atomicConstraints[idx] for idx in indices])

            self.atomLabelsList = np.array(atomLabelsList_new, dtype=object)
            if has_positions:
                self.atomPositions = np.array(atomPositions_new)
            if has_frac:
                self.atomPositions_fractional = np.array(atomPositions_frac_new)
            if has_constraints:
                self.atomicConstraints = np.array(atomicConstraints_new, dtype=int)
            
            self.uniqueAtomLabels = uniqueAtomLabels_new
        else:
            self.atomLabelsList = np.array(atomLabelsList, dtype=object)
            if atomPositions is not None:
                self.atomPositions = np.array(atomPositions)
            if hasattr(self, 'uniqueAtomLabels') and self.uniqueAtomLabels is not None:
                pass # keep existing
            else:
                self.uniqueAtomLabels = list(dict.fromkeys(atomLabelsList))  # keep first appearance order

        unique_labels, count = np.unique(self.atomLabelsList, return_counts=True)
        # We must preserve the order of self.uniqueAtomLabels
        ordered_counts = []
        if hasattr(self, 'uniqueAtomLabels') and self.uniqueAtomLabels is not None:
            label_to_count = dict(zip(unique_labels, count))
            for ul in self.uniqueAtomLabels:
                ordered_counts.append(label_to_count.get(ul, 0))
            self.atomCountByType = np.array(ordered_counts)
        else:
            self.atomCountByType = count

        return True


    def export_as_POSCAR(self, file_location:str=None, v:bool=False) -> bool:
        if file_location is None:
            if hasattr(self, 'file_location') and isinstance(self.file_location, str):
                file_location = self.file_location + '_export.POSCAR'
            else:
                file_location = 'POSCAR'

        self.group_elements_and_positions(regroup=True)

        with open(file_location, 'w') as file:
            # Comentario inicial
            comment = getattr(self, 'comment', 'POSCAR : JML code')
            file.write(f'{comment}\n')

            # Factor de escala
            scale = getattr(self, 'scaleFactor', [1.0])
            if isinstance(scale, (list, tuple, np.ndarray)):
                file.write(f"{' '.join(map(str, scale))}\n")
            else:
                file.write(f"{scale}\n")

            # Vectores de la celda unitaria
            for lv in getattr(self, 'latticeVectors', []):
                file.write('  {:>18.15f}  {:>18.15f}  {:>18.15f}\n'.format(*lv))

            # Tipos de átomos y sus números
            labels = getattr(self, 'uniqueAtomLabels', None)
            if labels is not None:
                file.write('  ' + '  '.join(labels) + '\n')
            counts = getattr(self, 'atomCountByType', [])
            file.write('  ' + '  '.join(map(str, counts)) + '\n')
            
            # Opción para dinámica selectiva (opcional)
            selective = getattr(self, 'selectiveDynamics', False)
            if selective:
                file.write('Selective dynamics\n')
                
            # Tipo de coordenadas (Direct o Cartesian)
            coord_type = getattr(self, 'atomCoordinateType', 'cartesian')
            is_cartesian = coord_type[0].upper() in ['C', 'K'] if coord_type else True
            aCT = 'Cartesian' if is_cartesian else 'Direct'
            file.write(f'{aCT}\n')

            # Coordenadas atómicas y sus restricciones
            positions = getattr(self, 'atomPositions', []) if is_cartesian else getattr(self, 'atomPositions_fractional', getattr(self, 'atomPositions', []))
            
            for i, atom in enumerate(positions):
                coords = '  '.join(['{:>18.15f}'.format(n) for n in atom])
                constr = ''
                if selective and hasattr(self, 'atomicConstraints') and self.atomicConstraints is not None:
                    if i < len(self.atomicConstraints):
                        constr = '  ' + '  '.join(['T' if n else 'F' for n in self.atomicConstraints[i]])
                file.write(f'  {coords}{constr}\n')

            # Comentario final (opcional)
            diff_positions = getattr(self, 'dynamical_eigenvector_diff', None)
            if diff_positions is not None:
                diff_pos_to_write = diff_positions if is_cartesian else getattr(self, 'dynamical_eigenvector_diff_fractional', diff_positions)
                if diff_pos_to_write is not None:
                    file.write('Comment_line\n')
                    for i, atom in enumerate(diff_pos_to_write):
                        coords = '  '.join(['{:>18.15f}'.format(n) for n in atom])
                        file.write(f'  {coords}\n')
        
        return True
                
    def read_POSCAR(self, file_location:str=None):
        file_location = file_location if type(file_location) == str else self.file_location
        # Read lines and filter out empty ones
        lines = [n for n in read_file(file_location) if n.strip() != '']
        
        self.comment = lines[0].strip()
        self.scaleFactor = list(map(float, lines[1].strip().split()))
        
        # Reading lattice vectors
        self.latticeVectors = np.array([list(map(float, line.strip().split())) for line in lines[2:5]])
        
        # Species names (optional)
        species_line = lines[5].strip().split()
        if is_number(species_line[0]):
            self.uniqueAtomLabels = None
            offset = 0
        else:
            self.uniqueAtomLabels = species_line
            offset = 1
  
        # Ions per species
        self.atomCountByType = np.array(list(map(int, lines[5+offset].strip().split())))
        self.atomCount = int(np.sum(self.atomCountByType))
        
        # Selective dynamics (optional)
        type_line = lines[6+offset].strip()
        if len(type_line) > 0 and not is_number(type_line[0]):
            if type_line[0].upper() == 'S':
                self.selectiveDynamics = True
                offset += 1
            else:
                self.selectiveDynamics = False
        else:
            self.selectiveDynamics = False
        
        # atomic coordinated system
        coord_sys_line = lines[6+offset].strip()
        if len(coord_sys_line) > 0 and coord_sys_line[0].upper() in ['C', 'K']:
            self.atomCoordinateType = 'cartesian'
        else:
            self.atomCoordinateType = 'direct'

        # Ion positions and constraints
        positions = []
        constraints = []
        for line in lines[7+offset : 7+offset+self.atomCount]:
            parts = line.strip().split()
            # Coordinates
            positions.append(list(map(float, parts[:3])))
            # Constraints
            if self.selectiveDynamics:
                # Expect T or F at index 3, 4, 5. Default to T if missing
                c = [parts[i].upper() == 'T' if i < len(parts) else True for i in range(3, 6)]
                constraints.append(c)

        if self.atomCoordinateType == 'cartesian':
            self.atomPositions = np.array(positions)
        else:
            self.atomPositions_fractional = np.array(positions)

        if self.selectiveDynamics:
            self.atomicConstraints = np.array(constraints, dtype=int)
        else:
            self.atomicConstraints = None

    def read_CONTCAR(self, file_location:str=None):
        self.read_POSCAR(file_location)

