# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/I__bld_i.ipynb (unless otherwise specified).

__all__ = ['w_conf', 'Test', 'phrase_ylim', 'main']

# Cell
import numpy as np
import h5py
import matplotlib.pyplot as plt
import mpl_interactions.ipyplot as iplt
import ipywidgets as widgets
from ipywidgets import interact

import re
import os
from glob import glob

from .core.baseline import get_baseline, sub_baseline
from .utils.misc import extend_Trues
from .utils.io import PolarMjdChan_to_MjdChanPolar

from scipy import ndimage

import warnings
warnings.filterwarnings("ignore", r'overflow encountered in exp')

# Cell
w_conf = {}
w_conf['continuous_update'] = False
w_conf['orientation'] = 'horizontal'
w_conf['readout'] = True
w_conf['layout'] = {}


def _IntSlider(readout_format='d', **kwargs):
    """
    kwargs:
    slider = (value=7, min=0, max=10, step=1)
    """
    res = {}
    for key in kwargs.keys():
        res[key] = widgets.IntSlider(*kwargs[key], description=f'{key}',
                                     **w_conf, readout_format=readout_format)
    return res


def _BoundedIntText(readout_format='d', **kwargs):
    """
    kwargs:
    slider = (value=7, min=0, max=10, step=1)
    """
    res = {}
    for key in kwargs.keys():
        res[key] = widgets.BoundedIntText(*kwargs[key], description=f'{key}',
                                          **w_conf, readout_format=readout_format)
    return res


def _IntText(readout_format='d', **kwargs):
    """
    kwargs:
    slider = (value=7, min=0, max=10, step=1)
    """
    res = {}
    for key in kwargs.keys():
        res[key] = widgets.IntText(*kwargs[key], description=f'{key}',
                                   **w_conf, readout_format=readout_format)
    return res


def _FloatSlider(readout_format='.2f', **kwargs):
    """
    kwargs:
    slider = (value=7, min=0, max=10, step=1)
    """
    res = {}
    for key in kwargs.keys():
        v = kwargs[key]
        res[key] = widgets.FloatSlider(value=v[0], min=v[1], max=v[2], step=v[3], description=f'{key}',
                                       **w_conf, readout_format=readout_format)
    return res


def _FloatLog10Slider(readout_format='.2f', **kwargs):
    """
    log10
    kwargs:
    slider = (value, min exponent of base, max exponent of base, exponent step)
    """
    res = {}
    for key in kwargs.keys():
        v = kwargs[key]
        res[key] = widgets.FloatLogSlider(value=v[0], base=10, min=v[1], max=v[2], step=v[3], description=f'{key}',
                                          **w_conf, readout_format=readout_format)
    return res


def _RadioButtons(**kwargs):
    """
    kwargs:
    slider = options
    """
    res = {}
    for key in kwargs.keys():
        v = kwargs[key]
        res[key] = widgets.RadioButtons(options=v, value=v[0], description=f'{key}',
                                        **w_conf,)
    return res
# widgets.RadioButtons(
#     options=['pepperoni', 'pineapple', 'anchovies'],
# #    value='pineapple', # Defaults to 'pineapple'
# #    layout={'width': 'max-content'}, # If the items' names are long


def _Dropdown(**kwargs):
    """
    kwargs:
    slider = options
    """
    res = {}
    for key in kwargs.keys():
        v = kwargs[key]
        res[key] = widgets.Dropdown(options=v, value=v[0], description=f'{key}',
                                    **w_conf,)
    return res

# Cell
class Test(object):
    def __init__(self, T2p, freq, frange=None):

        self.T2p = T2p
        freq = freq[:]
        if frange is not None:
            self.is_use = (freq > frange[0]) & (freq < frange[1])
            self.freq = freq[self.is_use]
        else:
            self.is_use = None
            self.freq = freq
        self.select()

    def select(self, start=0, length=20, polar=0):
        self.T2p_t = PolarMjdChan_to_MjdChanPolar(self.T2p[polar:polar+1, start:start+length])
        if self.is_use is not None:
            self.T2p_t = self.T2p_t[:, self.is_use]
    def sub(self, x, start=0, length=20, polar=0, **kwargs):
        self.select(start, length, polar)
        self.bld = sub_baseline(self.freq, self.T2p_t,
                                verbose=False, **kwargs)
        return np.full(x.shape, np.nan)

    def get_ori(self, x, i, **kwargs):
        res = self.T2p_t[i, :, 0]
        return np.vstack([res, ndimage.gaussian_filter1d(res, 3)]).T

    def get_bld(self, x, i, **kwargs):
        res = self.bld[i, :, 0]
        return np.vstack([res, ndimage.gaussian_filter1d(res, 3)]).T

    def get_bl(self, x, i,  **kwargs):
        return self.T2p_t[i, :, 0] - self.bld[i, :, 0]

    def get_bld_mean(self, x, start_stop,  **kwargs):
        start, stop = start_stop
        res = np.mean(self.bld[int(start):int(stop)+1, :, 0], axis=0)
        return np.vstack([res, ndimage.gaussian_filter1d(res, 3)]).T

# Cell
def phrase_ylim(ylim, vals):

    ylim = list(map(lambda x: float(x) if 'per' != str(x)[:3] else
                    np.nanpercentile(vals, float(str(x)[3:]), interpolation='nearest'), ylim))
    return ylim

# Cell
def main():
    global ylim

    tes = Test(T2p, freq, frange)

    sliders = {}
    sliders.update(_BoundedIntText(
        start=(0, 0, tes.T2p.shape[1]-1-length, 1), polar=(0, 0, 1, 1)))
    sliders.update(_IntSlider(njoin=(0, 1, length, 1)))
    sliders.update(_Dropdown(method=('arPLS', 'asPLS', 'masPLS', 'poly', 'Chebyshev'),
                             s_method_freq=('gaussian', 'boxcar', 'none'),
                             s_method_t=('none', 'gaussian', 'boxcar'),
                             ))
    sliders.update(_IntSlider(s_sigma_freq=(3, 1, 20, 1), s_sigma_t=(3, 1, 20, 1)))
    sliders.update(_FloatLog10Slider(lam=(1e8, 5, 13, 0.2), readout_format='.2e'))
    sliders.update(_IntSlider(deg=(2, 1, 10, 1),
                              niter=(100, 1, 200, 1)))



    plt.ioff()
    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=figsize)

    if isinstance(ylim, str):
        if ylim == 'auto':
            ylim = phrase_ylim(['per0.1', 'per99.5'], tes.T2p_t)
        elif ylim == 'full':
            ylim = 'fixed'
    else:
        ylim = phrase_ylim(ylim, tes.T2p_t)

    if isinstance(ylim, str):
        ylim2 = ylim
    else:
        ylim2 = [(ylim[0]-ylim[1])/2, (ylim[1]-ylim[0])/2]
    xlim = [tes.freq.min(), tes.freq.max()]

    # bld parameter
    # slider_formats = {'lam': "{:.2e}",
    #                   'njoin': "{:d}",
    #                   }

    controls = iplt.plot(tes.freq, tes.sub, nproc=nproc,
                         **sliders,
                         color='r', linewidth=2, xlim=xlim, ylim=ylim, ax=ax1, play_buttons=True, display_controls=False)

    controls2 = iplt.plot(tes.freq, tes.get_ori, i=range(length),
                          controls=controls, xlim=xlim, ylim=ylim, ax=ax1, display_controls=False, play_buttons=True)
    controls3 = iplt.plot(tes.freq, tes.get_bl,
                          controls=controls2, xlim=xlim, ylim=ylim, ax=ax1)

    iplt.plot(tes.freq, tes.get_bld,
              controls=controls2, xlim=xlim, ylim=ylim2, ax=ax2)
    plt.axhline(y=0,)

    controls4 = iplt.plot(tes.freq, tes.get_bld_mean, start_stop=('r', 0, length, 11),
                          controls=controls, xlim=xlim, ylim=ylim2, ax=ax3)

    plt.axhline(y=0, )
    ax1.set_title('origial spectra', fontsize=8)
    ax2.set_title('spectra after baseline removed', fontsize=8)

    fig.tight_layout()

    w = controls.controls
    hbs = [
        widgets.HBox([widgets.Label(value=f"1. Select spectra from $start$(max:{controls.controls['start'].max}) to $start+{length}$ to fit:"),
                      w['start'], w['polar']]),
        widgets.HBox([widgets.Label(value=f"2. Fitting:"), w['njoin'], ]),
        widgets.HBox([w['s_method_t'], w['s_sigma_t'],
                     w['s_method_freq'], w['s_sigma_freq']]),
        widgets.HBox([w['method'], w['lam'], w['deg'], w['niter']]),
        widgets.HBox(
            [widgets.Label(value=f"3. Show one spectra in top and middle panels:"), w['i']]),
        widgets.HBox([widgets.Label(
            value=f"4. Show the stack spectra of in the range in the bottom panel:"), w['start_stop']]),
    ]

    BOX = widgets.VBox(hbs)
    display(BOX)
    plt.show()
    display(BOX)
