from mantid.simpleapi import *
import matplotlib.pyplot as plt
import numpy as np


def rebin_data(datax, datay, xmin, xmax, xbin):
    """Rebin input data with specified binning params

    Parameters:
    datax (numpy.ndarray): Input data X array
    datay (numpy.ndarray): Input data Y array
    xmin (float): X Min for rebinning
    xmax (float): X Max for rebinning
    xbin (float): Bin size for rebinning

    Returns:
    tuple: Rebinned data X and Y arrays
    """
    wksp = CreateWorkspace(
        DataX=datax,
        DataY=datay,
        NSpec=1,
        Distribution=True
    )

    Rebin(
        InputWorkspace=wksp,
        OutputWorkspace="wksp_rebin",
        Params=f"{xmin}, {xbin}, {xmax}"
    )

    datax_rebinned = mtd["wksp_rebin"].extractX()[0]
    datay_rebinned = mtd["wksp_rebin"].extractY()[0]

    return (datax_rebinned, datay_rebinned)


if __name__ == "__main__":
    # Dummy data of a simple Gaussian distribution
    mu = 0
    sigma = 1
    amplitude = 1

    datax = np.linspace(-5, 5, 100)
    datay = amplitude * np.exp(-((datax - mu) ** 2) / (2 * sigma ** 2))

    rebinned_data = rebin_data(datax, datay, -4., 4., .2)

    # Plot the original and rebinned data together
    plt.plot(datax, datay, label="Original Data")
    plt.plot(rebinned_data[0], rebinned_data[1], label="Rebinned Data")
    plt.xlabel('x')
    plt.ylabel('y')
    plt.legend()
    plt.show()