#include <iostream>
#include <string>
#include <vector>
#include <cmath>
#include <armadillo>
#include "Rbm.hpp"
#include "AStack.hpp"
#include "DeepStack.hpp"
#include "StackCreator.hpp"

// Minimal, dependency-free smoke test suite: for each named project under
// prj/<name>/, load it, load its weights and training batch, run a full
// up-pass/down-pass reconstruction, and sanity-check the result. This
// replaces the old ad-hoc, since-bitrotted manual test program.

namespace
{

struct TestResult
{
	bool passed;
	std::string message;
};

// A model that hasn't learned anything shouldn't beat the trivial baseline
// of reconstructing every row as just the training batch's per-feature mean.
double meanBaselineError(const arma::mat &batch)
{
	arma::mat meanRow = arma::mean(batch, 0);
	arma::mat baseline = arma::repmat(meanRow, batch.n_rows, 1);
	return Rbm::rms_error_accu(batch - baseline);
}

TestResult testProject(const std::string &name)
{
	AStack *pStack = StackCreator::fromFile(".", name);
	if (!pStack)
	{
		return {false, "failed to parse project file"};
	}

	bool weightsOk = pStack->loadWeights(".");
	size_t numTraining = pStack->loadTrainingBatch(".");

	if (!weightsOk)
	{
		delete pStack;
		return {false, "failed to load weights"};
	}
	if (numTraining == 0)
	{
		delete pStack;
		return {false, "failed to load training batch"};
	}

	if (pStack->type() != AStack::StackType::Deep)
	{
		delete pStack;
		return {true, "skipped (reconstruction check only supports Deep stacks so far)"};
	}

	DeepStack *pDeep = static_cast<DeepStack*>(pStack);
	arma::mat batch = pStack->trainingBatch();
	arma::mat reconstruction = pDeep->upDownPass(pStack->getFirstLayer(), batch);
	delete pStack;

	if (reconstruction.n_rows != batch.n_rows || reconstruction.n_cols != batch.n_cols)
	{
		return {false, "reconstruction shape mismatch"};
	}

	double err = Rbm::rms_error_accu(batch - reconstruction);
	if (!std::isfinite(err))
	{
		return {false, "reconstruction error is not finite (NaN/Inf)"};
	}

	double baseline = meanBaselineError(batch);
	if (err > baseline * 1.5)
	{
		return {false, "reconstruction error " + std::to_string(err) +
			" is worse than 1.5x the trivial mean-baseline " + std::to_string(baseline)};
	}

	return {true, "reconstruction error = " + std::to_string(err) +
		" (mean-baseline = " + std::to_string(baseline) + ")"};
}

} // namespace

int main()
{
	std::vector<std::string> projects = {"1-hot", "prims", "norb_small_16h_v2", "norb_small_16h", "mnist_2", "prims_deep", "count", "many"};

	int numPassed = 0;
	for (const std::string &name : projects)
	{
		std::cout << "=== " << name << " ===" << std::endl;
		TestResult result = testProject(name);
		std::cout << (result.passed ? "PASS" : "FAIL") << ": " << result.message << std::endl << std::endl;
		if (result.passed)
		{
			numPassed++;
		}
	}

	std::cout << numPassed << "/" << projects.size() << " projects passed" << std::endl;
	return numPassed == (int)projects.size() ? 0 : 1;
}