#pragma once

#include <cassert>
#include <numeric>
#include <vector>

#include "../../geo/Point3.h"

enum struct SlicePlane { XY, XZ, YZ };

template <class TValue>
struct ImageView3D
{
    static_assert(std::is_arithmetic<TValue>::value, "ImageView3D only supports integers or floats");

    // forward declaration
    template<SlicePlane S> struct LineView;
    template<SlicePlane S> struct ConstLineView;

private:
    size_t width, height, depth, size; // TODO private
protected:
    TValue* values; // contains image data row-wise 
public:

    ImageView3D()
        : width(0), height(0), depth(0), size(0), values(nullptr)
    {
    }

    ImageView3D(size_t width, size_t height, size_t depth)
        : width(width), height(height), depth(depth), size(width*height*depth), values(nullptr)
    {
    }

    ImageView3D(size_t width, size_t height, size_t depth, TValue* data)
        : width(width), height(height), depth(depth), size(width*height*depth), values(data)
    {
    }

          TValue& operator() (size_t x, size_t y, size_t z)       { return values[indexFromCoord(x, y, z)]; }
    const TValue& operator() (size_t x, size_t y, size_t z) const { return values[indexFromCoord(x, y, z)]; }

    TValue* val_begin() { return values; }
    TValue* val_end  () { return values + size; }

    const TValue* val_begin() const { return values; }
    const TValue* val_end  () const { return values + size; }

	size_t getWidth()  const { return width; }
	size_t getHeight() const { return height; }
	size_t getDepth()  const { return depth; }
	
    size_t indexFromCoord(size_t x, size_t y, size_t z) const
    {
        assertMsg(x < width && y < height && z < depth, "(x,y,z) out of bounds");
        return z*width*height + y*width + x;
    }

	_Point3<size_t> coordFromIndex(size_t index) const
    {
        assertMsg(index < size, "Index out of bounds");

        size_t z = index / (width*height);
        size_t y = (index - z*width*height) / width;
        size_t x = index - y*width - z*width*height;

        return _Point3<size_t>(x, y, z);
    }

	_Point3<size_t> maximum() const
    {
        size_t maxValueIndex = std::distance(val_begin(), std::max_element(val_begin(), val_end()));
        return coordFromIndex(maxValueIndex);
    }

	TValue maximum(_Point3<size_t>& pt) const
	{
		size_t maxValueIndex = std::distance(val_begin(), std::max_element(val_begin(), val_end()));
		pt = coordFromIndex(maxValueIndex);
		return values[maxValueIndex];
	}

    void clear()
    {
        std::fill(val_begin(), val_end(), TValue(0));
    }

    void assign(const ImageView3D<TValue>& other)
    {
        assertMsg(size == other.size, "Other must be of same size as this");
        std::copy(other.val_begin(), other.val_end(), val_begin());
    }

    TValue sum() const
    {
        return std::accumulate(val_begin(), val_end(), TValue(0));
    }

    void swap(ImageView3D<TValue>& other)
    {
        TValue* tmp = other.values;
        other.values = this->values;
        this->values = tmp;
    }

    // Returns the value at (x,y). Throws if out of bounds.
    const TValue& at(size_t x, size_t y, size_t z) const
    {
        return values[indexFromCoord(x, y, z)];
    }

    // Returns the value at (x,y) but falls back to default if out of bounds.
    const TValue& get(size_t x, size_t y, size_t z, TValue dflt = 0) const
    {
        if (x >= width || y >= height || z >= depth)
            return dflt;
        else
            return values[indexFromCoord(x, y, z)];
    }

public:
    typedef LineView<SlicePlane::XY> LineViewXY;
    typedef LineView<SlicePlane::XZ> LineViewXZ;
    typedef LineView<SlicePlane::YZ> LineViewYZ;

    typedef ConstLineView<SlicePlane::XY> ConstLineViewXY;
    typedef ConstLineView<SlicePlane::XZ> ConstLineViewXZ;
    typedef ConstLineView<SlicePlane::YZ> ConstLineViewYZ;

    LineViewXY getViewXY(size_t fixedX, size_t fixedY) { return LineViewXY(*this, fixedX, fixedY); }
    LineViewXZ getViewXZ(size_t fixedX, size_t fixedZ) { return LineViewXZ(*this, fixedX, fixedZ); }
    LineViewYZ getViewYZ(size_t fixedY, size_t fixedZ) { return LineViewYZ(*this, fixedY, fixedZ); }

    ConstLineViewXY getViewXY(size_t fixedX, size_t fixedY) const { return ConstLineViewXY(*this, fixedX, fixedY); }
    ConstLineViewXZ getViewXZ(size_t fixedX, size_t fixedZ) const { return ConstLineViewXZ(*this, fixedX, fixedZ); }
    ConstLineViewYZ getViewYZ(size_t fixedY, size_t fixedZ) const { return ConstLineViewYZ(*this, fixedY, fixedZ); }

    template<SlicePlane S>
    struct LineView
    {
        size_t fixIdx1, fixIdx2;
        ImageView3D<TValue>& img;

        LineView(ImageView3D<TValue>& img, size_t fixedIndex1, size_t fixedIndex2)
            : fixIdx1(fixedIndex1), fixIdx2(fixedIndex2), img(img)
        {}

        TValue& operator() (size_t idx)
        {
            switch (S)
            {
            case SlicePlane::XY: return img(fixIdx1, fixIdx2, idx);
            case SlicePlane::XZ: return img(fixIdx1, idx, fixIdx2);
            case SlicePlane::YZ: return img(idx, fixIdx1, fixIdx2);
            default: assert(false); abort();
            }
        }
    };

    template<SlicePlane S>
    struct ConstLineView
    {
              size_t fixIdx1, fixIdx2;
        const ImageView3D<TValue>& img;

        ConstLineView(const ImageView3D<TValue>& img, size_t fixedIndex1, size_t fixedIndex2)
            : fixIdx1(fixedIndex1), fixIdx2(fixedIndex2), img(img)
        {}

        const TValue& operator() (size_t idx)
        {
            switch (S)
            {
            case SlicePlane::XY: return img(fixIdx1, fixIdx2, idx);
            case SlicePlane::XZ: return img(fixIdx1, idx, fixIdx2);
            case SlicePlane::YZ: return img(idx, fixIdx1, fixIdx2);
            default: assert(false); abort();
            }
        }

        const TValue& operator() (size_t idx) const
        {
            switch (S)
            {
            case SlicePlane::XY: return img(fixIdx1, fixIdx2, idx);
            case SlicePlane::XZ: return img(fixIdx1, idx, fixIdx2);
            case SlicePlane::YZ: return img(idx, fixIdx1, fixIdx2);
            default: assert(false); abort();
            }
        }
    };
};


template <class TValue>
struct Image3D : public ImageView3D<TValue>
{
    std::vector<TValue> values_vec;

public:
    Image3D(size_t width, size_t height, size_t depth)
        : ImageView3D<TValue>(width, height, depth), values_vec(width * height * depth)
    {
        this->values = values_vec.data();
    }

    Image3D(size_t width, size_t height, size_t depth, std::vector<TValue>& data)
        : ImageView3D<TValue>(width, height, depth), values_vec(data)
    {
        assertMsg(data.size() == width*height*depth, "Data must be of the same size as this");
        this->values = values_vec.data();
    }

          std::vector<TValue>& data()       { return values_vec; }
    const std::vector<TValue>& data() const { return values_vec; }
};

template struct ImageView3D<float>;
//template struct ImageView3D<double>;

template struct Image3D<float>;
//template struct Image3D<double>;