#include 
#include "model.hpp"
#include "constants.hpp"
#include "tile_bounds.hpp"
#include "project_gaussians.hpp"
#include "rasterize_gaussians.hpp"
#include "tensor_math.hpp"
#include "gsplat.hpp"
#include "utils.hpp"
#include 
#include "ply.hpp"

#ifdef USE_HIP
#include 
#elif defined(USE_CUDA)
#include 
#endif

namespace fs = std::filesystem;

torch::Tensor randomQuatTensor(long long n){
    torch::Tensor u = torch::rand(n);
    torch::Tensor v = torch::rand(n);
    torch::Tensor w = torch::rand(n);
    return torch::stack({
        torch::sqrt(1 - u) * torch::sin(2 * PI * v),
        torch::sqrt(1 - u) * torch::cos(2 * PI * v),
        torch::sqrt(u) * torch::sin(2 * PI * w),
        torch::sqrt(u) * torch::cos(2 * PI * w)
    }, -1);
}

torch::Tensor projectionMatrix(float zNear, float zFar, float fovX, float fovY, const torch::Device &device){
    // OpenGL perspective projection matrix
    float t = zNear * std::tan(0.5f * fovY);
    float b = -t;
    float r = zNear * std::tan(0.5f * fovX);
    float l = -r;
    return torch::tensor({
        {2.0f * zNear / (r - l), 0.0f, (r + l) / (r - l), 0.0f},
        {0.0f, 2 * zNear / (t - b), (t + b) / (t - b), 0.0f},
        {0.0f, 0.0f, (zFar + zNear) / (zFar - zNear), -1.0f * zFar * zNear / (zFar - zNear)},
        {0.0f, 0.0f, 1.0f, 0.0f}
    }, device);
}

torch::Tensor psnr(const torch::Tensor& rendered, const torch::Tensor& gt){
    torch::Tensor mse = (rendered - gt).pow(2).mean();
    return (10.f * torch::log10(1.0 / mse));
}

torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt){
    return torch::abs(gt - rendered).mean();
}

Model::Model(const InputData &inputData,
			 const ModelParams& modelParams,
	  int maxSteps,
			 std::array backgroundColour,
	  const torch::Device &device) :
	numCameras(inputData.cameras.size()),
	params(modelParams), 
	stopSplitAt(maxSteps / 2), 
	maxSteps(maxSteps),
	device(device), 
	ssim(11, 3)
{
	//	this will fail later with a mod %0, but not the craziest check anyway
	if ( numCameras < 1 )
		throw std::runtime_error("Model requires at least one camera");
	
	//scale = inputData.scale;
	//translation = torch::tensor( {inputData.translation.x, inputData.translation.y, inputData.translation.z } );
	
	torch::manual_seed(42);
	
	//	catch data errors before obscure torch error 
	auto numPoints = inputData.points.xyz.size(0);
	auto numPointColours = inputData.points.rgb.size(0);

	if ( numPoints != numPointColours )
	{
		std::stringstream Error;
		Error