import warnings
import star_privateer as sp
import pycwt 
import pywt
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors, ticker
from scipy.optimize import least_squares
from scipy.stats import chi2

'''
Copyright 2024 Sylvain Breton, Eva Panetier

This file is part of star-privateer, an open-source software distributed
under MIT License.
'''

def compute_periods_wavelet (s, dt) :
    '''
    Create a period vector for wavelet
    analysis.
    '''
    dt_days=dt/86400
    period_samples = 512
    minimum_period = 2*dt_days
    maximum_period = s.size*dt_days
    #geometric resolution to fit the instrinsic limitations of cwt resolution at high periods
    periods = np.geomspace(minimum_period, maximum_period, period_samples) 
    return periods

def zero_pad_vector (s, pad=0) :
    """
    Apply zero-padding to vector by completing
    its length up to upper power of 2.
    """
    n_padding = int (2**(np.ceil(np.log2 (s.size))+pad) - s.size)
    s = np.concatenate ((s, np.zeros(n_padding)))
    return s

def compute_wps (s, dt, periods=None, nbins=500,
                 normalise=False, mother=None, mode=None,
                 remove_coi_gwps=False, correct_scale=True,
                 zero_padding=False, pad=0, backend="pycwt") :
    '''
    Compute Wavelet Power Spectrum (WPS), Global Wavelet Power
    Spectrum (GWPS) and corresponding cone of influence (COI) from
    input time series.
    
    Parameters
    ----------
    
    s : ndarray
      time series to analyse
      
    dt : float
      sampling of the time series (in s)
      
    periods : ndarray
      Periods on which to compute the WPS. 
      Optional, if not given, WPS will be computed for
      periods ranging from 0.1 to 50 day. Must be
      given in days.
      
    nbins : int
      number of bins to consider if the periods vector
      is generated by the function. 

    normalise : bool 
      if set to ``True``, the computed GWPS will be 
      normalised by the maximum value. Optional,
      default ``False``.
 
    mother : object
      mother wavelet to consider. Optional, if set
      to ``None``, ``pycwt.Morlet (6)`` will be used
      if ``backend`` is ``"pycwt"``, ``"cmorl1.5-1.0"``
      if ``backend`` is ``"pywavelets"``.
    
    mode : str
      alternative ways to compute and filter the WPS.
      ``mm`` for a mathematical morphology filter trying to
      remove brief high frequency systematical noise.
      ``ssq`` to reassign the WPS according to the
      synchrosqueezing method (Daubechies et al. 2000).
      ``ssqmm`` for both methods combined.
      Default to ``None`` to use the standard wavelet approach.
      Will only use a Morlet wavelet with w=6 
  
    remove_coi_gwps : bool
      if set to ``True``, do not account for bins in the cone of 
      influence when computing GWPS. Optional, default ``False``.

    correct_scale : bool
      If set to ``True``, apply the scale correction from
      Liu et al. (2007). Optional, default ``True``.

    zero_padding_wavelet : bool
      If set to ``True``, time series will be zero padded
      up to the next power of two plus ``pad`` before computing
      wavelet transform. Optional, default ``False``.

    pad : int
      If ``zero_padding_wavelet`` is set to ``True``, 
      the time series will be padded to the power of two 
      corresponding to the first larger than the time series 
      length plus ``pad``. Optional, default ``0``.

    backend : str
      Backend to consider for wavelet computation when 
      ``mode="standard"``. Either ``"pycwt"`` (default)
      or ``"pywavelets"``.
     
    Returns
    -------
    tuple
      Tuple with WPS periods (in day), WPS, GWPS and COI. 
    '''
    if zero_padding :
        n_ref = s.size
        s = zero_pad_vector (s, pad=pad)

    if periods is None :
        periods = sp.compute_periods_wavelet (s, dt)

    if mode is None or mode=="standard":  
        freqs = 1 / (periods*86400)
        if backend=="pycwt" :
          if mother is None :
            mother = pycwt.Morlet(6)
          wave, scales, freqs, coi, fft, fftfreqs = pycwt.cwt(s, dt, 
                                                              freqs=freqs, wavelet=mother)
        elif backend=="pywavelets" :
          mother = "cmor1.5-1.0"
          scales = pywt.frequency2scale(mother, freqs*dt)
          wave, freqs = pywt.cwt (s, scales, sampling_period=dt, wavelet=mother,
                                  method="fft")
          coi = None # for now cone of influence seems not to be implemented in pywavelets
        else :
          raise Exception ("Unkown backend for wavelet computation. Allowed backend are 'pycwt' or 'pywavelets'.")
        # Recomputing periods after the new frequency
        # vector provided by the function
        periods = (1/freqs)/86400
        wps = (np.abs(wave)) ** 2
        if correct_scale :
          # Liu et al. 2007 normalisation
          wps[periods!=0] /= scales[periods!=0, None]

    elif mode in ["mm", "ssq", "ssqmm"]:
        #Liu et al. normalisation already implemented in ssq
        (wps, periods, 
         coi, scales) = sp.morphology.cwt_modes(s, dt, periods, mode) 
    else :
        raise Exception ("Unkown requested mode, availables modes are standard, mm, ssq and ssqmm.")

    if zero_padding :
      n_diff = s.size - n_ref 
      s, wps = s[:n_ref], wps[:,:n_ref],
      if coi is not None :
        coi = coi[n_diff//2:n_ref+n_diff//2]
    
    wps = wps / np.var (s)
    gwps = np.copy (wps)

    if coi is not None :
      coi = coi / 86400
      coi_grid, periods_grid = np.meshgrid (coi, periods)
      if remove_coi_gwps :
        gwps[periods_grid>coi_grid] = np.nan
    gwps = np.nanmean (gwps, axis=1)
    if normalise :
        gwps = gwps / np.nanmax (gwps) 
    gwps[np.isnan (gwps)] = 0
    
    return periods, wps, gwps, coi, scales
        

def find_prot_gwps (periods, gwps,
                    return_power=False) :
  '''
  Return rotation period obtained 
  from GWPS or CS (in current implementation,
  simply return the maximum).
  '''
  # The np.abs is there for the CS 
  # case where we want also to deal
  # with negative values. 
  index = np.argmax (np.abs (gwps))
  prot = periods[index]
  if return_power :
    power = gwps[index]
    return prot, power
  else :
    return prot

def fun_residual (param, x, y) :
  return np.abs (y - sp.gauss (x, param[0], param[1], param[2])) 

def fit_gaussian (periods, ps, p_init, a0, 
                  reduce_interval=False,
                  pmin=None, pmax=None) :
  '''
  Perform a least-square fit with a Gaussian profile.
  '''
  if pmin is None :
    pmin = 0 
  if pmax is None :
    pmax = np.inf 
  param0 = np.array ([a0,
                     p_init,
                     0.1*p_init])
  # For a0 bound, first case is for positive
  # guess, second case for negative guess
  bounds = (np.array([min (0.5*a0, 4*a0),
                      max (pmin, 0.99*p_init),
                      0]),
            np.array ([max (4*a0, 0.5*a0),
                       min (pmax, 1.01*p_init),
                       10*p_init])
             )
  if reduce_interval :
    mask = (periods>0.75*p_init)&(periods<1.25*p_init)
  else :
    mask = np.full (periods.size, True)
  result = least_squares(fun_residual, param0,
                         args=(periods[mask], ps[mask]), bounds=bounds)

  model = sp.gauss (periods, result.x[0], result.x[1], result.x[2])
  return result.x, model, result.status

def compute_prot_err_gaussian_fit (periods, ps,
                                   n_profile=5, threshold=0.1,
                                   verbose=False, pmin=None,
                                   pmax=None) :
  '''
  Fit a series of gaussian profiles on a power 
  spectrum (GWPS or CS) and use it to extract 
  the rotation period estimate and corresponding error.

  Parameters
  ----------
  periods : ndarray
    Period array, in days.

  ps : ndarray
    Power spectrum.

  n_profile : int
    Maximal number of Gaussian profile to fit. 

  threshold : float
    Peaks with height below this threshold (given as a fraction
    as the highest peak) will not be fitted.

  verbose : float
    Output verbosity. Optional, default ``False``.

  pmin : float
    Minimum period to fit. Optional, default ``None``.

  pmax : float
    Maximum period to fit. Optional, default ``None``.
 
  Returns
  -------
  tuple
    The rotation period, its uncertainty and the parameters
    fitted for the ``n_profile`` profiles.
  '''
  cond = sp.compute_cond (pmin, pmax, periods)
  ps, periods = ps[cond], periods[cond]

  param = np.full ((n_profile, 3), -1.)
  p_init, power_init = find_prot_gwps (periods, ps,
                                       return_power=True)
  max_init = np.amax (ps)
  ii = 0
  while ii < n_profile and np.amax (ps) > threshold*max_init :
    fitted_param, model, status = fit_gaussian (periods, ps, p_init, power_init,
                                                pmin=pmin, pmax=pmax)
    if verbose :
      print ('Fitted profile {}, param obtained:{}, status: {}'.format (ii, fitted_param, status))
    if status < 1 :
      break 
    ps = ps - model
    p_init, power_init  = find_prot_gwps (periods, ps,
                                          return_power=True)
    param[ii,:] = fitted_param
    ii += 1 
  
  param = param[param[:,1]!=-1,:]
  
  if param.size > 0 : 
    prot, E_p = param[0,1], param[0,2]
  else :
    prot, E_p = -1, -1

  return prot, E_p, param

def plot_wps (t, periods, wps, gwps, coi=None,
              scales=None, cmap='jet', shading='auto',
              color_coi='black', ylogscale=False,
              ax1=None, ax2=None, lw=1, normscale='log',
              vmin=None, vmax=None, filename=None, dpi=300,
              param_gauss=None, show_kepler_quarters=False,
              tref=0, figsize=(14, 8), show_contour=True,
              significance_contours=False, 
              significance_level=0.95, 
              cmap_contour='cividis', levels=None,
              contourf_plot=False, ylim=(0.1, 100),
              xlabel=None, xticks=None, xticklabels=None, 
              yticks=None, yticklabels=None, 
              gwps_logscale=False, gwps_xlim=None,
              renormalise_gwps_in_ylim=True, color_gwps = 'black',
              legend_gwps = None) :
    """
    Plot output from ``compute_wps``. 

    Parameters
    ----------
    t : ndarray
        Time stamp of the time series.
    periods : ndarray
        Period vector of the wavelet analysis.
    wps : ndarray
        2d wavelet power spectrum.
    gwps : ndarray
        1d global wavelet power spectrum.
    coi : ndarray
        Boolean array defining the cone of influence
        of the wavelet analysis. Optional, default ``None``.
    cmap : str or ``ColorMap`` object
        Colormap to use for the WPS.
    shading : str
        Shading of the WPS colormesh.
    color_coi : str
        Color to use for the cone of influence.
    ylogscale : bool
        If set to ``True``, the y-axis will be shown
        in log scale.
    ax1 : ``Axes``
        ``Axe`` of external figure to use for WPS plotting.
        If any of ``ax1`` or ``ax2`` is None, a new figure
        with corresponding axes will be generated. 
        Optional, default ``None``.
    ax2 : ``Axes``
        ax2
        ``Axe`` of external figure to use for GWPS plotting. 
        Optional, default ``None``.
    lw : float
        Linewidth
    normscale : str
        Colormap scaling: ``linear`` or ``log``.
    vmin : float
        Minimal value to consider for the colormap.
    vmax : float
        Maximal value to consider for the colormap.
    filename : str or ``Path``
        Path where to save the generated figure.
    dpi : int
        Dot-per-inch of the figure.
    param_gauss : array-like
        Parameters of the Gaussian profiles fitted 
        on the GWPS.
    show_kepler_quarters : bool
        For Kepler light curves, whether to show or
        not the boundary of the quarters.
    tref : float
        Reference time for the Kepler quarters.
    figsize : tuple
        Figure size.
    show_contour : bool
        Whether to show or not contour of the WPS.
    cmap_contour : str or ``Colormap``
        Colormap to consider for the contours.
    levels : array-like
        Level on which draw the contour of the WPS
        on the figure.
    contourf_plot : bool
        If set to ``True``, a contour-filled plot will
        be produced for the WPS panel.
    ylim : tuple
        Limit of the y-axis of the two panels.
    color_gwps: str
        Color to use for the gwps plot.
    legend_gwps: NoneType or str
        Label of the gwps plot.

    Returns
    -------
      The corresponding ``matplotlib.pyplot.Figure``.
    """
    # Working with a copy of the wps
    wps = wps.copy ()

    if vmax is None :
      vmax = np.amax (wps)
    if normscale=='linear':
      if vmin is None :
        vmin = 0.
      norm = colors.Normalize(vmin=vmin, vmax=vmax) 
    if normscale=='log' :
      if vmin is None :
        vmin = 1e-4*np.amax (wps)
      # Some cells seems to be 0 with the 
      # synchrosqueezing
      wps[wps==0] = vmin
      norm = colors.LogNorm(vmin=vmin, vmax=vmax)

    if ax1 is None or ax2 is None :
      gs_kw = dict(width_ratios=[3, 1])
      fig, axs = plt.subplots (1, 2, figsize=figsize, gridspec_kw=gs_kw)
    else :
      fig = None 
      axs = (ax1, ax2)

    if contourf_plot :
      axs[0].contourf (t, periods, wps, cmap=cmap,
                       levels=levels, norm=norm)
    else :
      axs[0].pcolormesh (t, periods, wps, cmap=cmap,
                         shading=shading, norm=norm,
                         rasterized=True)

    if show_contour :
      if levels is None :
        if normscale=='log' :
          locator = ticker.LogLocator(base=10, subs=(1., 5.),
                                      )
        else :
          locator = ticker.MaxNLocator()
      else :
        locator = None
      axs[0].contour (t, periods, wps, 
                      linestyles='-', cmap=cmap_contour,
                      locator=locator, levels=levels, 
                      )

    if significance_contours :
      if scales is None :
        warnings.warn ("Scales must be provided to plot significance contour.")
      else :
        level = chi2.ppf(significance_level, 2)
        axs[0].contour (t, periods, wps * scales[:, None], 
                        levels=[level], linewidths=2, colors="black")

    if coi is not None :
      axs[0].fill_between (t, np.minimum (coi, np.amax(periods)), 
                           np.maximum (coi, np.amax(periods)),
                           alpha=0.3, hatch='x', 
                           color=color_coi)

    if ylim is not None :
      ylim = (max (np.amin (periods), ylim[0]), 
              min (np.amax (periods), ylim[1]))

    if ylim is not None and renormalise_gwps_in_ylim :
      norm_gwps = np.nanmax (gwps[(periods>ylim[0])&(periods<ylim[1])])
      gwps_xlim = (0, 1)
      axs[1].set_xticks ([])
    else :
      norm_gwps = 1
    axs[1].plot (gwps / norm_gwps, periods, 
                 color=color_gwps, lw=lw, label=legend_gwps)

    if param_gauss is not None :
      n_gauss = param_gauss.shape[0] 
      model = np.zeros (periods.size)
      for ii in range (n_gauss) :
        if param_gauss[ii,0]!=-1 : #checking a fit has been performed
          model += sp.gauss (periods, *param_gauss[ii,:])
      axs[1].plot (model / norm_gwps, periods, color='blue', lw=lw)

    if xlabel is None :
      xlabel = 'Time (day)'
    axs[0].set_xlabel (xlabel)
    axs[0].set_ylabel ('Period (day)')
    axs[1].set_xlabel ('GWPS')

    if ylogscale :
      axs[0].set_yscale ('log')
      axs[1].set_yscale ('log')
      for ax in axs :
          ax.set_ylim (max (0.1, np.amin (periods)), np.amax(periods))
    else :
      for ax in axs :
          ax.set_ylim (np.amin (periods), np.amax(periods))
    if gwps_logscale :
      axs[1].set_xscale ('log')

    if show_kepler_quarters :
      start, _ = sp.get_kepler_quarters ()
      axs[0].vlines(start - tref, np.amin (periods), np.amax (periods), 
                colors='grey', linestyles='--', zorder=10)

    axs[0].set_xlim (t[0], t[-1])
    if gwps_xlim is not None :
      axs[1].set_xlim (gwps_xlim) 
    axs[0].set_ylim (ylim)
    axs[1].set_ylim (ylim)

    if xticks is not None :
      axs[0].set_xticks (xticks)
      if xticklabels is not None :
        axs[0].set_xticklabels (xticklabels)
    if yticks is not None :
      # Removing minor ticks first as it appears
      # that it is not always done automatically
      axs[0].set_yticklabels ([], minor=True)
      axs[1].set_yticklabels ([], minor=True)
      axs[0].set_yticks (yticks)
      axs[1].set_yticks (yticks)
      if yticklabels is not None :
        axs[0].set_yticklabels (yticklabels)
        axs[1].set_yticklabels (yticklabels)

    if filename is not None :
      plt.savefig (filename, dpi=dpi)
    
    return fig
