from IPython.display import display
import matplotlib.pyplot as plt
from visualization import visualize_results

def comparative_visualization(results_list):
  """
 Visualize optimized inputs and loss curves across architectures and layers.

    Args:
        results_list (list): A list of dictionaries with results for each layer.

    Returns:
        matplotlib.figure.Figure: The created figure object with visualizations.
  """

  if len(results_list)==1:
    result = results_list[0]
    return visualize_results(
      input_tensor=result["optimized_input"],
      loss_history=result["loss_history"],
      neuron_id=result["adjusted_neurons"],
      layer_name=result["layer_name"],
      )
    
  fig, ax = plt.subplots(nrows=len(results_list), ncols=2, figsize=(12,5*len(results_list)))


  for i, result in enumerate(results_list):
    input_image = result["optimized_input"].squeeze().permute(1,2,0).detach().cpu().numpy()
    input_image_normalized = (input_image - input_image.min())/(input_image.max()-input_image.min())
    ax[i,0].imshow(input_image_normalized)
    ax[i,0].set_title(f"{result['architecture']}-{result['layer_name']}")
    ax[i,0].axis("Off")

    loss = result["loss_history"]
    ax[i,1].plot(loss, marker='o')
    ax[i,1].set_title(f"Loss History({result['architecture']}-{result['layer_name']})")
    ax[i,1].set_xlabel("Steps")
    ax[i,1].set_ylabel("Loss")

  plt.tight_layout()
  return fig