import random
import math
import matplotlib.pyplot as plt
import matplotlib
from smolagents import tool

@tool
def generate_normal_distribution(mean: float, std_dev: float, count: int = 10000)->list:
    """Generate a list of random numbers from a normal distribution.

    This function generates a list of random numbers drawn from a normal 
    distribution specified by the mean and standard deviation.

    Args:
        mean: The mean (average) of the normal distribution.
        std_dev: The standard deviation of the normal distribution.
        count: The number of random samples to generate (default: 10000).

    Returns:
        list: A list of samples drawn from the specified normal distribution.
    """
    samples = []

    for _ in range(count // 2):  # Generate pairs of samples
        u1 = random.random()
        u2 = random.random()
        
        # Box-Muller transform
        z0 = math.sqrt(-2.0 * math.log(u1)) * math.cos(2.0 * math.pi * u2)
        z1 = math.sqrt(-2.0 * math.log(u1)) * math.sin(2.0 * math.pi * u2)

        # Scale and shift to the specified mean and standard deviation
        samples.append(z0 * std_dev + mean)
        samples.append(z1 * std_dev + mean)

    return samples

@tool
def create_histogram_and_theorical_pdf(mean: float, std_dev:float, random_numbers:list)->str:
    """Generate a histogram of random numbers and overlay the theoretical 
    probability density function (PDF) of a normal distribution. 
    Return the histogram as a base64-encoded string.

    Args:
        mean: The mean (average) of the normal distribution.
        std_dev: The standard deviation of the normal distribution.
        random_numbers: A list of random numbers generated from a 
                                         normal distribution.

    Returns:
        str: The graphics for the histogram and probability density function (PDF) on string format
    """
    # Prepare data for plotting
    hist_data = [0] * 50  # Create a list to hold histogram data
    min_value = min(random_numbers)
    max_value = max(random_numbers)
    bin_width = (max_value - min_value) / len(hist_data)

    # Fill histogram data
    for number in random_numbers:
        bin_index = int((number - min_value) / bin_width)
        if bin_index >= len(hist_data):
            bin_index = len(hist_data) - 1
        hist_data[bin_index] += 1

    # Normalize histogram data
    hist_data = [count / len(random_numbers) / bin_width for count in hist_data]

    # Prepare x values for the theoretical PDF
    x_values = [min_value + i * bin_width for i in range(len(hist_data))]

    # Calculate the corresponding y values for the theoretical normal distribution
    pdf_values = [
        (1 / (std_dev * math.sqrt(2 * math.pi))) * math.exp(-0.5 * ((x - mean) / std_dev) ** 2) \
            for x in x_values
    ]

     # Scale for ASCII output
    max_hist = max(hist_data) if hist_data else 1  # Avoid division by zero
    max_pdf = max(pdf_values) if pdf_values else 1  # Avoid division by zero
    max_height = 20  # Maximum height of the ASCII histogram

    # Building the ASCII graph as a string
    ascii_graph = "Histogram (|: Counts, -: PDF)\n"
    for i in range(len(hist_data)):
        hist_count = int((hist_data[i] / max_hist) * max_height)
        pdf_count = int((pdf_values[i] / max_pdf) * max_height)
        ascii_graph += f"{'|' * hist_count} {'-' * pdf_count} {x_values[i]:.2f}\n"
    
    return ascii_graph