import torch
from typing import Dict
from utils.regularizations import *
from optimization.ActivationExtractor import *


def optimize_activation_batch(model,
                        target_neurons:list,
                        input_data:torch.Tensor,
                        target_layer:torch.nn.Module,
                        lr:float=0.01,
                        steps:int=50,
                        regularizations_dict:Dict=None, # this dict is being declared in main.py
                        log_freq: int=10):

  """
  Optimize input to maximize activation of a specific neuron with dynamic regularizations.
  Args:
      model: Pre-trained model set to evaluation mode.
      target_neuron (list): list of the neurons to maximize.
      input_data (torch.Tensor): Input tensor to optimize. Must have requires_grad=True.
      lr (float): Learning rate for the optimizer.
      steps (int): Number of optimization steps.
      regularizations_dict (dict): Dictionary of regularization functions and their parameters.
          Supported keys:
          - 'l2': L2 regularization weight (float).
          - 'tv': Total Variation loss weight (float).
          - 'sparsity': Sparsity constraint weight (float).
          - 'clip': Tuple of (min_value, max_value) for clipping input values.
      log_freq (int): Frequency of logging optimization progress.
      reduction_method (str): Method to aggregate neuron activations ('mean' or 'sum').

  Returns:
      dict: A dictionary containing:
        - 'optimized_input': The input tensor optimized to maximize the target neuron's activation.
        - 'loss_history': A list of loss values recorded during the optimization process.
  """

  optimizer = torch.optim.Adam(params=[input_data], lr=lr)
  regularizations_dict = regularizations_dict or {}

  reg_functions = {
                  'l2': l2_regularization,
                   'tv': total_variance_loss,
                   'sparsity': sparsity_constraint,
                   'entropy':entropy_loss,
                   'l1': l1_regularization,
                   'linf':linf_regularization,
                   'feature_map_sparsity':feature_map_sparsity
                   }

  results = {"optimized_input": input_data, "loss_history": []}


  with ActivationExtractor(model=model, target_layer=target_layer) as extractor:
    for step in range(steps):

      _ = model(input_data)
      activation = extractor.activation


      if not target_neurons or max(target_neurons) >= activation.size(1):
        raise ValueError(f"Invalid target neurons:{target_neurons}. Layer has {activation.size(1)} neurons.")


      if len(target_neurons)==1:
        loss = - activation[0, target_neurons[0]].sum()

      else:

        reduction_method = regularizations_dict.get("reduction", "mean") # is used to dynamically retrieve a key ("reduction") from the regularizations_dict dictionary, with a default value of "mean" if the key is not provided.
        if reduction_method == "mean":
          loss = - activation[0,target_neurons].mean() #making it scalar
        elif reduction_method == "sum":
          loss = - activation[0,target_neurons].sum()
        else:
          raise ValueError(f"Unspported Reduction Method: {reduction_method}. Supported methods are 'mean' and 'sum'.")
        #print(f"reduction:{reduction_method}")


      for reg_name, reg_weight in regularizations_dict.items():
        if reg_name=='clip':
          continue

        elif reg_name in reg_functions:
          if reg_name == "feature_map_sparsity":
            loss += reg_functions[reg_name](activation = activation, weight = reg_weight)
          else:
            loss += reg_functions[reg_name](input_data = input_data, weight = reg_weight)

      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

      if 'clip' in regularizations_dict:
        with torch.no_grad():
          input_data.clamp(*regularizations_dict['clip'])

      results["loss_history"].append(loss.item())

      if log_freq and step % log_freq == 0:
        print(f"Step {step}/{steps} | Loss: {loss.item():.4f} | Neurons: {target_neurons}")

  results['optimized_input'] = input_data

  return results