# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
from __future__ import annotations

from functools import lru_cache
from numbers import Integral

import numpy as np

import sisl._array as _a

# Import the geometry object
from sisl import Atom, Geometry, Lattice
from sisl._indices import indices
from sisl._internal import set_module
from sisl.messages import deprecate, info, warn
from sisl.physics.distribution import fermi_dirac
from sisl.unit.siesta import unit_convert

# Import sile objects
from ..sile import SileWarning
from .sile import SileCDFTBtrans

__all__ = ["_ncSileTBtrans", "_devncSileTBtrans"]


Bohr2Ang = unit_convert("Bohr", "Ang")
Ry2eV = unit_convert("Ry", "eV")
Ry2K = unit_convert("Ry", "K")
eV2Ry = unit_convert("eV", "Ry")


@set_module("sisl.io.tbtrans")
class _ncSileTBtrans(SileCDFTBtrans):
    r"""Common TBtrans NetCDF file object due to a lot of the files having common entries

    This enables easy read of the Geometry and Lattices etc.
    """

    @lru_cache(maxsize=1)
    def read_lattice(self):
        """Returns `Lattice` object from this file"""
        cell = _a.arrayd(np.copy(self.cell))
        cell.shape = (3, 3)

        nsc = self._value("nsc")
        lattice = Lattice(cell, nsc=nsc)
        lattice.sc_off = self._value("isc_off")
        return lattice

    def read_geometry(self, *args, **kwargs):
        """Returns `Geometry` object from this file"""
        lattice = self.read_lattice()

        xyz = _a.arrayd(np.copy(self.xa))
        xyz.shape = (-1, 3)

        # Create list with correct number of orbitals
        lasto = _a.arrayi(np.copy(self.lasto) + 1)
        nos = np.diff(lasto, prepend=0)

        if "atom" in kwargs:
            # The user "knows" which atoms are present
            atms = kwargs["atom"]
            # Check that all atoms have the correct number of orbitals.
            # Otherwise we will correct them
            for i in range(len(atms)):
                if atms[i].no != nos[i]:
                    atms[i] = Atom(atms[i].Z, [-1] * nos[i], tag=atms[i].tag)

        else:
            # Default to Hydrogen atom with nos[ia] orbitals
            # This may be counterintuitive but there is no storage of the
            # actual species
            atms = [Atom("H", [-1] * o) for o in nos]

        # Create and return geometry object
        geom = Geometry(xyz, atms, lattice=lattice)

        return geom

    # This class also contains all the important quantities elements of the
    # file.

    @property
    @lru_cache(maxsize=1)
    def geometry(self):
        """The associated geometry from this file"""
        return self.read_geometry()

    @property
    def geom(self):
        """Same as `geometry`, but deprecated"""
        deprecate(
            f"{self.__class__.__name__}.geom is deprecated, please use '.geometry'."
            "0.14",
            "0.16",
        )
        return self.geometry

    @property
    @lru_cache(maxsize=1)
    def cell(self):
        """Unit cell in file"""
        return self._value("cell") * Bohr2Ang

    @property
    @lru_cache(maxsize=1)
    def na(self):
        """Returns number of atoms in the cell"""
        return len(self._dimension("na_u"))

    na_u = na

    @property
    @lru_cache(maxsize=1)
    def no(self):
        """Returns number of orbitals in the cell"""
        return len(self._dimension("no_u"))

    no_u = no

    @property
    @lru_cache(maxsize=1)
    def xyz(self):
        """Atomic coordinates in file"""
        return self._value("xa") * Bohr2Ang

    xa = xyz

    @property
    @lru_cache(maxsize=1)
    def lasto(self):
        """Last orbital of corresponding atom"""
        return self._value("lasto") - 1

    @property
    @lru_cache(maxsize=1)
    def k(self):
        """Sampled k-points in file"""
        return self._value("kpt")

    kpt = k

    @property
    @lru_cache(maxsize=1)
    def wk(self):
        """Weights of k-points in file"""
        return self._value("wkpt")

    wkpt = wk

    @property
    @lru_cache(maxsize=1)
    def nk(self):
        """Number of k-points in file"""
        return len(self.dimensions["nkpt"])

    nkpt = nk

    @property
    @lru_cache(maxsize=1)
    def E(self):
        """Sampled energy-points in file"""
        return self._value("E") * Ry2eV

    @property
    @lru_cache(maxsize=1)
    def ne(self):
        """Number of energy-points in file"""
        return len(self._dimension("ne"))

    nE = ne

    def Eindex(self, E):
        """Return the closest energy index corresponding to the energy ``E``

        Parameters
        ----------
        E : float or int or str
           if `int`, return it-self, else return the energy index which is
           closests to the energy.
           For a `str` it will be parsed to a float and treated as such.
        """
        if isinstance(E, Integral):
            return E
        elif isinstance(E, str):
            E = float(E)
        idxE = np.abs(self.E - E).argmin()
        ret_E = self.E[idxE]
        if abs(ret_E - E) > 5e-3:
            warn(
                f"{self.__class__.__name__} requesting energy "
                f"{E:.5f} eV, found {ret_E:.5f} eV as the closest energy!"
            )
        elif abs(ret_E - E) > 1e-3:
            info(
                f"{self.__class__.__name__} requesting energy "
                f"{E:.5f} eV, found {ret_E:.5f} eV as the closest energy!"
            )
        return idxE

    def _bias_window_integrator(self, elec_from=0, elec_to=1):
        r"""An integrator for the bias window between two electrodes

        Given two chemical potentials this returns an integrator (function) which
        returns weights for an input energy-point roughly equivalent to:

        .. math::
           dI = \mathrm dE\,[n_F(\mu_t, k_B T_t) - n_F(\mu_f, k_B T_f)]

        In this case :mathrm:`\mathrm dE` is the distance between two consecutive
        energy points in this file.

        Parameters
        ----------
        elec_from: str, int
           the originating electrode
        elec_to: str, int
           the absorbing electrode (different from `elec_from`)

        """
        elec_from = self._elec(elec_from)
        kt_from = self.kT(elec_from)
        mu_from = self.chemical_potential(elec_from)
        elec_to = self._elec(elec_to)
        kt_to = self.kT(elec_to)
        mu_to = self.chemical_potential(elec_to)
        # Get energies
        E = self._E_T_sorted(elec_from, elec_to)[0]
        dE = E[1] - E[0]

        def integrator(E):
            return dE * (
                fermi_dirac(E, kt_from, mu_from) - fermi_dirac(E, kt_to, mu_to)
            )

        return integrator

    def kindex(self, k):
        """Return the index of the k-point that is closests to the queried k-point (in reduced coordinates)

        Parameters
        ----------
        k : array_like of float or int
           the queried k-point in reduced coordinates :math:`]-0.5;0.5]`. If ``int``
           return it-self.
        """
        if isinstance(k, Integral):
            return k
        ik = np.sum(np.abs(self.k - _a.asarrayd(k)[None, :]), axis=1).argmin()
        ret_k = self.k[ik, :]
        if not np.allclose(ret_k, k, atol=0.0001):
            warn(
                SileWarning(
                    self.__class__.__name__
                    + " requesting k-point "
                    + "[{:.3f}, {:.3f}, {:.3f}]".format(*k)
                    + " found "
                    + "[{:.3f}, {:.3f}, {:.3f}]".format(*ret_k)
                )
            )
        return ik


@set_module("sisl.io.tbtrans")
class _devncSileTBtrans(_ncSileTBtrans):
    r"""Common TBtrans NetCDF file object due to a lot of the files having common entries

    This one also enables device region atoms and pivoting tables.
    """

    def read_geometry(self, *args, **kwargs):
        """Returns `Geometry` object from this file"""
        g = super().read_geometry(*args, **kwargs)
        try:
            g["Buffer"] = self.a_buf[:]
        except Exception:
            # Then no buffer atoms
            pass
        g["Device"] = self.a_dev[:]
        try:
            for elec in self.elecs:
                g[elec] = self._value("a", [elec]) - 1
                g[f"{elec}+"] = self._value("a_down", [elec]) - 1
        except Exception:
            pass
        return g

    @property
    @lru_cache(maxsize=1)
    def na_b(self):
        """Number of atoms in the buffer region"""
        return len(self._dimension("na_b"))

    na_buffer = na_b

    @property
    @lru_cache(maxsize=1)
    def a_buf(self):
        """Atomic indices (0-based) of device atoms"""
        return self._value("a_buf") - 1

    # Device atoms and other quantities
    @property
    @lru_cache(maxsize=1)
    def na_d(self):
        """Number of atoms in the device region"""
        return len(self._dimension("na_d"))

    na_dev = na_d

    @property
    @lru_cache(maxsize=1)
    def a_dev(self):
        """Atomic indices (0-based) of device atoms (sorted)"""
        return self._value("a_dev") - 1

    @lru_cache(maxsize=16)
    def a_elec(self, elec):
        """Electrode atomic indices for the full geometry (sorted)

        Parameters
        ----------
        elec : str or int
           electrode to retrieve indices for
        """
        return self._value("a", self._elec(elec)) - 1

    def a_down(self, elec, bulk=False):
        """Down-folding atomic indices for a given electrode

        Parameters
        ----------
        elec : str or int
           electrode to retrieve indices for
        bulk : bool, optional
           whether the returned indices are *only* in the pristine electrode,
           or the down-folding region (electrode + downfolding region, not in device)
        """
        if bulk:
            return self.a_elec(elec)
        return self._value("a_down", self._elec(elec)) - 1

    @property
    @lru_cache(maxsize=1)
    def o_dev(self):
        """Orbital indices (0-based) of device orbitals (sorted)

        See Also
        --------
        pivot : retrieve the device orbitals, non-sorted
        """
        return self.pivot(sort=True)

    @property
    @lru_cache(maxsize=1)
    def no_d(self):
        """Number of orbitals in the device region"""
        return len(self.dimensions["no_d"])

    def _elec(self, elec):
        """Converts a string or integer to the corresponding electrode name

        Parameters
        ----------
        elec : str or int
           if `str` it is the *exact* electrode name, if `int` it is the electrode
           index

        Returns
        -------
        str
            the electrode name
        """
        try:
            elec = int(elec)
            return self.elecs[elec]
        except Exception:
            return elec

    @property
    @lru_cache(maxsize=1)
    def elecs(self):
        """List of electrodes"""
        return list(self.groups.keys())

    @lru_cache(maxsize=16)
    def chemical_potential(self, elec):
        """Return the chemical potential associated with the electrode `elec`"""
        return self._value("mu", self._elec(elec))[0] * Ry2eV

    mu = chemical_potential

    @lru_cache(maxsize=16)
    def eta(self, elec=None):
        """The imaginary part used when calculating the self-energies in eV (or for the device

        Parameters
        ----------
        elec : str, int, optional
           electrode to extract the eta value from. If not specified (or None) the device
           region eta will be returned.
        """
        try:
            return self._value("eta", self._elec(elec))[0] * self._E2eV
        except Exception:
            return 0.0  # unknown!

    @lru_cache(maxsize=16)
    def electron_temperature(self, elec):
        """Electron bath temperature [Kelvin]

        Parameters
        ----------
        elec : str or int
           electrode to extract the temperature from

        See Also
        --------
        kT: bath temperature in [eV]
        """
        return self._value("kT", self._elec(elec))[0] * Ry2K

    @lru_cache(maxsize=16)
    def kT(self, elec):
        """Electron bath temperature [eV]

        Parameters
        ----------
        elec : str or int
           electrode to extract the temperature from

        See Also
        --------
        electron_temperature: bath temperature in [K]
        """
        return self._value("kT", self._elec(elec))[0] * Ry2eV

    @lru_cache(maxsize=16)
    def bloch(self, elec):
        """Bloch-expansion coefficients for an electrode

        Parameters
        ----------
        elec : str or int
           bloch expansions of electrode
        """
        try:
            return self._value("bloch", self._elec(elec))
        except Exception:
            return _a.onesi(3)

    @lru_cache(maxsize=16)
    def n_btd(self, elec=None):
        """Number of blocks in the BTD partioning

        Parameters
        ----------
        elec : str or int, optional
           if None the number of blocks in the device region BTD matrix. Else
           the number of BTD blocks in the electrode down-folding.
        """
        return len(self._dimension("n_btd", self._elec(elec)))

    @lru_cache(maxsize=16)
    def btd(self, elec=None):
        """Block-sizes for the BTD method in the device/electrode region

        Parameters
        ----------
        elec : str or int, optional
           the BTD block sizes for the device (if none), otherwise the downfolding
           BTD block sizes for the electrode
        """
        return self._value("btd", self._elec(elec))

    @lru_cache(maxsize=16)
    def na_down(self, elec):
        """Number of atoms in the downfolding region (without device downfolded region)

        Parameters
        ----------
        elec : str or int
           Number of downfolding atoms for electrode `elec`
        """
        return len(self._dimension("na_down", self._elec(elec)))

    @lru_cache(maxsize=16)
    def no_e(self, elec):
        """Number of orbitals in the downfolded region of the electrode in the device

        Parameters
        ----------
        elec : str or int
           Specify the electrode to query number of downfolded orbitals
        """
        return len(self._dimension("no_e", self._elec(elec)))

    @lru_cache(maxsize=16)
    def no_down(self, elec):
        """Number of orbitals in the downfolding region (plus device downfolded region)

        Parameters
        ----------
        elec : str or int
           Number of downfolding orbitals for electrode `elec`
        """
        return len(self._dimension("no_down", self._elec(elec)))

    @lru_cache(maxsize=16)
    def pivot_down(self, elec):
        """Pivoting orbitals for the downfolding region of a given electrode

        Parameters
        ----------
        elec : str or int
           the corresponding electrode to get the pivoting indices for
        """
        return self._value("pivot_down", self._elec(elec)) - 1

    @lru_cache(maxsize=32)
    def pivot(self, elec=None, in_device=False, sort=False):
        """Return the pivoting indices for a specific electrode (in the device region) or the device

        Parameters
        ----------
        elec : str or int
           the corresponding electrode to return the pivoting indices from
        in_device : bool, optional
           If ``True`` the pivoting table will be translated to the device region orbitals.
           If `sort` is also true, this would correspond to the orbitals directly translated
           to the geometry ``self.geometry.sub(self.a_dev)``.
        sort : bool, optional
           Whether the returned indices are sorted. Mostly useful if you want to handle
           the device in a non-pivoted order.

        Examples
        --------
        >>> se = tbtncSileTBtrans(...)
        >>> se.pivot()
        [3, 4, 6, 5, 2]
        >>> se.pivot(sort=True)
        [2, 3, 4, 5, 6]
        >>> se.pivot(0)
        [2, 3]
        >>> se.pivot(0, in_device=True)
        [4, 0]
        >>> se.pivot(0, in_device=True, sort=True)
        [0, 1]
        >>> se.pivot(0, sort=True)
        [2, 3]

        See Also
        --------
        pivot_down : for the pivot table for electrodes down-folding regions
        """
        if elec is None:
            if in_device and sort:
                return _a.arangei(self.no_d)
            pvt = self._value("pivot") - 1
            if in_device:
                # Count number of elements that we need to subtract from each orbital
                subn = _a.onesi(self.no)
                subn[pvt] = 0
                pvt -= _a.cumsumi(subn)[pvt]
            elif sort:
                pvt = np.sort(pvt)
            return pvt

        # Get electrode pivoting elements
        se_pvt = self._value("pivot", tree=self._elec(elec)) - 1
        if sort:
            # Sort pivoting indices
            # Since we know that pvt is also sorted, then
            # the resulting in_device would also return sorted
            # indices
            se_pvt = np.sort(se_pvt)

        if in_device:
            pvt = self._value("pivot") - 1
            if sort:
                pvt = np.sort(pvt)
            # translate to the device indices
            se_pvt = indices(pvt, se_pvt, 0)
        return se_pvt

    def a2p(self, atoms):
        """Return the pivoting orbital indices (0-based) for the atoms, possibly on an electrode

        This is equivalent to:

        >>> p = self.o2p(self.geometry.a2o(atom, True))

        Will warn if an atom requested is not in the device list of atoms.

        Parameters
        ----------
        atoms : array_like or int
           atomic indices (0-based)
        """
        return self.o2p(self.geometry.a2o(atoms, True))

    def o2p(self, orbitals, elec=None):
        """Return the pivoting indices (0-based) for the orbitals, possibly on an electrode

        Will warn if an orbital requested is not in the device list of orbitals.

        Parameters
        ----------
        orbitals : array_like or int
           orbital indices (0-based)
        elec : str or int or None, optional
           electrode to return pivoting indices of (if None it is the
           device pivoting indices).
        """
        # We need ravel(), otherwise taking len of an int will fail
        orbitals = self.geometry._sanitize_orbs(orbitals).ravel()
        porb = np.isin(self.pivot(elec), orbitals).nonzero()[0]
        d = len(orbitals) - len(porb)
        if d != 0:
            warn(
                f"{self.__class__.__name__}.o2p requesting an orbital outside the device region, "
                f"{d} orbitals will be removed from the returned list"
            )
        return porb
