#!/usr/bin/env python
""" This script takes self-energy of DFT+DMFT code,
and performs max-entropy on auxiliary Green's function of the form
Gc(iom) = 1/(iom-Sigma+s_oo)
Then it performs Kramars-Kronig on resulting Dos, to obtain auxiliary G(omega).
From G(omega), one can obtain Sigma(omega) by
Sigma(omega) = omega+s_oo-1/G(omega)
which is written to 'Sig.out'
"""
import sys, shutil, os
from scipy import *

if len(sys.argv)<2:
    print 'give input file Sigma(iom)'
    sys.exit(0)

Sfile = sys.argv[1]
fi = open(Sfile, 'r')
firstlines = [fi.next(),fi.next()]
fi.close()

Sdata = loadtxt(Sfile).transpose()
# number of baths
nb = (len(Sdata)-1)/2

# input to maxent code
execfile('maxent_params.dat')

# prepares input for fortran maxent executable
fi = open('input', 'w')
for k in params:
    print >> fi, k[1], '! '+k[0]
fi.close()

exepath = os.environ.get('WIEN_DMFT_ROOT')

iom = Sdata[0]
for b in range(nb):
    Gi = 1/(iom*1j-Sdata[1+2*b]-Sdata[2+2*b]*1j)
    Gdata = array([iom,real(Gi),imag(Gi)]).transpose()
    savetxt('Gf.'+str(b), Gdata)
    cmd = exepath+'/inverseF ' +str(params[1][1])+' Gf.'+str(b)+' > Gtau.'+str(b)
    print cmd
    print os.popen(cmd).read()
    
print 'Starting Maxent'

Sigt=[]
for b in range(nb):
    # Run maxentropy
    shutil.copy2('Gtau.'+str(b), 'Gtau1.dat')
    cmd = exepath+'/maxentropy < input >> nohup.out.me'
    print cmd
    print os.popen(cmd).read()
    shutil.copy2('dos', 'dos.'+str(b))
    # removes zero from mesh, because Kramars-Kronig can be be done with zero
    Ddata = loadtxt('dos').transpose()
    izero = Ddata[0].tolist().index(0)
    c1 = hstack([Ddata[0][:izero],Ddata[0][izero+1:]])
    c2 = hstack([Ddata[1][:izero],Ddata[1][izero+1:]])
    Data1 = array([c1,c2]).transpose()
    savetxt('dosn', Data1)
    # Performs Kramars-Kronig
    cmd = exepath + '/skrams -cn 2 -s -pi dosn > Gc'
    print cmd
    print os.popen(cmd).read()
    Gdata = loadtxt('Gc').transpose()
    om = Gdata[0]
    Sc = om-1/(Gdata[1]+Gdata[2]*1j)
    savetxt('sig.'+str(b), array([om,real(Sc),imag(Sc)]).transpose())
    if b==0: Sigt.append(om)
    Sigt.append(real(Sc))
    Sigt.append(imag(Sc))


Sigt = array(Sigt).transpose()
fo = open('Sig.out', 'w')
print >> fo, firstlines[0].strip()
print >> fo, firstlines[1].strip()
for sg in Sigt:
    for b in sg:
        print >> fo, b,
    print >> fo