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

#include "ImageAnalogy.h"
#include "SearchTree.h"

// Neighborhood from the paper for feature vector (F)
static const int LOW_RADIUS = 1;
static const int RADIUS = 2;

/**
 2D gaussian of a given radius, evaluated at  -radius <= offset.x,y <= radius.
 */
static float gaussian(int radius, Vector2int16 offset) {
    
    // Tables of 1D gaussian filters at different resolutions.
    // Radius is known to be small, so we use an array instead of
    // a hashtable for speed.
    static Array< Array<float> > coefficient;

    if (coefficient.size() < radius) {
        coefficient.resize(radius);
    }

    if (coefficient[radius].size() == 0) {
        // Initialize the table
        gaussian1D(coefficient[radius], 2 * radius + 1, radius);
    }

    return coefficient[radius][offset.x + radius] * coefficient[radius][offset.y + radius];
}

bool PRINT_DISTANCE = false;

void ImageAnalogy::computeB2() {

    // This is the CreateImageAnalogy function from the paper.  
    // Initialization is already complete at this point in the code.

    printf("Solving for B2 (intermediate results will be available from out/out.html after each stage.)\n");

    // Work through the levels from lowest-resolution to highest-resolution
    for (int L = 0; L < levelArray.size(); ++L) {
        Level& level = levelArray[L];

        int w = level.B->width();
        int h = level.B->height();

        printf("Level %d (%dx%d)...", L, w, h);
        // Flush so that the message shows up right away
        fflush(stdout);

        Vector2int16 q;
        for (q.y = 0; q.y < h; ++q.y) {
            for (q.x = 0; q.x < w; ++q.x) {

                // Find the location of the best matching pixel (adjusting coordinates
                // to highest-resolution ones)
                Vector2int16 p = bestMatch(L, q << level.shift);

                // Copy the best match over
                TargetFeature best = level.A2->get(p >> level.shift);
                if (settings.colorSource == Settings::SOURCE_B) {
                    // Override with the color from B
                    ColorYIQ bColor = level.B->get(q).color;
                    best.color.i = bColor.i;
                    best.color.q = bColor.q;
                }
                level.B2->set(q, best);

                // Keep track of where the pixel came from, for coherence
                level.s->set(q, p);
            }
            
            if ((h > 70) && (q.y % 7 == 0)) {
                printf("%d%%,", 100 * q.y / h);
                fflush(stdout);
                saveData();
            }
        }

        printf("done\n");
        // Flush so that the message shows up right away
        fflush(stdout);

        saveData();
    }
}


Vector2int16 ImageAnalogy::bestMatch(int L, Vector2int16 q) {
    const float k = settings.coherencePreference;

    // Positions and feature distance of best matches
    Vector2int16 p_approx;
    float d_approx;

    if (settings.exhaustiveSearch) {
        bestExhaustiveApproximationMatch(L, q, p_approx, d_approx);
    } else {
        bestANNApproximationMatch(L, q, p_approx, d_approx);
    }
    
    Vector2int16 p_coherence = UNINITIALIZED;
    float d_coherence = inf();
    if (k > 0) {
        bestCoherenceMatch(L, q, p_coherence, d_coherence);
    }
    
    // Favor coherence at high resolutions
    float approxWeight = 0.1 + k * pow(2.0f, (float)L - levelArray.size() + 1);

    // Choose the one with the least error, scaling the importance of approximation with 
    // resolution so that it is most important at high resolution
    if ((d_coherence <= d_approx * approxWeight) && (k != 0)) {
        levelArray[L].debug->set(q >> levelArray[L].shift, Color1uint8(127));
        return p_coherence;
    } else {
        levelArray[L].debug->set(q >> levelArray[L].shift, Color1uint8(255));
        return p_approx;
    }
}


void ImageAnalogy::bestCoherenceMatch(int L, Vector2int16 q, Vector2int16& bestP, float& bestDistance) {
    // Find r that minimizes featureDistance(s(r) + (q-r), q)
    // by looking around q at already synthesized pixels.

    bestP = Vector2int16(0, 0);
    bestDistance = inf();

    const Level& current = levelArray[L];

    // The paper is ambiguous on this point. Setting equal to 1 means that we only look at strictly
    // adjacent pixels, guaranteeing perfect feature coherence.  Larger than 1 means that we are
    // allowed to rescale an image feature when preserving coherece, since a "coherent" set of pixels
    // will then draw from, e.g., every-other pixel of the source.
    const int coherenceRadius = 1;

    // Look in the source neighborhood of all current-resolution neighbors
    // that have already been synthesized

    // offset is in the level L coordinate system
    Vector2int16 offset;
    for (offset.y = -coherenceRadius; offset.y <= 0; ++offset.y) {

        // Don't look at unsynthesized pixels
        int maxOffset = (offset.y < 0) ? coherenceRadius : -1;

        for (offset.x = -coherenceRadius; offset.x <= maxOffset; ++offset.x) {
            // Compute r in the current-resolution coordinate system
            Vector2int16 r = (q >> current.shift) + offset;

            // Don't bother looking at out of bounds locations; there is 
            // no corresponding pixel to draw from.
            if (current.s->inBounds(r)) {
            
                // We have to shift offset by current.shift because it is in the coordinate system
                // of the level L image.
                // Note: -offset = (q - r) from the paper.
                Vector2int16 p = current.s->get(r) - (offset << current.shift);
                
                float distance = featureDistance(L, p, q);
                
                if (distance < bestDistance) {
                    bestDistance = distance;
                    bestP = p;
                }
            }
        }
    }
}


void ImageAnalogy::updateSearchTree(int L) {
    if (searchTreeLevel == L) {
        // The tree is already up to date
        return;
    }

    searchTreeLevel = L;
    // Make a new tree
    debugAssert(L > 0);
    delete searchTree;

    searchTree = new SearchTree();

    const Level& current  = levelArray[L];
    const Level& previous = levelArray[L - 1];

    const int inc = 1 << current.shift;
    
    // Full resolution dimensions
    int w = levelArray.last().s->width();
    int h = levelArray.last().s->height();

    // Iterate through
    FeatureVector feature;
    Vector2int16 p;

    // Don't bother making feature points for areas on borders; we can never 
    // match them since they would have the wrong number of dimensions (and therefore
    // end up being clamped and padded in a way that makes little sense).
    int lowRad = LOW_RADIUS << previous.shift;

    for (p.y = lowRad; p.y < h - lowRad; p.y += inc) {
        for (p.x = lowRad; p.x < w - lowRad; p.x += inc) {
            makeFeatureVector(feature, p, current.A, current.A2, current.shift, previous.A, previous.A2);
            searchTree->addPoint(p, feature);
        }
    }

    searchTree->balance();
}


void ImageAnalogy::makeFeatureVector
(FeatureVector& feature, 
 const Vector2int16& p,
 SourceImage::Ref C,
 TargetImage::Ref C2,
 int shift,
 SourceImage::Ref lowC,
 TargetImage::Ref lowC2) {
    
    feature.fastClear();

    // Only C2 has ungenerated elements; all other images use the full
    // neighborhood at their resolution.
    
    // High resolution first
    extendFeatureVector(feature, p >> shift, RADIUS, C);               // 25 * 3 points
    extendFeatureVector(feature, p >> shift, RADIUS, C2, true);        // 12 * 3 points
    extendFeatureVector(feature, p >> (shift + 1), LOW_RADIUS, lowC);  // 9 * 3 points
    extendFeatureVector(feature, p >> (shift + 1), LOW_RADIUS, lowC2); // 9 * 3 points
} 


void ImageAnalogy::bestExhaustiveApproximationMatch(int L, Vector2int16 q, Vector2int16& bestP, float& bestDistance) {
    // For now, brute force search...an O(N^4) total algorithm!
    int w = levelArray.last().A->width();
    int h = levelArray.last().A->height();
    int shift = levelArray[L].shift;
    int inc = 1 << shift;

    bestP = Vector2int16(0, 0);
    bestDistance = inf();
    
    // Iterate over full res pixels of A, but step at the resolution of
    // level L
    Vector2int16 p;
    for (p.y = 0; p.y < h; p.y += inc) {
        for (p.x = 0; p.x < w; p.x += inc) {
            float distance = featureDistance(L, p, q);
            if (distance < bestDistance) {
                bestDistance = distance;
                bestP = p;
            }
        }
    }
}
    
void ImageAnalogy::bestANNApproximationMatch(int L, Vector2int16 q, Vector2int16& bestP, float& bestDistance) {

    const Level& current = levelArray[L];

    // Low-resolution B dimensions are half the current resolution (we
    // can look at this before we even know if L > 0)
    int lowShift = current.shift + 1;
    int lowBw = current.s->width() >> 1;
    int lowBh = current.s->height() >> 1;
    Vector2int16 lowQ = q >> lowShift;

    // Detect border cases 
    bool atLevelBorder = (L == 0);
    bool atXBorder = (lowQ.x < LOW_RADIUS) || (lowQ.x >= lowBw - LOW_RADIUS);
    bool atYBorder = (lowQ.y < LOW_RADIUS) || (lowQ.y >= lowBh - LOW_RADIUS);

    if (atLevelBorder || atYBorder || atXBorder) {
        // We're at a border case where the feature vectors have
        // non-standard length.  Defer to exhaustive search, which
        // should be relatively fast since we're on the
        // lowest-resolution image or on the perimeter.
        bestExhaustiveApproximationMatch(L, q, bestP, bestDistance);
        return;
    }
    // We are not at a border case, so the feature vectors must both have full size.

    // Make sure the search tree has been prepared for this resolution
    updateSearchTree(L);

    // Create the feature vector for this pixel
    FeatureVector feature;
    const Level& previous = levelArray[L - 1];

    makeFeatureVector(feature, q, current.B, current.B2, current.shift, previous.B, previous.B2);
    bestP = searchTree->findNearestNeighbor(feature, 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 ImageAnalogy::featureDistance(int L, const Vector2int16& pCenter, const Vector2int16& qCenter) const {
    
    float distance = 0.0f;

    // The order of computations is specifically ordered to match makeFeatureVector

    // High-resolution ||[A, A'] - [B, B']||^2
    distance += levelArray[L].featureDistance(RADIUS, pCenter, qCenter, settings.luminanceOnly);

    // Low-resolution ||[A, A'] - [B, B']||^2
    if (L > 0) {
        distance += levelArray[L - 1].featureDistance(LOW_RADIUS, pCenter, qCenter, settings.luminanceOnly);
    }

    return distance;
}


Image3Ref ImageAnalogy::compute(const Image3Ref& iA, const Image3Ref& iA2, const Image3Ref& iB, const Settings& settings) {
    
    ImageAnalogy analogy(settings);

    // Num levels
    int N = 3;
    analogy.createImagePyramid(N, iA, iA2, iB);

    analogy.computeB2();

    TargetImage::Ref fB2 = analogy.levelArray.last().B2;
    Image3Ref iB2 = reconstructImage(fB2);

    analogy.saveData();

    return iB2;
}



void ImageAnalogy::extendFeatureVector
(FeatureVector& feature, 
 const Vector2int16& centerP,
 int radius,
 const SourceImage::Ref& source,
 bool onlyGenerated) const {
    
    // In the case where we are limited to onlyGenerated points,
    // grab all pixels in raster order within radius up to but not
    // including the one at p, which will correspond to the pixel
    // being generated in B' when we issue search requests.
    int ycutoff = onlyGenerated ? 0 : radius;
    
    int first = feature.length();

    Vector2int16 offset;
    float normalize = 0.0f;
    for (offset.y = -radius; offset.y <= ycutoff; ++offset.y) {
        // Stop before the pixel being generated
        int xcutoff = (offset.y == 0 && onlyGenerated) ? -1 : radius;
        for (offset.x = -radius; offset.x <= xcutoff; ++offset.x) {
            Vector2int16 p = centerP + offset;

            // Extract the feature information at the pixel
            const SourceFeature& f = source->get(p);
            
            // Weigh by square root of gaussian (this will end up
            // weighing each component by a gaussian when the norm is
            // computed, since that involves a square).
            float weight = gaussian(radius, offset);

            (f * sqrt(weight)).extendFeatureVector(feature, settings.luminanceOnly);
            
            normalize += weight;
        }
    }

    // Normalize the elements that we just added based on the sum of all weights.
    // We want to divide the *squares* of the elements by the sum of the weights,
    // but we are multiplying the elements themselves so we must normalize
    // by the square root of that factor.
    normalize = sqrt(1.0f / normalize);
    for (int i = first; i < feature.length(); ++i) {
        feature[i] *= normalize;
    }
}


float ImageAnalogy::Level::featureDistance
(
 int             radius, 
 Vector2int16    pCenter,
 Vector2int16    qCenter, 
 bool            luminanceOnly) const {
    
    // Weighted sum squares of the feature differences
    float distance1 = 0.0f;
    float distance2 = 0.0f;

    // The normalizing factor is the total number of filter taps; it takes into account
    // the lack of sufficient points at borders.
    float normalize1 = 0.0f;
    float normalize2 = 0.0f;

    // Adjust to the resolution of the level
    pCenter = pCenter >> shift;
    qCenter = qCenter >> shift;

    Vector2int16 offset;
    for (offset.y = -radius; offset.y <= radius; ++offset.y) {
        for (offset.x = -radius; offset.x <= radius; ++offset.x) {
            Vector2int16 p = pCenter + offset;
            Vector2int16 q = qCenter + offset;
            if (A->inBounds(p) && B->inBounds(q)) {
                
                // Gaussian falloff
                float weight = gaussian(radius, offset);

                distance1 += A->get(p).distanceSquared(B->get(q), luminanceOnly) * weight;
                normalize1 += weight;

                // When comparing A' and B', only look at already synthesized neighborhoods.
                if (s->get(q) != UNINITIALIZED) {
                    distance2 += A2->get(p).distanceSquared(B2->get(q), luminanceOnly) * weight;
                    normalize2 += weight;
                }
            }
        }
    }
    

    if ((normalize1 > 0) && (normalize2 > 0)) {
        return distance1 / normalize1 + distance2 / normalize2;
    } else {
        // At least one term went to infinity
        return inf();
    }

}
