# $CVSHeader: pal8e/src/PDBRandomize.py,v 1.5 2005/02/15 01:10:26 indraneel Rel $

#create random coordinate PDB files with varying RMSD

from Bio.PDB.PDBParser import PDBParser
from Bio.PDB.Structure import Structure
from Bio.PDB.PDBIO import PDBIO
import math
import random
import os
from CoordCalc import dist_sqr
import Numeric

class Parm:
    pdbdir = '../../dat/cull100/'
    listfile = '../../dat/cull100extra'
    outdirbase = '/mnt/e12/data/indraneel/PDBNoise/'
    RMSDoutFname = '../dat/stat/PDBNoiseRMSD'
    maxdistorsions = 10
    maxversions = 100

class RunOne:
    def __init__(self, parm):
        self.parm = parm
        #self.AvRMSD = []
        #self.RMSD = None

    def mangle(self, structure, itn, n):
        sigma = itn * 0.2
        parm = self.parm
        models = structure.get_list()
        root3 = math.sqrt(3)
        for model in models:
            chains = model.get_list()
            #deviation = []
            #ctr = 0
            for chain in chains:
                residues = chain.get_list()
                for res in residues:
                    shift = []
                    for i in range(3):
                        shift.append((random.gauss(0.0, sigma))/root3)
                    atoms = res.get_list()
                    for atom in atoms:
                        oldcoord = []
                        for i in range(3):
                            if atom.name == 'CA':
                                oldcoord.append(atom.coord[i])
                            atom.coord[i] += shift[i]
                        if atom.name == 'CA':
                            oldcoord = Numeric.array(oldcoord)
                            #SD = dist_sqr(atom.coord, oldcoord)
                            #deviation.append(SD)
                            #ctr += 1
            #MSD = sum(deviation)/ctr
            #self.RMSD = math.sqrt(MSD)

    def writeOut(self, fname, structure, i, n):
        parm = self.parm
        outfname = parm.outdirbase+fname[:-4]
        outfname += '_'+str(i)+'_'+str(n)+'.pdb'
        pdbio = PDBIO()
        pdbio.set_structure(structure)
        body = pdbio.get_lines()
        outfh = file(outfname, 'w')
        outfh.writelines(self.header+body+self.trailer)
        outfh.close()

    def run(self, fname):
        parm = self.parm
        #random.seed()
        infname = parm.pdbdir + fname
        parser = PDBParser(PERMISSIVE=1)
        for i in range(1, parm.maxdistorsions+1):
            #RMSD = []
            for n in range(parm.maxversions):
                structure = parser.get_structure('xxx', infname)
                self.header = parser.get_header()
                self.trailer = parser.get_trailer()
                self.mangle(structure, i, n)
                #RMSD.append(self.RMSD)
                self.writeOut(fname, structure, i, n)
            #RMSD = sum(RMSD)/(n+1)
            #self.AvRMSD.append(RMSD)

class Runner:
    def __init__(self, parm):
        self.parm = parm

    def showRMSD(self, rmsdDict):
        lines = []
        for (fname, rmsd) in rmsdDict.items():
            lines += [fname+'\t'+str(rmsd)+'\n']
        outfname = parm.RMSDoutFname
        fh = file(outfname, 'w')
        fh.writelines(lines)
        fh.close()

    def getFileList(self):
        parm = self.parm
        fh = file(parm.listfile)
        lines = fh.readlines()
        fh.close()
        for (i, line) in enumerate(lines):
            lines[i] = line[:-1]
        return lines

    def forkJobs(self, fname):
        if len(self.children.keys()) > self.multiprocess:
            pid,stat = os.wait()
            del self.children[pid]
            
        newpid = os.fork()
        if newpid == 0:
            runone = RunOne(parm)
            runone.run(fname)
            os._exit(0)
        else:
            self.children[newpid] = 1

    def run(self, processes):
        parm = self.parm
        filelist = self.getFileList()
        self.children = {}
        self.multiprocess = processes
        for fname in filelist:
            self.forkJobs(fname)
        while len(self.children.keys()):
            pid,stat = os.wait()
            del self.children[pid]

if __name__ == '__main__':
    parm = Parm()
    runner = Runner(parm)
    runner.run(processes=40)
