const std = @import("std");

const coordModule = @import("coord.zig");
const areaModule = @import("area.zig");

const Coord = coordModule.Coord;
const Area = areaModule.Area;

pub const UnionFindNode = struct {
    parent: usize,
    size: usize,
};

pub const UnionFind  = struct {
    nodes: []UnionFindNode,
    nodesAllocator: std.mem.Allocator,

    pub fn Find(self: *UnionFind, nodeIndex: usize) usize {
        if (self.nodes[nodeIndex].parent == nodeIndex) {
            return nodeIndex;
        }

        self.nodes[nodeIndex].parent = self.Find(self.nodes[nodeIndex].parent);
        return self.nodes[nodeIndex].parent;
    }

    fn Union(self: *UnionFind, nodeAIndex: usize, nodeBIndex: usize) bool {
        var nodeARootIndex = self.Find(nodeAIndex);
        var nodeBRootIndex = self.Find(nodeBIndex);

        if (nodeARootIndex == nodeBRootIndex) {
            return false;
        }

        if (self.nodes[nodeARootIndex].size < self.nodes[nodeBRootIndex].size) {
            const t = nodeBRootIndex;
            nodeBRootIndex = nodeARootIndex;
            nodeARootIndex = t;
        }

        self.nodes[nodeBRootIndex].parent = nodeARootIndex;
        self.nodes[nodeARootIndex].size += self.nodes[nodeBRootIndex].size;
        return true;
    }

    fn UnionIteration(self: *UnionFind, area: *const Area) bool {
        var performedUnion = false;
        for (0 .. @intCast(area.rows)) |row| {
            for (0 .. @intCast(area.cols)) |col| {
                const nodeAIndex:usize = @intCast(row*@abs(area.cols) + col);
                if (area.getCell(@intCast(row), @intCast(col)) == '.') {
                    if (row > 0 and area.getCell(@intCast(row-1), @intCast(col)) == '.') {
                        const nodeBIndex: usize = @intCast((row-1)*@abs(area.cols) + col);
                        if (self.Union(nodeAIndex, nodeBIndex)) {
                            performedUnion = true;
                        }
                    }

                    if (col > 0 and area.getCell(@intCast(row), @intCast(col-1)) == '.') {
                        const nodeBIndex: usize = @intCast(row*@abs(area.cols) + col - 1);
                        if (self.Union(nodeAIndex, nodeBIndex)) {
                            performedUnion = true;
                        }
                    }

                    if (row < area.rows - 1 and area.getCell(@intCast(row+1), @intCast(col)) == '.') {
                        const nodeBIndex: usize = @intCast((row+1)*@abs(area.cols) + col);
                        if (self.Union(nodeAIndex, nodeBIndex)) {
                            performedUnion = true;
                        }
                    }

                    if (col < area.cols - 1 and area.getCell(@intCast(row), @intCast(col+1)) == '.') {
                        const nodeBIndex: usize = @intCast(row*@abs(area.cols) + col + 1);
                        if (self.Union(nodeAIndex, nodeBIndex)) {
                            performedUnion = true;
                        }
                    }
                }
            }
        }

        return performedUnion;
    }

    pub fn init(allocator: std.mem.Allocator, area: *const Area) !UnionFind {
        var unionFind: UnionFind  = .{
            .nodes = try allocator.alloc(UnionFindNode, @intCast(area.rows*area.cols)),
            .nodesAllocator = allocator,
        };

        for (0..unionFind.nodes.len) |i| {
            unionFind.nodes[i].parent = i;
            unionFind.nodes[i].size = 1;
        }

        while (true) {
          if (!unionFind.UnionIteration(area)) {
            break;
          }
        }

        return unionFind;
    }

    pub fn deinit(self: *const UnionFind) void {
        self.nodesAllocator.free(self.nodes);
    }
};