#!/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)

#######################################################################################
## ADJUSTABLE PARAMETERS                                                             ##
#######################################################################################

# set specific system names to visualize
#visualize=[]

#show=["ALL", "no_PHQ"]

#######################################################################################
## JOB CONTROL                                                                       ##
#######################################################################################

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

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

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

# 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])
structure_select = structure.select("resnum 88:203")

# load trajectory
if os.path.exists(out_file + '_prody.dcd'):
  trajectory = parseDCD(out_file + '_prody.dcd')
  trajectory.setCoords(structure_select)
  trajectory.setAtoms(structure_select)
else:
  combined_dcd = glob.glob('1-concat/eda/combined-*.dcd')
  trajectory = parseDCD(combined_dcd[0])
  trajectory.setCoords(structure)
  trajectory.setAtoms(structure_select)
  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 = structure_select.getResnums()
first_res = res_nums[0]
last_res = res_nums[-1]

# calculate covariance and modes
if os.path.exists(out_file + '.eda.npz'):
  eda_trajectory = loadModel(out_file + '.eda.npz')
else:
  eda_trajectory = EDA(combined_dcd[0])
  eda_trajectory.buildCovariance(trajectory, aligned=True)
  eda_trajectory.calcModes(n_modes=10)
  saveModel(eda_trajectory, filename=out_file)
  
for mode in eda_trajectory[:4]:
  print(calcFractVariance(mode))

# write VMD files
writeNMD(out_file, eda_trajectory[:10], structure_select)

# create EDA rmsf figures
def my_showSqFlucts(modes, *args, **kwargs):
    """Show square fluctuations using :func:`~matplotlib.pyplot.plot`.  See
    also :func:`.calcSqFlucts`."""

    import matplotlib.pyplot as plt
    sqf = calcSqFlucts(modes)
    if not 'label' in kwargs:
        kwargs['label'] = str(modes)
    show = plt.plot(res_nums,sqf, *args, **kwargs)
    plt.xlabel('Residue #')
    plt.ylabel('Square fluctuations')
    plt.xlim([first_res,last_res])
    plt.title(str(modes))
    return show
for i in range(0, 4):
  j = i + 1
  plt.figure(figsize=(8, 6))
  my_showSqFlucts(eda_trajectory[i])
  title('RMSF (EDA Mode ' + str(j) + ')')
  plt.savefig(out_file + '_rmsf_mode_' + str(j) + '.png', dpi=300, format='png')
  plt.close('all[')

# create cross-corr figure
def my_showCrossCorr(modes, *args, **kwargs):
    import matplotlib.pyplot as plt
    arange = np.arange(first_res - 1 ,last_res)
    cross_correlations = np.zeros((arange[-1]+2, arange[-1]+2))
    cross_correlations[arange[0]+1:,
                       arange[0]+1:] = calcCrossCorr(modes)
    kwargs['interpolation'] = 'bilinear'
    kwargs['origin'] = 'lower'
    show = plt.imshow(cross_correlations, *args, **kwargs), plt.colorbar()
    plt.axis([arange[0]+0.5, arange[-1]+1.5, arange[0]+0.5, arange[-1]+1.5])
    plt.title('Cross-correlations for {0}'.format(str(modes)))
    plt.xlabel('Residue')
    plt.ylabel('Residue')
    return show

plt.figure(figsize=(8, 6))
my_showCrossCorr(eda_trajectory)
title('Cross-correlation matrix')
plt.savefig(out_file + '_corr.png', dpi=300, format='png')
plt.close('all')

# create projection figures
plt.figure(figsize=(8, 6))
showProjection(trajectory[:frame_incr], eda_trajectory[0,1], color='red', marker='.', label=systems[0])
showProjection(trajectory[frame_incr + 1:frame_incr * 2 + 1], eda_trajectory[0,1], color='green', marker='.', label=systems[1]);
legend(markerscale=2, numpoints=1, handletextpad=0.25)
title('EDA')
plt.savefig(out_file + '_proj1_2.png', dpi=300, format='png')
plt.close('all')

plt.figure(figsize=(8, 6))
showProjection(trajectory[:frame_incr], eda_trajectory[2,3], color='red', marker='.', label=systems[0])
showProjection(trajectory[frame_incr + 1:frame_incr * 2 + 1], eda_trajectory[2,3], color='green', marker='.', label=systems[1]);
legend(markerscale=2, numpoints=1, handletextpad=0.25)
title('EDA')
plt.savefig(out_file + '_proj3_4.png', dpi=300, format='png')
plt.close('all')