#!/usr/bin/python
import sys
import numpy as np
import argparse
from scipy.special import xlogy

# Entropy is in units of kB


parser = argparse.ArgumentParser()
parser.add_argument('Nspecies', type=int, help='number of atomic species')
parser.add_argument('L', type=int, help='length input for N=2(L^3) atom configuration')
parser.add_argument('Nrun',type=int,help='number of data points (from last output) to use')
#parser.add_argument('-a', action='store_true', help="when selected, averages occupations over a-b and over c-d sublattices") 
parser.add_argument('-p', metavar='PATH', help='path of directory with ZABGD correlation files;\
        default: `%(default)s`', default='./correlations_data')
parser.add_argument('-d', metavar='PATH', help='path of directory to place output files;\
        default: `%(default)s`', default='.')
args = parser.parse_args()

N=2*(args.L**3)
nspecies=args.Nspecies

with open('Tlist','r') as TlistFile:
    Tlist = np.loadtxt(TlistFile, dtype=int)
    TlistFile.close()

xe=[]; xo=[]
y0g=[]; y1g=[]; #y2g=[]; y3g=[]

for T in Tlist:
    with open(args.p+'/ZABGD_T_%04d'%T,'r') as ZFile:
        Zall=np.loadtxt(ZFile, dtype=int)
        Z0=np.sum(Zall[len(Zall)-args.Nrun:], axis=0)
        ZFile.close()

    ZABGD=Z0/float(len(Z0)*6*N)
    ZABGD=ZABGD.reshape((nspecies,nspecies,nspecies,nspecies))
    UAGD=np.sum(ZABGD,1)
    UBGD=np.sum(ZABGD,0)
    UGAB=np.swapaxes(np.swapaxes(np.sum(ZABGD,3),0,2),1,2)
    UDAB=np.swapaxes(np.swapaxes(np.sum(ZABGD,2),0,2),1,2)
    YAG=np.swapaxes(np.sum(UGAB,2),0,1)
    YAD=np.swapaxes(np.sum(UDAB,2),0,1)
    YBG=np.swapaxes(np.sum(UGAB,1),0,1)
    YBD=np.swapaxes(np.sum(UDAB,1),0,1)
    VAB=np.sum(UDAB,0)
    VGD=np.sum(UAGD,0)
    XA=np.sum(YAG,1); XB=np.sum(YBG,1)
    XG=np.sum(YAG,0); XD=np.sum(YAD,0)
    Xa=(XA+XB)/2.0
    Xg=(XG+XD)/2.0
    Yag=(YAG+YAD+YBG+YBD)/4.0
    y0g.append([T, Yag[0,0], Yag[0,1]])#, Yag[0,2], Yag[0,3]])
    y1g.append([T, Yag[1,0], Yag[1,1]])#, Yag[1,2], Yag[1,3]])
    #y2g.append([T, Yag[2,0], Yag[2,1], Yag[2,2], Yag[2,3]])
    #y3g.append([T, Yag[3,0], Yag[3,1], Yag[3,2], Yag[3,3]])
    xe.append([T, Xa[0], Xa[1]])#, Xa[2], Xa[3]])
    xo.append([T, Xg[0], Xg[1]])#, Xg[2], Xg[3]])

xe = np.array(xe)
xo = np.array(xo)
y0g=np.array(y0g); y1g=np.array(y1g)
#y2g=np.array(y2g); y3g=np.array(y3g)

with open(args.d+'/xe_ALL.dat','w') as out_xe:
    np.savetxt(out_xe, xe, fmt='%4d %6.12f %6.12f')
with open(args.d+'/xo_ALL.dat','w') as out_xo:
    np.savetxt(out_xo, xo, fmt='%4d %6.12f %6.12f')
with open(args.d+'/y1g_ALL.dat','w') as out_y0g:
    np.savetxt(out_y0g, y0g, fmt='%4d %1.8f %1.8f')
with open(args.d+'/y2g_ALL.dat','w') as out_y1g:
    np.savetxt(out_y1g, y1g, fmt='%4d %1.8f %1.8f')
#with open(args.d+'/y3g_ALL.dat','w') as out_y2g:
    #np.savetxt(out_y2g, y2g, fmt='%4d %1.8f %1.8f %1.8f %1.8f')
#with open(args.d+'/y4g_ALL.dat','w') as out_y3g:
    #np.savetxt(out_y3g, y3g, fmt='%4d %1.8f %1.8f %1.8f %1.8f')