import numpy as np

from ase import Atoms
from ase.units import Bohr, Ry


def read_scf(filename):
    try:
        f = open(filename + '.scf', 'r')
        pip = f.readlines()
        ene = []
        for line in pip:
            if line[0:4] == ':ENE':
                ene.append(float(line[43:59]) * Ry)
        f.close()
        return ene
    except:
        return None

        
def read_struct(filename, ase=True):
    f = open(filename, 'r')
    pip = f.readlines()
    lattice = pip[1][0:3]
    nat = int(pip[1][27:30])
    cell = np.zeros(6)
    for i in range(6):
        cell[i] = float(pip[3][0 + i * 10:10 + i * 10])
    cell[0:3] = cell[0:3] * Bohr
    if lattice == 'P  ':
        lattice = 'P'
    elif lattice == 'H  ':
        lattice = 'P'
        cell[3:6] = [90.0, 90.0, 120.0]
    elif lattice == 'R  ':
        lattice = 'R'
    elif lattice == 'F  ':
        lattice = 'F'
    elif lattice == 'B  ':
        lattice = 'I'
    elif lattice == 'CXY':
        lattice = 'C'
    elif lattice == 'CXZ':
        lattice = 'B'
    elif lattice == 'CYZ':
        lattice = 'A'
    else:
        raise RuntimeError('TEST needed')
    pos = np.array([])
    atomtype = []
    rmt = []
    neq = np.zeros(nat)
    iline = 4
    indif = 0
    for iat in range(nat):
        indifini = indif
        if len(pos) == 0:
            pos = np.array([[float(pip[iline][12:22]),
                             float(pip[iline][25:35]),
                             float(pip[iline][38:48])]])
        else:
            pos = np.append(pos, np.array([[float(pip[iline][12:22]),
                                            float(pip[iline][25:35]),
                                            float(pip[iline][38:48])]]),
                            axis=0)
        indif += 1
        iline += 1
        neq[iat] = int(pip[iline][15:17])
        iline += 1
        for ieq in range(1, int(neq[iat])):
            pos = np.append(pos, np.array([[float(pip[iline][12:22]),
                                            float(pip[iline][25:35]),
                                            float(pip[iline][38:48])]]),
                            axis=0)
            indif += 1
            iline += 1
        for i in range(indif - indifini):
            atomtype.append(pip[iline][0:2].replace(' ', ''))
            rmt.append(float(pip[iline][43:48]))
        iline += 4
    if ase:
        cell2 = coorsys(cell)
        atoms = Atoms(atomtype, pos, pbc=True)
        atoms.set_cell(cell2, scale_atoms=True)
        cell2 = np.dot(c2p(lattice), cell2)
        if lattice == 'R':
            atoms.set_cell(cell2, scale_atoms=True)
        else:
            atoms.set_cell(cell2)
        return atoms
    else:
        return cell, lattice, pos, atomtype, rmt

        
def write_struct(filename, atoms2=None, rmt=None, lattice='P', zza=None):
    atoms = atoms2.copy()
    atoms.wrap()
    f = open(filename, 'w')
    f.write('ASE generated\n')
    nat = len(atoms)
    if rmt is None:
        rmt = [2.0] * nat
        f.write(lattice +
                '   LATTICE,NONEQUIV.ATOMS:%3i\nMODE OF CALC=RELA\n' % nat)
    cell = atoms.get_cell()
    metT = np.dot(cell, np.transpose(cell))
    cell2 = cellconst(metT)
    cell2[0:3] = cell2[0:3] / Bohr
    f.write(('%10.6f' * 6) % tuple(cell2) + '\n')
    if zza is None:
        zza = atoms.get_atomic_numbers()
    for ii in range(nat):
        f.write('ATOM %3i: ' % (ii + 1))
        pos = atoms.get_scaled_positions()[ii]
        f.write('X=%10.8f Y=%10.8f Z=%10.8f\n' % tuple(pos))
        f.write('          MULT= 1          ISPLIT= 1\n')
        zz = zza[ii]
        if zz > 71:
            ro = 0.000005
        elif zz > 36:
            ro = 0.00001
        elif zz > 18:
            ro = 0.00005
        else:
            ro = 0.0001
        f.write('%-10s NPT=%5i  R0=%9.8f RMT=%10.4f   Z:%10.5f\n' %
                (atoms.get_chemical_symbols()[ii], 781, ro, rmt[ii], zz))
        f.write('LOCAL ROT MATRIX:    %9.7f %9.7f %9.7f\n' % (1.0, 0.0, 0.0))
        f.write('                     %9.7f %9.7f %9.7f\n' % (0.0, 1.0, 0.0))
        f.write('                     %9.7f %9.7f %9.7f\n' % (0.0, 0.0, 1.0))
    f.write('   0\n')

    
def cellconst(metT):
    """ metT=np.dot(cell,cell.T) """
    aa = np.sqrt(metT[0, 0])
    bb = np.sqrt(metT[1, 1])
    cc = np.sqrt(metT[2, 2])
    gamma = np.arccos(metT[0, 1] / (aa * bb)) / np.pi * 180.0
    beta = np.arccos(metT[0, 2] / (aa * cc)) / np.pi * 180.0
    alpha = np.arccos(metT[1, 2] / (bb * cc)) / np.pi * 180.0
    return np.array([aa, bb, cc, alpha, beta, gamma])

    
def coorsys(latconst):
    a = latconst[0]
    b = latconst[1]
    c = latconst[2]
    cal = np.cos(latconst[3] * np.pi / 180.0)
    cbe = np.cos(latconst[4] * np.pi / 180.0)
    cga = np.cos(latconst[5] * np.pi / 180.0)
    sga = np.sin(latconst[5] * np.pi / 180.0)
    return np.array([[a, b * cga, c * cbe],
                     [0, b * sga, c * (cal - cbe * cga) / sga],
                     [0, 0, c * np.sqrt(1 - cal**2 - cbe**2 - cga**2 +
                                        2 * cal * cbe * cga) / sga]
                     ]).transpose()

    
def c2p(lattice):
    """ apply as eg. cell2 = np.dot(c2p('F'), cell)"""
    if lattice == 'P':
        cell = np.eye(3)
    elif lattice == 'F':
        cell = np.array([[0.0, 0.5, 0.5], [0.5, 0.0, 0.5], [0.5, 0.5, 0.0]])
    elif lattice == 'I':
        cell = np.array([[-0.5, 0.5, 0.5], [0.5, -0.5, 0.5], [0.5, 0.5, -0.5]])
    elif lattice == 'C':
        cell = np.array([[0.5, 0.5, 0.0], [0.5, -0.5, 0.0], [0.0, 0.0, -1.0]])
    elif lattice == 'B':
        cell = np.array([[0.5, 0.0, 0.5], [0.0, 1.0, 0.0], [0.5, 0.0, -0.5]])
    elif lattice == 'A':
        cell = np.array([[-1.0, 0.0, 0.0], [0.0, -0.5, 0.5], [0.0, 0.5, 0.5]])
    elif lattice == 'R':
        cell = np.array([[2.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0],
                         [-1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0],
                         [-1.0 / 3.0, -2.0 / 3.0, 1.0 / 3.0]])

    else:
        raise ValueError('lattice is ' + lattice + '!')
    return cell
