///////////////////////////////////////////////////////////////////////////
//
// NAME
//  BlendImages.cpp -- blend together a set of overlapping images
//
// DESCRIPTION
//  This routine takes a collection of images aligned more or less horizontally
//  and stitches together a mosaic.
//
//  The images can be blended together any way you like, but I would recommend
//  using a soft halfway blend of the kind Steve presented in the first lecture.
//
//  Once you have blended the images together, you should crop the resulting
//  mosaic at the halfway points of the first and last image.  You should also
//  take out any accumulated vertical drift using an affine warp.
//  Lucas-Kanade Taylor series expansion of the registration error.
//
// SEE ALSO
//  BlendImages.h       longer description of parameters
//
// Copyright ?Richard Szeliski, 2001.  See Copyright.h for more details
// (modified for CSE455 Winter 2003)
//
///////////////////////////////////////////////////////////////////////////

#include "../ImageLib/ImageLib.h"
#include "BlendImages.h"
#include <math.h>

/******************* TO DO 4 *********************
 * AccumulateBlend:
 *	INPUT:
 *		img: a new image to be added to acc
 *		acc: portion of the accumulated image where img is to be added
 *		blendWidth: width of the blending function (horizontal hat function;
 *	    try other blending functions for extra credit)
 *	OUTPUT:
 *		add a weighted copy of img to the subimage specified in acc
 *		the first 3 band of acc records the weighted sum of pixel colors
 *		the fourth band of acc records the sum of weight
 */
static void AccumulateBlend(CByteImage& img, CFloatImage& acc, float blendWidth)
{
	// *** BEGIN TODO ***
	// fill in this routine..
	CShape sh = img.Shape();
	CShape sh1 = acc.Shape();
	int x, y, k;
	float slope = 1 / (blendWidth/2 - 1);
	float alpha;
	for (y = 0; y < sh.height; y ++)
	{
		// Piecewise horizontal hat function
		// from 0 to (blendWidth/2 - 1): ascending
		for (x = 0; x < blendWidth/2; x ++)
		{
			alpha = x * slope;
			for (k = 0; k < 4; k ++)
			{
				acc.Pixel(x, y, k) += alpha * img.Pixel(x, y, k);
			}
		}
		// from (blendWidth/2) to (sh.width - blendWidth/2 - 1): flat
		for (x = blendWidth/2; x < sh.width - blendWidth/2; x ++)
		{
			for (k = 0; k < 4; k ++)
			{
				acc.Pixel(x, y, k) += img.Pixel(x, y, k);
			}
		}
		// from (sh.width - blendWidth/2) to (sh.width - 1): descending
		for (x = sh.width - blendWidth/2; x < sh.width; x ++)
		{
			alpha = (sh.width - x - 1) * slope;
			for (k = 0; k < 4; k ++)
			{
				acc.Pixel(x, y, k) += alpha * img.Pixel(x, y, k);
			}
		}
	}
	// *** END TODO ***
}

/******************* TO DO 5 *********************
 * NormalizeBlend:
 *	INPUT:
 *		acc: input image whose alpha channel (4th channel) contains
 *		     normalizing weight values
 *		img: where output image will be stored
 *	OUTPUT:
 *		normalize r,g,b values (first 3 channels) of acc and store it into img
 */
static void NormalizeBlend(CFloatImage& acc, CByteImage& img)
{
	// *** BEGIN TODO ***
	// fill in this routine..
	CShape sh = img.Shape();
	int x, y, k;
	for (y = 0; y < sh.height; y ++)
	{
		for (x = 0; x < sh.width; x ++)
		{
			for (k = 0; k < 4; k ++)
			{
				// Normalize to between 0 and 255
				img.Pixel(x, y, k) = (char)(255 * acc.Pixel(x, y, k) / acc.Pixel(x, y, 3));
			}
		}
	}	
	// *** END TODO ***
}

/******************* TO DO 6 *********************
 * BlendImages:
 *	INPUT:
 *		ipv: list of input images and their relative positions in the mosaic
 *		blendWidth: width of the blending function
 *	OUTPUT:
 *		create & return final mosaic by blending all images
 *		and correcting for any vertical drift
 */
CByteImage BlendImages(CImagePositionV& ipv, float blendWidth)
{
    // Assume all the images are of the same shape (for now)
    CByteImage& img0 = ipv[0].img;
    CShape sh        = img0.Shape();
    int width        = sh.width;
    int height       = sh.height;
    int nBands       = sh.nBands;
    int dim[2]       = {width, height};

    // Compute the bounding box for the mosaic
    int n = ipv.size();
    float min_x = 0, min_y = 0;
    float max_x = 0, max_y = 0;
    int i;
    for (i = 0; i < n; i++)
    {
        float* pos = ipv[i].position;
        
		// *** BEGIN TODO #1 ***
		// add some code here to update min_x, ..., max_y
		// Find the limits of x and y
		min_x = pos[0] < min_x ? pos[0] : min_x;
		max_x = pos[0] + width > max_x ? pos[0] + width : max_x;
		min_y = pos[1] < min_y ? pos[1] : min_y; 
		max_y = pos[1] + height > max_y ? pos[1] + height : max_y;
		// *** END TODO #1 ***
    }

    // Create a floating point accumulation image
    CShape mShape((int)(ceil(max_x) - floor(min_x)),
                  (int)(ceil(max_y) - floor(min_y)), nBands);
    CFloatImage accumulator(mShape);
    accumulator.ClearPixels();

    // Add in all of the images
    for (i = 0; i < n; i++)
    {
        // Compute the sub-image involved
        float* pos = ipv[i].position;
        int start_x = 0, start_y = 0;       

		// *** BEGIN TODO #2 ***
		// fill in with the correct values for start_x and start_y,
		// the relative position of i-th image in the mosaic
		// hint: we only require you to use pixel accuracy to start_x and start_y;
		// Pixel accuracy:
		start_x = (int)(pos[0] - min_x);
		start_y = (int)(pos[1] - min_y);
		// *** END TODO #2 ***

        CByteImage& img = ipv[i].img;
        CFloatImage sub = accumulator.SubImage(start_x, start_y, width, height);

        // Perform the accumulation
        AccumulateBlend(img, sub, blendWidth);
    }

    // Normalize the results
    CByteImage compImage(mShape);
    NormalizeBlend(accumulator, compImage);
    bool debug_comp = false;
    if (debug_comp)
        WriteFile(compImage, "tmp_comp.tga");

    // Allocate the final image shape
    CShape cShape(mShape.width - width, height, nBands);
    CByteImage croppedImage(cShape);

    // Compute the affine deformation
    CTransform3x3 A;
    
	// *** BEGIN TODO #3 ***
    // fill in the right entries in A to trim the left edge and
    // to take out the vertical drift
	float* pos_start = ipv[0].position;
	float* pos_end = ipv[n-1].position;
	/* *** */
	// Corrected here: y-shift should be pos_end[1] - min_y instead of pos_end[1]
	A = CTransform3x3::Translation(0.5*width, pos_end[1] - min_y);
	/* *** */
	A[1][0] = (pos_start[1] - pos_end[1]) / (pos_start[0] - pos_end[0]);
	// *** END TODO #3 ***

    // Warp and crop the composite
    WarpGlobal(compImage, croppedImage, A, eWarpInterpLinear);

    return croppedImage;
}

/******************* Extra -- Subpixel accuracy *********************
 * BlendImages_Subpixel:
 *	INPUT:
 *		ipv: list of input images and their relative positions in the mosaic
 *		blendWidth: width of the blending function
 *	OUTPUT:
 *		create & return final mosaic by blending all images
 *		and correcting for any vertical drift
 */
CByteImage BlendImages_Subpixel(CImagePositionV& ipv, float blendWidth)
{
    // Assume all the images are of the same shape (for now)
    CByteImage& img0 = ipv[0].img;
    CShape sh        = img0.Shape();
    int width        = sh.width;
    int height       = sh.height;
    int nBands       = sh.nBands;
    int dim[2]       = {width, height};

    // Compute the bounding box for the mosaic
    int n = ipv.size();
    float min_x = 0, min_y = 0;
    float max_x = 0, max_y = 0;
    int i;
    for (i = 0; i < n; i++)
    {
        float* pos = ipv[i].position;
        
		// *** BEGIN TODO #1 ***
		// add some code here to update min_x, ..., max_y
		// Find the limits of x and y
		min_x = pos[0] < min_x ? pos[0] : min_x;
		max_x = pos[0] + width > max_x ? pos[0] + width : max_x;
		min_y = pos[1] < min_y ? pos[1] : min_y; 
		max_y = pos[1] + height > max_y ? pos[1] + height : max_y;
		// *** END TODO #1 ***
    }

    // Create a floating point accumulation image
    CShape mShape((int)(ceil(max_x) - floor(min_x)),
                  (int)(ceil(max_y) - floor(min_y)), nBands);
    CFloatImage accumulator(mShape);
    accumulator.ClearPixels();

    // Add in all of the images
    for (i = 0; i < n; i++)
    {
        // Compute the sub-image involved
        float* pos = ipv[i].position;
        int start_x = 0, start_y = 0;       

		// *** BEGIN TODO #2 ***
		// fill in with the correct values for start_x and start_y,
		// the relative position of i-th image in the mosaic
		// Subpixel accuracy: use the translated image to compensate for rounding
		// Transform the image, such that imgt(x,y) = img(x+t[0],y+t[1]);
		start_x = (int)(pos[0] - min_x);
		start_y = (int)(pos[1] - min_y);
		CByteImage& img = ipv[i].img;		// image to be added
		CByteImage imgt(img.Shape());       // the translated image
		float t[2];
		// Subpixel shifts to compensate for rounding
		t[0] = (pos[0] - min_x) - start_x;
		t[1] = (pos[1] - min_y) - start_y;
		CTransform3x3 T = CTransform3x3::Translation(t[0], t[1]);
		WarpGlobal(img, imgt, T, eWarpInterpLinear);	

        CFloatImage sub = accumulator.SubImage(start_x, start_y, width, height);

        // Perform the accumulation, use translated image instead
        AccumulateBlend(imgt, sub, blendWidth);

		// *** END TODO #2 ***

    }

    // Normalize the results
    CByteImage compImage(mShape);
    NormalizeBlend(accumulator, compImage);
    bool debug_comp = false;
    if (debug_comp)
        WriteFile(compImage, "tmp_comp.tga");

    // Allocate the final image shape
    CShape cShape(mShape.width - width, height, nBands);
    CByteImage croppedImage(cShape);

    // Compute the affine deformation
    CTransform3x3 A;
    
	// *** BEGIN TODO #3 ***
    // fill in the right entries in A to trim the left edge and
    // to take out the vertical drift
	float* pos_start = ipv[0].position;
	float* pos_end = ipv[n-1].position;
	A = CTransform3x3::Translation(0.5*width, pos_end[1]);
	A[1][0] = (pos_start[1] - pos_end[1]) / (pos_start[0] - pos_end[0]);
	// *** END TODO #3 ***

    // Warp and crop the composite
    WarpGlobal(compImage, croppedImage, A, eWarpInterpLinear);

    return croppedImage;
}