# -*- coding: utf-8 -*-
#  Copyright (c) 2020. Distributed under the terms of the MIT License.
import os
from argparse import Namespace
from copy import deepcopy
from pathlib import Path

import yaml
from mp_api.client import MPRester
from pymatgen.core import Structure
from pymatgen.io.vasp import Vasprun, Outcar
from tabulate import tabulate
from vise.analyzer.dielectric_function import DieleFuncData
from vise.analyzer.plot_band import BandMplPlotter
from vise.analyzer.plot_diele_func_data import DieleFuncMplPlotter, \
    DieleFuncPlotType
from vise.analyzer.plot_dos import DosPlotter
from vise.analyzer.vasp.band_edge_properties import VaspBandEdgeProperties
from vise.analyzer.vasp.dos_data import DosDataFromVasp
from vise.analyzer.vasp.make_diele_func import make_diele_func
from vise.analyzer.vasp.make_effective_mass import make_effective_mass
# from vise.analyzer.vasp.make_irreps import special_points_from_kpoints, \
#     make_irreps_from_wavecar
from vise.analyzer.vasp.plot_band import BandPlotInfoFromVasp
from vise.cli.main_tools import potcar_str2dict, list2dict
from vise.defaults import defaults
from vise.input_set.datasets.dataset_util import all_incar_flags
from vise.input_set.input_options import CategorizedInputOptions, \
    assignable_option_set
from vise.input_set.kpoints_mode import KpointsMode
from vise.input_set.prior_info import prior_info_from_calc_dir, PriorInfo
from vise.input_set.task import Task
from vise.input_set.vasp_input_files import VaspInputFiles
from vise.util.file_transfer import FileTransfers
from vise.util.logger import get_logger
from vise.util.structure_symmetrizer import StructureSymmetrizer

logger = get_logger(__name__)


def structure_info(args: Namespace) -> None:
    input_structure = Structure.from_file(args.poscar)
    symmetrizer = StructureSymmetrizer(input_structure,
                                       symprec=args.symprec,
                                       angle_tolerance=args.angle_tolerance)
    if args.show_primitive:
        primitive = symmetrizer.primitive
        if primitive != input_structure:
            print(primitive.to(fmt="poscar"))
        else:
            logger.info("Input structure is a primitive cell.")
    elif args.show_conventional:
        conventional = symmetrizer.conventional
        if conventional != input_structure:
            print(conventional.to(fmt="poscar"))
        else:
            logger.info("Input structure is a conventional cell.")
    else:
        la = input_structure.lattice
        print(symmetrizer)
        print(f"Volume: {input_structure.volume}")
        print("abc   : " + " ".join([f"{i:0.6f}".rjust(10) for i in la.abc]))
        print("angles: " + " ".join([f"{i:0.6f}".rjust(10) for i in la.angles]))


def get_most_stable_mp_id_from_formula(formula: str):
    # API key is parsed via .pmgrc.yaml
    with MPRester() as m:
        # Due to mp_decode=True by default, class objects are restored.
        candidates = m.query(criteria=formula,
                             properties=["task_id", "e_above_hull",
                                         "spacegroup.symbol", "band_gap"])

    sorted_candidates = sorted(candidates, key=lambda x: x["e_above_hull"])
    x = []
    for c in sorted_candidates:
        x.append([c["task_id"], round(c["e_above_hull"], 3),
                  c["spacegroup.symbol"], round(c["band_gap"], 3)])

    if not x:
        logger.info(f"Formula {formula} does not exist in the materials "
                    f"project.")
        return None
    print(tabulate(x, headers=["mp_id", "e_above_hull", "space group",
                               "band gap"]))
    return x[0][0]


def get_most_stable_mp_id_from_formula_w_new_mprester(formula: str):
    # API key is parsed via .pmgrc.yaml
    with MPRester() as m:
        logger.info("Note that you're using the newer MPRester.")
        candidates = m.summary.search(
            formula=formula,
            fields=["material_id", "energy_above_hull",
                    "symmetry", "band_gap"])
    sorted_candidates = sorted(candidates, key=lambda x: x.energy_above_hull)
    table_info = []
    for c in sorted_candidates:
        table_info.append([str(c.material_id), round(c.energy_above_hull, 3),
                  c.symmetry.symbol, round(c.band_gap, 3)])

    if not table_info:
        logger.info(f"Formula {formula} does not exist in the materials "
                    f"project.")
        return None
    headers = ["mp_id", "energy_above_hull", "space group", "band gap"]
    print(tabulate(table_info, headers=headers))
    return table_info[0][0]


def get_poscar_from_mp(args: Namespace) -> None:
    if args.mpid is None and args.formula is None:
        raise ValueError("Either mp-id or formula needs to be specified.")

    if args.mpid is None and args.formula:
        args.mpid \
            = get_most_stable_mp_id_from_formula_w_new_mprester(args.formula)
        if args.mpid is None:
            return

    s = MPRester().get_structure_by_material_id(args.mpid)
    s.to(fmt="poscar", filename="POSCAR")
    query = MPRester().summary.search(material_ids=[args.mpid],
                                      fields=["total_magnetization",
                                              "band_gap"])
    data = query[0]
    d = {"total_magnetization": data.total_magnetization,
         "band_gap": data.band_gap,
         "data_source": f"new MPRester {args.mpid}"}

    Path("prior_info.yaml").write_text(yaml.dump(d), None)


class VaspSet:
    def __init__(self, args: Namespace):
        self.args = args
        if args.prev_dir:
            self.args.prev_dir = args.prev_dir.absolute()
        if args.poscar:
            self.args.poscar = args.poscar.absolute()

        for _dir in [d.absolute() for d in args.dirs]:
            if _dir.is_file():
                logger.info(f"{_dir} is a file, so skipped.")
                continue

            os.chdir(_dir)
            logger.info(f"Creating VASP input set in {_dir}...")
            try:
                self._prior_info = PriorInfo.load_yaml()
            except FileNotFoundError:
                self._prior_info = PriorInfo()
            self.task = Task.cluster_opt if self._prior_info.is_cluster \
                else args.task

            options = CategorizedInputOptions(
                structure=self._structure(),
                task=self.task,
                xc=args.xc,
                kpt_density=args.kpt_density,
                overridden_potcar=self._overridden_potcar(),
                **self._option_kwargs())

            vif = VaspInputFiles(options, self._overridden_incar_settings())
            vif.create_input_files(Path.cwd())
            if hasattr(self, "_file_transfers"):
                self._file_transfers.transfer()
            vif.vise_log.to_yaml_file()

    def _structure(self):
        # avoid overlapping structure for e.g., phonon forces.
        if self.args.prev_dir and self.task.fine_to_inherit_structure_from_prev:
            return Structure.from_file(self.args.prev_dir / defaults.contcar)
        if self.args.poscar:
            return Structure.from_file(self.args.poscar)
        return Structure.from_file("POSCAR")

    def _overridden_incar_settings(self):
        result = deepcopy(defaults.user_incar_settings)
        result.update(self._prior_info.incar)

        if self.args.user_incar_settings:
            args = list2dict(self.args.user_incar_settings, all_incar_flags)
            result.update(args)

        return result

    def _overridden_potcar(self):
        result = {}
        if self.args.overridden_potcar:
            result.update(potcar_str2dict(self.args.overridden_potcar))
        return result

    def _option_kwargs(self):
        result = deepcopy(defaults.options)
        if self._prior_info:
            result.update(self._prior_info.input_options_kwargs)

        if self.args.prev_dir:
            pi = prior_info_from_calc_dir(prev_dir_path=self.args.prev_dir,
                                          vasprun=self.args.vasprun,
                                          outcar=self.args.outcar)
            result.update(pi.input_options_kwargs)

            self._file_transfers = FileTransfers(self._file_transfer(),
                                                 path=self.args.prev_dir)

        if self.args.options:
            args = list2dict(self.args.options, assignable_option_set)
            if "kpt_density" in args and self.args.kpt_density:
                logger.warning(
                    f"Because kpt_density is set at kpt_density argument and "
                    f"optional argument, kpt_density argument value "
                    f"{self.args.kpt_density} is used.")
                args.pop("kpt_density")
            result.update(args)

            if "set_spin_orbit" in args:
                logger.info("Spin-orbit coupling is switched on.")
                result["time_reversal"] = False

        if self.args.uniform_kpt_mode:
            result["kpt_mode"] = KpointsMode.uniform

        return result

    def _file_transfer(self):
        if not self.args.file_transfer_type:
            return {}
        else:
            result = {}
            for filename, transfer in zip(self.args.file_transfer_type[0::2],
                                          self.args.file_transfer_type[1::2]):
                result[filename] = transfer
            return result


def plot_band(args: Namespace):
    band_plot_info_from_vasp = BandPlotInfoFromVasp(
        vasprun=Vasprun(args.vasprun), kpoints_filename=args.kpoints_filename)
    plot_info = band_plot_info_from_vasp.make_band_plot_info()

    # if args.wavecar_filename:
    #     try:
    #         with open(args.kpoints_filename) as f:
    #             first_line = f.readline()
    #         sg_num = int(first_line.split()[-1])
    #         logger.info(f"Space group number is set to {sg_num} read from "
    #                     f"{args.kpoints_filename} header.")
    #     except ValueError:
    #         logger.info(f"Failed to get space group number from "
    #                     f"{args.kpoints_filename} header.")
    #         sg_num = None

        # special_points, kpt_indices = \
        #     special_points_from_kpoints(args.kpoints_filename)
        # irreps = make_irreps_from_wavecar(
        #     special_points, kpt_indices, sg_num, args.wavecar_filename,
        #     args.poscar)
        # plot_info.band_energy_infos["1"].irreps = irreps

    plot_info.to_json_file()

    # if args.plotly:
    #     plotter = BandDosPlotlyPlotter(band_plot_info=plot_info)
    #     plotter.fig.show()

#         bz_info = band_plot_info_from_vasp.make_bz_plot_info()
#         bz_plotter = BZPlotlyPlotter(bz_info)
#         fig = bz_plotter.create_figure(set_k123=False)
#         fig.show()
#         fig.write_image("brillouin_zone.pdf")
#         return

    plotter = BandMplPlotter(plot_info, energy_range=args.y_range)
    plotter.construct_plot()
    plotter.plt.savefig(args.filename, format="pdf")


def plot_dos(args: Namespace):
    vasprun = Vasprun(args.vasprun)
    outcar = Outcar(args.outcar)
    band_edge = VaspBandEdgeProperties(vasprun, outcar)

    if band_edge.band_gap:
        vertical_lines = [band_edge.vbm_info.energy, band_edge.cbm_info.energy]
    else:
        vertical_lines = [vasprun.efermi]

    if args.base_energy is None:
        base = vertical_lines[0]
    else:
        base = args.base_energy

    dos_data_from_vasp = \
        DosDataFromVasp(vasprun, vertical_lines, base, args.crop_first_value)
    dos_data = dos_data_from_vasp.make_dos_data()

    ylim_set = None
    if args.y_max_ranges:
        if dos_data.spin:
            ylim_set = [[-y_max, y_max] for y_max in args.y_max_ranges]
        else:
            ylim_set = [[0, y_max] for y_max in args.y_max_ranges]

    structure = vasprun.final_structure
    grouped_atom_indices = args.type.grouped_atom_indices(structure,
                                                          args.target)
    logger.info(f"Grouped atom indices: {grouped_atom_indices}")
    plot_data = dos_data.dos_plot_data(grouped_atom_indices,
                                       energy_range=args.x_range,
                                       dos_ranges=ylim_set,
                                       title=args.title)
    plot_data.to_json_file()
    plotter = DosPlotter(plot_data, args.legend)
    plotter.construct_plot()
    plotter.plt.savefig(args.filename, format="pdf")


def plot_diele_func(args: Namespace):
    if args.input_csv_name:
        diele_func_data = DieleFuncData.from_csv_file(args.input_csv_name)
    else:
        diele_func_data = make_diele_func(Vasprun(args.vasprun),
                                          Outcar(args.outcar),
                                          use_vasp_real=not args.calc_kk,
                                          ita=args.ita)
    diele_func_data.to_json_file()
    if args.to_csv:
        diele_func_data.to_csv_file()
    if args.y_range:
        if args.plot_type is DieleFuncPlotType.absorption_coeff:
            y_range = [10 ** args.y_range[0], 10 ** args.y_range[1]]
        else:
            y_range = args.y_range
    else:
        y_range = None

    diele_func_data.title = args.title
    plotter = DieleFuncMplPlotter(diele_func_data)
    plotter.construct_plot(directions=args.directions,
                           plot_type=args.plot_type, y_range=y_range)
    filename = args.filename or str(args.plot_type) + ".pdf"
    plotter.plt.savefig(filename, format="pdf")


def calc_effective_mass(args: Namespace):
    vasprun, outcar = Vasprun(args.vasprun), Outcar(args.outcar)
    band_edge_prop = VaspBandEdgeProperties(vasprun, outcar)
    try:
        vbm, cbm = band_edge_prop.vbm_cbm
    except TypeError:
        logger.warning("Band gap does not exist, so not suited for effective"
                       "mass calculation.")
        return
    effective_mass = make_effective_mass(vasprun,
                                         args.temperature,
                                         args.concentrations,
                                         vbm, cbm)
    effective_mass.to_json_file()
    print(effective_mass)


def band_edge_properties(args: Namespace):
    vasprun = Vasprun(args.vasprun)
    outcar = Outcar(args.outcar)
    print(VaspBandEdgeProperties(vasprun, outcar))
