using System.Diagnostics;
using System.IO;
using System.Linq.Expressions;
using System.Numerics;
using System.Runtime.InteropServices;
using System.Text;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Data;
using System.Windows.Documents;
using System.Windows.Input;
using System.Windows.Media;
using System.Windows.Media.Imaging;
using System.Windows.Media.Media3D;
using System.Windows.Navigation;
using System.Windows.Shapes;
using MathNet.Numerics;
using MathNet.Numerics.IntegralTransforms;
using Microsoft.Win32;
using static DSPS_Lab3.Ex;

namespace DSPS_Lab3;
/// <summary>
/// Interaction logic for MainWindow.xaml
/// </summary>
public partial class MainWindow : System.Windows.Window
{
    BitmapImage source;
    WriteableBitmap _spatial, _fft;
    CancellationTokenSource _cts;
    string path = "";

    public MainWindow()
    {
        _cts = new CancellationTokenSource();
        InitializeComponent();
        MathNet.Numerics.Control.UseNativeMKL();
    }

    private void ChangeParams(object sender, RoutedPropertyChangedEventArgs<double> e)
    {
        if (this.sigmaValue == null) return;
        int kernelSize = 0, imgWnew = 0, imgHnew = 0, imgW = 0, imgH = 0;
        float sigmaValue = 0;

        Application.Current.Dispatcher.Invoke(() => {
            kernelSize = (int)this.kernelSize.Value;
            sigmaValue = (float)this.sigmaValue.Value;
        });

        DrawGaussianGraph(kernelSize, sigmaValue);

        if (source == null) return;

        _cts.Cancel();
        _cts = new CancellationTokenSource();

        _spatial = new WriteableBitmap(source);
        _fft = new WriteableBitmap(source);
        SlowImage.Source = _spatial;
        FastImage.Source = _fft;
        Application.Current.Dispatcher.Invoke(() => {
            imgWnew = _fft.PixelWidth;
            imgHnew = _fft.PixelHeight;
            imgW = source.PixelWidth;
            imgH = source.PixelHeight;
        });

        Stopwatch timer1 = Stopwatch.StartNew(), timer2 = Stopwatch.StartNew();
        Task.Run(() => SpatialBlur(GenerateSpatialKernel(kernelSize, sigmaValue), _spatial, _cts.Token), _cts.Token)
            .ContinueWith(t => {
                Application.Current.Dispatcher.Invoke(() => {
                    timer1.Stop();
                    spatialProgress.Content = $"{timer1.ElapsedMilliseconds}ms";
                });
            });
        Task.Run(() => FrequencyBlur(PadAndShiftKernel(GenerateSpatialKernel(kernelSize, sigmaValue), imgWnew, imgHnew), _fft, _cts.Token), _cts.Token)
            .ContinueWith(t => { 
                Application.Current.Dispatcher.Invoke(() => {
                    timer2.Stop();
                    fftProgress.Content = $"{timer2.ElapsedMilliseconds}ms";
                });
            });
    }

    private void LoadImageClick(object sender, RoutedEventArgs e)
    {
        OpenFileDialog openFileDialog = new OpenFileDialog();
        openFileDialog.Filter = "Image Files (*.jpg;*.jpeg;*.png;*.bmp;*.gif;*.tiff;*.tif;*.ico)|*.jpg;*.jpeg;*.png;*.bmp;*.gif;*.tiff;*.tif;*.ico|" +
                                "JPEG / JPG (*.jpg;*.jpeg)|*.jpg;*.jpeg|" +
                                "PNG (*.png)|*.png|" +
                                "BMP (*.bmp)|*.bmp|" +
                                "GIF (*.gif)|*.gif|" +
                                "TIFF (*.tiff;*.tif)|*.tiff;*.tif|" +
                                "ICO (*.ico)|*.ico";
        openFileDialog.FilterIndex = 1;
        if (openFileDialog.ShowDialog() == true)
        {
            _cts.Cancel();
            path = openFileDialog.FileName;
            source = new BitmapImage(new Uri(path));
            MainImage.Source = source;
        }
    }

    private void SaveImagesClick(object sender, RoutedEventArgs e)
    {
        string directory = System.IO.Path.GetDirectoryName(path);
        string fileNameNoExt = System.IO.Path.GetFileNameWithoutExtension(path);
        string targetFolder = System.IO.Path.Combine(directory, fileNameNoExt);

        if (!Directory.Exists(targetFolder))
        {
            Directory.CreateDirectory(targetFolder);
        }

        string originalPath = System.IO.Path.Combine(targetFolder, $"{fileNameNoExt}_original.png");
        string spatialPath = System.IO.Path.Combine(targetFolder, $"{fileNameNoExt}_spatial.png");
        string fftPath = System.IO.Path.Combine(targetFolder, $"{fileNameNoExt}_fft.png");

        SaveToPng(source, originalPath);
        SaveToPng(_spatial, spatialPath);
        SaveToPng(_fft, fftPath);
    }

    public void SpatialBlur(float[,] kernel, WriteableBitmap bitmap, CancellationToken ct)
    {
        int width = 0, height = 0, stride = 0;
        IntPtr buffer = IntPtr.Zero;

        Application.Current.Dispatcher.Invoke(() =>
        {
            width = bitmap.PixelWidth;
            height = bitmap.PixelHeight;
            stride = bitmap.BackBufferStride;
            buffer = bitmap.BackBuffer;
        });

        byte[] srcPixels = new byte[height * stride];
        Marshal.Copy(buffer, srcPixels, 0, srcPixels.Length);

        int kWidth = kernel.GetLength(0);
        int kHeight = kernel.GetLength(1);
        int radX = kWidth / 2;
        int radY = kHeight / 2;

        Application.Current.Dispatcher.Invoke(() =>
        {
            progressBar1.Maximum = height;
            progressBar1.Value = 0;
        });

        try
        {
            unsafe
            {
                byte* pDst = (byte*)buffer;

                fixed (byte* pSrc = srcPixels)
                {
                    for (int y = 0; y < height; y++)
                    {
                        ct.ThrowIfCancellationRequested();

                        for (int x = 0; x < width; x++)
                        {
                            float b = 0, g = 0, r = 0, a = 0;

                            // Kernel Window
                            for (int ky = -radY; ky <= radY; ky++)
                            {
                                // Edge handling: Clamp to image boundaries
                                int py = Math.Clamp(y + ky, 0, height - 1);
                                byte* rowPtr = pSrc + (py * stride);

                                for (int kx = -radX; kx <= radX; kx++)
                                {
                                    int px = Math.Clamp(x + kx, 0, width - 1);
                                    float weight = kernel[kx + radX, ky + radY];

                                    // 4 bytes per pixel (BGRA)
                                    int pixelIdx = px * 4;
                                    b += rowPtr[pixelIdx] * weight;
                                    g += rowPtr[pixelIdx + 1] * weight;
                                    r += rowPtr[pixelIdx + 2] * weight;
                                    a += rowPtr[pixelIdx + 3] * weight;
                                }
                            }

                            int destIdx = (y * stride) + (x * 4);
                            pDst[destIdx] = (byte)Math.Clamp(b, 0, 255);
                            pDst[destIdx + 1] = (byte)Math.Clamp(g, 0, 255);
                            pDst[destIdx + 2] = (byte)Math.Clamp(r, 0, 255);
                            pDst[destIdx + 3] = (byte)Math.Clamp(a, 0, 255);
                        }

                        if (y % 10 == 0 || y == height - 1)
                        {
                            Application.Current.Dispatcher.Invoke(() =>
                            {
                                progressBar1.Value = y;
                                bitmap.Lock();
                                bitmap.AddDirtyRect(new Int32Rect(0, 0, width, height));
                                bitmap.Unlock();
                            }, System.Windows.Threading.DispatcherPriority.Render);
                        }
                    }
                }
            }
        }
        catch (OperationCanceledException)
        {
            // Blurring was cancelled, which is expected when changing kernel size rapidly.
            // No action needed here.
            return;
        }

        Application.Current.Dispatcher.Invoke(() =>
        {
            progressBar1.Value = progressBar1.Maximum;
        });
    }

    public void FrequencyBlur(float[,] paddedKernel, WriteableBitmap bitmap, CancellationToken ct)
    {
        int width = 0, height = 0, stride = 0;
        IntPtr buffer = IntPtr.Zero;

        Application.Current.Dispatcher.Invoke(() => {
            width = bitmap.PixelWidth;
            height = bitmap.PixelHeight;
            stride = bitmap.BackBufferStride;
            buffer = bitmap.BackBuffer;
        });

        int pixelCount = width * height;

        Complex[] kernelFFT = new Complex[pixelCount];
        for (int y = 0; y < height; y++)
            for (int x = 0; x < width; x++)
                kernelFFT[y * width + x] = new Complex(paddedKernel[y, x], 0);

        Fourier.Forward2D(kernelFFT, height, width, FourierOptions.Matlab);

        byte[] srcPixels = new byte[height * stride];
        Marshal.Copy(buffer, srcPixels, 0, srcPixels.Length);

        try
        {
            for (int c = 0; c < 4; c++)
            {
                ct.ThrowIfCancellationRequested();
                Complex[] channel = new Complex[pixelCount];

                for (int y = 0; y < height; y++)
                    for (int x = 0; x < width; x++)
                        channel[y * width + x] = new Complex(srcPixels[y * stride + x * 4 + c], 0);

                Fourier.Forward2D(channel, height, width, FourierOptions.Matlab);

                for (int i = 0; i < pixelCount; i++)
                    channel[i] *= kernelFFT[i];

                Fourier.Inverse2D(channel, height, width, FourierOptions.Matlab);

                for (int y = 0; y < height; y++)
                    for (int x = 0; x < width; x++)
                        srcPixels[y * stride + x * 4 + c] = (byte)Math.Clamp(channel[y * width + x].Real, 0, 255);
            }
        }
        catch (OperationCanceledException)
        {
            // Blurring was cancelled, which is expected when changing kernel size rapidly.
            // No action needed here.
            return;
        }

        Application.Current.Dispatcher.Invoke(() => {
            bitmap.Lock();
            Marshal.Copy(srcPixels, 0, bitmap.BackBuffer, srcPixels.Length);
            bitmap.AddDirtyRect(new Int32Rect(0, 0, width, height));
            bitmap.Unlock();
        });
    }

    private void SaveToPng(BitmapSource bitmapSource, string filePath)
    {
        if (bitmapSource == null) return;

        using (var fileStream = new FileStream(filePath, FileMode.Create))
        {
            PngBitmapEncoder encoder = new PngBitmapEncoder();
            encoder.Frames.Add(BitmapFrame.Create(bitmapSource));
            encoder.Save(fileStream);
        }
    }

    public void DrawGaussianGraph(int kernelSize, double sigma)
    {
        GraphCanvas.Children.Clear();

        if (kernelSize <= 0 || sigma <= 0) return;

        Polyline curve = new Polyline
        {
            Stroke = Brushes.DodgerBlue,
            StrokeThickness = 2,
            StrokeLineJoin = PenLineJoin.Round // Ensures smooth corners
        };

        double center = (kernelSize - 1) / 2.0;
        PointCollection points = new PointCollection();

        for (int i = 0; i < kernelSize; i++)
        {
            // 1. Calculate Gaussian Value
            double xOffset = i - center;
            double exponent = -(xOffset * xOffset) / (2 * sigma * sigma);
            double yValue = Math.Exp(exponent); // Normalized height is 1.0

            // 2. Map to Canvas Coordinates (100x100)
            // X maps the index to 0-100
            double xPos = (i / (double)(kernelSize - 1)) * 100;

            // Y maps 1.0 (peak) to 0 (top of canvas) and 0.0 to 100 (bottom)
            double yPos = (25 - (yValue * 20)) ; // 90 to leave a small margin at top

            points.Add(new Point(xPos, yPos));
        }

        curve.Points = points;
        GraphCanvas.Children.Add(curve);
    }

    public WriteableBitmap PreparePaddedBitmap(BitmapImage source)
    {
        int newWidth = NextPowerOfTwo(source.PixelWidth);
        int newHeight = NextPowerOfTwo(source.PixelHeight);

        // Create the larger bitmap
        WriteableBitmap paddedBitmap = new WriteableBitmap(
            newWidth,
            newHeight,
            source.DpiX,
            source.DpiY,
            PixelFormats.Bgra32,
            null);

        // Copy the original image into the top-left corner
        // The rest of the bitmap remains black (0), which is "Zero Padding"
        byte[] pixels = new byte[source.PixelWidth * 4 * source.PixelHeight];
        source.CopyPixels(pixels, source.PixelWidth * 4, 0);

        paddedBitmap.WritePixels(
            new Int32Rect(0, 0, source.PixelWidth, source.PixelHeight),
            pixels,
            source.PixelWidth * 4,
            0);

        return paddedBitmap;
    }

    private int NextPowerOfTwo(int n)
    {
        if (n < 1) return 1;
        return (int)Math.Pow(2, Math.Ceiling(Math.Log2(n)));
    }
}