/**
 * Original code from https://github.com/KinglittleQ/SuperPoint_SLAM
 */

#include <superpoint_torch/SuperPoint.h>
#include <rtabmap/utilite/ULogger.h>
#include <rtabmap/utilite/UDirectory.h>
#include <rtabmap/utilite/UFile.h>
#include <rtabmap/utilite/UConversion.h>


namespace rtabmap
{

const int c1 = 64;
const int c2 = 64;
const int c3 = 128;
const int c4 = 128;
const int c5 = 256;
const int d1 = 256;



SuperPoint::SuperPoint()
      : conv1a(torch::nn::Conv2dOptions( 1, c1, 3).stride(1).padding(1)),
        conv1b(torch::nn::Conv2dOptions(c1, c1, 3).stride(1).padding(1)),

        conv2a(torch::nn::Conv2dOptions(c1, c2, 3).stride(1).padding(1)),
        conv2b(torch::nn::Conv2dOptions(c2, c2, 3).stride(1).padding(1)),

        conv3a(torch::nn::Conv2dOptions(c2, c3, 3).stride(1).padding(1)),
        conv3b(torch::nn::Conv2dOptions(c3, c3, 3).stride(1).padding(1)),

        conv4a(torch::nn::Conv2dOptions(c3, c4, 3).stride(1).padding(1)),
        conv4b(torch::nn::Conv2dOptions(c4, c4, 3).stride(1).padding(1)),

        convPa(torch::nn::Conv2dOptions(c4, c5, 3).stride(1).padding(1)),
        convPb(torch::nn::Conv2dOptions(c5, 65, 1).stride(1).padding(0)),

        convDa(torch::nn::Conv2dOptions(c4, c5, 3).stride(1).padding(1)),
        convDb(torch::nn::Conv2dOptions(c5, d1, 1).stride(1).padding(0))

{
    register_module("conv1a", conv1a);
    register_module("conv1b", conv1b);

    register_module("conv2a", conv2a);
    register_module("conv2b", conv2b);

    register_module("conv3a", conv3a);
    register_module("conv3b", conv3b);

    register_module("conv4a", conv4a);
    register_module("conv4b", conv4b);

    register_module("convPa", convPa);
    register_module("convPb", convPb);

    register_module("convDa", convDa);
    register_module("convDb", convDb);
}

std::vector<torch::Tensor> SuperPoint::forward(torch::Tensor x)
{
    x = torch::relu(conv1a->forward(x));
    x = torch::relu(conv1b->forward(x));
    x = torch::max_pool2d(x, 2, 2);

    x = torch::relu(conv2a->forward(x));
    x = torch::relu(conv2b->forward(x));
    x = torch::max_pool2d(x, 2, 2);

    x = torch::relu(conv3a->forward(x));
    x = torch::relu(conv3b->forward(x));
    x = torch::max_pool2d(x, 2, 2);

    x = torch::relu(conv4a->forward(x));
    x = torch::relu(conv4b->forward(x));

    auto cPa = torch::relu(convPa->forward(x));
    auto semi = convPb->forward(cPa);  // [B, 65, H/8, W/8]

    auto cDa = torch::relu(convDa->forward(x));
    auto desc = convDb->forward(cDa);  // [B, d1, H/8, W/8]

    auto dn = torch::norm(desc, 2, 1);
    desc = desc.div(torch::unsqueeze(dn, 1));

    semi = torch::softmax(semi, 1);
    semi = semi.slice(1, 0, 64);
    semi = semi.permute({0, 2, 3, 1});  // [B, H/8, W/8, 64]


    int Hc = semi.size(1);
    int Wc = semi.size(2);
    semi = semi.contiguous().view({-1, Hc, Wc, 8, 8});
    semi = semi.permute({0, 1, 3, 2, 4});
    semi = semi.contiguous().view({-1, Hc * 8, Wc * 8});  // [B, H, W]


    std::vector<torch::Tensor> ret;
    ret.push_back(semi);
    ret.push_back(desc);

    return ret;
}

SPDetector::SPDetector(const std::string & modelPath, float threshold, bool nms, int minDistance, bool cuda) :
		threshold_(threshold),
		nms_(nms),
		minDistance_(minDistance),
		detected_(false)
{
	UDEBUG("modelPath=%s thr=%f nms=%d cuda=%d", modelPath.c_str(), threshold, nms?1:0, cuda?1:0);
	if(modelPath.empty())
	{
		UERROR("Model's path is empty!");
		return;
	}
	std::string path = uReplaceChar(modelPath, '~', UDirectory::homeDir());
	if(!UFile::exists(path))
	{
		UERROR("Model's path \"%s\" doesn't exist!", path.c_str());
		return;
	}
	model_ = std::make_shared<SuperPoint>();
	torch::load(model_, uReplaceChar(path, '~', UDirectory::homeDir()));

	if(cuda && !torch::cuda::is_available())
	{
		UWARN("Cuda option is enabled but torch doesn't have cuda support on this platform, using CPU instead.");
	}
	cuda_ = cuda && torch::cuda::is_available();
	torch::Device device(cuda_?torch::kCUDA:torch::kCPU);
	model_->to(device);
}

SPDetector::~SPDetector()
{
}

std::vector<cv::KeyPoint> SPDetector::detect(const cv::Mat &img, const cv::Mat & mask)
{
	UASSERT(img.type() == CV_8UC1);
	UASSERT(mask.empty() || (mask.type() == CV_8UC1 && img.cols == mask.cols && img.rows == mask.rows));
	detected_ = false;
	if(!model_)
	{
		UERROR("No model is loaded!");
		return std::vector<cv::KeyPoint>();
	}
	try
	{
		torch::NoGradGuard no_grad_guard;
		auto x = torch::from_blob(img.data, {1, 1, img.rows, img.cols}, torch::kByte);
		x = x.to(torch::kFloat) / 255;

		torch::Device device(cuda_?torch::kCUDA:torch::kCPU);
		x = x.set_requires_grad(false);
		auto out = model_->forward(x.to(device));

		auto scores = out[0];       // [1, H, W]
		desc_ = out[1];             // [1, 256, H/8, W/8]

		if(nms_)
		{
			auto options = torch::nn::functional::MaxPool2dFuncOptions(minDistance_*2+1).stride(1).padding(minDistance_);
			auto options_r1 = torch::nn::functional::MaxPool2dFuncOptions(3).stride(1).padding(1);

			auto zeros = torch::zeros_like(scores);
			auto max_mask = scores == torch::nn::functional::max_pool2d(scores, options);
			auto max_mask_r1 = scores == torch::nn::functional::max_pool2d(scores, options_r1);
			for(size_t i=0; i<2; i++)
			{
				auto supp_mask = torch::nn::functional::max_pool2d(max_mask.to(torch::kF32), options) > 0;
				auto supp_scores = torch::where(supp_mask, zeros, scores);
				auto new_max_mask = supp_scores == torch::nn::functional::max_pool2d(supp_scores, options);
				max_mask = max_mask | (new_max_mask & (~supp_mask) & max_mask_r1);
			}
			prob_ = torch::where(max_mask, scores, zeros).squeeze(0);
		}
		else
		{
			prob_ = scores.squeeze(0);
		}

		auto kpts = (prob_ > threshold_);
		kpts = torch::nonzero(kpts);  // [n_keypoints, 2]  (y, x)

		//convert back to cpu if in gpu
		auto kpts_cpu = kpts.to(torch::kCPU);
		auto prob_cpu = prob_.to(torch::kCPU);

		std::vector<cv::KeyPoint> keypoints;
		for(int i=0; i<kpts_cpu.size(0); i++)
		{
			if(mask.empty() || mask.at<unsigned char>(kpts_cpu[i][0].item<int>(), kpts_cpu[i][1].item<int>()) != 0)
			{
				float response = prob_cpu[kpts_cpu[i][0]][kpts_cpu[i][1]].item<float>();
				keypoints.emplace_back(cv::KeyPoint(kpts_cpu[i][1].item<float>(), kpts_cpu[i][0].item<float>(), 8, -1, response));
			}
		}

		detected_ = true;
		return keypoints;
	}
	catch(const std::exception & e)
	{
		UERROR("SPDetector::detect() threw: %s", e.what());
		return std::vector<cv::KeyPoint>();
	}
}

cv::Mat SPDetector::compute(const std::vector<cv::KeyPoint> &keypoints)
{
	if(!detected_)
	{
		UERROR("SPDetector has been reset before extracting the descriptors! detect() should be called before compute().");
		return cv::Mat();
	}
	if(keypoints.empty())
	{
		return cv::Mat();
	}
	if(model_.get())
	{
		cv::Mat kpt_mat(keypoints.size(), 2, CV_32F);  // [n_keypoints, 2]  (y, x)

		// Based on sample_descriptors() of SuperPoint implementation in SuperGlue:
		// https://github.com/magicleap/SuperGluePretrainedNetwork/blob/45a750e5707696da49472f1cad35b0b203325417/models/superpoint.py#L80-L92
		float s = 8;
		for (size_t i = 0; i < keypoints.size(); i++) {
			kpt_mat.at<float>(i, 0) = (float)keypoints[i].pt.y - s/2 + 0.5;
			kpt_mat.at<float>(i, 1) = (float)keypoints[i].pt.x - s/2 + 0.5;
		}

		auto fkpts = torch::from_blob(kpt_mat.data, {(long int)keypoints.size(), 2}, torch::kFloat);

		float w = desc_.size(3); //W/8
		float h = desc_.size(2); //H/8

		torch::Device device(cuda_?torch::kCUDA:torch::kCPU);
		auto grid = torch::zeros({1, 1, fkpts.size(0), 2}).to(device);  // [1, 1, n_keypoints, 2]
		grid[0][0].slice(1, 0, 1) = 2.0 * fkpts.slice(1, 1, 2) / (w*s - s/2 - 0.5) - 1;  // x
		grid[0][0].slice(1, 1, 2) = 2.0 * fkpts.slice(1, 0, 1) / (h*s - s/2 - 0.5) - 1;  // y

		auto desc = torch::grid_sampler(desc_, grid, 0, 0, true);  // [1, 256, 1, n_keypoints]

		// normalize to 1
		desc = torch::nn::functional::normalize(desc.reshape({desc_.size(1), -1}), torch::nn::functional::NormalizeFuncOptions().dim(0)); //[256, n_keypoints]
		desc = desc.transpose(0, 1).contiguous(); //[n_keypoints, 256]

		if(cuda_)
			desc = desc.to(torch::kCPU);

		cv::Mat desc_mat(cv::Size(desc.size(1), desc.size(0)), CV_32FC1, desc.data_ptr<float>());

		return desc_mat.clone();
	}
	else
	{
		UERROR("No model is loaded!");
		return cv::Mat();
	}
}

}
