from .wave import jonswap
import numpy as np
from scipy.optimize import curve_fit

def fit_jonswap(S, w, sigma=[0.07, 0.09], initial_values=None):
    """
    Fit the JONSWAP spectrum to a given numeric array describing spectral density.

    Parameters
    ----------
    S : array_like
        Spectral density values.
    w : array_like
        Angular frequency values corresponding to `S`.
    sigma : list of float, optional
        Sigma parameters for the JONSWAP spectrum. Default is [0.07, 0.09].
    initial_values : dict, optional
        Dictionary of initial values for the parameters 'Hs', 'Tp', and 'gamma'.
        If not provided, default values are used.

    Returns
    -------
    out : dict
        Dictionary containing the fitted parameters:
        - 'Hs': Significant wave height.
        - 'Tp': Peak period.
        - 'gamma': Peak enhancement factor.
        - 'p0': Initial parameter values used for fitting.

    Notes
    -----
    Docstring is generated by GitHub Copilot.
    """
    if initial_values is None:
        initial_values = {}
        
    Hs0 = np.sqrt(np.trapz(S, w))*4

    p0 = {'Hs': Hs0, 'Tp': 1, 'gamma': 1}     #standard values of curve_fit are 1
    p0.update(**initial_values)

    if w[0]==0:
        w = w[1:]
        S = S[1:]

    fun = lambda om, Hs, Tp, gamma: jonswap(Hs, Tp, gamma, sigma=sigma)(om)
    popt,__ = curve_fit(fun, w, S)
    out = dict(Hs=popt[0], Tp=popt[1], gamma=popt[2], p0=[p0['Hs'], p0['Tp'], p0['gamma']])

    return out

def onesided_to_twosided(omega, S, axis=-1):
    """
    Convert a one-sided spectrum to a two-sided spectrum.

    Parameters
    ----------
    omega : array_like
        One-sided angular frequency array.
    S : array_like
        One-sided spectral density array.
    axis : int, optional
        Axis along which to concatenate. Default is -1.

    Returns
    -------
    omega2 : ndarray
        Two-sided angular frequency array.
    S2 : ndarray
        Two-sided spectral density array.

    Notes
    -----
    Docstring is generated by GitHub Copilot.
    """
    S2 = 0.5*np.concatenate([np.flip(S, axis=axis), S], axis=axis)
    omega2 = np.hstack([np.flip(-omega), omega])
    
    return omega2, S2

def twosided_to_onesided(omega, S):
    """
    Convert a two-sided spectrum to a one-sided spectrum.

    Parameters
    ----------
    omega : array_like
        Two-sided angular frequency array.
    S : array_like
        Two-sided spectral density array.

    Returns
    -------
    omega1 : ndarray
        One-sided angular frequency array.
    S1 : ndarray
        One-sided spectral density array.

    Notes
    -----
    Docstring is generated by GitHub Copilot.
    """
    n_samples = len(omega)    
    return omega[:n_samples//2], S[:,:,:n_samples//2]

def ramp_up(Nramp, Ntot):
    """
    Create a ramp-up scaling array.

    Parameters
    ----------
    Nramp : int
        Number of samples for ramp-up.
    Ntot : int
        Total number of samples.

    Returns
    -------
    t_scale : ndarray
        Array with ramp-up scaling applied.

    Notes
    -----
    Docstring is generated by GitHub Copilot.
    """
    t_scale = np.ones(Ntot)
    t_scale[:Nramp] = np.linspace(0, 1, Nramp)
    return t_scale

def ramp_up_t(t, t0):
    """
    Create a ramp-up scaling array based on time threshold.

    Parameters
    ----------
    t : array_like
        Time array.
    t0 : float
        Time threshold for ramp-up.

    Returns
    -------
    t_scale : ndarray
        Array with ramp-up scaling applied.

    Notes
    -----
    Docstring is generated by GitHub Copilot.
    """
    Nramp = np.sum(t<t0)
    Ntot = len(t)
    
    return ramp_up(Nramp, Ntot)