const Allocator = std.mem.Allocator;
const rl = @import("raylib");
const std = @import("std");
const assert = std.debug.assert;

const SeamCarver = @This();

const seam_window_size = 3;
comptime {
    assert(seam_window_size % 2 == 1);
}

allocator: Allocator,
source_image: rl.Image,

work_thread: ?std.Thread = null,
resized_image: ?rl.Image = null,

resizing_progress: f32 = 0,

pub fn init(allocator: Allocator, source_image: rl.Image) SeamCarver {
    return SeamCarver{
        .source_image = source_image,
        .allocator = allocator
    };
}

pub fn startResizing(self: *SeamCarver, new_width: u32, new_height: u32) !void {
    if (self.work_thread != null) {
        return;
    }

    const thread = try std.Thread.spawn(.{ .allocator = self.allocator }, resizeImageThread, .{self, new_width, new_height});
    thread.detach(); // TODO: Because of this, thread might leak memory when closing program. Small issue.

    self.work_thread = thread;
}

fn rgbToLuminocity(r: u8, g: u8, b: u8) f32 {
    return 0.2126*@as(f32, @floatFromInt(r)) + 0.7152*@as(f32, @floatFromInt(g)) + 0.0722*@as(f32, @floatFromInt(b));
}

fn applyLuminance(allocator: Allocator, image: rl.Image) ![]f32 {
    assert(image.format == .pixelformat_uncompressed_r8g8b8);

    const pixel_count: usize = @intCast(image.width * image.height);
    const luminance = try allocator.alloc(f32, pixel_count);
    errdefer allocator.free(luminance);

    const image_data: [*]u8 = @ptrCast(image.data);
    for (0..pixel_count) |i| {
        const r = image_data[i*3 + 0];
        const g = image_data[i*3 + 1];
        const b = image_data[i*3 + 2];
        luminance[i] = rgbToLuminocity(r, g, b);
    }

    return luminance;
}

fn applySobelPixel(image: []f32, width: u32, height: u32, x: u32, y: u32) f32 {
    const sobel_center_x = 1;
    const sobel_center_y = 1;
    const sobel_size_x = 3;
    const sobel_size_y = 3;
    const sobel_x: []const []const f32 = &[_][]const f32{
        &[_]f32{ -1, 0, 1 },
        &[_]f32{ -2, 0, 2 },
        &[_]f32{ -1, 0, 1 },
    };
    const sobel_y: []const []const f32 = &[_][]const f32{
        &[_]f32{ -1, -2, -1 },
        &[_]f32{  0,  0,  0 },
        &[_]f32{  1,  2,  1 },
    };

    var sum_x: f32 = 0;
    var sum_y: f32 = 0;

    for (0..sobel_size_y) |oy| {
        for (0..sobel_size_x) |ox| {
            var pixel_x: i32 = @as(i32, @intCast(x + ox)) - sobel_center_x;
            var pixel_y: i32 = @as(i32, @intCast(y + oy)) - sobel_center_y;

            pixel_x = std.math.clamp(pixel_x, 0, @as(i32, @intCast(width))-1);
            pixel_y = std.math.clamp(pixel_y, 0, @as(i32, @intCast(height))-1);

            const pixel_index = @as(usize, @intCast(pixel_y)) * width + @as(usize, @intCast(pixel_x));

            sum_x += image[pixel_index] * sobel_x[oy][ox];
            sum_y += image[pixel_index] * sobel_y[oy][ox];

            // sum_x += @as(f32, @floatFromInt(image_data[3*pixel_index + 0])) * sobel_x[oy][ox];
            // sum_x += @as(f32, @floatFromInt(image_data[3*pixel_index + 1])) * sobel_x[oy][ox];
            // sum_x += @as(f32, @floatFromInt(image_data[3*pixel_index + 2])) * sobel_x[oy][ox];

            // sum_y += @as(f32, @floatFromInt(image_data[3*pixel_index + 0])) * sobel_y[oy][ox];
            // sum_y += @as(f32, @floatFromInt(image_data[3*pixel_index + 1])) * sobel_y[oy][ox];
            // sum_y += @as(f32, @floatFromInt(image_data[3*pixel_index + 2])) * sobel_y[oy][ox];
        }
    }

    return @sqrt(sum_x*sum_x + sum_y*sum_y);
}

fn applySobel(allocator: Allocator, image: []f32, width: u32, height: u32) ![]f32 {
    const image_sobel = try allocator.alloc(f32, width * height);
    errdefer allocator.free(image_sobel);

    for (0..height) |y| {
        for (0..width) |x| {
            image_sobel[y * width + x] = applySobelPixel(image, width, height, @intCast(x), @intCast(y));
        }
    }

    return image_sobel;
}

fn calculateEnergy(allocator: Allocator, image: rl.Image) ![]f32 {
    const luminance = try applyLuminance(allocator, image);
    defer allocator.free(luminance);

    return try applySobel(allocator, luminance, @intCast(image.width), @intCast(image.height));
}

const ClampedIterator = struct {
    min: i32,
    max: i32,

    value: i32,
    count: usize,

    fn initCentered(min: i32, max: i32, value: i32, radius: u32) ClampedIterator {
        return ClampedIterator{
            .min = min,
            .max = max,
            .value = value - @as(i32, @intCast(radius / 2)),
            .count = radius
        };
    }

    fn next(self: *ClampedIterator) ?i32 {
        if (self.count == 0) {
            return null;
        }

        defer {
            self.value += 1;
            self.count -= 1;
        }
        return std.math.clamp(self.value, self.min, self.max);
    }
};

fn removeVerticalSeam(image: *rl.Image, vertical_seam: []u32) void {
    const image_width: u32 = @intCast(image.width);
    const image_height: u32 = @intCast(image.height);
    const image_data: [*]u8 = @ptrCast(image.data);

    assert(image_height == vertical_seam.len);

    var new_index: usize = 0;
    for (0..image_height) |y| {
        for (0..image_width) |x| {
            if (vertical_seam[y] == x) {
                continue;
            }

            image_data[3*new_index+0] = image_data[3*(y * image_width + x)+0];
            image_data[3*new_index+1] = image_data[3*(y * image_width + x)+1];
            image_data[3*new_index+2] = image_data[3*(y * image_width + x)+2];
            new_index += 1;
        }
    }

    image.width -= 1;
}

fn findVerticalSeam(result_buffer: []u32, energy_field: []f32, image_width: u32, image_height: u32) []u32 {
    var seam_x: u32 = 0;
    for (0..image_width) |x| {
        if (energy_field[(image_height-1) * image_height + seam_x] > energy_field[(image_height-1) * image_height + x]) {
            seam_x = @intCast(x);
        }
    }

    var vertical_seam = result_buffer[0..image_height];

    var seam_y: usize = image_height-1;
    while (true) {
        vertical_seam[seam_y] = seam_x;

        if (seam_y == 0) break;
        seam_y -= 1;

        var x_iter = ClampedIterator.initCentered(0, @as(i32, @intCast(image_width))-1, @intCast(seam_x), seam_window_size);
        while (x_iter.next()) |signed_next_x| {
            const next_x: usize = @intCast(signed_next_x);
            const next_energy = energy_field[seam_y * image_width + next_x];
            const current_energy = energy_field[seam_y * image_width + seam_x];
            if (current_energy > next_energy) {
                seam_x = @intCast(next_x);
            }
        }
    }

    return vertical_seam;
}

fn removeHorizontalSeam(image: *rl.Image, horizontal_seam: []u32) void {
    const image_width: u32 = @intCast(image.width);
    const image_height: u32 = @intCast(image.height);
    const old_image_data: [*]u8 = @ptrCast(image.data);
    const new_image_data: [*]u8 = @ptrCast(std.c.malloc(image_width * (image_height-1) * 3) orelse @panic("OOM"));

    assert(image_width == horizontal_seam.len);

    for (0..image_width) |x| {
        var new_y: u32 = 0;
        for (0..image_height) |y| {
            if (horizontal_seam[x] == y) {
                continue;
            }

            new_image_data[3*(new_y * image_width + x)+0] = old_image_data[3*(y * image_width + x)+0];
            new_image_data[3*(new_y * image_width + x)+1] = old_image_data[3*(y * image_width + x)+1];
            new_image_data[3*(new_y * image_width + x)+2] = old_image_data[3*(y * image_width + x)+2];
            new_y += 1;
        }
    }

    image.height -= 1;
    image.data = new_image_data;
    std.c.free(old_image_data);
}

fn findHorizontalSeam(result_buffer: []u32, energy_field: []f32, image_width: u32, image_height: u32) []u32 {
    var seam_y: u32 = 0;
    for (0..image_height) |y| {
        if (energy_field[seam_y * image_height + (image_width-1)] > energy_field[y * image_height + (image_width-1)]) {
            seam_y = @intCast(y);
        }
    }

    const horizontal_seam = result_buffer[0..image_width];
    @memset(horizontal_seam, image_width/2);

    var seam_x: usize = image_width-1;
    while (true) {
        horizontal_seam[seam_x] = seam_y;

        if (seam_x == 0) break;
        seam_x -= 1;

        var y_iter = ClampedIterator.initCentered(0, @as(i32, @intCast(image_height))-1, @intCast(seam_y), seam_window_size);
        while (y_iter.next()) |signed_next_y| {
            const next_y: usize = @intCast(signed_next_y);
            const next_energy = energy_field[next_y * image_width + seam_x];
            const current_energy = energy_field[seam_y * image_width + seam_x];
            if (current_energy > next_energy) {
                seam_y = @intCast(next_y);
            }
        }
    }

    return horizontal_seam;
}

fn resizeImageThread(self: *SeamCarver, new_width: u32, new_height: u32) !void {
    var image = self.source_image.copy();
    errdefer rl.unloadImage(image);

    self.resizing_progress = 0;

    const horizontal_reductions = @as(u32, @intCast(image.width)) - new_width;
    const vertical_reductions = @as(u32, @intCast(image.height)) - new_height;
    const total_progress = @as(f32, @floatFromInt(horizontal_reductions)) + @as(f32, @floatFromInt(vertical_reductions));

    var arena = std.heap.ArenaAllocator.init(self.allocator);
    defer arena.deinit();

    const vertical_seam_buffer = try self.allocator.alloc(u32, @intCast(image.height));
    defer self.allocator.free(vertical_seam_buffer);

    for (0..horizontal_reductions) |iter| {
        _ = arena.reset(.retain_capacity);

        const image_width: u32 = @intCast(image.width);
        const image_height: u32 = @intCast(image.height);

        const energy_field = try calculateEnergy(arena.allocator(), image);

        // Calculate top down energy field
        for (1..image_height) |y| {
            for (0..image_width) |x| {
                var min_top_value: f32 = std.math.floatMax(f32);

                var seam_x_iter = ClampedIterator.initCentered(0, @as(i32, @intCast(image_width))-1, @intCast(x), seam_window_size);
                while (seam_x_iter.next()) |seam_x| {
                    const index = (y - 1) * image_width + @as(usize, @intCast(seam_x));
                    min_top_value = @min(min_top_value, energy_field[index]);
                }

                energy_field[y * image_width + x] += min_top_value;
            }
        }

        const vertical_seam = findVerticalSeam(vertical_seam_buffer, energy_field, image_width, image_height);
        removeVerticalSeam(&image, vertical_seam);

        self.resizing_progress = @as(f32, @floatFromInt(iter)) / total_progress;
    }

    const horizontal_seam_buffer = try self.allocator.alloc(u32, @intCast(image.width));
    defer self.allocator.free(horizontal_seam_buffer);

    for (0..vertical_reductions) |iter| {
        _ = arena.reset(.retain_capacity);

        const image_width: u32 = @intCast(image.width);
        const image_height: u32 = @intCast(image.height);

        const energy_field = try calculateEnergy(arena.allocator(), image);

        // Calculate left to right energy field
        for (1..image_width) |x| {
            for (0..image_height) |y| {
                var min_left_value: f32 = std.math.floatMax(f32);

                var seam_y_iter = ClampedIterator.initCentered(0, @as(i32, @intCast(image_height))-1, @intCast(y), seam_window_size);
                while (seam_y_iter.next()) |seam_y| {
                    const index = @as(usize, @intCast(seam_y)) * image_width + (x - 1);
                    min_left_value = @min(min_left_value, energy_field[index]);
                }

                energy_field[y * image_width + x] += min_left_value;
            }
        }

        const horizontal_seam = findHorizontalSeam(horizontal_seam_buffer, energy_field, image_width, image_height);
        removeHorizontalSeam(&image, horizontal_seam);

        self.resizing_progress = @as(f32, @floatFromInt(horizontal_reductions + iter)) / total_progress;
    }

    self.resized_image = image;
    self.work_thread = null;
}

pub fn deinit(self: SeamCarver) void {
    if (self.resized_image) |image| {
        rl.unloadImage(image);
    }
}