const std = @import("std");

const shapeModule = @import("shape.zig");
const inputModule = @import("input.zig");
const areaModule = @import("area.zig");
const coordModule = @import("coord.zig");
const unionFindModule = @import("unionfind.zig");
const areacacheModule = @import("areacache.zig");

const Shape = shapeModule.Shape;
const Area = areaModule.Area;
const Coord = coordModule.Coord;
const UnionFind = unionFindModule.UnionFind;
const AreaCache = areacacheModule.AreaCache;
const AreaCacheResult = areacacheModule.AreaCacheResult;
const AreaHashContext = areacacheModule.AreaHashContext;


const kShapeSize = shapeModule.kShapeSize;
const readInputIntoString = inputModule.readInputIntoString;
const parseInput = inputModule.parseInput;

pub fn calculateShapeSizes(allocator: std.mem.Allocator, shapeTransforms: []const []const Shape) ![]usize{
    var shapeSizes = try allocator.alloc(usize, shapeTransforms.len);
    for (0..shapeTransforms.len) |i| {
        shapeSizes[i] = 0;
        for (shapeTransforms[i][0].data) |d| {
            if (d != '.') {
                shapeSizes[i] += 1;
            }
        }
    }
    return shapeSizes;
}

pub fn calcShapesTotalSize(shapeCounts: []i32, shapeSizes: []usize) usize {
    var totalSize: usize = 0;

    for (0..shapeCounts.len) |i| {
        totalSize += shapeSizes[i] * @as(usize, @abs(shapeCounts[i]));
    }
    return totalSize;
}


pub fn canContainShapes(allocator: std.mem.Allocator, area: *const Area,
                        shapeCounts: []i32, shapeTransforms: []const []const Shape,
                        shapeSizes: []usize, minShapeSize: usize,
                        areaCache: *AreaCache) !bool {

    const cacheCheckResult = areaCache.check(area, shapeCounts);
    if (cacheCheckResult == AreaCacheResult.kPossible) {
        return true;
    } else if (cacheCheckResult == AreaCacheResult.kPossible) {
        return false;
    }

    if (calcShapesTotalSize(shapeCounts, shapeSizes) > area.totalFreeSize()) {
        try areaCache.addImpossibleCacheEntry(area, shapeCounts);
        return false;
    }
    
    var allEmpty = true;
    for (shapeCounts) |shapeCount| {
        if (shapeCount != 0) {
            allEmpty = false;
            break;
        }
    }
    if (allEmpty) {
        return true;
    }

    for (0 .. shapeCounts.len) |shapeIndex| {
        if (shapeCounts[shapeIndex] == 0) {
            continue;
        }

        for (0..shapeTransforms[shapeIndex].len) |shapeVariant| {
            for (0 .. @intCast(area.rows)) |row| {
                for (0 .. @intCast(area.cols)) |col|  {
                    if (try area.appendShape(
                          allocator,
                          &shapeTransforms[shapeIndex][shapeVariant],
                          @intCast(row), @intCast(col), '#')) |*newArea| {
                        defer newArea.deinit();
                    
                        var mutableArea: Area = newArea.*;
                        var unionFind = try UnionFind.init(allocator, &mutableArea);
                        for (0 .. @intCast(mutableArea.rows)) |unionRow| {
                            for (0 .. @intCast(mutableArea.cols)) |unionCol| {
                                if (mutableArea.getCell(@intCast(unionRow), @intCast(unionCol)) == '.') {
                                    const nodeAIndex:usize = @intCast(unionRow*@abs(mutableArea.cols) + unionCol);
                                    const nodeARootIndex = unionFind.Find(nodeAIndex);
                                    if (unionFind.nodes[nodeARootIndex].size < minShapeSize) {
                                        mutableArea.setCell(@intCast(unionRow), @intCast(unionCol), '#');
                                    }
                                }
                            }
                        }
                        unionFind.deinit();
                    
                    
                        var newShapeCounts = try allocator.alloc(i32, shapeCounts.len);
                        defer allocator.free(newShapeCounts);
                        std.mem.copyForwards(i32, newShapeCounts, shapeCounts);
                        newShapeCounts[shapeIndex] -= 1;
                    
                        if (try canContainShapes(allocator, newArea, newShapeCounts,
                                                 shapeTransforms, shapeSizes, minShapeSize,
                                                 areaCache)) {
                            try areaCache.addPossibleCacheEntry(area, newShapeCounts);
                            return true;
                        }
                    }
                }
            }
        }
    }

    try areaCache.addImpossibleCacheEntry(area, shapeCounts);
    return false;
}

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const allocator = gpa.allocator();

    const inputBuffer = try readInputIntoString(allocator, "input.txt");
    const input = try parseInput(allocator, inputBuffer);

    var shapeTransforms: [][]Shape = try allocator.alloc([]Shape, input.shapes.len);
    for (0..input.shapes.len) |i| {
        var allTransforms = try input.shapes[i].getAllTransforms(allocator);
        defer allTransforms.deinit();

        shapeTransforms[i] = try allocator.alloc(Shape, allTransforms.count());

        var j: usize = 0;
        var it = allTransforms.keyIterator();
        while (it.next()) |shapeTransform| {
            shapeTransforms[i][j] = shapeTransform.*;
            j += 1;
        }
    }

    const shapeSizes = try calculateShapeSizes(allocator, shapeTransforms);
    defer allocator.free(shapeSizes);

    var minShapeSize = shapeSizes[0];
    for (shapeSizes) |shapeSize| {
        if (shapeSize < minShapeSize) {
            minShapeSize = shapeSize;
        }
    }

    var areaCache = try AreaCache.init(allocator);
    defer areaCache.deinit();

    var counter: usize = 0;
    for (0..input.areas.len) |i| {
        const area = try Area.init(allocator, input.areas[i].rows, input.areas[i].cols);
        defer area.deinit();
        
        if(try canContainShapes(allocator, &area, input.areas[i].shapesCount,
                                shapeTransforms, shapeSizes, minShapeSize, &areaCache)) {
            counter += 1;
            try std.fs.File.stdout().writeAll("can contain shapes\n");
        } else {
            try std.fs.File.stdout().writeAll("CAN't contain shapes\n");
        }
    }

    try std.fs.File.stdout().writeAll(
      try std.fmt.allocPrint(allocator, "Answer: {d}\n", .{counter}));
}