import sodetlib as sdl
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle


def compute_tracking_quality(S, f, df, sync):
    """
    Computes the tracking quality parameter from tracking_setup results.
    Tracking quality is a measure of how periodic the tracking response is with
    respect to the flux-ramp. It is defined as the r-squared value between the
    measured frequency response (f+df), and the (f+df) response averaged across
    all flux-ramp periods. If the signal is regular with respect to flux-ramp,
    the true freq response will be similar to the average freq response,
    resulting in a high r-squared value. If the signal is irregular, the true
    and average freq response will differ, resulting in a lower r-squared
    value. This can be used in conjunction with the f and df peak-to-peak
    values to determine which channels are tracking correctly.

    Args
    ------
    S : SmurfControl
        Pysmurf instance
    f : np.ndarray
        Array of the tracked frequency for each channelk (kHz), as returned by
        tracking_setup
    df : np.ndarray
        Array of the tracked frequency error for each channel (kHz), as
        returned by tracking_setup
    sync : np.ndarray
        Array containing tracking sync flags, as returned by tracking_setup
    """
    sync_idxs = S.make_sync_flag(sync)
    seg_size = np.min(np.diff(sync_idxs))
    nstacks = len(sync_idxs) - 1
    nchans = len(f[0])
    sig = f + df

    fstack = np.zeros((seg_size, nchans))
    for i in range(nstacks):
        si = sync_idxs[i]
        fstack[:seg_size] += (sig)[si:si+seg_size] / nstacks

    # calculates quality of estimate wrt real data
    y_real = (sig)[sync_idxs[0]:sync_idxs[0] + nstacks * seg_size, :]
    y_est = np.vstack([fstack for _ in range(nstacks)])

    # Force these to be the same len in case all segments are not the same size
    y_est = y_est[:len(y_real)]

    with np.errstate(invalid='ignore'):
        sstot = np.sum((y_real - np.nanmean(y_real, axis=0))**2, axis=0)
        ssres = np.sum((y_real - y_est)**2, axis=0)
        r2 = 1 - ssres/sstot

    return r2


class TrackingResults:
    """
    Class for storing, saving, and interpreting results from tracking_setup.
    This class can store tracking results for multiple bands at a time.
    When created, all results array are initialized to be empty. To add
    results from individual bands one at a time, use the ``add_band_data``
    function.

    Args
    ----
    S : SmurfControl
        Pysmurf Instance
    cfg : DetConfig
        DetConfig

    Attributes
    --------------
    meta : dict
        Metadata dictionary
    f_ptp_range : np.ndarray
        Array of len(2) containing the min and max allowed f_ptp. This will
        be pulled from the dev cfg.
    df_ptp_range : np.ndarray
        Array of len(2) containing the min and max allowed df_ptp. This will
        be pulled from the dev cfg.
    r2_min : float
        Min value of r2 for a channel to be considered good
    bands : np.ndarray
        Array of len (nchans) containing the band of each channel
    channels : np.ndarray
        Array of shape (nchans) containing the smurf-channel of each channel
    f : np.ndarray
        Array of shape (nchans, nsamps) containing the tracked freq response
        (kHz) throughout the tracking setup call
    df : np.ndarray
        Array of shape (nchans, nsamps) containing the untracked freq response
        (kHz) throughout the tracking setup call
    sync_idx : np.ndarray
        Array of shape (nchans, num_fr_periods) containing the indices where
        the flux ramp resetted
    r2 : np.ndarray
        Array of shape (nchans) containing the r-squared value computed by
        tracking-quality for each channel
    f_ptp : np.ndarray
        Array of shape (nchans) containing the f_ptp of each channel
    df_ptp : np.ndarray
        Array of shape (nchans) containing the df_ptp of each channel
    is_good : np.ndarray
        Array of shape (nchans) containing True if the channel passes cuts
        and False otherwise
    """
    def __init__(self, *args, **kwargs):
        if len(args) > 0:
            self.initialize(*args, **kwargs)

    def initialize(self, S, cfg):
        self._S = S
        self._cfg = cfg
        self.meta = sdl.get_metadata(S, cfg)
        self.f_ptp_range = np.array(cfg.dev.exp['f_ptp_range'])
        self.df_ptp_range = np.array(cfg.dev.exp['df_ptp_range'])
        self.r2_min = cfg.dev.exp['r2_min']

        self.bands = np.array([], dtype=int)
        self.channels = np.array([], dtype=int)
        self.subband_centers = np.array([], dtype=float)
        self.nchans = 0
        self.ngood = 0

        self.f = None
        self.df = None
        self.tracking_kwargs = [None for _ in range(8)]
        self.sync_idxs = [None for _ in range(8)]
        self.r2 = np.array([], dtype=float)
        self.f_ptp = np.array([], dtype=float)
        self.df_ptp = np.array([], dtype=float)
        self.is_good = np.array([], dtype=bool)

    def add_band_data(self, band, f, df, sync, tracking_kwargs=None):
        """
        Computes tracking-related data based on the tracking response,
        and updates the arrays described in the attributes with channels
        from a new band

        Args
        -----
        band : int
            Band of the data you're adding
        f : np.ndarray
            Tracked freq response as returned by tracking_setup
        df : np.ndarray
            Untracked freq response as returned by tracking_setup
        sync : np.ndarray
            sync arrayas returned by tracking_setup
        """
        if band in self.bands:
            raise ValueError(f"Data for band {band} has already been added!")

        dfptp = np.ptp(df, axis=0)
        m = dfptp != 0

        channels = np.where(m)[0]
        nchans = len(channels)
        bands = np.array([band for _ in channels])
        sb_centers = np.array(
            self._S.get_subband_centers(band, as_offset=False)[1]
        )
        self.channels = np.concatenate((self.channels, channels))
        self.bands = np.concatenate((self.bands, bands))
        self.subband_centers = np.concatenate((
                self.subband_centers, sb_centers[channels]))

        self.tracking_kwargs[band] = tracking_kwargs
        r2 = compute_tracking_quality(self._S, f, df, sync)

        if self.f is not None:
            # It's possible for different tracking setup calls to have a
            # slightly different number of samples, so we have to make sure
            # we cut/elongate results so they can fit into the f/df arrays
            nsamps = len(self.f[0])
            _f = np.full((nchans, nsamps), np.nan)
            _df = np.full((nchans, nsamps), np.nan)
            fi = min(nsamps, len(f))
            _f[:, :fi] = f.T[m, :fi] * 1000
            _df[:, :fi] = df.T[m, :fi] * 1000
            if fi < nsamps:
                # Fill with the last data point to not mess up ptp calcs....
                _f[:, fi:] = _f[:, fi-1][:, None]
                _df[:, fi:] = _df[:, fi-1][:, None]
        else:
            _f = f.T[m] * 1000
            _df = df.T[m] * 1000

        self.r2 = np.concatenate((self.r2, r2[m]))

        f_ptp = np.ptp(_f, axis=1)
        self.f_ptp = np.concatenate((self.f_ptp, f_ptp))

        df_ptp = np.ptp(_df, axis=1)
        self.df_ptp = np.concatenate((self.df_ptp, df_ptp))

        is_good = np.ones_like(r2, dtype=bool)
        self.is_good = np.concatenate((self.is_good, is_good))

        if self.f is None:
            self.f = _f
            self.df = _df
        else:
            nsamps = len(self.f[0])
            self.f = np.concatenate((self.f, _f[:, :nsamps]), axis=0)
            self.df = np.concatenate((self.df, _df[:, :nsamps]), axis=0)

        self.sync_idxs[band] = self._S.make_sync_flag(sync)
        self.nchans += len(_f)
        self.find_bad_chans()

    def find_bad_chans(self, f_ptp_range=None, df_ptp_range=None, r2_min=None):
        """
        Recomputes the ``is_good`` array based on cuts ranges.
        """
        if f_ptp_range is None:
            f_ptp_range = self.f_ptp_range
        if df_ptp_range is None:
            df_ptp_range = self.df_ptp_range
        if r2_min is None:
            r2_min = self.r2_min

        f0, f1 = f_ptp_range
        df0, df1 = df_ptp_range

        self.is_good = np.ones_like(self.r2, dtype=bool)
        self.is_good[self.r2 < r2_min] = 0
        self.is_good[~((f0 < self.f_ptp) & (self.f_ptp < f1))] = 0
        self.is_good[~((df0 < self.df_ptp) & (self.df_ptp < df1))] = 0
        self.ngood = np.sum(self.is_good)

    def save(self, path=None):
        saved_fields = [
            'meta', 'bands', 'channels',
            'f_ptp_range', 'df_ptp_range', 'r2_min',
            'f', 'df', 'sync_idxs', 'r2', 'is_good',
            'tracking_kwargs', 'f_ptp', 'df_ptp',
        ]
        data = {}
        for field in saved_fields:
            data[field] = getattr(self, field)
        if path is None:
            path = sdl.make_filename(self._S, 'tracking_results.npy')

        self.filepath = path
        np.save(path, data, allow_pickle=True)
        if self._S is not None:
            self._S.pub.register_file(path, 'tracking_data', format='npy')

    @classmethod
    def load(cls, path):
        self = cls()
        for k, v in np.load(path, allow_pickle=True).item().items():
            setattr(self, k, v)
        return self


def plot_tracking_channel(tr, idx, show_text=True):
    """
    Plots single tracking channel from results
    """
    f = tr.f[idx]
    df = tr.df[idx]
    fig, ax = plt.subplots(figsize=(12, 4))
    band = int(tr.bands[idx])
    channel = int(tr.channels[idx])

    ax.set_title(f"Band {band} Channel {channel}")
    ax.plot(df+f, color='grey', alpha=0.8, label='Freq Response')
    ax.plot(f, label='Tracked Freq')
    ax.set_xticks([])
    if hasattr(tr, 'subband_centers'):
        ax.set_ylabel(f"Freq offset from {tr.subband_centers[idx]:0.2f} [kHz]")
    else:
        ax.set_ylabel(f"Freq offset [kHz]")
    ax.legend(loc='upper left')

    if show_text:
        txt = '\n'.join([
            f'fptp = {tr.f_ptp[idx]:0.2f}',
            f'dfptp = {tr.df_ptp[idx]:0.2f}',
            f'r2 = {tr.r2[idx]:0.2f}',
            f'is_good = {tr.is_good[idx]}',
        ])
        bbox = dict(facecolor='white', alpha=0.8)
        ax.text(0.02, 0.1, txt, transform=ax.transAxes, bbox=bbox)
    for s in tr.sync_idxs[band]:
        ax.axvline(s, color='grey', ls='--')
    return fig, ax


def plot_tracking_summary(tr):
    """
    Plots summary of tracking results
    """
    # fig, axes = plt.subplots(figsize=(12, 8))
    mosaic = """
    BBBB.
    AAAAC
    AAAAC
    AAAAC
    AAAAC
    """
    gs = dict(hspace=0, wspace=0)
    fig, axes = plt.subplot_mosaic(mosaic, figsize=(12, 8), gridspec_kw=gs)
    f0, f1 = tr.f_ptp_range
    df0, df1 = tr.df_ptp_range

    m = tr.is_good
    ax = axes['A']
    ax.scatter(tr.f_ptp[m], tr.df_ptp[m], marker='.', alpha=0.8)
    ax.scatter(tr.f_ptp[~m], tr.df_ptp[~m], marker='.', color='red', alpha=0.2)
    xlim = (-5, f1 * 1.3)
    ylim = (-5, df1 * 1.3)
    ax.set_xlim(*xlim)
    ax.set_ylim(*ylim)
    rect = Rectangle((f0, df0), f1-f0, df1-df0, fc='green', alpha=0.2,
                     linestyle='--', ec='black', lw=3)
    ax.add_patch(rect)
    ax.set_ylabel(r"$df_\mathrm{ptp}$ (kHz)", fontsize=20)
    ax.set_xlabel(r"$f_\mathrm{ptp}$ (kHz)", fontsize=20)

    text = f"{tr.ngood} / {tr.nchans} good tracking channels"
    bbox = dict(fc='wheat', alpha=0.8)
    ax.text(0.02, 0.9, text, transform=ax.transAxes, bbox=bbox, fontsize=16)

    ax = axes['B']
    ax.set_xticks([])
    ax.hist(tr.f_ptp, range=xlim, bins=40)
    ax.set_xlim(*xlim)
    ax.axvspan(f0, f1, color='green', alpha=0.2)

    ax = axes['C']
    ax.set_yticks([])
    ax.hist(tr.df_ptp, range=ylim, bins=40, orientation='horizontal')
    ax.set_ylim(*ylim)
    ax.axhspan(df0, df1, color='green', alpha=0.2)

    return fig, ax


def disable_bad_chans(S, tr, bands=None, **kwargs):
    """
    Disables cut channels based on a TrackingResults object.
    """
    tr.find_bad_chans(**kwargs)
    if bands is None:
        bands = np.arange(8)

    bands = np.atleast_1d()
    for b in bands:
        m = tr.bands == b
        asa = S.get_amplitude_scale_array(b)
        fea = S.get_feedback_enable_array(b)
        for good, c in zip(tr.is_good[m], tr.channels[m]):
            if not good:
                asa[c] = 0
                fea[c] = 0
        S.set_amplitude_scale_array(b, asa)
        S.set_feedback_enable_array(b, asa)


@sdl.set_action()
def setup_tracking_params(S, cfg, bands, update_cfg=True, show_plots=False):
    """
    Sets up tracking parameters by determining correct frac-pp and lms-freq
    for each band.

    Args
    -----
    S : SmurfControl
        Pysmurf instance
    cfg : DetConfig
        DetConfig instance
    bands : np.ndarray, int
        Band or list of bands to run on
    update_cfg : bool
        If true, will update the device cfg and save the file.
    show_plots : bool
        If true, will show summary plots
    """

    bands = np.atleast_1d(bands)
    exp = cfg.dev.exp

    res = TrackingResults(S, cfg)

    tk = {
        'reset_rate_khz': exp['flux_ramp_rate_khz'],
        'make_plot': False, 'show_plot': False, 'channel': [],
        'nsamp': 2**18, 'return_data': True,
        'feedback_start_frac': exp['feedback_start_frac'],
        'feedback_end_frac': exp['feedback_end_frac'],
    }
    for band in bands:
        sdl.pub_ocs_log(S, f"Setting up trackng params: band {band}")
        bcfg = cfg.dev.bands[band]
        tk.update({
            'lms_freq_hz':         None,
            'meas_lms_freq':       True,
            'fraction_full_scale': exp['init_frac_pp'],
            'reset_rate_khz':      exp['flux_ramp_rate_khz'],
            'lms_gain': bcfg['lms_gain'],
            'feedback_gain': bcfg['feedback_gain'],
        })

        f, df, sync = S.tracking_setup(band, **tk)
        tr = TrackingResults(S, cfg)
        tr.add_band_data(band, f, df, sync)
        asa_init = S.get_amplitude_scale_array(band)
        disable_bad_chans(S, tr, bands=band, r2_min=0.95)

        # Calculate trracking parameters
        S.tracking_setup(band, **tk)
        lms_meas = S.lms_freq_hz[band]
        lms_freq = exp['nphi0'] * tk['reset_rate_khz'] * 1e3
        frac_pp = tk['fraction_full_scale'] * lms_freq / lms_meas

        # Re-enables all channels and re-run tracking setup with correct params
        S.set_amplitude_scale_array(band, asa_init)
        tk.update({
            'meas_lms_freq': False,
            'fraction_full_scale': frac_pp,
            'lms_freq_hz': lms_freq,
        })

        f, df, sync = S.tracking_setup(band, **tk)
        res.add_band_data(band, f, df, sync, tracking_kwargs=tk)

        # Print tracking params
        d = {
            'frac_pp': frac_pp,
            'lms_freq_hz': lms_freq,
        }
        S.log(f"Tracking params found for band {band}: {d}")

        # Update det config
        if update_cfg:
            cfg.dev.update_band(band, d, update_file=True)

    res.save()

    S.log("Running relock tracking setup...")
    res = relock_tracking_setup(S, cfg, bands=bands, show_plots=show_plots)

    return res


def _get_fixed_tone_channels(S, cfg, bands=None):
    """
    Identify channels assigned to fixed tones by checking state of amplitude and
    feedback arrays.

    Args
    ----
    S : SmurfControl
        Pysmurf instance
    cfg : DetConfig
        Det config instance
    bands : list of int
        Bands to operate on. Get from config by default.
    """

    if bands is None:
        bands = cfg.dev.exp['active_bands']

    # infer from amplitude and feedback state
    ft_ind = {}
    for band in bands:
        amp = S.get_amplitude_scale_array(band)
        feedback = S.get_feedback_enable_array(band)
        ft_ind[band] = np.where((amp != 0) & (feedback == 0))[0]

    return ft_ind


def _restore_fixed_tone_channels(S, chan_per_band: dict):
    """
    Ensure feedback is disabled on channels that are assigned fixed tones.

    Args
    ----
    S : SmurfControl
        Pysmurf instance
    chan_per_band : dict
        List of channels to disable feedback on, given as a dictionary indexed by band number.
    """

    # infer from amplitude and feedback state
    for band, chan in chan_per_band.items():
        feedback = S.get_feedback_enable_array(band)
        feedback[chan] = 0
        S.set_feedback_enable_array(band, feedback)


@sdl.set_action()
def relock_tracking_setup(S, cfg, bands=None, reset_rate_khz=None, nphi0=None,
                          feedback_gain=None, lms_gain=None, show_plots=False,
                          frac_pp=None): 
    """
    Sets up tracking for smurf. This assumes you already have optimized
    lms_freq and frac-pp for each bands in the device config. This function
    will chose the flux-ramp fraction-full-scale by averaging the optimized
    fractions across the bands you're running on.

    This function also allows you to set reset_rate_khz and nphi0. The
    fraction-full-scale, and lms frequencies of each band will be automatically
    adjusted based on their pre-existing optimized values.

    Additional keyword args specified will be passed to S.tracking_setup.

    Args
    -----
    S : SmurfControl
        Pysmurf instance
    cfg : DetConfig
        Det config instance
    reset_rate_khz : float, optional
        Flux Ramp Reset Rate to set (kHz), defaults to the value in the dev cfg
    nphi0 : int, optional
        Number of phi0's to ramp through. Defaults to the value that was used
        during setup.
    disable_bad_chans : bool
        If true, will disable tones for bad-tracking channels

    Returns
    --------
    res : dict
        Dictionary of results of all tracking-setup calls, with the bands number
        as key.
    """
    if bands is None:
        bands = cfg.dev.exp['active_bands']
    bands = np.atleast_1d(bands)

    nbands = len(bands)
    exp = cfg.dev.exp

    # check if any channels are assigned fixed tones
    fixed_tones = _get_fixed_tone_channels(S, cfg, bands)

    # Arrays containing the optimized tracking parameters for each band
    frac_pp0 = np.zeros(nbands)
    lms_freq0 = np.zeros(nbands)  # Hz
    reset_rate_khz0 = exp['flux_ramp_rate_khz']
    init_nphi0 = exp['nphi0']

    for i, b in enumerate(bands):
        bcfg = cfg.dev.bands[b]
        frac_pp0[i] = bcfg['frac_pp']
        lms_freq0[i] = bcfg['lms_freq_hz']

    # Choose frac_pp to be the mean of all running bands.
    # This is the frac-pp at the flux-ramp-rate used for optimization
    fpp0 = np.mean(frac_pp0)
    lms_freq0 *= fpp0 / frac_pp0
    
    S.log(f"Using frac-pp={fpp0}")

    # Adjust fpp, lms_freq, and flux-ramp-rate depending on desired
    # flux-ramp-rate and nphi0
    fpp, lms_freqs = fpp0, lms_freq0
    if nphi0 is not None:
        fpp *= nphi0 / init_nphi0
        lms_freqs *= fpp / fpp0
    if reset_rate_khz is not None:
        lms_freqs *= reset_rate_khz / reset_rate_khz0
    else:
        reset_rate_khz = reset_rate_khz0

    if frac_pp is not None:
        lms_freqs *=  frac_pp / fpp
        fpp = frac_pp

    res = TrackingResults(S, cfg)
    tk = {
        'reset_rate_khz': reset_rate_khz, 'fraction_full_scale': fpp,
        'make_plot': False, 'show_plot': False, 'channel': [],
        'nsamp': 2**18, 'return_data': True,
        'feedback_start_frac': exp['feedback_start_frac'],
        'feedback_end_frac': exp['feedback_end_frac'],
    }


    for i, band in enumerate(bands):
        bcfg = cfg.dev.bands[band]
        tk.update({
            'lms_freq_hz': lms_freqs[i],
            'lms_gain': bcfg['lms_gain'],
            'feedback_gain': bcfg['feedback_gain'],
        })

        if lms_gain is not None:
            tk['lms_gain'] = lms_gain
        if feedback_gain is not None:
            tk['feedback_gain'] = feedback_gain

        f, df, sync = S.tracking_setup(band, **tk)
        res.add_band_data(band, f, df, sync, tracking_kwargs=tk)

    res.find_bad_chans()
    res.save()

    # ensure channels with assigned fixed tones have feedback disabled
    _restore_fixed_tone_channels(S, fixed_tones)

    is_interactive = plt.isinteractive()
    try:
        if not show_plots:
            plt.ioff()
        fig, ax = plot_tracking_summary(res)
        path = sdl.make_filename(S, 'tracking_results.png', plot=True)
        fig.savefig(path)
        S.pub.register_file(path, 'tracking_summary', plot=True, format='png')
        if not show_plots:
            plt.close(fig)
    finally:
        if is_interactive:
            plt.ion()

    return res
