#!/usr/bin/env python

'''
Usage:
orient_object.py filename output.txt

- filename = path and name of the stack file containing the object (background=0)

Output:
the image with the object reoriented:
- CM in the middle of the image
- first principal axis along y
- second principal axis along x

'''

from sys import argv
import os
import numpy as np
from scipy.ndimage.morphology import binary_closing, binary_opening

from titk_tools.io import imread, imsave, tfread, SpatialImage
from titk_tools.registration import isotropic_resampling_seg, isotropic_resampling
from titk_tools.registration import apply_transfo_on_seg1, apply_transfo_on_image

from bib_sepals import *

print('============== orient_sepal.py =================')

filename = argv[1]
outputdir=argv[2]+'/'

segmented=True

if len(argv)>3:
	imagetype=argv[3]
	if imagetype=='grey':
		segmented=False

print(argv)
print('segmented=',segmented)

imname = filename.split('/')[-1]
imnameroot = imname.split('.')[0]

outfilename = outputdir+imnameroot+'_oriented.inr.gz'

# read the image
img=imread(filename)
# resample isotropically
print("----------------------")
print("isotropic resampling of : filename=", filename)
print("----------------------")
voxelsize = img.voxelsize
print("voxelsize=",voxelsize)
s = img.shape
imgtype=img.dtype
print("shape=",s)

radius=int(voxelsize[2]/voxelsize[0])
isovs=(voxelsize[0]*voxelsize[1]*voxelsize[2])**(1.0/3)
#isovs=voxelsize[0]
print("new voxelsize=",isovs)

sepale=SpatialImage(img>0, dtype='uint8', voxelsize=voxelsize, origin=[0, 0, 0])
if segmented:
	sepale=isotropic_resampling_seg(sepale,isovs)
else:
	img=isotropic_resampling(img,isovs)
	sepale=isotropic_resampling_seg(sepale,isovs)

s = sepale.shape
voxelsize = sepale.voxelsize
print('new shape', s)


# alongate the target image so that a sepal originally along the diagonal fits into the target image
print("----------------------")
print("elongating field : filename=", filename)
print("----------------------")
s_addy=int((np.sqrt(s[0]*s[1]*2)-np.sqrt(s[0]*s[1]))/2)
s_addz=int(s[2]*0.20)
sepale=add_side_slices(sepale, 0, 0, s_addy, s_addy, s_addz, s_addz)
if not segmented :
	img=add_side_slices(img, 0, 0, s_addy, s_addy, s_addz, s_addz)

'''
# smooth it with opening
#sepale=imread(isoname)
s = sepale.shape
voxelsize = sepale.voxelsize
struct_el=struct_element_sphere(radius)
sepale = binary_closing(sepale, structure=struct_el, iterations=int(radius))
struct_el=struct_element_sphere(radius/2)
sepale = binary_closing(sepale, structure=struct_el, iterations=int(radius/2))
sepale=SpatialImage(sepale*255, dtype='uint8', voxelsize=voxelsize, origin=[0, 0, 0])
#imsave(sname, sepale)
'''

# put the center of mass in the center of the image
sepale_centered=center_on_cm(sepale, volumic=True)
if not segmented:
	img=center_on_cm(img, volumic=True)

# compute principal directions
pv_flo=principal_directions(sepale_centered, volumic=True)
print(pv_flo)

# sepals originally are oriented with top in more-or less "up" direction on the images
# while the eigenvectors' orientation is arbitrary
# --> do not change the original rough orientation

# check logitudinal direction:
if pv_flo[0][1]<0 :
	pv_flo[0] = -pv_flo[0]

# check the transversal direction
if pv_flo[1][0]<0 :
	pv_flo[1] = -pv_flo[1]

print("New pv_flo",pv_flo)

'''
pv_ref = [np.array([0,-1,0]),
np.array([-1,0,0])]

'''
pv_ref = [np.array([0,1,0]),
np.array([1,0,0])]


# compute superposing transformation
trsf_file=outputdir+imnameroot+'_rigid.txt'
write_superposing_transfo(pv_flo, pv_ref, sepale_centered, trsf_file)

trsf=tfread(trsf_file)
os.system("rm "+trsf_file)

# apply transformation
if segmented :
	oriented=apply_transfo_on_seg1(trsf, sepale_centered)
else :
	oriented=apply_transfo_on_image(trsf, img)

# smooth the transformed image
#sepale = binary_closing(oriented, structure=struct_el, iterations=radius)

# and save the oriented sepal
sepale_filename = outputdir+imnameroot+'_oriented.inr.gz'
if segmented :
	sepale=SpatialImage(oriented*255, dtype='uint8', voxelsize=voxelsize, origin=[0, 0, 0])
else:
	sepale=SpatialImage(oriented, dtype=imgtype, voxelsize=voxelsize, origin=[0, 0, 0])

imsave(sepale_filename, sepale)