using System.Diagnostics;
using System.Numerics;
using System.Runtime.CompilerServices;
using ImageMagick;
using float2 = System.Numerics.Vector2;
using float3 = System.Numerics.Vector3;
using float4 = System.Numerics.Vector4;

/*
public struct float2(float x, float y) {
    public float x = x, y = y;
}

public struct float3(float r, float g, float b) {
    public float r = r, g = g, b = b;
    public float3(float value) : this(value, value, value) {}
    public float3(float r, float g) : this(r, g, 0f) {}
}

public struct float4(float r, float g, float b, float a) {
    public float r = r, g = g, b = b, a = a;
    public float4(float value) : this(value, value, value, value) {}
    public float4(float r, float g, float b) : this(r, g, b, 1f) {}
}*/

public record ImageData(MagickImage Image, float3[,] Pixels, List<float2> Edges);
public record MaskData(float3[,] Mask, ImageData Image, List<float2> Edges);
public record TransitionMaskData(float3[,] Mask, ImageData ImageA, ImageData ImageB);

public record SDFData(float3[,] SDF);

public static class ArrayExt {
    public static float3[,] To2DFloat3(this float[] array, uint width, uint height) {
        float3[,] result = new float3[width, height];
        for(int i = 0; i < width*height; i++) {
            uint x = (uint)(i % width);
            uint y = (uint)(i / width);
            result[y, x] = new (array[i*3], array[i*3+1], array[i*3+2]);
        }
        return result;
    }
    
    public static float[] ToFloatArray(this float3[,] array) {
        float[] result = new float[array.GetLength(0) * array.GetLength(1) * 3];
        for(int x = 0; x < array.GetLength(0); x++) {
            for(int y = 0; y < array.GetLength(1); y++) {
                result[x*array.GetLength(1)*3 + y*3] = array[x, y].X;
                result[x*array.GetLength(1)*3 + y*3+1] = array[x, y].Y;
                result[x*array.GetLength(1)*3 + y*3+2] = array[x, y].Z;
            }
        }
        return result;
    }
}

public class  Program {
    private const float MAX = 65535f;
    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<ImageData> Images = new();
    static List<TransitionMaskData> TransitionMasks = new();
    static List<MaskData> ImageMasks = new();
    static List<SDFData> SDFs = new();
    static List<float3[,]> Gradients = new();
    
    static void ConsoleUpdateLine(string s) => Console.Write("\r"+s);
    
    static void LoadImage(string imgPath) {
        var image = new MagickImage(imgPath);
        float3[,] pixels;
        if(image.Channels.Count() ==4)
            //skip ever 4th value if we have an alpha channel
            pixels = image.GetPixels().ToArray()!
                .Where((_, i) => (i + 1) % 4 != 0).ToArray()
                .To2DFloat3(image.Width, image.Height);
        else
            pixels = image.GetPixels().ToArray()!
                .To2DFloat3(image.Width, image.Height);
        Images.Add(new (image, pixels, new()));
        Console.WriteLine($"Loaded image: {imgPath}");
        ImageData(image, image.GetPixels());
    }
    
    public static void Main(string[] args) {
        Console.WriteLine("Reading images...");
        //foreach image in arguments load the image
        
        var imagesPath = "images";

        
        for (int i = 0; i < 8; i++)
            LoadImage($"./{imagesPath}{Path.DirectorySeparatorChar}{i+1:00}.png");
        
        
        //LoadImage("spherecut.png");
        //LoadImage("spherefull.png");
        
        //LoadImage("1.png");
        //LoadImage("2.png");
        //check if all the images in Images are the same resolution
        if (Images.Select(img => (img.Image.Width, img.Image.Height)).Distinct().Count() > 1) {
            Console.WriteLine("Error: Not all images have the same resolution.");
            Environment.Exit(1);
        }
        
        Console.WriteLine("Creating masks...");
        //for each image pair, create a mask
        var width = Images[0].Image.Width;
        var height = Images[0].Image.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++) {
                var mask = new MagickImage(MagickColors.Black, (uint)TransitionMasks[i].Mask.GetLength(0), (uint)TransitionMasks[i].Mask.GetLength(1));
                mask.GetPixels().SetPixels(TransitionMasks[i].Mask.ToFloatArray());
                mask.Write($"mask{i}.png", MagickFormat.Png24);
            }
        }
        
        Console.WriteLine("Creating SDFs...");
        for (var i = 0; i < ImageMasks.Count; i++) {
            var mask = ImageMasks[i];
            SDFs.Add(SDF(mask));
            if(!outputSDFs) continue;
            var sdf = new MagickImage(MagickColors.Black, (uint)mask.Mask.GetLength(0), (uint)mask.Mask.GetLength(1));
            sdf.GetPixels().SetPixels(SDFs[i].SDF.ToFloatArray());
            sdf.Write($"sdf{i}.png", MagickFormat.Png48);
        }
        
        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) continue;
            var gradient = new MagickImage(MagickColors.Black, (uint)TransitionMasks[i].Mask.GetLength(0), (uint)TransitionMasks[i].Mask.GetLength(1));
            gradient.GetPixels().SetPixels(gradientData.ToFloatArray());
            gradient.Write($"gradient{i}.png", MagickFormat.Png24);
        }
        
        // generate final image
        var finalImage = new float3[width, height];
        var currStep = 0f;
        var stepIncrement = MAX / (Gradients.Count + 1);
        for (var i = 0; i < Gradients.Count; i++) {
            var mask = ImageMasks[i + 1];
            var gradient = Gradients[i];
            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 || mask.Mask[x, y].Y > 0) continue;
                    finalImage[x, y] = new(Remap(gradient[x,y].X, MAX, MIN, currStep, currStep + stepIncrement));
                }
            }
            currStep += stepIncrement;
        }
        
        var finalImageMagick = new MagickImage(MagickColors.Black, (uint)width, (uint)height);
        finalImageMagick.GetPixels().SetPixels(finalImage.ToFloatArray());
        finalImageMagick.Write("final.png", MagickFormat.Png24);
        
        Console.WriteLine("Done!");
    }

    private static void EdgeDetect(MaskData maskData) {
        uint width = maskData.Image.Image.Width;
        uint height = maskData.Image.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($"Edge pixels: {maskData.Edges.Count} | {maskData.Edges.Count/sw.ElapsedMilliseconds} pixels/s\n Time: {sw.Elapsed.TotalSeconds:F4}s");
    }

    static float3[,] Gradient(TransitionMaskData mask, SDFData sdfA, SDFData sdfB) {
        uint width = mask.ImageA.Image.Width;
        uint height = mask.ImageA.Image.Height;
        int iterCount = 0;
        var sw = new Stopwatch();
        sw.Start();
        float3[,] temp = new float3[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 float3[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);
            float2 p = new(x, y);
            
            float minDist = MAX; //initialize the minimum distance to the maximum possible value
            
            //loop through all the pixels in the mask
            foreach (var edge in mask.Edges) {
                float dist = float2.DistanceSquared(p, edge);
                if (dist < minDist) minDist = mask.Mask[x, y].X == 0 ? dist : dist;
            }
            
            temp[x, y] = new(minDist);
            if (minDist > AbsMax) AbsMax = minDist;
            iterCount++;
            if (iterCount % (width * height / 100) == 0) {
                ConsoleUpdateLine($"Progress: {iterCount/(width*height):P}% | {iterCount/(sw.Elapsed.TotalSeconds):N0} pixels/s");
            }
        });

        Console.WriteLine($"SDF 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(float2 a, float2 b) =>
        MathF.Sqrt(MathF.Pow(a.X - b.X, 2) + MathF.Pow(a.Y - b.Y, 2));
    
    private static void ImageData(MagickImage image1, IPixelCollection<float> pixels1) {
        Console.WriteLine(
            $"""
             Image file: {image1.Format.ToString()}
             Resolution: {image1.Width}x{image1.Height}
             Total Pixels: {pixels1.Count()} |{pixels1.Channels} channels, {image1.Depth} bits per channel
             
             """);
    }
    
    static float3[,] GetABMask(float3[,] A, float3[,] B, uint resX, uint resY) {
        var temp = new float3[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 float3[,] SelfMask(float3[,] A, uint resX, uint resY) {
        var temp = new float3[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 ? MAX : MIN;
            temp[x, y] = new(resultPixel, 0, 0);
        });
        return temp;
    }

    static bool EdgeKernel(float3[,] 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 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;
}
