import argparse
import enum
import pathlib
import re
import warnings
from dataclasses import dataclass
from multiprocessing import Process, Queue
from typing import Optional, Tuple, Union

import numpy as np
import scipy.fft
from numpy.typing import NDArray

from fpfind import NP_PRECISEFLOAT
from fpfind.lib import logging
from fpfind.lib.constants import TSRES, PeakFindingFailed
from fpfind.lib.parse_epochs import date2epoch, epoch2int, int2epoch, read_T1, read_T2
from fpfind.lib.parse_timestamps import read_a1_kth_timestamp
from fpfind.lib.typing import (
    Complex_,
    DetectorArray,
    Float,
    Integer,
    PathLike,
    TimestampArray,
)

# Create alias to external dependency for fast histogramming
try:
    from S15lib.g2lib.g2lib import histogram  # noqa: F401 # type: ignore
except ModuleNotFoundError:
    pass

logger, log = logging.get_logger("fpfind")


def get_overlap(*arrays):
    """Returns right-truncated arrays of largest possible common length.

    Used for comparing timestamps of different length, e.g. raw timestamps
    vs chopper-generated individual epochs.
    """
    overlap = min(map(len, arrays))
    arrays = [a[:overlap] for a in arrays]
    return overlap, arrays


@dataclass
class PeakStatistics:
    signal: NDArray[np.number]
    background: NDArray[np.number]

    def has_signal(self):
        return len(self.signal) != 0

    def has_background(self):
        return len(self.background) != 0

    def has_data(self):
        return self.has_signal() and self.has_background()

    @property
    def max(self):
        if not self.has_signal():
            return None
        return np.max(self.signal)

    @property
    def mean(self):
        if not self.has_background():
            return None
        return np.mean(self.background)  # type: ignore (not checked, TODO)

    @property
    def stdev(self):
        if not self.has_background():
            return None
        return np.std(self.background)  # type: ignore (not checked, TODO)

    @property
    def total(self):
        if not self.has_data():
            return None
        return sum(self.signal) - len(self.signal) * self.mean  # type: ignore (None accounted for)

    @property
    def significance(self):
        if self.stdev == 0:
            return 0
        if not self.has_data():
            return 0
        return round((self.max - self.mean) / self.stdev, 3)  # type: ignore (None accounted for)

    @property
    def significance_raw(self):
        if self.stdev == 0:
            return None
        full = np.hstack(self.signal, self.background)  # type: ignore (not checked)
        return (np.max(full) - np.mean(full)) / np.std(full)

    @property
    def significance2(self):
        if self.stdev == 0:
            return None
        # Estimate stdev after grouping in bins of 'len(signal)'
        if not self.has_signal():
            return None
        length = (len(self.background) // len(self.signal)) * len(self.signal)
        rebinned = np.sum(
            self.background[:length].reshape(-1, len(self.signal)), axis=1
        )
        stdev = np.std(rebinned)
        if stdev == 0:
            return None
        return round(self.total / stdev, 3)  # type: ignore (None accounted for)

    @property
    def g2(self):
        if self.mean is None or self.max is None:
            return None
        if self.mean == 0:
            return None
        return self.max / self.mean


def get_statistics(
    hist: NDArray[np.number],
    resolution: Optional[Float] = None,
    center: Optional[Float] = None,
    window: Float = 0.0,
) -> PeakStatistics:
    """Returns statistics of histogram, after performing cross-correlation.

    Args:
        hist: Timing histogram to analyze.
        resolution: Resolution of the histogram.
        center: Timing center of the peak, if known beforehand.
        window: Desired timing window width to exclude from background mean calculation.
    """
    # Fallback to simple statistics, if not other arguments supplied
    if resolution is None:
        if center is None:
            return PeakStatistics(hist, hist)
        else:
            raise ValueError("Resolution must be supplied if 'center' is supplied.")

    # Guess non-negative center bin position, assuming aligned at zero
    if center is None:
        bin_center = np.argmax(hist)
    else:
        bin_center = np.abs(center) // resolution
        if center < 0:
            bin_center = len(hist) - bin_center

    # Retrieve size of symmetrical window
    num_windowbins_onesided = int(np.ceil(window / 2 / resolution))
    bin_offset_left = max(0, bin_center - num_windowbins_onesided)
    bin_offset_right = min(len(hist), bin_center + num_windowbins_onesided)

    # Avoid tails of the cross-correlation by taking only half of the spectrum
    # Use-case when same timestamp is used to obtain histogram, resulting in deadtime
    bin_offset_left_bg = bin_offset_left // 2
    bin_offset_right_bg = (len(hist) + bin_offset_right) // 2

    # Retrieve signal
    signal = hist[bin_offset_left : bin_offset_right + 1]
    background = np.hstack(
        (
            hist[bin_offset_left_bg:bin_offset_left],
            hist[bin_offset_right + 1 : bin_offset_right_bg + 1],
        )
    )
    return PeakStatistics(signal, background)


def generate_fft(
    arr: TimestampArray,
    num_bins: Integer,
    time_res: Float,
) -> NDArray[np.complex128]:
    """Returns the FFT and frequency resolution for the set of timestamps.

    Assumes the inputs are real-valued, i.e. the FFT output is symmetrical.

    Args:
        arr: The timestamp series.
        num_bins: The number of bins in the time/frequency domain.
        bin_size: The size of each timing bin, in ns.

    Note:
        This function is technically not cacheable due to the mutability of
        np.ndarray.
    """
    if len(arr) == 0:
        raise ValueError("Array is empty!")
    bin_arr = np.bincount(np.int32((arr // time_res) % num_bins), minlength=num_bins)
    return scipy.fft.rfft(bin_arr), num_bins  # type: ignore (dispatchable)


def get_xcorr(
    afft: NDArray[Complex_],
    bfft: NDArray[Complex_],
    filter: Optional[NDArray[np.number]] = None,
) -> NDArray[np.float64]:
    """Returns the cross-correlation.

    Note:
        The conjugation operation on an FFT is essentially a time-reversal
        operation on the original time-series data.
    """
    fft = np.conjugate(afft) * bfft
    if filter is not None:
        fft = fft * filter
    result = scipy.fft.irfft(fft)
    return np.abs(result)  # type: ignore (dispatchable)


def convert_histogram_fft(hist: NDArray[np.number], time_bins: NDArray[np.number]):
    """Adjust axes to estimate position."""
    hlen = len(hist) // 2
    hist = np.hstack((hist[hlen:], hist[:hlen]))
    time_bins = np.hstack((-np.flip(time_bins[1 : hlen + 1]), time_bins[:hlen]))
    return hist, time_bins


def get_timing_delay_fft(
    hist: NDArray[np.number],
    time_bins: NDArray[np.number],
    include_negative: bool = False,
) -> Tuple[np.signedinteger, np.signedinteger]:
    """Returns the timing delay.

    Args:
        hist: Timing histogram
        time_bins: Time bins

    Example:
        >>> get_timing_delay_fft([1,3,0,1], [2,4,6,8])
        (4, -8)

    """
    # [0, 1, 2, 3]  -->  (0, -4)
    ppos = np.argmax(hist)
    ptime = time_bins[ppos]
    if ppos == 0:
        npos = 0
    else:
        npos = len(hist) - ppos
    ntime = -time_bins[npos]
    result = (ptime, ntime) if np.abs(ppos) < np.abs(npos) else (ntime, ptime)
    return result


def get_delay_at_index_fft(
    time_bins: NDArray[np.number], pos: Integer
) -> Tuple[np.number, np.number]:
    assert pos >= 0
    pos2 = len(time_bins) - pos if pos != 0 else 0  # corner case
    ptime = time_bins[pos]
    ntime = -time_bins[pos2]
    result = (ptime, ntime) if np.abs(pos) < np.abs(pos2) else (ntime, ptime)
    return result


def get_top_k_delays_fft(
    hist: NDArray[np.number], time_bins: NDArray[np.number], k: Integer
):
    assert k >= 1
    if k == 1:
        return [get_delay_at_index_fft(time_bins, np.argmax(hist))]

    xs_raw = np.argpartition(hist, -k)[-k:]
    ys_raw = hist[xs_raw]
    sort = ys_raw.argsort()[::-1]  # descending order
    result = [
        (y, *get_delay_at_index_fft(time_bins, x))
        for x, y in zip(xs_raw[sort], ys_raw[sort])
    ]
    return result


def slice_timestamps(
    ts: TimestampArray,
    start: Union[Float, None] = None,
    duration: Union[Float, None] = None,
) -> TimestampArray:
    dtype = ts.dtype
    if start is not None:
        ts = ts - start  # note: 'ts -= start' is in-place
        ts = ts[ts >= 0]
    else:
        ts = ts - ts[0]
    if duration is not None:
        if len(ts) == 0:
            warnings.warn("No data available.")
            return np.array([], dtype=dtype)
        if duration >= ts[-1]:
            warnings.warn(
                f"Desired duration ({duration * 1e-9:.3f} s) is longer "
                f"than available data ({ts[-1] * 1e-9:.3f} s)"
            )
        ts = ts[ts < duration]
    return ts


def histogram_fft(  # noqa: PLR0913
    alice: TimestampArray,
    bob: TimestampArray,
    num_bins: Integer,
    resolution: Float = 1,
    num_wraps: Integer = 1,
    acq_start: Optional[Float] = None,
    freq_corr: Float = 0.0,
    filter: Optional[NDArray[np.number]] = None,
    statistics: bool = False,
    center: Optional[Float] = None,
    window: Float = 0.0,
):
    """Returns the cross-correlation histogram.

    Args:
        acq_start: Starting timing, relative to first common timestamp.
        filter: Optional filter in frequency-space.
    """
    if not isinstance(num_wraps, (int, np.integer)):
        warnings.warn(
            "Number of wraps is not an integer - "
            "statistical significance will be lower."
        )

    duration = num_wraps * num_bins * resolution
    first_timestamp = max(alice[0], bob[0])
    last_timestamp = min(alice[-1], bob[-1])
    if first_timestamp + duration > last_timestamp:
        warnings.warn(
            f"Desired duration of timestamps ({duration} ns) "
            f"exceeds available data ({last_timestamp - first_timestamp} ns)."
        )

    # Normalize timestamps
    acq_start = 0 if acq_start is None else acq_start
    alice -= first_timestamp + acq_start
    bob -= first_timestamp + acq_start
    bob = bob * (1 + freq_corr)

    # Generate FFT
    afft, alen = generate_fft(alice, num_bins, resolution)  # TODO: Check if correct
    bfft, blen = generate_fft(bob, num_bins, resolution)  # TODO: Check if correct
    result = get_xcorr(afft, bfft, filter)
    bins = np.arange(num_bins) * resolution
    if statistics:
        return (
            result,
            bins,
            alen,
            blen,
            get_statistics(result, resolution, center, window),
        )
    return result, bins


# https://stackoverflow.com/a/23941599
class ArgparseCustomFormatter(argparse.RawDescriptionHelpFormatter):
    RAW_INDICATOR = "rawtext|"

    def _format_action_invocation(self, action):
        if not action.option_strings:
            _ = self._metavar_formatter(action, action.dest)(1)
            print(action, _)
            (metavar,) = _
            return metavar
        else:
            parts = []
            # if the Optional doesn't take a value, format is:
            #    -s, --long
            if action.nargs == 0:
                parts.extend(action.option_strings)

            # if the Optional takes a value, format is:
            #    -s ARGS, --long ARGS
            # change to
            #    -s, --long ARGS
            else:
                default = action.dest.upper()
                args_string = self._format_args(action, default)
                for option_string in action.option_strings:
                    # parts.append('%s %s' % (option_string, args_string))
                    parts.append("%s" % option_string)
                parts[-1] += " %s" % args_string
            return ", ".join(parts)

    def _split_lines(self, text, width):
        marker = ArgparseCustomFormatter.RAW_INDICATOR
        if text.startswith(marker):
            return text[len(marker) :].splitlines()
        return super()._split_lines(text, width)


def get_first_overlapping_epoch(
    dir1: PathLike,
    dir2: PathLike,
    first_epoch: Optional[Union[str, bytes]] = None,
    return_length: bool = False,
):
    """Get epoch name of smallest overlapping epoch.

    If 'return_length' is True, the return value is a tuple of the epoch name
    and the number of continguous overlapping epochs starting from said epoch.
    """
    epochints1 = [epoch2int(fp.name) for fp in pathlib.Path(dir1).glob("*")]
    epochints2 = [epoch2int(fp.name) for fp in pathlib.Path(dir2).glob("*")]
    epochints = set(epochints1).intersection(epochints2)

    # Exclude epochs smaller than 'first_epoch', if supplied
    if first_epoch is not None:
        epochint_first = epoch2int(first_epoch)
        epochints = set([v for v in epochints if v >= epochint_first])

    # Calculate number of overlapping epochs
    if len(epochints) == 0:
        min_epoch = None
        num_epochs = 0
    else:
        min_epochint = min(epochints)
        min_epoch = int2epoch(min_epochint)
        num_epochs = 1
        while (min_epochint + num_epochs) in epochints:
            num_epochs += 1

    # Return result
    if return_length:
        return min_epoch, num_epochs
    return min_epoch


def iterate_epochs(
    epoch: Union[str, bytes], length: Optional[int] = None, step: int = 1
):
    """Stream incremental epoch names, starting from specified epoch.

    Mainly as a convenience function. This is a generator, so it
    should be fed into a consumer, e.g. `list(iterate_epochs(...))`.
    If 'length' is None, this becomes an infinite stream of epochs.
    The total number of epochs emitted does not change, if 'step' is specified.

    Args:
        epoch: Starting epoch, in hex.
        length: Total number of epochs to emit.
        step: Separation between consecutive epochs.

    Example:
        >>> for epoch in iterate_epochs("bbbbaaa0", length=3, step=2):
        ...     print(epoch)
        bbbbaaa0
        bbbbaaa2
        bbbbaaa4
    """
    epochint = epoch2int(epoch)
    if length is not None:
        for i in range(0, length * step, step):
            yield int2epoch(epochint + i)
    else:
        while True:
            yield int2epoch(epochint)
            epochint += step


def get_timestamp(
    dirname: PathLike,
    file_type: str,
    first_epoch: Union[str, bytes],
    skip_epoch: int,
    num_of_epochs: int,
):
    epochdir = pathlib.Path(dirname)
    timestamp = np.array([], dtype=NP_PRECISEFLOAT)
    for i in range(num_of_epochs):
        epoch_name = epochdir / int2epoch(epoch2int(first_epoch) + skip_epoch + i)
        reader = read_T1 if file_type == "T1" else read_T2
        timestamp = np.append(timestamp, reader(epoch_name)[0])
    return timestamp


def get_timestamp_pattern(
    dirname: PathLike,
    file_type: str,
    first_epoch: Union[str, bytes],
    skip_epoch: int,
    num_of_epochs: int,
) -> Tuple[TimestampArray, DetectorArray]:
    epochdir = pathlib.Path(dirname)
    timestamp = np.array([], dtype=NP_PRECISEFLOAT)
    patterns = np.array([], dtype=np.uint32)
    reader = read_T1 if file_type == "T1" else read_T2
    for i in range(num_of_epochs):
        epoch_name = epochdir / int2epoch(epoch2int(first_epoch) + skip_epoch + i)
        ts, ps = reader(epoch_name)
        timestamp = np.append(timestamp, ts)
        patterns = np.append(patterns, ps)
    return timestamp, patterns


def normalize_timestamps(
    *T: TimestampArray, skip: float = 0.0, preserve_relative: bool = True
):
    """Shifts timestamp arrays to reference zero.

    Args:
        T: List of timestamp arrays.
        skip: Duration to skip, in seconds.
        preserve_relative: Preserve relative durations between arrays.
    """
    if not preserve_relative:
        T = tuple(slice_timestamps(ts) for ts in T)

    start_time = max([ts[0] for ts in T]) + skip * 1e9  # units of ns
    T = tuple(slice_timestamps(ts, start_time) for ts in T)
    return T


def parse_docstring_description(docstring: Optional[str]) -> str:
    if docstring is None:
        return ""

    placeholder = "~~~PLACEHOLDER~~~"
    # Remove all annotated sections, including changelog
    d, *_ = re.split(r"\n[a-zA-Z0-9\s]+:\n", docstring)

    # Replace all newlines except the first
    d = re.sub(r"\n+", placeholder, d, count=1)
    d = re.sub(r"\n+", " ", d)
    d = re.sub(placeholder, "\n\n", d)
    return d


def timestamp2epoch(
    filename: PathLike,
    resolution: TSRES = TSRES.PS4,
    legacy: bool = False,
    full: bool = False,
) -> str:
    (t,), _ = read_a1_kth_timestamp(
        filename,
        [0],
        legacy,
        resolution=resolution,
        fractional=False,
    )
    epoch = epoch2int(date2epoch())  # get current epoch (32-bit)
    # epoch_header = epoch >> 17  # increments every ~20h (first 15 bits)

    if full:
        epochint = read_a1_kth_timestamp(
            filename, [0], legacy=True, resolution=TSRES.PS125, fractional=False
        )
    epochint = (epoch & ~((1 << 17) - 1)) | (int(t) >> 37)  # 54 - 17
    return int2epoch(epochint)


def xcorr(abs: NDArray[np.number], bbs: NDArray[np.number]):
    """Performs the main cross-correlation routine."""
    afft: NDArray[Complex_] = scipy.fft.rfft(abs)  # type: ignore (dispatchable)
    bfft: NDArray[Complex_] = scipy.fft.rfft(bbs)  # type: ignore (dispatchable)
    ys = get_xcorr(afft, bfft)
    return ys


def xcorr_process(abs: NDArray[np.number], bbs: NDArray[np.number]):
    """Runs 'xcorr()' in a separate process for interruptible FFT.

    Note:
        Activate this with caution, because spawning a separate process incurs a penalty
        of pickling and piping inputs/outputs between processes. Rough estimates yield
        20% increase in runtime, e.g. 17.6(2)s -> 20.2(2) for N = 2^27 bins.
    """
    q = Queue()
    args = (q, abs, bbs)

    def _xcorr(q, *args, **kwargs):
        result = xcorr(*args, **kwargs)
        q.put(result)

    p = Process(target=_xcorr, args=args)
    p.start()
    try:
        return q.get()
    except KeyboardInterrupt:  # SIGINT
        p.terminate()
        raise


def generate_fft_bins(
    arr: TimestampArray,
    num_bins: Integer,
    time_res: Float,
) -> NDArray[np.signedinteger]:
    if len(arr) == 0:
        raise ValueError("Array is empty!")
    bin_arr = np.bincount(np.int32((arr // time_res) % num_bins), minlength=num_bins)
    return bin_arr


def histogram_fft2(
    ats: TimestampArray,
    bts: TimestampArray,
    start: Optional[Float],
    duration: Optional[Float],
    N: Integer,
    r: Float,
    interruptible: bool = False,
):
    """Convenience function that wraps histogram routines."""
    ats_early: TimestampArray = slice_timestamps(ats, start, duration)
    bts_early: TimestampArray = slice_timestamps(bts, start, duration)
    abs = generate_fft_bins(ats_early, N, r)
    bbs = generate_fft_bins(bts_early, N, r)
    _xcorr = xcorr_process if interruptible else xcorr
    ys = _xcorr(abs, bbs)

    _dtype = np.int32  # signed number needed for negative delays
    if N * r > 2147483647:  # int32 max
        _dtype = np.int64
    xs = np.arange(N, dtype=_dtype) * r
    return xs, ys


def fold_histogram(
    xs: NDArray[np.number], ys: NDArray[np.number], binning_factor: int = 2
) -> Tuple[NDArray[np.number], NDArray[np.number]]:
    xs = xs[::binning_factor]
    ys = np.sum(ys.reshape(-1, binning_factor), axis=1)
    return xs, ys


class CoarseHistogramStrategy(enum.Enum):
    RESOLUTION = enum.auto()
    BINS = enum.auto()


def histogram_fft3(
    ats: TimestampArray,
    bts: TimestampArray,
    start: Optional[Float],
    duration: Float,
    num_bins: Integer,
    resolution: Float,
    max_duration: Float,
    strategy: CoarseHistogramStrategy = CoarseHistogramStrategy.RESOLUTION,
    interruptible: bool = False,
):
    """Histogram with max duration limit."""
    ceil = lambda v: int(np.ceil(v))  # noqa: E731 (using lambda cleaner)
    factor = ceil(resolution * num_bins / max_duration)
    factor = 1 << ceil(np.log2(factor))  # better equal to 2^k for FFT/reshape
    duration = min(duration, max_duration)
    args = (ats, bts, start, duration, num_bins, resolution, interruptible)

    if factor == 1:
        xs, ys = histogram_fft2(*args)

    elif strategy is CoarseHistogramStrategy.RESOLUTION:
        resolution = resolution / factor
        xs, ys = histogram_fft2(*args)
        xs, ys = fold_histogram(xs, ys, factor)

    elif strategy is CoarseHistogramStrategy.BINS:
        num_bins = int(num_bins // factor)
        xs, ys = histogram_fft2(*args)

    else:
        raise NotImplementedError

    return xs, ys


def match_dts(dt1s_early, dt1s_late, ddt_window):
    dt1s_early = sorted(dt1s_early)
    dt1s_late = sorted(dt1s_late)

    i = j = 0  # two-pointer method
    best = (None, (0, 0, 0), (0, 0, 0))  # min_ddt, early, late
    while i < len(dt1s_early) and j < len(dt1s_late):
        dt1_early = dt1s_early[i]
        dt1_late = dt1s_late[j]
        ddt = abs(dt1_early[0] - dt1_late[0])
        if best[0] is None or (ddt < best[0] and ddt < ddt_window):
            best = (ddt, dt1_early, dt1_late)

        if dt1_early[0] < dt1_late[0]:
            i += 1
        else:
            j += 1

    if best[0] is None:
        raise PeakFindingFailed("No coincident dt")

    # Valid point found
    # Approximate the resolvable resolution by finding the
    # time differences found using the smallest resolution
    def resolve(best_pt, pts) -> Tuple[int, int, float]:
        left = best_pt[0] - best_pt[1]
        right = best_pt[0] + best_pt[1]
        assert len(pts) > 0
        for pt in pts:
            if left <= pt[0] <= right:
                break
        return pt  # pyright: ignore[reportPossiblyUnboundVariable]

    dt1_early = resolve(best[1], dt1s_early)
    dt1_late = resolve(best[2], dt1s_late)

    # Pass control back with expected variables
    r = min([dt1_early[1], dt1_late[1]])
    dt1 = dt1_early[0]
    _dt1 = dt1_late[0]
    return dt1, _dt1, r
