using System.Diagnostics;
using System.Numerics;
using System.Runtime.CompilerServices;
using SDFMapCreator;

public class Program {
    private const float MAX = 1f;
    private const float MIN = 0f;
    private static readonly int MAX_THREADS = Environment.ProcessorCount - 2;
    private const bool outputMasks = true;
    private const bool outputSDFs = true;
    private const bool outputGradients = true;
    static readonly ParallelOptions parallelOptions = new() { MaxDegreeOfParallelism = MAX_THREADS };
    static List<Image> Images = new();
    static List<MaskData> ImageMasks = new();
    static List<TransitionMaskData> TransitionMasks = new();
    static List<SDFData> SDFs = new();
    static List<Vector3[,]> Gradients = new();

    static void ConsoleUpdateLine(string s) => Console.Write("\r" + s);

    public static void Main(string[] args) {
        Console.WriteLine("Reading images...");
        
        var imagesPath = "images";

        
        for (int i = 0; i < 8; i++) {
            var pixels = ImageUtil.LoadImage<Vector3>($"./{imagesPath}{Path.DirectorySeparatorChar}{i + 1:00}.png");
            Images.Add(new(pixels, pixels.GetLength(0), pixels.GetLength(1)));
        }
        
        /*
        var pixels = ImageUtil.LoadImage<Vector3>($"./sphereempty.png");
        Images.Add(new(pixels, pixels.GetLength(0), pixels.GetLength(1)));
        pixels = ImageUtil.LoadImage<Vector3>($"./spherehalf.png");
        Images.Add(new(pixels, pixels.GetLength(0), pixels.GetLength(1)));
        pixels = ImageUtil.LoadImage<Vector3>($"./spherecut.png");
        Images.Add(new(pixels, pixels.GetLength(0), pixels.GetLength(1)));
        pixels = ImageUtil.LoadImage<Vector3>($"./spherefull.png");
        Images.Add(new(pixels, pixels.GetLength(0), pixels.GetLength(1)));
        */
        
        //check if all the images in Images are the same resolution
        if (Images.Select(img => (img.Width, img.Height)).Distinct().Count() > 1) {
            Console.WriteLine("Error: Not all images have the same resolution.");
            Environment.Exit(1);
        }

        for (int i = 1; i < Images.Count; i++)
        {
            for (int x = 0; x < Images[i].Width; x++)
            {
                for (int y = 0; y < Images[i].Height; y++)
                {
                    Images[i].Pixels[x, y].X = MathF.Min(Images[i - 1].Pixels[x, y].X + Images[i].Pixels[x, y].X, MAX);
                    Images[i].Pixels[x, y].Y = MathF.Min(Images[i - 1].Pixels[x, y].Y + Images[i].Pixels[x, y].X, MAX);
                    Images[i].Pixels[x, y].Z = MathF.Min(Images[i - 1].Pixels[x, y].Z + Images[i].Pixels[x, y].X, MAX);
                }
            }
            
            Images[i].Pixels.SaveImage($"./Sum{i}.png");
        }

        Console.WriteLine("Creating masks...");
        //for each image pair, create a mask
        var width = (uint)Images[0].Width;
        var height = (uint)Images[0].Height;
        for (int i = 0; i < Images.Count; i++) {
            ImageMasks.Add(new(SelfMask(Images[i].Pixels, width, height), Images[i], new()));
            if (i < Images.Count - 1) {
                Console.WriteLine($"Creating mask {i}...");
                var mask = GetABMask(Images[i].Pixels, Images[i + 1].Pixels, width, height);
                TransitionMasks.Add(new(mask, Images[i], Images[i + 1]));
            }
        }

        Console.WriteLine("Edge detecting masks...");
        //EdgeDetect all masks
        foreach (var t in ImageMasks) EdgeDetect(t);

        if (outputMasks) {
            Console.WriteLine("Writing masks...");
            for (int i = 0; i < TransitionMasks.Count; i++) ImageMasks[i].Mask.SaveImage($"mask{i}.png");
        }

        Console.WriteLine("Creating SDFs...");
        for (var i = 0; i < ImageMasks.Count; i++) {
            var mask = ImageMasks[i];
            SDFs.Add(SDF(mask));
            if (outputSDFs) SDFs[i].SDF.SaveImage($"sdf{i}.png");
        }

        Console.WriteLine("Creating gradients...");
        for (var i = 0; i < TransitionMasks.Count; i++) {
            Console.WriteLine($"Generating gradient {i}...");
            var gradientData = Gradient(TransitionMasks[i], SDFs[i], SDFs[i + 1]);
            Gradients.Add(gradientData);
            if (outputGradients) gradientData.SaveImage($"gradient{i}.png");
        }

        // generate final image
        var finalImage = new Vector3[width, height];
        var currStep = 0f;
        var stepIncrement = MAX / (Gradients.Count);
        
        for(int x=0; x<width; x++)
        for (int y = 0; y < height; y++)
            finalImage[x, y] = new Vector3(Images[0].Pixels[x, y].X);
        
        for (var i = 0; i < Gradients.Count; i++) {
            var mask = ImageMasks[i + 1];
            var gradient = Gradients[i];
            Console.WriteLine($"Applying gradient {i}..., {currStep} -> {currStep + stepIncrement}");
            for (var x = 0; x < mask.Mask.GetLength(0); x++) {
                for (var y = 0; y < mask.Mask.GetLength(1); y++) {
                    if(gradient[x, y].X == 0) continue;
                    var pixel = new Vector3(Remap(gradient[x, y].X, MIN, MAX, 1.0f-currStep, 1.0f-(currStep + stepIncrement)));
                    if(pixel.X > finalImage[x, y].X) finalImage[x, y] = pixel;
                }
            }
            currStep += stepIncrement;
        }
        
        finalImage.SaveImage("final.png");

        
        // apply directional blur
        var iterations = 10;
        var finalImageMask = new Vector3[width, height];
        finalImageMask = SelfMask(Images[^1].Pixels, width, height);
        
        for (int i = 0; i < iterations; i++){
            Console.WriteLine($"Applying directional blur {i}...");
            var radius = 2f;
            finalImage = DirectionalBlur(finalImage, finalImageMask, radius, 0.5f, 1f);
        }
        
        finalImage.SaveImage("finalBlur.png");
        Console.WriteLine("Done!");
    }

    private static void EdgeDetect(MaskData maskData) {
        uint width = (uint)maskData.Image.Width;
        uint height = (uint)maskData.Image.Height;
        int iterCount = 0;
        var sw = new Stopwatch();
        sw.Start();
        Console.WriteLine("Running edge detection...");
        Parallel.For(0, width * height, parallelOptions, (i) =>
        {
            int x = (int)(i % width);
            int y = (int)(i / width);

            if (!EdgeKernel(maskData.Mask, x, y, width, height)) return;
            var color = maskData.Mask[x, y];
            color.Y = MAX;
            maskData.Mask[x, y] = color;
            lock (maskData.Edges) maskData.Edges.Add(new(x, y));
            iterCount++;
            if (iterCount % (width * height / 100) == 0) {
                ConsoleUpdateLine(
                    $"Progress: {iterCount / (float)(width * height):P} | {iterCount / (sw.Elapsed.TotalSeconds):N0} pixels/s");
            }
        });
        sw.Stop();
        Console.WriteLine(
            $"\nEdge pixels: {maskData.Edges.Count} | {maskData.Edges.Count / (sw.ElapsedMilliseconds+1)} pixels/s\n Time: {sw.Elapsed.TotalSeconds:F4}s");
    }

    static Vector3[,] Gradient(TransitionMaskData mask, SDFData sdfA, SDFData sdfB) {
        uint width = (uint)mask.ImageA.Width;
        uint height = (uint)mask.ImageA.Height;
        int iterCount = 0;
        var sw = new Stopwatch();
        sw.Start();
        Vector3[,] temp = new Vector3[width, height];

        var min = MAX;
        var max = MIN;

        Console.WriteLine("Running gradient generation...");
        Parallel.For(0, width * height, parallelOptions, (i) =>
        {
            int x = (int)(i % width);
            int y = (int)(i / width);

            if (mask.Mask[x, y].X == 0 || mask.Mask[x, y].Y > 0) return;

            var a = sdfA.SDF[x, y].X;
            var b = sdfB.SDF[x, y].X;

            var gradient = a / (a + b);
            temp[x, y] = new(Remap(gradient, 0, 1, MIN, MAX));

            if (gradient < min) min = gradient;
            if (gradient > max) max = gradient;
        });

        Console.WriteLine(
            $"Gradient Time: {sw.Elapsed.TotalSeconds:N4}s ({iterCount / (float)sw.Elapsed.TotalSeconds:N0} pixels/s)");
        Console.WriteLine($"Min: {min} | Max: {max}");
        return temp;
    }

    static SDFData SDF(MaskData mask) {
        var width = (uint)mask.Mask.GetLength(0);
        var height = (uint)mask.Mask.GetLength(1);
        var temp = new Vector3[width, height];
        float AbsMax = MIN;
        int iterCount = 0;
        var sw = new Stopwatch();
        sw.Start();
        Parallel.For(0, width * height, parallelOptions, (i) =>
        {
            //convert 1D index to 2D index
            var x = (int)(i % width);
            var y = (int)(i / width);
            Vector2 p = new(x/(float)width, y/(float)height); //get the pixel position as a Vector2

            float minDist = MAX + 1; //initialize the minimum distance to the maximum possible value

            //loop through all the pixels in the mask
            foreach (var edge in mask.Edges) {
                Vector2 edgeNorm = new(edge.X / (float)width, edge.Y / (float)height);
                float dist = Vector2.DistanceSquared(p, edgeNorm);
                if (dist < minDist) minDist = dist;
            }

            if (minDist > MAX)
                minDist = MAX;

            temp[x, y] = new(minDist);
            if (minDist > AbsMax) AbsMax = minDist;
            iterCount++;
            if (iterCount % (width * height / 100) == 0) {
                ConsoleUpdateLine(
                    $"Progress: {iterCount / (float)(width * height):P}% | {iterCount / (sw.Elapsed.TotalSeconds):N0} pixels/s");
            }
        });

        Console.WriteLine(
            $"\nSDF Generation Time: {sw.Elapsed.TotalSeconds:N4}s ({iterCount / sw.Elapsed.TotalSeconds:N0} pixels/s)");
        sw.Restart();
        Parallel.For(0, width * height, parallelOptions, (i) =>
        {
            //convert 1D index to 2D index
            var x = (int)(i % width);
            var y = (int)(i / width);
            temp[x, y] = new(Remap(temp[x, y].X, 0, AbsMax, MIN, MAX));
        });
        sw.Stop();

        Console.WriteLine(
            $"SDF Normalization Time: {sw.Elapsed.TotalSeconds:N4}s ({iterCount / sw.Elapsed.TotalSeconds:N0} pixels/s)");
        Console.WriteLine("AbsMax: " + AbsMax);
        return new(temp);
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static float EuclideanDistance(Vector2 a, Vector2 b) =>
        MathF.Sqrt(MathF.Pow(a.X - b.X, 2) + MathF.Pow(a.Y - b.Y, 2));
    
    static Vector3[,] GetABMask(Vector3[,] A, Vector3[,] B, uint resX, uint resY) {
        var temp = new Vector3[resX, resY];
        Parallel.For(0, resX * resY, parallelOptions, (i) =>
        {
            uint x = (uint)(i % resX);
            uint y = (uint)(i / resX);
            var pixelA = A[x, y];
            var pixelB = B[x, y];
            float lumaA = (pixelA.X + pixelA.Y + pixelA.Z) / 3;
            float lumaB = (pixelB.X + pixelB.Y + pixelB.Z) / 3;
            float resultPixel = lumaB > lumaA ? MAX : MIN;
            temp[x, y] = new(resultPixel, 0, 0);
        });
        return temp;
    }

    static Vector3[,] SelfMask(Vector3[,] A, uint resX, uint resY) {
        var temp = new Vector3[resX, resY];
        Parallel.For(0, resX * resY, parallelOptions, (i) =>
        {
            uint x = (uint)(i % resX);
            uint y = (uint)(i / resX);
            var pixelA = A[x, y];
            float lumaA = (pixelA.X + pixelA.Y + pixelA.Z) / 3;
            float resultPixel = lumaA > 0.99 ? MAX : MIN;
            temp[x, y] = new(resultPixel, 0, 0);
        });
        return temp;
    }

    static bool EdgeKernel(Vector3[,] mask, int x, int y, uint width, uint height) {
        //if we are already empty, return false
        if (mask[x, y].X == 0) return false;
        //if we are on the edge of the image, return false
        if (x == 0 || y == 0 || x == width - 1 || y == height - 1) return false;
        //check the 3x3 kernel
        for (int xi = x - 1; xi <= x + 1; xi++) {
            for (int yi = y - 1; yi <= y + 1; yi++) {
                if (xi < 0 || xi >= width || yi < 0 || yi >= height)
                    continue; //skip out of bounds pixels
                if (mask[xi, yi].X == 0)
                    return true; //if we find a black pixel, return true
            }
        }
        //if we didn't find any black pixels, return false
        return false;
    }

    static Vector3 SampleBilinear(Vector3[,] image, Vector2 pixelCoord) {
        var width = (uint)image.GetLength(0);
        var height = (uint)image.GetLength(1);

        var x = pixelCoord.X;
        var y = pixelCoord.Y;
        
        var x0 = (int)x;
        var y0 = (int)y;
        
        var x1 = Math.Min(x0 + 1, width - 1);
        var y1 = Math.Min(y0 + 1, height - 1);
        
        var dx = x - x0;
        var dy = y - y0;
        
        return Lerp(
            Lerp(image[x0, y0], image[x1, y0], dx),
            Lerp(image[x0, y1], image[x1, y1], dx),
            dy
        );
    }
    
    static Vector3[,] DirectionalBlur(Vector3[,] image, Vector3[,] mask, float radius = 3f, float step = .5f, float sigma = 1f)
    {
        var width = (uint)image.GetLength(0);
        var height = (uint)image.GetLength(1);
        var output = new Vector3[width, height];
        
        Parallel.For(0, (int)width * height, parallelOptions, (i) =>
        {
            var x = i % width;
            var y = i / width;
            var gradient = new Vector2(0, 0);
            
            if(mask[x, y].X == 0)
            {
                output[x, y] = image[x, y];
                return;
            }
            
            for (int dx = -1; dx <= 1; dx++)
            {
                for (int dy = -1; dy <= 1; dy++)
                {
                    if (x + dx < 0 || x + dx >= width || y + dy < 0 || y + dy >= height) continue;
                    gradient += new Vector2(dx, dy) * image[x + dx, y + dy].X;
                }
            }

            if (gradient == Vector2.Zero)
            {
                output[x, y] = image[x, y];
                return;
            }

            gradient = Vector2.Normalize(gradient);
            float sum = 0;
            
            // now we follow the direction line and sample the image for length;
            for(float l = -radius; l <= radius; l += step)
            {
                var pixel = new Vector2(x, y) + gradient * l;
                // check if the pixel is out of bounds
                if (pixel.X < 0 || pixel.X >= width || pixel.Y < 0 || pixel.Y >= height) continue;
                var sample = SampleBilinear(image, pixel);
                float weight = (float)Math.Exp(-(l * l) / (2 * sigma * sigma));
                output[x, y] += sample * weight;
                sum += weight;
            }
            
            output[x, y] /= sum;
        });
        
        return output;
    }
    
    static T Lerp<T>(T a, T b, float t)
        where T : INumber<T>, IMultiplyOperators<T, float, T>, IAdditionOperators<T, T, T>
        => a * (1 - t) + b * t;
    
    static Vector3 Lerp(Vector3 a, Vector3 b, float t) => a * (1 - t) + b * t;

    static T Remap<T>(T value, T min, T max, T newMin, T newMax)
        where T : INumber<T>, ISubtractionOperators<T, T, T>, IMultiplyOperators<T, T, T>, IAdditionOperators<T, T, T>
        => (value - min) / (max - min) * (newMax - newMin) + newMin;
}