import cv2
import numpy as np
from itertools import product
import pickle

from matrices_reader import *

PICKLE_PATH = "/tmp/pickle.truc"

VOXEL_SIZE = 8e-2
X_MIN, X_MAX = -2.0, 2.0
Y_MIN, Y_MAX = -2.0, 2.0
Z_MIN, Z_MAX = -2.0, 2.0

with open(PICKLE_PATH, 'rb') as file:
    projection_matrices = pickle.load(file)
nb_frame = len(projection_matrices)

points = np.array([[x, y, z, 1.0] for x, y, z in product(
    np.arange(X_MIN, X_MAX, VOXEL_SIZE),
    np.arange(Y_MIN, Y_MAX, VOXEL_SIZE),
    np.arange(Z_MIN, Z_MAX, VOXEL_SIZE))])

background = np.array([255, 255, 255])
print('ok')
for k in range(nb_frame):
    frame = cv2.imread(f'/tmp/masks/Image{k+1:04}.png')
    proj_mat = projection_matrices[k]
    
    cam_points = proj_mat @ points.T
    cam_points /= cam_points[2,:]
    cam_points = np.round(cam_points, 0).astype(np.int32)

    visible = np.logical_and.reduce((0 <= cam_points[0,:], cam_points[0,:] < frame.shape[1], 0 <= cam_points[1,:], cam_points[1,:] < frame.shape[0]))
    cam_points = cam_points[:,visible]
    points = points[visible,:]

    solid = np.invert(((frame[cam_points[1,:],cam_points[0,:]] == background)).all(1))
    cam_points = cam_points[:,solid]
    points = points[solid,:]
print('ok')
for k in range(nb_frame):
    frame = cv2.imread(f'/tmp/images/Image{k+1:04}.png')

    proj_mat = projection_matrices[k]
    
    cam_points = proj_mat @ points.T
    cam_points /= cam_points[2,:]
    cam_points = np.round(cam_points, 0).astype(np.int32)

    for cam_point in cam_points.T:
        cv2.circle(frame, (cam_point[0], cam_point[1]), 2, (255, 0, 0))
            

    cv2.imshow('Frame', frame)
    cv2.waitKey(0)