# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

import os
import unittest
import warnings

import numpy as np
from monty.serialization import loadfn

from pymatgen.electronic_structure.core import OrbitalType, Spin
from pymatgen.io.vasp import Vasprun
from pymatgen.util.testing import PymatgenTest

try:
    from pymatgen.electronic_structure.boltztrap2 import (
        BandstructureLoader,
        BztInterpolator,
        BztPlotter,
        BztTransportProperties,
        VasprunBSLoader,
        VasprunLoader,
    )

    BOLTZTRAP2_PRESENT = True
except Exception:
    BOLTZTRAP2_PRESENT = False

# BOLTZTRAP2_PRESENT = False

test_dir = os.path.join(PymatgenTest.TEST_FILES_DIR, "boltztrap2")


vrunfile = os.path.join(test_dir, "vasprun.xml")
vrun = Vasprun(vrunfile, parse_projected_eigen=True)

vrunfile_sp = os.path.join(test_dir, "vasprun_spin.xml")
vrun_sp = Vasprun(vrunfile_sp, parse_projected_eigen=True)
bs = loadfn(os.path.join(test_dir, "PbTe_bandstructure.json"))
bs_sp = loadfn(os.path.join(test_dir, "N2_bandstructure.json"))

bztinterp_fn = os.path.join(test_dir, "bztInterp.json.gz")
bzttransp_fn = os.path.join(test_dir, "bztTranspProps.json.gz")


@unittest.skipIf(not BOLTZTRAP2_PRESENT, "No boltztrap2, skipping tests...")
class VasprunBSLoaderTest(unittest.TestCase):
    def setUp(self):
        self.loader = VasprunBSLoader(vrun)
        self.assertIsNotNone(self.loader)
        self.loader = VasprunBSLoader(bs, vrun.final_structure)
        self.assertIsNotNone(self.loader)
        self.loader = VasprunBSLoader.from_file(vrunfile)
        self.assertIsNotNone(self.loader)

        warnings.simplefilter("ignore")

        self.loader_sp = VasprunBSLoader(vrun_sp)
        self.assertIsNotNone(self.loader_sp)
        self.loader_sp = VasprunBSLoader(bs_sp, vrun_sp.final_structure)
        self.assertIsNotNone(self.loader_sp)
        self.loader_sp = VasprunBSLoader.from_file(vrunfile_sp)
        self.assertIsNotNone(self.loader_sp)

        warnings.simplefilter("ignore")

    def tearDown(self):
        warnings.simplefilter("default")

    def test_properties(self):
        self.assertEqual(self.loader.is_spin_polarized, False)
        self.assertAlmostEqual(self.loader.fermi, 0.185266535678, 5)
        self.assertAlmostEqual(self.loader.structure.lattice.a, 4.64303565932548, 5)
        self.assertEqual(self.loader.nelect_all, 20.0)
        self.assertEqual(self.loader_sp.nelect_all, 10.0)

        self.assertTupleEqual(self.loader.ebands_all.shape, (20, 120))
        self.assertAlmostEqual(self.loader.ebands_all[10, 100], 0.2708057, 5)
        self.assertEqual(len(self.loader.proj_all), 1)
        self.assertTupleEqual(self.loader.proj_all[Spin.up].shape, (120, 20, 2, 9))

        self.assertEqual(self.loader_sp.is_spin_polarized, True)
        self.assertTupleEqual(self.loader_sp.ebands_all.shape, (24, 198))
        self.assertAlmostEqual(self.loader_sp.ebands_all[10, 100], 0.2543788, 4)
        self.assertAlmostEqual(self.loader_sp.ebands_all[22, 100], 0.2494617, 4)
        self.assertEqual(len(self.loader_sp.proj_all), 2)
        self.assertTupleEqual(self.loader_sp.proj_all[Spin.down].shape, (198, 12, 2, 9))

    def test_get_volume(self):
        self.assertAlmostEqual(self.loader.get_volume(), 477.6256714925874, 5)


@unittest.skipIf(not BOLTZTRAP2_PRESENT, "No boltztrap2, skipping tests...")
class BandstructureLoaderTest(unittest.TestCase):
    def setUp(self):
        self.loader = BandstructureLoader(bs, vrun.structures[-1])
        self.assertIsNotNone(self.loader)

        self.loader_sp = BandstructureLoader(bs_sp, vrun_sp.structures[-1])
        self.assertIsNotNone(self.loader_sp)
        self.assertTupleEqual(self.loader_sp.ebands_all.shape, (24, 198))

        warnings.simplefilter("ignore")

    def tearDown(self):
        warnings.simplefilter("default")

    def test_properties(self):
        self.assertTupleEqual(self.loader.ebands_all.shape, (20, 120))
        self.assertAlmostEqual(self.loader.fermi, 0.185266535678, 5)
        self.assertAlmostEqual(self.loader.structure.lattice.a, 4.64303565932548, 5)

    def test_get_volume(self):
        self.assertAlmostEqual(self.loader.get_volume(), 477.6256714925874, 5)

    # def test_set_upper_lower_bands(self):
    #     min_bnd = min(self.loader_sp_up.ebands.min(),
    #                   self.loader_sp_dn.ebands.min())
    #     max_bnd = max(self.loader_sp_up.ebands.max(),
    #                   self.loader_sp_dn.ebands.max())
    #     self.loader_sp_up.set_upper_lower_bands(min_bnd, max_bnd)
    #     self.loader_sp_dn.set_upper_lower_bands(min_bnd, max_bnd)
    #     self.assertTupleEqual(self.loader_sp_up.ebands.shape, (14, 198))
    #     self.assertTupleEqual(self.loader_sp_dn.ebands.shape, (14, 198))


@unittest.skipIf(not BOLTZTRAP2_PRESENT, "No boltztrap2, skipping tests...")
class VasprunLoaderTest(unittest.TestCase):
    def setUp(self):
        self.loader = VasprunLoader(vrun)
        self.assertTupleEqual(self.loader.proj.shape, (120, 20, 2, 9))
        self.assertIsNotNone(self.loader)
        warnings.simplefilter("ignore")

    def tearDown(self):
        warnings.simplefilter("default")

    def test_properties(self):
        self.assertTupleEqual(self.loader.ebands.shape, (20, 120))
        self.assertAlmostEqual(self.loader.fermi, 0.185266535678, 5)
        self.assertAlmostEqual(self.loader.structure.lattice.a, 4.64303565932548, 5)

    def test_get_volume(self):
        self.assertAlmostEqual(self.loader.get_volume(), 477.6256714925874, 5)

    def test_from_file(self):
        self.loader = VasprunLoader().from_file(vrunfile)
        self.assertIsNotNone(self.loader)


@unittest.skipIf(not BOLTZTRAP2_PRESENT, "No boltztrap2, skipping tests...")
class BztInterpolatorTest(unittest.TestCase):
    def setUp(self):
        self.loader = VasprunBSLoader(vrun)
        self.bztInterp = BztInterpolator(self.loader, lpfac=2)
        self.assertIsNotNone(self.bztInterp)
        self.bztInterp = BztInterpolator(self.loader, lpfac=2, save_bztInterp=True, fname=bztinterp_fn)
        self.assertIsNotNone(self.bztInterp)
        self.bztInterp = BztInterpolator(self.loader, load_bztInterp=True, fname=bztinterp_fn)
        self.assertIsNotNone(self.bztInterp)

        warnings.simplefilter("ignore")

        self.loader_sp = VasprunBSLoader(vrun_sp)
        self.bztInterp_sp = BztInterpolator(self.loader_sp, lpfac=2)
        self.assertIsNotNone(self.bztInterp_sp)
        self.bztInterp_sp = BztInterpolator(self.loader_sp, lpfac=2, save_bztInterp=True, fname=bztinterp_fn)
        self.assertIsNotNone(self.bztInterp_sp)
        self.bztInterp_sp = BztInterpolator(self.loader_sp, lpfac=2, load_bztInterp=True, fname=bztinterp_fn)
        self.assertIsNotNone(self.bztInterp_sp)

        warnings.simplefilter("ignore")

    def tearDown(self):
        warnings.simplefilter("default")

    def test_properties(self):
        self.assertTupleEqual(self.bztInterp.cband.shape, (6, 3, 3, 3, 29791))
        self.assertTupleEqual(self.bztInterp.eband.shape, (6, 29791))
        self.assertTupleEqual(self.bztInterp.coeffs.shape, (6, 322))
        self.assertEqual(self.bztInterp.data.nelect, 6.0)
        self.assertEqual(self.bztInterp.data.nelect_all, 20.0)
        self.assertTupleEqual(self.bztInterp.data.ebands.shape, (6, 120))

        self.assertTupleEqual(self.bztInterp_sp.cband.shape, (10, 3, 3, 3, 23275))
        self.assertTupleEqual(self.bztInterp_sp.eband.shape, (10, 23275))
        self.assertTupleEqual(self.bztInterp_sp.coeffs.shape, (10, 519))
        self.assertEqual(self.bztInterp_sp.data.nelect, 6.0)
        self.assertEqual(self.bztInterp_sp.data.nelect_all, 10.0)
        self.assertTupleEqual(self.bztInterp_sp.data.ebands.shape, (10, 198))

    def test_get_band_structure(self):
        sbs = self.bztInterp.get_band_structure()
        self.assertIsNotNone(sbs)
        self.assertTupleEqual(sbs.bands[Spin.up].shape, (6, 137))
        kpaths = [["L", "K"]]
        kp_lbl = {"L": np.array([0.5, 0.5, 0.5]), "K": np.array([0.375, 0.375, 0.75])}
        sbs = self.bztInterp.get_band_structure(kpaths, kp_lbl)
        self.assertIsNotNone(sbs)
        self.assertTupleEqual(sbs.bands[Spin.up].shape, (6, 20))

        sbs = self.bztInterp_sp.get_band_structure()
        self.assertIsNotNone(sbs)
        self.assertTupleEqual(sbs.bands[Spin.up].shape, (6, 143))
        self.assertTupleEqual(sbs.bands[Spin.down].shape, (4, 143))

    def test_tot_dos(self):
        tot_dos = self.bztInterp.get_dos(T=200, npts_mu=100)
        self.assertIsNotNone(tot_dos)
        self.assertEqual(len(tot_dos.energies), 100)
        self.assertAlmostEqual(tot_dos.densities[Spin.up][0], 1.35371715, 5)

        tot_dos = self.bztInterp_sp.get_dos(T=200, npts_mu=100)
        self.assertIsNotNone(tot_dos)
        self.assertEqual(len(tot_dos.energies), 100)
        self.assertAlmostEqual(tot_dos.densities[Spin.up][75], 88.034456, 5)
        self.assertAlmostEqual(tot_dos.densities[Spin.down][75], 41.421367, 5)

    def test_tot_proj_dos(self):
        tot_proj_dos = self.bztInterp.get_dos(partial_dos=True, T=200, npts_mu=100)
        self.assertIsNotNone(tot_proj_dos)
        self.assertEqual(len(tot_proj_dos.get_spd_dos().values()), 3)
        pdos = tot_proj_dos.get_spd_dos()[OrbitalType.s].densities[Spin.up][75]
        self.assertAlmostEqual(pdos, 2490.169396, 5)

        tot_proj_dos = self.bztInterp_sp.get_dos(partial_dos=True, T=200, npts_mu=100)
        self.assertIsNotNone(tot_proj_dos)
        self.assertEqual(len(tot_proj_dos.get_spd_dos().values()), 3)
        pdos = tot_proj_dos.get_spd_dos()[OrbitalType.s].densities[Spin.up][75]
        self.assertAlmostEqual(pdos, 166.4933305, 5)
        pdos = tot_proj_dos.get_spd_dos()[OrbitalType.s].densities[Spin.down][75]
        self.assertAlmostEqual(pdos, 272.194174, 5)


@unittest.skipIf(not BOLTZTRAP2_PRESENT, "No boltztrap2, skipping tests...")
class BztTransportPropertiesTest(unittest.TestCase):
    def setUp(self):
        loader = VasprunBSLoader(vrun)
        bztInterp = BztInterpolator(loader, lpfac=2)
        self.bztTransp = BztTransportProperties(bztInterp, temp_r=np.arange(300, 600, 100))
        self.assertIsNotNone(self.bztTransp)
        warnings.simplefilter("ignore")

        self.bztTransp = BztTransportProperties(
            bztInterp, doping=10.0 ** np.arange(20, 22), temp_r=np.arange(300, 600, 100)
        )
        self.assertIsNotNone(self.bztTransp)
        self.assertEqual(self.bztTransp.contain_props_doping, True)

        warnings.simplefilter("ignore")

        bztInterp = BztInterpolator(loader, lpfac=2)
        self.bztTransp = BztTransportProperties(
            bztInterp,
            temp_r=np.arange(300, 600, 100),
            save_bztTranspProps=True,
            fname=bzttransp_fn,
        )
        self.assertIsNotNone(self.bztTransp)
        warnings.simplefilter("ignore")

        bztInterp = BztInterpolator(loader, lpfac=2)
        self.bztTransp = BztTransportProperties(bztInterp, load_bztTranspProps=True, fname=bzttransp_fn)
        self.assertIsNotNone(self.bztTransp)
        warnings.simplefilter("ignore")

        loader_sp = VasprunBSLoader(vrun_sp)
        bztInterp_sp = BztInterpolator(loader_sp, lpfac=2)
        self.bztTransp_sp = BztTransportProperties(bztInterp_sp, temp_r=np.arange(300, 600, 100))
        self.assertIsNotNone(self.bztTransp_sp)
        warnings.simplefilter("ignore")

        bztInterp_sp = BztInterpolator(loader_sp, lpfac=2)
        self.bztTransp_sp = BztTransportProperties(
            bztInterp_sp,
            temp_r=np.arange(300, 600, 100),
            save_bztTranspProps=True,
            fname=bzttransp_fn,
        )
        self.assertIsNotNone(self.bztTransp_sp)
        warnings.simplefilter("ignore")

        bztInterp_sp = BztInterpolator(loader_sp, lpfac=2)
        self.bztTransp_sp = BztTransportProperties(bztInterp_sp, load_bztTranspProps=True, fname=bzttransp_fn)
        self.assertIsNotNone(self.bztTransp_sp)
        warnings.simplefilter("ignore")

    def tearDown(self):
        warnings.simplefilter("default")

    def test_properties(self):
        for p in [
            self.bztTransp.Conductivity_mu,
            self.bztTransp.Seebeck_mu,
            self.bztTransp.Kappa_mu,
            self.bztTransp.Effective_mass_mu,
            self.bztTransp.Power_Factor_mu,
        ]:
            self.assertTupleEqual(p.shape, (3, 3686, 3, 3))

        for p in [
            self.bztTransp.Carrier_conc_mu,
            self.bztTransp.Hall_carrier_conc_trace_mu,
        ]:
            self.assertTupleEqual(p.shape, (3, 3686))

        for p in [
            self.bztTransp_sp.Conductivity_mu,
            self.bztTransp_sp.Seebeck_mu,
            self.bztTransp_sp.Kappa_mu,
            self.bztTransp_sp.Effective_mass_mu,
            self.bztTransp_sp.Power_Factor_mu,
        ]:
            self.assertTupleEqual(p.shape, (3, 3252, 3, 3))

        for p in [
            self.bztTransp_sp.Carrier_conc_mu,
            self.bztTransp_sp.Hall_carrier_conc_trace_mu,
        ]:
            self.assertTupleEqual(p.shape, (3, 3252))

    def test_compute_properties_doping(self):
        self.bztTransp.compute_properties_doping(doping=10.0 ** np.arange(20, 22))
        for p in [
            self.bztTransp.Conductivity_doping,
            self.bztTransp.Seebeck_doping,
            self.bztTransp.Kappa_doping,
            self.bztTransp.Effective_mass_doping,
            self.bztTransp.Power_Factor_doping,
        ]:
            self.assertTupleEqual(p["n"].shape, (3, 2, 3, 3))
            self.assertEqual(self.bztTransp.contain_props_doping, True)

        self.bztTransp_sp.compute_properties_doping(doping=10.0 ** np.arange(20, 22))
        for p in [
            self.bztTransp_sp.Conductivity_doping,
            self.bztTransp_sp.Seebeck_doping,
            self.bztTransp_sp.Kappa_doping,
            self.bztTransp_sp.Effective_mass_doping,
            self.bztTransp_sp.Power_Factor_doping,
        ]:
            self.assertTupleEqual(p["n"].shape, (3, 2, 3, 3))
            self.assertEqual(self.bztTransp_sp.contain_props_doping, True)


@unittest.skipIf(not BOLTZTRAP2_PRESENT, "No boltztrap2, skipping tests...")
class BztPlotterTest(unittest.TestCase):
    def test_plot(self):
        loader = VasprunBSLoader(vrun)
        bztInterp = BztInterpolator(loader, lpfac=2)
        bztTransp = BztTransportProperties(bztInterp, temp_r=np.arange(300, 600, 100))
        self.bztPlotter = BztPlotter(bztTransp, bztInterp)
        self.assertIsNotNone(self.bztPlotter)
        fig = self.bztPlotter.plot_props("S", "mu", "temp", temps=[300, 500])
        self.assertIsNotNone(fig)
        fig = self.bztPlotter.plot_bands()
        self.assertIsNotNone(fig)
        fig = self.bztPlotter.plot_dos()
        self.assertIsNotNone(fig)


if __name__ == "__main__":
    unittest.main()
