#!/usr/struct/bin/python
# This script will go through all the individual tomogram directories
# get the subtomo/tomo/ctfsubtomo/coords and put them in a data.star file for a combined relion run.
# By Tanmay Bharat, MRC-LMB 05.11.2013

import os
import sys
import commands
import numpy as np

######## INPUT ###########################

#starname = workdir + 'Refine3D/wbp1_it009_data.star'
starname = 'run2_ct15_it020_data.star'

# binning
binfactor = 1.0

# Output star name
outfilename = '../particles_subtomo_good.star'

######## INPUT ###########################


######## FUNCTIONS #######################

def read_relion_header(filename):

  print 'Reading the header of the relion star file'
  ifile = open(filename, 'r')

  # first get the header
  i=-1
  XextCol=0
  YextCol=0
  ZextCol=0
  XaliCol=0
  YaliCol=0
  ZaliCol=0
  RotCol=0
  TiltCol=0
  PsiCol=0
  AlignOrNot=0
  MaxLogValCol=0
  MicNameCol=0
  ImNameCol=0
  CtfImNameCol=0
  ClassNumCol=0

  for line in ifile:
    
    emptycheck = line.isspace()
    if(emptycheck):
      continue

    fields = line.split()
    if fields[0] == 'data_' or fields[0] == 'loop_':
      continue

    i= i+1

    firstcol = fields[0]

    if firstcol == '_rlnCoordinateX':
      XextCol = i
    if firstcol == '_rlnCoordinateY':
      YextCol = i
    if firstcol == '_rlnCoordinateZ':
      ZextCol = i
    if firstcol == '_rlnOriginX':
      XaliCol = i
    if firstcol == '_rlnOriginY':
      YaliCol = i
    if firstcol == '_rlnOriginZ':
      ZaliCol = i
    if firstcol == '_rlnAngleRot':
      RotCol = i
    if firstcol == '_rlnAngleTilt':
      TiltCol = i
    if firstcol == '_rlnAnglePsi':
      PsiCol = i
    if firstcol == '_rlnMaxValueProbDistribution':
      MaxLogValCol = i
    if firstcol == '_rlnImageName':
      ImNameCol = i
    if firstcol == '_rlnCtfImage':
      CtfImNameCol = i
    if firstcol == '_rlnMicrographName':
      MicNameCol = i
    if firstcol == '_rlnClassNumber':
      ClassNumCol = i

    if firstcol[0] != '_':
      break
  
    if MaxLogValCol > 0:
      AlignOrNot = 1

  print 'XextCol,YextCol,ZextCol,XaliCol,YaliCol,ZaliCol,RotCol,TiltCol,PsiCol,AlignOrNot,MaxLogValCol,ImNameCol,CtfImNameCol,MicNameCol'
  print XextCol,YextCol,ZextCol,XaliCol,YaliCol,ZaliCol,RotCol,TiltCol,PsiCol,AlignOrNot,MaxLogValCol
  return(XextCol,YextCol,ZextCol,XaliCol,YaliCol,ZaliCol,RotCol,TiltCol,PsiCol,AlignOrNot,MaxLogValCol,ImNameCol,CtfImNameCol,MicNameCol,ClassNumCol)
      



######## FUNCTIONS #######################



######## RUNNING THE SCRIPT #################

grepline = 'grep ' + 'Tomograms ' + starname + '| grep -v "#" | grep -v "loop" | grep -v "data" | awk "NF"  > temp.txt'
print grepline
os.system(grepline)

infile=open('temp.txt')
data=infile.readlines()

# Getting the column Numbers
XextCol,YextCol,ZextCol,XaliCol,YaliCol,ZaliCol,RotCol,TiltCol,PsiCol,AlignOrNot,MaxLogValCol,ImNameCol,CtfImNameCol,MicNameCol,ClassNumCol = read_relion_header(starname)

######

######
ofile = open(outfilename, 'w')
ofile.write('data_' + '\n' + '\n')
ofile.write('loop_' + '\n')
ofile.write('_rlnMicrographName #1' + '\n')
ofile.write('_rlnCoordinateX #2' + '\n')
ofile.write('_rlnCoordinateY #3'+ '\n')
ofile.write('_rlnCoordinateZ #4' + '\n')
ofile.write('_rlnImageName #5' + '\n')
ofile.write('_rlnCtfImage #6' +'\n')

for i in range(0, len(data)):

  pair = data[i].split()
  
  ### X,Y,Z extraction
  Xextraction = float(pair[XextCol])
  Yextraction = float(pair[YextCol])
  Zextraction = float(pair[ZextCol])

  ### X,Y,Z alignment
  colorclass = 1
  if AlignOrNot == 1:
    Xali = float(pair[XaliCol])
    Yali = float(pair[YaliCol])
    Zali = float(pair[ZaliCol])
  if AlignOrNot == 0:
    Xali = 0.0
    Yali = 0.0 
    Zali = 0.0

  ### New X,Y,Z
  X = Xextraction
  Y = Yextraction
  Z = Zextraction

  ### Euler angles
  Rot = float(pair[RotCol])
  Tilt = float(pair[TiltCol])
  Psi = float(pair[PsiCol])

  # Image and CTF image
  ImageName = str(pair[ImNameCol])
  CtfImage = str(pair[CtfImNameCol])
  MicrographName = str(pair[MicNameCol])

  # Class Number
  ClassNumber = float(pair[ClassNumCol])

  if ClassNumber == 2 or ClassNumber ==3:
    outline = MicrographName + '\t' + str("%.2f" % X) + '\t' + str("%.2f" % Y) + '\t' + str("%.2f" % Z) + '\t' + ImageName + '\t' +CtfImage + '\n' 
    ofile.write(outline)
  
sys.exit()