// Drop-in replacement for OpticsFEM external complex linear solver.
// Reads CSR matrix + RHS from binary .dat files (same protocol as solver/interface.cpp)
// and writes x_real_0.dat / x_complex_0.dat for the main program.

#include <Eigen/Sparse>
#include <Eigen/SparseLU>
#include <Eigen/IterativeLinearSolvers>

#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>

using SpMat = Eigen::SparseMatrix<std::complex<double>>;
using VecX = Eigen::VectorXcd;

static std::vector<int> readIntFile(const char* path)
{
	std::ifstream in(path, std::ios::binary);
	if (!in)
		throw std::runtime_error(std::string("cannot open ") + path);
	in.seekg(0, std::ios::end);
	const std::streamsize bytes = in.tellg();
	in.seekg(0, std::ios::beg);
	std::vector<int> data(static_cast<size_t>(bytes / sizeof(int)));
	if (!data.empty())
		in.read(reinterpret_cast<char*>(data.data()), bytes);
	return data;
}

static std::vector<double> readDoubleFile(const char* path)
{
	std::ifstream in(path, std::ios::binary);
	if (!in)
		throw std::runtime_error(std::string("cannot open ") + path);
	in.seekg(0, std::ios::end);
	const std::streamsize bytes = in.tellg();
	in.seekg(0, std::ios::beg);
	std::vector<double> data(static_cast<size_t>(bytes / sizeof(double)));
	if (!data.empty())
		in.read(reinterpret_cast<char*>(data.data()), bytes);
	return data;
}

static void writeDoubleFile(const char* path, const double* data, int n)
{
	std::ofstream out(path, std::ios::binary | std::ios::trunc);
	if (!out)
		throw std::runtime_error(std::string("cannot write ") + path);
	out.write(reinterpret_cast<const char*>(data), static_cast<std::streamsize>(n) * sizeof(double));
}

static SpMat loadCsrMatrix(const std::vector<int>& roffsets,
	const std::vector<int>& cindices,
	const std::vector<double>& aReal,
	const std::vector<double>& aImag)
{
	if (roffsets.size() < 2)
		throw std::runtime_error("invalid Roffsets.dat");
	const int n = static_cast<int>(roffsets.size()) - 1;
	const int nnz = roffsets.back();
	if (static_cast<int>(cindices.size()) != nnz
		|| static_cast<int>(aReal.size()) != nnz
		|| static_cast<int>(aImag.size()) != nnz)
		throw std::runtime_error("CSR dimensions mismatch");

	std::vector<Eigen::Triplet<std::complex<double>>> triplets;
	triplets.reserve(static_cast<size_t>(nnz));
	for (int i = 0; i < n; ++i)
	{
		for (int k = roffsets[static_cast<size_t>(i)]; k < roffsets[static_cast<size_t>(i + 1)]; ++k)
		{
			triplets.emplace_back(i, cindices[static_cast<size_t>(k)],
				std::complex<double>(aReal[static_cast<size_t>(k)], aImag[static_cast<size_t>(k)]));
		}
	}

	SpMat A(n, n);
	A.setFromTriplets(triplets.begin(), triplets.end());
	A.makeCompressed();
	return A;
}

static VecX loadRhs(const std::vector<double>& bReal, const std::vector<double>& bImag, int n)
{
	if (static_cast<int>(bReal.size()) != n || static_cast<int>(bImag.size()) != n)
		throw std::runtime_error("RHS dimensions mismatch");
	VecX b(n);
	for (int i = 0; i < n; ++i)
		b(i) = std::complex<double>(bReal[static_cast<size_t>(i)], bImag[static_cast<size_t>(i)]);
	return b;
}

static VecX solveDirect(const SpMat& A, const VecX& b)
{
	Eigen::SparseLU<SpMat> lu;
	lu.compute(A);
	if (lu.info() != Eigen::Success)
		throw std::runtime_error("SparseLU factorization failed");
	VecX x = lu.solve(b);
	if (lu.info() != Eigen::Success)
		throw std::runtime_error("SparseLU solve failed");
	return x;
}

static VecX solveIterative(const SpMat& A, const VecX& b, double relTol, double absTol, int& iterations)
{
	Eigen::BiCGSTAB<SpMat, Eigen::IncompleteLUT<std::complex<double>>> solver;
	solver.preconditioner().setDroptol(1e-4);
	solver.setMaxIterations(2000);
	solver.setTolerance(relTol > 0.0 ? relTol : 1e-8);
	solver.compute(A);
	if (solver.info() != Eigen::Success)
		throw std::runtime_error("BiCGSTAB setup failed");

	VecX x = solver.solve(b);
	iterations = static_cast<int>(solver.iterations());
	const double bn = b.norm();
	const double rn = (A * x - b).norm();
	if (solver.info() != Eigen::Success && rn > absTol && rn > relTol * bn)
		throw std::runtime_error("BiCGSTAB did not converge");
	return x;
}

int main()
{
	try
	{
		const std::vector<int> roffsets = readIntFile("Roffsets.dat");
		const std::vector<int> cindices = readIntFile("Cindices.dat");
		const std::vector<double> aReal = readDoubleFile("A_real.dat");
		const std::vector<double> aImag = readDoubleFile("A_imag.dat");
		const std::vector<double> bReal = readDoubleFile("B_real.dat");
		const std::vector<double> bImag = readDoubleFile("B_imag.dat");
		const std::vector<double> opt = readDoubleFile("opt.dat");

		const int n = static_cast<int>(roffsets.size()) - 1;
		const double relTol = opt.size() > 3 ? opt[3] : 1e-3;
		const double absTol = opt.size() > 4 ? opt[4] : 1e-6;

		std::cerr << "[complexsolver] n=" << n << " nnz=" << roffsets.back()
			<< " relTol=" << relTol << std::endl;

		const SpMat A = loadCsrMatrix(roffsets, cindices, aReal, aImag);
		const VecX b = loadRhs(bReal, bImag, n);

		VecX x;
		std::string method;
		int iterations = 0;

		const char* mode = std::getenv("OPTICSFEM_COMPLEXSOLVER");
		const bool useIterative = mode && std::strcmp(mode, "iter") == 0;

		if (useIterative && n >= 5000)
		{
			x = solveIterative(A, b, relTol, absTol, iterations);
			method = "BiCGSTAB+ILUT";
		}
		else
		{
			x = solveDirect(A, b);
			method = "SparseLU";
		}

		std::vector<double> xr(static_cast<size_t>(n)), xi(static_cast<size_t>(n));
		for (int i = 0; i < n; ++i)
		{
			xr[static_cast<size_t>(i)] = x(i).real();
			xi[static_cast<size_t>(i)] = x(i).imag();
		}

		writeDoubleFile("x_real_0.dat", xr.data(), n);
		writeDoubleFile("x_complex_0.dat", xi.data(), n);

		const double residual = (A * x - b).norm();
		const double bnorm = b.norm();
		writeDoubleFile("res.dat", &residual, 1);

		std::ofstream log("log.dat", std::ios::trunc);
		log << "method=" << method << "\n";
		log << "n=" << n << " nnz=" << roffsets.back() << "\n";
		log << "|b|=" << bnorm << " |r|=" << residual << "\n";
		if (iterations > 0)
			log << "iterations=" << iterations << "\n";

		std::cerr << "[complexsolver] done " << method << " |x|=" << x.norm()
			<< " |r|=" << residual << std::endl;
		return 0;
	}
	catch (const std::exception& ex)
	{
		std::cerr << "[complexsolver] ERROR: " << ex.what() << std::endl;
		return 1;
	}
}