/**
 * Copyright (c) 2011, The University of Southampton and the individual contributors.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 *   *  Redistributions of source code must retain the above copyright notice,
 *      this list of conditions and the following disclaimer.
 *
 *   *  Redistributions in binary form must reproduce the above copyright notice,
 *      this list of conditions and the following disclaimer in the documentation
 *      and/or other materials provided with the distribution.
 *
 *   *  Neither the name of the University of Southampton nor the names of its
 *      contributors may be used to endorse or promote products derived from this
 *      software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 * ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
package org.openimaj.image.analysis.watershed;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.List;

import org.openimaj.image.FImage;
import org.openimaj.image.analysis.watershed.event.ComponentStackMergeListener;
import org.openimaj.image.analysis.watershed.feature.ComponentFeature;
import org.openimaj.image.pixel.IntValuePixel;

/**
 * Maximally Stable Extremal Region watershed algorithm, implemented as
 * described in the Microsoft paper of Nister and Stewenius.
 * 
 * @author David Dupplaw (dpd@ecs.soton.ac.uk)
 * @author Jonathon Hare (jsh2@ecs.soton.ac.uk)
 * 
 */
public class WatershedProcessorAlgorithm
{
        /**
         * A sorted heap of pixels. When {@link #pop()} is called the lowest value
         * pixel is returned first.
         * 
         * @author David Dupplaw (dpd@ecs.soton.ac.uk)
         * @author Jonathon Hare (dpd@ecs.soton.ac.uk)
         * 
         */
        private class BoundaryHeap
        {
                private BitSet availablePixels;
                private ArrayDeque<IntValuePixel>[] stacks;

                /**
                 * Construct a boundary heap object with a given number of levels (i.e.
                 * max value of a pixel).
                 * 
                 * @param sz
                 *            number of levels.
                 */
                @SuppressWarnings("unchecked")
                public BoundaryHeap(int sz) {
                        availablePixels = new BitSet(sz);
                        stacks = new ArrayDeque[sz];

                        for (int i = 0; i < sz; i++)
                                stacks[i] = new ArrayDeque<IntValuePixel>();
                }

                /**
                 * Pushes the pixel onto the heap.
                 * 
                 * @param p
                 *            The {@link IntValuePixel} to push onto the heap.
                 */
                public void push(IntValuePixel p)
                {
                        final ArrayDeque<IntValuePixel> l = stacks[p.value];
                        l.push(p);
                        availablePixels.set(p.value);
                }

                /**
                 * Returns the lowest available pixel off the heap (removing it from the
                 * heap). Pixels are returned in sorted order (lowest value first). The
                 * method will return null if the heap is empty.
                 * 
                 * @return The lowest value available pixel or NULL if no pixels are
                 *         available.
                 */
                public IntValuePixel pop()
                {
                        final int l = availablePixels.nextSetBit(0);
                        if (l == -1)
                                return null; // Null means no available pixels (empty heap)

                        final IntValuePixel xx = this.stacks[l].pop();
                        if (this.stacks[l].size() == 0)
                                availablePixels.set(l, false);
                        return xx; // lowest and newest pixel
                }
        }

        /** The pixel where the pour will start */
        private IntValuePixel startPixel = null;

        /** The mask that shows which pixels have been visited */
        private BitSet accessibleMask = null;

        /** The current pixel being investigated */
        private IntValuePixel currentPixel = null;

        /** The stack of components during processing */
        private ArrayDeque<Component> componentStack = null;

        /** The boundary heap during processing */
        private BoundaryHeap boundaryHeap = null;

        /** The image being processed */
        private int[][] greyscaleImage = null;

        /**
         * The listeners for this watershed process. They will be called as regions
         * are detected
         */
        private List<ComponentStackMergeListener> csmListeners = null;

        private Class<? extends ComponentFeature>[] featureClasses;

        /**
         * Default constructor
         * 
         * @param greyscaleImage
         *            the image as a 2d array of integer values
         * @param startPixel
         *            The pixel to start the process at
         * @param featureClasses
         *            the features that should be created for each detected
         *            component
         */
        public WatershedProcessorAlgorithm(int[][] greyscaleImage, IntValuePixel startPixel,
                        Class<? extends ComponentFeature>... featureClasses)
        {
                this.greyscaleImage = greyscaleImage;
                this.startPixel = startPixel;
                this.csmListeners = new ArrayList<ComponentStackMergeListener>();

                this.featureClasses = featureClasses;
        }

        /**
         * Default constructor
         * 
         * @param bGreyscaleImage
         *            the image to apply the watershed transform too
         * @param startPixel
         *            The pixel to start the process at
         * @param featureClasses
         *            the features that should be created for each detected
         *            component
         */
        public WatershedProcessorAlgorithm(FImage bGreyscaleImage, IntValuePixel startPixel,
                        Class<? extends ComponentFeature>... featureClasses)
        {
                this(new int[bGreyscaleImage.getHeight()][bGreyscaleImage.getWidth()], startPixel, featureClasses);

                for (int j = 0; j < bGreyscaleImage.getHeight(); j++) {
                        for (int i = 0; i < bGreyscaleImage.getWidth(); i++) {
                                greyscaleImage[j][i] = (int) (bGreyscaleImage.pixels[j][i] * 255);
                        }
                }
        }

        /**
         * Start the detection process by pouring on water at the pour point. (part
         * 1 and 2)
         * 
         */
        public void startPour()
        {
                // For each step on the downhill stream is created as
                // a component.
                this.currentPixel = startPixel;

                // Store the grey level of the current pixel
                this.currentPixel.value = greyscaleImage[this.startPixel.y][this.startPixel.x];

                // Create the mask the shows where the water has access to
                this.accessibleMask = new BitSet(this.greyscaleImage.length * this.greyscaleImage[0].length);

                // Create the stack of components
                this.componentStack = new ArrayDeque<Component>();

                // Create the heap of boundary pixels
                this.boundaryHeap = new BoundaryHeap(256);

                // Create a dummy component with a greylevel higher than
                // any allowed and push it onto the stack
                final Component dummyComponent = new Component(new IntValuePixel(-1, -1, Integer.MAX_VALUE), featureClasses);
                this.componentStack.push(dummyComponent);

                // Continue the processing at the first pixel
                this.processNeighbours();

                // System.err.println("Component Stack: "+componentStack );
        }

        /**
         * Process the current pixel's neighbours (part 4, 5, 6 and 7).
         */
        private void processNeighbours()
        {
                // Push an empty component with the current level
                // onto the component stack
                Component currentComponent = new Component(this.currentPixel, featureClasses);
                componentStack.push(currentComponent);

                // System.err.println( "Processing neighbours of "+currentPixel );

                final boolean processNeighbours = true;
                while (processNeighbours)
                {
                        boolean toContinue = false;

                        // Get all the neighbours of the current pixel
                        final IntValuePixel[] neighbours = getNeighbourPixels_4(this.currentPixel);
                        // System.err.println("Neighbours: "+outputArray( neighbours ) );

                        // For each of the neighbours, check if the the neighbour
                        // is already accessible.
                        for (final IntValuePixel neighbour : neighbours)
                        {
                                if (neighbour == null)
                                        break; // neighbours array is packed, so nulls only occur at
                                                        // the end

                                final int idx = neighbour.x + neighbour.y * this.greyscaleImage[0].length;

                                // If the neighbour is not accessible...
                                if (!this.accessibleMask.get(idx))
                                {
                                        // Mark it as accessible...
                                        this.accessibleMask.set(idx);
                                        // System.err.println("Making "+neighbour+" accessible" );

                                        // If its greylevel is not lower than the current one...
                                        if (neighbour.value >= currentPixel.value)
                                        {
                                                // Push it onto the heap of boundary pixels
                                                this.boundaryHeap.push(neighbour);
                                                // System.err.println("1. Push "+neighbour+", = "+boundaryHeap
                                                // );
                                        }
                                        // If, on the other hand, the greylevel is lower
                                        // than the current one, enter the current pixel
                                        // back into the queue of boundary pixels for later
                                        // processing (with the next edge number), consider
                                        // the new pixel and process it
                                        // (this is the water pouring into the local minimum)
                                        else
                                        {
                                                this.boundaryHeap.push(currentPixel);
                                                // System.err.println("2. Push "+currentPixel+", = "+boundaryHeap
                                                // );
                                                this.currentPixel = neighbour;
                                                currentComponent = new Component(this.currentPixel, featureClasses);
                                                componentStack.push(currentComponent);
                                                toContinue = true;
                                                break;
                                        }
                                }
                        }

                        if (toContinue)
                                continue;

                        // Accumulate the current pixel to the component at the top of the
                        // stack. (part 5)
                        this.componentStack.peek().accumulate(this.currentPixel);
                        // System.err.println("Added "+currentPixel+" to top component "+componentStack.peek()
                        // );

                        // Pop the heap of boundary pixels. (part 6)
                        final IntValuePixel p = this.boundaryHeap.pop();
                        // System.err.println("Popped "+p+", = "+boundaryHeap );

                        // If the heap is empty, then we're done
                        if (p == null)
                                return;

                        // If it's at the same grey-level we process its neighbours (part 6)
                        if (p.value == currentPixel.value)
                        {
                                this.currentPixel = p;
                        }
                        // If it's at a higher grey-level we must process the components in
                        // the stack (part 7)
                        else
                        {
                                this.currentPixel = p;
                                processComponentStack();
                        }
                }
        }

        private void processComponentStack()
        {
                while (this.currentPixel.value > this.componentStack.peek().pivot.value)
                {
                        // System.err.println( "Processing stack: "+componentStack );

                        // If the second component on the stack has a greater
                        // grey-level than the pixel, we set the component's grey-level
                        // to that of the pixel and quit...
                        final Component topOfStack = this.componentStack.pop();

                        // System.err.println( "Top of stack gl: "+topOfStack.greyLevel );
                        // System.err.println(
                        // "Second stack gl: "+componentStack.peek().greyLevel );
                        // System.err.println( "Pixel greylevel: "+currentPixel.value );

                        if (this.currentPixel.value < this.componentStack.peek().pivot.value)
                        {
                                topOfStack.pivot = this.currentPixel;
                                this.componentStack.push(topOfStack);

                                fireComponentStackMergeListener(componentStack.peek());

                                return;
                        }

                        fireComponentStackMergeListener(componentStack.peek(), topOfStack);

                        // Otherwise...
                        // Join the pixel lists
                        this.componentStack.peek().merge(topOfStack);

                        // TODO: histories of components
                }
        }

        /**
         * Returns the neighbouring pixels for a given pixel with 4-connectedness.
         * If the pixel lies outside of the image the result will be null.
         * 
         * @param pixel
         *            The pixel to find the neighbours of
         * @return An array of pixels some of which may be null if they lie outside
         *         of the image boundary.
         */
        private IntValuePixel[] getNeighbourPixels_4(IntValuePixel pixel)
        {
                final IntValuePixel[] p = new IntValuePixel[4];
                final int x = pixel.x;
                final int y = pixel.y;

                final int height = this.greyscaleImage.length;
                final int width = this.greyscaleImage[0].length;

                // Find the pixels
                int c = 0;

                if (x < width - 1)
                        p[c++] = new IntValuePixel(x + 1, y, greyscaleImage[y][x + 1]);

                if (x > 0)
                        p[c++] = new IntValuePixel(x - 1, y, greyscaleImage[y][x - 1]);

                if (y < height - 1)
                        p[c++] = new IntValuePixel(x, y + 1, greyscaleImage[y + 1][x]);

                if (y > 0)
                        p[c++] = new IntValuePixel(x, y - 1, greyscaleImage[y - 1][x]);

                return p;
        }

        /**
         * Add a component stack merge listener
         * 
         * @param csml
         *            The {@link ComponentStackMergeListener} to add
         */
        public void addComponentStackMergeListener(ComponentStackMergeListener csml)
        {
                csmListeners.add(csml);
        }

        /**
         * Removes the given {@link ComponentStackMergeListener} from the listeners
         * list.
         * 
         * @param csml
         *            The {@link ComponentStackMergeListener} to remove
         */
        public void removeComponentStackMergeListener(ComponentStackMergeListener csml)
        {
                csmListeners.remove(csml);
        }

        /**
         * Fire the component stack merge listener event for the merging of two
         * components.
         * 
         * @param c1
         *            The first component
         * @param c2
         *            The second component
         */
        private void fireComponentStackMergeListener(Component c1, Component c2)
        {
                for (final ComponentStackMergeListener csm : csmListeners)
                        csm.componentsMerged(c1, c2);
        }

        /**
         * Fire the component stack merge listener event for the upward merge of a
         * component.
         * 
         * @param c1
         *            The component that has been promoted to a higher intensity
         *            level
         */
        private void fireComponentStackMergeListener(Component c1)
        {
                for (final ComponentStackMergeListener csm : csmListeners)
                        csm.componentPromoted(c1);
        }

        /**
         * Helper function for debugging arrays
         * 
         * @param o
         * @return
         */
        @SuppressWarnings("unused")
        private String outputArray(Object[] o)
        {
                final StringBuilder sb = new StringBuilder();
                sb.append("[");
                boolean first = true;
                for (final Object obj : o)
                {
                        if (!first)
                                sb.append(",");
                        if (obj == null)
                                sb.append("null");
                        else
                                sb.append(obj.toString());
                        first = false;
                }
                sb.append("]");
                return sb.toString();
        }
}