#!/usr/bin/python
import os, pickle, sys
import numpy as np
from math import sqrt, sin
from lib.libSAXSPM import SAXSPM, calc_f,read_pdb_xyz

#==========================================================================
def check_oddatoms(xyz):
    allatoms = ["P", "S", "Ab", "Ub", "Gb", "Cb"]
    for i in range(len(xyz)):
        if (SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][0] not in allatoms):
            print "atom %s not in list: check list" % SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][0]
            sys.exit(1)


def get_globs(pdb_name):
    #f1 = open('thf_ribo_deform.pdb','r')
    f2 = open("rna_ini.pdb", 'w')
    xyz = read_pdb_xyz(pdb_name)
    #print(xyz)
    check_oddatoms(xyz)
    num_lines = len(xyz)
    serial = 1
    fmt_pdb = '%-6s%5d %-4s %-4s%1s%4d%1s%3s%8.3f%8.3f%8.3f%18s%-4s'
    nx,ny,nz,n_tot=0,0,0,0  # initiatlize
    glob = []
    for i in range(0,num_lines):
        count_begin = 0;
        if i != num_lines -1 and xyz[i][2] == xyz[i+1][2] : #check the resids are the same
             nx += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][3]
             ny += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][4]
             nz += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][5]
             n_tot += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]
        elif i == num_lines:
             nx += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][3]
             ny += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][4]
             nz += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][5]
             n_tot += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]
             print >> f2, fmt_pdb % ('ATOM', serial,'CA',str(xyz[i][1]),'X',xyz[i][2],\
                                     "","",nx/n_tot,ny/n_tot,nz/n_tot,'0.00','')
             glob.append([str(xyz[i][1]), xyz[i][2], nx/n_tot,ny/n_tot,nz/n_tot])
             #print xyz[i][2], nx/n_tot, ny/n_tot, nz/n_tot
        else:
             nx += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][3]
             ny += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][4]
             nz += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]*xyz[i][5]
             n_tot += SAXSPM.atom_map[xyz[i][1]][xyz[i][0]][2]
             print >> f2, fmt_pdb % ('ATOM', serial,'CA',str(xyz[i][1]),'X',xyz[i][2],"",\
                                     "",nx/n_tot,ny/n_tot,nz/n_tot,'0.00','')
             glob.append([str(xyz[i][1]), xyz[i][2], nx/n_tot,ny/n_tot,nz/n_tot])
             #print xyz[i][2], nx/n_tot, ny/n_tot, nz/n_tot
             nx,ny,nz,n_tot=0,0,0,0  # reinitiatlize
             serial += 1
    #return glob

def replace_by_word(origin_word, replace_word):
    nLines1 = []
    nLines2 = []
    f_name = file_name
    f_id = open(f_name, 'r')
    nLines1 = f_id.readlines()
    f_id.close()
    os.remove(f_name)
    nLen = len(nLines1)
    kk=0
    while kk < nLen:
        line = nLines1[kk]
        nLines2.append(line.replace(origin_word, replace_word))
        kk = kk+1
        #print line
    f_id = open(f_name, 'w')
    f_id.writelines(nLines2)
    f_id.close()


#==========================================================================
# start program
#==========================================================================

# check if argument is there
if len(sys.argv) == 2:
        file_name = sys.argv[1]
        replace_by_word("O1P", "OP1")
        replace_by_word("O2P", "OP2")
        replace_by_word("*", "'")
        get_globs(file_name)
        
else:
    message = """
ERROR: missing pdb filename as argument
usage: %s file.pdb""" %(sys.argv[0])
    sys.exit(message)
#==========================================================================

