#!/usr/bin/env python3

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")

filenam = input('Please input the pdb filename: ')
n_tot = input('Please input the total number of atoms: ')
n_tot = int(n_tot)
readfile = open(filenam,'r')
sepfile = readfile.read().split('\n')
readfile.close()
count2 = 0
count=0
xf=[]
yf=[]
zf=[]
for pair in sepfile:
	if pair == '' or pair=='END':
		break
	else:
		matrix = pair.split()
		if matrix[0]=='CONECT' or matrix[0]=='MASTER':
			print(matrix[0])
		else:
			if count>1:
				xf.append(float(matrix[5]))
				yf.append(float(matrix[6]))
				zf.append(float(matrix[7]))
		count=count+1
print(len(xf))
filename2="initial2.pdb"
writefile = open(filename2,'w')
count=1
cst=7
cst2=0
j2=0
for j in range(n_tot):
	if j2+1<7:
		writefile.write("ATOM")
		#writefile.write("      ")
		writefile.write("%7s"%(str(j+1)))
		writefile.write("  ")
		writefile.write("C"+str(j2+1))
		writefile.write("  ")
		writefile.write("BEC")
		writefile.write("%6s"%(str(cst2+1)))
		count=count+1
		writefile.write("%10s%8s%8s"%(str(round(xf[j2+cst2*12],2)),str(round(yf[j2+cst2*12],2)),str(round(zf[j2+cst2*12],2))))
		writefile.write("  1.00")
		writefile.write("  ")
		writefile.write("0.00")
		writefile.write("\n")
	elif j2+1>=7:
		writefile.write("ATOM")
		#writefile.write("     ")
		writefile.write("%7s"%(str(j+1)))
		writefile.write("  ")
		writefile.write("H"+str(j2+1-6))
		writefile.write("  ")
		writefile.write("BEC")
		writefile.write("%6s"%(str(cst2+1)))
		count=count+1
		writefile.write("%10s%8s%8s"%(str(round(xf[j2+cst2*12],2)),str(round(yf[j2+cst2*12],2)),str(round(zf[j2+cst2*12],2))))
		writefile.write("  1.00")
		writefile.write("  ")
		writefile.write("0.00")
		writefile.write("\n")	
	if (j+1)%12==0:
		j2=0
		cst2=cst2+1
	else:
		j2=j2+1

