from numpy.linalg import norm
import numpy as np

from .mag_atom import ELEMENT_TO_NUMBER,ELEMENT_TABLE
# ************************ function read_poscar ****************************
# 
# > Input POSCAR, can with MAGMOM   x  y  z  mz  my  mz
# > (lattice, positions, numbers, mag) 
# > (lattice, positions, numbers, mag) = read_poscar() 
# 
# ***********************************************************************
def read_poscar(file_name = 'POSCAR'): 
    lattice  = np.zeros((3, 3))
    deg_of_points = []
    numbers = []
    mag = []
    # file_name = 'POSCAR_test'
    with open(file_name, 'r') as f:
        f_lines = f.readlines()
        for cnt, line in enumerate(f_lines):
            # cnt = 0,1 are meaningless
            # cnt = 2,3,4 show the basis vectors
            if cnt in [2, 3, 4]:
                line = line.strip().split()
                lattice[cnt-2, :] = [float(c) for c in line]
            if cnt  == 5:
                elements = line.strip().split()
            if cnt  == 6:
                line = line.strip().split()
                deg_of_points = [int(c) for c in line]
                num_of_points = sum(deg_of_points)
                num_of_types = np.size(deg_of_points)
                # get the numbers of cell using atomic numbers
                for m in range(num_of_types):
                    element_symbol = elements[m]
                    atomic_number = ELEMENT_TO_NUMBER.get(element_symbol)
                    if atomic_number is None:
                        raise ValueError(f'Unknown element symbol: {element_symbol}')
                    for n in range(deg_of_points[m]):
                        numbers.append(atomic_number)
                positions = np.zeros((num_of_points, 3))
            if cnt > 7 and cnt < (7 + num_of_points + 1):
                line = line.strip().split()
                if len(line) == 3:
                    positions[cnt-8, :] = [float(c) for c in line]
                    mag.append([0, 0, 0])
                elif len(line) == 6:
                    positions[cnt-8, :] = [float(c) for c in line[:3]]
                    mag.append([float(c) for c in line[3:]])
                else:
                    # print('line', cnt+1 ,'not valid POSCAR input !!!')
                    raise ValueError('Invalid input, please check POSCAR line ' + str(cnt+1) )
        cell = (lattice, positions, numbers, np.array(mag)) 
    return cell



# ************************ function write_poscar ****************************
# 
# > input cell = (lattice, positions, numbers, mag) 
# > write a POSCAR file using the cell 
# 
# ***********************************************************************
def write_poscar(cell, file_name='SPOSCAR'):

    lattice, positions, numbers, mag = cell
    positions = np.array(positions)
    mag = np.array(mag)

    # get unique atomic numbers and their counts
    unique_numbers = []
    for n in numbers:
        if n not in unique_numbers:
            unique_numbers.append(n)
    
    # Convert atomic numbers to element symbols
    elements = []
    for atomic_num in unique_numbers:
        element_symbol = ELEMENT_TABLE.get(atomic_num)
        if element_symbol is None:
            raise ValueError(f'Unknown atomic number: {atomic_num}')
        elements.append(element_symbol)

    # Count atoms of each type
    element_counts = [np.sum(numbers == atomic_num) for atomic_num in unique_numbers]

    # write POSCAR
    with open(file_name, 'w') as f:
        f.write("Generated by write_poscar\n")
        f.write("1.0\n")
        for vec in lattice:
            f.write("  {:.16f}  {:>.16f}  {:>.16f}\n".format(*vec))
        f.write("  " + "  ".join(t for t in elements) + "\n")
        f.write("  " + "  ".join(str(c) for c in element_counts) + "\n")
        f.write("Direc\n")

        for i in range(len(positions)):
            pos_str = "  {: .16f}  {: .16f}  {: .16f}".format(*positions[i])
            if mag is not None and len(mag) > 0:
                mag_str = "  {: .6f}  {: .6f}  {: .6f}".format(*mag[i])
                if norm(mag[i]) > 1e-5:
                    f.write(pos_str + mag_str + "\n")
                else:
                    f.write(pos_str + "\n")
            else:
                f.write(pos_str + "\n")