# -*- coding: utf-8 -*-

# from framesplit import trajectory
# too slow using this module

import matplotlib.pyplot as plt

dirName = r"F:\simulations\asphaltenes\na-mont\TMBO-oil\water\373-continue/"
xyzName = 'all.xyz'
hetero = 'O'  # 'oh' 'N' 'sp'      'O' 'Np' 'sp'

with open(dirName + xyzName, 'r') as foo:
    coords = foo.readlines()
    nAtoms = int(coords[0])
    nFrames = int(len(coords) / (nAtoms + 2))
    pos = []
    for i in range(nFrames):
        istart = i * (nAtoms + 2)
        iend = (i + 1) * (nAtoms + 2)
        pos.append(coords[istart:iend])

# for i in range(200):
#     print coords[i]

heteroatom = 0
# all of my molecules have atoms less than 200
for i in range(200):
    x = pos[0][i].split()[0]
    if x == hetero:
        heteroatom = i
        break

heteroZ = []
for p in pos:
    # print p[heteroatom].split()[0]
    zx = float(p[heteroatom].split()[3])
    if zx < 10:
        zx = zx + 80
    heteroZ.append(zx)

with open(dirName + 'heteroZ.dat', 'w') as foo:
    for i, z in enumerate(heteroZ):
        print >> foo, "%3d %8.5f" % (i, z)

# energy plot
plt.figure(0, figsize=(8, 4))
figName = dirName + 'heteroZ.png'
plt.title('z of heteroatom', fontsize=20)
plt.plot(range(len(heteroZ)-1), heteroZ[1:], linewidth=2)
plt.grid(True)
plt.xlabel('steps')
plt.ylabel('Z')
plt.axis([0, len(heteroZ)*1.1, 0, max(heteroZ)*1.1])
plt.savefig(figName, format='png', dpi=300)
plt.close()
