using OpenCvSharp;

public static class Segmentation
{
    // Otsu thresholding workflow: convert to gray, blur, Otsu threshold
    public static Mat OtsuThreshold(Mat src)
    {
        Mat gray = new();
        Cv2.CvtColor(src, gray, ColorConversionCodes.BGR2GRAY);
        Mat blur = new();
        Cv2.GaussianBlur(gray, blur, new Size(5, 5), 0);
        Mat mask = new();
        Cv2.Threshold(blur, mask, 0, 255, ThresholdTypes.Binary | ThresholdTypes.Otsu);
        return mask;
    }

    // K-means color segmentation
    public static Mat KMeansSegmentation(Mat src, int k = 2, int attempts = 5)
    {
        // Reshape to a 2D samples matrix where each row is a pixel and columns are channels
        Mat samples = src.Reshape(1, src.Rows * src.Cols);
        Mat samples32f = new();
        samples.ConvertTo(samples32f, MatType.CV_32F);

        Mat labels = new();
        TermCriteria criteria = new TermCriteria(CriteriaTypes.Eps | CriteriaTypes.MaxIter, 10, 1.0);
        Mat centers = new();

        Cv2.Kmeans(samples32f, k, labels, criteria, attempts, KMeansFlags.PpCenters, centers);

        // Build segmented image by mapping each pixel to its cluster center color
        Mat result = new Mat(src.Rows, src.Cols, src.Type());
        int cols = src.Cols;
        for (int r = 0; r < src.Rows; r++)
        {
            for (int c = 0; c < cols; c++)
            {
                int idx = r * cols + c;
                int clusterIdx = labels.At<int>(idx, 0);
                byte b = (byte)Clamp(centers.At<float>(clusterIdx, 0));
                byte g = (byte)Clamp(centers.At<float>(clusterIdx, 1));
                byte rcol = (byte)Clamp(centers.At<float>(clusterIdx, 2));
                result.Set(r, c, new Vec3b(b, g, rcol));
            }
        }

        return result;
    }

    // Simple morphological cleaning (opening then closing)
    public static Mat MorphologicalClean(Mat mask, int kernelSize = 3)
    {
        Mat kernel = Cv2.GetStructuringElement(MorphShapes.Rect, new Size(kernelSize, kernelSize));
        Mat opened = new();
        Cv2.MorphologyEx(mask, opened, MorphTypes.Open, kernel);
        Mat cleaned = new();
        Cv2.MorphologyEx(opened, cleaned, MorphTypes.Close, kernel);
        return cleaned;
    }

    // Create a binary mask by HSV range. lower and upper are HSV in OpenCV order: H:0-180, S:0-255, V:0-255
    public static Mat ColorRangeHSV(Mat src, Scalar lowerHsv, Scalar upperHsv)
    {
        Mat hsv = new();
        Cv2.CvtColor(src, hsv, ColorConversionCodes.BGR2HSV);
        Mat mask = new();
        Cv2.InRange(hsv, lowerHsv, upperHsv, mask);
        return mask;
    }

    // Create a binary mask by Euclidean distance in color space to a target color.
    // If useLab is true, convert BGR->Lab before distance to get perceptual distance.
    // targetColor is Scalar(b, g, r). thresh is distance threshold.
    public static Mat ColorDistanceMask(Mat src, Scalar targetColor, double thresh, bool useLab = false)
    {
        Mat working = new();
        if (useLab)
        {
            Cv2.CvtColor(src, working, ColorConversionCodes.BGR2Lab);
        }
        else
        {
            src.CopyTo(working);
        }

        Mat f32 = new();
        working.ConvertTo(f32, MatType.CV_32F);

        // Prepare a Mat filled with the target color (in same color space)
        Mat target = new Mat(f32.Size(), MatType.CV_32FC3, new Scalar(targetColor.Val0, targetColor.Val1, targetColor.Val2));

        Mat diff = new();
        Cv2.Subtract(f32, target, diff);

        Mat sq = new();
        Cv2.Multiply(diff, diff, sq);

        // sum channels
        Mat[] ch = Cv2.Split(sq);
        Mat sum = new();
        Cv2.Add(ch[0], ch[1], sum);
        Cv2.Add(sum, ch[2], sum);

        Mat dist = new();
        Cv2.Sqrt(sum, dist);

        Mat mask = new();
        // pixels with distance <= thresh -> 255 in mask
        Cv2.Threshold(dist, mask, thresh, 255, ThresholdTypes.BinaryInv);

        // ensure mask is CV_8U
        if (mask.Type() != MatType.CV_8U)
        {
            Mat tmp = new();
            mask.ConvertTo(tmp, MatType.CV_8U);
            return tmp;
        }

        return mask;
    }

    // Canny edge based segmentation.
    // Returns a binary edge mask (0/255). Parameters threshold1 and threshold2 are the Canny low/high thresholds.
    // If makeRegions is true, the function will invert edges and apply closing to produce coarse region masks.
    public static Mat CannySegmentation(Mat src, double threshold1 = 100.0, double threshold2 = 200.0,
        bool blur = true, int blurKernel = 5, int apertureSize = 3, bool L2gradient = false,
        bool makeRegions = false, int morphKernel = 3)
    {
        Mat gray = new();
        Cv2.CvtColor(src, gray, ColorConversionCodes.BGR2GRAY);

        if (blur)
        {
            Cv2.GaussianBlur(gray, gray, new Size(blurKernel, blurKernel), 0);
        }

        Mat edges = new();
        Cv2.Canny(gray, edges, threshold1, threshold2, apertureSize, L2gradient);

        if (!makeRegions)
        {
            return edges;
        }

        // Convert edge map to coarse regions by inverting edges and closing gaps
        Mat inv = new();
        Cv2.BitwiseNot(edges, inv);

        Mat kernel = Cv2.GetStructuringElement(MorphShapes.Rect, new Size(morphKernel, morphKernel));
        Mat closed = new();
        Cv2.MorphologyEx(inv, closed, MorphTypes.Close, kernel);

        // Optional: threshold to ensure binary
        Mat mask = new();
        Cv2.Threshold(closed, mask, 128, 255, ThresholdTypes.Binary);
        return mask;
    }

    // Sobel gradient magnitude thresholding segmentation.
    // thresh is the threshold applied to gradient magnitude (0-255 after normalization).
    public static Mat SobelMagnitudeSegmentation(Mat src, double thresh = 50.0, bool blur = true, int blurKernel = 3)
    {
        Mat gray = new();
        Cv2.CvtColor(src, gray, ColorConversionCodes.BGR2GRAY);

        if (blur)
        {
            Cv2.GaussianBlur(gray, gray, new Size(blurKernel, blurKernel), 0);
        }

        Mat gx = new();
        Mat gy = new();
        Cv2.Sobel(gray, gx, MatType.CV_32F, 1, 0, ksize: 3);
        Cv2.Sobel(gray, gy, MatType.CV_32F, 0, 1, ksize: 3);

        Mat mag = new();
        Cv2.Magnitude(gx, gy, mag);

        // Normalize magnitude to 0-255
        Mat mag8u = new();
        Cv2.Normalize(mag, mag, 0, 255, NormTypes.MinMax);
        mag.ConvertTo(mag8u, MatType.CV_8U);

        Mat mask = new();
        Cv2.Threshold(mag8u, mask, thresh, 255, ThresholdTypes.Binary);
        return mask;
    }

    // --- Second-derivative (Laplacian) based methods ---

    // Simple Laplacian magnitude thresholding.
    // ksize: aperture size for the Laplacian operator (1,3,5...). thresh: threshold on normalized magnitude (0-255).
    public static Mat LaplacianSegmentation(Mat src, int ksize = 3, double thresh = 30.0, bool blur = true, int blurKernel = 3)
    {
        Mat gray = new();
        Cv2.CvtColor(src, gray, ColorConversionCodes.BGR2GRAY);

        if (blur)
        {
            Cv2.GaussianBlur(gray, gray, new Size(blurKernel, blurKernel), 0);
        }

        Mat lap = new();
        Cv2.Laplacian(gray, lap, MatType.CV_32F, ksize: ksize);

        Mat absLap = new();
        absLap = Cv2.Abs(lap);

        Mat norm = new();
        Cv2.Normalize(absLap, norm, 0, 255, NormTypes.MinMax);

        Mat mask = new();
        norm.ConvertTo(norm, MatType.CV_8U);
        Cv2.Threshold(norm, mask, thresh, 255, ThresholdTypes.Binary);
        return mask;
    }

    // Laplacian of Gaussian: blur with Gaussian (sigma), then Laplacian, then detect zero-crossings.
    // zeroCrossThreshold is applied to absolute Laplacian magnitude to avoid weak crossings caused by noise.
    public static Mat LoGSegmentation(Mat src, double sigma = 1.4, double zeroCrossThreshold = 5.0, bool useKernelSize = false, int ksize = 0)
    {
        // choose kernel size from sigma if requested
        if (!useKernelSize && ksize == 0)
        {
            ksize = (int)(sigma * 6) | 1; // ensure odd
        }

        Mat gray = new();
        Cv2.CvtColor(src, gray, ColorConversionCodes.BGR2GRAY);

        Mat blurred = new();
        Cv2.GaussianBlur(gray, blurred, new Size(ksize, ksize), sigma);

        Mat lap = new();
        Cv2.Laplacian(blurred, lap, MatType.CV_32F, ksize: 3);

        Mat absLap = new();
        absLap = Cv2.Abs(lap);

        Mat mask = ZeroCrossingMask(absLap, zeroCrossThreshold);
        return mask;
    }

    // Difference of Gaussians as approximation of LoG.
    // sigma1 < sigma2. thresh is applied to absolute DoG response after normalization.
    public static Mat DoGSegmentation(Mat src, double sigma1 = 1.0, double sigma2 = 2.0, double thresh = 10.0, bool normalize = true)
    {
        Mat gray = new();
        Cv2.CvtColor(src, gray, ColorConversionCodes.BGR2GRAY);

        int k1 = (int)(sigma1 * 6) | 1;
        int k2 = (int)(sigma2 * 6) | 1;

        Mat g1 = new();
        Mat g2 = new();
        Cv2.GaussianBlur(gray, g1, new Size(k1, k1), sigma1);
        Cv2.GaussianBlur(gray, g2, new Size(k2, k2), sigma2);

        Mat dog = new();
        Cv2.Subtract(g1, g2, dog);

        Mat absDog = new();
        absDog = Cv2.Abs(dog);

        if (normalize)
        {
            Cv2.Normalize(absDog, absDog, 0, 255, NormTypes.MinMax);
        }

        Mat mask = new();
        absDog.ConvertTo(absDog, MatType.CV_8U);
        Cv2.Threshold(absDog, mask, thresh, 255, ThresholdTypes.Binary);
        return mask;
    }

    // Compute Otsu variance profiles for thresholds 0..254 (we evaluate t as threshold separating [0..t] and [t+1..255]).
    // Returns best thresholds (max between-class variance and min within-class variance) and the per-threshold arrays.
    public static (int bestByBetween, int bestByWithin, double[] betweenVars, double[] withinVars) ComputeOtsuVarianceProfiles(Mat src)
    {
        // ensure grayscale CV_8U
        Mat gray = new();
        if (src.Type() == MatType.CV_8U && src.Channels() == 1)
        {
            src.CopyTo(gray);
        }
        else
        {
            Cv2.CvtColor(src, gray, ColorConversionCodes.BGR2GRAY);
        }

        int rows = gray.Rows;
        int cols = gray.Cols;
        int total = rows * cols;

        // histogram
        double[] hist = new double[256];
        for (int y = 0; y < rows; y++)
        {
            for (int x = 0; x < cols; x++)
            {
                byte v = gray.At<byte>(y, x);
                hist[v] += 1.0;
            }
        }

        // convert to probabilities
        for (int i = 0; i < 256; i++) hist[i] /= total;

        // cumulative arrays
        double[] P = new double[256]; // cumulative probability
        double[] S = new double[256]; // cumulative first moment (i * p(i))
        double[] M2 = new double[256]; // cumulative second moment (i^2 * p(i))

        double pAcc = 0, sAcc = 0, m2Acc = 0;
        for (int i = 0; i < 256; i++)
        {
            pAcc += hist[i];
            sAcc += i * hist[i];
            m2Acc += i * i * hist[i];
            P[i] = pAcc;
            S[i] = sAcc;
            M2[i] = m2Acc;
        }

        double muT = S[255];
        double m2T = M2[255];

        // total variance (scalar)
        double totalVar = m2T - muT * muT;

        int maxIndexBetween = 0;
        int minIndexWithin = 0;
        double maxBetween = double.MinValue;
        double minWithin = double.MaxValue;

        // arrays to hold per-threshold values (we'll fill 256 but only meaningful up to 254)
        double[] betweenVars = new double[256];
        double[] withinVars = new double[256];

        for (int t = 0; t < 255; t++)
        {
            double w0 = P[t];
            double w1 = 1.0 - w0;

            if (w0 <= 0.0 || w1 <= 0.0)
            {
                betweenVars[t] = 0.0;
                withinVars[t] = double.PositiveInfinity;
                continue;
            }

            double s0 = S[t];
            double s1 = S[255] - s0;
            double m20 = M2[t];
            double m21 = m2T - m20;

            double mu0 = s0 / w0;
            double mu1 = s1 / w1;

            // between-class variance
            double sigmaB = w0 * w1 * (mu0 - mu1) * (mu0 - mu1);
            betweenVars[t] = sigmaB;

            // within-class variance via second moments: var0 = M2_0/w0 - mu0^2
            double var0 = m20 / w0 - mu0 * mu0;
            double var1 = m21 / w1 - mu1 * mu1;
            double sigmaW = w0 * var0 + w1 * var1; // weighted within-class variance
            withinVars[t] = sigmaW;

            if (sigmaB > maxBetween)
            {
                maxBetween = sigmaB;
                maxIndexBetween = t;
            }

            if (sigmaW < minWithin)
            {
                minWithin = sigmaW;
                minIndexWithin = t;
            }
        }

        return (maxIndexBetween, minIndexWithin, betweenVars, withinVars);
    }

    // Helper: detect zero-crossings in a floating-point Laplacian image.
    // A pixel is considered edge if any 4-neighbor has opposite sign and the absolute difference exceeds 'threshold'.
    private static Mat ZeroCrossingMask(Mat lapFloat, double threshold)
    {
        // ensure lapFloat is CV_32F
        Mat lap = new();
        if (lapFloat.Type() != MatType.CV_32F)
        {
            lapFloat.ConvertTo(lap, MatType.CV_32F);
        }
        else
        {
            lap = lapFloat.Clone();
        }

        int rows = lap.Rows;
        int cols = lap.Cols;
        Mat mask = new Mat(rows, cols, MatType.CV_8U, Scalar.All(0));

        for (int y = 1; y < rows - 1; y++)
        {
            for (int x = 1; x < cols - 1; x++)
            {
                float v = lap.At<float>(y, x);
                // check 8 neighbors
                bool isZeroCross = false;
                float maxDiff = 0f;
                for (int ny = -1; ny <= 1; ny++)
                {
                    for (int nx = -1; nx <= 1; nx++)
                    {
                        if (ny == 0 && nx == 0) continue;
                        float vn = lap.At<float>(y + ny, x + nx);
                        if ((v > 0 && vn < 0) || (v < 0 && vn > 0))
                        {
                            float diff = MathF.Abs(v - vn);
                            if (diff > maxDiff) maxDiff = diff;
                            isZeroCross = true;
                        }
                    }
                }

                if (isZeroCross && maxDiff >= threshold)
                {
                    mask.Set(y, x, (byte)255);
                }
            }
        }

        return mask;
    }

    private static float Clamp(float v)
    {
        if (v < 0) return 0;
        if (v > 255) return 255;
        return v;
    }
}