#!/usr/bin/perl -w

# hdrprep: register exposure-bracketed digicam images
# Copyright (C) 2006  Axel Jacobs
#
# Version: 0.1.1, 4 Aug 2006
# 
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor,
# Boston, MA  02110-1301, USA.

require 5.001;
use strict;

# Prerequisites
use Image::Magick;
use Getopt::Long;
use Image::ExifTool;
use Math::Trig;
use Tie::File;

# What are we doing?
my $ALIGN='';
my $REDO='';
my $EXIF;
my $HELP;

# Variables that may we overwritten by command-line args
# are all in CAPITALS.
my $VERBOSE;
my $DIR = "aligned";
my $QUAL = "80";
my $KEEP;
my $NOORDER=0;
my $transFile;
my $panofile="$DIR/panofile.pto";

sub init
{
	# Process command line arguments
	Getopt::Long::Configure ('bundling');
	GetOptions ("q|quality=i" => \$QUAL,
			"d|directory=s" => \$DIR,
			"v|verbose" => \$VERBOSE,
			"r|redo=s" => \$REDO,
			"a|align=s" => \$ALIGN,
			"k|keep" => \$KEEP,
			"e|exif" => \$EXIF,
			"n|noorder" => \$NOORDER,
			"h|help" => \$HELP );

	if ($ALIGN) {
		if ($ALIGN eq "ale") {
			print "using aligning engine: ale \n"
		}
		elsif ($ALIGN eq "sift") {
			print "using aligning engine: sift \n"
		}
	}

	$HELP && usage();
	if (! ($ALIGN || $REDO || $EXIF)) {
		print STDERR "Error: Tell me what to do [-a|-r] [-e].\n";
		hint();
	}
	$transFile = "$DIR/ale_align.trans";

	my @imgs;
	if ($ARGV[0]) {;
		foreach (@ARGV) {
			if (! -e $_) {
				print STDERR "Error: File $_ doesn't exists.\n";
				hint();
			}
			push @imgs, $_;
		} 
	} else {
		print STDERR "Error: Please supply at least two JPG files.\n";
		hint();
	}

	if (! -e $DIR) {
		if ($REDO) {
			print STDERR "Error: You need to run me with -a first.\n";
			hint();
		} else {
			print "Creating directory $DIR.\n";
			mkdir $DIR, 0700 or
				die ("Can\'t create directory $DIR: $!\n");
		}
	}

	my $allJpegs = join(', ', @imgs);
	print "Processing files: $allJpegs.\n";
	($ALIGN || $REDO) &&
			print "Setting quality of aligned images to $QUAL.\n";
	print "Using directory $DIR to store new images.\n";

	return @imgs;
}

sub hint
{
	print "Type hdrprep -h or --help for usage instructions.\n";
	exit;
}

sub usage
{
	print STDERR <<EndOfUsage;

       hdrprep : Align digicam images and fix EXIF information.

 Prerequisites : ALE, Anti-Lamenessing Engine (only for ale engine),
                 hugin and autopano-sift (for sift engine),
                 Perl::Image::ExifTool

         hdrprep [-h|--help] [-v|--verbose] \
	             [-a|--align] [-r|--redo] [-e|--exif] [-k|--keep] \
	             [-d|--directory dir] [-q|--quality qual] images

-h|--help                           : print this help message and exit
-v|--verbose                        : verbose output
-a|--align <ale|sift>               : align images
-r|--redo  <ale|sift>               : re-create aligned images, requires -a to be run first
-e|--exif                           : fix exposure tags in EXIF header
-k|--keep                           : don\'t remove intermediate ALE output
-d|--directory                      : directory for storing the aligned images in (default: aligned)
-q|--quality                        : quality setting for aligned images (default: 80)
-n|--noorder                        : do not order my exposure (default: order)
EndOfUsage
	exit;
}

sub sortByExposure
{
	# Determine the exposure for each image
	my @jpgs = @_;
	my %exposure;
	for (my $i=0 ; $i<=$#jpgs ; $i++) {
		my $exif_data = getExifTags ($jpgs[$i]);
		my $shutter_speed = getShutter ($exif_data);
		($shutter_speed == 0) &&
				die ("Can\'t determine Shutter Speed for file $jpgs[$i].\n");
		my $aperture = getAperture ($exif_data);
		($aperture == 0) &&
				die ("Can\'t determine Aperture setting for file $jpgs[$i].\n");
		my $iso_speed = getFilmSpeed ($exif_data);
		if ($iso_speed == 0) {
			print STDERR "Can\'t determine ISO film speed for file $jpgs[$i]. " .
					"Assuming ISO 100.\n";
			$iso_speed = 100;
		}
		my $expo = doExposure ($shutter_speed, $aperture, $iso_speed);
		$exposure{$expo} = $i;
	}

	my $key;
	my @imgs;
	# Sort by exposure
	foreach $key (sort {$a <=> $b} keys %exposure) {
		my $keyRnd = sprintf "%.2f", $key;
		$VERBOSE && print "$jpgs[$exposure{$key}]: $keyRnd EV\n";
		push @imgs, $jpgs[$exposure{$key}];
	}
	return @imgs;
}

sub runALE
{
	my @imgs = @_;
	# Align images in pairs of descending exposure
	my $aleOpts = "--euclidean --mc 10 --exp-extend";
	for ( my $i=0 ; $i<$#imgs ; $i++ ) {
		my $outFile = "$DIR/ale_out" . $i . ".jpg";
		`ale $aleOpts --trans-save=$transFile$i $imgs[$i] $imgs[$i + 1] $outFile`;
		unlink $outFile unless $KEEP;
	}
}

sub runAUTOPANOCOMPLETE
{
	my @imgs = @_;
	#launch autopano-complete to find control points.
	# TODO make it so the user can change -s 700 an -p 10
	my $autopanoOpts = "-o $panofile -p 3 2>&1";
	my $autopanoOutput=`./autopano-complete.sh $autopanoOpts @imgs`;
	$VERBOSE && print "$autopanoOutput  \n";
}

# fix some values that autopano.exe (i.e. autopano-complete.sh) sets... your mileage may vary
sub changeHorizV
{
	my @array;
	tie @array, 'Tie::File', "$panofile" or die "$panofile not found, exiting...";
	for (@array) {
# 		m/^i/ && print "before:" . $_ . "\n";
		m/^i/ && s/v\d+/v100/; #change horizontal field of view so that autooptimizer does not complain
		m/^i/ && s/b-?\d+\.\d+/b0/;# && print "after:" . $_ . "\n";
	}
	untie @array;
}

sub runAutoOptimiser
{
	print "optimizing variables \n";
	`autooptimiser -o $panofile $panofile`;
	`rm -f hugin_debug_optim_results.txt`
}

sub getExifTags
{
	my $tmpFile = shift;

	# Print the info necessary to determine the exposure 
	# from the EXIF header
	my $exifTool = new Image::ExifTool;
	my $exifData = $exifTool->ImageInfo($tmpFile);

	return $exifData;
}

sub getShutter
{
	my $exifTags = shift;
	my $shutterSpeed;

	if ($$exifTags{ExposureTime} or $$exifTags{ShutterSpeed}) {
		$shutterSpeed = $$exifTags{ExposureTime} or
				$shutterSpeed = $$exifTags{ShutterSpeed};
		$shutterSpeed =~ m/(\d+)([\.\/])?(\d+)?/;
		$shutterSpeed = $1;
		if (defined $3) {
			my $tmp = $3;
			if ($2 =~ m/\//) {
				# e.g. 6/10
				$shutterSpeed = $shutterSpeed / $tmp;
			} else {
				# e.g 0.6
				$shutterSpeed = $shutterSpeed . '.' . $tmp;
			}
		}
		$shutterSpeed = sprintf "%.5f", $shutterSpeed;
	} else {
		$shutterSpeed = 0;
	}
	return $shutterSpeed;
}

sub getAperture
{
	my $exifTags = shift;
	my $aperture;

	if ($$exifTags{FNumber} or $$exifTags{ApertureValue}) {
		$aperture = $$exifTags{FNumber} or
				$aperture = $$exifTags{ApertureValue};
		$aperture =~ m/(\d+)\.(\d+)?/;
		$aperture = $1;
		if (defined $2) {
			$aperture = $aperture . '.' . $2;
		}
	} else {
		$aperture = 0;
	}
	return $aperture;
}

sub getFilmSpeed
{
	my $exifTags = shift;
	my $filmSpeed;

	if ($$exifTags{ISO} or $$exifTags{CCDISOSensitivity}) {
		$filmSpeed = $$exifTags{ISO} or
				$filmSpeed = $$exifTags{CCDISOSensitivity};
		$filmSpeed =~ m/(\d+)/;
		$filmSpeed = $1;
	} else {
		$filmSpeed = 0;
	}
	return $filmSpeed;
}

sub doExposure
{
	my $shutterSpeed = shift;
	my $aperture = shift;
	my $filmSpeed = shift;
	my $exposure;

	# Exposure = f(Shutter_speed, Aperture, ISO speed):
	# EV = log2(aperture2 x (1/shutter speed) x (ISO sensitivity/100))
	# log base2 (x) = log (x) / log (2)
	if ($shutterSpeed > 0) {
		$exposure = log( $aperture*$aperture / $shutterSpeed * $filmSpeed/100) / log(2);
	} else {
		$exposure = 0;
	}
	return $exposure;
}

sub copyFiles
{
	# If neither -a nor -r, but only -e was requested, copy the original files
	# to the new directory for fixing the EXIF info.
	use File::Copy;
	my @imgs = @_;
	for (my $i=0 ; $i<=$#imgs ; $i++) {
		if (! -e "$DIR/$imgs[$i]") {
			copy("$imgs[$i]", "$DIR/$imgs[$i]");
		}
	}
}

sub fixEXIF
{
	# Correct exposure information of the new images
	# We don't need to check for correctness again.
	my @imgs = @_;
	for (my $i=0 ; $i<=$#imgs ; $i++) {
		my $exifData = getExifTags ($imgs[$i]);

		my $shutterSpeed = getShutter ($exifData);
		my $aperture = getAperture ($exifData);
		my $isoSpeed;
		$isoSpeed = getFilmSpeed ($exifData) or $isoSpeed = 100;
		$VERBOSE && printf "%s: %.4fs, F%.1f, ISO%d\n",
				$imgs[$i], $shutterSpeed, $aperture, $isoSpeed;

		my $newExif = new Image::ExifTool;
		$newExif->SetNewValue(ISO => $isoSpeed);
		$newExif->SetNewValue(FNumber => $aperture);
		$newExif->SetNewValue(ExposureTime => $shutterSpeed);
		$newExif->WriteInfo("$DIR/$imgs[$i]");
	}
}

MAIN:
{
	my @jpegs = init();

	my @images = $NOORDER ? @jpegs : sortByExposure( @jpegs );

	# Only fix the exposure info, don't register images
	if ($EXIF && !$ALIGN && !$REDO) {
		copyFiles (@images);
		fixEXIF (@images);
		exit;
	}

	my $prefixImgOut;

	my @xShift;
	my @yShift;
	my @rRotation;

	if ($ALIGN eq "ale" || $REDO eq "ale") {
		$ALIGN && runALE( @images );
		$prefixImgOut="ale-";
		
		my @xShiftRel;
		my @yShiftRel;
		my @RelRotate;
		for ( my $i=0 ; $i<$#images ; $i++ ) {
			open(TRANS_FH, "<$transFile$i") or
				die ("Can\'t open $transFile$i: $!\n");
			# Read the x and y offset from the ALE trans file
			while (<TRANS_FH>) {
				next unless m/^E/;
				# E x-dim y-dim x-shift y-shift rotation
				m/^E(\s-?\d*\.\d*){2}\s(-?\d*\.\d*)\s(-?\d*\.\d*)\s(-?\d*\.\d*)/;
				push @xShiftRel, sprintf "%.3f", $2;
				push @yShiftRel, sprintf "%.3f", $3;
				push @RelRotate, sprintf "%.3f", $4;
			}
			close(TRANS_FH);
		}
		$VERBOSE && print "xShiftRelative: " . join(', ', @xShiftRel) . "\n";
		$VERBOSE && print "yShiftRelative: " . join(', ', @yShiftRel) . "\n";
		$VERBOSE && print "Rotat_Relative: " . join(', ', @RelRotate) . "\n";
	
		# The translation is done with reference to the first image.
		# x,y shift for image $i is sum of x,y shift for all previous images.
		# same for rotation, i.e. rotation for image $i is sum of rotation for all previous images.
		$xShift[0]=0;
		$yShift[0]=0;
		$rRotation[0]=0;
		for ( my $i=1 ; $i<$#images+1 ; $i++ ) {
			$xShift[$i] = 0;
			for ( my $j=0 ; $j<$i ; $j++ ) {
				$xShift[$i] += $xShiftRel[$j];
			}
# 			$xShift[$i] = sprintf "%+d", $xShift[$i];
	
			$yShift[$i] = 0;
			for ( my $j=0 ; $j<$i ; $j++ ) {
				$yShift[$i] += $yShiftRel[$j];
			}
# 			$yShift[$i] = sprintf "%+d", $yShift[$i];
	
			$rRotation[$i] = 0;
			for ( my $j=0 ; $j<$i ; $j++ ) {
				$rRotation[$i] += $RelRotate[$j];
			}
# 			$rRotation[$i] = sprintf "%+d", $rRotation[$i];
		}
	}
	elsif ($ALIGN eq "sift" || $REDO eq "sift") {
		$ALIGN && runAUTOPANOCOMPLETE( @images );
		$ALIGN && changeHorizV();
		$ALIGN && runAutoOptimiser();

		$prefixImgOut="sift-";
		my $hfov=100; # the value we modified
		my $rotation;
		my $pitch;
		my $yaw;
		my $width;
		my $height;
	
		my @array;
		tie @array, 'Tie::File', "$panofile" or die "$panofile not found, exiting...";
		for (@array) {
			if(m/^i\ w(\d+)\ h(\d+).*p(-?\d+\.?\d*)\ r(-?\d+\.?\d*).*y(-?\d+\.?\d*)/) {
				$width=$1;
				$height=$2;
				$pitch=$3;
				$rotation=$4;
				$yaw=$5;
				print "sift found: width $width height $height pitch $pitch yaw $yaw rotation $rotation \n";
				push @rRotation, $rotation;
# 				push @xShift, sprintf "%+d", ($yaw*$width/100.0) ;
# 				push @yShift, sprintf "%+d", (-$pitch*$width/100.0) ;
				push @xShift, ($yaw*$width/100.0) ;
				push @yShift, (-$pitch*$width/100.0) ;
			}
		}
		untie @array;
	}

	$VERBOSE && print "xShift        : " . join(', ', @xShift) . "\n";
	$VERBOSE && print "yShift        : " . join(', ', @yShift) . "\n";
	$VERBOSE && print "Rotation      : " . join(', ', @rRotation) . "\n";

	# Get width and height of the original images
	my $img = Image::Magick->new(magick=>'JPEG');
	my $err = $img->Read("$images[0]");
	my($height, $width) = $img->Get('rows', 'columns');
	$VERBOSE && print "Old dimensions: $width x $height\n";

	# finding crop values
	my $numimages=$#images+1;
	print "calling... ./find_offset.m $numimages $width $height @xShift @yShift @rRotation \n";
	my $crop_values=`./find_offset.m $numimages $width $height @xShift @yShift @rRotation `;
	$VERBOSE && print "./find_offset.m returned: $crop_values\n";
	my @script_output=($crop_values =~  m/(-?\d+\.?\d*\ )/g);
	my $crop_width=$script_output[0];
	my $crop_height=$script_output[1];

	my $ale;
	my $cos_alpha;
	my $sin_alpha;
	# Transform images
	for (my $i=0 ; $i<=$#images ; $i++) {
		$img = Image::Magick->new(magick=>'JPEG');
		$err = $img->Read("$images[$i]");
		$ale = $img->Copy;
		$cos_alpha=cos($rRotation[$i]*pi/180);
		$sin_alpha=sin($rRotation[$i]*pi/180);
		if ($rRotation[$i]<0) {
			$ale ->AffineTransform([$cos_alpha, $sin_alpha, -$sin_alpha, $cos_alpha, 0, 2*$width*(-$sin_alpha), 1]);
		}
		else {
			$ale ->AffineTransform([$cos_alpha, $sin_alpha, -$sin_alpha, $cos_alpha, 2*$height*$sin_alpha, 0, 1]);
		}
		$ale ->Crop(width => $crop_width, height => $crop_height,
				x => $script_output[2+2*$i], y => $script_output[3+2*$i]);
		
		$ale ->Set(quality => $QUAL);
		$ale ->Write("$DIR/$prefixImgOut$images[$i]");
	}

	$EXIF && fixEXIF( @images );
} # MAIN

#EOF