#!/bin/python
# Read R1R2 files produced by fit_all.py and COLUMNS file (see below),
# and produce means and std-devs for each repeat residue by:
# 1) averaging over all corresponding residues for a given sample (e.g., columns 2 4 5 9, averaged over residues 262, 281, 300, etc)
# 2) computing std-deviations per repeat residue based on the different samples
#
# Sample columns are extracted from the COLUMNS file generated by fit_all.py

##### IMPORTS #####
import glob
import os.path
import numpy as np
import warnings

###### PARAMS ######
PERIODICITY=125 # repeat length (or 125 for no averaging)
RESIDS=range(251,251+PERIODICITY+1,1) # resids to analyze
OFFSET=0 # offset for periodicity. For example, first value is for (resid + OFFSET) % PERIODICITY + OFFSET == 0 ; 0 for no averaging

##### Main: #####
if __name__=='__main__':
  # Process columns for each specified residue
  Rs=  np.array(PERIODICITY*[0.0],dtype=float)
  Rs2= np.array(PERIODICITY*[0.0],dtype=float)
  Ns=  np.array(PERIODICITY*[0.0],dtype=float)
  IDs= np.array(PERIODICITY*[0.0],dtype=float)
  F=open('COLUMNS','r');
  for line in F:
    columns=[int(x) for x in line.split()]   # convert to int to verify proper input
    for resid in RESIDS:
      r1r2file= '%d.txt.%d.%d.%d.%d.fit7.R1R2' % (resid,columns[0],columns[1],columns[2],columns[3])
      if os.path.exists(r1r2file):
        FR= open(r1r2file,'r')
        FRlines= FR.readlines()
        R2R1= float(FRlines[2])
        id= (resid+OFFSET) % PERIODICITY 
        if id>0:
          IDs[id]=id;
        else:
          IDs[id]=id+PERIODICITY
        Rs[id]=Rs[id]+R2R1
        Rs2[id]=Rs2[id]+R2R1**2
        Ns[id]=  Ns[id]+1.0
        FR.close()
  F.close()
  E=  np.divide(Rs,Ns)
  E2= np.divide(Rs2,Ns)
#  print "Ns: "
#  print [int(x) for x in Ns]
#  print "Means: "
#  for x in E:
#    print "%5.2f " % x,
#  print
#  print "Std-devs:"
  S=np.sqrt(np.subtract(E2,np.multiply(E,E)))
  with warnings.catch_warnings():
    warnings.simplefilter("ignore")
#    for x in S:
#      if(np.isnan(x)): 
#        print "  nan ",
#      else:
#        print "%5.2f " % x,
#    print
  for i,id in enumerate(IDs):
    if(np.isnan(S[i])):
      print "%5d %6.2f    nan" % (id,E[i])
    else:
      print "%5d %6.2f %6.2f" % (id,E[i],S[i])
    

  
