from __future__ import annotations

from ase.io.espresso import Namelist

PH_KEYS = Namelist(
    {
        "INPUTPH": [
            "amass",
            "outdir",
            "prefix",
            "niter_ph",
            "tr2_ph",
            "alpha_mix",
            "nmix_ph",
            "verbosity",
            "reduce_io",
            "max_seconds",
            "dftd3_hess",
            "fildyn",
            "fildrho",
            "fildvscf",
            "epsil",
            "lrpa",
            "lnoloc",
            "trans",
            "lraman",
            "eth_rps",
            "eth_ns",
            "dek",
            "recover",
            "low_directory_check",
            "only_init",
            "qplot",
            "q2d",
            "q_in_band_form",
            "electron_phonon",
            "el_ph_nsigma",
            "el_ph_sigma",
            "ahc_dir",
            "ahc_nbnd",
            "ahc_nbndskip",
            "skip_upperfan",
            "lshift_q",
            "zeu",
            "zue",
            "elop",
            "fpol",
            "ldisp",
            "nogg",
            "asr",
            "ldiag",
            "lqdir",
            "search_sym",
            "nq1",
            "nq2",
            "nq3",
            "nk1",
            "nk2",
            "nk3",
            "k1",
            "k2",
            "k3",
            "diagonalization",
            "read_dns_bare",
            "ldvscf_interpolate",
            "wpot_dir",
            "do_long_range",
            "do_charge_neutral",
            "start_irr",
            "last_irr",
            "nat_todo",
            "modenum",
            "start_q",
            "last_q",
            "dvscf_star",
            "drho_star",
        ]
    }
)

NEB_KEYS = Namelist(
    {
        "PATH": [
            "string_method",
            "restart_mode",
            "nstep_path",
            "num_of_images",
            "opt_scheme",
            "CI_scheme",
            "first_last_opt",
            "minimum_image",
            "temp_req",
            "ds",
            "k_max",
            "k_min",
            "path_thr",
            "uses_masses",
            "use_freezing",
            "lfcp",
            "fcp_mu",
            "fcp_thr",
            "fcp_scheme",
        ]
    }
)

Q2R_KEYS = Namelist({"INPUT": ["fildyn", "flfrc", "zasr", "loto_2d", "write_lr"]})

MATDYN_KEYS = Namelist(
    {
        "INPUT": [
            "flfrc",
            "asr",
            "huang",
            "dos",
            "nk1",
            "nk2",
            "nk3",
            "deltaE",
            "ndos",
            "degauss",
            "fldos",
            "flfrq",
            "flvec",
            "fleig",
            "fldyn",
            "at",
            "l1",
            "l2",
            "l3",
            "ntyp",
            "amass",
            "readtau",
            "fltau",
            "la2F",
            "q_in_band_form",
            "q_in_cryst_coord",
            "eigen_similarity",
            "fd",
            "na_ifc",
            "nosym",
            "loto_2d",
            "loto_disable",
            "read_lr",
            "write_frc",
        ]
    }
)

DYNMAT_KEYS = Namelist(
    {
        "INPUT": [
            "fildyn",
            "q",
            "amass",
            "asr",
            "remove_interaction_blocks",
            "axis",
            "lperm",
            "lplasma",
            "filout",
            "fileig",
            "filmol",
            "filxsf",
            "loto_2d",
            "el_ph_nsig",
            "el_ph_sigma",
        ]
    }
)

PROJWFC_KEYS = Namelist(
    {
        "PROJWFC": [
            "prefix",
            "outdir",
            "ngauss",
            "degauss",
            "emin",
            "emax",
            "deltae",
            "lsymdiag_basis",
            "pawproj",
            "filpdos",
            "filproj",
            "lwrite_overlaps",
            "lbinary_data",
            "kresolveddos",
            "tdosinboxes",
            "n_proj_boxes",
            "irmin",
            "irmax",
            "plot_boxes",
        ]
    }
)


# Updated pw.x keys from my own branch

PW_KEYS = Namelist(
    (
        (
            "CONTROL",
            [
                "calculation",
                "title",
                "verbosity",
                "restart_mode",
                "wf_collect",
                "nstep",
                "iprint",
                "tstress",
                "tprnfor",
                "dt",
                "outdir",
                "wfcdir",
                "prefix",
                "lkpoint_dir",
                "max_seconds",
                "etot_conv_thr",
                "forc_conv_thr",
                "disk_io",
                "pseudo_dir",
                "tefield",
                "dipfield",
                "lelfield",
                "nberrycyc",
                "lorbm",
                "lberry",
                "gdir",
                "nppstr",
                "gatelfcp",
                "trism",
            ],
        ),
        (
            "SYSTEM",
            [
                "ibrav",
                "celldm",
                "a",
                "b",
                "c",
                "cosAB",
                "cosAC",
                "cosBCnat",
                "ntyp",
                "nbnd",
                "tot_charge",
                "tot_magnetization",
                "starting_magnetization",
                "ecutwfc",
                "ecutrho",
                "ecutfock",
                "nr1",
                "nr2",
                "nr3",
                "nr1s",
                "nr2s",
                "nr3s",
                "nosym",
                "nosym_evc",
                "noinv",
                "no_t_rev",
                "force_symmorphic",
                "use_all_frac",
                "occupations",
                "one_atom_occupations",
                "starting_spin_angle",
                "degauss",
                "smearing",
                "nspin",
                "noncolin",
                "ecfixed",
                "qcutz",
                "q2sigma",
                "input_dft",
                "exx_fraction",
                "screening_parameter",
                "exxdiv_treatment",
                "x_gamma_extrapolation",
                "ecutvcut",
                "nqx1",
                "nqx2",
                "nqx3",
                "lda_plus_u",
                "lda_plus_u_kind",
                "Hubbard_U",
                "Hubbard_J0",
                "Hubbard_occ",
                "Hubbard_alpha",
                "Hubbard_beta",
                "Hubbard_J",
                "starting_ns_eigenvalue",
                "U_projection_type",
                "edir",
                "emaxpos",
                "eopreg",
                "eamp",
                "angle1",
                "angle2",
                "constrained_magnetization",
                "fixed_magnetization",
                "lambda",
                "report",
                "lspinorb",
                "assume_isolated",
                "esm_bc",
                "esm_w",
                "esm_efield",
                "esm_nfit",
                "fcp_mu",
                "vdw_corr",
                "london",
                "london_s6",
                "london_c6",
                "london_rvdw",
                "london_rcut",
                "ts_vdw_econv_thr",
                "ts_vdw_isolated",
                "xdm",
                "xdm_a1",
                "xdm_a2",
                "space_group",
                "uniqueb",
                "origin_choice",
                "rhombohedral",
                "zmon",
                "realxz",
                "block",
                "block_1",
                "block_2",
                "block_height",
            ],
        ),
        (
            "ELECTRONS",
            [
                "electron_maxstep",
                "scf_must_converge",
                "conv_thr",
                "adaptive_thr",
                "conv_thr_init",
                "conv_thr_multi",
                "mixing_mode",
                "mixing_beta",
                "mixing_ndim",
                "mixing_fixed_ns",
                "diagonalization",
                "ortho_para",
                "diago_thr_init",
                "diago_cg_maxiter",
                "diago_david_ndim",
                "diago_full_acc",
                "efield",
                "efield_cart",
                "efield_phase",
                "startingpot",
                "startingwfc",
                "tqr",
            ],
        ),
        (
            "IONS",
            [
                "ion_dynamics",
                "ion_positions",
                "pot_extrapolation",
                "wfc_extrapolation",
                "remove_rigid_rot",
                "ion_temperature",
                "tempw",
                "tolp",
                "delta_t",
                "nraise",
                "refold_pos",
                "upscale",
                "bfgs_ndim",
                "trust_radius_max",
                "trust_radius_min",
                "trust_radius_ini",
                "w_1",
                "w_2",
            ],
        ),
        (
            "CELL",
            [
                "cell_dynamics",
                "press",
                "wmass",
                "cell_factor",
                "press_conv_thr",
                "cell_dofree",
            ],
        ),
        (
            "RISM",
            [
                "nsolv",
                "closure",
                "tempv",
                "ecutsolv",
                "solute_lj",
                "solute_epsilon",
                "solute_sigma",
                "starting1d",
                "starting3d",
                "smear1d",
                "smear3d",
                "rism1d_maxstep",
                "rism3d_maxstep",
                "rism1d_conv_thr",
                "rism3d_conv_thr",
                "mdiis1d_size",
                "mdiis3d_size",
                "mdiis1d_step",
                "mdiis3d_step",
                "rism1d_bond_width",
                "rism1d_dielectric",
                "rism1d_molesize",
                "rism1d_nproc",
                "rism3d_conv_level",
                "rism3d_planar_average",
                "laue_nfit",
                "laue_expand_right",
                "laue_expand_left",
                "laue_starting_right",
                "laue_starting_left",
                "laue_buffer_right",
                "laue_buffer_left",
                "laue_both_hands",
                "laue_wall",
                "laue_wall_z",
                "laue_wall_rho",
                "laue_wall_epsilon",
                "laue_wall_sigma",
                "laue_wall_lj6",
            ],
        ),
    )
)

ALL_KEYS = {
    "pw": PW_KEYS,
    "ph": PH_KEYS,
    "neb": NEB_KEYS,
    "q2r": Q2R_KEYS,
    "dynmat": DYNMAT_KEYS,
    "projwfc": PROJWFC_KEYS,
}
