using LindtLeerformPlugin.Models;
using OpenCvSharp;

namespace LindtLeerformPlugin.Services;

public static class CameraCalibrationService
{
    public static CalibrationData Calibrate(string imageDirectory, Size patternSize)
    {
        var imageFiles = Directory.GetFiles(imageDirectory, "*.png");
        if (imageFiles.Length == 0)
            throw new InvalidOperationException($"No PNG images found in {imageDirectory}");

        var objectPointsList = new List<Mat>();
        var imagePointsList = new List<Mat>();
        Size imageSize = default;

        int cornerCount = patternSize.Width * patternSize.Height;
        var objPts = new Point3f[cornerCount];
        for (int row = 0; row < patternSize.Height; row++)
            for (int col = 0; col < patternSize.Width; col++)
                objPts[row * patternSize.Width + col] = new Point3f(col, row, 0);

        int found = 0;
        foreach (var file in imageFiles)
        {
            using var image = Cv2.ImRead(file);
            using var gray = new Mat();
            Cv2.CvtColor(image, gray, ColorConversionCodes.BGR2GRAY);
            imageSize = new Size(image.Width, image.Height);

            if (Cv2.FindChessboardCorners(gray, patternSize, out var corners,
                    ChessboardFlags.AdaptiveThresh | ChessboardFlags.FastCheck))
            {
                Cv2.CornerSubPix(gray, corners,
                    new Size(11, 11), new Size(-1, -1),
                    new TermCriteria(CriteriaTypes.Eps | CriteriaTypes.MaxIter, 30, 0.001));

                var objMat = new Mat(cornerCount, 1, MatType.CV_32FC3);
                for (int i = 0; i < cornerCount; i++)
                    objMat.Set(i, 0, objPts[i]);
                objectPointsList.Add(objMat);

                var imgMat = new Mat(corners.Length, 1, MatType.CV_32FC2);
                for (int i = 0; i < corners.Length; i++)
                    imgMat.Set(i, 0, corners[i]);
                imagePointsList.Add(imgMat);

                found++;
            }
        }

        if (found == 0)
            throw new InvalidOperationException("Checkerboard corners not found in any image");

        var cameraMatrix = new Mat();
        var distCoeffs = new Mat();
        double rms = Cv2.CalibrateCamera(
            objectPointsList, imagePointsList, imageSize,
            cameraMatrix, distCoeffs,
            out _, out _);

        var data = new CalibrationData
        {
            CameraMatrix = MatToArray(cameraMatrix),
            DistCoeffs = MatToArray(distCoeffs),
            RmsError = rms,
            ImageWidth = imageSize.Width,
            ImageHeight = imageSize.Height
        };

        cameraMatrix.Dispose();
        distCoeffs.Dispose();
        foreach (var m in objectPointsList) m.Dispose();
        foreach (var m in imagePointsList) m.Dispose();

        return data;
    }

    public static (Mat cameraMatrix, Mat distCoeffs) LoadCalibrationMats(CalibrationData data)
    {
        var cameraMatrix = new Mat(3, 3, MatType.CV_64F);
        for (int i = 0; i < 9; i++)
            cameraMatrix.Set(i / 3, i % 3, data.CameraMatrix[i]);

        var distCoeffs = new Mat(1, data.DistCoeffs.Length, MatType.CV_64F);
        for (int i = 0; i < data.DistCoeffs.Length; i++)
            distCoeffs.Set(0, i, data.DistCoeffs[i]);

        return (cameraMatrix, distCoeffs);
    }

    private static double[] MatToArray(Mat mat)
    {
        var total = mat.Rows * mat.Cols;
        var arr = new double[total];
        for (int i = 0; i < total; i++)
            arr[i] = mat.At<double>(i / mat.Cols, i % mat.Cols);
        return arr;
    }
}