#!/usr/bin/env python3

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from numpy import linalg
from numpy.linalg import inv
import warnings
warnings.filterwarnings("ignore")
from scipy.optimize import curve_fit
from scipy.special import erf
import sys

filenam = input('Please input the pdb filename: ')
n=input('Please input number of carbon: ')
n=int(n)
n1=input('Please input number of carbon per chain: ')
n1=int(n1)
readfile = open(filenam,'r')
sepfile = readfile.read().split('\n')
readfile.close()
count = 0
count2 = 0
xf=[]
yf=[]
zf=[]
for pair in sepfile:
	if pair == '' or pair=='END':
		break
	else:
		matrix = pair.split()
		if count>4:
			xf.append(float(matrix[6]))
			yf.append(float(matrix[7]))
			zf.append(float(matrix[8]))
		count=count+1

filename2="initial2.pdb"
writefile = open(filename2,'w')
count=1
for j in range(n):
	if j+1<10:
		writefile.write("ATOM")
		#writefile.write("      ")
		writefile.write("%7s"%(str(j+1)))
		writefile.write("  ")
		if j==0 or j==n-1:
			writefile.write("C1")
			writefile.write("  ")
		elif j%2==0 and j!=0 and j!=(n-1):
			writefile.write("C1")
			writefile.write("  ")
		else:
			writefile.write("C1")
			writefile.write("  ")
		if (j)%n1==0:
			writefile.write("PEM")
		elif (j+1)%n1==0:
			writefile.write("PEN")
		else:
			writefile.write("PEE")
		#writefile.write("     ")
		writefile.write("%6s"%(str(count)))
		count=count+1
		writefile.write("%10s%8s%8s"%(str(round(xf[j],2)),str(round(yf[j],2)),str(round(zf[j],2))))
		writefile.write("  1.00")
		writefile.write("  ")
		writefile.write("0.00")
		writefile.write("\n")
	elif j+1>=10 and j+1<100:
		writefile.write("ATOM")
		#writefile.write("     ")
		writefile.write("%7s"%(str(j+1)))
		writefile.write("  ")
		if j==0 or j==n-1:
			writefile.write("C1")
			writefile.write("  ")
		elif j%2==0 and j!=0 and j!=(n-1):
			writefile.write("C1")
			writefile.write("  ")
		else:
			writefile.write("C1")
			writefile.write("  ")
		if j>=9:
			if (j)%n1==0:
				writefile.write("PEM")
			elif (j+1)%n1==0:
				writefile.write("PEN")
			else:
				writefile.write("PEE")
			#if count<10:
			#	writefile.write("     ")
			#if count>=10:
			#	writefile.write("    ")
			writefile.write("%6s"%(str(count)))
			count=count+1
		writefile.write("%10s%8s%8s"%(str(round(xf[j],2)),str(round(yf[j],2)),str(round(zf[j],2))))
		writefile.write("  1.00")
		writefile.write("  ")
		writefile.write("0.00")
		writefile.write("\n")
	elif j+1>=100:
		writefile.write("ATOM")
		#writefile.write("    ")
		writefile.write("%7s"%(str(j+1)))
		writefile.write("  ")
		if j==0 or j==n-1:
			writefile.write("C1")
			writefile.write("  ")
		elif j%2==0 and j!=0 and j!=(n-1):
			writefile.write("C1")
			writefile.write("  ")
		else:
			writefile.write("C1")
			writefile.write("  ")
		if j>=99:
			if (j)%n1==0:
				writefile.write("PEM")
			elif (j+1)%n1==0:
				writefile.write("PEN")
			else:
				writefile.write("PEE")
			#if count<10:
			#	writefile.write("     ")
			#if count>=10:
			#	writefile.write("    ")
			#if count>=100:
			#	writefile.write("")
			writefile.write("%6s"%(str(count)))
			count=count+1
		writefile.write("%10s%8s%8s"%(str(round(xf[j],2)),str(round(yf[j],2)),str(round(zf[j],2))))
		writefile.write("  1.00")
		writefile.write("  ")
		writefile.write("0.00")
		writefile.write("\n")

