
import sqlite3
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
import os
import random
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # Not strictly needed in newer versions, but safe to include
from sklearn.neighbors import NearestNeighbors
import view_graph_utils
import torch
import sys
from preprocess.read_write_model import read_next_bytes


def getWorld2View2(R, t, translate=np.array([.0, .0, .0]), scale=1.0):
    Rt = np.zeros((4, 4))
    Rt[:3, :3] = R.transpose()
    Rt[:3, 3] = t
    Rt[3, 3] = 1.0

    C2W = np.linalg.inv(Rt)
    cam_center = C2W[:3, 3]
    cam_center = (cam_center + translate) * scale
    C2W[:3, 3] = cam_center
    Rt = np.linalg.inv(C2W)
    return np.float32(Rt)

def qvec2rotmat(qvec):
    return np.array([
        [1 - 2 * qvec[2]**2 - 2 * qvec[3]**2,
         2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
         2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]],
        [2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
         1 - 2 * qvec[1]**2 - 2 * qvec[3]**2,
         2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]],
        [2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
         2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
         1 - 2 * qvec[1]**2 - 2 * qvec[2]**2]])

def construct_distance_graph(images_file):
    text = images_file.ends_with(".txt") 
    with open(images_file, "r" if text else "rb") as file:
        if text:
            lines = [line.rstrip() for line in file]
            number_of_images = len(lines) //2 - 2
            positions = np.zeros((number_of_images, 3))
            quats = np.zeros((number_of_images, 4))
            names = []
            for i, line in enumerate(lines[4::2]):
                print(line)
                split = line.split(" ")
                positions[i, 0] = split[5]
                positions[i, 1] = split[6]
                positions[i, 2] = split[7]
                quats[i, 0] = split[1]
                quats[i, 1] = split[2]
                quats[i, 2] = split[3]
                quats[i, 3] = split[4]
                names.append(split[-1])
        else:
            num_reg_images = read_next_bytes(file, 8, "Q")[0]
            positions = np.zeros((num_reg_images, 3))
            quats = np.zeros((num_reg_images, 4))
            names = []
            for _ in range(num_reg_images):
                binary_image_properties = read_next_bytes(
                    file, num_bytes=64, format_char_sequence="idddddddi"
                )
                image_id = binary_image_properties[0]
                quats[i] = np.array(binary_image_properties[1:5])
                positions[i] = np.array(binary_image_properties[5:8])
                camera_id = binary_image_properties[8]
                image_name = ""
                current_char = read_next_bytes(file, 1, "c")[0]
                while current_char != b"\x00":  # look for the ASCII 0 entry
                    image_name += current_char.decode("utf-8")
                    current_char = read_next_bytes(file, 1, "c")[0]
                names.append(image_name)
                
        print(positions.shape)
        
        # For some bullshit reason, camera objects are sorted in alphabetical order
        sorted_indices = sorted(range(len(names)), key=lambda i: names[i])
        sorted_indices = np.array(sorted_indices)
        positions = positions[sorted_indices]
        quats = quats[sorted_indices]
        
        
        
        for i in range(number_of_images):
            positions[i] = qvec2rotmat(quats[i]).transpose() @ positions[i]
        positions[i, 2] #/= 8
        
        points = positions

        k = 100
        nbrs = NearestNeighbors(n_neighbors=k+1, algorithm='ball_tree').fit(points)
        distances, indices = nbrs.kneighbors(points)

        G_knn = nx.DiGraph()

        # Add nodes
        for i in range(points.shape[0]):
            G_knn.add_node(i, pos=tuple(points[i]))

        # Add edges (skip the first neighbor since it's the point itself)
        for i in range(points.shape[0]):
            for j in range(1, k+1):
                neighbor_idx = indices[i][j]
                dist = distances[i][j]
                G_knn.add_edge(i, neighbor_idx, weight= float(dist))#float(1000.0/(np.sqrt(dist) + 15)))
                
        nx.write_edgelist(G_knn, "/home/<redacted>-<redacted>/Datasets/BIGCITY/camera_calibration/aligned/sparse/0/consistency_graph.edge_list")
        print(G_knn.nodes)
        node = 0
        N = 5000
        colors = np.zeros((N, 3))
        walk = np.zeros(N, dtype=np.int32)
        current_color = np.random.rand(3)
        for i in range(N):
            colors[i] = current_color
            node = view_graph_utils.metropolis_hastings_walk(G_knn, node)
            walk[i] = node
            if i % 100 == 0:
                i = random.randint(0, points.shape[0])
                current_color = np.random.rand(3)
        neighbours = walk #np.array(list(nx.all_neighbors(G_knn, 5000)))
        print(neighbours)
        print(G_knn.number_of_edges())
        x = positions[neighbours, 0]
        y = positions[neighbours, 1]
        z = positions[neighbours, 2]#*10
        # Create 3D scatter plot
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
       #ax.plot(x, y, z, color='blue', label='Path')
        ax.scatter(x, y, z, c=colors, marker='o')
        # Optional: Add labels
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.axis('equal')
        plt.show()

DATASET_DIR = sys.argv[1] #"/home/<redacted>-<redacted>/Datasets/BIGCITY"

construct_distance_graph(DATASET_DIR + "/camera_calibration/aligned/sparse/0_txt/images.txt")


