# ==========================================================================
#
#   Copyright NumFOCUS
#
#   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
#
#          https://www.apache.org/licenses/LICENSE-2.0.txt
#
#   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 itk
import sys

if len(sys.argv) < 3:
    print("Usage: " + sys.argv[0] + " inputFile outputFile [numberOfDimensions]")
    sys.exit(1)
input_file = sys.argv[1]
output_file = sys.argv[2]
Dimension = 2
if len(sys.argv) > 3:
    Dimension = int(sys.argv[3])

# Testing GPU Neighborhood Operator Image Filter
InputPixelType = itk.F
OutputPixelType = itk.F

InputImageType = itk.Image[InputPixelType, Dimension]
OutputImageType = itk.Image[OutputPixelType, Dimension]
InputGPUImageType = itk.GPUImage[InputPixelType, Dimension]
OutputGPUImageType = itk.GPUImage[OutputPixelType, Dimension]

input_image = itk.imread(input_file, InputPixelType)
input_gpu_image = itk.cast_image_filter(
    input_image, in_place=False, ttype=(InputImageType, InputGPUImageType)
)
input_gpu_image.UpdateBuffers()

RealOutputPixelType = OutputPixelType

NeighborhoodFilterType = itk.NeighborhoodOperatorImageFilter[
    InputImageType, OutputImageType, RealOutputPixelType
]
GPUNeighborhoodFilterType = itk.GPUNeighborhoodOperatorImageFilter[
    InputGPUImageType, OutputGPUImageType, RealOutputPixelType
]

# Create 1D Gaussian operator
OperatorType = itk.GaussianOperator[RealOutputPixelType, Dimension]

oper = OperatorType()
oper.SetDirection(0)
oper.SetVariance(8.0)
oper.CreateDirectional()

# test 1~8 work units for CPU
for number_of_work_units in range(1, 9):
    cpu_filter = NeighborhoodFilterType.New()
    cpu_timer = itk.TimeProbe()

    cpu_timer.Start()

    cpu_filter.SetNumberOfWorkUnits(number_of_work_units)

    cpu_filter.SetInput(input_image)
    cpu_filter.SetOperator(oper)
    cpu_filter.Update()

    cpu_timer.Stop()

    print(
        "CPU NeighborhoodFilter took {} seconds with {} work units.\n".format(
            cpu_timer.GetMean(), cpu_filter.GetNumberOfWorkUnits()
        )
    )


gpu_filter = GPUNeighborhoodFilterType.New()

gpu_timer = itk.TimeProbe()
gpu_timer.Start()

gpu_filter.SetInput(input_gpu_image)
gpu_filter.SetOperator(oper)
gpu_filter.Update()

gpu_filter.GetOutput().UpdateBuffers()  # synchronization point (GPU->CPU memcpy)

gpu_timer.Stop()

print(f"GPU NeighborhoodFilter took {gpu_timer.GetMean()} seconds.\n")

output_image = itk.cast_image_filter(
    gpu_filter.GetOutput(), ttype=(OutputGPUImageType, OutputImageType)
)
output_gpu_image = gpu_filter.GetOutput()
itk.imwrite(output_image, output_file)
