#!/usr/bin/env python

from Dacapo import Dacapo
from ASE import Atom, ListOfAtoms

from Numeric import *
import sys

if len(sys.argv)<2:
  print """
  Makes an estimation of the magnetic moment of the individual atoms
  based on the atom projected local density of states. 
  
  USAGE: python LMagMom.py [nc.file]
  
  NOTE: Works only if 
   !calc.CalculateAtomicDOS()!
  is specified for the dacapo calculation
  """
  sys.exit()  
atoms=Dacapo.ReadAtoms(sys.argv[1])
calc=atoms.GetCalculator()


spin0=calc.GetDensityArray(spin=0)
spin1=calc.GetDensityArray(spin=1)
ngx,ngy,ngz=shape(spin0)
gridpoints=ngx*ngy*ngz
VolumeElement=atoms.GetUnitCellVolume()/float(gridpoints)
Nup=sum(sum(sum(spin0)))*VolumeElement
Ndown=sum(sum(sum(spin1)))*VolumeElement
print "Nup", Nup,"Ndown", Ndown,"Diff", Nup-Ndown

for i in range(len(atoms)):
  dosu=calc.GetLDOS(atoms=[i+1],angularchannels=["s","p","d"],spin=[0])
  dosus=calc.GetLDOS(atoms=[i+1],angularchannels=["s","p","d"],spin=[0],cutoffradius = 'short')
  dosd=calc.GetLDOS(atoms=[i+1],angularchannels=["s","p","d"],spin=[1])
  dosds=calc.GetLDOS(atoms=[i+1],angularchannels=["s","p","d"],spin=[1],cutoffradius = 'short')
  print "Atom_"+str(i)+" "+atoms[i].GetChemicalSymbol()
  print "N_inf=",dosd.GetIntegratedDOS()+dosu.GetIntegratedDOS(),"N_short=", dosus.GetIntegratedDOS()+dosds.GetIntegratedDOS()  
  print "M_inf=",dosu.GetIntegratedDOS()-dosd.GetIntegratedDOS(), "M_short=",dosus.GetIntegratedDOS()-dosds.GetIntegratedDOS()