import textwrap
from pathlib import Path
from subprocess import Popen


def writeIncar(directory, **kwargs):
    """VASP输入文件
    :kwargs 字典形式的vasp参数，忽略大小写
    """
    path = Path(directory)
    INCAR = path / "INCAR" if path.name != "INCAR" else path
    INCAR.parent.mkdir(parents=True, exist_ok=True)
    with open(INCAR, "w+") as fd:
        for k, v in kwargs.items():
            if isinstance(v, bool):
                if v is True:
                    kwargs[k] = ".TRUE."
                elif v is False:
                    kwargs[k] = ".FALSE."
            if isinstance(v, (list, tuple)):
                v = [str(i) for i in v]
                kwargs[k] = " ".join(v)
            if type(v) is int or type(v) is float:
                kwargs[k] = str(v)
            fd.write(f"{k.upper().ljust(8)} = {v}\n")


def writeKpoints(directory, mode="Monkhorst-Pack", subdivisions=None, kpaths=None, specialKpoints=None, density=0,
                 vaspkit=None):
    """VASP KPOINTS Writer
    :param specialKpoints: Special Kpoints
    :param vaspkit: generate KPOINTS by vaspkit
    :param directory: Write Path
    :param mode: Monkhorst-Pack | BandStructure
    :param subdivisions: subdivisions N_1, N_2 and N_3 along recipr. latt. vectors
    :param kpaths: BandStructure KPath
    :param density: Number of k-points
    Example:
    For structure optimizition:
        writeKpoints(
            directory=".",
            mode="Monkhorst-Pack",
            subdivisions=(6,6,6),
            density=10
        )
        You will get a KPOINTS file like this:
            KPOINTS Generated by HTSCT
            10
            Gamma
            6 1 1
            0 0 0
    For bandstructure:
        writeKpoints(
            directory=".",
            mode="BandStructure",
            kpaths=["GAMMA", "X", "GAMMA"],
            specialKpoints={"GAMMA": (0.0, 0.0, 0.0), "X": (0.5, 0.0, 0.0)},
            density=20
        )
        You will get a KPOINTS file like this:
            K-PATH Generated by HTSCT
               20
            Line-Mode
            Reciprocal
               0.0000000000   0.0000000000   0.0000000000      GAMMA
               0.5000000000   0.0000000000   0.0000000000      X

               0.5000000000   0.0000000000   0.0000000000      X
               0.0000000000   0.0000000000   0.0000000000      GAMMA
    """
    path = Path(directory)
    KPOINTS = path / "KPOINTS" if path.name != "KPOINTS" else path
    KPOINTS.parent.mkdir(parents=True, exist_ok=True)
    if vaspkit:
        if vaspkit == "1D":
            Popen("vaspkit -task 301\n", cwd=directory).wait()
        elif vaspkit == "2D":
            Popen("vaspkit -task 302\n", cwd=directory).wait()
        elif vaspkit == "3D":
            Popen("vaspkit -task 303\n", cwd=directory).wait()
        else:
            raise Exception("vaspkit can only be set to 1D|2D|3D")
        (directory / "KPATH.in").replace(directory / "KPOINTS")
        return

    if mode == "Monkhorst-Pack":
        subdivisions = map(lambda x: str(x), subdivisions)
        content = textwrap.dedent(f"""
        KPOINTS Generated by HTSCT
        {density}
        Gamma
        {" ".join(subdivisions)}
        0 0 0
        """).lstrip()
        with open(KPOINTS, "w+") as fd:
            fd.write(content)
    if mode == "BandStructure":
        _kpaths = ""
        for i in range(len(kpaths)):
            if i < len(kpaths) - 1:
                k1 = kpaths[i]
                k2 = kpaths[i + 1]
                v1 = map(lambda x: str(x).ljust(12, "0").rjust(15), specialKpoints[k1])
                v2 = map(lambda x: str(x).ljust(12, "0").rjust(15), specialKpoints[k2])
                _kpaths += f"{''.join(v1)}{' ' * 6}{k1}\n{''.join(v2)}{' ' * 6}{k2}\n\n"
        content = textwrap.dedent(f"""
        K-PATH Generated by HTSCT
           {density}
        Line-Mode
        Reciprocal
        """).lstrip() + _kpaths
        with open(KPOINTS, "w+") as fd:
            fd.write(content)


def writeOptcell(directory, a=None, b=None, c=None):
    """
    Help to write the vasp OPTCELL file.
    :param directory: Write Path
    :param a:a axis
    :param b:b axis
    :param c:c axis
    """
    default = ("0", "0", "0")
    path = Path(directory)
    OPTCELL = path / "OPTCELL" if path.name != "OPTCELL" else path
    OPTCELL.parent.mkdir(parents=True, exist_ok=True)
    _a = list(map(lambda x: str(x), a)) if a else default
    _b = list(map(lambda x: str(x), b)) if b else default
    _c = list(map(lambda x: str(x), c)) if c else default
    content = f"{' '.join(_a)}\n{' '.join(_b)}\n{' '.join(_c)}"
    with open(OPTCELL, "w+") as fd:
        fd.write(content)


def writePotcar(directory):
    Popen("vaspkit -task 103\n", cwd=directory, shell=True).wait()  # 生成POTCAR


if __name__ == '__main__':
    writeKpoints(
        directory=".",
        subdivisions=(1, 1, 1)
    )
