__all__ = ["build_template"]

import numpy as np

import os

import shutil

from tempfile import mktemp

import ants

def build_template(

    initial_template=None,

    image_list=None,

    iterations=3,

    gradient_step=0.2,

    blending_weight=0.75,

    weights=None,

    useNoRigid=True,

    output_dir=None,

    **kwargs

):

    """

    Estimate an optimal template from an input image_list

    ANTsR function: N/A

    Arguments

    ---------

    initial_template : ANTsImage

        initialization for the template building

    image_list : ANTsImages

        images from which to estimate template

    iterations : integer

        number of template building iterations

    gradient_step : scalar

        for shape update gradient

    blending_weight : scalar

        weight for image blending

    weights : vector

        weight for each input image

    useNoRigid : boolean

        equivalent of -y in the script. Template update

        step will not use the rigid component if this is True.

    output_dir : path

        directory name where intermediate transforms are written

    kwargs : keyword args

        extra arguments passed to ants registration

    Returns

    -------

    ANTsImage

    Example

    -------

    >>> import ants

    >>> image = ants.image_read( ants.get_ants_data('r16') )

    >>> image2 = ants.image_read( ants.get_ants_data('r27') )

    >>> image3 = ants.image_read( ants.get_ants_data('r85') )

    >>> timage = ants.build_template( image_list = ( image, image2, image3 ) ).resample_image( (45,45))

    >>> timagew = ants.build_template( image_list = ( image, image2, image3 ), weights = (5,1,1) )

    """

    work_dir = mktemp() if output_dir is None else output_dir

    def make_outprefix(k: int):

        os.makedirs(os.path.join(work_dir, f"img{k:04d}"), exist_ok=True)

        return os.path.join(work_dir, f"img{k:04d}", "out")

    if "type_of_transform" not in kwargs:

        type_of_transform = "SyN"

    else:

        type_of_transform = kwargs.pop("type_of_transform")

    if weights is None:

        weights = np.repeat(1.0 / len(image_list), len(image_list))

    weights = [x / sum(weights) for x in weights]

    if initial_template is None:

        initial_template = image_list[0] * 0

        for i in range(len(image_list)):

            temp = image_list[i] * weights[i]

            temp = ants.resample_image_to_target(temp, initial_template)

            initial_template = initial_template + temp

    xavg = initial_template.clone()

    for i in range(iterations):

        affinelist = []

        for k in range(len(image_list)):

            w1 = ants.registration(

                xavg, image_list[k], type_of_transform=type_of_transform, outprefix=make_outprefix(k), **kwargs

            )

            L = len(w1["fwdtransforms"])

            # affine is the last one

            affinelist.append(w1["fwdtransforms"][L-1])

            if k == 0:

                if L == 2:

                    wavg = ants.image_read(w1["fwdtransforms"][0]) * weights[k]

                xavgNew = w1["warpedmovout"] * weights[k]

            else:

                if L == 2:

                    wavg = wavg + ants.image_read(w1["fwdtransforms"][0]) * weights[k]

                xavgNew = xavgNew + w1["warpedmovout"] * weights[k]

        if useNoRigid:

            avgaffine = ants.average_affine_transform_no_rigid(affinelist)

        else:

            avgaffine = ants.average_affine_transform(affinelist)

        afffn = os.path.join(work_dir, "avgAffine.mat")

        ants.write_transform(avgaffine, afffn)

        if L == 2:

            print(wavg.abs().mean())

            wscl = (-1.0) * gradient_step

            wavg = wavg * wscl

            # apply affine to the nonlinear?

            # need to save the average

            wavgA = ants.apply_transforms(fixed=xavgNew, moving=wavg, imagetype=1, transformlist=afffn, whichtoinvert=[1])

            wavgfn = os.path.join(work_dir, "avgWarp.nii.gz")

            ants.image_write(wavgA, wavgfn)

            xavg = ants.apply_transforms(fixed=xavgNew, moving=xavgNew, transformlist=[wavgfn, afffn], whichtoinvert=[0, 1])

        else:

            xavg = ants.apply_transforms(fixed=xavgNew, moving=xavgNew, transformlist=[afffn], whichtoinvert=[1])

        if blending_weight is not None:

            xavg = xavg * blending_weight + ants.iMath(xavg, "Sharpen") * (

                1.0 - blending_weight

            )

    if output_dir is None:

        shutil.rmtree(work_dir)

    return xavg