from scipy import *
from pylab import *
import sys, os, re

styles = ('s-', 'o-', 'v-', '^-', '<-', '>-')
colors = ('r', 'g', 'y', 'm', 'c', 'b', 'y')

def fit_cuadratic(parameters, values, ys):
    x = values
    s0, s1, s2 = parameters
    return (s0 + s1*x + s2*x**2) - ys
    #x0, a, b, c = parameters
    #return (a *(x-x0)**2 + b*(x-x0) + c ) - ys
    #return reduce(lambda x, dx: x+dx**2, dchi)

def fitSig(cols, ns, Sig, xs):
    pars=[]
    for col in cols:
        yt = Sig[:,col]
        ys = yt[:ns]
    
        parameters = (0.1, 0.3, 0.4)
        (par,success) = optimize.leastsq(fit_cuadratic, parameters, args=(xs, ys))
        if not success: print 'Did not succeed!!!'
        pars.append(par)
        
    return pars

if __name__ == '__main__':

    Nom = 10
    
    if len(sys.argv)<3 or not os.path.isfile(sys.argv[1]):
        print 'Can not find input file', sys.argv[1:]
        print 'give two arguments: filename  number-of-baths'
        sys.exit(0)
        
    fsig = sys.argv[1]
    nbaths = int(sys.argv[2])
    print 'Number of baths=', nbaths
    cols  = [2*i+1 for i in range(nbaths)]
    colsr = [2*i   for i in range(nbaths)]
    #legends=["1","","2","","3","","4"]
    legends = array([[str(i+1),""] for i in range(nbaths)]).flatten().tolist()
    print 'legends=', legends
    
    fs = open(fsig, 'r')
    data = fs.next()
    #mus=[]
    #Tmp=[]
    while ( re.match('\s*#', data) is not None ):
    #    m = re.search(r'mu=(\d+\.\d+)', data)
    #    if m is not None:
    #        mus.append(float(m.groups()[0]))
    #    m = re.search(r'T=(\d+\.\d+)', data)
    #    if m is not None:
    #        Tmp.append(float(m.groups()[0]))
    #    
        data = fs.next()
    #dmu = sum(array(mus)-Edc)/len(mus)
    #Temperature = sum(array(Tmp))/len(Tmp)
    #print 'dmu=', dmu, 'T=', Temperature
    
    om=[]
    Sig=[]
    for i in range(Nom):
        dt = data.split()
        om.append(float(dt[0]))
        Sig.append(map(float,dt[1:]))
        data = fs.next()
    
    Sig=array(Sig)
    om=array(om)
    
    Temperature = om[0]/pi
    print 'Temperature=', Temperature
    
    ns = 4
    xs = om[:ns]
    
    pars = fitSig(cols, ns, Sig, xs)
    
    parsr = fitSig(colsr, ns, Sig, xs)
    
    for icol,col in enumerate(cols):
        print parsr[icol][0], pars[icol][0], 1/(1-pars[icol][1])
    
    fo = open(fsig+'_t', 'w')
    print >> fo, '# Temperature'
    print >> fo, Temperature
    print >> fo, '# Self-energy at zero'
    for icol,col in enumerate(cols):
        print >> fo, parsr[icol][0], pars[icol][0],
    print >> fo
    print >> fo, '# Z - quasiparticle renormlization amplitude'
    for icol,col in enumerate(cols):
        print >> fo, 1/(1-pars[icol][1]),
    print >> fo
    fo.close()
    
    x0 = arange(0, om[-1], 1e-3)
    for icol,col in enumerate(cols):
        yt = Sig[:,col]
        plot(om, yt, styles[icol]+colors[icol])
        
        y0 = fit_cuadratic(pars[icol], x0, zeros(len(x0)))
        plot(x0, y0, colors[icol])

    v = axis()
    axis([v[0], v[1], v[2], 0])
    legend(legends, loc=0, numpoints=2, shadow=True)
    
    show()