# -*- coding: utf-8 -*-
"""
Created on Sat Dec 28 09:56:33 2024

@author: tomoh
"""

import numpy as np
import matplotlib.pyplot as plt

def gaussian(x, center, width):
    """Generate a Gaussian function."""
    return np.exp(-((x - center) ** 2) / (2 * width ** 2))

def smooth(y, width):
    """Smooth the signal using a Gaussian convolution."""
    kernel = gaussian(np.linspace(-3*width, 3*width, 6*width), 0, width)
    kernel /= np.sum(kernel)
    return np.convolve(y, kernel, mode='same')

def fourier_self_deconvolution(x, y, deconvwidth, smoothwidth):
    """
    Perform Fourier self-deconvolution to sharpen signal peaks.

    Parameters:
        x: array-like, the x-axis values of the signal.
        y: array-like, the y-axis values of the signal.
        deconvwidth: float, the width of the Gaussian used for deconvolution.
        smoothwidth: float, the width of the smoothing kernel.

    Returns:
        x: array-like, the x-axis values of the deconvoluted signal.
        y_deconv_smoothed: array-like, the deconvoluted and smoothed y-axis values.
    """
    n = len(y)
    y_fft = np.fft.fft(y)

    freq = np.fft.fftfreq(n, d=(x[1] - x[0]))
    gaussian_filter = np.exp(-2 * (np.pi * deconvwidth * freq) ** 2)

    # Perform the deconvolution in the Fourier domain
    y_deconv_fft = y_fft / (gaussian_filter+.01)

    # Inverse FFT to transform back to the time domain
    y_deconv = np.real(np.fft.ifft(y_deconv_fft))

    # Apply smoothing
    y_deconv_smoothed = smooth(y_deconv, smoothwidth)

    return x, y_deconv_smoothed

# Generate a test signal: two overlapping Gaussian peaks
x = np.linspace(-10, 10, 1000)
y1 = gaussian(x, -2, 1)
y2 = gaussian(x, 2, 1.5)
y = y1 + y2 + 0.00 * np.random.normal(size=len(x))

# Apply Fourier self-deconvolution
deconvwidth = .9
smoothwidth = 6
x_deconv, y_deconv = fourier_self_deconvolution(x, y, deconvwidth, smoothwidth)

# Plot the results
plt.figure(figsize=(10, 6))
plt.plot(x, y, label='Original Signal', color='blue')
plt.plot(x_deconv, y_deconv, label='Deconvoluted Signal', color='red')
plt.legend()
plt.title('Fourier Self-Deconvolution')
plt.xlabel('x')
plt.ylabel('y')
plt.grid()
plt.show()