from BeamDiffusionModel.models.clip.clip import clip
import numpy as np

class Node:
    def __init__(self, step_text, image_path, latents, heat_map, step, generated_from_step, generated_from_latent, parent=None):
        self.step_text = step_text
        self.step = step
        self.generated_from_step = generated_from_step
        self.generated_from_latent = generated_from_latent
        self.image = image_path
        self.latents = latents
        self.heat_map = heat_map
        self.parent = parent
        self.softmax = 0.
        self.children = []


    def get_ancestors(self, n):
        ancestors = [self]
        current_node = self.parent
        while current_node is not None and current_node.parent is not None and n > 0:
            ancestors.append(current_node)
            current_node = current_node.parent
            n -= 1
        return ancestors

    def set_softmax(self, softmax, n_latents, n_max_latents, epsilon=1e-6):
        self.softmax = np.log((softmax.cpu() * (1 / n_max_latents)) / (1 / n_latents) + epsilon)

    def get_softmax(self):
        return self.softmax

    def add_child(self, node):
        self.children.append(node)

    def find_all_paths(self):
        all_paths = []
        stack = [(self, [self])]  # Stack holds tuples of (current_node, path_to_current_node)

        while stack:
            current_node, path = stack.pop()

            # If the current node is a leaf, add the path to all_paths
            if not current_node.children:
                all_paths.append(path)
            else:
                # Add all children to the stack with the updated path
                for child in current_node.children:
                    stack.append((child, path + [child]))

        return all_paths

    def get_latent(self, idx):
        return self.latents[idx]

    def get_features(self):
        text_embedding = clip.generate_embedding(self.step_text)
        img_embedding = clip.generate_embedding(self.image)
        return text_embedding, img_embedding