#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Feb 17 15:24:29 2021

@author: villa
"""
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import pandas as pd

from .defects import Defect, format_legend_with_charge_number, get_defect_from_string


def plot_formation_energies(entries,
                            chemical_potentials,
                            vbm,
                            band_gap,
                            temperature=0,
                            xlim=None,
                            ylim=None,
                            title=None,
                            fermi_level=None,
                            grid=True,
                            figsize=(6,6),
                            fontsize=12,
                            colors=None,
                            show_legend=True,
                            format_legend=True,
                            ax=None,
                            **eform_kwargs):
    """
    Produce defect Formation energy vs Fermi energy plot.

    Parameters
    ----------
    entries : list 
        List of entries to calculate.
    chemical_potentials : dict
        Dictionary with chemical potentials of the elements {'element':chempot}.
    vbm : float
        Valence band maximum in eV.
    band_gap : float
        Band gap of bulk material in eV.
    temperature : float
        Temperature in K. If no custom formation energy is provided, this arg has no effect. 
    xlim : tuple
        Tuple giving the range of the x (fermi energy) axis.
    ylim : tuple
        Tuple giving the range for the formation energy axis.
    title : str
        Title of the figure.
    fermi_level : float
        Plot Fermi energy position with a vertical line.
    grid : bool
        Show grid.
    figsize : float or tuple
        Figure size.
    fontsize : float
        Font size.
    colors : list
        List of colors for line plot.
    show_legend  : bool
        Show legend.
    format_legend : bool
        Get latex-like legend based on the name of defect entries.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.
    eform_kwargs : dict
        Kwargs to pass to `entry.formation_energy`.

    Returns
    -------
    matplotlib.axes.Axes

    """
    from .analysis import DefectsAnalysis
    import matplotlib.pyplot as plt

    formation_energies = DefectsAnalysis(
                        entries=entries,
                        band_gap=band_gap,
                        vbm=vbm,
                        sort_entries=False).formation_energies(
                            chemical_potentials=chemical_potentials,
                            fermi_level=0,
                            temperature=temperature,
                            entries=entries,
                            **eform_kwargs)

    if xlim is None:
        xlim = (-0.5, band_gap + 0.5)

    npoints = 200
    step = abs(xlim[1] + 0.1 - xlim[0]) / npoints
    x = np.arange(xlim[0], xlim[1] + 0.1, step)

    ax = _get_ax(ax, figsize=figsize)

    for idx, name in enumerate(formation_energies):
        energy = np.zeros(len(x))
        emin = np.zeros(len(x))
        x_star = []
        y_star = []
        q_previous = None

        for i in range(npoints):
            emin[i] = 1e40
            for d in formation_energies[name]:
                q = d[0]
                e0 = d[1]
                energy[i] = e0 + q * x[i]
                if energy[i] < emin[i]:
                    emin[i] = energy[i]
                    q_stable = q
            if q_stable != q_previous:
                if q_previous is not None:
                    x_star.append(x[i])
                    y_star.append(emin[i])
                q_previous = q_stable

        label_txt = get_defect_from_string(name).symbol if format_legend else name
        color = colors[idx] if colors else None
        ax.plot(x, emin, label=label_txt, linewidth=3, color=color)
        ax.scatter(x_star, y_star, s=120, marker='*', color=color)

    ax.axvline(x=0.0, linestyle='-', color='k', linewidth=2)
    ax.axvline(x=band_gap, linestyle='-', color='k', linewidth=2)
    if fermi_level:
        ax.axvline(x=fermi_level, linestyle='dashed', color='k', linewidth=1.5, label='$\\mu _{e}$')
    ax.axvspan(xlim[0], 0, facecolor='k', alpha=0.2)
    ax.axvspan(band_gap, xlim[1] + 0.1, facecolor='k', alpha=0.2)
    ax.hlines(0, xlim[0], xlim[1] + 0.1, colors='k', linestyles='dashed', alpha=0.5)

    ax.set_xlim(xlim)
    if ylim:
        ax.set_ylim(ylim)

    ax.set_xlabel('Fermi level (eV)')
    ax.set_ylabel('Formation energy (eV)')

    if title:
        ax.set_title(title)

    #xlim = ax.get_xlim()
    ax.text(xlim[0] + 0.25, 0.5, 'VB', fontsize=fontsize * 1.35, horizontalalignment='center')
    ax.text(xlim[1] - 0.25, 0.5, 'CB', fontsize=fontsize * 1.35, horizontalalignment='center')

    _style_ax(ax=ax, fontsize=fontsize, legend=show_legend, grid=grid)

    return ax



def plot_binding_energies(entries,
                          vbm,
                          band_gap,
                          temperature=0,
                          names=None,
                          xlim=None,
                          ylim=None,
                          figsize=(6,6),
                          fontsize=None,
                          colors=None,
                          format_legend=True,
                          ax=None,
                          **eform_kwargs):
    """
    Plot binding energies for complex of defects as a function of the fermi level

    Parameters
    ----------
    entries : list 
        List of entries to calculate.
    vbm : float
        Valence band maximum in eV.
    band_gap : float
        Band gap of bulk material in eV.
    temperature : float
        Temperature in K. If no custom formation energy is provided, this arg has no effect.
    names : list
        List of strings with names of DefectEntry. If None, all defect complexes are plotted.
    xlim : tuple
        Tuple giving the range of the x (fermi energy) axis
    ylim : tuple
        Tuple giving the range for the formation energy axis
    figsize : tuple
        Figure size.
    fontsize : float
        Font size.
    colors : list
        List of colors for line plot.
    format_legend : bool
        Bool for getting latex-like legend based on the name of defect entries.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.
    eform_kwargs : dict
        Kwargs to pass to `entry.formation_energy`.

    Returns
    -------
    matplotlib.axes.Axes

    """   
    ax = _get_ax(ax, figsize=figsize)

    from .analysis import DefectsAnalysis
    da = DefectsAnalysis(entries=entries,band_gap=band_gap,vbm=vbm,sort_entries=False)
        

    if xlim is None:
        xlim = (-0.5, da.band_gap + 0.5)

    ef = np.arange(xlim[0], xlim[1] + 0.1, (xlim[1] - xlim[0]) / 200)
    binding_energy = np.zeros(len(ef))

    if not names:
        names = []
        for e in da.entries:
            if e.defect_type == 'DefectComplex' and e.name not in names:
                names.append(e.name)

    if not names:
        raise ValueError('No DefectComplex entries found')

    for idx, name in enumerate(names):
        label = da.select_entries(names=[name])[0].symbol if format_legend else name

        for i in range(len(ef)):
            binding_energy[i] = da.binding_energy(
                name=name,
                fermi_level=ef[i],
                temperature=temperature,
                **eform_kwargs)

        color = colors[idx] if colors else None
        ax.plot(ef, binding_energy,
                linewidth=2.5*(figsize[1]/figsize[0]),
                label=label,
                color=color)

    ax.axvline(x=0.0, linestyle='-', color='k', linewidth=2)
    ax.axvline(x=da.band_gap, linestyle='-', color='k', linewidth=2)

    ax.axvspan(xlim[0], 0, facecolor='k', alpha=0.2)
    ax.axvspan(da.band_gap, xlim[1], facecolor='k', alpha=0.2)
    ax.hlines(0, xlim[0], xlim[1], colors='k', linestyles='dashed', alpha=0.5)

    ax.set_xlim(xlim)

    if ylim:
        ax.set_ylim(ylim)

    ax.set_xlabel('Fermi level (eV)')
    ax.set_ylabel('Binding energy (eV)')

    _style_ax(ax=ax, fontsize=fontsize, legend=True, grid=False)

    return ax



def plot_charge_transition_levels(entries,
                                  vbm,
                                  band_gap,
                                  temperature=0,
                                  ylim=None,
                                  figsize=(6,6),
                                  fontsize=None,
                                  colors=None,
                                  fermi_level=None,
                                  format_legend=True,
                                  get_integers=True,
                                  ax=None,
                                  **eform_kwargs):
    """
    Plotter for the charge transition levels.

    Parameters
    ----------
    entries : list
        List of entries to calculate.
    vbm : float
        Valence band maximum in eV.
    band_gap : float
        Band gap of bulk material in eV.
    temperature : float
        Temperature in K. If no custom formation energy is provided, this arg has no effect.
    ylim : tuple
        y-axis limits.
    figsize : tuple
        Figure size.
    fontsize : float
        Font size.
    colors : list or str
        Colors for CTL lines and charge labels. Provide a string to use the same color for all species.
    fermi_level : float
        Plot Fermi energy position.
    format_legend : bool
        Bool for getting latex-like legend based on the name of defect entries.
    get_integers : bool
        Get charge transition levels as integers.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.
    eform_kwargs : dict
        Kwargs to pass to `entry.formation_energy`.

    Returns
    -------
    matplotlib.axes.Axes

    """  
    from .analysis import DefectsAnalysis
    
    ax = _get_ax(ax, figsize=figsize)
    
    da = DefectsAnalysis(entries=entries,band_gap=band_gap,vbm=vbm,sort_entries=False)
    
    default_colors = matplotlib.color_sequences['tab10'] + \
                     matplotlib.color_sequences['tab20'] + \
                     matplotlib.color_sequences['Pastel1']
    if type(colors) == str:
        colors = [colors] * len(default_colors)
    elif not colors:
        colors = default_colors

    if ylim == None:
        ylim = (-0.5,da.band_gap +0.5)        
    charge_transition_levels = da.charge_transition_levels(
                                                temperature=temperature,
                                                entries=entries,
                                                get_integers=get_integers,
                                                **eform_kwargs)
    
    number_defects = len(charge_transition_levels)   
    x_max = 10
    interval = x_max/(number_defects + 1)
    x = np.arange(0,x_max,interval)

    x_ticks_positions = []
    for i in range(0,len(x)-1):
        x_ticks_positions.append((x[i+1]-x[i])/2 + x[i])            

    x_ticks_labels = []
    for name in charge_transition_levels:
        x_ticks_labels.append(name)        

    for i in x:
        ax.axvline(x=i, linestyle='-', color='k', linewidth=1.2, alpha=1, zorder=1)

    xlim = (x[0],x[-1])        

    ax.axhspan(ylim[0], 0, facecolor='grey', alpha=0.5, zorder=2)
    ax.axhspan(da.band_gap,ylim[1], facecolor = 'grey', alpha=0.5, zorder=2)

    for i in range(0,len(x_ticks_labels)):
        name = x_ticks_labels[i]
        color = colors[i] if colors else 'black'
        for ctl in charge_transition_levels[name]:
            energy = ctl[2]
            ax.hlines(energy,x[i],x[i+1],colors=color,linewidth=2.25, zorder=3)

            charge1 = '+' + str(ctl[1]) if ctl[1] > 0 else str(ctl[1])
            charge2 = '+' + str(ctl[0]) if ctl[0] > 0 else str(ctl[0])
            label_charge = '(' + charge2 + '/' + charge1 + ')'

            font_space = abs(ylim[1]-ylim[0]) / 100

            if energy < ylim[1] and energy > ylim[0]:
                ax.text(
                        x[i]+(interval/2)*2/number_defects,
                        energy+font_space,
                        label_charge,
                        fontsize=fontsize,
                        color=color)

    if format_legend:    
         for name in x_ticks_labels:            
            x_ticks_labels[x_ticks_labels.index(name)] = get_defect_from_string(name).symbol               

    if fermi_level:
        ax.axhline(y=fermi_level, linestyle='dashed', color='k', linewidth=1.5, label='$\\mu _{e}$')   
    
    ax.text(x[-1]-0.75,-0.3,'VB',fontsize=25*(fontsize/16))
    ax.text(x[-1]-0.75,da.band_gap+0.2,'CB',fontsize=25*(fontsize/16))

    ax.set_xticks(x_ticks_positions)
    ax.set_xticklabels(x_ticks_labels,fontsize=(25-number_defects)*(fontsize/16))

    ax.tick_params(axis='x',length=0,width=0)
    ax.tick_params(axis='y',labelsize=16*(fontsize/16))

    ax.set_xlim(xlim)  
    ax.set_ylim(ylim)

    ax.set_ylabel('Energy(eV)',fontsize=20*(fontsize/16))  

    _style_ax(ax=ax, fontsize=fontsize, legend=False, grid=False)

    return ax
    


def plot_pO2_vs_concentrations(
                            thermodata,
                            output='total',
                            figsize=(8,8),
                            fontsize=22,
                            xlim=(1e-20,1e10),
                            ylim=None,
                            colors=None,
                            ax=None,
                            **kwargs):
    """
    Plot defect and carrier concentrations in a range of oxygen partial pressure.

    Parameters
    ----------
    thermodata: ThermoData
        ThermoData object that contains the thermodynamic data:

        partial_pressures : (list)
            list with partial pressure values.
        defect_concentrations : (list or dict)
            Defect concentrations in the same format as the output of DefectsAnalysis. 
        carrier_concentrations : (list)
            List of tuples with intrinsic carriers concentrations (holes,electrons).

    output : str
        Type of output for defect concentrations:

        - "all": The output is the concentration of every defect entry.
        - "stable": The output is the concentration of the stable charge for every defect at each fermi level point.
        - "total": The output is the sum of the concentration in every charge for each specie.

    figsize : tuple
        Size of the matplotlib figure.
    fontsize : float
        Size of font for matplotlib rcParams.
    xlim : tuple
        Range of x-axis.
    ylim : tuple
        Range of y-axis. 
    colors : list
        List of colors to use for plotting with matplotlib. If None the defaults are used.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.
    kwargs : dict
        Kwargs to pass to `DefectConcentrations.filter_concentrations(**kwargs)`.
        If provided, only the filtered concentrations will be plotted. If output
        is set to "total", only the filtered concentrations will be used to 
        compute the total concentration.

    Returns
    -------
    matplotlib.axes.Axes

    """
    td = thermodata
    p,dc,cc = td.partial_pressures,td.defect_concentrations,td.carrier_concentrations
    xlabel = 'Oxygen partial pressure (atm)'
    ax = plot_x_vs_concentrations(
                                x=p,
                                xlabel=xlabel,
                                defect_concentrations=dc,
                                carrier_concentrations=cc,
                                output=output,
                                figsize=figsize,
                                fontsize=fontsize,
                                xlim=xlim,
                                ylim=ylim,
                                colors=colors,
                                ax=ax,
                                **kwargs)
    return ax


def plot_pO2_vs_conductivity(
                            partial_pressures,
                            conductivities,
                            label=None,
                            figsize=(8,8),
                            fontsize=22,
                            xlim=(1e-20,1e10),
                            ylim=None,
                            ax=None):
    """
    Plot conductivity as a function of the oxygen partial pressure.

    Parameters
    ----------
    partial_pressures : list
        list with partial pressure values.
    conductivities : dict or list
        If is a dict multiples lines will be plotted, with labels as keys and conductivity list
        as values. If is a list only one line is plotted with label taken from the "label" argument.
    label : str
        Label for the data.
    figsize : tuple
        Size of the matplotlib figure.
    fontsize : float
        Size of font for matplotlib rcParams. 
    xlim : tuple
        Range of x-axis.
    ylim : tuple
        Range of y-axis.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.
        
    Returns
    -------
    matplotlib.axes.Axes

    """
    xlabel = 'Oxygen partial pressure (atm)'
    ax = plot_x_vs_conductivity(
                                x=partial_pressures,
                                xlabel=xlabel,
                                conductivities=conductivities,
                                label=label,
                                figsize=figsize,
                                fontsize=fontsize,
                                xlim=xlim,
                                ylim=ylim,
                                ax=ax)
    return ax    



def plot_pO2_vs_fermi_level(
                            partial_pressures,
                            fermi_levels,
                            band_gap,
                            label=None,
                            figsize=(8,8),
                            fontsize=22,
                            xlim=(1e-20,1e10),
                            ylim=None,
                            colors=None,
                            ax=None):
    """
    Plot Fermi level as a function of the oxygen partial pressure.

    Parameters
    ----------
    partial_pressures : list
        list with partial pressure values.
    fermi_levels : dict or list
        If is a dict multiples lines will be plotted, with labels as keys and fermi level list
        as values. If is a list only one line is plotted with label taken from the "label" argument.
    band_gap : float
        Band gap of the bulk material.
    label : str
        Label for the data.
    figsize : tuple
        Size of the matplotlib figure.
    fontsize : float
        Size of font for matplotlib rcParams. 
    xlim : tuple
        Range of x-axis.
    ylim : tuple
        Range of y-axis.
    colors : list
        List with colors for Fermi level data.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.

    Returns
    -------
    matplotlib.axes.Axes

    """
    xlabel = 'Oxygen partial pressure (atm)'
    ax = plot_x_vs_fermi_level(
                                x=partial_pressures,
                                xlabel=xlabel,
                                fermi_levels=fermi_levels,
                                band_gap=band_gap,
                                label=label,
                                figsize=figsize,
                                fontsize=fontsize,
                                xlim=xlim,
                                ylim=ylim,
                                colors=colors,
                                ax=ax)
    return ax


def plot_variable_species_vs_concentrations(
                                        thermodata,
                                        output='total',
                                        figsize=(8,8),
                                        fontsize=22,
                                        xlim=(1e-20,1e10),
                                        ylim=None,
                                        colors=None,
                                        ax=None,
                                        **kwargs):
    """
    Plot defect and carrier concentrations in a range of oxygen partial pressure.

    Parameters
    ----------
    thermodata: ThermoData
        ThermoData object that contains the thermodynamic data:

        partial_pressures : (list)
            list with partial pressure values.
        defect_concentrations : (list or dict)
            Defect concentrations in the same format as the output of DefectsAnalysis. 
        carrier_concentrations : (list)
            List of tuples with intrinsic carriers concentrations (holes,electrons).

    output : str
        Type of output for defect concentrations:

        - "all": The output is the concentration of every defect entry.
        - "stable": The output is the concentration of the stable charge for every defect at each fermi level point.
        - "total": The output is the sum of the concentration in every charge for each specie.

    figsize : tuple
        Size of the matplotlib figure.
    fontsize : float
        Size of font for matplotlib rcParams.
    xlim : tuple
        Range of x-axis.
    ylim : tuple
        Range of y-axis.
    colors : list
        List of colors to use for plotting with matplotlib. If None the defaults are used.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.
    kwargs : dict
        Kwargs to pass to `DefectConcentrations.filter_concentrations(**kwargs)`.
        If provided, only the filtered concentrations will be plotted. If output
        is set to "total", only the filtered concentrations will be used to 
        compute the total concentration.

    Returns
    -------
    matplotlib.axes.Axes
    """
    td = thermodata
    c,dc,cc = td.variable_concentrations,td.defect_concentrations,td.carrier_concentrations
    dname = _get_variable_defect_specie_label(td.variable_defect_specie)
    xlabel = '[%s] (cm$^{-3})$' %dname
    
    ax = plot_x_vs_concentrations(
                                x=c,
                                xlabel=xlabel,
                                defect_concentrations=dc,
                                carrier_concentrations=cc,
                                output=output,
                                figsize=figsize,
                                fontsize=fontsize,
                                xlim=xlim,
                                ylim=ylim,
                                colors=colors,
                                ax=ax,
                                **kwargs)
    return ax


def plot_variable_species_vs_conductivity(
                                        xlabel,
                                        variable_concentrations,
                                        conductivities,
                                        label=None,
                                        figsize=(8,8),
                                        fontsize=22,
                                        xlim=(1e-20,1e10),
                                        ylim=None,
                                        ax=None):
    """
    Plot conductivity as a function of the oxygen partial pressure.

    Parameters
    ----------
    xlabel : str
        Label for concentration axis (cm^-3 is added).
    variable_concentrations : list
        List of concentrations of variable species. 
    conductivities : dict or list
        If is a dict multiples lines will be plotted, with labels as keys and conductivity list
        as values. If is a list only one line is plotted with label taken from the "label" argument.
    label : str
        Label for the data.
    figsize : tuple
        Size of the matplotlib figure.
    fontsize : float
        Size of font for matplotlib rcParams. 
    xlim : tuple
        Range of x-axis.
    ylim : tuple
        Range of y-axis.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.

    Returns
    -------
    matplotlib.axes.Axes
    """
    xlabel = '%s (cm$^{-3})$' %xlabel
    ax = plot_x_vs_conductivity(
                                x=variable_concentrations,
                                xlabel=xlabel,
                                conductivities=conductivities,
                                label=label,
                                figsize=figsize,
                                fontsize=fontsize,
                                xlim=xlim,
                                ylim=ylim,
                                ax=ax)
    return ax
    

def plot_variable_species_vs_fermi_level(
                                        xlabel,
                                        variable_concentrations,
                                        fermi_levels,
                                        band_gap,
                                        label=None,
                                        figsize=(8,8),
                                        fontsize=22,
                                        xlim=(1e-20,1e10),
                                        ylim=None,
                                        colors=None,
                                        ax=None):
    """
    Plot Fermi level as a function of the oxygen partial pressure.

    Parameters
    ----------
    xlabel : str
        Label for concentration axis (cm^-3 is added).
    variable_concentrations : list
        List of concentrations of variable species. 
    fermi_levels : dict or list
        If is a dict multiples lines will be plotted, with labels as keys and fermi level list
        as values. If is a list only one line is plotted with label taken from the "label" argument.
    band_gap : float
        Band gap of the bulk material.
    label : str
        Label for the data.
    figsize : tuple
        Size of the matplotlib figure.
    fontsize : float
        Size of font for matplotlib rcParams. 
    xlim : tuple
        Range of x-axis.
    ylim : tuple
        Range of y-axis.
    colors : list
        List with colors for Fermi level data.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.

    Returns
    -------
    matplotlib.axes.Axes

    """
    xlabel = '[%s] (cm$^{-3})$' %xlabel
    ax = plot_x_vs_fermi_level(
                                x=variable_concentrations,
                                xlabel=xlabel,
                                fermi_levels=fermi_levels,
                                band_gap=band_gap,
                                label=label,
                                figsize=figsize,
                                fontsize=fontsize,
                                xlim=xlim,
                                ylim=ylim,
                                colors=colors,
                                ax=ax)
    return ax



def plot_x_vs_concentrations(
                            x,
                            xlabel,
                            defect_concentrations,
                            carrier_concentrations,
                            output='total',
                            figsize=(8,8),
                            fontsize=14,
                            xlim=(1e-20,1e10),
                            ylim=None,
                            colors=None,
                            ax=None,
                            **kwargs):
    """
    Plot defect concentrations as a function of generic data on the x-axis.        
    
    Parameters
    ----------
    x : list
        List with data on x-axis.
    xlabel : str
        Label for x-axis.
    defect_concentrations : list
        List of DefectConcentrations objects.
    carrier_concentrations : list
        List of tuples with carrier concentrations (holes,electrons).
    output : str
        Type of output for defect concentrations:

        - "all": The output is the concentration of every defect entry.
        - "stable": The output is the concentration of the stable charge for every defect at each fermi level point.
        - "total": The output is the sum of the concentration in every charge for each specie.
        
    figsize : tuple
        Size of the matplotlib figure.
    fontsize : float
        Size of font for matplotlib rcParams. 
    xlim : tuple
        Range of x-axis. 
    ylim : tuple
        Range of y-axis.
    colors : list
        List of colors to use for plotting with matplotlib. If None the defaults are used.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.
    kwargs : dict
        Kwargs to pass to `DefectConcentrations.filter_concentrations(**kwargs)`.
        If provided, only the filtered concentrations will be plotted. If output
        is set to "total", only the filtered concentrations will be used to 
        compute the total concentration.

    Returns
    -------
    matplotlib.axes.Axes

    """
    ax = _get_ax(ax,figsize=figsize)        
    if output == 'all' or output == 'stable':
        ax = _plot_conc(
                        x=x,
                        defect_concentrations=defect_concentrations,
                        carrier_concentrations=carrier_concentrations,
                        output=output,figsize=figsize,colors=colors,
                        ax=ax,**kwargs)
        
    elif output == 'total':
        ax = _plot_conc_total(
                            x=x,
                            defect_concentrations=defect_concentrations,
                            carrier_concentrations=carrier_concentrations,
                            figsize=figsize,colors=colors,ax=ax,**kwargs)
    else:
        raise ValueError('The options for plot output are "all", "stable" or "total".')
        
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_xlim(xlim)
    if ylim:
        ax.set_ylim(ylim)

    ax.set_xlabel(xlabel)
    ax.set_ylabel('Concentrations (cm$^{-3})$')
    _style_ax(ax=ax,fontsize=fontsize)
    
    return ax   

    

def plot_x_vs_conductivity(
                        x,
                        xlabel,
                        conductivities,
                        label=None,
                        figsize=(8,8),
                        fontsize=22,
                        xlim=(1e-20,1e10),
                        ylim=None,
                        ax=None):
    """
    Plot conductivity as a function of the oxygen partial pressure.

    Parameters
    ----------
    x : list
        List with data on x-axis.
    xlabel : str
        Label for x-axis.
    conductivities : dict or list
        If is a dict multiples lines will be plotted, with labels as keys and conductivity list
        as values. If is a list only one line is plotted with label taken from the "label" argument.
    label : str
        Label for the data.
    figsize : tuple
        Size of the matplotlib figure.
    fontsize : float
        Size of font for matplotlib rcParams. 
    xlim : tuple
        Range of x-axis.
    ylim : tuple
        Range of y-axis.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.

    Returns
    -------
    matplotlib.axes.Axes

    """
    ax = _get_ax(ax, figsize=figsize)

    if isinstance(conductivities,dict):
        for name,sigma in conductivities.items():
            ax.plot(x,sigma,linewidth=4,marker='s',label=name)
    else:
        sigma = conductivities
        ax.plot(x,sigma,linewidth=4,marker='s',label=label)

    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_xlim(xlim)
    if ylim:
        ax.set_ylim(ylim)

    ax.set_xlabel(xlabel)
    ax.set_ylabel('Conductivity (S/m)')

    legend = True if isinstance(conductivities, dict) else False
    _style_ax(ax=ax, fontsize=fontsize, legend=legend, grid=False)

    return ax



def plot_x_vs_fermi_level(
                        x,
                        xlabel,
                        fermi_levels,
                        band_gap,
                        label=None,
                        figsize=(8,8),
                        fontsize=20,
                        xlim=(1e-20,1e10),
                        ylim=None,
                        colors=None,
                        ax=None):
    """
    Parameters
    ----------
    x : list
        List with data on x-axis.
    xlabel : str
        Label for x-axis.
    fermi_levels : dict or list
        If is a dict multiples lines will be plotted, with labels as keys and fermi level list
        as values. If is a list only one line is plotted with label taken from the "label" argument.
    band_gap : float
        Band gap of the bulk material.
    label : str
        Label for the data.
    figsize : tuple
        Size of the matplotlib figure.
    fontsize : float
        Size of font for matplotlib rcParams. 
    xlim : tuple
        Range of x-axis. 
    ylim : tuple
        Range of y-axis.
    colors : list,
        List with colors for Fermi level data.
    ax : matplotlib.axes.Axes
        Axis to plot into. If None, a new figure and axis are created.

    Returns
    -------
    matplotlib.axes.Axes
    """
    ax = _get_ax(ax, figsize=figsize)

    ylim = ylim if ylim else (-0.5, band_gap+0.5)

    if isinstance(fermi_levels, dict):
        for name, mue in fermi_levels.items():
            clr = colors[list(fermi_levels.keys()).index(name)] if colors else None
            ax.plot(x, mue, linewidth=4, label=name, color=clr)
    else:
        mue = fermi_levels
        clr = colors[0] if colors else None
        ax.plot(x, mue, linewidth=4, label=label, color=clr)

    ax.set_xscale('log')
    ax.set_xlim(xlim)
    if ylim:
        ax.set_ylim(ylim)

    ax.hlines(0, xlim[0], xlim[1], color='k')
    ax.text(xlim[1] + (xlim[1]-xlim[0]), -0.05, 'VB',fontsize=fontsize*1.1)
    ax.text(xlim[1] + (xlim[1]-xlim[0]), band_gap - 0.05, 'CB',fontsize=fontsize*1.1)
    ax.hlines(band_gap, xlim[0], xlim[1], color='k')
    ax.axhspan(ylim[0], 0, alpha=0.2, color='k')
    ax.axhspan(band_gap, ylim[1], alpha=0.2, color='k')

    ax.set_xlabel(xlabel)
    ax.set_ylabel('Electron chemical potential (eV)')

    legend = True if isinstance(fermi_levels, dict) else False
    _style_ax(ax, fontsize=fontsize, grid=False, legend=legend)

    return ax


def _get_variable_defect_specie_label(variable_defect_specie):
    try:
        defect = get_defect_from_string(variable_defect_specie)
        return defect.symbol
    except:
        return variable_defect_specie



def _plot_conc(
            x,
            defect_concentrations,
            carrier_concentrations,
            output,
            figsize,
            colors,
            ax,
            **kwargs):
    
    width, height = ax.figure.get_size_inches()
        
    dc = defect_concentrations if output != 'stable' else [c.stable for c in defect_concentrations] 
    if kwargs:
        dc = [c.filter_concentrations(**kwargs) for c in dc] #filter concentrations based on kwargs

    h = [cr[0] for cr in carrier_concentrations] 
    n = [cr[1] for cr in carrier_concentrations]
    previous_charge = None
    for i in range(0,len(dc[0])):
            conc = [c[i].conc for c in dc]
            charges = [c[i].charge for c in dc]
            try:
                label_txt = get_defect_from_string(dc[0][i].name).symbol
            except:
                label_txt = dc[0][i].name
            if output == 'all':
                label_txt = format_legend_with_charge_number(label_txt,dc[0][i].charge)
            elif output == 'stable':
                for q in charges:
                    if q != previous_charge:
                        previous_charge = q
                        label_charge = '+' + str(int(q)) if q > 0 else str(int(q))
                        index = charges.index(q)
                        ax.text(x[index],conc[index],label_charge,clip_on=True,fontsize=width*2.5)
            color = colors[i] if colors else None
            ax.plot(x,conc,label=label_txt,linewidth=4,color=color)
    ax.plot(x,h,label='$n_{h}$',linestyle='--',color='r',linewidth=4)
    ax.plot(x,n,label='$n_{e}$',linestyle='--',color='b',linewidth=4)
    return ax


def _plot_conc_total(
                    x,
                    defect_concentrations,
                    carrier_concentrations,
                    figsize,
                    colors,
                    ax,
                    **kwargs):

    dc = defect_concentrations
    if kwargs:
        dc = [c.filter_concentrations(**kwargs) for c in dc]

    h = [cr[0] for cr in carrier_concentrations] 
    n = [cr[1] for cr in carrier_concentrations]
    for name in dc[0].names:
        conc = [c.total[name] for c in dc]
        try:
            label_txt = get_defect_from_string(name).symbol
        except:
            label_txt = name
        color = colors[dc[0].names.index(name)] if colors else None
        ax.plot(x,conc,label=label_txt,linewidth=4,color=color)

    ax.plot(x,h,label='$n_{h}$',linestyle='--',color='r',linewidth=4)
    ax.plot(x,n,label='$n_{e}$',linestyle='--',color='b',linewidth=4)

    return ax


def _get_unstable_bool(defect_concentrations):
    slist = []
    for dc in defect_concentrations:
        unstable = False
        for d in dc:
            if 'stable' in vars(d).keys() and d.stable == False:
                unstable = True
        slist.append(unstable)
    
    return slist
        

def _get_ax(ax=None, **fig_kwargs):
    if ax is None:
        _, ax = plt.subplots(**fig_kwargs)
    return ax


def _style_ax(ax, fontsize=None, legend=True, grid=True):

    if fontsize is not None:
        ax.set_xlabel(ax.get_xlabel(), fontsize=fontsize)
        ax.set_ylabel(ax.get_ylabel(), fontsize=fontsize)
        ax.tick_params(axis='both', labelsize=fontsize*0.9)
        if legend:
            ax.legend(fontsize=fontsize*0.9)
    else:
        if legend:
            ax.legend()
    if grid:
        ax.grid()
    return