import os
from ase.geometry import get_layers
from ase.constraints import FixAtoms
from ase.io.espresso import construct_namelist
from xespresso import Espresso
from xespresso.xio import build_atomic_species_str
from ase.dft.bandgap import bandgap
import pickle
import multiprocessing
import numpy as np

# ====================================================


def get_nbnd(atoms=None, scale=1.2, pseudopotentials={}, nspin=1, input_data={}):
    input_parameters = construct_namelist(input_data)
    atomic_species_str, species_info, total_valence = build_atomic_species_str(
        atoms, input_parameters, pseudopotentials
    )
    nbnd = int(total_valence / (2.0 / nspin))
    nbnd_scale = int(nbnd * scale)
    print(
        " total valence: %s\n nbnd: %s\n scaled nbnd: %s"
        % (total_valence, nbnd, nbnd_scale)
    )
    return nbnd


def merge_slab(slab1, slab2, index=2):
    """ """
    slab2.cell[index] = slab1.cell[index]
    slab2.set_cell(slab1.cell, scale_atoms=True)
    slab1 = slab1 + slab2
    slab1.cell[2][2] = max(slab1.positions[:, 2] + 15)
    slab1.wrap()
    return slab1


def qeinp(
    calculation,
    ecutwfc=30,
    mixing_beta=0.5,
    conv_thr=1.0e-8,
    edir=False,
    input_ntyp={},
    atoms=None,
):
    #
    inp = {
        # control
        "calculation": calculation,
        "max_seconds": 78000,
        "verbosity": "high",
        "tprnfor": True,
        # system
        "ecutwfc": ecutwfc,
        "ecutrho": ecutwfc * 8,
        "occupations": "smearing",
        "degauss": 0.01,
        "input_ntyp": input_ntyp,
        # electrons
        "mixing_beta": mixing_beta,
        "conv_thr": conv_thr,
        "electron_maxstep": 400,
    }
    #
    if edir:
        inp.update(dipole_correction(atoms, edir))

    return inp


def dipole_correction(atoms, edir=3):
    inp = {
        "dipfield": True,
        "tefield": True,
        "edir": edir,
        "eamp": 0.001,
        "eopreg": 0.05,
    }
    emaxpos = (
        max(atoms.positions[:, edir - 1] + atoms.cell[edir - 1][edir - 1])
        / 2.0
        / atoms.cell[edir - 1][edir - 1]
    )
    inp["emaxpos"] = round(emaxpos, 2)
    return inp


# tools


def build_oer(atoms):
    """ """
    from ase.atoms import Atoms

    # ----------------------------------
    ooh = Atoms("O2H", positions=[[0, 0, 0], [1.4, 0, 0], [1.4, 0, 1.0]])
    ooh.rotate("z", np.pi / 4)
    mols = {
        "o": Atoms("O"),
        "oh": Atoms("OH", positions=[[0, 0, 0], [0, 0, 1.0]]),
        "ooh": ooh,
    }
    #
    jobs = {}
    maxz = max(atoms.positions[:, 2])
    indm = [atom.index for atom in atoms if atom.z > maxz - 1.0 and atom.symbol != "O"][
        0
    ]
    for job, mol in mols.items():
        # print(job, mol)
        ads = mol.copy()
        natoms = atoms.copy()
        ads.translate(atoms[indm].position - ads[0].position + [0, 0, 1.9])
        natoms = natoms + ads
        jobs[job] = natoms
    return jobs


def fix_layers(atoms, miller=(0, 0, 1), tol=1.0, n=[0, 4]):
    """ """
    layers = get_layers(atoms, miller, tol)[0]
    index = [j for j in range(len(atoms)) if layers[j] in range(n[0], n[1])]
    constraint = FixAtoms(indices=index)
    atoms.set_constraint(constraint)
    return atoms


def mypool(jobs, func, showInfo=False):
    """ """
    from random import random
    from time import sleep

    pool = multiprocessing.Pool(len(jobs))
    results = []
    images = []
    for job, atoms in jobs.items():
        if showInfo:
            print(job, len(atoms), atoms)
        sleep(random() * 2)
        results.append(pool.apply_async(func, (job, atoms)))
    for r in results:
        r.get()
    pool.close()
    pool.join()


def dwubelix(updates=[]):
    file = "pw err out dos pdos projwfc int xyz path a.xml txt png"
    print("Downloading.....")
    cwd = os.getcwd()
    for update in updates:
        os.chdir(update)
        os.system("dwubelix-sc.py %s" % file)
        os.chdir(cwd)
    print("Finished")


def ana(dire, calc):
    atoms = calc.results["atoms"]
    results = [dire, atoms, atoms.cell, atoms.positions]
    for prop in ["energy", "forces", "stress", "magmoms"]:
        if prop in calc.results:
            prop = calc.results[prop]
        else:
            prop = None
        results.append(prop)
    return results


# tools


def summary(updates=[], prefix="datas"):
    import pandas as pd

    columns = ["label", "atoms", "cell", "positions", "energy", "forces", "stress"]
    file = "%s.pickle" % prefix
    db = "%s.db" % prefix
    if os.path.exists(file):
        with open(file, "rb") as f:
            datas, df = pickle.load(f)
    else:
        datas = {}
        df = pd.DataFrame(columns=columns)
    calc = Espresso()
    print("Reading.....")
    for update in updates:
        cwd = os.getcwd()
        for i, j, y in os.walk(update):
            output = is_espresso(i)
            if output:
                os.chdir(i)
                print("=" * 30)
                print("Reading dire:", i)
                calc.directory = cwd + "/" + i
                calc.prefix = output[0:-4]
                try:
                    calc.results = {}
                    calc.read_results()
                    datas[i] = calc.results
                    atoms = calc.results["atoms"]
                    atoms.write(os.path.join(calc.directory, "%s.cif" % calc.prefix))
                except Exception as e:
                    print("error: %s \n" % e)
            os.chdir(cwd)
    with open(file, "wb") as f:
        pickle.dump([datas, df], f)
    print("Finished")


def is_espresso(path):
    """
    check espresso
    """
    dirs = os.listdir(path)
    # print(dirs)
    # flag = True

    for qefile in [".pwi"]:
        flag = False
        for file in dirs:
            if qefile in file:
                return file
        if not flag:
            return False
    # return flag


def grep_valence_configuration(pseudopotential):
    """
    Given a UPF pseudopotential file, find the valence configuration.

    Valence configuration:
    nl pn  l   occ       Rcut    Rcut US       E pseu
    3S  1  0  2.00      0.700      1.200    -6.910117

    """
    orbitals = ["S", "P", "D", "F"]
    valence = {}
    with open(pseudopotential) as psfile:
        lines = psfile.readlines()
        for i in range(len(lines)):
            if "valence configuration:" in lines[i].lower():
                j = i + 2
                ob = lines[j].split()[0]
                while ob[1] in orbitals:
                    valence[ob] = lines[j].split()[3]
                    j += 1
                    ob = lines[j].split()[0]
                return valence
    if not valence:
        raise ValueError("Valence configuration missing in {}".format(pseudopotential))
