#!/usr/bin/env python


import sys
import string
import random
from EMAN2 import *
import os
import math
import numpy as np


def main():
        if len(sys.argv) < 9:
                print "usage :" + sys.argv[0] + " <input refined data star> <output symmetry expanded data star> <string for C symmetry> <x coordinate> <y coordinate> <z coordinate> <boxsize> <debug level-1 for plus; 0 for minus> <bining level>"
                sys.exit()

	data_star = sys.argv[1]
	output_star = sys.argv[2]
	c_symmetry = sys.argv[3]
	x_coord = float(sys.argv[4])
	y_coord = float(sys.argv[5])
	z_coord = float(sys.argv[6])
	box = float(sys.argv[7])
	debug = sys.argv[8]
	bin_level = float(sys.argv[9])

	PI = 3.14159265359
	scale = 1.0000000000000
	tap = '\t'
	n = '\n'

	data_line = (open(data_star, 'r').read()).splitlines()
	fw = open(output_star, 'w+')

	fw.write(n + 'data_' + n + n + 'loop_' + n + '_rlnMicrographName #1' + n + '_rlnCoordinateX #2' + n + '_rlnCoordinateY #3' + n + '_rlnVoltage #4' + n + '_rlnDefocusU #5' + n + '_rlnDefocusV #6' + n + '_rlnDefocusAngle #7' + n + '_rlnSphericalAberration #8' + n + '_rlnDetectorPixelSize #9' + n + '_rlnCtfFigureOfMerit #10' + n + '_rlnMagnification #11' + n + '_rlnAmplitudeContrast #12' + n + '_rlnAutopickFigureOfMerit #13' + n + '_rlnPhaseShift #14' + n + '_rlnImageName #15' + n + '_rlnOriginX #16' + n + '_rlnOriginY #17' + n + '_rlnAngleTilt #18' + n + '_rlnGroupNumber #19' + n + '_rlnAngleRot #20' + n + '_rlnAnglePsi #21' + n + '_rlnClassNumber #22' + n + '_rlnNormCorrection #23' + n + '_rlnLogLikeliContribution #24' + n + '_rlnMaxValueProbDistribution #25' + n + '_rlnNrOfSignificantSamples #26' + n + '_rlnRandomSubset #27' + n + '_rlnCtfMaxResolution #28' + n)

	Cn = []
	for k in range(0, int(c_symmetry)):
		Cn.append([])
		Cn[k] = k * 360./float(c_symmetry)
	for i in range(32, len(data_line)):
		data_str = data_line[i].split()
	#read in irrelevant parameters
		mic = data_str[6]
		img = data_str[5]
		defa = data_str[14]		
		ps = data_str[16]
		volt = data_str[11]
		cs = data_str[15]
		ctf_fom = data_str[10]
		amp = data_str[17]
		autopick_fom = data_str[4]
		group = data_str[20]
		cls = data_str[2]
		norm = data_str[23]
		rand = data_str[27]
		log_likeli = data_str[24]
		prob_distribu = data_str[25]
		nr_sample = data_str[26]
		ctf_res = data_str[9]

		defu = float(data_str[12])
		defv = float(data_str[13])
		rot = float(data_str[21])
		tilt = float(data_str[22])
		psi = float(data_str[3])

		mic_x = float(data_str[0])
		mic_y = float(data_str[1])

		x_shift = float(data_str[18])
		y_shift = float(data_str[19])

		center_x = float(mic_x) - bin_level * float(x_shift)
                center_y = float(mic_y) - bin_level * float(y_shift)

		d_size = data_str[8]
		mag = data_str[7]
		apix = float(d_size) / float(mag) * 10000

		t = Transform()
		t = Transform({"type":"spider","phi":float(rot),"theta":float(tilt),"psi":float(psi)})
		t_eman = t.get_rotation("eman")
		alt, az, phi = t_eman['alt'], t_eman['az'], t_eman['phi']
		#print("rot:{} - tilt:{} - psi:{}".format(rot, tilt, psi))
		#print("alt:{} - az:{} - phi:{}".format(alt, az, phi))

		euler1 = alt / 180 * PI
		euler2 = az / 180 * PI
		euler3 = phi / 180 * PI

		mx0=scale*(math.cos(euler3)*math.cos(euler2)-math.cos(euler1)*math.sin(euler2)*math.sin(euler3))
		mx3=-scale*(math.sin(euler3)*math.cos(euler2)+math.cos(euler1)*math.sin(euler2)*math.cos(euler3))
		mx6=scale*math.sin(euler1)*math.sin(euler2)
		mx1=scale*(math.cos(euler3)*math.sin(euler2)+math.cos(euler1)*math.cos(euler2)*math.sin(euler3))
		mx4=scale*(-1*math.sin(euler3)*math.sin(euler2)+math.cos(euler1)*math.cos(euler2)*math.cos(euler3))
		mx7=-scale*math.sin(euler1)*math.cos(euler2)
		mx2=scale*math.sin(euler1)*math.sin(euler3)
		mx5=scale*math.sin(euler1)*math.cos(euler3)
		mx8=scale*math.cos(euler1)

		for j in range(0, int(c_symmetry)):
			c_euler1=0.
			c_euler2=float(str(Cn[j]))/180*PI
			c_euler3=0.

			c_mx0=(math.cos(c_euler3)*math.cos(c_euler2)-math.cos(c_euler1)*math.sin(c_euler2)*math.sin(c_euler3))
			c_mx1=-1*(math.sin(c_euler3)*math.cos(c_euler2)+math.cos(c_euler1)*math.sin(c_euler2)*math.cos(c_euler3))
			c_mx2=math.sin(c_euler1)*math.sin(c_euler2)
			c_mx3=(math.cos(c_euler3)*math.sin(c_euler2)+math.cos(c_euler1)*math.cos(c_euler2)*math.sin(c_euler3))
			c_mx4=(-1*math.sin(c_euler3)*math.sin(c_euler2)+math.cos(c_euler1)*math.cos(c_euler2)*math.cos(c_euler3))
			c_mx5=-1*math.sin(c_euler1)*math.cos(c_euler2)
			c_mx6=math.sin(c_euler1)*math.sin(c_euler3)
			c_mx7=math.sin(c_euler1)*math.cos(c_euler3)
			c_mx8=math.cos(c_euler1)

			c_mult_mx0 = mx0*c_mx0 + mx1*c_mx3 + mx2*c_mx6
			c_mult_mx1 = mx0*c_mx1 + mx1*c_mx4 + mx2*c_mx7
			c_mult_mx2 = mx0*c_mx2 + mx1*c_mx5 + mx2*c_mx8
			c_mult_mx3 = mx3*c_mx0 + mx4*c_mx3 + mx5*c_mx6
			c_mult_mx4 = mx3*c_mx1 + mx4*c_mx4 + mx5*c_mx7
			c_mult_mx5 = mx3*c_mx2 + mx4*c_mx5 + mx5*c_mx8
			c_mult_mx6 = mx6*c_mx0 + mx7*c_mx3 + mx8*c_mx6
			c_mult_mx7 = mx6*c_mx1 + mx7*c_mx4 + mx8*c_mx7
			c_mult_mx8 = mx6*c_mx2 + mx7*c_mx5 + mx8*c_mx8
			
			new_euler1=math.acos(c_mult_mx8)
			new_euler2=math.atan2(c_mult_mx6,-1*c_mult_mx7)
			new_euler3=math.atan2(c_mult_mx2,c_mult_mx5)

			if(c_mult_mx8>0.999999):
				new_euler1=0
				new_euler2=math.atan2(c_mult_mx1,c_mult_mx4)
				new_euler3=0
			if(c_mult_mx8<-0.999999):
				new_euler1=PI
				new_euler2=math.atan2(c_mult_mx1,-1*c_mult_mx4)
				new_euler3=0


			new_x=c_mult_mx0*(x_coord-box/2)+c_mult_mx1*(y_coord-box/2)+c_mult_mx2*(z_coord-box/2)
			new_y=c_mult_mx3*(x_coord-box/2)+c_mult_mx4*(y_coord-box/2)+c_mult_mx5*(z_coord-box/2)
			new_z=c_mult_mx6*(x_coord-box/2)+c_mult_mx7*(y_coord-box/2)+c_mult_mx8*(z_coord-box/2)

			new_mic_x = int(math.floor(bin_level * new_x + center_x + 0.5))
			new_mic_y = int(math.floor(bin_level * new_y + center_y + 0.5))

			if debug == str(1):
				new_z_defu = new_z * apix + defu
				new_z_defv = new_z * apix + defv

			else:
				new_z_defu = -1 * new_z * apix + defu
				new_z_defv = -1 * new_z * apix + defv


			new_alt = new_euler1/PI*180
			new_az = new_euler2/PI*180
			new_phi = new_euler3/PI*180

			w = Transform()
			w = Transform({"type":"eman","alt":new_alt,"az":new_az,"phi":new_phi})
			w_relion = w.get_rotation("spider")
			new_rot, new_tilt, new_psi = w_relion['phi'], w_relion['theta'], w_relion['psi']


			newline = mic + tap + str(new_mic_x) + tap + str(new_mic_y) + tap + volt + tap + str(new_z_defu) + tap + str(new_z_defv) + tap + defa + tap + cs + tap + d_size + tap + ctf_fom + tap + mag + tap + amp + tap + autopick_fom + tap + ps + tap + img + tap + '0' + tap + '0' + tap + str(new_tilt) + tap + group + tap + str(new_rot) + tap + str(new_psi) + tap + cls + tap + norm + tap + log_likeli + tap + prob_distribu + tap + nr_sample + tap + rand + tap	+ ctf_res + n
			fw.write(newline)



if __name__ == "__main__":
    main()