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

namespace SDFMapCreator;

public static class Program {
    const float MAX = 1f;
    const float MIN = 0f;
    static readonly int MAX_THREADS = Environment.ProcessorCount - 2;
    static bool debug = 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 readonly SdfKernels kernels = new();
    static uint width;
    static uint height;

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

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

        if (debug) {
            if (!Directory.Exists("Debug")) Directory.CreateDirectory("Debug");
            Console.WriteLine("Debug mode enabled.");
        }

        
        var imagesPath = "images";
        for (var 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);
        }

        width = (uint)Images[0].Width;
        height = (uint)Images[0].Height;
        
        // sum all images together
        if(debug) Images[0].Pixels.SaveImage("Debug/Sum0.png");

        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);
                }
            }

            if(debug)Images[i].Pixels.SaveImage($"Debug/Sum{i}.png");
        }
        

        Console.WriteLine("Creating masks...");
        for (var i = 0; i < Images.Count; i++) { //for each image pair, create a mask
            var selfMask = SelfMask(Images[i].Pixels);
            ImageMasks.Add(new(selfMask, Images[i], new()));
            if (debug) selfMask.SaveImage($"Debug/selfMask{i}.png");
            if (i >= Images.Count - 1) continue;
            ConsoleUpdateLine($"Creating mask {i}...");
            var mask = GetABMask(Images[i].Pixels, Images[i + 1].Pixels);
            TransitionMasks.Add(new(mask, Images[i], Images[i + 1]));
        }

        //EdgeDetect all masks
        Console.WriteLine("\nEdge detecting masks...");
        foreach (var mask in ImageMasks) {
            ConsoleUpdateLine($"Edge detecting mask {ImageMasks.IndexOf(mask)}...");
            EdgeDetect(mask);
        }
        if (debug)
            for (var i = 0; i < TransitionMasks.Count; i++)
                ImageMasks[i].Mask.SaveImage($"Debug/mask{i}.png");

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

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

        // generate final image
        var finalImage = new Vector3[width, height];
        var currStep = 0f;
        var stepIncrement = MAX / Gradients.Count;

        
        for (var x = 0; x < width; x++)
            for (var y = 0; y < height; y++)
                finalImage[x, y] = new(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 (mask.Mask[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("Debug/Final.png");

        // apply directional blur
        var iterations = 1;
        var radius = 100f;
        var step = .5f;
        var sigma = 1f;
        var totalMask = SelfMask(Images[^1].Pixels);
        totalMask.SaveImage("Debug/TotalMask.png");
        for (var i = 0; i < iterations; i++) {
            Console.WriteLine($"Applying directional blur {i + 1}/{iterations}...");
            finalImage = DirectionalBlur(finalImage, totalMask, radius, step, sigma);
        }

        finalImage.SaveImage("finalBlur.png");
        Console.WriteLine("Done!");
    }

    static void EdgeDetect(MaskData maskData) {
        var sw = new Stopwatch();
        sw.Start();
        Console.WriteLine("Running edge detection...");
        kernels.EdgeDetect(maskData.Mask, out var temp);
        maskData.Mask = temp;
        Parallel.For(0, width * height, parallelOptions, (i) => {
            if (maskData.Mask[i % width, i / width].Y == 0) return;
            // ReSharper disable once PossibleLossOfFraction
            lock (maskData.Edges) maskData.Edges.Add(new(i % width, i / width));
        });

        sw.Stop();
        Console.WriteLine(
            $"\nEdge pixels: {maskData.Edges.Count} | {width * height / (sw.ElapsedMilliseconds + 1)} pixels/s\n Time: {sw.Elapsed.TotalSeconds:F4}s");
    }

    static Vector3[,] Gradient(TransitionMaskData mask, SDFData sdfA, SDFData sdfB) {
        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;
        });*/

        kernels.Gradient(mask.Mask, sdfA.SDF, sdfB.SDF, out var temp);
        Console.WriteLine($"Gradient Time: {sw.Elapsed.TotalSeconds:N4}s ({width * height / (float)sw.Elapsed.TotalSeconds:N0} pixels/s)");
        return temp;
    }

    static SDFData SDF(MaskData mask) {
        var sw = new Stopwatch();
        sw.Start();
        kernels.Sdf(mask.Edges.ToArray(), (int)width, (int)height, out var temp);
        Console.WriteLine($"\nSDF Generation Time: {sw.Elapsed.TotalSeconds:N4}s ({width * height / sw.Elapsed.TotalSeconds:N0} pixels/s)");

        sw.Restart();
        var absMax = 0f;
        foreach (var pixel in temp) {
            if (pixel.X > absMax) absMax = pixel.X;
        }

        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 ({width * height / sw.Elapsed.TotalSeconds:N0} pixels/s)");

        return new(temp);
    }

    static Vector3[,] DirectionalBlur(Vector3[,] image, Vector3[,] mask, float radius = 3f, float step = .5f, float sigma = 1f) {
        var sw = new Stopwatch();
        sw.Start();
        kernels.DirectionalBlur(image, mask, out var temp, radius, step, sigma);
        Console.WriteLine(
            $"Directional Blur Time: {sw.Elapsed.TotalSeconds:N4}s ({width * height / sw.Elapsed.TotalSeconds:N0} pixels/s)");
        return temp;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    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) {
        kernels.ABMask(A, B, out var temp);
        return temp;
    }

    static Vector3[,] SelfMask(Vector3[,] A) {
        kernels.SelfMask(A, out var temp);
        return temp;
    }

    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 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;
}