#
# Copyright 2025 Universidad Complutense de Madrid
#
# This file is part of Numina
#
# SPDX-License-Identifier: GPL-3.0-or-later
# License-Filename: LICENSE.txt
#

"""Auxiliary script to perform binary image arithmetic with 3D cubes.

The computation is performed casting to np.float32."""

import argparse
from rich import print
from rich_argparse import RichHelpFormatter
import sys

from astropy.io import fits
import numpy as np
import pathlib

from .add_script_info_to_fits_history import add_script_info_to_fits_history


def compute_operation(file1, file2, operation, extname1, extname2, dtype):
    """Compute output = file1 operation file2.

    Parameters
    ----------
    file1 : file object
        First FITS file.
    file2 : file object or float
        Second FITS file or float number.
    operation : string
        Mathematical operation.
    extname1 : str
        Extension name of the first FITS file (default: 'PRIMARY').
    extname2 : str
        Extension name of the second FITS file (default: 'PRIMARY').
    dtype : str
        Data type of the output image.

    Returns
    -------
    solution : `numpy.ndarray`
        Resulting 3D image after applying the operation.
    """

    # read first FITS file
    with fits.open(file1) as hdulist:
        image_header1 = hdulist[extname1].header
        image1 = hdulist[extname1].data.astype(dtype)
    naxis = image_header1['naxis']
    if naxis != 3:
        raise ValueError("Input file must be a 3D cube (NAXIS=3).")
    naxis1 = image_header1['naxis1']
    naxis2 = image_header1['naxis2']
    naxis3 = image_header1['naxis3']

    # read second FITS file or number
    try:
        with fits.open(file2) as hdulist:
            image_header2 = hdulist[extname2].header
            image2 = hdulist[extname2].data.astype(dtype)
            naxis_ = image_header2['naxis']
            naxis1_ = image_header2['naxis1']
            naxis2_ = image_header2['naxis2']
            naxis3_ = image_header2['naxis3']
    except FileNotFoundError:
        image2 = np.zeros((naxis3, naxis2, naxis1), dtype=dtype)
        if 'int' in dtype:
            file2 = int(file2)
        elif 'float' in dtype:
            file2 = float(file2)
        else:
            raise ValueError(f"Unsupported dtype: {dtype}. Use an integer or float type.")
        image2 += file2  # if file2 is a number
        naxis_ = naxis
        naxis1_ = naxis1
        naxis2_ = naxis2
        naxis3_ = naxis3

    # additional dimension checks
    if naxis != naxis_:
        raise ValueError("NAXIS values are different.")
    if naxis1 != naxis1_:
        raise ValueError("NAXIS1 values are different.")
    if naxis2 != naxis2_:
        raise ValueError("NAXIS2 values are different.")
    if naxis3 != naxis3_:
        raise ValueError("NAXIS3 values are different.")

    # compute operation
    if operation == "+":
        solution = image1 + image2
    elif operation == "-":
        solution = image1 - image2
    elif operation == "x":
        solution = image1 * image2
    elif operation == "/":
        solution = image1 / image2
    elif operation == "=":
        solution = image2.copy()
    else:
        raise ValueError("Unexpected operation=" + str(operation))

    return solution


def main(args=None):

    # parse command-line options
    parser = argparse.ArgumentParser(
        description="description: binary image arithmetic",
        formatter_class=RichHelpFormatter
    )
    # positional parameters
    parser.add_argument("file1",
                        help="First FITS image",
                        type=argparse.FileType('rb'))
    parser.add_argument("operation",
                        help="Arithmetic operation",
                        type=str,
                        choices=['+', '-', 'x', '/', '='])
    parser.add_argument("file2",
                        help="Second FITS image or number",
                        type=str)
    parser.add_argument("output",
                        help="Output FITS image",
                        type=str)
    # optional arguments
    parser.add_argument("--extname1", type=str,
                        help="Extension name of the first FITS file (default: 'PRIMARY').",
                        default='PRIMARY')
    parser.add_argument("--extname2", type=str,
                        help="Extension name of the second FITS file (default: 'PRIMARY').",
                        default='PRIMARY')
    parser.add_argument("--overwrite",
                        help="Overwrite output file if already exists",
                        action="store_true")
    parser.add_argument("--dtype",
                        help="Data type of the output image (default: float32)",
                        type=str,
                        choices=['uint8', 'int8', 'uint16', 'int16', 'uint32', 'int32', 'uint64', 'int64',
                                 'float32', 'float64'],
                        default='float32')
    parser.add_argument("--echo",
                        help="Display full command line",
                        action="store_true")

    args = parser.parse_args(args=args)

    if args.echo:
        print('\033[1m\033[31mExecuting: ' + ' '.join(sys.argv) + '\033[0m\n')

    if not args.overwrite and pathlib.Path(args.output).exists():
        print(f"Output file {args.output} already exists. Use --overwrite to overwrite it.")
        sys.exit(1)

    # compute operation
    solution = compute_operation(
        file1=args.file1,
        file2=args.file2,
        operation=args.operation,
        extname1=args.extname1,
        extname2=args.extname2,
        dtype=args.dtype
    )

    # save output file
    header = fits.getheader(args.file1, args.extname1)
    if header['naxis1'] != solution.shape[2] or \
       header['naxis2'] != solution.shape[1] or \
       header['naxis3'] != solution.shape[0]:
        raise ValueError("Output image dimensions do not match the input image dimensions.")
    hdu = fits.PrimaryHDU(solution.astype(args.dtype), header=header)
    add_script_info_to_fits_history(hdu.header, args)
    hdu.writeto(args.output, overwrite=args.overwrite)


if __name__ == "__main__":

    main()
