import warnings

import numba
import numpy as np

import strax
import straxen

DEFAULT_MAX_SAMPLES = 20_000


@straxen.mini_analysis(requires=('records',),
                       warn_beyond_sec=10,
                       default_time_selection='touching')
def records_matrix(records, time_range, seconds_range, config, to_pe,
                   max_samples=DEFAULT_MAX_SAMPLES,
                   ignore_max_sample_warning=False):
    """Return (wv_matrix, times, pms)
      - wv_matrix: (n_samples, n_pmt) array with per-PMT waveform intensity in PE/ns
      - times: time labels in seconds (corr. to rows)
      - pmts: PMT numbers (corr. to columns)
    Both times and pmts have one extra element.

    :param max_samples: Maximum number of time samples. If window and dt
    conspire to exceed this, waveforms will be downsampled.
    :param ignore_max_sample_warning: If True, suppress warning when this happens.

    Example:
        wvm, ts, ys = st.records_matrix(run_id, seconds_range=(1., 1.00001))
        plt.pcolormesh(ts, ys, wvm.T,
                       norm=matplotlib.colors.LogNorm())
        plt.colorbar(label='Intensity [PE / ns]')
    """
    if len(records):
        dt = records[0]['dt']
        samples_per_record = len(records[0]['data'])
    else:
        # Defaults here do not matter, nothing will be plotted anyway
        dt = 10, 110
    record_duration = samples_per_record * dt

    window = time_range[1] - time_range[0]
    if window / dt > max_samples:
        with np.errstate(divide='ignore', invalid='ignore'):
            # Downsample. New dt must be
            #  a) multiple of old dt
            dts = np.arange(0, record_duration + dt, dt).astype(np.int)
            #  b) divisor of record duration
            dts = dts[record_duration / dts % 1 == 0]
            #  c) total samples < max_samples
            dts = dts[window / dts < max_samples]
            if len(dts):
                # Pick lowest dt that satisfies criteria
                dt = dts.min()
            else:
                # Records will be downsampled to single points
                dt = max(record_duration, window // max_samples)
        if not ignore_max_sample_warning:
            warnings.warn(f"Matrix would exceed max_samples {max_samples}, "
                          f"downsampling to dt = {dt} ns.")

    wvm = _records_to_matrix(
        records,
        t0=time_range[0],
        n_channels=config['n_tpc_pmts'],
        dt=dt,
        window=window)
    wvm = wvm.astype(np.float32) * to_pe.reshape(1, -1) / dt

    # Note + 1, so data for sample 0 will range from 0-1 in plot
    ts = (np.arange(wvm.shape[0] + 1) * dt / int(1e9) + seconds_range[0])
    ys = np.arange(wvm.shape[1] + 1)

    return wvm, ts, ys


@straxen.mini_analysis(requires=('raw_records',),
                       warn_beyond_sec=3e-3,
                       default_time_selection='touching')
def raw_records_matrix(context, run_id, raw_records, time_range,
                       ignore_max_sample_warning=False,
                       max_samples=DEFAULT_MAX_SAMPLES,
                       **kwargs):
    # Convert raw to records. We may not be able to baseline correctly
    # at the start of the range due to missing zeroth fragments
    records = strax.raw_to_records(raw_records)
    strax.baseline(records, allow_sloppy_chunking=True)
    strax.zero_out_of_bounds(records)

    return context.records_matrix(run_id=run_id,
                                  records=records,
                                  time_range=time_range,
                                  max_samples=max_samples,
                                  ignore_max_sample_warning=ignore_max_sample_warning,
                                  **kwargs)


@numba.njit
def _records_to_matrix(records, t0, window, n_channels, dt=10):
    n_samples = (window // dt) + 1
    # Use 32-bit integers, so downsampling saturated samples doesn't
    # cause wraparounds
    # TODO: amplitude bit shift!
    y = np.zeros((n_samples, n_channels),
                 dtype=np.int32)

    if not len(records):
        return y
    samples_per_record = len(records[0]['data'])

    for r in records:
        if r['channel'] > n_channels:
            continue

        if dt >= samples_per_record * r['dt']:
            # Downsample to single sample -> store area
            idx = (r['time'] - t0) // dt
            if idx >= len(y):
                print(len(y), idx)
                raise IndexError('Despite n_samples = window // dt + 1, our '
                                 'idx is too high?!')
            y[idx, r['channel']] += r['area']
            continue

        # Assume out-of-bounds data has been zeroed, so we do not
        # need to do r['data'][:r['length']] here.
        # This simplifies downsampling.
        w = r['data'].astype(np.int32)

        if dt > r['dt']:
            # Downsample
            duration = samples_per_record * r['dt']
            assert duration % dt == 0, "Cannot downsample fractionally"
            # .astype here keeps numba happy ... ??
            w = w.reshape(duration // dt, -1).sum(axis=1).astype(np.int32)

        elif dt < r['dt']:
            raise ValueError("Upsampling not yet implemented")

        (r_start, r_end), (y_start, y_end) = strax.overlap_indices(
            r['time'] // dt, len(w),
            t0 // dt, n_samples)
        # += is paranoid, data in individual channels should not overlap
        # but... https://github.com/AxFoundation/strax/issues/119
        y[y_start:y_end, r['channel']] += w[r_start:r_end]

    return y

[(12, 'straxen.mini_analysis', 'straxen.mini_analysis', False, 'import straxen\n'), (77, 'straxen.mini_analysis', 'straxen.mini_analysis', False, 'import straxen\n'), (72, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (86, 'strax.raw_to_records', 'strax.raw_to_records', False, 'import strax\n'), (87, 'strax.baseline', 'strax.baseline', False, 'import strax\n'), (88, 'strax.zero_out_of_bounds', 'strax.zero_out_of_bounds', False, 'import strax\n'), (104, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (44, 'numpy.errstate', 'np.errstate', True, 'import numpy as np\n'), (59, 'warnings.warn', 'warnings.warn', False, 'import warnings\n'), (71, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (47, 'numpy.arange', 'np.arange', True, 'import numpy as np\n')]