#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#   Copyright 2016-2026 Blaise Frederick
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
#
import argparse
from typing import Any

import numpy as np

import rapidtide.correlate as tide_corr
import rapidtide.fit as tide_fit
import rapidtide.io as tide_io
import rapidtide.miscmath as tide_math
import rapidtide.resample as tide_resample
import rapidtide.workflows.parser_funcs as pf


def _get_parser() -> Any:
    """
    Construct and return an argument parser for aligning two time series.

    This function sets up an `argparse.ArgumentParser` with required and optional
    arguments for resampling and aligning two time series datasets. It supports
    specifying input files, sample rates, output file, and various processing options
    such as plotting and verbosity.

    Returns
    -------
    argparse.ArgumentParser
        Configured argument parser object with all necessary arguments for
        time series alignment.

    Notes
    -----
    The function uses a custom helper `pf.is_float` to validate sample rate inputs.
    It also adds search range and filter options via `pf.addsearchrangeopts` and
    `pf.addfilteropts`.

    Examples
    --------
    >>> parser = _get_parser()
    >>> args = parser.parse_args()
    """
    # get the command line parameters
    parser = argparse.ArgumentParser(
        prog="aligntcs",
        description="Resample and align two time series.",
        allow_abbrev=False,
    )

    # Required arguments
    parser.add_argument(
        "infile1",
        metavar="infile1[:COLNUM]",
        type=str,
        help="text file containing a timeseries.  Select column COLNUM if multicolumn file.",
    )
    parser.add_argument(
        "insamplerate1",
        type=lambda x: pf.is_float(parser, x),
        help="The input data file (BOLD fmri file or NIRS text file)",
    )
    parser.add_argument(
        "infile2",
        metavar="infile2[:COLNUM]",
        type=str,
        help="text file containing a timeseries.  Select column COLNUM if multicolumn file.",
    )
    parser.add_argument(
        "insamplerate2",
        type=lambda x: pf.is_float(parser, x),
        help="The input data file (BOLD fmri file or NIRS text file)",
    )
    parser.add_argument("outputfile", help="The name of the output file")

    parser.add_argument(
        "--nodisplay",
        dest="displayplots",
        action="store_false",
        help=("Do not plot the data (for noninteractive use)"),
        default=True,
    )
    parser.add_argument(
        "--verbose",
        dest="verbose",
        action="store_true",
        help=("Print out more debugging information"),
        default=False,
    )
    pf.addsearchrangeopts(parser)

    # Filter arguments
    pf.addfilteropts(parser, filtertarget="timecourses")

    return parser


def aligntcs(args: Any) -> None:
    """
    Align two time series using cross-correlation and resampling.

    This function reads two input time series from text files, aligns them based on
    cross-correlation, and writes the aligned second time series to an output file.
    Optional plotting of cross-correlation and aligned signals can be enabled via
    the `displayplots` argument in `args`.

    Parameters
    ----------
    args : Any
        An object containing the following attributes:
        - infile1 : str
            Path to the first input text file.
        - infile2 : str
            Path to the second input text file.
        - insamplerate1 : float
            Sampling rate of the first input signal.
        - insamplerate2 : float
            Sampling rate of the second input signal.
        - outputfile : str
            Path to the output file where the aligned second signal will be written.
        - lagmin : float
            Minimum lag for cross-correlation search.
        - lagmax : float
            Maximum lag for cross-correlation search.
        - displayplots : bool
            If True, displays cross-correlation and aligned signals using matplotlib.

    Returns
    -------
    None
        This function does not return a value but writes the aligned data to a file
        and optionally displays plots.

    Notes
    -----
    - The function applies a prefilter to the input data before alignment.
    - The second time series is resampled to match the timing of the first.
    - Cross-correlation is performed using a fast correlation method.
    - If `displayplots` is True, the function will use the 'TkAgg' backend for matplotlib.

    Examples
    --------
    >>> import argparse
    >>> args = argparse.Namespace(
    ...     infile1='signal1.txt',
    ...     infile2='signal2.txt',
    ...     insamplerate1=100.0,
    ...     insamplerate2=100.0,
    ...     outputfile='aligned_signal2.txt',
    ...     lagmin=-0.1,
    ...     lagmax=0.1,
    ...     displayplots=False
    ... )
    >>> aligntcs(args)
    """
    if args.displayplots:
        import matplotlib as mpl

        mpl.use("TkAgg")
        import matplotlib.pyplot as plt

    args = pf.postprocesssearchrangeopts(args)
    args, theprefilter = pf.postprocessfilteropts(args)

    intimestep1 = 1.0 / args.insamplerate1
    intimestep2 = 1.0 / args.insamplerate2

    inputdata1 = tide_io.readcolfromtextfile(args.infile1)
    inputdata2 = tide_io.readcolfromtextfile(args.infile2)

    # determine waveform lengths
    time1 = args.insamplerate1 * (len(inputdata1) - 1)
    time2 = args.insamplerate2 * (len(inputdata2) - 1)

    fulltime = np.max([time1, time2])
    # pad waveform1 if it's shorter than waveform2
    if time1 < fulltime:
        paddeddata1 = np.zeros(int(np.ceil(fulltime // intimestep1)), dtype=float)
        paddeddata1[0 : len(inputdata1) + 1] = tide_math.corrnormalize(
            theprefilter.apply(args.insamplerate1, inputdata1)
        )
    else:
        paddeddata1 = tide_math.corrnormalize(theprefilter.apply(args.insamplerate1, inputdata1))

    timeaxisfull = np.linspace(
        0.0, intimestep1 * len(paddeddata1), num=len(paddeddata1), endpoint=False
    )
    timeaxis1 = np.linspace(
        0.0, intimestep1 * len(inputdata1), num=len(inputdata1), endpoint=False
    )
    timeaxis2 = np.linspace(
        0.0, intimestep2 * len(inputdata2), num=len(inputdata2), endpoint=False
    )
    paddeddata2 = tide_resample.doresample(
        timeaxis2,
        tide_math.corrnormalize(theprefilter.apply(args.insamplerate2, inputdata2)),
        timeaxisfull,
    )

    # now paddeddata1 and 2 are on the same timescales
    thexcorr = tide_corr.fastcorrelate(paddeddata1, paddeddata2)
    xcorrlen = len(thexcorr)
    xcorr_x = (
        np.r_[0.0:xcorrlen] * intimestep1 - (xcorrlen * intimestep1) / 2.0 + intimestep1 / 2.0
    )

    (
        maxindex,
        maxdelay,
        maxval,
        maxsigma,
        maskval,
        failreason,
        peakstart,
        peakend,
    ) = tide_fit.findmaxlag_gauss(
        xcorr_x,
        thexcorr,
        args.lagmin,
        args.lagmax,
        1000.0,
        refine=True,
        useguess=False,
        fastgauss=False,
        displayplots=False,
    )

    print("Crosscorrelation_Rmax:\t", maxval)
    print("Crosscorrelation_maxdelay:\t", maxdelay)

    # now align the second timecourse to the first

    aligneddata2 = tide_resample.doresample(timeaxis2, inputdata2, timeaxis1 - maxdelay)
    tide_io.writevec(aligneddata2, args.outputfile)

    if args.displayplots:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        # ax.set_title('GCC')
        plt.plot(xcorr_x, thexcorr, "k")
        plt.show()
        fig = plt.figure()
        plt.plot(timeaxis1, inputdata1)
        plt.plot(timeaxis1, aligneddata2)
        plt.plot()
        plt.show()
