#pragma once

#include <sstream>

#include "DataStructures.h"
#include "Image2D.h"

template<class T>
struct Grid2D
{
	static_assert(std::is_floating_point<T>::value, "Grid2D only supports float values");

    BoundingBox<T> bb;
    size_t numBinsX, numBinsY;
    T binSizeX, binSizeY;
private:
	      Image2D<T> data;

public:

    Grid2D(){   } //TODO: fast hack

    Grid2D(BoundingBox<T> bb, size_t numBins)
        : Grid2D(bb, numBins, numBins)
    {
    }

	Grid2D(BoundingBox<T> bb, size_t numBinsX, size_t numBinsY)
		: bb(bb),
          numBinsX(numBinsX),
          numBinsY(numBinsY),
          binSizeX((bb.width()) / (numBinsX-1)),
          binSizeY((bb.heigth()) / (numBinsY-1)),
          data(numBinsX, numBinsY)
	{
	}

          Image2D<T>& image()       { return data; }
    const Image2D<T>& image() const { return data; }

          T& operator() (size_t x, size_t y)       { return data(x, y); }
    const T& operator() (size_t x, size_t y) const { return data(x, y); }

    void clear() { data.clear(); }

	void fill(const std::vector<Point2D<T>>& samples)
	{
		assertMsg(!samples.empty(), "Samples must be non-empty");

		data.clear();

		const T weight = T(1.0) / samples.size();
		for (const auto& pt : samples)
		{
			add(pt.X, pt.Y, weight);
		}
	}

	void fill(const std::vector<Point2D<T>>& samples, const std::vector<T>& weights)
	{
        assertMsg(!samples.empty(), "Samples must be non-empty");
        assertMsg(weights.size() == samples.size(), "Weights must have the same size as samples");

		data.clear();

		for (size_t i = 0; i < samples.size(); i++)
		{
			add(samples[i].X, samples[i].Y, weights[i]);
		}
	}

	Point2D<size_t> add(T x, T y, T w)
	{
        if (assertCond(bb.isInside(x,y)))
        {
            std::stringstream ss;
            ss << "Point " << Point2D<T>(x, y) << " is out of bounds. " << bb;
            assertThrow(ss.str());
        }

		return add_simple_bin(x, y, w);
		//return add_linear_bin(x, y, w);
	}

    void add_linear(T x, T y, T w)
    {
        if (assertCond(bb.isInside(x, y)))
        {
            std::stringstream ss;
            ss << "Point " << Point2D<T>(x, y) << " is out of bounds. " << bb;
            assertThrow(ss.str());
        }

        add_linear_bin(x, y, w);
    }

    T fetch(T x, T y) const
    {
        size_t bin_x = (size_t)((x - bb.MinX) / binSizeX);
        size_t bin_y = (size_t)((y - bb.MinY) / binSizeY);

        //if (bin_x == data.width)  bin_x--;
        //if (bin_y == data.height) bin_y--;

        return data(bin_x, bin_y);
    }

    // Returns the summation of all bin values
    T sum() const
    {
        return data.sum();
    }

    // Takes a point in input space and converts it to grid space coordinates
    Point2D<size_t> asGridSpace(T x, T y) const
    {
        size_t bin_x = (size_t)((x - bb.MinX) / binSizeX);
        size_t bin_y = (size_t)((y - bb.MinY) / binSizeY);

        if (bin_x == data.width)  bin_x--;
        if (bin_y == data.height) bin_y--;

        return Point2D<size_t>(bin_x, bin_y);
    }

    // Takes a Size2D in input space and converts it to grid space coordiantes
    Size2D<size_t> asGridSpace(Size2D<T> sz) const
    {
        return Size2D<size_t>(sz.sX / binSizeX, sz.sY / binSizeY);
    }

    // Takes a point in grid space and converts it to input space coordinates
    Point2D<T> asInputSpace(Point2D<size_t> pt) const
    {
        return asInputSpace(pt.X, pt.Y);
    }

    // Takes a point in grid space and converts it to input space coordinates
    Point2D<T> asInputSpace(size_t x, size_t y) const
    {
        return Point2D<T>(x * binSizeX + bb.MinX + T(0.5)*binSizeX,
                          y * binSizeY + bb.MinY + T(0.5)*binSizeY);
    }


    T maximum(Point2D<T>& pt) const
    {
        Point2D<size_t> gridPt;

        T maxValue = image().maximum(gridPt);
        pt = asInputSpace(gridPt);
        return maxValue;
    }

private:
    Point2D<size_t> add_simple_bin(T x, T y, T w)
	{
		size_t bin_x = (size_t)((x - bb.MinX) / binSizeX);
		size_t bin_y = (size_t)((y - bb.MinY) / binSizeY);

        if (bin_x == data.width)  bin_x--;
        if (bin_y == data.height) bin_y--;

		data(bin_x, bin_y) += w;

        return Point2D<size_t>(bin_x, bin_y);
	}

    void add_linear_bin(T x, T y, T w)
	{
        T xpos1 = (x - bb.MinX) / binSizeX;
        T xpos2 = (y - bb.MinY) / binSizeY;
        size_t ix1 = (size_t)floor(xpos1);
        size_t ix2 = (size_t)floor(xpos2);
        T fx1 = xpos1 - ix1;
        T fx2 = xpos2 - ix2;
        
        const size_t ixmin1 = 0;
        const size_t ixmax1 = numBinsX - 2;
        const size_t ixmin2 = 0;
        const size_t ixmax2 = numBinsY - 2;

        if ( ixmin1 <= ix1 && ixmin2 <= ix2 && ix2 <= ixmax2) 
        {
            data(ix1, ix2) += w*(1 - fx1)*(1 - fx2);
            data(ix1+1, ix2) += w*fx1*(1 - fx2);
            data(ix1, ix2+1) += w*(1 - fx1)*fx2;
            data(ix1 + 1, ix2 + 1) += w*fx1*fx2;
        }
        else if (ix1 == ixmax1 + 1 && ixmin2 <= ix2 && ix2 <= ixmax2) 
        {
            // rechts
            data(ix1, ix2) += w*(1 - fx1)*(1 - fx2);
            data(ix1, ix2+1) += w*(1 - fx1)*fx2;
        }
        else if (ixmin1 <= ix1 && ix1 <= ixmax1 && ix2 == ixmax2 + 1) 
        {
            // unten
            data(ix1, ix2) += w*(1 - fx1)*(1 - fx2);
            data(ix1+1, ix2) += w*fx1*(1 - fx2);
        }
        else if (ix1 == ixmax1 + 1 && ix2 == ixmax2 + 1) 
        {
            // rechts-unten
            data(ix1, ix2) += w*(1 - fx1)*(1 - fx2);
        }
	}
};

template struct Grid2D<float>;
template struct Grid2D<double>;