tensor/io/fits.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package io

import (
	"encoding/binary"
	"fmt"
	"maps"
	"math"
	"os"
	"slices"
	"strconv"
	"strings"

	"sourcedock.dev/petrbalvin/tensor/internal/base"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// FITS image I/O. The Flexible Image Transport System is astronomy's
// archival format: a self-describing header of 80-character ASCII
// cards in 2880-byte blocks, followed by big-endian binary data
// padded to the same block size. Every observation archived by a
// telescope in the last four decades reads back with the same parser,
// which is the property that makes a format worth speaking.
//
// This implementation covers the primary HDU image: BITPIX -64
// (float64) and -32 (float32), any rank with positive axes. FITS
// orders axes Fortran-style with NAXIS1 varying fastest, the opposite
// of Go's row-major convention, so NAXISj is declared from the shape
// in reverse and the flat payload needs no permutation. Extensions
// (XTENSION), tables and the integer bit depths are refused with an
// error rather than half-read.

// fitsCardsPerBlock is the number of 80-byte cards in one 2880-byte
// FITS block.
const fitsCardsPerBlock = 2880 / 80

// SaveFITS writes a float64 or float32 array as a FITS primary image
// with the given header entries. Keywords are uppercased and must be
// 1 to 8 characters from A-Z, 0-9, '-' and '_' (the format's reserved
// SIMPLE, BITPIX, NAXIS, NAXISn, EXTEND and END are refused); values
// are written as FITS strings, at most 68 characters after the
// format's quote escaping.
func SaveFITS(path string, a *core.Array, headers map[string]string) error {
	var bitpix int
	switch a.Dtype() {
	case core.Float:
		bitpix = -64
	case core.Float32:
		bitpix = -32
	default:
		return base.Errf("SaveFITS: supports float64 and float32 arrays, got dtype %s", a.Dtype())
	}
	if a.NDim() == 0 {
		return base.Errf("SaveFITS: the image needs at least one axis")
	}
	for i, d := range a.Shape() {
		if d <= 0 {
			return base.Errf("SaveFITS: axis %d has extent %d, every axis must be positive", i+1, d)
		}
	}

	cards := []string{
		fitsBoolCard("SIMPLE", true),
		fitsIntCard("BITPIX", bitpix),
		fitsIntCard("NAXIS", a.NDim()),
	}
	// NAXIS1 is the fastest-varying axis, the last one in Go's
	// row-major order.
	for j := range a.NDim() {
		cards = append(cards, fitsIntCard("NAXIS"+strconv.Itoa(j+1), a.Shape()[a.NDim()-1-j]))
	}
	cards = append(cards, fitsBoolCard("EXTEND", true))
	userCards, err := fitsUserCards(headers)
	if err != nil {
		return base.Errf("SaveFITS: %w", err)
	}
	cards = append(cards, userCards...)
	cards = append(cards, fitsEndCard())

	elem := a.Len()
	width := 8
	if a.Dtype() == core.Float32 {
		width = 4
	}
	// The final size is known up front: both the header and the payload
	// pad to whole blocks, so one allocation serves the whole file.
	out := make([]byte, 0, fitsBlockSize(len(cards))+fitsBlockSize(elem*width))
	out = fitsAppendCards(out, cards)
	if a.Dtype() == core.Float {
		raw := a.RawFloats()
		for i := range elem {
			out = binary.BigEndian.AppendUint64(out, math.Float64bits(raw[i]))
		}
	} else {
		raw := a.RawFloat32s()
		for i := range elem {
			out = binary.BigEndian.AppendUint32(out, math.Float32bits(raw[i]))
		}
	}
	// Zero bytes pad the data to the block boundary.
	out = fitsAppendZeroPad(out)
	return os.WriteFile(path, out, 0o644)
}

// LoadFITS reads a FITS primary image into a float64 (BITPIX -64) or
// float32 (BITPIX -32) array, returning every non-structural header
// entry alongside it. String values are unquoted and unescaped,
// logical values come back as "T" or "F", numbers as their literal
// text; COMMENT, HISTORY and blank cards carry no value and are
// skipped.
func LoadFITS(path string) (*core.Array, map[string]string, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, nil, base.Errf("LoadFITS: %w", err)
	}
	return parseFITS(data)
}

// parseFITS decodes a FITS primary image from raw bytes.
func parseFITS(data []byte) (*core.Array, map[string]string, error) {
	if len(data) < 80 {
		return nil, nil, base.Errf("LoadFITS: file is shorter than one header card")
	}
	if string(data[:9]) == "XTENSION " {
		return nil, nil, base.Errf("LoadFITS: extensions are not supported, only the primary image")
	}
	if strings.TrimRight(string(data[:8]), " ") != "SIMPLE" {
		return nil, nil, base.Errf("LoadFITS: the first card must be SIMPLE")
	}

	cards, dataAt, err := scanFITSCards(data, 0)
	if err != nil {
		return nil, nil, base.Errf("LoadFITS: %w", err)
	}

	var (
		headers  = map[string]string{}
		bitpix   int
		naxis    = -1
		axisVals = map[int]int{}
		dims     []int
	)
	for _, c := range cards {
		switch {
		case c.key == "SIMPLE":
			if c.value != "T" {
				return nil, nil, base.Errf("LoadFITS: SIMPLE = F marks a non-conformant file")
			}
		case c.key == "EXTEND":
			// Structural; EXTEND still reports itself to the caller.
			headers[c.key] = c.value
		case c.key == "BITPIX":
			v, cerr := strconv.Atoi(c.value)
			if cerr != nil {
				return nil, nil, base.Errf("LoadFITS: %s = %q is not an integer", c.key, c.value)
			}
			bitpix = v
		case c.key == "NAXIS":
			v, cerr := strconv.Atoi(c.value)
			if cerr != nil {
				return nil, nil, base.Errf("LoadFITS: %s = %q is not an integer", c.key, c.value)
			}
			naxis = v
		default:
			// NAXISn is an axis only when a number follows the prefix,
			// the rule fitsCheckKeyword applies on write kept symmetric;
			// NAXISREF and every other spelling is a user keyword. The
			// value is stored under its axis number, so the card order
			// cannot re-bind the axes.
			if fitsAxisKeyword(c.key) {
				v, cerr := strconv.Atoi(c.value)
				if cerr != nil {
					return nil, nil, base.Errf("LoadFITS: %s = %q is not an integer", c.key, c.value)
				}
				j, _ := strconv.Atoi(c.key[len("NAXIS"):])
				if j < 1 {
					return nil, nil, base.Errf("LoadFITS: %s is not an axis keyword", c.key)
				}
				if _, dup := axisVals[j]; dup {
					return nil, nil, base.Errf("LoadFITS: %s repeats", c.key)
				}
				axisVals[j] = v
				continue
			}
			headers[c.key] = c.value
		}
	}
	// A zero-axis primary HDU (the standard container for extension
	// files) answers an empty image.
	if naxis == 0 {
		return core.New(core.Float, 0), headers, nil
	}
	if bitpix != -64 && bitpix != -32 {
		return nil, nil, base.Errf("LoadFITS: BITPIX %d is not supported (want -64 or -32)", bitpix)
	}
	// NAXISn values are bound by their axis number, not by card order:
	// a header writing NAXIS2 before NAXIS1, or omitting an axis, is
	// malformed and must be refused rather than silently re-bound.
	if naxis < 0 {
		return nil, nil, base.Errf("LoadFITS: NAXIS = %d is negative or the card is missing", naxis)
	}
	// The card count gates the allocation: a hostile NAXIS far beyond
	// the NAXISn cards the file actually carries is refused here, not
	// turned into a slice of that length.
	if len(axisVals) != naxis {
		return nil, nil, base.Errf("LoadFITS: NAXIS = %d with %d NAXISn cards", naxis, len(axisVals))
	}
	dims = make([]int, naxis)
	for j := 1; j <= naxis; j++ {
		d, ok := axisVals[j]
		if !ok {
			return nil, nil, base.Errf("LoadFITS: NAXIS%d is missing under NAXIS = %d", j, naxis)
		}
		dims[j-1] = d
	}
	width := -bitpix / 8
	if dataAt > len(data) {
		return nil, nil, base.Errf("LoadFITS: data is truncated (%d bytes present)", len(data))
	}
	// Every axis extent and running product is bounded by the bytes
	// the file actually holds, so a hostile header cannot overflow the
	// int product before the truncation check rejects it.
	avail := (len(data) - dataAt) / width
	shape := make([]int, naxis)
	total := 1
	for i, d := range dims {
		if d <= 0 {
			return nil, nil, base.Errf("LoadFITS: NAXIS%d = %d, every axis must be positive", i+1, d)
		}
		if d > avail || total > avail/d {
			return nil, nil, base.Errf("LoadFITS: data is truncated (%d bytes present, more needed)",
				len(data)-dataAt)
		}
		shape[naxis-1-i] = d
		total *= d
	}

	payload := data[dataAt : dataAt+total*width]
	var out *core.Array
	if bitpix == -64 {
		out = core.New(core.Float, shape...)
		raw := out.RawFloats()
		for i := range total {
			raw[i] = math.Float64frombits(binary.BigEndian.Uint64(payload[i*8:]))
		}
	} else {
		out = core.New(core.Float32, shape...)
		raw := out.RawFloat32s()
		for i := range total {
			raw[i] = math.Float32frombits(binary.BigEndian.Uint32(payload[i*4:]))
		}
	}
	// BSCALE/BZERO scaling: physical = raw*scale + zero (FITS 4.1).
	// A silent skip would hand back storage values as physical ones, so
	// the affine map is applied whenever the keywords deviate from the
	// identity; float32 values are computed in float64 and rounded once.
	scale, serr := fitsScaledHeader(headers, "BSCALE", 1)
	if serr != nil {
		return nil, nil, base.Errf("LoadFITS: %w", serr)
	}
	zero, zerr := fitsScaledHeader(headers, "BZERO", 0)
	if zerr != nil {
		return nil, nil, base.Errf("LoadFITS: %w", zerr)
	}
	if scale != 1 || zero != 0 {
		if out.Dtype() == core.Float {
			raw := out.RawFloats()
			for i := range raw {
				raw[i] = raw[i]*scale + zero
			}
		} else {
			raw := out.RawFloat32s()
			for i := range raw {
				raw[i] = float32(float64(raw[i])*scale + zero)
			}
		}
	}
	return out, headers, nil
}

// fitsScaledHeader parses a floating-point header entry, falling back
// to def when the keyword is absent. A present but malformed value is
// an error, never silently the default.
func fitsScaledHeader(headers map[string]string, key string, def float64) (float64, error) {
	v, ok := headers[key]
	if !ok || v == "" {
		return def, nil
	}
	f, err := strconv.ParseFloat(strings.TrimSpace(v), 64)
	if err != nil {
		return 0, base.Errf("%s = %q is not a number", key, v)
	}
	return f, nil
}

// fitsValue extracts the value field of a card: everything after the
// "= " indicator, minus any trailing comment, with FITS string
// quoting resolved. The second return says whether the value was a
// quoted string, which the long-string CONTINUE convention keys on.
func fitsValue(card string) (string, bool, error) {
	field := card[10:]
	if strings.HasPrefix(strings.TrimLeft(field, " "), "'") {
		// A quoted string: '' inside escapes one quote.
		var b strings.Builder
		in := field[strings.Index(field, "'"):]
		i := 1
		for i < len(in) {
			if in[i] == '\'' {
				if i+1 < len(in) && in[i+1] == '\'' {
					b.WriteByte('\'')
					i += 2
					continue
				}
				return strings.TrimRight(b.String(), " "), true, nil
			}
			b.WriteByte(in[i])
			i++
		}
		return "", false, base.Errf("card %q has an unterminated string value", card[:min(20, len(card))])
	}
	// Free-format value: cut at the comment slash and trim.
	if slash := strings.IndexByte(field, '/'); slash >= 0 {
		field = field[:slash]
	}
	return strings.TrimSpace(field), false, nil
}

// fitsContinueString extracts the quoted segment of a CONTINUE card.
// The keyword occupies columns 1-8 and no "= " indicator follows, so
// the string opens at the first quote anywhere in the card.
func fitsContinueString(card string) (string, error) {
	q := strings.IndexByte(card, '\'')
	if q < 0 {
		return "", base.Errf("CONTINUE card %q has no quoted segment", card[:min(20, len(card))])
	}
	var b strings.Builder
	i := q + 1
	for i < len(card) {
		if card[i] == '\'' {
			if i+1 < len(card) && card[i+1] == '\'' {
				b.WriteByte('\'')
				i += 2
				continue
			}
			return b.String(), nil
		}
		b.WriteByte(card[i])
		i++
	}
	return "", base.Errf("CONTINUE card %q has an unterminated string", card[:min(20, len(card))])
}

// fitsCheckKeyword validates a header keyword the caller supplies.
// The format's commentary keywords COMMENT and HISTORY carry no
// value, so they are refused like the structural ones.
func fitsCheckKeyword(kw string) error {
	if len(kw) < 1 || len(kw) > 8 {
		return base.Errf("keyword %q must be 1 to 8 characters", kw)
	}
	for _, r := range kw {
		switch {
		case r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '-', r == '_':
		default:
			return base.Errf("keyword %q may only contain A-Z, 0-9, '-' and '_'", kw)
		}
	}
	switch kw {
	case "SIMPLE", "BITPIX", "NAXIS", "EXTEND", "END", "COMMENT", "HISTORY":
		return base.Errf("keyword %q is reserved by the format", kw)
	case "BSCALE", "BZERO":
		// SaveFITS writes physical values directly; a user scaling card
		// would make conforming readers scale them a second time.
		return base.Errf("keyword %q is reserved by the format (values are written unscaled)", kw)
	}
	if rest, ok := strings.CutPrefix(kw, "NAXIS"); ok && rest != "" {
		if _, err := strconv.Atoi(rest); err == nil {
			return base.Errf("keyword %q is reserved by the format", kw)
		}
	}
	return nil
}

// fitsAxisKeyword reports whether key is a structural NAXISn card: the
// prefix followed by a number, the same rule fitsCheckKeyword applies
// when it refuses a reserved keyword on write. Any other spelling of
// the prefix (NAXISREF and the like) is a user keyword, and counting it
// as an axis used to make the reader answer "NAXIS = 1 with 2 NAXISn
// cards" for a file that carries one.
func fitsAxisKeyword(key string) bool {
	rest, ok := strings.CutPrefix(key, "NAXIS")
	if !ok {
		return false
	}
	_, err := strconv.Atoi(rest)
	return err == nil
}

// fitsLookupN reads headers[prefix+index] without building the key on
// the heap: the digits are formatted into a stack buffer, and the map
// lookup over that buffer compiles to a lookup over the bytes with no
// string conversion.
func fitsLookupN(headers map[string]string, prefix string, index int) string {
	var kb [24]byte
	b := append(kb[:0], prefix...)
	b = strconv.AppendInt(b, int64(index), 10)
	return headers[string(b)]
}

// fitsIntCard renders an integer-valued card with the value
// right-justified in columns 11 to 30.
func fitsIntCard(keyword string, v int) string {
	return fitsPadCard(fmt.Sprintf("%-8s= %20d", keyword, v))
}

// fitsBoolCard renders a logical-valued card.
func fitsBoolCard(keyword string, v bool) string {
	t := "F"
	if v {
		t = "T"
	}
	return fitsPadCard(fmt.Sprintf("%-8s= %20s", keyword, t))
}

// fitsStringCard renders a string-valued card; the format pads the
// quoted value to at least eight characters.
func fitsStringCard(keyword, v string) (string, error) {
	escaped := strings.ReplaceAll(v, "'", "''")
	inner := escaped
	if len(inner) < 8 {
		inner += strings.Repeat(" ", 8-len(inner))
	}
	body := fmt.Sprintf("%-8s= '%s'", keyword, inner)
	if len(body) > 80 {
		return "", base.Errf("value for %q does not fit a card after quote escaping (%d characters)",
			keyword, len(escaped))
	}
	return fitsPadCard(body), nil
}

// fitsEndCard renders the header terminator.
func fitsEndCard() string {
	return fitsPadCard("END")
}

// fitsPadCard right-pads a card body with spaces to the full 80 bytes.
func fitsPadCard(body string) string {
	return body + strings.Repeat(" ", 80-len(body))
}

// fitsUserCards renders the caller's header entries as cards, keywords
// uppercased and sorted so the output is deterministic.
func fitsUserCards(headers map[string]string) ([]string, error) {
	cards := make([]string, 0, len(headers))
	for _, key := range slices.Sorted(maps.Keys(headers)) {
		kw := strings.ToUpper(key)
		if err := fitsCheckKeyword(kw); err != nil {
			return nil, err
		}
		card, err := fitsStringCard(kw, headers[key])
		if err != nil {
			return nil, err
		}
		cards = append(cards, card)
	}
	return cards, nil
}

// fitsCard is one parsed header card: its keyword and the value field
// with quoting resolved.
type fitsCard struct {
	key, value string
}

// scanFITSCards walks the 80-byte cards of one header starting at off
// and returns every value card up to the END terminator together with
// the block-aligned header length in bytes, which is where the data
// block begins. Blank, COMMENT and HISTORY cards carry no value and
// are skipped before the value-indicator check, because a commentary
// card may legitimately carry an "= " sequence in columns 9-10; a
// header that never terminates is an error.
func scanFITSCards(data []byte, off int) ([]fitsCard, int, error) {
	var cards []fitsCard
	seen := 0
	for pos := off; pos+80 <= len(data); pos += 80 {
		line := string(data[pos : pos+80])
		seen++
		key := strings.TrimRight(line[:8], " ")
		if key == "END" {
			return cards, (seen + fitsCardsPerBlock - 1) / fitsCardsPerBlock * 2880, nil
		}
		// Commentary cards have no value regardless of what follows
		// columns 9-10; check them before the "= " indicator.
		if key == "" || key == "COMMENT" || key == "HISTORY" {
			continue
		}
		if key == "CONTINUE" {
			// A CONTINUE card only makes sense after an open long
			// string; reaching one here means the base card never
			// ended in '&', and silently dropping it would lose the
			// caller's value.
			return nil, 0, base.Errf("CONTINUE card without an open long string")
		}
		if line[8:10] != "= " {
			continue
		}
		value, wasString, err := fitsValue(line)
		if err != nil {
			return nil, 0, err
		}
		// Long-string convention: a string value ending in '&' is
		// continued by the following CONTINUE cards, each carrying the
		// next quoted segment. Dropping them would truncate the value
		// at the card boundary.
		for wasString && strings.HasSuffix(value, "&") && pos+160 <= len(data) {
			next := string(data[pos+80 : pos+160])
			if strings.TrimSpace(next[:8]) != "CONTINUE" {
				break
			}
			pos += 80
			seen++
			cont, cerr := fitsContinueString(next)
			if cerr != nil {
				return nil, 0, cerr
			}
			value = strings.TrimSuffix(value, "&") + strings.TrimRight(cont, " ")
		}
		cards = append(cards, fitsCard{key, value})
	}
	return nil, 0, base.Errf("no END card terminates the header")
}

// fitsBlockSize returns n rounded up to a whole FITS block.
func fitsBlockSize(n int) int {
	return (n + 2879) / 2880 * 2880
}

// fitsAppendCards appends the rendered cards as one header block,
// blank-padded to the block boundary.
func fitsAppendCards(dst []byte, cards []string) []byte {
	for _, card := range cards {
		dst = append(dst, card...)
	}
	// Blank cards pad the header to the block boundary.
	if rem := len(dst) % 2880; rem != 0 {
		blank := fitsPadCard("")
		for i := 0; i < 2880-rem; i += 80 {
			dst = append(dst, blank...)
		}
	}
	return dst
}

// fitsAppendZeroPad appends zero bytes up to the block boundary.
func fitsAppendZeroPad(dst []byte) []byte {
	if rem := len(dst) % 2880; rem != 0 {
		dst = append(dst, make([]byte, 2880-rem)...)
	}
	return dst
}