#!python
# -*- coding: latin-1 -*-
#
#   Copyright 2016-2021 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
import sys

import numpy as np

import rapidtide.io as tide_io
import rapidtide.workflows.parser_funcs as pf


def summarize(thevoxels, method="mean"):
    theshape = thevoxels.shape
    if len(theshape) > 1:
        numtimepoints = theshape[1]
    else:
        numtimepoints = 1

    if method == "mean":
        themethod = np.mean
    elif method == "sum":
        themethod = np.sum
    elif method == "median":
        themethod = np.median
    else:
        print(f"illegal summary method {method} in summarize")
        sys.exit()

    if numtimepoints > 1:
        regionsummary = np.nan_to_num(themethod(thevoxels, axis=0))
    else:
        regionsummary = np.nan_to_num(themethod(thevoxels))
    return regionsummary


def _get_parser():
    # get the command line parameters
    parser = argparse.ArgumentParser(
        prog="atlasaverage",
        description="Average data within atlas regions.",
        usage="%(prog)s datafile templatefile outputfile [options]",
    )
    parser.add_argument(
        "datafile",
        type=lambda x: pf.is_valid_file(parser, x),
        help="The name of the 3 or 4D nifti file with the data to be averaged over atlas regions.",
    )
    parser.add_argument(
        "templatefile",
        type=lambda x: pf.is_valid_file(parser, x),
        help="The name of the atlas region NIFTI file",
    )
    parser.add_argument("outputfile", help="The name of the output NIFTI file.")

    # add optional arguments
    parser.add_argument(
        "--normmethod",
        dest="normmethod",
        action="store",
        type=str,
        choices=["none", "pct", "var", "std", "p2p"],
        help=(
            "Normalization to apply to input timecourses (in addition to demeaning) prior to "
            "combining. Choices are 'none' (no normalization, default), 'pct' (divide by mean), "
            "'var' (unit variance), 'std' (unit standard deviation), and 'p2p' (unit range)."
        ),
        default="none",
    )
    parser.add_argument(
        "--summarymethod",
        dest="summarymethod",
        action="store",
        type=str,
        choices=["mean", "median"],
        help=("Method to summarize a region.  Choices are 'mean' (default), 'median', and 'sum'."),
        default="mean",
    )
    parser.add_argument(
        "--debug",
        dest="debug",
        action="store_true",
        help="Output additional debugging information.",
        default=False,
    )

    return parser


def main():
    # get the command line parameters
    try:
        args = _get_parser().parse_args()
    except SystemExit:
        _get_parser().print_help()
        raise

    if args.normmethod == "none":
        print("will not normalize timecourses")
    elif args.normmethod == "pct":
        print("will normalize timecourses to percentage of mean")
    elif args.normmethod == "std":
        print("will normalize timecourses to standard deviation of 1.0")
    elif args.normmethod == "var":
        print("will normalize timecourses to variance of 1.0")
    elif args.normmethod == "p2p":
        print("will normalize timecourses to p-p deviation of 1.0")

    print("loading fmri data")
    input_img, input_data, input_hdr, thedims, thesizes = tide_io.readfromnifti(args.datafile)
    if args.debug:
        print("loading template data")
    (
        template_img,
        template_data,
        template_hdr,
        templatedims,
        templatesizes,
    ) = tide_io.readfromnifti(args.templatefile)

    print("checking dimensions")
    if not tide_io.checkspacematch(input_hdr, template_hdr):
        print("template file does not match spatial coverage of input fmri file")
        sys.exit()

    print("reshaping")
    xsize = thedims[1]
    ysize = thedims[2]
    numslices = thedims[3]
    numtimepoints = thedims[4]
    numvoxels = int(xsize) * int(ysize) * int(numslices)
    templatevoxels = np.reshape(template_data, numvoxels).astype(int)
    inputvoxels = np.reshape(input_data, (numvoxels, numtimepoints))
    numregions = np.max(templatevoxels)
    timecourses = np.zeros((numregions, numtimepoints), dtype="float")

    print("reshaped inputdata shape", inputvoxels.shape)

    if numtimepoints > 1:
        print("processing 4D input file")
        for theregion in range(1, numregions + 1):
            theregionvoxels = inputvoxels[np.where(templatevoxels == theregion)[0], :] + 0.0
            print(
                "extracting",
                theregionvoxels.shape[1],
                "voxels from region",
                theregion,
                "of",
                numregions,
            )

            # demean
            themeans = np.mean(theregionvoxels, axis=1)
            theregionvoxels -= themeans[:, None]

            if args.normmethod == "none":
                thenormfac = themeans * 0.0 + 1.0
            elif args.normmethod == "pct":
                thenormfac = themeans
            elif args.normmethod == "var":
                thenormfac = np.var(theregionvoxels, axis=1)
            elif args.normmethod == "std":
                thenormfac = np.std(theregionvoxels, axis=1)
            elif args.normmethod == "p2p":
                thenormfac = np.max(theregionvoxels, axis=1) - np.min(theregionvoxels, axis=1)
            else:
                print("illegal normalization method", args.normmethod)
                _get_parser().print_help()
                raise
            if args.debug:
                print(theregionvoxels.shape, thenormfac.shape)
            for theloc in range(theregionvoxels.shape[0]):
                if thenormfac[theloc] != 0.0:
                    theregionvoxels[theloc, :] /= thenormfac[theloc]
            if theregionvoxels.shape[1] > 0:
                timecourses[theregion - 1, :] = summarize(
                    theregionvoxels, method=args.summarymethod
                )
        if args.debug:
            print("timecourses shape:", timecourses.shape)
        tide_io.writenpvecs(timecourses, args.outputfile)
    else:
        print("processing 3D input file")
        outputvoxels = np.reshape(input_data, (numvoxels, numtimepoints))
        for theregion in range(1, numregions + 1):
            theregionvoxels = inputvoxels[np.where(templatevoxels == theregion)]
            regionval = summarize(theregionvoxels, method=args.summarymethod)
            outputvoxels[np.where(templatevoxels == theregion)] = regionval
        template_hdr["dim"][4] = numregions
        tide_io.savetonifti(
            outputvoxels.reshape((xsize, ysize, numslices, numregions)),
            template_hdr,
            args.outputfile,
        )


if __name__ == "__main__":
    main()
