#!/usr/bin/env python
############################################################################
# split_models.py splits a multi-model PDB files into several tmp files
# author: Yifeng Liu
# yifeng@cs.ualberta.ca
# April 2011 
############################################################################
import sys, os

def print_model(model_lines, file_prefix, num):
    ''' model_lines contains all lines for a given protein model 
    this method will print the lines into a single file '''
    outfile = file_prefix + "-m" + str(num) + ".pdb"
    output = open(outfile, 'w')
    for line in model_lines:
        print >>output, line.rstrip()
    output.close()

if __name__=='__main__':
    if len(sys.argv)<2:
        print "Usage: python %prog [input.pdb]"
        sys.exit(-1)
    infile = sys.argv[1]
        
    file_prefix = infile[0:-4]

    input = open(infile, 'r')
    lines = input.readlines()
    input.close()
    model_num = 0
    model_lines = []
    INMODEL = False
    for line in lines:
        if line.startswith("ENDMDL"):
            print_model(model_lines, infile, model_num)
            model_lines = []
            INMODEL = False
            continue
        if INMODEL:
            model_lines.append(line)
        elif line.startswith("MODEL"):
            model_lines.append(line)
            INMODEL = True
            model_num += 1
        else:
            continue
            
    print "detected", model_num, "models"
