/*
	mg2imod.c
	Converts between IMOD files and a micrograph parameter file
	Authors:  Bernard Heymann
	Created: 20070501	Modified: 20081031
	Compile with:
		g++ mg2imod.c -o mg2imod -I$BSOFT/include -L$BSOFT/lib -lbsoft
*/

#include "mg_processing.h"
#include "mg_img_proc.h"
#include "mg_tomography.h"
#include "rwmg.h"
#include "rwimg.h"
#include "linked_list.h"
#include "utilities.h"
#include "options.h"
#include "timer.h"

// Declaration of global variables
extern int 	verbose;		// Level of output to the screen
extern long memory;			// Total memory allocated

Bproject*	project_from_IMOD(Bstring* file_list, int flag);
int			write_IMOD_parameters(Bstring& imodfile, Bproject* project);

// Usage assistance
const char* use[] = {
" ",
"Usage: mg2imod [options] input.mrc input.xf input.tlt input.xyz",
"---------------------------------------------------------------",
"Converts between IMOD files and a micrograph parameter file.",
" ",
"Actions:",
"-topif                   Flag to convert the image file to multi-image PIF (default not).",
" ",
"Parameters:",
"-verbose 7               Verbosity of output.",
"-sampling 1.5            Sampling (A/pixel).",
" ",
"Output:",
"-output file.star        Output micrograph parameter file.",
"-imod file.xf            Output IMOD parameter files.",
" ",
NULL
};


int			main(int argc, char** argv)
{
	// Initializing variables
	int				flag = 0;					// Flag to rewrite the image file
	float			sampling;					// A/pixel
	Bstring			outfile;					// Output parameter file
	Bstring			imodfile;					// IMOD output parameter file

	int				optind;
	Option*			option = get_option_list(use, argc, argv, &optind);
	Option*			curropt;
	for ( curropt = option; curropt; curropt = curropt->next ) {
		if ( strcmp(curropt->tag, "topif") == 0 ) flag = 1;
		if ( strcmp(curropt->tag, "sampling") == 0 )
			if ( sscanf(curropt->value, "%f", &sampling) < 1 )
				fprintf(stderr, "-sampling: A pixel size must be specified\n");
		if ( strcmp(curropt->tag, "output") == 0 )
			outfile = get_option_filename(curropt->value);
		if ( strcmp(curropt->tag, "imod") == 0 )
			imodfile = get_option_filename(curropt->value);
    }
	option_kill(option);

	timeval			ti;
	if ( verbose & VERB_TIME )
		ti = timer_start();

	// Read all the parameter files
	Bstring*		file_list = NULL;
	while ( optind < argc ) string_add(&file_list, argv[optind++]);
	if ( !file_list ) {
		fprintf(stderr, "Error: No parameter files specified!\n");
		bexit(-1);
	}

	Bproject*		project = NULL;
	Bstring			ext = file_list->extension();
	
	if ( ext.contains("star") || ext.contains("xml") )
		project = read_project(file_list);
	else
		project = project_from_IMOD(file_list, flag);

	string_kill(file_list);
	
	if ( project == NULL )  {
		fprintf(stderr, "Error: No project generated!\n");
		bexit(-1);
	}

	if ( sampling > 0.1 )
		project_set_pixel_size(project, sampling);

	project_update_comment(project, argc, argv);

	if ( outfile.length() )
		write_project(outfile, project, 0, 0);
	
	if ( imodfile.length() )
		write_IMOD_parameters(imodfile, project);
	
	project_kill(project);

	if ( verbose & VERB_TIME )
		timer_report(ti);

	bexit(0);
}

/************************************************************************
@Function: project_from_IMOD
@Description:
	Creates a project structure using IMOD parameters.
@Algorithm:
	Requirements:
		Tilt series micrograph image (if 3D then converted to multi-2D)
		2D transform file (.xf)
		Tilt angle file (.tlt)
	Calculations:
		Tilt angle: from tilt angle file
		Tilt axis: ta = (atan2(-A11, A12) + atan2(-A22, -A21))/2
		Mg origin: o = on - Ad
	where
		A:  2x2 transformation matrix
		ta: tilt axis angle
		on: nominal micrograph origin (center of mg)
		o:  aligned micrograph origin
		d:  shift for micrograph
@Arguments:
	Bstring* file_list	list of IMOD parameter file names.
	int flag			flag to indicate conversion of the image file.
@Returns:
	Bproject*			new project parameter structure.
*************************************************************************/
Bproject*	project_from_IMOD(Bstring* file_list, int flag)
{
	if ( file_list->empty() ) {
		fprintf(stderr, "No filenames given to create a project!\n");
		exit(-1);
	}

	Bimage* 		p = NULL;
	Bstring			base, imgfile;
	Bstring*		thisfile = NULL;
	Bstring			ext;

	for ( thisfile = file_list; thisfile; thisfile = thisfile->next ) {
		ext = thisfile->extension();
		if ( ext.contains("pif") || ext.contains("mrc") || ext.contains("st") ) {
			imgfile = *thisfile;
			base = imgfile.base();
			if ( ext.contains("st") ) imgfile += ":mrc";
			p = read_img(imgfile, 0, -1);
			if ( !p ) {
				fprintf(stderr, "Error: Image file %s not read!\n", imgfile.c_str());
				error_show("project_from_IMOD", __FILE__, __LINE__);
				return(NULL);
			}
			if ( p->z > 1 ) {
				img_slices_to_images(p);
				if ( flag ) {
					kill_img(p);
					p = read_img(imgfile, 1, 0);
					img_slices_to_images(p);
					imgfile = base + ".pif";
					write_img(imgfile, p);
				}
			}
		}
	}

	if ( !p ) {
		fprintf(stderr, "Error: No image file provided!\n");
		error_show("project_from_IMOD", __FILE__, __LINE__);
		return(NULL);
	}
	
	int					j, n;
	Bstring				mg_id, rec_id, rec_name;
	Bproject*			project = project_init();
	Bfield*				field = field_add(&project->field, base);
	Bmicrograph*		mg = NULL;
	Breconstruction*	rec = NULL;
	Bmarker*			mark = NULL;
	
	for ( n=0; n<p->n; n++ ) {
		mg_id = base + Bstring(n+1, "_%03d");
		mg = micrograph_add(&mg, mg_id);
		if ( !field->mg ) field->mg = mg;
		mg->block = n;
		mg->img_num = n;
		mg->fmg = imgfile;
		mg->pixel_size = p->ux;
		mg->origin[0] = p->x/2;
		mg->origin[1] = p->y/2;
		mg->origin[2] = p->z/2;
		if ( verbose & VERB_FULL )
			printf("Creating field \"%s\", micrograph \"%s\"\n", 
						field->id.c_str(), mg->id.c_str());
	}

	kill_img(p);
	
	FILE*			fd = NULL;
	char			aline[100], *aptr;
	int				mark_id, dum;
	float			x, y, z;
	double			a[4], dx, dy, ta;
	
	for ( thisfile = file_list; thisfile; thisfile = thisfile->next ) {
		ext = thisfile->extension();
		if ( ext.contains("xf") ) {
			if ( ( fd = fopen(thisfile->c_str(), "r") ) == NULL ) {
				fprintf(stderr, "Error: A text file specifying 2D transformations must be given!\n");
				return(NULL);
			}
			if ( verbose & VERB_PROCESS )
				printf("Reading 2D transformations from \"%s\"\n", thisfile->c_str());
			for ( mg = field->mg; mg; mg = mg->next ) {
				fscanf(fd, "%lf %lf %lf %lf %lf %lf", &a[0], &a[1], &a[2], &a[3], &dx, &dy);
				mg->tilt_axis = (atan2(-a[0], a[1]) + atan2(-a[3], -a[2]))/2.0;
				mg->origin[0] -= a[0]*dx + a[2]*dy;
				mg->origin[1] -= a[1]*dx + a[3]*dy;
			}
			fclose(fd);
		} else if ( ext.contains("tlt") ) {
			if ( ( fd = fopen(thisfile->c_str(), "r") ) == NULL ) {
				fprintf(stderr, "Error: A text file specifying tilt angles must be given!\n");
				return(NULL);
			}
			if ( verbose & VERB_PROCESS )
				printf("Reading tilt angles from \"%s\"\n", thisfile->c_str());
			for ( mg = field->mg; mg; mg = mg->next ) {
				fscanf(fd, "%lf", &ta);
				mg->tilt_angle = ta*M_PI/180.0;
//				printf("%s: %g\n", mg->id.c_str(), ta);
			}
			fclose(fd);
		} else if ( ext.contains("xyz") ) {
			if ( ( fd = fopen(thisfile->c_str(), "r") ) == NULL ) {
				fprintf(stderr, "Error: A text file specifying the 3D marker model must be given!\n");
				return(NULL);
			}
			if ( verbose & VERB_PROCESS )
				printf("Reading the 3D marker model from \"%s\"\n", thisfile->c_str());
			rec_id = base;
			rec = reconstruction_add(&project->rec, rec_id);
			rec_name = base + ".rec";
			if ( access(rec_name.c_str(), R_OK) ) {
				fgets(aline, 100, fd);
				aptr = strstr(aline, "Pix");
				sscanf(aptr, "Pix: %f Dim: %f %f\n", &rec->voxel_size, &rec->origin[0], &rec->origin[1]);
				rec->origin /= 2;
				rewind(fd);
			} else {
				rec->frec = rec_name;
				rec_name += ":mrc";
				p = read_img(rec_name, 0, -1);
				rec->voxel_size = p->ux;
				rec->origin[0] = p->x/2;
				rec->origin[1] = p->y/2;
				rec->origin[2] = p->z/2;
				kill_img(p);
			}
			while ( fgets(aline, 100, fd) ) {
				if ( sscanf(aline, "%d %f %f %f %d %d", &mark_id, &x, &y, &z, &dum, &dum) > 5 ) {
					mark = (Bmarker *) add_item((char **) &mark, sizeof(Bmarker));
					if ( !rec->mark ) rec->mark = mark;
					mark->id = mark_id;
					mark->loc = Vector3<float>(x - rec->origin.x(), y - rec->origin.y(), z);
					mark->fom = 1;
					mark->sel = 1;
				}
			}
			fclose(fd);
		}
	}

	project_mg_tilt_to_matrix(project);
	
	return(project);
}

/************************************************************************
@Function: write_IMOD_parameters
@Description:
	Creates IMOD parameter files from a project structure.
@Algorithm:
	Writes both .xf and .tlt files.
	Calculations:
		Tilt angle: from micrograph tilt angle
		Matrix: A = {cos(ta), -sin(ta), sin(ta), cos(ta)}
		Shift:  d = o - on
	where
		A:  2x2 transformation matrix
		ta: tilt axis angle
		on: nominal micrograph origin (center of mg)
		o:  aligned micrograph origin
		d:  shift for micrograph
@Arguments:
	Bstring& imodfile	IMOD parameter file name.
	Bproject* project	project structure.
@Returns:
	int					0.
*************************************************************************/
int			write_IMOD_parameters(Bstring& imodfile, Bproject* project)
{
	Bstring				paramfile = imodfile.base() + ".xf";
	Bmicrograph*		mg = project->field->mg;
	Breconstruction*	rec = project->rec;
	Bmarker*			mark = NULL;

	Vector3<float>		d, origin = micrograph_get_nominal_origin(mg);

	FILE*				fd = NULL;

	if ( ( fd = fopen(paramfile.c_str(), "w") ) == NULL ) {
		error_show("write_IMOD_parameters", __FILE__, __LINE__);
		return(NULL);
	}
	
	if ( verbose & VERB_PROCESS )
		printf("Writing 2D transformations to \"%s\"\n", paramfile.c_str());
	
	double				a[4];
	
	for ( mg = project->field->mg; mg; mg = mg->next ) {
		a[0] = a[3] = -sin(mg->tilt_axis);
		a[1] = cos(mg->tilt_axis);
		a[2] = -cos(mg->tilt_axis);
		d = origin - mg->origin;
		fprintf(fd, "%12.7lf%12.7lf%12.7lf%12.7lf%12.3lf%12.3lf\n", a[0], a[1],
			a[2], a[3], a[0]*d.x() + a[1]*d.y(), a[2]*d.x() + a[3]*d.y());
	}
	
	fclose(fd);

	paramfile = imodfile.base() + ".tlt";
	
	if ( ( fd = fopen(paramfile.c_str(), "w") ) == NULL ) {
		error_show("write_IMOD_parameters", __FILE__, __LINE__);
		return(NULL);
	}

	if ( verbose & VERB_PROCESS )
		printf("Writing tilt angles to \"%s\"\n", paramfile.c_str());
	
	for ( mg = project->field->mg; mg; mg = mg->next )
		fprintf(fd, "%g\n", mg->tilt_angle*180.0/M_PI);
	
	fclose(fd);

	if ( rec && rec->mark ) {
		paramfile = imodfile.base() + "fid.xyz";
	
		if ( ( fd = fopen(paramfile.c_str(), "w") ) == NULL ) {
			error_show("write_IMOD_parameters", __FILE__, __LINE__);
			return(NULL);
		}

		if ( verbose & VERB_PROCESS )
			printf("Writing the 3D marker model to \"%s\"\n", paramfile.c_str());
	
		for ( mark = rec->mark; mark; mark = mark->next ) {
			fprintf(fd, "%4d%10.2f%10.2f%10.2f %6d %4d", mark->id, 
				mark->loc.x() + rec->origin.x(), mark->loc.y() + rec->origin.y(), mark->loc.z(), 1, mark->id);
			if ( mark == rec->mark )
				fprintf(fd, " Pix: %11.5f Dim: %5.0f %5.0f\n", 
					rec->voxel_size, 2*rec->origin.x(), 2*rec->origin.y());
			else
				fprintf(fd, "\n");
		}
	
		fclose(fd);
	}

	return(0);
}