from typing import Dict, List, Tuple, Any, Optional
import warnings

import numpy as np
from scipy.optimize import brentq # type: ignore

from py_sc_fermi.dos import DOS
from py_sc_fermi.defect_species import DefectSpecies
from py_sc_fermi.inputs import InputSet


class DefectSystem(object):
    """This class is used to calculate the self consistent Fermi energy for
    a defective material, observing the condition of charge neutrality and
    therefore, point defect and carrier concentrations under equilibrium
    conditions.

    Args:
        defect_species (List[DefectSpecies]): List of ``DefectSpecies`` objects
          which are present in the ``DefectSystem``.
        volume (float): volume of the unit cell in Angstroms cubed
        dos (DOS): the ``DOS`` object associated with the unit cell
        temperature (float): temperature at which self-consistent Fermi energy
          will be solved for.
        convergence_tolerance (float, optional): Tolerance for the Fermi energy
          convergence in eV. If not specified, uses scipy's default.
        n_trial_steps (int, optional): Deprecated. Previously set the maximum
          number of solver iterations. The solver now uses Brent's method
          which converges reliably without this parameter.
    """

    def __init__(
        self,
        defect_species: List[DefectSpecies],
        dos: DOS,
        volume: float,
        temperature: float,
        convergence_tolerance: Optional[float] = None,
        n_trial_steps: Optional[int] = None, # deprecated
    ):
        self.defect_species = defect_species
        self.volume = volume
        self.dos = dos
        self.temperature = temperature
        self.convergence_tolerance = convergence_tolerance
        if n_trial_steps is not None:
            warnings.warn(
                "n_trial_steps is deprecated and will be removed in a future version. "
                "The solver now uses Brent's method which converges reliably without "
                "this parameter.",
                DeprecationWarning,
                stacklevel=2,
            )
        self.n_trial_steps = n_trial_steps

    def __repr__(self):
        to_return = [
            f"DefectSystem\n",
            f"  nelect: {self.dos.nelect} e\n",
            f"  bandgap: {self.dos.bandgap} eV\n",
            f"  volume: {self.volume} A^3\n",
            f"  temperature: {self.temperature} K\n",
            f"\nContains defect species:\n",
        ]
        for ds in self.defect_species:
            to_return.append(str(ds))
        return "".join(to_return)

    @property
    def defect_species_names(self) -> List[str]:
        """list of the names of all ``DefectSpecies`` considered in the
        ``DefectSystem``.

        Returns:
            List[str]: list of names of ``DefectSpecies`` objects
        """
        return [ds.name for ds in self.defect_species]

    @classmethod
    def from_input_set(cls, input_set: InputSet) -> "DefectSystem":
        """generate ``DefectSystem`` from ``InputSet``

        Args:
            input_set (InputSet): ``InputSet`` defining input parameters

        Returns:
            DefectSystem: ``DefectSystem`` corresponding to provided ``InputSet``
        """

        return cls(
            defect_species=input_set.defect_species,
            dos=input_set.dos,
            volume=input_set.volume,
            temperature=input_set.temperature,
            convergence_tolerance=input_set.convergence_tolerance,
            n_trial_steps=input_set.n_trial_steps,
        )

    @classmethod
    def from_yaml(cls, filename: str, structure_file="", dos_file="") -> "DefectSystem":
        """generate ``DefectSystem`` via a yaml file.

        Args:
            filename (str): path to yaml file containing the ``DefectSystem``
              data
            structure_file (str): path to file containing volume information.
              Defaults to an empty string.
            dos_file (str): path to file containing dos information. Defaults
              to an empty string.

        Returns:
            DefectSystem: ``DefectSystem`` corresponding to provided yaml file
        """

        input_set = InputSet.from_yaml(filename, structure_file, dos_file)
        return cls(
            defect_species=input_set.defect_species,
            dos=input_set.dos,
            volume=input_set.volume,
            temperature=input_set.temperature,
            convergence_tolerance=input_set.convergence_tolerance,
            n_trial_steps=input_set.n_trial_steps,
        )

    @classmethod
    def from_dict(cls, dictionary: dict) -> "DefectSystem":
        """Generate a DefectSystem from a dictionary.
    
        Args:
            dictionary (dict): Dictionary containing the DefectSystem data.
    
        Returns:
            DefectSystem: DefectSystem corresponding to the provided dictionary.
        """
        return cls(
            dos=DOS.from_dict(dictionary["dos"]),
            volume=dictionary["volume"],
            temperature=dictionary["temperature"],
            convergence_tolerance=dictionary.get("convergence_tolerance"),
            n_trial_steps=dictionary.get("n_trial_steps"),
            defect_species=[
                DefectSpecies.from_dict(defect_species)
                for defect_species in dictionary["defect_species"]
            ],
        )
        
    def defect_species_by_name(self, name: str) -> DefectSpecies:
        """return a ``DefectSpecies`` contained within the ``DefectSystem``
        via its name.

        Args:
            name (str): name of the ``DefectSpecies`` to return

        Returns:
            DefectSpecies: ``DefectSpecies`` where ``DefectSpecies.name == name``
        """
        return [ds for ds in self.defect_species if ds.name == name][0]
    
    def get_sc_fermi(self) -> Tuple[float, float]:
        """Calculate the self-consistent Fermi energy.
        
        Finds the Fermi energy at which charge neutrality is achieved,
        using Brent's method for root finding.
        
        Returns:
            Tuple[float, float]: The self-consistent Fermi energy and the
            absolute residual charge density at that energy.
        
        Raises:
            RuntimeError: If no solution is found within the DOS energy range.
        """
        emin = self.dos.emin()
        emax = self.dos.emax()
        
        try:
            kwargs = {}
            if self.convergence_tolerance is not None:
                kwargs['xtol'] = self.convergence_tolerance
            if self.n_trial_steps is not None:
                kwargs['maxiter'] = self.n_trial_steps
            
            e_fermi = brentq(
                self.q_tot,
                emin,
                emax,
                **kwargs,
            )  # type: ignore[call-overload]
        except ValueError:
            raise RuntimeError(f"No solution found between {emin} and {emax}")
        
        residual = abs(self.q_tot(e_fermi))
        return e_fermi, residual

    def report(self) -> None:
        """print a report in the style of `SC-Fermi <https://github.com/jbuckeridge/sc-fermi>`_
        which summarises key properties of the defect system."""
        print(self._get_report_string())

    def _get_report_string(self) -> str:
        """generate string to facilitate self.report()"""
        string = ""
        e_fermi = self.get_sc_fermi()[0]
        string += f"Temperature :      {self.temperature}  (K)\n"
        string += f"SC Fermi level :      {e_fermi}  (eV)\n"
        p0, n0 = self.dos.carrier_concentrations(e_fermi, self.temperature)
        string += "Concentrations:\n"
        string += f"n (electrons)  : {n0 * 1e24 / self.volume} cm^-3\n"
        string += f"p (holes)      : {p0 * 1e24 / self.volume} cm^-3\n"
        for ds in self.defect_species:
            concall = ds.get_concentration(e_fermi, self.temperature)
            if ds.fixed_concentration == None:
                string += f"{ds.name:9}      : {concall * 1e24 / self.volume} cm^-3, (percentage of defective sites: {(concall / ds.nsites) * 100:.3} %)\n"
            else:
                string += (
                    f"{ds.name:9}      : {concall * 1e24 / self.volume} cm^-3 [fixed]\n"
                )
        string += "\nBreakdown of concentrations for each defect charge state:\n"
        for ds in self.defect_species:
            concall = ds.get_concentration(e_fermi, self.temperature)
            string += "---------------------------------------------------------\n"
            if concall == 0.0:
                string += f"{ds.name:11}: Zero total - cannot give breakdown\n"
                continue
            string += f"{ds.name:11}: Charge Concentration(cm^-3) Total\n"
            for q, conc in ds.charge_state_concentrations(
                e_fermi, self.temperature
            ).items():
                if ds.charge_states[q].fixed_concentration:
                    fix_str = " [fixed]"
                else:
                    fix_str = ""

                string += f"           : {q: 1}  {conc * 1e24 / self.volume:5e}          {(conc * 100 / concall):.2f} {fix_str}\n"
        return string

    def total_defect_charge_contributions(self, e_fermi: float) -> Tuple[float, float]:
        """
        Calculate the charge contributions from each ``DefectSpecies`` in all charge
        states to the total charge density

        Args:
            e_fermi (float): Fermi energy

        Returns:
            Tuple[float, float]: charge contributions of positive (lhs) and
            negative (rhs) charge states of all defects
        """
        contrib = np.array(
            [
                ds.defect_charge_contributions(e_fermi, self.temperature)
                for ds in self.defect_species
            ]
        )
        lhs = np.sum(contrib[:, 0])
        rhs = np.sum(contrib[:, 1])
        return lhs, rhs

    def q_tot(self, e_fermi: float) -> float:
        """for a given Fermi energy, calculate the net charge density of the
        ``DefectSystem`` as the difference between charge contributions from all
        positive species (including holes) and all negative species (including
        electrons).

        Args:
            e_fermi (float): Fermi energy

        Returns:
            float: net charge density of the ``DefectSystem`` at ``e_fermi``
        """
        p0, n0 = self.dos.carrier_concentrations(e_fermi, self.temperature)
        lhs_def, rhs_def = self.total_defect_charge_contributions(e_fermi)
        lhs = p0 + lhs_def
        rhs = n0 + rhs_def
        diff = rhs - lhs
        return diff

    def get_transition_levels(self) -> Dict[str, List[List]]:
        """Return transition_levels transition levels profiles of all ``DefectSpecies``
        all defects as dictionary of ``{DefectSpecies.name : [e_fermi, e_formation]}``
        over the whole density of states energy range.

        Returns:
            Dict[str, List[List]]: Dictionary giving per-defect transition-level
            profiles.
        """
        transition_levels = {}
        for defect_species in self.defect_species_names:
            transition_level = self.defect_species_by_name(defect_species).tl_profile(
                self.dos.emin(), self.dos.emax()
            )
            x = [[x_value][0][0] for x_value in transition_level]
            y = [[y_value][0][1] for y_value in transition_level]
            transition_levels.update({defect_species: [x, y]})
        return transition_levels

    def concentration_dict(
        self,
        decomposed: bool = False,
        per_volume: bool = True,
    ) -> Dict[str, Any]:
        """Returns a dictionary of the properties of the ``DefectSystem`` object
        after solving for the self-consistent Fermi energy.

        Args:
            decomposed (bool, optional): if True, return a dictionary in which the
              concentration of each ``DefectChargeState`` is given explicitly,
              rather than as a sum over all ``DefectChargeState`` objects in the
              each ``DefectSpecies``. Defaults to False.
            per_volume (bool, optional): if True, return concentrations in units
              of cm^-3, else returns concentration per unit cell. Defaults to True.

        Returns:
            Dict[str, Any]: dictionary specifying the Fermi Energy,
            hole concentration (``"p0"``), electron concentration
            (``"n0"``), temperature, and the defect concentrations.
        """
        if per_volume == True:
            scale = 1e24 / self.volume
        else:
            scale = 1

        e_fermi = self.get_sc_fermi()[0]
        p0, n0 = self.dos.carrier_concentrations(e_fermi, self.temperature)
        run_stats = {
            "Fermi Energy": float(e_fermi),
            "p0": float(p0 * scale),
            "n0": float(n0 * scale),
        }
        if decomposed == False:
            sum_concs = {
                str(ds.name): float(
                    ds.get_concentration(e_fermi, self.temperature) * scale
                )
                for ds in self.defect_species
            }
            return {**run_stats, **sum_concs}
        else:
            decomp_concs = {
                str(ds.name): {
                    int(q): float(
                        ds.charge_state_concentrations(e_fermi, self.temperature)[q]
                        * scale
                    )
                    for q in ds.charge_states
                }
                for ds in self.defect_species
            }
            return {**run_stats, **decomp_concs}

    def site_percentages(
        self, 
    ) -> Dict[str, float]:
        """Returns a dictionary of the DefectSpecies in the DefectSystem which
        giving the percentage of the sites in the structure that will host that 
        defect.

        Returns:
            Dict[str, Any]: dictionary specifying the per-DefectSpecies site
            concentrations.
        """

        e_fermi = self.get_sc_fermi()[0]

        sum_concs = {
                str(ds.name): float(
                    (ds.get_concentration(e_fermi, self.temperature) / ds.nsites) * 100
                )
                for ds in self.defect_species
            }
        return sum_concs

    def as_dict(self) -> dict:
        """

        Returns:
            dict: _description_
        """

        defect_system_dict = dict(
            volume=float(self.volume),
            temperature=float(self.temperature),
            defect_species=[
                defect_species.as_dict() for defect_species in self.defect_species
            ],
            dos=self.dos.as_dict(),
        )
        if self.convergence_tolerance is not None:
            defect_system_dict["convergence_tolerance"] = self.convergence_tolerance
        if self.n_trial_steps is not None:
            defect_system_dict["n_trial_steps"] = self.n_trial_steps
        return defect_system_dict
