import numpy as np
import opt_einsum as oe

def analyze_tensors(inputs: list, output: list, shapes: list, optimizer='optimal', bytes_per_element: int = 16):
    """
    Simulates the contraction using integer edge lists to support infinite scale.
    """
    tensors = [np.broadcast_to(np.int8(0), shape) for shape in shapes]
    
    operands = []
    for tensor, idxs in zip(tensors, inputs):
        operands.extend([tensor, idxs])
    operands.append(output)
    
    path, path_info = oe.contract_path(*operands, optimize=optimizer)
    
    dim_map = {idx: dim for idx_list, shape in zip(inputs, shapes) for idx, dim in zip(idx_list, shape)}
            
    working_indices = list(inputs)
    num_initial = len(shapes)
    nodes, edges = [], []
    working_ids = list(range(num_initial))
    working_lineages = [[i] for i in range(num_initial)]
    
    for i in range(num_initial):
        nodes.append({
            "id": i, "label": f"T{i}", "group": "initial", 
            "shape_str": str(shapes[i]), "indices": str(inputs[i]),
            "lineage": [i], "memory_mb": 0
        })
        
    current_node_id = num_initial
    
    for step_idx, (contraction, size) in enumerate(zip(path, path_info.size_list)):
        inputs_popped_indices, inputs_popped_ids, inputs_popped_lineages = [], [], []
        
        for i in sorted(contraction, reverse=True):
            inputs_popped_indices.append(working_indices.pop(i))
            inputs_popped_ids.append(working_ids.pop(i))
            inputs_popped_lineages.append(working_lineages.pop(i))
            
        chars_in_inputs = set()
        for idx_list in inputs_popped_indices: chars_in_inputs.update(idx_list)
            
        chars_needed_later = set(output)
        for wi in working_indices: chars_needed_later.update(wi)
            
        new_indices_set = chars_in_inputs.intersection(chars_needed_later)
        new_indices_list = sorted(list(new_indices_set))
        new_shape = tuple(dim_map.get(c, 1) for c in new_indices_list)
        
        new_lineage = set()
        for lin in inputs_popped_lineages: new_lineage.update(lin)
        for input_id in inputs_popped_ids: new_lineage.add(input_id)
        new_lineage_list = sorted(list(new_lineage))
        
        leaf_lineage = [loc for loc in new_lineage_list if loc < num_initial]
        lineage_str = ", ".join([f"T{loc}" for loc in leaf_lineage])
        
        memory_mb = (size * bytes_per_element) / (1024**2)
        
        inputs_str = ", ".join([str(idx) for idx in inputs_popped_indices])
        eq_step = f"{inputs_str} -> {new_indices_list}"
        
        nodes.append({
            "id": current_node_id, "label": f"Step {step_idx}", "group": "intermediate", 
            "shape_str": str(new_shape), "equation_step": eq_step,
            "memory_mb": memory_mb, "lineage": new_lineage_list
        })
        
        for from_id in inputs_popped_ids: edges.append({"from": from_id, "to": current_node_id})
            
        working_indices.append(new_indices_list)
        working_ids.append(current_node_id)
        working_lineages.append(new_lineage_list)
        current_node_id += 1
        
    return nodes, edges, path, path_info