/**
   @file ImageAnalogy.h
   @author Morgan McGuire, morgan@cs.williams.edu
 */

#ifndef IMAGEANALOGY_H
#define IMAGEANALOGY_H

#include <G3D/G3DAll.h>
#include <GLG3D/GLG3D.h>
#include "Feature.h"

/**
   Implementation of Hertzmann et al.'s Image Analogies algorithm.

 */
class ImageAnalogy {
public:
    //
    // Note: p and q are always in the scale of the highest-resolution image when 
    // passed between methods.  p always indexes into A or A' and q into B or B'.
    //

    class Settings {
    public:

        std::string title;

        /**
           The paper recommends 0.5 < k < 25; I've found that that 
           over stresses coherence and values like 0.1 give better
           results.
         */
        float coherencePreference;

        enum ColorSource {SOURCE_A, SOURCE_B};

        /**
           If true, matching is only performed on luminance channels.
         */
        bool luminanceOnly;

        /**
           Should I and Q channels come from A'(p) or B(q)?
           Section 3.5 of the paper recommends using IQ from B
           for toy filters, super-resolution, and NPR.
         */
        ColorSource colorSource;

        /** If true, do not use ANN algorithm.  Saves space, consumes more time. */
        bool exhaustiveSearch;

        Settings() : coherencePreference(5.0f), luminanceOnly(true), colorSource(SOURCE_B), exhaustiveSearch(false) {}

        void deserialize(TextInput& t);
    };

private:
    /**
       Used to mark pixels in Level.s that have not yet been
       set.  Helps to identify out-of-bounds neighborhoods as well.
    */
    static const Vector2int16 UNINITIALIZED;

    Settings   settings;

    /**
       One level of a multi-scale representation of the entire analogy problem
     */
    class Level {
    public:
        /** bit shift right by this factor when looking at coordinates to account for the scale change.*/
        int shift;

        SourceImage::Ref A;
        TargetImage::Ref A2;

        /** s(q) = p; q = location in B, p = origin of that pixel in A.*/
        Map2D<Vector2int16>::Ref s;

        SourceImage::Ref B;
        TargetImage::Ref B2;

        /** Used for debugging purposes */
        Image1uint8::Ref debug;

        /**  ||[A, A'] - [B, B']||^2, looking only at in-bounds and already generated pixels.
             Examines at most (2 * radius + 1)^2 pixels 

             Semantically equivalent to using makeFeatureVector on A(p) and B(q) and computing their norm,
             but faster.

             @param luminanceOnly if true, only the luminance channel of the color is compared.
        */
        float featureDistance(int radius, Vector2int16 pCenter, Vector2int16 qCenter, bool luminanceOnly) const;
    };

    /** Initialized by createImagePyramid */
    Array<Level> levelArray;

    /** 
        Created by updateSearchTree, deleted by destructor.
     */
    class SearchTree* searchTree;

    /** Value of L to which the current searchTree corresponds */
    int searchTreeLevel;


    /** Computes the image analogy B' such that A:A' :: B:B'.  Assumes that
        the levelArray has already been filled out. 
      */
    void computeB2();

    /**
       Initializes the levelArray
     */
    void createImagePyramid(int N, Image3Ref iA, Image3Ref iA2, Image3Ref iB);

    /**
       Write out PNG files for intermediate results and an HTML file.
       This is useful for debugging (and demoing) and isn't that slow
       compared to the cost of computing the images.
     */
    void saveData();

    /** 
        Implementation of the BestMatch routine from the paper.
        operates at level L.
     */
    Vector2int16 bestMatch(int L, Vector2int16 q);

    /**
       Match feature vectors by exhaustive search. O(N^2)
     */
    void bestExhaustiveApproximationMatch(int L, Vector2int16 q, Vector2int16& bestP, float& bestDistance);

    /**
       Match feature vectors using Approximate-Nearest-Neighbor
       search. O(log N)
     */
    void bestANNApproximationMatch(int L, Vector2int16 q, Vector2int16& bestP, float& bestDistance);

    /**
       Called from bestANNApproximationMatch to update the ANN search
       tree to use level L (and L-1) features.
     */
    void updateSearchTree(int L);

    /**
       Examining only already-synthesized pixels, find the p that is
       most coherent with those pixels.
     */
    void bestCoherenceMatch(int L, Vector2int16 q, Vector2int16& bestP, float& bestDistance);

    /**
       Returns the Gaussian-weighted sum squares of all features in a
       neighborhood of A'(p) and B'(q) at levels L and L - 1,
       normalized to account for partial neighborhoods near boundaries.
    */
    float featureDistance(int L, const Vector2int16& p, const Vector2int16& q) const;

    /**
       Converts Image3 into a source feature vector (see SourceFeature) at each pixel.
     */
    static SourceImage::Ref extractSourceFeatures(const Image3Ref& A);

    /**
       Converts Image3 into a target feature vector (see SourceFeature) at each pixel.
     */
    static TargetImage::Ref extractTargetFeatures(const Image3Ref& A2);

    /**
       Convertes the B' result of the analogy back into an Image3 for
       display.  May produce some speckled colors that are over
       saturated because the IQ components of color are completely
       separted from Y during processing.
    */
    static Image3Ref reconstructImage(const TargetImage::Ref& B2);

    /**
       Gathers all feature pixels of source within the box from
       (centerP.x - radius, centerP.y - radius) to (centerP.x + radius, centerP.y +
       radius), inclusive of those bounds and excluding elements corresponding to ungenerated
       B' locations, and appends them to the feature.

       Elements are weighted by the <i>square root</i> of 2D gaussian falloff
       from centerP, normalized so that all weights sum to 1.

       Luminance-only feature vectors have 165 dimensions: (25 + 12 + 9 + 9) pixels * 3 features/pix

       Called from makeFeatureVector
       @param centerP <b>in the resolution of source</b>, unlike other calling conventions.
       @param onlyGenerated if true, source is A2 or B2 and we can only add values from pixels that have already been generated.
     */
    void extendFeatureVector(FeatureVector& feature, const Vector2int16& centerP,
                             int radius, const SourceImage::Ref& source, bool onlyGenerated = false) const;

    /**
       Creates a feature vector containing all of the features of C
       and lowC serialized, cutting out so as to correspond to only
       already generated pixels from B':

       <pre>
       .........
       .........
       ....X
       </pre>


       @param shift The level.shift value corresponding to A @param
       centerP In the resolution of the highest image.  Assumed to be
       far enough in bounds that stepping by RADIUS at the resolution
       of A will not trigger out of bounds access.

       Elements for each of the four input images are individual weighted so that
       all of their weights sum to 1.0 (for that image) and the weights are 
       distributed in a 2D gaussian falloff.

       Note that annDist on two feature vectors should return the same value as
       featureDistance() for the points corresponding to those features.
     */
    void makeFeatureVector(FeatureVector& feature, const Vector2int16& centerP,
                           SourceImage::Ref C, TargetImage::Ref C2, int shift, 
                           SourceImage::Ref lowC, TargetImage::Ref lowC2);

    
    /** 
        Called from compute()
     */
    ImageAnalogy(const Settings& settings);

public:

    ~ImageAnalogy();

    /** Computes the image analogy B' such that A:A' :: B:B' */
    static Image3Ref compute(const Image3Ref& A, const Image3Ref& A2, const Image3Ref& B, 
                             const Settings& settings = Settings());

};


#endif
