#include "knn_img.h"
#include <cmath>
#include <cassert>

#define MIN(x,y) ((x <= y)? x : y)
#define MAX(x,y) ((x >= y)? x : y)

#include <iostream>
using namespace std;

// From external; used in actual algorithm as well.
int compute_weight(int i,int j,int x,int y, int range_nx,int range_px,
                                int range_ny,int range_py,Img& img)
{
    int w=0;
    for(int d = -range_nx; d <= range_px; ++d)
    {
        for(int e = -range_ny; e <= range_py; ++e)
        {
            for(int f = 0; f < img.get_Bpp(); ++f)
            {
                int wtemp = (int)img.get_data(i+d,j+e,f)-
                                    (int)img.get_data(x+d,y+e,f);
                if(wtemp < 0)
                    wtemp = -wtemp;
                w += wtemp;
            }
        }
    }
    return w;
}




void NearestNeighbors::addNeighbor(int i, int j, double weight)
{
    double maxweight=neighbors[0].weight;
    int maxind=0;
    for(int d = 0; d < k; ++d)
    {
        if(neighbors[d].weight >= maxweight)
        {
            maxweight = neighbors[d].weight;
            maxind = d;
        }
    }
    if(weight < maxweight)
    {
        neighbors[maxind].i = i;
        neighbors[maxind].j = j;
        neighbors[maxind].weight = weight;
    }
}


void KNN::compute_nns(Img& img)
{
    // Boundary case decision. If a pixel is close enough to a boundary
    // that it can't use the full nradius, the comparison window is resized.
    // I.e. at the topleft corner, the window with which it compares is of range
    //  (pixelx + [0,nradius], pixely + [-nradius,0]). Where the normal
    // comparison window is (pixelx + [-nradius,nradius], (same for pixely))
    // Other boundary conditions get sized accordingly.
    // Then for pixels that cannot contain the same window; no comparison 
    // is done. I.e., top left pixel never compared to bottom right pixel.

    // Also, never compare a pixel to itself.

    // Weight is computed as Sum[abs[r1-r2] + ... g, b] * C
    //  where C = num_pixels_tested/num_orig_pixels. This is so that
    //  if a very small window is checked due to boundaries, it won't be
    //  as important as a full window.

    // Require the img to be >= 2*nradius + 1 in width and height.
    assert( 2*nradius + 1 <= img.get_w());
    assert( 2*nradius + 1 <= img.get_h());

    for(int i = 0; i < w; ++i)
    {
        for(int j = 0; j < h; ++j)
        {
            int mini = MAX(0, i-nradius);
            int minj = MAX(0, j-nradius);
            int maxi = MIN(w-1, i+nradius);
            int maxj = MIN(h-1, j+nradius);

            for(int x = 0; x < w; ++x)
            {
                for(int y = 0; y < h; ++y)
                {
                    if( (i != x) || (j != y) )
                    {
                        int range_nx = i-mini;
                        int range_px = maxi-i;
                        int range_ny = j-minj;
                        int range_py = maxj-j;

                        // Late addition; forces all nearest neighbors to
                        // be on the interior, unless the original pixel had
                        // a distorted neighborhood.
                        if( (x-range_nx >= 0) && (x+range_px < w) &&
                            (y-range_ny >= 0) && (y+range_py < h) )
                        {

                            int rawweight = compute_weight(i,j,x,y,
                                   range_nx,range_px,range_ny,range_py,img);

                            double num_pixels_tested=(range_px + range_nx + 1)*
                                                (range_py + range_ny + 1);
                            double num_orig_pixels= (2*nradius+1)*(2*nradius+1);

                            double C = num_pixels_tested/num_orig_pixels;
                            double weight = C*(double)(rawweight);

                            nns[ind(i,j)].addNeighbor(x,y,weight);
                        }
                    }
                }
            }
        }
    }
}