import logging
from pathlib import Path

import h5py
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
from scipy import io as sio

from scipy.signal import argrelmin
from skimage.filters import (
    threshold_minimum,
)

logger = logging.getLogger("parsing")
logger.addHandler(logging.NullHandler())


def loadmat(filepath):
    config = {}
    try:
        f = h5py.File(filepath)
        for k, v in f.items():
            config[k] = np.array(v).squeeze()
    except OSError:
        config = sio.matlab.loadmat(filepath, squeeze_me=True)
    finally:
        return config


class WantedValueError(Exception):
    def __init__(self, actual, wanted):
        self.actual = actual
        self.wanted = wanted

    def __str__(self):
        return f"Error: {self.actual} <~ {self.wanted}"


def parse(
    measured_dict,
    parameters_dict,
    generated_dict,
    simple_threshold=True,
    remove_min=False,
    resolution=np.power(2, 8, dtype=int),
    awg_default=np.array(5.0e8),
    pico_default=np.array(8.0e-10),
    *,
    debug_output=Path(),
):
    debug = logger.isEnabledFor(logging.DEBUG)

    data = measured_dict["out"]
    awg_sample_rate = parameters_dict.get("AWG_sample_rate", awg_default)
    pico_sample_time = parameters_dict.get("PICO_sampling_time", pico_default)
    bit_width = np.reciprocal(awg_sample_rate * pico_sample_time).squeeze()

    img_size = generated_dict["img_size"]
    pad_size = generated_dict["pad_size"]
    pause_size = generated_dict["pause_size"]

    if debug:
        logger.debug("bit_width: %d", bit_width)
        logger.debug("img_size: %d", img_size)
        logger.debug("pad_size: %d", pad_size)
        logger.debug("pause_size: %d", pause_size)

    initial_pause = generated_dict.get("initial_pause", None)
    if initial_pause is None:
        signal = generated_dict.get("signal", None)
        signal_max = np.nonzero(signal == np.iinfo(signal.dtype).max)[0]
        first_bit = signal_max[0]
        first_space = pause_size - first_bit
    else:
        initial_pause = initial_pause.flatten()
        first_space = (
            pause_size - np.nonzero(np.min(initial_pause) != initial_pause)[0][0]
        )

    final_pause = generated_dict.get("final_pause", None)
    if initial_pause is None:
        assert final_pause is None
        last_bit = signal_max[-1]
        last_space = pause_size - (signal.size - last_bit)
    else:
        final_pause = final_pause.flatten()
        last_space = np.nonzero(np.min(final_pause) != final_pause)[0][0]

    n_samples = generated_dict["features"].shape[0]
    data_size = n_samples * (img_size + pad_size)

    # create a window such that is a bit smaller than the distance between two images.
    window = int((pad_size - 1) * bit_width)
    # do a moving maximum. Each img will be separated by some white space.

    if remove_min:
        min_data = (
            pd.Series(data)
            .rolling(window=window, center=True, min_periods=1)
            .min()
            .values  # .astype(dtype=data.dtype)
        )

        max_data = (
            pd.Series(data - min_data)
            .rolling(window=window, center=True, min_periods=1)
            .max()
            .values  # .astype(dtype=data_type)
        )
    else:
        max_data = (
            pd.Series(data)
            .rolling(window=window, center=True, min_periods=1)
            .max()
            .values  # .astype(dtype=data.dtype)
        )

    # create an histogram of all the possible values of the signal.
    hist, bins = np.histogram(
        max_data,
        bins=np.arange(np.min(max_data), np.max(max_data) + resolution, resolution),
    )

    show = None
    if simple_threshold:
        # find the first minimum in the histogram. This will be the threshold
        threshold = bins[argrelmin(hist)[0][0] + 2]

        if debug:
            f, ax = plt.subplots()
            ax.plot(bins[:-1], hist)
            ax.axvline(threshold)

            f.savefig(debug_output / "histogram.png")
            plt.close(f)
    else:
        _bins = np.arange(len(hist))
        _mean = np.average(_bins, weights=np.square(hist))
        thr = threshold_minimum(hist=hist[: min(int(2 * _mean), len(hist))])

        threshold = bins[int(thr)]
        logger.debug("threshold: %d", threshold)

        if debug:
            if show is None:
                show = np.ones_like(hist)

            f, axs = plt.subplots(2, 1, sharex=True, height_ratios=[3, 1])
            ax = axs[0]
            ax.plot(bins[:-1], hist)
            # ax.axvline(bins[_min], color="C1", label="min")
            ax.axvline(bins[int(_mean)], color="C2", label="mean")
            ax.axvline(threshold, color="C3", label="thr")
            ax.legend()

            ax = axs[1]
            # ax.plot(bins[:-1], _area)
            # ax.set_yscale("log")

            f.savefig(debug_output / "histogram.png")
            plt.close(f)

    signal_threshold = max_data > threshold

    if debug:
        f, ax = plt.subplots()
        debug_last = pause_size + 5 * (pad_size + img_size)
        ax.plot(data[:debug_last])
        ax.plot(max_data[:debug_last] * signal_threshold[:debug_last])
        ax.plot(max_data[:debug_last] * (~signal_threshold[:debug_last]))

        f.savefig(debug_output / "data.png")
        plt.close(f)

        logger.debug("image saved at %s", debug_output.resolve())

    initial_bits = (
        np.diff(signal_threshold, prepend=signal_threshold[0]) & signal_threshold
    )

    initial_bits_index = window // 2 - 1 + np.where((initial_bits))[0]

    if debug:
        logger.debug("initial_bits_index: %s", initial_bits_index)

    samples_initial_bits = (
        np.diff(initial_bits_index, prepend=initial_bits_index[0])
        > (pause_size / 2 - pad_size) * bit_width
    )

    if debug:
        logger.debug(
            "diff: %s",
            np.diff(initial_bits_index, prepend=initial_bits_index[0]),
        )
        logger.debug("const: %d", (pause_size / 2 - pad_size) * bit_width)
        logger.debug("samples_initial_bits: %s", np.any(samples_initial_bits))

    interesting_index = initial_bits_index[samples_initial_bits]

    index_of_interesting_index = np.where(
        np.diff(interesting_index, append=interesting_index[-1])
        > (data_size) * bit_width
    )[0]

    initial_index = interesting_index[index_of_interesting_index - 1] + int(
        bit_width * first_space
    )
    final_index = interesting_index[index_of_interesting_index + 1] - int(
        bit_width * last_space
    )

    if debug:
        logger.debug(
            "diff: %s", np.diff(interesting_index, append=interesting_index[-1])
        )
        logger.debug("const: %d", (data_size) * bit_width)
        logger.debug(
            "interesting_index: %s and %s",
            interesting_index,
            index_of_interesting_index,
        )
        for index in interesting_index:
            f, ax = plt.subplots()
            debug_ii = max(index - int(bit_width * 5 * pause_size), 0)
            debug_fi = min(index + int(bit_width * 5 * pause_size), max_data.size)
            lin = np.arange(debug_ii, debug_fi)
            ax.plot(lin, data[debug_ii:debug_fi], "C0")

            if remove_min:
                ax.plot(
                    lin, data[debug_ii:debug_fi] - min_data[debug_ii:debug_fi], "C3"
                )
            ax.plot(
                lin,
                np.ma.masked_where(
                    signal_threshold[debug_ii:debug_fi], max_data[debug_ii:debug_fi]
                ),
                "C1",
            )
            ax.plot(
                lin,
                np.ma.masked_where(
                    ~signal_threshold[debug_ii:debug_fi], max_data[debug_ii:debug_fi]
                ),
                "C2",
            )

            for i in interesting_index:
                ax.axvline(i)

            ax.set_xlim(lin[0], lin[-1])

            f.savefig(debug_output / f"interesting_{index}.png")
            plt.close(f)

        logger.debug("initial: %s", initial_index)
        logger.debug("final: %s", final_index)

    for ii, fi in zip(initial_index, final_index):
        index = np.linspace(
            ii,
            fi,
            num=n_samples,
            dtype=int,
            endpoint=False,
        )
        w = (fi - ii) / (data_size)
        if (not np.isclose(w, bit_width, rtol=0.001)) or (w > bit_width):
            yield WantedValueError((fi - ii) / (data_size), bit_width)
        else:
            reconstructed_img = data[
                index[:, None] + np.arange(int(img_size * bit_width))
            ]
            yield reconstructed_img