#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" Methos for making beam profile realted plots, e.g. from ionisation profile monitor (IPM) data

"""

__author__ = "Philipp Niedermayer"
__contact__ = "p.niedermayer@gsi.de"
__date__ = "2022-09-29"



from itertools import cycle

from .common import *






def plot_beam_size(ax, ipm_data, xy, time_range=None, smoothing=None, fit=fit_lorenzian, **kwargs):
    """Plot beam size as function of time from fitted IPM data
    
    :param ipm_data: instance of IPMData
    :param xy: plane to consider, either 'x' or 'y'
    :param time_range: range of (start, stop) time in s to plot
    :param smoothing: if specified, apply a moving average smoothing filter of this width to the data
    :param fit: fit function
    """
    mask = slice(None, None) if time_range is None else irng(ipm_data.t, *time_range)
    data = getattr(ipm_data, xy)[mask]    
    
    #if smoothing: data = scipy.signal.savgol_filter(data, smoothing, 0, axis=0)
    
    n = data.shape[0]
    v, ve = np.nan*np.zeros(n), np.nan*np.zeros(n)
    for j in range(n):
        try:
            (*_, v[j]), (*_, ve[j]), _, fitted = fit(ipm_data.w, data[j,:])
        except RuntimeError: # fit failed
            ...
    
    if smoothing:
        v, ve = [scipy.signal.savgol_filter(_, smoothing, 0) for _ in (v, ve)] # savgol filter with order 0 is moving average
    
    ls, = ax.plot(ipm_data.t[mask], v, **kwargs)
    ax.fill_between(ipm_data.t[mask], v-ve, v+ve, alpha=0.3, color=ls.get_color(), zorder=-10)    
    ax.set(ylabel=f'Beam size $\\sigma_{xy}$ / mm', xlabel='Time / s')



def plot_beam_profile(ax, ipm_data, xy, at_time, smoothing=None, fit=fit_gaussian, cut_lower=None, **kwargs):
    """Plot IPM beam profile at given time with fits
    
    :param ax: The axis to plot onto
    :param ipm_data: data object with IPM profile data
    :param xy: what to plot, 'x' or 'y'
    :param at_time: time in sec of profile to plot
    :param smoothing: smooth profile over every this many bins
    :param fit: function or list of functions to use for fitting, see fitting.py
    """
    profile = np.copy(getattr(ipm_data, xy)[np.argmin(np.abs(ipm_data.t-at_time)),:]) # copy in order not to modify in-place
    if smoothing: profile = smooth(profile, n=smoothing)[0]
    
    # cut MCP noise away
    if cut_lower:
        ax.stairs(profile[1:-1], (ipm_data.w[1:]+ipm_data.w[:-1])/2, color='gray', alpha=.5, **kwargs)
        profile[profile<cut_lower*max(profile)] = 0
    ax.stairs(profile[1:-1], (ipm_data.w[1:]+ipm_data.w[:-1])/2, #color=cmap_petroff_10(0 if xy=='x' else 2),
              label=f'Profile ($t={at_time*SI.s:~L}$)', **kwargs)
    
    # Fit
    if not hasattr(fit, '__iter__'): fit = (fit,)
    ls = cycle(["--",":","-."])
    fitresults = []
    for fit in fit:
        p, pe, _, plot = fit(ipm_data.w, profile)
        if fit == fit_gaussian:
            s = SI.Measurement(p[3], pe[3], ipm_data.w_unit)
            label = f'Gaussian fit\n$\\sigma_{xy}={s:~L}$'
        elif fit == fit_lorenzian:
            s = SI.Measurement(p[3], pe[3], ipm_data.w_unit)
            label = f'Lorenzian fit\n$\\gamma_{xy}={s:~L}$'
        else:
            s = (p, pe)
            label = fit.__name__
        fitresults.append(s)
        ax.plot(*plot, ls=next(ls), label=label)
        ax.legend()
    
    ax.set(xlabel=f'{xy} / {ipm_data.w_unit}', ylabel=f'Intensity / {ipm_data.unit}', xlim=(min(ipm_data.w), max(ipm_data.w)), ylim=(0, ax.get_ylim()[1]))
    return fitresults[0] if len(fitresults)==1 else fitresults