#include <librapid/array/multiarray.hpp>
#include <librapid/array/multiarray_operations.hpp>

namespace librapid
{
	void Array::transpose(const Extent& order)
	{
		Extent tempOrder;

		if (order.ndim() == 0) {
			// Order was not provided, so initialize it to be n, n-1, ... , 2, 0
			tempOrder = Extent(m_extent.ndim());
			for (int64_t i = 0; i < m_extent.ndim(); ++i)
				tempOrder[i] = m_extent.ndim() - i - 1;
		}
		else {
			tempOrder = order;
		}

		// Check the dimensions are correct
		if (tempOrder.ndim() != m_extent.ndim()) {
			throw std::invalid_argument(std::to_string(tempOrder.ndim()) + " indices "
				"were passed to Array transpose, though "
				+ std::to_string(m_extent.ndim()) +
				" indices are required");
		}

		bool valid = true;
		int64_t count = 0;
		int64_t missing[LIBRAPID_MAX_DIMS]{};
		for (int64_t i = 0; i < ndim(); ++i) {
			if (std::count(&(*tempOrder.begin()), &(*tempOrder.end()), i) != 1) {
				missing[count++] = i;
				valid = false;
			}
		}

		if (!valid) {
			auto stream = std::stringstream();
			for (int64_t i = 0; i < count; ++i) {
				if (i == m_stride.ndim() - 1) stream << missing[i];
				else stream << missing[i] << ", ";
			}
			std::string missing_str = "(" + stream.str() + ")";

			throw std::runtime_error("Transpose requires that each index is passed "
				"exactly once, but indices " + missing_str +
				" were passed more than once or not at all");
		}

		m_extent.reorder(tempOrder.toVec());
		m_stride.reorder(tempOrder.toVec());
	}

	void Array::reshape(const Extent& newShape) {
		if (m_isChild)
			throw std::runtime_error("Cannot reshape child array. Either "
				"assign it to a variable or copy the value first");

		if (!m_isScalar && (newShape.ndim() == 0 || (newShape.ndim() == 1 && newShape[0] == 0)))
		{
			m_isScalar = true;
			m_extent = Extent({ 1 });
			m_stride = Stride({ 1 });
		}

		auto tempShape = newShape.fixed(m_extent.size());

		if (tempShape.size() != m_extent.size())
			throw std::domain_error("Cannot reshape array with "
				+ m_extent.str() + " (" + std::to_string(m_extent.size())
				+ " elements) to array with " + tempShape.str()
				+ " (" + std::to_string(m_extent.size()) + " elements");

		if (!m_stride.isTrivial() || !m_stride.isContiguous()) {
			Array res(tempShape, m_dtype, m_location);
			Array::applyUnaryOp(res, *this, ops::Copy(), true);
			m_extent = res.m_extent;
			m_stride = res.m_stride;
			decrement();
			m_dataOrigin = res.m_dataOrigin;
			m_dataStart = res.m_dataStart;
			res.increment();
			return;
		}

		m_extent = tempShape;
		m_stride = Stride::fromExtent(tempShape);
	}

	Array concatenate(const std::vector<Array>& arrays, int64_t axis) {
		if (arrays.empty()) return Array();
		if (arrays.size() == 1) return arrays[0].clone();

		// Check all arrays are the same size and compute the new dimension for the
		// resulting array
		int64_t index = 0;
		int64_t dims = 0;
		Datatype resDtype = arrays[0].dtype();
		Accelerator resLocn = arrays[0].location();

		for (int64_t i = 0; i < arrays.size(); ++i)
		{
			if (arrays[i].ndim() > dims) {
				dims = arrays[i].ndim();
				index = i;
			};
			if (arrays[i].dtype() > resDtype) resDtype = arrays[i].dtype();
			if (arrays[i].location() > resLocn) resLocn = arrays[i].location();
		}

		if (axis < 0 || axis > dims)
			throw std::range_error("Axis " + std::to_string(axis)
				+ " is out of range for array with "
				+ std::to_string(dims) + " dimensions");

		int64_t newDim = arrays[index].extent()[axis];
		const Extent& dim0 = arrays[index].extent();

		for (uint64_t i = 1; i < arrays.size(); ++i) {
			bool adjustCheck = false;
			if (arrays[i].ndim() != dims && !(adjustCheck = arrays[i].ndim() == dims - 1))
				throw std::invalid_argument("To concatenate arrays, all "
					"values must have the same number of "
					"dimensions, however (at least) one "
					"array failed this condition. "
					+ std::to_string(arrays[i].ndim())
					+ " dimension(s) is not compatible with "
					+ std::to_string(dims) + " dimension(s)");

			// Ensure every dimension other than <axis> is equal
			// (concatenating on <axis> allows for different size in that dimension)
			for (int64_t j = 0; j < dims - adjustCheck; ++j) {
				if (j != axis && arrays[i].extent()[j] != dim0[j + (adjustCheck * (j > axis))])
					throw std::invalid_argument("To concatenate arrays, all "
						"dimensions other than index <axis> "
						"must be equal, however (at least) "
						"one array failed this condition. "
						+ arrays[i].extent().str()
						+ " is not compatible with "
						+ dim0.str());
			}

			newDim += adjustCheck ? 1 : arrays[i].extent()[axis];
		}

		int64_t step = 1;
		for (int64_t d = axis; d < dims; ++d)
			step *= dim0[d];

		Extent resShape = dim0;
		resShape[axis] = newDim;
		// resShape.update();

		Array res(resShape, resDtype, resLocn);

		int64_t offset = 0;
		if (axis) res._offsetData(0);

		Extent fixed = dim0;
		fixed[axis] = -1;
		for (const auto& arr : arrays)
		{
			if (arr.ndim() == dims)
				Array::applyUnaryOp(res, arr, ops::Copy(), true, offset);
			else if (arr.ndim() == dims - 1)
				Array::applyUnaryOp(res, arr.reshaped(fixed), ops::Copy(), true, offset);
			offset += step;
		}
		res._setScalar(false);

		return res;
	}

	Array stack(const std::vector<Array>& arrays, int64_t axis) {
		int64_t newDim = arrays.size();
		int64_t dims = arrays[0].ndim();

		// Perform some checks to make sure everything is going to work
		if (axis < 0 || axis > arrays[0].ndim() + 1)
			throw std::invalid_argument("Axis " + std::to_string(axis)
				+ " is out of range for arrays with "
				+ std::to_string(arrays[0].ndim()) + " dimensions");

		if (newDim == 0) return Array();
		if (newDim == 1) return arrays[0].clone();

		Extent dim0 = arrays[0].extent();
		for (int64_t i = 0; i < newDim; ++i) {
			if (arrays[i].extent() != dim0)
				throw std::invalid_argument("Array with " + arrays[i].extent().str()
					+ " cannot be stacked with array with " + dim0.str()
					+ ". All arrays must have the same extent");
		}

		Extent resShape(arrays[0].isScalar() ? 1 : dims + 1);
		if (arrays[0].isScalar()) {
			resShape[0] = arrays.size();
		}
		else {
			for (int64_t i = 0; i < dims + 1; ++i) {
				if (i < axis) resShape[i] = dim0[i];
				else if (i == axis) resShape[i] = newDim;
				else resShape[i] = dim0[i - 1];
			}
		}
		resShape.update();

		// If we're adding to the last dimension, reshape the arrays to ensure
		// that everything will turn out correctly
		if (axis == dims) {
			Extent tempShape(dims + 1);
			for (int64_t i = 0; i < dims + 1; ++i) {
				if (i == axis) tempShape[i] = 1;
				else tempShape[i] = dim0[i];
			}
			tempShape.update();

			std::vector<Array> tempArrays(arrays.size());
			for (int64_t i = 0; i < arrays.size(); ++i)
				tempArrays[i] = arrays[i].reshaped(tempShape);

			return concatenate(tempArrays, axis).reshaped(resShape);
		}

		return concatenate(arrays, axis).reshaped(resShape);
	}
}