import torch
import torch.nn as nn
from smolagents import tool

@tool
def ml_layer_analyzer(layer_type: str, input_shape: list, **kwargs) -> str:
    """A tool that dynamically instantiates a PyTorch layer, simulates a forward pass, 
    and reports output shape and parameter count. Ideal for planning ML code.
    
    Args:
        layer_type: The exact string name of the PyTorch layer (e.g., 'Conv2d', 'Linear', 'LSTM').
        input_shape: A list of integers representing the input tensor dimensions (e.g., [1, 3, 224, 224]).
        **kwargs: Arbitrary keyword arguments needed to configure the layer (e.g., in_features=10, out_features=20, out_channels=64, kernel_size=3).
    """
    try:
        # Resolve layer class from torch.nn
        if not hasattr(nn, layer_type):
            return f"Error: '{layer_type}' is not a valid layer type in torch.nn"
            
        layer_cls = getattr(nn, layer_type)
        layer_instance = layer_cls(**kwargs)
        
        # Create a mock tensor based on input shape
        mock_input = torch.randn(*input_shape)
        
        # Simulate forward pass
        with torch.no_grad():
            output = layer_instance(mock_input)
            
        # Handle unpacking if the output is a tuple (like RNNs/LSTMs)
        if isinstance(output, tuple):
            out_shape = str([list(o.shape) if hasattr(o, 'shape') else type(o) for o in output])
        else:
            out_shape = list(output.shape)
            
        param_count = sum(p.numel() for p in layer_instance.parameters())
        
        return f"Success! Layer: {layer_type} | Output Shape: {out_shape} | Total Parameters: {param_count}"
        
    except Exception as e:
        return f"Failed to analyze layer configuration. Error encountered: {str(e)}"