# -*- coding: utf-8 -*-
#  Copyright (c) 2020. Distributed under the terms of the MIT License.

from argparse import Namespace
import unittest
import numpy as np
import os
import shutil
from pathlib import Path
from unittest.mock import patch

from pymatgen.core.composition import Composition
from pymatgen.core.sites import Element
from pymatgen.analysis.phase_diagram import (
    PDEntry, PhaseDiagram, CompoundPhaseDiagram, PDPlotter)

from vise.chempotdiag.chem_pot_diag import ChemPotDiag, sort_coords
from vise.util.testing import ViseTest
from vise.custodian_extension.jobs import ViseVaspJob
from vise.cli.main import (
    simple_override, parse_args)
from vise.util.main_tools import get_default_args
from vise.util.mp_tools import make_poscars_from_mp
from vise.input_set.input_set import ViseInputSet
from vise.config import (
    SYMMETRY_TOLERANCE, ANGLE_TOL, KPT_DENSITY, KPT_INIT_DENSITY, TIMEOUT)


parent_dir = Path(__file__).parent


class SimpleOverrideTest(ViseTest):
    def setUp(self) -> None:
        vise_yaml = parent_dir / ".." / ".." / "test_files" / "vise.yaml"
        shutil.copy(str(vise_yaml), "vise.yaml")

    def test(self):
        d = {"xc": "pbesol", "symprec": 1}
        simple_override(d, ["xc"])
        expected = {'xc': 'hse', 'symprec': 1}

        self.assertEqual(expected, d)

    def tearDown(self) -> None:
        os.remove("vise.yaml")


vis_defaults = get_default_args(ViseInputSet.make_input)
vis_defaults.update(ViseInputSet.TASK_OPTIONS)
vis_defaults.update(ViseInputSet.XC_OPTIONS)

kpt_conv_defaults = get_default_args(ViseVaspJob.kpt_converge)

# func is a pointer so need to point the same address.
# The directly passed args from argpaser.
default_vasp_args = {
    "potcar_set": None,
    "potcar_set_name": vis_defaults["potcar_set_name"],
    "xc": vis_defaults["xc"],
    "task": vis_defaults["task"],
    "vise_opts": None,
    "user_incar_settings": None,
    "additional_user_incar_settings": None,
    "ldauu": vis_defaults["ldauu"],
    "ldaul": vis_defaults["ldaul"],
    "charge": None}

symprec_args = {"symprec": SYMMETRY_TOLERANCE,  "angle_tolerance": ANGLE_TOL}


class MainGetPoscarsTest(ViseTest):

    def test_get_poscars_wo_options(self):
        parsed_args = parse_args(["gp"])

        d = get_default_args(make_poscars_from_mp)
        # func is a pointer so need to point the same address.
        expected = Namespace(
            poscar=None,
            number=None,
            elements=None,
            e_above_hull=d["e_above_hull"],
            molecules=d["molecules"],
            func=parsed_args.func)
        self.assertEqual(expected, parsed_args)

    def test_get_poscars_w_options(self):
        parsed_args = parse_args(["gp",
                                  "-p", "POSCAR",
                                  "-n", "123",
                                  "-e", "Mg", "O",
                                  "--e_above_hull", "0.1",
                                  "--molecules", "F"])

        expected = Namespace(
            poscar="POSCAR",
            number=123,
            elements=["Mg", "O"],
            e_above_hull=0.1,
            molecules=False,
            func=parsed_args.func)
        self.assertEqual(expected, parsed_args)


class MainVaspSetTest(ViseTest):
    def test_vasp_set_wo_options(self):
        parsed_args = parse_args(["vs"])
        # func is a pointer so need to point the same address.
        expected = Namespace(
            json=None,
            poscar="POSCAR",
            kpt_density=KPT_DENSITY,
            standardize_structure=True,
            prior_info=True,
            dirs=None,
            prev_dir=None,
            func=parsed_args.func,
            **default_vasp_args, **symprec_args)
        self.assertEqual(expected, parsed_args)

    def test_vasp_set_yaml(self):
        vise_yaml = parent_dir / ".." / ".." / "test_files" / "vise.yaml"
        shutil.copy(str(vise_yaml), str("vise.yaml"))
        parsed_args = parse_args(["vs",
                                  "--ldauu", "Mg", "5",
                                  "-d", "c"])
        os.remove("vise.yaml")
        expected = Namespace(
            potcar_set=default_vasp_args["potcar_set"],
            potcar_set_name="gw",
            xc="hse",
            task=default_vasp_args["task"],
            vise_opts=default_vasp_args["vise_opts"],
            user_incar_settings=default_vasp_args["user_incar_settings"],
            additional_user_incar_settings=
            default_vasp_args["additional_user_incar_settings"],
            ldauu=["Mg", "5"],
            ldaul=["Mn", "2", "O", "1"],
            charge=default_vasp_args["charge"],
            symprec=2,
            angle_tolerance=symprec_args["angle_tolerance"],
            json=None,
            poscar="POSCAR",
            kpt_density=KPT_DENSITY,
            standardize_structure=True,
            prior_info=True,
            dirs=None,
            prev_dir="c",
            func=parsed_args.func)

        self.assertEqual(expected, parsed_args)

    # Here, test vasp_parser and symprec_parser as well.
    def test_vasp_set_w_options(self):
        parsed_args = parse_args(["vs",
                                  "--potcar", "Mg_pv", "O_h",
                                  "--potcar_set_name", "gw",
                                  "-x", "pbesol",
                                  "-t", "band",
                                  "--vise_opts", "encut", "800",
                                  "--user_incar_settings", "LREAD", "F",
                                  "-auis", "ALGO", "D",
                                  "--ldauu", "Mg", "5",
                                  "--ldaul", "Mg", "1",
                                  "-c", "10",
                                  "--symprec", "0.0001",
                                  "--angle_tolerance", "10",
                                  "--json", "tmp.json",
                                  "--poscar", "POSCAR-tmp",
                                  "-k", "4.2",
                                  "-s", "F",
                                  "-pi", "T",
                                  "--dirs", "a", "b",
                                  "-d", "c"])

        expected = Namespace(
            potcar_set=["Mg_pv", "O_h"],
            potcar_set_name="gw",
            xc="pbesol",
            task="band",
            vise_opts=["encut", "800"],
            user_incar_settings=["LREAD", "F"],
            additional_user_incar_settings=["ALGO", "D"],
            ldauu=["Mg", "5"],
            ldaul=["Mg", "1"],
            charge=10.0,
            symprec=0.0001,
            angle_tolerance=10.0,
            json="tmp.json",
            poscar="POSCAR-tmp",
            kpt_density=4.2,
            standardize_structure=False,
            prior_info=True,
            dirs=["a", "b"],
            prev_dir="c",
            func=parsed_args.func)

        self.assertEqual(expected, parsed_args)


class MainVaspRunTest(ViseTest):
    def test_vasp_run_wo_options(self):
        parsed_args = parse_args(["vr", "-v", "vasp", "cmd"])
        # func is a pointer so need to point the same address.
        expected = Namespace(
            json_file=None,
            vasp_cmd=["vasp", "cmd"],
            initial_kpt_density=KPT_INIT_DENSITY,
            handler_name="default",
            timeout=TIMEOUT,
            remove_wavecar=False,
            max_relax_num=kpt_conv_defaults["max_relax_num"],
            convergence_criterion=kpt_conv_defaults["convergence_criterion"],
            left_files=kpt_conv_defaults["left_files"],
            kpoint_conv=False,
            func=parsed_args.func,
            **default_vasp_args, **symprec_args)
        self.assertEqual(expected, parsed_args)

    def test_vasp_set_w_options(self):
        parsed_args = parse_args(["vr",
                                  "-kc",
                                  "--json_file", "test.json",
                                  "-v", "vasp", "cmd",
                                  "-ikd", "3.5",
                                  "--handler_name", "dielectric",
                                  "--timeout", "1000",
                                  "--remove_wavecar",
                                  "--max_relax_num", "10",
                                  "--criterion", "0.01",
                                  "--left_files", "POSCAR", "PCDAT",
                                  ])
        expected = Namespace(
            kpoint_conv=True,
            json_file="test.json",
            vasp_cmd=["vasp", "cmd"],
            initial_kpt_density=3.5,
            handler_name="dielectric",
            timeout=1000,
            remove_wavecar=True,
            max_relax_num=10,
            convergence_criterion=0.01,
            left_files=["POSCAR", "PCDAT"],
            func=parsed_args.func,
            **default_vasp_args, **symprec_args)
        self.assertEqual(expected, parsed_args)


class MainChemPotDiagTest(ViseTest):
    def test_vasp_run_wo_options(self):
        parsed_args = parse_args(["cpd"])
        expected = Namespace(
            draw_phase_diagram=False,
            vasp_dirs=None,
            vasprun="vasprun.xml",
            elements=None,
            target_comp=None,
            filename=None,
            parse_gas=True,
            partial_pressures=None,
            temperature=0.0,
            func=parsed_args.func)
        self.assertEqual(expected, parsed_args)

    def test_vasp_run_w_options(self):
        parsed_args = parse_args(["cpd",
                                  "-dpd",
                                  "-d", "a_dir", "b_dir",
                                  "-v", "vasprun.xml.finish",
                                  "-e", "Mg", "O",
                                  "-c", "MgO",
                                  "-f", "cpd.pdf",
                                  "-pg", "F",
                                  "-pp", "O", "1e+5",
                                  "-t", "1000"])

        expected = Namespace(
            draw_phase_diagram=True,
            vasp_dirs=["a_dir", "b_dir"],
            vasprun="vasprun.xml.finish",
            elements=["Mg", "O"],
            target_comp="MgO",
            filename="cpd.pdf",
            parse_gas=False,
            partial_pressures=["O", "1e+5"],
            temperature=1000.0,
            func=parsed_args.func)
        self.assertEqual(expected, parsed_args)


# class MainPlotBandTest(ViseTest):
#     def
