#ifndef LAYERS_
#define LAYERS_
#include<vector>
#include<cmath>
// 3D Matrix
struct Mat {
	int dim, row, col;
	std::vector< double > mat;
	Mat(int dim, int row, int col, std::vector< double > v) : dim(dim), row(row), col(col), mat(v) {};
};

// Mat 일차원으로 계산 + dimension 변수
class Layer {
public:
	virtual std::vector<Mat> forward(std::vector<Mat> &x) { return std::vector<Mat>(); }
};

// padding
class Convolution : public Layer {
private:
	std::vector<Mat> W;
	int stride, pad;
public:
	Convolution() {};
	~Convolution() {};
	Convolution(std::vector<Mat> W, int stride=1, int pad=0) : W(W), stride(stride), pad(pad) {};

	// Each image x conv with each w
	virtual std::vector<Mat> forward(std::vector<Mat> &x) {
		std::vector< Mat > out;
		int n, nw;
		for (n = 0; n < x.size(); n++) {
			std::vector<double>rev;
			for (nw = 0; nw < W.size(); nw++) {
				auto e = Convolution_(x[n], W[nw]);
				rev.insert(rev.end(), e.begin(), e.end());
			}
			int out_r = (x[n].row + 2 * pad - W[0].row) / stride + 1;
			int out_c = (x[n].col + 2 * pad - W[0].col) / stride + 1;
			out.push_back(Mat(nw, out_r, out_c, rev));
		}
		return out;
	}

	// Convolution x and W (both are 3-D Mat)
	std::vector<double> Convolution_(const Mat& x,const Mat& w) {
		std::vector<double> ret;
		int ndim = x.dim - w.dim + 1;
		for (int d = 0; d < x.dim - w.dim + 1; d++) {
			for (int r = -pad; r < x.row - w.row + 1 + pad; r++) {
				for (int c = -pad; c < x.col - w.col +1 +pad; c++) {
					ret.push_back(Convolution_(x, w, d, r, c));
				}
			}
		}
		return ret;
	}

	double Convolution_(const Mat& x, const Mat& w, int d, int r,int c) {
		double ret = 0, xx=0;
		int ds = w.col * w.row, rs = w.col;
		int dxs = x.col * x.row, rxs = x.col;
		for (int dd = 0; dd < w.dim; dd++) {
			for (int rr = 0; rr < w.row; rr++) {
				for (int cc = 0; cc < w.col; cc++) {
					if ((pad > 0) && (r + rr < 0 || c + cc < 0 || r + rr >= x.row || c + cc >= x.col))
						xx = 0;					
					else
						xx = x.mat[(d + dd)*(dxs)+(r + rr)*rxs + (c + cc)];
					ret += xx * w.mat[dd*(ds)+rr*(rs)+cc];					
				}
			}
		}
		return ret;
	}
};

// Depthwise Conv
class DW_Convolution : public Layer {
private:
	std::vector<Mat> W;
	int stride, pad;
public:
	DW_Convolution() {};
	~DW_Convolution() {};
	DW_Convolution(std::vector<Mat> W, int stride=1, int pad=0) : W(W), stride(stride), pad(pad) {};

	virtual std::vector<Mat> forward(std::vector<Mat> &x) {
		std::vector<Mat> out;
		int n, d;
		for (n = 0; n < x.size(); n++) {
			// Each dimension Conv with each filter
			std::vector<double> rev;			
			for (d = 0; d < x[n].dim; d++) {
				std::vector<double> e = Convolution_(x[n], W[d], d);
				rev.insert(rev.end(), e.begin(), e.end());
			}
			int out_r = (x[n].row + 2 * pad - W[0].row) / stride + 1;
			int out_c = (x[n].col + 2 * pad - W[0].col) / stride + 1;
			out.push_back(Mat(d, out_r, out_c, rev));
		}
		return out;
	}

	std::vector<double> Convolution_(const Mat& x, const Mat& w, int d) {
		std::vector<double> out;
		int dd = d * x.col * x.row;
		for (int r = -pad; r < x.row - w.row + 1 + pad; r++) { // r+=stride
			for (int c = -pad; c < x.col - w.col + 1 + pad; c++) {
				out.push_back(Convolution_(x, w, dd, r, c));
			}
		}
		return out;
	}

	double Convolution_(const Mat& x, const Mat& w, int dd, int r, int c) {
		double ret = 0, xx=0;
		for (int rr = 0; rr < w.row; rr++) {
			for (int cc = 0; cc < w.col; cc++) {
				if ((pad > 0) && (r + rr < 0 || c + cc < 0 || r + rr >= x.row || c + cc >= x.col))
					xx = 0;
				else
					xx = x.mat[dd + (r + rr)*x.col + (c+cc)];					
				ret += xx * w.mat[rr*w.col + cc];
			}
		}
		return ret;
	}
};

// n개의 이미지 같은 위치의 Row, col 값을 Normalization
class LightNormalization : public Layer{
public:
	virtual std::vector<Mat> forward(std::vector<Mat>& x) {
		std::vector<Mat> out;
		int dim = x[0].dim, row = x[0].row, col = x[0].col, nx = x.size();
		int ds = row*col;
		for (int d = 0; d < dim; d++) {	
			double mu = 0, var=0, std, tmp; // mu : mean of x img each dim
			for (int r = 0; r < row; r++) 
				for (int c = 0; c < col; c++) 
					for (int n = 0; n < nx; n++) 
						mu += x[n].mat[d*ds + r*col + c];

			mu = mu / (double)(row*col*nx);

			for (int r = 0; r < row; r++)
				for (int c = 0; c < col; c++) 
					for (int n = 0; n < nx; n++) {
						tmp = x[n].mat[d*ds + r*col + c] - mu;
						var += (tmp*tmp);
					}

			var = var / (double)(row*col*nx);
			std = sqrt(var+10e-7);
			for (int r = 0; r < row; r++) 
				for (int c = 0; c < col; c++) 
					for (int n = 0; n < nx; n++) 
						x[n].mat[d*ds + r*col + c] = (x[n].mat[d*ds + r*col + c] - mu) / std;
		}		
		return x;
	}
};

class Relu : public Layer {
public:
	virtual std::vector<Mat> forward(std::vector<Mat> &x) {
		int nx = x.size(), nm = x[0].dim * x[0].row * x[0].col;
		for (int n = 0; n < nx; n++)
			for (int i = 0; i < nm; i++)
				if (x[n].mat[i] < 0)
					x[n].mat[i] = 0;
		return x;
	}
};

class Pooling : public Layer{
private:
	int pool_h, pool_w, stride, pad;
public:
	Pooling() { pad = 0; };
	~Pooling() {};
	Pooling(int pool_h, int pool_w, int stride=1, int pad=0) :pool_h(pool_h), pool_w(pool_w), stride(stride), pad(pad) {};

	virtual std::vector<Mat> forward(std::vector<Mat>& x) {
		std::vector<Mat> out;
		int n, d, nx = x.size();
		for (n = 0; n < nx; n++) {
			std::vector<double> rev;
			for (d = 0; d < x[n].dim; d++) {
				std::vector<double> e = MaxPooling_(x[n], d);
				rev.insert(rev.end(), e.begin(), e.end());
			}
			int out_h = (x[n].row + 2 * pad - pool_h) / stride + 1;
			int out_w = (x[n].col + 2 * pad - pool_w) / stride + 1;
			out.push_back(Mat(d, out_h, out_w, rev));
		}
		return out;
	}

	// Pooling each image
	std::vector<double> MaxPooling_(Mat& x, int d) {
		std::vector<double> out;
		int row = x.row, col = x.col;
		int dd = d * col * row;
		for (int r = -pad; r < row - pool_h + 1 + pad; r+=stride) {
			for (int c = -pad; c < col - pool_w + 1 + pad; c+=stride) {
				out.push_back(MaxPooling_(x, dd, r, c));
			}
		}
		return out;
	}

	// Pooling pool_w * pool_h 
	double MaxPooling_(Mat& x, int dd, int r, int c) {
		double ret = 0, xx = 0;
		for (int rr = 0; rr < pool_h; rr++) {
			for (int cc = 0; cc < pool_w; cc++) {
				if ((pad > 0) && (r + rr < 0 || c + cc < 0 || r + rr >= x.row || c + cc >= x.col))
					xx = 0;
				else
					xx = x.mat[dd + (r + rr)*x.col + (c + cc)];
				if(ret < xx)
					ret = xx;
			}
		}
		return ret;
	}
};

class Affine : public Layer{
private:
	std::vector<Mat> W;
public:
	Affine() {}
	~Affine() {}
	Affine(std::vector<Mat>& W) : W(W){}
	
	virtual std::vector<Mat> forward(std::vector<Mat>& x) {
		std::vector<Mat> out;
		int nx = x.size();
		for (int n = 0; n < nx; n++) {
			Mat e = Dot_(x[n]);
			out.push_back(e);
		}
		return out;
	}

	Mat Dot_(const Mat& x) {		
		int dim = W[0].dim, row = W[0].row, col = W[0].col, nw = W.size();
		int size = dim*row*col;
		std::vector<double> ret(col);

		for (int c = 0; c < col; c++) {
			for (int n = 0; n < nw; n++) {
				ret[c] += W[n].mat[c] * x.mat[n];
			}

		}
		return Mat(col, 1, 1, ret);
	}
};
#endif