from typing import Any, Optional
from smolagents.tools import Tool

class InfoTheoryExplainerTool(Tool):
    name = "information_theory_explainer"
    description = """
    Five basic information theory concepts in simple terms.
    Given a concept name, it returns a high-school explanation with an example if applicable.
    """
    inputs = {'concept': {'type': 'string', 'description': "The information theory concept to explain (e.g., 'entropy', 'mutual information', 'KL divergence')."}}
    output_type = "string"

    def forward(self, concept: str):
        explanations = {
            "entropy": "Entropy measures the uncertainty or unpredictability in a set of outcomes. Example: A fair coin flip has 1 bit of entropy.",
            "mutual information": "Mutual information quantifies how much information one random variable tells you about another. Example: Knowing the weather can reduce uncertainty about umbrella usage.",
            "kl divergence": "KL Divergence measures how one probability distribution differs from another. It's not symmetric and often used in model evaluation.",
            "cross entropy": "Cross-entropy measures the difference between the true distribution and the predicted one, often used as a loss function in classification tasks.",
            "channel capacity": "Channel capacity is the maximum rate at which information can be reliably transmitted over a communication channel."
        }

        key = concept.lower().strip()
        return explanations.get(key, "Concept not found. Try 'entropy', 'mutual information', 'KL divergence', 'cross entropy', or 'channel capacity'.")

    def __init__(self, *args, **kwargs):
        self.is_initialized = False