import logging
import pandas as pd
import numpy as np

from tardis.plasma.properties.base import ProcessingPlasmaProperty

logger = logging.getLogger(__name__)

__all__ = ['LevelNumberDensity', 'LevelNumberDensityHeNLTE']


class LevelNumberDensity(ProcessingPlasmaProperty):
    """
    Attributes:
    level_number_density : Pandas DataFrame, dtype float
                           Index atom number, ion number, level number.
                           Columns are zones.
    """
    outputs = ('level_number_density',)
    latex_name = ('N_{i,j,k}',)
    latex_formula = ('N_{i,j}\\dfrac{bf_{i,j,k}}{Z_{i,j}}',)

    def __init__(self, plasma_parent):
        """
        Calculates the level populations with the Boltzmann equation in LTE.
        """
        super(LevelNumberDensity, self).__init__(plasma_parent)
        self._update_inputs()
        self.initialize_indices = True

    def _initialize_indices(self, levels, partition_function):
        indexer = pd.Series(np.arange(partition_function.shape[0]),
                            index=partition_function.index)
        self._ion2level_idx = indexer.ix[levels.droplevel(2)].values

    def calculate(
            self, level_boltzmann_factor, ion_number_density,
            levels, partition_function):
        """
        Reduces non-metastable level populations by a factor of W
        compared to LTE in the case of dilute-lte excitation.
        """
        if self.initialize_indices:
            self._initialize_indices(levels, partition_function)
            self.initialize_indices = False
        partition_function_broadcast = partition_function.values[
            self._ion2level_idx]
        level_population_fraction = (level_boltzmann_factor.values /
                                     partition_function_broadcast)
        ion_number_density_broadcast = ion_number_density.values[
            self._ion2level_idx]
        level_number_density = (level_population_fraction *
                                ion_number_density_broadcast)
        return pd.DataFrame(level_number_density,
                            index=level_boltzmann_factor.index)


class LevelNumberDensityHeNLTE(LevelNumberDensity):
    """
    Attributes:
    level_number_density : Pandas DataFrame, dtype float
                           Index atom number, ion number, level number.
                           Columns are zones.
    """

    def calculate(
            self, level_boltzmann_factor,
            ion_number_density, levels, partition_function,
            helium_population_updated):
        """
        If one of the two helium NLTE methods is used, this updates
        the helium level populations to the appropriate
        values.
        """
        level_number_density = super(LevelNumberDensityHeNLTE, self).calculate(
            level_boltzmann_factor, ion_number_density, levels,
            partition_function)
        if helium_population_updated is not None:
            level_number_density.ix[2].update(helium_population_updated)
        return level_number_density

@@ -269,32 +269,70 @@
 perty):%0A    %22%22%22%0A
+    Calculates the level populations%0A%0A
     Attributes:%0A
@@ -323,32 +323,32 @@
     Attributes:%0A
-
     level_number
@@ -691,109 +691,8 @@
 t):%0A
-        %22%22%22%0A        Calculates the level populations with the Boltzmann equation in LTE.%0A        %22%22%22%0A
     
@@ -1082,32 +1082,33 @@
 values%0A%0A    def 
+_
 calculate(%0A     
@@ -1092,32 +1092,43 @@
   def _calculate
+_dilute_lte
 (%0A            se
@@ -1240,65 +1240,72 @@
     
-Reduces non-metastable level populations by a 
+Calculate the level populations from the level_boltzmann_
 factor
- of W
+,
 %0A   
@@ -1313,61 +1313,49 @@
     
-compared to LTE in the case of dilute-lte excita
+ion_number_density and partition_func
 tion
-.
 %0A   
@@ -2079,16 +2079,55 @@
 index)%0A%0A
+    calculate = _calculate_dilute_lte%0A%0A
 %0Aclass L
@@ -2738,47 +2738,15 @@
 ty =
- super(LevelNumberDensityHeNLTE,
  self
-)
 .
+_
 calc
@@ -2742,32 +2742,43 @@
  self._calculate
+_dilute_lte
 (%0A            le

Fix docstrings