#!./env/bin/python3
# -*- coding: utf-8 -*-
# Time-stamp: "2024-04-26 20:56:18 (ywatanabe)"


"""
This script does XYZ.
"""


"""
Imports
"""
import os
import sys

import matplotlib.pyplot as plt
import mngs
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F

"""
Config
"""
# CONFIG = mngs.gen.load_configs()


"""
Functions & Classes
"""
import numpy as np


def crop(sig_2d, window_length, overlap_factor=0.0, axis=-1, time=None):
    """
    Crops the input signal into overlapping windows of a specified length,
    allowing for an arbitrary axis and considering a time vector.

    Parameters:
    - sig_2d (numpy.ndarray): The input sig_2d array to be cropped. Can be multi-dimensional.
    - window_length (int): The length of each window to crop the sig_2d into.
    - overlap_factor (float): The fraction of the window that consecutive windows overlap. For example, an overlap_factor of 0.5 means 50% overlap.
    - axis (int): The time axis along which to crop the sig_2d.
        - time (numpy.ndarray): The time vector associated with the signal. Its length should match the signal's length along the cropping axis.

    Returns:
    - cropped_windows (numpy.ndarray): The cropped signal windows. The shape depends on the input shape and the specified axis.
    """
    # Ensure axis is in a valid range
    if axis < 0:
        axis += sig_2d.ndim
    if axis >= sig_2d.ndim or axis < 0:
        raise ValueError(
            "Invalid axis. Axis out of range for sig_2d dimensions."
        )

    if time is not None:
        # Validate the length of the time vector against the signal's dimension
        if sig_2d.shape[axis] != len(time):
            raise ValueError(
                "Length of time vector does not match signal's dimension along the specified axis."
            )

    # Move the target axis to the last position
    axes = np.arange(sig_2d.ndim)
    axes[axis], axes[-1] = axes[-1], axes[axis]
    sig_2d_permuted = np.transpose(sig_2d, axes)

    # Compute the number of windows and the step size
    seq_len = sig_2d_permuted.shape[-1]
    step = int(window_length * (1 - overlap_factor))
    n_windows = max(
        1, ((seq_len - window_length) // step + 1)
    )  # Ensure at least 1 window

    # Crop the sig_2d into windows
    cropped_windows = []
    cropped_times = []
    for i in range(n_windows):
        start = i * step
        end = start + window_length
        cropped_windows.append(sig_2d_permuted[..., start:end])
        if time is not None:
            cropped_times.append(time[start:end])

    # Convert list of windows back to numpy array
    cropped_windows = np.array(cropped_windows)
    cropped_times = np.array(cropped_times)

    # Move the last axis back to its original position if necessary
    if axis != sig_2d.ndim - 1:
        # Compute the inverse permutation
        inv_axes = np.argsort(axes)
        cropped_windows = np.transpose(cropped_windows, axes=inv_axes)

    if time is None:
        return cropped_windows
    else:
        return cropped_windows, cropped_times


def main():
    import random

    FS = 128
    N_CHS = 19
    RECORD_S = 13
    WINDOW_S = 2
    FACTOR = 0.5

    # To pts
    record_pts = int(RECORD_S * FS)
    window_pts = int(WINDOW_S * FS)

    # Demo signal
    sig2d = np.random.rand(N_CHS, record_pts)
    time = np.arange(record_pts) / FS

    # Main
    xx, tt = crop(sig2d, window_pts, overlap_factor=FACTOR, time=time)

    print(f"sig2d.shape: {sig2d.shape}")
    print(f"xx.shape: {xx.shape}")

    # Validation
    i_seg = random.randint(0, len(xx) - 1)
    start = int(i_seg * window_pts * FACTOR)
    end = start + window_pts
    assert np.allclose(sig2d[:, start:end], xx[i_seg])


if __name__ == "__main__":
    # # Argument Parser
    # import argparse
    # parser = argparse.ArgumentParser(description='')
    # parser.add_argument('--var', '-v', type=int, default=1, help='')
    # parser.add_argument('--flag', '-f', action='store_true', default=False, help='')
    # args = parser.parse_args()

    # Main
    CONFIG, sys.stdout, sys.stderr, plt, CC = mngs.gen.start(
        sys, plt, verbose=False
    )
    main()
    mngs.gen.close(CONFIG, verbose=False, notify=False)

# EOF
