import numpy as np
import networkx as nx

class DiffusionModelGraph:
    def __init__(self, models, architectures, outputs):
        """
        Initialize the graph of diffusion models.
        
        Args:
            models (list): List of diffusion models.
            architectures (list): List of architectural descriptions for each model.
            outputs (list): List of generated outputs (e.g., images) for each model.
        """
        self.models = models
        self.architectures = architectures
        self.outputs = outputs
        self.graph = nx.Graph()
        self._build_graph()

    def _build_graph(self):
        """
        Construct the graph by defining nodes and edges based on architectural and output similarity.
        """
        n = len(self.models)
        for i in range(n):
            self.graph.add_node(i, model=self.models[i], architecture=self.architectures[i], output=self.outputs[i])

        for i in range(n):
            for j in range(i + 1, n):
                arch_sim = self._architectural_similarity(i, j)
                out_sim = self._output_similarity(i, j)
                weight = (arch_sim * out_sim) / max(arch_sim, out_sim)
                if weight > 0:  # Add edge only if similarity exists
                    self.graph.add_edge(i, j, weight=weight)

    def _architectural_similarity(self, i, j):
        """
        Compute architectural similarity between two models.
        
        Args:
            i, j (int): Indices of the models.
        
        Returns:
            float: Architectural similarity score.
        """
        arch_i = self.architectures[i]
        arch_j = self.architectures[j]
        shared_components = sum(1 for a, b in zip(arch_i, arch_j) if a == b)
        return shared_components / len(arch_i)

    def _output_similarity(self, i, j):
        """
        Compute output similarity between two models using Mean Squared Error (MSE).
        
        Args:
            i, j (int): Indices of the models.
        
        Returns:
            float: Output similarity score.
        """
        output_i = self.outputs[i]
        output_j = self.outputs[j]
        mse = np.mean((output_i - output_j) ** 2)
        return 1 / (1 + mse)  # Convert MSE to similarity score