# Path: tilde/berlinium/dos.py
# class TotalDos(Dos):
#     def __init__(self, eigenvalues, sigma=None):
#         Dos.__init__(self, eigenvalues, sigma)
# 
#     def get_density_of_states_at_omega(self, omega):
#         return np.sum( self.smearing_function.calc(self.eigenvalues - omega))
# 
#     def calculate(self):
#         omega = self.omega_min
#         dos = []
# 
#         if self.omega_pitch == 0: self.omega_pitch = 1 # beware of endless loop if omega_pitch=0
#         while omega < self.omega_max + self.omega_pitch/10 :
#             p = self.get_density_of_states_at_omega(omega)
#             dos.append( [round(omega, 3), round(p, 3)] ) # round to reduce output
#             omega += self.omega_pitch
#         return dos
# 
# class PartialDos(Dos):
#     # eigenvalues and impacts must be unique and one-to-one correspondent
#     def __init__(self, eigenvalues, impacts, sigma=None):
#         Dos.__init__(self, eigenvalues, sigma)
#         self.impacts = np.array(impacts)
# 
#     def get_partial_dos_impact_at_omega(self, omega):
#         # function for obtaining scaled partial impacts in any omega point using a linear equation interpolation
#         if omega in self.eigenvalues:
#             return self.impacts[ np.where(self.eigenvalues == omega)[0][0] ]
#         if omega < self.eigenvalues[0]:
#             return np.zeros( len(self.impacts[0]) )
#         elif omega > self.eigenvalues[-1]:
#             return self.impacts[-1]
#         else:
#             for n in range(len(self.eigenvalues)):
#                 if n < len(self.eigenvalues) and self.eigenvalues[n] < omega < self.eigenvalues[n+1]:
#                     result = (omega - self.eigenvalues[n])*(self.impacts[n+1] - self.impacts[n])/(self.eigenvalues[n+1] - self.eigenvalues[n]) + self.impacts[n]
#                     return result
# 
#     def calculate(self, types, labels):
#         omega = self.omega_min
#         pdos = []
#         omegas = []
# 
#         plots = []
# 
#         if self.omega_pitch == 0: self.omega_pitch = 1 # beware of endless loop if omega_pitch=0
#         while omega < self.omega_max + self.omega_pitch/10:
#             partial = self.get_partial_dos_impact_at_omega(omega) * np.sum( self.smearing_function.calc(self.eigenvalues - omega) )
#             omegas.append( omega )
#             pdos.append( partial )
#             omega += self.omega_pitch
#         partial_dos = np.array(pdos).transpose()
#         omegas = np.array(omegas)
# 
#         for n, set_for_sum in enumerate(types):
#             atom = [k for k, v in labels.iteritems() if v == n][0] # find k by v
# 
#             multdos = 1
#             #if atom != 'Fe': continue
#             #else: multdos = 3
# 
#             pdos_sum = np.zeros(omegas.shape, dtype=float)
#             for i in set_for_sum:
#                 pdos_sum += multdos * partial_dos[(i-1):i].sum(axis=0)
#             plots.append( {'label': atom, 'data': [ [ round(omegas[n], 3), round(i, 3) ] for n, i in enumerate(pdos_sum.tolist()) ]} ) # round to reduce output
#         return plots


import math
from numpy import dot, array, linalg, arange
from tilde.berlinium.cubicspline import NaturalCubicSpline
from tilde.berlinium.dos import TotalDos, PartialDos
from ase.data.colors import jmol_colors
from ase.data import chemical_symbols
                    d += linalg.norm( bz_vec_dir )
                    x.append(d)
                    results[-1]['ticks'].append( [d, bz.replace(' ', '')] )

                # end in nullstand point (normally, Gamma)
                #y.append(kwargs['values'][nullstand][N])
                #if d == 0: d+=0.5
                #else: d += linalg.norm( bz_vec_ref )
                #x.append(d)
                #results[-1]['ticks'].append( [d, nullstand.replace(' ', '')] )

                divider = 10 if len(order)<10 else 1.5
                step = (max(x)-min(x)) / len(kwargs['values']) / divider

                xnew = arange(min(x), max(x)+step/2, step).tolist()
                ynew = []
                f = NaturalCubicSpline( array(x), array(y) )
                for i in xnew:
                    results[-1]['data'].append([ round(  i, 3  ), round(  f(i), 3  ) ]) # round to reduce output

        return results

    elif task == 'dos': # CRYSTAL, VASP, EXCITING

        results = []

        if 'precomputed' in kwargs:
            total_dos = [[i, kwargs['precomputed']['total'][n]] for n, i in enumerate(kwargs['precomputed']['x'])]

        else:

# Path: tilde/berlinium/dos.py
# class TotalDos(Dos):
#     def __init__(self, eigenvalues, sigma=None):
#         Dos.__init__(self, eigenvalues, sigma)
# 
#     def get_density_of_states_at_omega(self, omega):
#         return np.sum( self.smearing_function.calc(self.eigenvalues - omega))
# 
#     def calculate(self):
#         omega = self.omega_min
#         dos = []
# 
#         if self.omega_pitch == 0: self.omega_pitch = 1 # beware of endless loop if omega_pitch=0
#         while omega < self.omega_max + self.omega_pitch/10 :
#             p = self.get_density_of_states_at_omega(omega)
#             dos.append( [round(omega, 3), round(p, 3)] ) # round to reduce output
#             omega += self.omega_pitch
#         return dos
# 
# class PartialDos(Dos):
#     # eigenvalues and impacts must be unique and one-to-one correspondent
#     def __init__(self, eigenvalues, impacts, sigma=None):
#         Dos.__init__(self, eigenvalues, sigma)
#         self.impacts = np.array(impacts)
# 
#     def get_partial_dos_impact_at_omega(self, omega):
#         # function for obtaining scaled partial impacts in any omega point using a linear equation interpolation
#         if omega in self.eigenvalues:
#             return self.impacts[ np.where(self.eigenvalues == omega)[0][0] ]
#         if omega < self.eigenvalues[0]:
#             return np.zeros( len(self.impacts[0]) )
#         elif omega > self.eigenvalues[-1]:
#             return self.impacts[-1]
#         else:
#             for n in range(len(self.eigenvalues)):
#                 if n < len(self.eigenvalues) and self.eigenvalues[n] < omega < self.eigenvalues[n+1]:
#                     result = (omega - self.eigenvalues[n])*(self.impacts[n+1] - self.impacts[n])/(self.eigenvalues[n+1] - self.eigenvalues[n]) + self.impacts[n]
#                     return result
# 
#     def calculate(self, types, labels):
#         omega = self.omega_min
#         pdos = []
#         omegas = []
# 
#         plots = []
# 
#         if self.omega_pitch == 0: self.omega_pitch = 1 # beware of endless loop if omega_pitch=0
#         while omega < self.omega_max + self.omega_pitch/10:
#             partial = self.get_partial_dos_impact_at_omega(omega) * np.sum( self.smearing_function.calc(self.eigenvalues - omega) )
#             omegas.append( omega )
#             pdos.append( partial )
#             omega += self.omega_pitch
#         partial_dos = np.array(pdos).transpose()
#         omegas = np.array(omegas)
# 
#         for n, set_for_sum in enumerate(types):
#             atom = [k for k, v in labels.iteritems() if v == n][0] # find k by v
# 
#             multdos = 1
#             #if atom != 'Fe': continue
#             #else: multdos = 3
# 
#             pdos_sum = np.zeros(omegas.shape, dtype=float)
#             for i in set_for_sum:
#                 pdos_sum += multdos * partial_dos[(i-1):i].sum(axis=0)
#             plots.append( {'label': atom, 'data': [ [ round(omegas[n], 3), round(i, 3) ] for n, i in enumerate(pdos_sum.tolist()) ]} ) # round to reduce output
#         return plots


# Path: tilde/berlinium/plotter.py
import math
from numpy import dot, array, linalg, arange
from tilde.berlinium.cubicspline import NaturalCubicSpline
from tilde.berlinium.dos import TotalDos, PartialDos
from ase.data.colors import jmol_colors
from ase.data import chemical_symbols

                    d += linalg.norm( bz_vec_dir )
                    x.append(d)
                    results[-1]['ticks'].append( [d, bz.replace(' ', '')] )

                # end in nullstand point (normally, Gamma)
                #y.append(kwargs['values'][nullstand][N])
                #if d == 0: d+=0.5
                #else: d += linalg.norm( bz_vec_ref )
                #x.append(d)
                #results[-1]['ticks'].append( [d, nullstand.replace(' ', '')] )

                divider = 10 if len(order)<10 else 1.5
                step = (max(x)-min(x)) / len(kwargs['values']) / divider

                xnew = arange(min(x), max(x)+step/2, step).tolist()
                ynew = []
                f = NaturalCubicSpline( array(x), array(y) )
                for i in xnew:
                    results[-1]['data'].append([ round(  i, 3  ), round(  f(i), 3  ) ]) # round to reduce output

        return results

    elif task == 'dos': # CRYSTAL, VASP, EXCITING

        results = []

        if 'precomputed' in kwargs:
            total_dos = [[i, kwargs['precomputed']['total'][n]] for n, i in enumerate(kwargs['precomputed']['x'])]

        else:

            tdos = TotalDos( kwargs['eigenvalues'], sigma=kwargs['sigma'] )