"""velph command line tool / velph-phelel."""

from __future__ import annotations

import pathlib
from typing import Optional

import click
import tomli

import phelel
from phelel.cui.phelel_script import finalize_phelel
from phelel.velph.cli import cmd_root
from phelel.velph.cli.phelel.differentiate import run_derivatives
from phelel.velph.cli.phelel.generate import write_supercell_input_files
from phelel.velph.cli.phelel.init import run_init
from phelel.velph.cli.phelel.phonopy import create_phonopy_yaml
from phelel.velph.utils.vasp import CutoffToFFTMesh


@cmd_root.group("phelel")
@click.help_option("-h", "--help")
def cmd_phelel():
    """Choose supercell options."""
    pass


#
# velph phelel init
#
@cmd_phelel.command("init")
@click.argument("toml_filename", nargs=1, type=click.Path(), default="velph.toml")
@click.option(
    "--dir-name",
    "dir_name",
    type=str,
    default="phelel",
    help=(
        'Used for backward compatibility, for which set "supercell". '
        '(dir_name: str, default="phelel")'
    ),
)
@click.help_option("-h", "--help")
def cmd_init(toml_filename: str, dir_name: str):
    """Generate displacements and write phelel_disp.yaml."""
    with open(toml_filename, "rb") as f:
        toml_dict = tomli.load(f)

    phe = run_init(toml_dict)

    phelel_yaml_filename = pathlib.Path(f"{dir_name}/phelel_disp.yaml")
    phelel_yaml_filename.parent.mkdir(parents=True, exist_ok=True)
    finalize_phelel(
        phe,
        displacements_mode=True,
        filename=phelel_yaml_filename,
        sys_exit_after_finalize=False,
    )

    click.echo(f'"{phelel_yaml_filename}" was generated. ')
    click.echo('VASP input files will be generated by "velph phelel generate".')


#
# velph phelel generate
#
@cmd_phelel.command("generate")
@click.argument("toml_filename", nargs=1, type=click.Path(), default="velph.toml")
@click.option(
    "--dir-name",
    "dir_name",
    type=str,
    default="phelel",
    help=(
        'Used for backward compatibility, for which set "supercell". '
        '(phelel_dirname: str, default="phelel")'
    ),
)
@click.help_option("-h", "--help")
def cmd_generate(toml_filename: str, dir_name: str):
    """Generate phelel input files."""
    if not pathlib.Path("POTCAR").exists():
        click.echo('"POTCAR" not found in current directory.')

    yaml_filename = pathlib.Path(f"{dir_name}/phelel_disp.yaml")
    write_supercell_input_files(
        pathlib.Path(toml_filename), yaml_filename, dir_name=dir_name
    )


#
# velph supercell differentiate
#
@cmd_phelel.command("differentiate")
@click.argument("toml_filename", nargs=1, type=click.Path(), default="velph.toml")
@click.option(
    "--dir-name",
    "dir_name",
    type=str,
    default="phelel",
    help=(
        'Used for backward compatibility, for which set "supercell". '
        '(phelel_dirname: str, default="phelel")'
    ),
)
@click.option(
    "--encut",
    nargs=1,
    type=float,
    default=None,
    help=(
        "Cutoff energy corresponding to FFT mesh of local potential grid. "
        "(encut: float, default=None)"
    ),
)
@click.help_option("-h", "--help")
def cmd_differentiate(
    toml_filename: str,
    encut: Optional[float],
    dir_name: str,
) -> None:
    """Calculate derivatives and write phelel_params.hdf5."""
    hdf5_filename = pathlib.Path(f"{dir_name}/phelel_params.hdf5")
    yaml_filename = pathlib.Path(f"{dir_name}/phelel_disp.yaml")

    with open(toml_filename, "rb") as f:
        toml_dict = tomli.load(f)

    if "phelel" not in toml_dict or "fft_mesh" not in toml_dict["phelel"]:
        click.echo('"fft_mesh" has to be specified in [phelel] section.', err=True)
        click.echo("See velph-hints.", err=True)
        return None

    is_symmetry = True
    try:
        if toml_dict["phelel"]["nosym"] is True:
            click.echo(
                'Found "nosym = true" in [phelel] section. '
                "Symmetrization is turned off."
            )
            is_symmetry = False
    except KeyError:
        pass

    phe = phelel.load(
        yaml_filename,
        fft_mesh=toml_dict["phelel"]["fft_mesh"],
        is_symmetry=is_symmetry,
    )

    if encut is not None:
        try:
            prec = toml_dict["vasp"]["selfenergy"]["incar"]["prec"]
        except KeyError:
            click.echo(f'[vasp.selfenergy.incar] not found in "{toml_filename}".')
            click.echo('prec = "accurate" is assumed.')
            prec = "accurate"
        click.echo(f"FFT mesh is generated for encut={encut}.")
        phe.fft_mesh = CutoffToFFTMesh.get_FFTMesh(encut, phe.primitive.cell, prec)

    if phe.fft_mesh is None:
        click.echo("FFT mesh is not specified.", err=True)
    else:
        if encut is None:
            click.echo(f"FFT mesh: {phe.fft_mesh}.")
        else:
            click.echo(f"FFT mesh: {phe.fft_mesh} (encut={encut}).")

    if run_derivatives(phe, dir_name=dir_name):
        pathlib.Path(hdf5_filename).parent.mkdir(parents=True, exist_ok=True)
        phe.save_hdf5(filename=hdf5_filename)
        click.echo(f'"{hdf5_filename}" has been made.')


#
# velph phelel phonopy
#
@cmd_phelel.command("phonopy")
@click.argument("toml_filename", nargs=1, type=click.Path(), default="velph.toml")
@click.option(
    "--dir-name",
    "dir_name",
    type=str,
    default="phelel",
    help=(
        'Used for backward compatibility, for which set "supercell". '
        '(phelel_dirname: str, default="phelel")'
    ),
)
@click.help_option("-h", "--help")
def cmd_phonopy(toml_filename: str, dir_name: str):
    """Create phonopy_params.yaml."""
    create_phonopy_yaml(
        pathlib.Path(toml_filename),
        pathlib.Path(f"{dir_name}/phelel_disp.yaml"),
        dir_name=dir_name,
    )
