using SkiaSharp;

namespace Scry.Core.Imaging;

public static class ImagePreprocessor
{
    public static SKBitmap ApplyClahe(SKBitmap source, int tileSize = 8, float clipLimit = 2.0f)
    {
        var width = source.Width;
        var height = source.Height;
        var result = new SKBitmap(width, height, SKColorType.Rgba8888, SKAlphaType.Premul);
        
        var labImage = ConvertToLab(source);
        
        var tilesX = (width + tileSize - 1) / tileSize;
        var tilesY = (height + tileSize - 1) / tileSize;
        
        var tileMappings = new byte[tilesY, tilesX, 256];
        
        for (var ty = 0; ty < tilesY; ty++)
        {
            for (var tx = 0; tx < tilesX; tx++)
            {
                var startX = tx * tileSize;
                var startY = ty * tileSize;
                var endX = Math.Min(startX + tileSize, width);
                var endY = Math.Min(startY + tileSize, height);
                
                var histogram = new int[256];
                var pixelCount = 0;
                
                for (var y = startY; y < endY; y++)
                {
                    for (var x = startX; x < endX; x++)
                    {
                        var l = (int)labImage[y, x, 0];
                        histogram[l]++;
                        pixelCount++;
                    }
                }
                
                var clipCount = (int)(clipLimit * pixelCount / 256);
                var excess = 0;
                
                for (var i = 0; i < 256; i++)
                {
                    if (histogram[i] > clipCount)
                    {
                        excess += histogram[i] - clipCount;
                        histogram[i] = clipCount;
                    }
                }
                
                var redistribution = excess / 256;
                for (var i = 0; i < 256; i++)
                {
                    histogram[i] += redistribution;
                }
                
                var cdf = new int[256];
                cdf[0] = histogram[0];
                for (var i = 1; i < 256; i++)
                {
                    cdf[i] = cdf[i - 1] + histogram[i];
                }
                
                var cdfMin = cdf.FirstOrDefault(c => c > 0);
                if (cdfMin == 0) cdfMin = 1;
                var denominator = Math.Max(1, pixelCount - cdfMin);
                
                for (var i = 0; i < 256; i++)
                {
                    tileMappings[ty, tx, i] = (byte)Math.Clamp((cdf[i] - cdfMin) * 255 / denominator, 0, 255);
                }
            }
        }
        
        for (var y = 0; y < height; y++)
        {
            for (var x = 0; x < width; x++)
            {
                var l = (int)labImage[y, x, 0];
                var a = labImage[y, x, 1];
                var b = labImage[y, x, 2];
                
                var txFloat = (float)x / tileSize - 0.5f;
                var tyFloat = (float)y / tileSize - 0.5f;
                
                var tx1 = Math.Clamp((int)Math.Floor(txFloat), 0, tilesX - 1);
                var ty1 = Math.Clamp((int)Math.Floor(tyFloat), 0, tilesY - 1);
                var tx2 = Math.Clamp(tx1 + 1, 0, tilesX - 1);
                var ty2 = Math.Clamp(ty1 + 1, 0, tilesY - 1);
                
                var xRatio = txFloat - tx1;
                var yRatio = tyFloat - ty1;
                xRatio = Math.Clamp(xRatio, 0, 1);
                yRatio = Math.Clamp(yRatio, 0, 1);
                
                var v1 = tileMappings[ty1, tx1, l];
                var v2 = tileMappings[ty1, tx2, l];
                var v3 = tileMappings[ty2, tx1, l];
                var v4 = tileMappings[ty2, tx2, l];
                
                var top = v1 * (1 - xRatio) + v2 * xRatio;
                var bottom = v3 * (1 - xRatio) + v4 * xRatio;
                var newL = top * (1 - yRatio) + bottom * yRatio;
                
                var (r, g, bl) = LabToRgb(newL, a, b);
                result.SetPixel(x, y, new SKColor(r, g, bl));
            }
        }
        
        return result;
    }
    
    private static float[,,] ConvertToLab(SKBitmap bitmap)
    {
        var result = new float[bitmap.Height, bitmap.Width, 3];
        
        for (var y = 0; y < bitmap.Height; y++)
        {
            for (var x = 0; x < bitmap.Width; x++)
            {
                var pixel = bitmap.GetPixel(x, y);
                var (l, a, b) = RgbToLab(pixel.Red, pixel.Green, pixel.Blue);
                result[y, x, 0] = l;
                result[y, x, 1] = a;
                result[y, x, 2] = b;
            }
        }
        
        return result;
    }
    
    private static (float L, float A, float B) RgbToLab(byte r, byte g, byte b)
    {
        var rf = PivotRgb(r / 255.0);
        var gf = PivotRgb(g / 255.0);
        var bf = PivotRgb(b / 255.0);
        
        var x = rf * 0.4124564 + gf * 0.3575761 + bf * 0.1804375;
        var y = rf * 0.2126729 + gf * 0.7151522 + bf * 0.0721750;
        var z = rf * 0.0193339 + gf * 0.1191920 + bf * 0.9503041;
        
        x /= 0.95047;
        z /= 1.08883;
        
        x = PivotXyz(x);
        y = PivotXyz(y);
        z = PivotXyz(z);
        
        var l = (float)(116 * y - 16);
        var a = (float)(500 * (x - y));
        var bVal = (float)(200 * (y - z));
        
        l = l * 255 / 100;
        
        return (l, a, bVal);
    }
    
    private static (byte R, byte G, byte B) LabToRgb(double l, double a, double b)
    {
        l = l * 100 / 255;
        
        var y = (l + 16) / 116;
        var x = a / 500 + y;
        var z = y - b / 200;
        
        var x3 = x * x * x;
        var y3 = y * y * y;
        var z3 = z * z * z;
        
        x = x3 > 0.008856 ? x3 : (x - 16.0 / 116) / 7.787;
        y = y3 > 0.008856 ? y3 : (y - 16.0 / 116) / 7.787;
        z = z3 > 0.008856 ? z3 : (z - 16.0 / 116) / 7.787;
        
        x *= 0.95047;
        z *= 1.08883;
        
        var rf = x * 3.2404542 + y * -1.5371385 + z * -0.4985314;
        var gf = x * -0.9692660 + y * 1.8760108 + z * 0.0415560;
        var bf = x * 0.0556434 + y * -0.2040259 + z * 1.0572252;
        
        rf = rf > 0.0031308 ? 1.055 * Math.Pow(rf, 1 / 2.4) - 0.055 : 12.92 * rf;
        gf = gf > 0.0031308 ? 1.055 * Math.Pow(gf, 1 / 2.4) - 0.055 : 12.92 * gf;
        bf = bf > 0.0031308 ? 1.055 * Math.Pow(bf, 1 / 2.4) - 0.055 : 12.92 * bf;
        
        return (
            (byte)Math.Clamp(rf * 255, 0, 255),
            (byte)Math.Clamp(gf * 255, 0, 255),
            (byte)Math.Clamp(bf * 255, 0, 255)
        );
    }
    
    private static double PivotRgb(double n)
    {
        return n > 0.04045 ? Math.Pow((n + 0.055) / 1.055, 2.4) : n / 12.92;
    }
    
    private static double PivotXyz(double n)
    {
        return n > 0.008856 ? Math.Pow(n, 1.0 / 3) : 7.787 * n + 16.0 / 116;
    }
    
    public static SKBitmap Resize(SKBitmap source, int width, int height)
    {
        var info = new SKImageInfo(width, height, SKColorType.Rgba8888, SKAlphaType.Premul);
        var sampling = new SKSamplingOptions(SKFilterMode.Linear, SKMipmapMode.Linear);
        return source.Resize(info, sampling);
    }
    
    public static SKBitmap ToGrayscale(SKBitmap source)
    {
        var result = new SKBitmap(source.Width, source.Height, SKColorType.Gray8, SKAlphaType.Opaque);
        
        for (var y = 0; y < source.Height; y++)
        {
            for (var x = 0; x < source.Width; x++)
            {
                var pixel = source.GetPixel(x, y);
                var gray = (byte)(0.299 * pixel.Red + 0.587 * pixel.Green + 0.114 * pixel.Blue);
                result.SetPixel(x, y, new SKColor(gray, gray, gray));
            }
        }
        
        return result;
    }
}
