import numpy as np


class GenerativeSignalSynthesis:
    """
    Generative Signal Synthesis for creating synthetic signal data using various methods.

    This class provides methods to generate synthetic signals using techniques such as random noise, Gaussian processes,
    autoregressive (AR) models, Markov chains, and custom-defined functions.

    Methods
    -------
    generate(method="random_noise", length=100, **kwargs) :
        Generates synthetic signals using the specified method.

    Example Usage
    -------------
    signal_synthesizer = GenerativeSignalSynthesis()

    # Generate random noise
    random_signal = signal_synthesizer.generate(method="random_noise", length=100)
    print("Random Noise Signal:", random_signal)

    # Generate Gaussian process
    gp_signal = signal_synthesizer.generate(method="gaussian_process", length=100, mean=0, std_dev=1, correlation=0.9)
    print("Gaussian Process Signal:", gp_signal)

    # Generate AR model signal
    ar_signal = signal_synthesizer.generate(method="autoregressive", length=100, coeffs=[0.9, -0.5])
    print("AR Model Signal:", ar_signal)

    # Generate Markov chain signal
    markov_signal = signal_synthesizer.generate(method="markov_chain", length=100, states=[-1, 1], transition_matrix=[[0.9, 0.1], [0.1, 0.9]])
    print("Markov Chain Signal:", markov_signal)

    # Generate custom function signal
    custom_signal = signal_synthesizer.generate(method="custom_function", length=100, func=lambda x: np.sin(x))
    print("Custom Function Signal:", custom_signal)
    """

    def __init__(self):
        """Initialize the GenerativeSignalSynthesis class."""
        pass

    def generate(self, method="random_noise", length=100, **kwargs):
        """
        Generate synthetic signals.

        Parameters
        ----------
        method : str, optional
            The method to use for generating the signal. Options include "random_noise", "gaussian_process",
            "autoregressive", "markov_chain", "custom_function". Default is "random_noise".
        length : int, optional
            The length of the generated signal. Default is 100.
        kwargs : dict
            Additional parameters depending on the generation method. These include:
            - mean (float): Mean of the noise or process (for "random_noise" and "gaussian_process").
            - std_dev (float): Standard deviation (for "random_noise" and "gaussian_process").
            - correlation (float): Correlation coefficient between successive points (for "gaussian_process").
            - coeffs (list of float): Coefficients for the AR model (for "autoregressive").
            - states (list): Possible states of the Markov chain (for "markov_chain").
            - transition_matrix (list of list of float): State transition matrix (for "markov_chain").
            - func (callable): Function to generate the signal (for "custom_function").

        Returns
        -------
        numpy.ndarray
            The generated synthetic signal.

        Raises
        ------
        ValueError
            If the specified generation method is unknown.
        """
        if method == "random_noise":
            return self._generate_random_noise(length, **kwargs)
        elif method == "gaussian_process":
            return self._generate_gaussian_process(length, **kwargs)
        elif method == "autoregressive":
            return self._generate_autoregressive(length, **kwargs)
        elif method == "markov_chain":
            return self._generate_markov_chain(length, **kwargs)
        elif method == "custom_function":
            return self._generate_custom_function(length, **kwargs)
        else:
            raise ValueError(f"Unknown generation method: {method}")

    def _generate_random_noise(self, length, mean=0, std_dev=1):
        """
        Generate a random noise signal.

        Parameters
        ----------
        length : int
            Length of the generated signal.
        mean : float, optional
            Mean of the noise. Default is 0.
        std_dev : float, optional
            Standard deviation of the noise. Default is 1.

        Returns
        -------
        numpy.ndarray
            The generated random noise signal.

        Examples
        --------
        >>> signal_synthesizer = GenerativeSignalSynthesis()
        >>> random_signal = signal_synthesizer._generate_random_noise(length=100, mean=0, std_dev=1)
        >>> print(random_signal)
        """
        return np.random.normal(mean, std_dev, length)

    def _generate_gaussian_process(self, length, mean=0, std_dev=1, correlation=0.9):
        """
        Generate a synthetic signal using a Gaussian process.

        Parameters
        ----------
        length : int
            The length of the generated signal.
        mean : float, optional
            Mean of the Gaussian process. Default is 0.
        std_dev : float, optional
            Standard deviation of the Gaussian process. Default is 1.
        correlation : float, optional
            Correlation coefficient between successive points. Default is 0.9.

        Returns
        -------
        numpy.ndarray
            The generated Gaussian process signal.

        Examples
        --------
        >>> signal_synthesizer = GenerativeSignalSynthesis()
        >>> gp_signal = signal_synthesizer._generate_gaussian_process(length=100, mean=0, std_dev=1, correlation=0.9)
        >>> print(gp_signal)
        """
        signal = np.zeros(length)
        signal[0] = np.random.normal(mean, std_dev)

        for i in range(1, length):
            signal[i] = correlation * signal[i - 1] + np.random.normal(mean, std_dev)

        return signal

    def _generate_autoregressive(self, length, coeffs):
        """
        Generate a synthetic signal using an autoregressive (AR) model.

        Parameters
        ----------
        length : int
            Length of the generated signal.
        coeffs : list of float
            Coefficients for the AR model.

        Returns
        -------
        numpy.ndarray
            The generated AR model signal.

        Examples
        --------
        >>> signal_synthesizer = GenerativeSignalSynthesis()
        >>> ar_signal = signal_synthesizer._generate_autoregressive(length=100, coeffs=[0.9, -0.5])
        >>> print(ar_signal)
        """
        signal = np.zeros(length)
        p = len(coeffs)

        # Initialize the signal with random noise
        signal[:p] = np.random.normal(0, 1, p)

        for i in range(p, length):
            signal[i] = np.sum(
                [coeffs[j] * signal[i - j - 1] for j in range(p)]
            ) + np.random.normal(0, 1)

        return signal

    def _generate_markov_chain(self, length, states, transition_matrix):
        """
        Generate a synthetic signal using a Markov chain.

        Parameters
        ----------
        length : int
            Length of the generated signal.
        states : list
            Possible states of the Markov chain.
        transition_matrix : list of list of float
            State transition matrix.

        Returns
        -------
        numpy.ndarray
            The generated Markov chain signal.

        Examples
        --------
        >>> signal_synthesizer = GenerativeSignalSynthesis()
        >>> markov_signal = signal_synthesizer._generate_markov_chain(length=100, states=[-1, 1], transition_matrix=[[0.9, 0.1], [0.1, 0.9]])
        >>> print(markov_signal)
        """
        n_states = len(states)
        signal = np.zeros(length, dtype=int)
        signal[0] = np.random.choice(n_states)

        for i in range(1, length):
            signal[i] = np.random.choice(n_states, p=transition_matrix[signal[i - 1]])

        return np.array([states[s] for s in signal])

    def _generate_custom_function(self, length, func):
        """
        Generate a synthetic signal using a custom function.

        Parameters
        ----------
        length : int
            Length of the generated signal.
        func : callable
            Function to generate the signal.

        Returns
        -------
        numpy.ndarray
            The generated signal.

        Examples
        --------
        >>> signal_synthesizer = GenerativeSignalSynthesis()
        >>> custom_signal = signal_synthesizer._generate_custom_function(length=100, func=lambda x: np.sin(x))
        >>> print(custom_signal)
        """
        return np.array([func(x) for x in range(length)])
