using DelimitedFiles, LinearAlgebra

function calculate_all_data(N::Int)
    ks = -N÷2:N÷2-1
    lattice = hcat((collect.(Iterators.product(ks, ks, ks)))...)
    labels = reshape(collect(1:N^3), (N, N, N))

    unique_i = collect(1:N^3)
    multiplicity = fill(1, length(unique_i))
    unique_point = transpose(lattice)

    return unique_i, unique_point, multiplicity, labels
end

function calculate_A1_data(N::Int)
    rotations = readdlm("rotations.mat", ',', Int, '\n')
    rotations = reshape(rotations, (24, 3, 3))
    
    ks = -N÷2:N÷2-1
    lattice = hcat((collect.(Iterators.product(ks, ks, ks)))...)
    labels = reshape(collect(1:N^3), (N, N, N))
    
    for r in 1:24
        rotated_lattice = Matrix(rotations[r, :, :]) * lattice
        for index in 1:N^3
            rotated_lattice_point = rotated_lattice[:, index]
            (i, j, k) = mod1.(rotated_lattice_point .+ (N÷2 + 1), N)
            old_label = max(labels[index], labels[i, j, k])
            new_label = min(labels[index], labels[i, j, k])
            if old_label != new_label
                for o in findall(isequal(old_label), labels)
                    labels[o] = new_label
                end
            end
        end
    end

    unique_i = unique(labels)
    multiplicity = [count(labels.==i) for i in unique_i]
    unique_point = transpose(lattice[:, unique_i])

    return unique_i, unique_point, multiplicity, labels
end

function sym_reduce(s, K_partial)
    I = one(K_partial)

    K_partial_x1 = kron(kron(K_partial, I), I)
    K_partial_y1 = kron(kron(I, K_partial), I)
    K_partial_z1 = kron(kron(I, I), K_partial)

    K_partial_x2 = kron(kron(K_partial, I), I)
    K_partial_y2 = kron(kron(I, K_partial), I)
    K_partial_z2 = kron(kron(I, I), K_partial)

    for (i, label) in enumerate(s.labels)
        if i != label
            K_partial_x2[:, i] .+= K_partial_x2[:, label]
            K_partial_y2[:, i] .+= K_partial_y2[:, label]
            K_partial_z2[:, i] .+= K_partial_z2[:, label]
        end
    end

    K_partial_x1 = K_partial_x1[s.unique_i, :]
    K_partial_y1 = K_partial_y1[s.unique_i, :]
    K_partial_z1 = K_partial_z1[s.unique_i, :]

    K_partial_x2 = K_partial_x2[:, s.unique_i]
    K_partial_y2 = K_partial_y2[:, s.unique_i]
    K_partial_z2 = K_partial_z2[:, s.unique_i]

    return (K_partial_x1, K_partial_y1, K_partial_z1), (K_partial_x2, K_partial_y2, K_partial_z2)
end