/*
 * Content-adaptive Lenticular Prints
 * Method (c) Disney Research 2013. 
 * Patent pending.
 * 
 * No commercial licence is given whatsoever; all inquiries for commercial use 
 * should be directed to Disney Research. 
 * This code is intended for educational use only.
 * 
 * 
 * Implementation (c) James Tompkin 2013
 * Developed independently of Disney Research and derived from the published material:
 * 'Content-adaptive Lenticular Prints' @ SIGGRAPH 2013, itself (c) ACM.
 * http://dl.acm.org/citation.cfm?doid=2461912.2462011
 * 
 * 
 * Please see the project webpage for more details:
 * http://vecg.cs.ucl.ac.uk/Projects/ContentAdaptiveLenticularPrints/
 * 
 * Or contact James personally: james.tompkin@gmail.com
 * 
 * 
 * As this is a personal reimplementation, this is a subset of the features of the 
 * code implemented for Disney Research. At the moment, the code computes an adaptive
 * light field sampling solve (method 'solve', one per scanline/EPI image), but this 
 * code doesn't yet optimize individual lenses or output STL geometry. Time permitting, 
 * my goal is to implement these two features, and e-mails bugging me to do it will 
 * be well-received.
 * 
 * All feedback is welcome!
 */


/*
 * Usage:
 * 1.  Set up your desired lenses and display (currently hard-coded in the main method).
 * 2.  Run the main method in the ContentAdaptiveLenticularPrints_20130805 class.
 * 3.  A dialog box will appear; select a folder containing a 1D light field stored as 
 *     a series of JPEG/PNG/BMP images.
 *     
 *     Wait a few minutes while it processes...
 *     
 * 4a. Adaptive result will be held in the DPEntry[] object 's' (one per scanline).
 * 4b. Adaptive image output will appear in a 'results' folder in the directory of execution.
 * 
 */
package main;

import java.awt.Color;
import java.awt.Graphics2D;
import java.awt.RenderingHints;
import java.awt.Transparency;
import java.awt.image.BufferedImage;
import java.awt.image.DataBufferInt;
import java.io.File;
import java.io.FileFilter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;

import javax.imageio.ImageIO;
import javax.swing.JFileChooser;
import javax.swing.UIManager;
import javax.swing.UnsupportedLookAndFeelException;

public class ContentAdaptiveLenticularPrints_20130805
{
    public static void setup()
    {
	// Set UI to match system L&F
	try
	{
	    UIManager.setLookAndFeel(UIManager.getSystemLookAndFeelClassName());
	} catch (ClassNotFoundException | InstantiationException
		| IllegalAccessException | UnsupportedLookAndFeelException e1)
	{
	    e1.printStackTrace();
	}
    }

    public static BufferedImage generateMeanSlice( LightField lf, DPEntry s )
    {
	return generateMeanSlice( lf.usSlices[0].getWidth(), lf.usSlices[0].getHeight(), lf.usSlices[0].getType(), s );
    }
    
    public static BufferedImage generateMeanSlice( int width, int height, int type, DPEntry s)
    {
	// Create output slice
	BufferedImage ret = new BufferedImage( width, height, type );
	
	int[] a = ((DataBufferInt) ret.getRaster().getDataBuffer()).getData();
	int scansize = ret.getWidth();

	// Cycle through lens applications and set colours in output slice
	int xPos = 0;
	// Per lens
	for (int j = 0; j < s.config.size(); j++)
	{
	    LensApplication la = s.config.get(j);
	    Lens l = la.lens;

	    // Per pixel in lens
	    for (int k = 0; k < la.lens.pixelWidth; k++)
	    {
		int yPos = Math.round(k * l.uHeight);
		int h = Math.round(l.uHeight);
		// Check bounds
		if (yPos + h > ret.getHeight())
		    h = (ret.getHeight() - yPos) - 1;

		// Per pixel in output image within extent of lens pixel
		for (int m = yPos; m < yPos + h; m++)
		{
		    for (int n = xPos; n < xPos + l.pixelWidth; n++)
		    {
			a[m * scansize + n] = la.meanPixels[k].getRGB();
		    }
		}
	    }

	    xPos += la.lens.pixelWidth;
	}

	return ret;
    }

    public static LightField generateImageResult(LightField lf, DPEntry[] s)
    {
	BufferedImage[] meanSlices = new BufferedImage[lf.usSlices.length];

	// For each image
	System.out.println("\nGenerating image result...");
	int pprev = 0;
	for (int i = 0; i < meanSlices.length; i++)
	{
	    // Create output slice
	    meanSlices[i] = generateMeanSlice( lf, s[i] );

	    // Display output
	    int pval = Math.round(100 * (i / (float) meanSlices.length));
	    if (pval > pprev)
	    {
		Tools.printProgressBar(pval);
		pprev = pval;
	    }
	}

	return new LightField(null, meanSlices);
    }

    public static LensApplication applyLens(BufferedImage usSlice, Lens l,
	    int position)
    {
	LensApplication la = new LensApplication(l);

	int[] a = ((DataBufferInt) usSlice.getRaster().getDataBuffer())
		.getData();
	int scansize = usSlice.getWidth();

	// For each pixel in the lens
	for (int i = 0; i < l.pixelWidth; i++)
	{
	    // Each pixel has a rectangular extent in usSlice.
	    // This rectangle is of constant area.
	    //
	    // Imagine we place a lens of width 1 at position.
	    // Its area will be the 'height' of usSlice - the u extent -
	    // as this pixel will need to cover all angular samples at this
	    // point.
	    // Thus, it spans 1 pixel in x and height pixels in y in usSlice.
	    //
	    // Conversely, imagine we place a lens of width 10.
	    // Each of the pixels has constant area in the output.
	    // Now, to fit 10 pixels in the same height, we must expand
	    // the extent in 'x'. In this case, each pixel spans 10 pixels in x
	    // and height/10 pixels in y in usSlice.

	    // So, let's define the region of interest...
	    //
	    // The area of each pixel in the lens has a float extent, and here
	    // we need to choose some means to round.
	    // TODO Ideally, we would upsample usSlice and subpixel sample over
	    // the extent, depending on the smallest fp decimal needed.
	    int yPos = Math.round(i * l.uHeight);
	    int h = Math.round(l.uHeight);
	    // Check bounds
	    if (yPos + h > usSlice.getHeight())
		h = (usSlice.getHeight() - yPos) - 1;

	    int area = l.pixelWidth * h;

	    // Sample from usSlice
	    //
	    // Now, let's compute the mean of this array
	    // array format is ARGB
	    float meanR = 0, meanG = 0, meanB = 0;
	    for (int j = yPos; j < yPos + h; j++)
	    {
		for (int k = position; k < position + l.pixelWidth; k++)
		{
		    Color rgb = new Color(a[j * scansize + k]);
		    meanR += rgb.getRed();
		    meanG += rgb.getGreen();
		    meanB += rgb.getBlue();
		}
	    }

	    meanR /= area;
	    meanG /= area;
	    meanB /= area;
	    try
	    {
		la.meanPixels[i] = new Color((int) meanR, (int) meanG,
			(int) meanB);
	    } catch (IllegalArgumentException iae)
	    {
		System.out.println((int) meanR + " " + (int) meanG + " "
			+ (int) meanB);
	    }

	    // Now, let's compare the input colour to this mean,
	    // pixel-by-pixel, and accumulate an error.
	    la.pixelErrors[i] = 0;
	    for (int j = yPos; j < yPos + h; j++)
	    {
		for (int k = position; k < position + l.pixelWidth; k++)
		{
		    Color rgb = new Color(a[j * scansize + k]);
		    float distR = rgb.getRed() - meanR;
		    float distG = rgb.getGreen() - meanG;
		    float distB = rgb.getBlue() - meanB;

		    la.pixelErrors[i] += distR * distR + distB * distB + distG
			    * distG;
		}
	    }

	    // Normalize error by area
	    //
	    // Each pixel 'technically' has the same area, but with
	    // rounding errors in sampling this can lead to large
	    // differences in total error.
	    // Thus, we normalize by the actual sampled error to
	    // remove this factor.
	    la.pixelErrors[i] /= area;
	}

	// Sum error over all pixels
	la.error = 0;
	for (int i = 0; i < la.pixelErrors.length; i++)
	    la.error += la.pixelErrors[i];

	return la;
    }

    public static DPEntry solve(Display d, BufferedImage usSlice, DPEntry[] dpt, int subset, int depth)
    {
	// For each lens
	for (int i = 0; i < d.lenses.length; i++)
	{
	    Lens l = d.lenses[i];

	    // Place the lens at the right-hand side of this subset of the
	    // usSlice
	    int remainingWidth = subset - l.pixelWidth;

	    DPEntry result = null;
	    if (remainingWidth < 0)
	    {
		// Error; this lens cannot be placed here.
		continue;
	    } 
	    else if (remainingWidth == 0)
	    {
		// Base case
		// Compute the error directly
		LensApplication la = applyLens(usSlice, l, remainingWidth);
		result = new DPEntry(la.error, new ArrayList<LensApplication>(
			Arrays.asList(la)));
	    } 
	    else if (remainingWidth > 0)
	    {
		// Look up in table
		DPEntry subsetResult = null;
		if (dpt[remainingWidth] == null)
		    // Recurse on the subset
		    subsetResult = solve(d, usSlice, dpt, remainingWidth,
			    depth + 1);
		else
		    subsetResult = dpt[remainingWidth];

		// subsetResult + this lens applied
		LensApplication la = applyLens(usSlice, l, remainingWidth);

		// Combine into new result
		ArrayList<LensApplication> combined = new ArrayList<LensApplication>(
			subsetResult.config);
		combined.add(la);
		result = new DPEntry(subsetResult.error + la.error, combined);
	    }

	    // Compare result to existing entry in the table for this subset.
	    // Is there an entry for this width? Is the error in the table
	    // greater?
	    if (dpt[subset] == null || dpt[subset].error > result.error)
	    {
		// Add error to table
		dpt[subset] = new DPEntry(result.error, result.config);
	    }
	}

	return dpt[subset];
    }

    // TODO Parallelize per scanline
    public static DPEntry[] optimizeDisplay(LightField lf, Display d)
    {
	// Set lens u extents
	for (int i = 0; i < d.lenses.length; i++)
	    d.lenses[i].uHeight = lf.usSlices[0].getHeight() / (float) d.lenses[i].pixelWidth;

	DPEntry[] perScanline = new DPEntry[lf.usSlices.length];

	// For each u,s slice in the LightField
	System.out.println("\nSolving scanlines...");
	int pprev = 0;
	for (int i = 0; i < perScanline.length; i++)
	{
	    // The width is an index into the arrays.
	    // So, the error and config for a LF width of '1'
	    // is at entries[1].
	    DPEntry[] dpt = new DPEntry[lf.usSlices[i].getWidth() + 1];
	    
	    int subset = lf.usSlices[i].getWidth();
	    perScanline[i] = solve(d, lf.usSlices[i], dpt, subset, 0);

	    // Display output
	    int pval = Math.round(100 * (i / (float) perScanline.length));
	    if (pval > pprev)
	    {
		Tools.printProgressBar(pval);
		pprev = pval;
	    }
	}

	return perScanline;
    }

    public static void main(String[] args)
    {
	setup();

	// Steps
	// 0. Define display
	// 1. Load light field
	// 2. Resample light field given display properties
	// 3. Slice light field in u,s space
	// 4. For each slice, fit lenses using dynamic programming approach
	// 5. Optimize lenses for given widths
	// 6. Create geometry

	// Step 0: Define display
	// ======================
	// Here, pick your output lenses (defined in pixel widths)
	Lens[] ls = new Lens[] { new Lens(1), new Lens(2), new Lens(4), new Lens(5), new Lens(10), new Lens(20) };
	// Also, here pick your pixel size (horiz & vert, first two terms) and display output size (second two terms)
	// All in mm.
	Display d = new Display(1f, 1f, 800, 400, ls);

	// Steps 1,2,3:
	// ============
	// 1. Load light field
	JFileChooser jf = new JFileChooser(System.getProperty("user.home"));
	jf.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY);
	int ret = jf.showOpenDialog(null);

	File selected = null;
	if (ret == JFileChooser.APPROVE_OPTION)
	{
	    selected = jf.getSelectedFile();
	} else
	{
	    System.exit(0);
	}

	// 2. Resample light field given display properties
	// 3. Slice light field in u,s space
	//
	// Both 2 & 3 are completed at load.
	LightField lf = LightField.load(selected, d, false);

	// Step 4: Fit lenses using DP approach
	// ====================================
	DPEntry[] s = optimizeDisplay(lf, d);

	// Step 4a: Write output LF as images
	LightField lfR = generateImageResult(lf, s);
	LightField.save(lfR, new File("./result"));

	// TODO Step 5: Optimize fixed lenses
	// =============================
	// for( int i = 0; i < d.lenses.length; i++ )
	// 	optimizeLens( d.lenses[i] );

	// TODO Step 6: Create geometry
	// =======================

    }
}

class DPEntry
{
    // The dynamic programming table stores, for each
    // width of the display (in widths), an error
    // and a corresponding lens configuration.
    //
    float error;
    ArrayList<LensApplication> config;

    public DPEntry(float error, ArrayList<LensApplication> config)
    {
	this.error = error;
	this.config = config;
    }
}

class Display
{
    float pixelWidthMM, pixelHeightMM;
    float displayWidthMM, displayHeightMM;
    int nPixelsWidth, nPixelsHeight;
    float fov;

    Lens[] lenses;

    public Display(float pw, float ph, float dw, float dh, Lens[] ls)
    {
	pixelWidthMM = pw;
	pixelHeightMM = ph;
	displayWidthMM = dw;
	displayHeightMM = dh;

	nPixelsWidth = Math.round(dw / pw);
	nPixelsHeight = Math.round(dh / ph);

	lenses = ls;
    }

    public Lens getLargestLens()
    {
	Lens largest = lenses[0];
	for (int i = 1; i < lenses.length; i++)
	{
	    if (lenses[i].pixelWidth > largest.pixelWidth)
		largest = lenses[i];
	}
	return largest;
    }
}

class Lens
{
    int pixelWidth;

    // Extent of this lens in u space within the usSlice.
    // Specific to a particular light field! Assigned later
    // in optimizeDisplay once the LF is input.
    float uHeight;

    public Lens(int pw)
    {
	pixelWidth = pw;
    }

    public String toString()
    {
	return "Lens[pixelWidth:" + pixelWidth + ",uHeight:" + uHeight + "]";
    }
}

class LensApplication
{
    Lens lens;
    float error;
    float[] pixelErrors;
    Color[] meanPixels;

    public LensApplication(Lens l)
    {
	this.lens = l;
	meanPixels = new Color[l.pixelWidth];
	pixelErrors = new float[l.pixelWidth];
    }
}

class LightField
{
    // In our case, we represent a light field as a 4D structure, captured
    // by a series of images.
    File directory;

    // However, most of the operations that we will perform will be on EPI slices
    // of this 4D structure in u,s space (each row of the image across angles).
    // As such, we're going to store a separate structure in this space.
    BufferedImage[] usSlices;

    public LightField(File dir, BufferedImage[] usSlices)
    {
	this.directory = dir;
	this.usSlices = usSlices;
    }

    // TODO Load directly into EPI slices instead of current two-stage process.
    public static LightField load(File directory, Display d,
	    boolean resampleAngular)
    {
	File[] imageFiles = directory.listFiles(new ImageFileFilter());

	// Storage for resampled images
	BufferedImage[] resampled = new BufferedImage[imageFiles.length];

	System.out.print("Loading images:");
	int pprev = 0;
	for (int i = 0; i < imageFiles.length; i++)
	{
	    BufferedImage b = null;
	    try
	    {
		b = ImageIO.read(imageFiles[i]);
	    } catch (IOException e)
	    {
		System.err.println("Error reading file: " + imageFiles[i]);
		System.exit(0);
	    }

	    // Resample to fit d if d had no lenses and just showed a single
	    // angle from the light field
	    float resampleX = b.getWidth() / (float) d.nPixelsWidth;
	    float resampleY = b.getHeight() / (float) d.nPixelsHeight;
	    int newWidth = Math.round(b.getWidth() * resampleX);
	    int newHeight = Math.round(b.getHeight() * resampleY);

	    resampled[i] = Tools.getScaledInstance(b, newWidth, newHeight,
		    RenderingHints.VALUE_INTERPOLATION_BILINEAR, true);

	    int pval = Math.round(100 * (i / (float) imageFiles.length));
	    if (pval > pprev)
	    {
		Tools.printProgressBar(pval);
		pprev = pval;
	    }
	}

	// Form slice structure
	BufferedImage[] usSlices = new BufferedImage[resampled[0].getHeight()];

	// Slice images into u,s scanlines
	System.out.print("\nSlicing images:");
	pprev = 0;
	// i is scanline
	for (int i = 0; i < resampled[0].getHeight(); i++)
	{
	    // Create slice image
	    int w = resampled[0].getWidth();
	    usSlices[i] = new BufferedImage(w, resampled.length,
		    resampled[0].getType());

	    // Take a slice from each input image
	    // j is angle
	    for (int j = 0; j < resampled.length; j++)
	    {
		// Grab slice from resampled
		int[] array = resampled[j].getRGB(0, i, w, 1, null, 0, w);
		usSlices[i].setRGB(0, j, w, 1, array, 0, w);
	    }

	    // If desired, we can angularly resample here
	    // to fit the maximum number of pixels under
	    // the largest possible display lens
	    if (resampleAngular)
	    {
		usSlices[i] = Tools.getScaledInstance(usSlices[i], w,
			d.getLargestLens().pixelWidth,
			RenderingHints.VALUE_INTERPOLATION_BILINEAR, true);
	    }

	    // // DEBUG
	    // try
	    // {
	    // ImageIO.write( usSlices[i], "PNG", new File( "slice_" +
	    // Tools.zeroPad(i,4) + ".png") );
	    // }
	    // catch (IOException e)
	    // {
	    // e.printStackTrace();
	    // }
	    int pval = Math.round(100 * (i / (float) resampled[0].getHeight()));
	    if (pval > pprev)
	    {
		Tools.printProgressBar(pval);
		pprev = pval;
	    }
	}

	// Form lf
	return new LightField(directory, usSlices);
    }

    public static void save(LightField lfR, File dir)
    {
	dir.mkdirs();

	// Convert usSlices in lfR into x,y slices and write out
	BufferedImage us = lfR.usSlices[0];
	int w = us.getWidth();

	System.out.println("\nWriting output images...");
	int pprev = 0;
	for (int i = 0; i < us.getHeight(); i++)
	{
	    BufferedImage output = new BufferedImage(w, lfR.usSlices.length,
		    us.getType());

	    // Cycle through each output image and copy across
	    for (int j = 0; j < lfR.usSlices.length; j++)
	    {
		// Get horizontal slice
		int[] array = lfR.usSlices[j].getRGB(0, i, w, 1, null, 0, w);
		output.setRGB(0, j, w, 1, array, 0, w);
	    }

	    // Output image
	    try
	    {
		ImageIO.write(output, "PNG", new File(dir + File.separator
			+ "output_" + Tools.zeroPad(i, 4) + ".png"));
	    } catch (IOException e)
	    {
		e.printStackTrace();
	    }

	    int pval = Math.round(100 * (i / (float) us.getHeight()));
	    if (pval > pprev)
	    {
		Tools.printProgressBar(pval);
		pprev = pval;
	    }
	}
    }
}


class Tools
{
    /**
     * From http://today.java.net/pub/a/today/2007/04/03/perils-of-image-
     * getscaledinstance.html , with modifications.
     * 
     * Convenience method that returns a scaled instance of the provided
     * {@code BufferedImage}.
     * 
     * @param img
     *            the original image to be scaled
     * @param targetWidth
     *            the desired width of the scaled instance, in pixels. Set to -1
     *            to auto-compute based on aspect ratio.
     * @param targetHeight
     *            the desired height of the scaled instance, in pixels. Set to
     *            -1 to auto-compute based on aspect ratio.
     * @param hint
     *            one of the rendering hints that corresponds to
     *            {@code RenderingHints.KEY_INTERPOLATION} (e.g.
     *            {@code RenderingHints.VALUE_INTERPOLATION_NEAREST_NEIGHBOR},
     *            {@code RenderingHints.VALUE_INTERPOLATION_BILINEAR},
     *            {@code RenderingHints.VALUE_INTERPOLATION_BICUBIC})
     * @param higherQuality
     *            if true, this method will use a multi-step scaling technique
     *            that provides higher quality than the usual one-step technique
     *            (only useful in downscaling cases, where {@code targetWidth}
     *            or {@code targetHeight} is smaller than the original
     *            dimensions, and generally only when the {@code BILINEAR} hint
     *            is specified)
     * @return a scaled version of the original {@code BufferedImage}
     */
    public static BufferedImage getScaledInstance(BufferedImage img,
	    int targetWidth, int targetHeight, Object hint,
	    boolean higherQuality)
    {
	int type = (img.getTransparency() == Transparency.OPAQUE) ? BufferedImage.TYPE_INT_RGB
		: BufferedImage.TYPE_INT_ARGB;
	BufferedImage ret = (BufferedImage) img;
	int w, h;
	if (targetHeight == -1 && targetWidth == -1)
	    return ret;
	else if (targetWidth == -1)
	{
	    // Compute ratio
	    float ratio = ret.getHeight() / targetHeight;
	    targetWidth = (int) (ret.getWidth() / ratio);
	} else if (targetHeight == -1)
	{
	    // Compute ratio
	    float ratio = ret.getWidth() / targetWidth;
	    targetHeight = (int) (ret.getHeight() / ratio);
	}

	if (higherQuality)
	{
	    // Use multi-step technique: start with original size, then
	    // scale down in multiple passes with drawImage()
	    // until the target size is reached
	    w = img.getWidth();
	    h = img.getHeight();
	} else
	{
	    // Use one-step technique: scale directly from original
	    // size to target size with a single drawImage() call
	    w = targetWidth;
	    h = targetHeight;
	}

	do
	{
	    if (higherQuality && w > targetWidth)
	    {
		w /= 2;
		if (w < targetWidth)
		{
		    w = targetWidth;
		}
	    } else
		// higherQuality and we are enlarging; do nothing
		w = targetWidth;

	    if (higherQuality && h > targetHeight)
	    {
		h /= 2;
		if (h < targetHeight)
		{
		    h = targetHeight;
		}
	    } else
		// higherQuality and we are enlarging; do nothing
		h = targetHeight;

	    BufferedImage tmp = new BufferedImage(w, h, type);
	    Graphics2D g2 = tmp.createGraphics();
	    g2.setRenderingHint(RenderingHints.KEY_INTERPOLATION, hint);
	    g2.drawImage(ret, 0, 0, w, h, null);
	    g2.dispose();

	    ret = tmp;
	} while (w != targetWidth || h != targetHeight);

	return ret;
    }

    /**
     * From http://nakkaya.com/2009/11/08/command-line-progress-bar/
     * Doesn't work in Eclipse output : )
     * 
     * @param percent
     */
    public static void printProgressBar(int percent)
    {
	StringBuilder bar = new StringBuilder("[");

	for (int i = 0; i < 50; i++)
	{
	    if (i < (percent / 2))
	    {
		bar.append("=");
	    } else if (i == (percent / 2))
	    {
		bar.append(">");
	    } else
	    {
		bar.append(" ");
	    }
	}

	bar.append("]   " + percent + "%     ");
	System.out.print("\r" + bar.toString());
    }

    public static String zeroPad(int n, int width)
    {
	String num = String.valueOf(n);
	while (num.length() < width)
	    num = "0" + num;
	return num;
    }
}

class ImageFileFilter implements FileFilter
{
    @Override
    public boolean accept(File file)
    {
	String name = file.getName();
	String suffix = name.substring( name.lastIndexOf(".")+1 ).toLowerCase();
	
	// Whatever file types you want that can be loaded by ImageIO...
	return suffix.equals( "png" ) || suffix.equals( "jpg" ) 
		|| suffix.equals( "jpeg" ) || suffix.equals( "bmp" );
    }
    
}