import netCDF4 as netcdf

import timeseries, pymbar

import analyze

import pylab

filename = 'p-xylene/complex.nc'
#filename = '1-methylpyrrole/complex.nc'
ncfile = netcdf.Dataset(filename, 'r')

niterations = ncfile.variables['positions'].shape[0]                                                                                                                                                                                              
nstates = ncfile.variables['positions'].shape[1]                                                                                                                                                                                                  
natoms = ncfile.variables['positions'].shape[2]                                                                                                                                                                                                   
print "Read %(niterations)d iterations, %(nstates)d states, %(natoms)d atoms" % vars() 

u_n = analyze.extract_u_n(ncfile)

figsize(20,6)
pylab.clf()
pylab.plot(u_n, 'k.')
pylab.xlabel('iteration')
pylab.ylabel('u_n')
pylab.show()

T = len(u_n)
T = min(T, 2500) # just analyze beginning of simulation if it's really long
g_t = numpy.ones([T-1], numpy.float32)
Neff_t = numpy.ones([T-1], numpy.float32)
for t in range(T-1):
    g_t[t] = timeseries.statisticalInefficiency(u_n[t:T], method='fast')
    Neff_t[t] = (T-t+1) / g_t[t]

figsize(20,6)
pylab.clf()
pylab.subplot(211)
pylab.plot(g_t, 'k.')
pylab.ylabel('$g$')
pylab.xlabel('origin $t_0$')
pylab.subplot(212)
pylab.plot(Neff_t, 'k.')
pylab.xlabel('origin $t_0$')
pylab.ylabel('$N_{eff}$')
pylab.show()

T = len(u_n)
nskip = 10
indices = range(0, T-1, nskip)
N = len(indices)
t0_t = numpy.ones([N], numpy.float32)
g_t = numpy.ones([N], numpy.float32)
Neff_t = numpy.ones([N], numpy.float32)
for n in range(N):
    t0 = nskip*n
    t0_t[n] = t0
    g_t[n] = timeseries.statisticalInefficiency(u_n[t0:T], method='fast')
    Neff_t[n] = (T-t0) / g_t[n]

figsize(20,6)
pylab.clf()
pylab.subplot(211)
pylab.plot(t0_t, g_t, 'k.')
pylab.ylabel('$g$')
pylab.xlabel('origin $t_0$')
pylab.subplot(212)
pylab.plot(t0_t, Neff_t, 'k.')
pylab.xlabel('origin $t_0$')
pylab.ylabel('$N_{eff}$')
pylab.show()

[nequil, g, Neff_max] = analyze.detect_equilibration(u_n, nskip=10, method='fast')
print "The first %d samples should be discarded to equilibration, leaving %.1f effectively uncorrelated samples in the remainder (statistical inefficiency g = %.1f)." % (nequil, Neff_max, g)

analyze.show_mixing_statistics(ncfile, cutoff=0.05, nequil=nequil)

(Deltaf_ij, dDeltaf_ij) = analyze.estimate_free_energies(ncfile, ndiscard=nequil, g=g)

print "Free energy difference (estimated using all replicas): %8.3f +- %8.3f kT" % (Deltaf_ij[0,-1], dDeltaf_ij[0,-1])

replica_estimates = list()
for replica in range(nstates):
    (Deltaf_ij_replica, dDeltaf_ij_replica) = analyze.estimate_free_energies(ncfile, ndiscard=nequil, g=g, replicas=[replica])
    replica_estimates.append( (Deltaf_ij_replica, dDeltaf_ij_replica) )
    print "Free energy difference (estimated using replica %5d): %8.3f +- %8.3f kT" % (replica, Deltaf_ij_replica[0,-1], dDeltaf_ij_replica[0,-1])

figsize(14,6)
pylab.clf()
pylab.hold(True)
# State indices for free energy difference.
(i,j) = (0,nstates-1)
nreplicas = nstates
# Plot global estimate and confidence bands.
DeltaF = Deltaf_ij[i,j]
dDeltaF = dDeltaf_ij[i,j]
pylab.plot([0, nreplicas+1], [DeltaF, DeltaF], 'k-')
pylab.plot([0, nreplicas+1], [DeltaF+2*dDeltaF, DeltaF+2*dDeltaF], 'r-')
pylab.plot([0, nreplicas+1], [DeltaF-2*dDeltaF, DeltaF-2*dDeltaF], 'r-')
# Plot replica estimates.
for replica in range(nreplicas):
    (Deltaf_ij_replica, dDeltaf_ij_replica) = replica_estimates[replica]
    DeltaF = Deltaf_ij_replica[i,j]
    dDeltaF = dDeltaf_ij_replica[i,j]
    pylab.plot([replica], [DeltaF], 'ko')
    pylab.plot([replica,replica], [DeltaF+2*dDeltaF,DeltaF-2*dDeltaF], 'k-')
pylab.xlabel('replica index')
pylab.ylabel('$\Delta F$ (kT)')
pylab.show()

Zscores = numpy.zeros([nreplicas], numpy.float64)
DeltaF_global = Deltaf_ij_replica[i,j]
for replica in range(nreplicas):
    (Deltaf_ij_replica, dDeltaf_ij_replica) = replica_estimates[replica]
    DeltaF_replica = Deltaf_ij_replica[i,j]
    dDeltaF_replica = dDeltaf_ij_replica[i,j]        
    Zscores[replica] = (DeltaF_replica - DeltaF_global) / dDeltaF_replica

figsize(14,6)
pyplot.clf()
pyplot.hold(True)
pyplot.plot([0, nreplicas], [0, 0], 'k-')
pyplot.plot(Zscores, 'ko')
for replica in range(nreplicas):
    pyplot.plot([replica, replica], [0, Zscores[replica]], 'k-')
pyplot.show()

figsize(14,6)
earlier_replica_estimates = list()
nuse = int((niterations - nequil)/4.0)
for replica in range(nstates):
    (Deltaf_ij_replica, dDeltaf_ij_replica) = analyze.estimate_free_energies(ncfile, ndiscard=nequil, g=g, replicas=[replica], nuse=nuse)
    earlier_replica_estimates.append( (Deltaf_ij_replica, dDeltaf_ij_replica) )
    print "Free energy difference (estimated using replica %5d): %8.3f +- %8.3f kT" % (replica, Deltaf_ij_replica[0,-1], dDeltaf_ij_replica[0,-1])
pylab.clf()
pylab.hold(True)
# State indices for free energy difference.
(i,j) = (0,nstates-1)
nreplicas = nstates
# Plot global estimate and confidence bands.
DeltaF = Deltaf_ij[i,j]
dDeltaF = dDeltaf_ij[i,j]
pylab.plot([0, nreplicas+1], [DeltaF, DeltaF], 'k-')
pylab.plot([0, nreplicas+1], [DeltaF+2*dDeltaF, DeltaF+2*dDeltaF], 'r-')
pylab.plot([0, nreplicas+1], [DeltaF-2*dDeltaF, DeltaF-2*dDeltaF], 'r-')
# Plot replica estimates.
for replica in range(nreplicas):
    (Deltaf_ij_replica, dDeltaf_ij_replica) = earlier_replica_estimates[replica]
    DeltaF = Deltaf_ij_replica[i,j]
    dDeltaF = dDeltaf_ij_replica[i,j]
    pylab.plot([replica], [DeltaF], 'ro')
    pylab.plot([replica,replica], [DeltaF+2*dDeltaF,DeltaF-2*dDeltaF], 'r-')

    (Deltaf_ij_replica, dDeltaf_ij_replica) = replica_estimates[replica]
    DeltaF = Deltaf_ij_replica[i,j]
    dDeltaF = dDeltaf_ij_replica[i,j]
    pylab.plot([replica], [DeltaF], 'ko')
    pylab.plot([replica,replica], [DeltaF+2*dDeltaF,DeltaF-2*dDeltaF], 'k-')
pylab.xlabel('replica index')
pylab.ylabel('$\Delta F$ (kT)')
pylab.show()

Zscore = 0.0
DeltaF_global = Deltaf_ij_replica[i,j]
for replica in range(nreplicas):
    (Deltaf_ij_replica, dDeltaf_ij_replica) = earlier_replica_estimates[replica]
    DeltaF_replica = Deltaf_ij_replica[i,j]
    dDeltaF_replica = dDeltaf_ij_replica[i,j]
    Zscore += ((DeltaF_replica - DeltaF_global) / dDeltaF_replica)**2
print numpy.sqrt(Zscore / (nreplicas-1))

state_nk = ncfile.variables['states'][:,:]

figsize(24,24)
pylab.clf()
pylab.hold(True)
for k in range(nstates):
    pylab.subplot(nstates,1,k+1)
    pylab.plot(state_nk[:,k], 'k.')
    pylab.axis([0, niterations, 0, nstates])
    pylab.yticks([])
    if k < (nstates-1): pylab.xticks([])
    pylab.ylabel('%d' % k)
pylab.show()

phi_kt = analyze.compute_torsion_trajectories(ncfile)

figsize(24,24)
pylab.clf()
pylab.hold(True)
for k in range(nstates):
    pylab.subplot(nstates,1,k+1)
    pylab.plot(phi_kt[k,:], 'k.')
    pylab.yticks([])
    pylab.ylabel('%d' % k)
    if k < (nstates-1): pylab.xticks([])
    pylab.axis([0, niterations, -180, +180])
pylab.show()