# Copyright: 2021-2024 Paul Obermeier (obermeier@poSoft.de)
# Distributed under BSD license.

namespace eval pawt {
    namespace eval fits {
        variable ns [namespace current]

        namespace export ReadImageFile
        namespace export ReadImageHeader
        namespace export WriteImageFile

        #
        # Utility procedures.
        #

        proc _GetNativeByteOrder {} {
            if { $::tcl_platform(byteOrder) eq "littleEndian" } {
                return "Intel"
            } else {
                return "Motorola"
            }
        }

        proc _GetPixelSize { pixelType } {
            if { $pixelType eq "double" } {
                return 8
            } elseif { $pixelType eq "float" } {
                return 4
            } elseif { $pixelType eq "int" } {
                return 4
            } elseif { $pixelType eq "short" } {
                return 2
            } elseif { $pixelType eq "byte" } {
                return 1
            } else {
                error "Read FITS file: Invalid PixelType value: $pixelType (must be byte, short, int, float or double)"
            }
        }

        proc _GetScanFormat { pixelType byteOrder } {
            if { $pixelType eq "double" } {
                if { $byteOrder eq "Intel" } {
                    set scanFmt "q"
                } else {
                    set scanFmt "Q"
                }
            } elseif { $pixelType eq "float" } {
                if { $byteOrder eq "Intel" } {
                    set scanFmt "r"
                } else {
                    set scanFmt "R"
                }
            } elseif { $pixelType eq "int" } {
                if { $byteOrder eq "Intel" } {
                    set scanFmt "iu"
                } else {
                    set scanFmt "Iu"
                }
            } elseif { $pixelType eq "short" } {
                if { $byteOrder eq "Intel" } {
                    set scanFmt "su"
                } else {
                    set scanFmt "Su"
                }
            } elseif { $pixelType eq "byte" } {
                set scanFmt "cu"
            } else {
                error "Read FITS file: Invalid PixelType value: $pixelType (must be byte, short, int, float or double)"
            }
            return $scanFmt
        }

        proc _GetConversionParams { fitsObj } {
            set fmtString ""
            set pixelType ""

            # If the image has BZERO or BSCALE keywords in the header, fitsTcl will
            # do the appropriate thing with them automatically, but the datatype
            # returned will be floating point doubles (isn't FITS fun:)
            if { ( [catch { $fitsObj get keyword BZERO}]  == 0) ||
                 ( [catch { $fitsObj get keyword BSCALE}] == 0) } {
                set fmtString "f*"
                set pixelType "float"
            } else {
                set imgType [$fitsObj info imgType]
                # Note, that 64 bit integer values are not handled by the RAW
                # image format, so the values are interpreted as 32-bit integers
                # and may give wrong results.
                switch -exact -- $imgType {
                      8 { set fmtString "cu*" ; set pixelType "byte" }
                     16 { set fmtString "su*" ; set pixelType "short" }
                     32 { set fmtString "iu*" ; set pixelType "int"   }
                     64 { set fmtString "iu*" ; set pixelType "int" }
                    -32 { set fmtString "f*"  ; set pixelType "float" }
                    -64 { set fmtString "d*"  ; set pixelType "double" }
                }
            }
            return [list $fmtString $pixelType]
        }

        #
        # Image header read/write procedures.
        #

        proc _GetImageHeader { fitsObj fileName } {
            set headerDict [dict create]

            lassign [$fitsObj info imgdim] width height numChan
            if { ! [info exists numChan] || $numChan eq "" } {
                set numChan 1
            }
            if { $numChan == 4 } {
                set numChan 3
            }
            lassign [_GetConversionParams $fitsObj] fmtString pixelType
            set byteOrder [_GetNativeByteOrder]

            # Store the values contained in the header.
            dict set headerDict Magic     "FITS"
            dict set headerDict Width     $width
            dict set headerDict Height    $height
            dict set headerDict NumChan   $numChan
            dict set headerDict ByteOrder $byteOrder
            dict set headerDict ScanOrder "BottomUp"
            dict set headerDict PixelType $pixelType

            # Calculate derived information from the header values.
            set pixelSize [_GetPixelSize  $pixelType]
            set scanFmt   [_GetScanFormat $pixelType $byteOrder]
            dict set headerDict PixelSize  $pixelSize
            dict set headerDict ScanFormat $scanFmt
            dict set headerDict NumPixel   [expr { $width * $height }]
            dict set headerDict NumByte    [expr { $width * $height * $numChan * $pixelSize }]
            dict set headerDict FileName   $fileName

            return $headerDict
        }

        proc ReadImageHeader { fitsImgFile args } {
            # Read the header of a FITS image file.
            #
            # fitsImgFile - File name of the FITS image.
            # args        - Optional format specific arguments. Currently none.
            #
            # See [::pawt::ReadImageHeader] for the default entries of a
            # header dictionary.
            #
            # Returns the header information as a dictionary.
            # If the file does not exist, the header contains invalid information,
            # or if the fitstcl extension is not available, an error is thrown.
            #
            # See also: ReadImageFile WriteImageFile

            if { ! [pawt::HavePkg "fitstcl"] } {
                error "Read FITS header: Missing fitstcl extension"
            }

            set retVal [catch { fits open $fitsImgFile } fitsObj]
            if { $retVal != 0 } {
                error "Read FITS file: Cannot open file $fitsImgFile"
            }

            set headerDict [_GetImageHeader $fitsObj [file tail $fitsImgFile]]
            $fitsObj close
            return $headerDict
        }

        #
        # Image data read/write procedures.
        #

        proc _GetImageData { fitsObj headerDict dataDict } {
            upvar 1 $dataDict myDict

            set blankMapValue 0

            lassign [_GetConversionParams $fitsObj] fmtString pixelType
            set numChan [dict get $headerDict NumChan]
            set width   [dict get $headerDict Width]
            set height  [dict get $headerDict Height]

            # Retrieve the FITS image data as a Tcl list and convert into a byte array.
            set imgAsList [$fitsObj get image]

            # Check, if image has BLANK pixels, which are encoded as "NULL" in the Tcl list.
            # Replace those pixels with replacement value defined in blankMapValue.
            if { [lsearch -exact -ascii $imgAsList "NULL"] >= 0 } {
                set imgAsList [lmap a $imgAsList { expr { $a eq "NULL" ? $blankMapValue : $a }}]
            }

            if { $pixelType eq "short" } {
                # FITS uses signed short values, so map the pixel values to unsigned short for RAW images.
                set imgAsList [lmap p $imgAsList { expr { $p + 32768 }}]
            }
            if { $pixelType eq "int" } {
                # FITS uses signed int values, so map the pixel values to unsigned int for RAW images.
                set imgAsList [lmap p $imgAsList { expr { $p + 2147483648 }}]
            }

            if { $numChan == 3 || $numChan == 4 } {
                # FITS returns RGB images channel by channel, while RAW images need RGB.
                set chanSize [expr { $width * $height }]
                set sr 0
                set er [expr { $chanSize - 1 }]
                set sg [expr { $er + 1 }]
                set eg [expr { 2 * $chanSize - 1 }]
                set sb [expr { $eg + 1 }]
                set eb [expr { 3 * $chanSize - 1 }]
                foreach r [lrange $imgAsList $sr $er] \
                        g [lrange $imgAsList $sg $eg] \
                        b [lrange $imgAsList $sb $eb] {
                    lappend rgbList $r $g $b
                }
                set imgAsList $rgbList
                set numChan 3
            }

            set imgAsByteArray [binary format $fmtString $imgAsList]

            dict set myDict Data $imgAsByteArray
        }

        proc ReadImageFile { fitsImgFile args } {
            # Read a FITS image file into a dictionary.
            #
            # fitsImgFile - File name of the FITS image.
            # args       - Optional format specific arguments. Currently none.
            #
            # Get the image data as a dictionary containing 2 keys: `Header` and `Data`.
            # * See [::pawt::ReadImageHeader] for the entries of the `Header` key.
            # * See [::pawt::ReadImageFile] for the entries of the `Data` key.
            #
            # Returns the image information as a dictionary.
            # If the file does not exist, no supported image format can be detected,
            # or if the fitstcl extension is not available, an error is thrown.
            #
            # See also: ReadImageHeader WriteImageFile

            if { ! [pawt::HavePkg "fitstcl"] } {
                error "Read FITS file: Missing fitstcl extension"
            }
            set retVal [catch { fits open $fitsImgFile } fitsObj]
            if { $retVal != 0 } {
                error "Read FITS file: Cannot open file $fitsImgFile"
            }

            set imgDict [dict create]
            set headerDict [_GetImageHeader $fitsObj [file tail $fitsImgFile]]
            _GetImageData $fitsObj $headerDict imgDict
            dict append imgDict Header $headerDict

            $fitsObj close
            return $imgDict
        }

        proc WriteImageFile { imgDict fitsImgFile args } {
            # Write the values of an image dict into a FITS image file.
            #
            # imgDict     - Image dictionary.
            # fitsImgFile - File name of the image.
            # args        - Optional format specific arguments. Currently none.
            #
            # See [::pawt::WriteImageFile] for general informations regarding writing images.
            #
            # Returns no value. If invalid options are specified or the file could not
            # be written, an error is thrown.
            #
            # See also: ReadImageFile ReadImageHeader

            if { ! [pawt::HavePkg "fitstcl"] } {
                error "Write FITS file: Missing fitstcl extension"
            }

            set headerDict [dict get $imgDict Header]
            set width      [dict get $headerDict Width]
            set height     [dict get $headerDict Height]
            set numChan    [dict get $headerDict NumChan]
            set pixelType  [dict get $headerDict PixelType]
            set pixelSize  [_GetPixelSize  $pixelType]

            # Create new FITS file (mode=2).
            set fitsObj [fits open $fitsImgFile 2]

            if { $pixelType eq "double" } {
                # 64-bit floating point image.
                set bitDepth -64
            } elseif { $pixelType eq "float" } {
                # 32-bit floating point image.
                set bitDepth -32
            } else {
                set bitDepth [expr { $pixelSize * 8 }]
            }
            if { $numChan == 1 } {
                set chansAxes 2
                set chanAxesSizes [list $width $height]
            } elseif { $numChan == 3 } {
                set chansAxes 3
                set chanAxesSizes [list $width $height 3]
            } else {
                error "Write FITS file: Invalid number of channels: $numChan (must be 1 or 3)"
            }
            $fitsObj insert image $bitDepth $chansAxes $chanAxesSizes

            set imgAsList [pawt GetImageAsList imgDict -scanorder "BottomUp"]
            if { $pixelType eq "short" } {
                # Map unsigned short values to signed short.
                set imgAsList [lmap a $imgAsList { expr { $a - 32768 }}]
            }
            if { $pixelType eq "int" } {
                # Map unsigned int values to signed int.
                set imgAsList [lmap a $imgAsList { expr { $a - 2147483648 }}]
            }

            if { $numChan == 3 || $numChan == 4 } {
                # FITS returns RGB images channel by channel, while RAW images need RGB.
                set len [llength $imgAsList]
                for { set r 0 } { $r < $len } { incr r 3 } {
                    lappend chanList [lindex $imgAsList $r]
                }
                for { set g 1 } { $g < $len } { incr g 3 } {
                    lappend chanList [lindex $imgAsList $g]
                }
                for { set b 2 } { $b < $len } { incr b 3 } {
                    lappend chanList [lindex $imgAsList $b]
                }
                set imgAsList $chanList
            }
            $fitsObj put image 1 $imgAsList
            $fitsObj flush
            $fitsObj close
        }
    }
}
