#!/bin/env python3
import warnings
from Bio import PDB
from Bio import BiopythonWarning
import numpy as np
import argparse
import sys
import os
from pathlib import Path
parser = argparse.ArgumentParser(
                    prog='pdb_move_origin.py',
                    description='Takes the protein and moves it to the center of the coordinate system',
                    epilog='(c) Stefan Richter, HITS gGmbH')

parser.add_argument('pdbfile', help="filename to process", metavar="pdbfile")
parser.add_argument('--warnings', help="show warnings from biopython", default=False,action=argparse.BooleanOptionalAction)

args = parser.parse_args()

if not args.warnings:
    warnings.simplefilter('ignore', BiopythonWarning)


print("Processing "+args.pdbfile)
f = Path(args.pdbfile)
if not f.is_file():
    print("File "+args.pdbfile+" is not existing")
    sys.exit(1)

parser = PDB.PDBParser()
io = PDB.PDBIO()


struct = parser.get_structure(args.pdbfile,args.pdbfile)

rotation_matrix = PDB.rotmat(PDB.Vector([0, 0, 0]), PDB.Vector([0, 0, 0]))

for atom in struct.get_atoms():
    atom_C1 = atom.coord.copy()
    break

for model in struct:
    for chain in model:
        for residue in chain:
            for atom in residue:
                atom.transform(rotation_matrix, -atom_C1)

io.set_structure(struct)
newfile=os.path.join(f.parent,f.stem+"-moved-origin"+f.suffix)
io.save(newfile)
print("Written to "+newfile)