#!/usr/bin/env python3

from prody import *
from pylab import *
import matplotlib.pyplot as plt
import itertools
import os
import glob
import re
ion()
prody.confProDy(auto_show=False)

# set output dir
out_dir = os.path.splitext(os.path.basename(sys.argv[0]))[0]
out_file = out_dir + "/combined"

# create output dir
if not os.path.exists(out_dir):
    os.makedirs(out_dir)

# create palette
colors = itertools.cycle(["red", "blue", "green", "orange"])

# parse files
in_files=os.listdir("1-concat/eda/")
out_files=os.listdir(out_dir)
systems = []
get_systems=os.listdir("1-concat/combined/")
for filename in get_systems:
  if filename.endswith(".pdb"):
     systems.append(re.sub('-.*','',filename))
my_pdb=glob.glob("1-concat/eda/combined*.pdb")

# if system_names file is present, set names from there
if os.path.exists('system_names'):
  exec(open('system_names').read())
else:
  systems=sorted(systems)

# load pdb
structure = parsePDB(my_pdb[0])

# load trajectory
if os.path.exists(out_file + '_prody.dcd'):
  trajectory = parseDCD(out_file + '_prody.dcd')
  trajectory.setCoords(structure)
  trajectory.setAtoms(structure.all)
else:
  dcd = glob.glob('1-concat/eda/combined-*.dcd')
  trajectory = parseDCD(dcd[0])
  trajectory.setCoords(structure)
  trajectory.setAtoms(structure.all)
  trajectory.superpose()
  # save aligned dcd
  writeDCD(out_file + '_prody.dcd', trajectory)

# calculate frame increments
num_frames = DCDFile.numFrames(trajectory)
num_systems = len(systems)
frame_incr = int(num_frames / num_systems - 1)
res_nums = AtomGroup.getResnums(structure)
first_res = res_nums[0]
last_res = res_nums[-1]

# create rmsf figure
first_frame = 0
last_frame = frame_incr
plt.figure(figsize=(12, 6))
for system in systems:
  rmsf = calcRMSF(trajectory[first_frame:last_frame])
  print(rmsf)
  plt.plot(res_nums, rmsf, label=system, color=next(colors), linewidth=1.5)
  first_frame = last_frame + 1
  last_frame = first_frame + frame_incr
title('RMSF', size=20)
plt.xlabel('Residue #', size=16)
plt.ylabel('Angstroms', size=16)
plt.xlim([first_res,last_res])
ax = subplot(1,1,1)
handles, labels = ax.get_legend_handles_labels()
labels, handles = zip(*sorted(zip(labels, handles), key=lambda t: t[0]))
plt.legend(handles, labels, loc='upper left')
ax.grid(True)
plt.savefig(out_file + '_rmsf.png', dpi=300, format='png')
plt.close('all[')