#!/usr/bin/env python3


import numpy as np
import MDAnalysis as MDA
from matplotlib import pyplot as plt
#from MDAnalysis import analysis as MDAa
from MDAnalysis.analysis.dihedrals import Dihedral as MDAaD
import MDAnalysis.lib.distances as MDAld


## EDIT HERE the resname and traj-file-name(no dcd)
resn='CBP'  # MCB # MCN # PCZ
trajfn='1atm_300k_CBP_512_mainrun_2'   #'pCBP_shorttraj' ## change for others
# number of frames to skip 
nskip=500

# loading trajectory
u=MDA.Universe(resn+'-512-new.psf',trajfn+'.dcd')
#u=MDA.Universe('MCB-512-new.psf','mCBP_shorttraj.dcd')

nmols=512
nframes=len(u.trajectory)

#print('u.trajectory')
#for i, ttr in enumerate(u.trajectory[::2]):
#    print(i,ttr)

if resn=='CBP':
    dh_atoms=[["C39", "N32", "C36", "C35"], ["C8", "N1", "C5", "C4"], ["C34", "C33", "C2", "C3"]]
elif resn=='MCB':
    dh_atoms=[["C43", "N42", "C35", "C33"], ["C7", "N8", "C22", "C23"], ["C23", "C25", "C32", "C33"]]
elif resn=='MCN':
    dh_atoms=[["C7", "N8", "C22", "C23"], ["C43", "N42", "C35", "C33"], ["C23", "C25", "C32", "C33"] ]
# add PCZ case here
    
ndhs=len(dh_atoms)
print(ndhs)

fig, ax = plt.subplots(ndhs,1)

print('qua')
print(len(u.trajectory[::nskip]))

dhs=np.zeros( (len(u.trajectory[::nskip]),nmols,ndhs) )

bindh=np.linspace(-180, 180, 361)


for kkd, a4dh in enumerate(dh_atoms):

    dA1=[res.atoms.select_atoms("name "+a4dh[0]) for res in u.residues[:nmols]]
    dA2=[res.atoms.select_atoms("name "+a4dh[1]) for res in u.residues[:nmols]]
    dA3=[res.atoms.select_atoms("name "+a4dh[2]) for res in u.residues[:nmols]]
    dA4=[res.atoms.select_atoms("name "+a4dh[3]) for res in u.residues[:nmols]]

    for i1, ts in enumerate(u.trajectory[::nskip]):
        #for d11 in d1:
            #xa1,xa2,xa3,xa4=np.split(d11.positions,4)
            #print(xa1[0])
            #print(dC39)
            #x=Dihedrals(xa4[0],xa1[0],xa3[0],xa2[0])
        for jj in range(nmols):
            x1=MDAld.calc_dihedrals(dA1[jj],dA2[jj],dA3[jj],dA4[jj])*180./np.pi
            dhs[i1,jj,kkd]=x1
        #MDAdh(d11.positions[:3],d11.positions[3:6], d11.positions[6:9], d11.positions[9:])

    #print(dhs[:,:,kkd])
    
    ax[kkd].hist(dhs[:,:,kkd].flatten(), label=" ".join(a4dh), bins=bindh, alpha=0.6)
    ax[kkd].legend()




np.savetxt('dhs_'+resn+'_long_3.dat',dhs.reshape(-1,ndhs))

print('fffffff')
print(dhs)
print()
print('qqqqqq')
print(dhs.reshape(-1,ndhs))
dhsnew = []
dhsresh = dhs.reshape(-1,ndhs)
print(len(dhsresh[:,0]))
for i in range(0,len(dhsresh[:,0])):
    if dhsresh[i,0] != 0.0:
        dhsnew = np.append(dhsnew,dhsresh[i,:])

dhsnew = np.reshape(dhsnew,(int(len(dhsnew)/ndhs),ndhs))
#print(len(dhsnew))

fig.suptitle(resn, fontsize=16)
fig.savefig('dhs_'+resn+'_long_3.png')