tensor/io/fits_table_pins_test.go

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

package io

import (
	"encoding/binary"
	"fmt"
	"math"
	"path/filepath"
	"strings"
	"testing"

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

// Regression pins for the FITS table codecs: the binary and ASCII
// column round trips, the header cards that size them, and the hostile
// counts and widths the reader must refuse by name.

// card renders one 80-byte FITS header card from a full-line body.
func card(body string) []byte {
	return []byte(body + strings.Repeat(" ", 80-len(body)))
}

// cardBlock concatenates cards and pads to the FITS block size.
func cardBlock(cards ...[]byte) []byte {
	var b []byte
	for _, c := range cards {
		b = append(b, c...)
	}
	if rem := len(b) % 2880; rem != 0 {
		b = append(b, make([]byte, 2880-rem)...)
	}
	return b
}

// TestBinaryStringColumnRoundTrip pins the nA binary column contract:
// the reader must accept the repeat-count width the writer emits, so
// SaveFITSTable to LoadFITSTable survives character columns.
func TestBinaryStringColumnRoundTrip(t *testing.T) {
	names := []string{"M31", "NGC 1275", "SMC"}
	cols := []FITSTableColumn{
		{Name: "NAME", Form: "8A", Text: names},
		{Name: "FLUX", Form: "D", Data: mustFloats(t, []float64{1, 2, 3}, 3)},
	}
	path := filepath.Join(t.TempDir(), "names.fits")
	if err := SaveFITSTable(path, false, cols, nil); err != nil {
		t.Fatalf("SaveFITSTable: %v", err)
	}
	table, err := LoadFITSTable(path)
	if err != nil {
		t.Fatalf("LoadFITSTable: %v", err)
	}
	for i, want := range names {
		if table.Text[0][i] != want {
			t.Fatalf("NAME[%d] = %q, want %q", i, table.Text[0][i], want)
		}
	}
}

// TestTableTSCALTZERO pins the scaling keywords: a binary int column
// with TSCAL/TZERO must read back through the affine map, the standard
// unsigned-integer convention included.
func TestTableTSCALTZERO(t *testing.T) {
	raw := []int64{1, 2, 3}
	ints := core.New(core.Int, len(raw))
	copy(ints.RawInts(), raw)
	cols := []FITSTableColumn{
		{Name: "U16", Form: "K", Data: ints},
	}
	path := filepath.Join(t.TempDir(), "scaled.fits")
	headers := map[string]string{"TSCAL1": "1", "TZERO1": "32768"}
	if err := SaveFITSTable(path, false, cols, headers); err != nil {
		t.Fatalf("SaveFITSTable: %v", err)
	}
	table, err := LoadFITSTable(path)
	if err != nil {
		t.Fatalf("LoadFITSTable: %v", err)
	}
	col := table.Columns[0]
	for i, v := range raw {
		want := float64(v) + 32768
		got := col.FloatAt(i)
		if got != want {
			t.Fatalf("U16[%d] = %g, want %g", i, got, want)
		}
	}
	// A fractional scale promotes the column to float.
	if err := SaveFITSTable(path, false, cols, map[string]string{"TSCAL1": "0.5", "TZERO1": "0"}); err != nil {
		t.Fatalf("SaveFITSTable: %v", err)
	}
	table, err = LoadFITSTable(path)
	if err != nil {
		t.Fatalf("LoadFITSTable: %v", err)
	}
	if table.Columns[0].Dtype() != core.Float {
		t.Fatalf("scaled int column dtype %s, want float", table.Columns[0].Dtype())
	}
}

// TestLoadFITSTableSkips16BitImage pins the image-HDU skip arithmetic
// on the case that used to break: a positive BITPIX (16) with a
// non-square shape, where |BITPIX|/8 · Π NAXISn ≠ Σ NAXISn.
func TestLoadFITSTableSkips16BitImage(t *testing.T) {
	// A 5×3 16-bit image (NAXIS1 = 5 is the fastest axis): 15 pixels,
	// 30 bytes of data.
	hdr := cardBlock(
		card("SIMPLE  =                    T"),
		card("BITPIX  =                   16"),
		card("NAXIS   =                    2"),
		card("NAXIS1  =                    5"),
		card("NAXIS2  =                    3"),
		card("END"),
	)
	data := make([]byte, 0, 30)
	for range 15 {
		data = binary.BigEndian.AppendUint16(data, 42)
	}
	image := append(hdr, data...)
	if rem := len(image) % 2880; rem != 0 {
		image = append(image, make([]byte, 2880-rem)...)
	}

	cols := []FITSTableColumn{
		{Name: "V", Form: "K", Data: mustInts(t, []int64{7, 8}, 2)},
	}
	tablePath := filepath.Join(t.TempDir(), "table.fits")
	if err := SaveFITSTable(tablePath, false, cols, nil); err != nil {
		t.Fatalf("SaveFITSTable: %v", err)
	}
	tableBytes, err := osReadFile(tablePath)
	if err != nil {
		t.Fatalf("read table: %v", err)
	}
	// Drop the table file's empty primary HDU (first block), keep the
	// extension behind the hand-built image.
	path := filepath.Join(t.TempDir(), "im_then_table.fits")
	if err := osWriteFile(path, image, tableBytes[2880:]); err != nil {
		t.Fatalf("write: %v", err)
	}
	table, err := LoadFITSTable(path)
	if err != nil {
		t.Fatalf("LoadFITSTable after a 16-bit image: %v", err)
	}
	if table.Rows != 2 || table.Columns[0].RawInts()[1] != 8 {
		t.Fatalf("table landed wrong: rows %d, V[1] = %d", table.Rows, table.Columns[0].RawInts()[1])
	}
}

// TestLoadFITSBSCALE pins image scaling: physical = raw·BSCALE + BZERO.
func TestLoadFITSBSCALE(t *testing.T) {
	hdr := cardBlock(
		card("SIMPLE  =                    T"),
		card("BITPIX  =                  -64"),
		card("NAXIS   =                    1"),
		card("NAXIS1  =                    2"),
		card("BSCALE  =                  2.0"),
		card("BZERO   =                 10.0"),
		card("END"),
	)
	var data []byte
	data = binary.BigEndian.AppendUint64(data, math.Float64bits(1))
	data = binary.BigEndian.AppendUint64(data, math.Float64bits(2))
	path := filepath.Join(t.TempDir(), "scaled.fits")
	var payload []byte
	payload = append(payload, data...)
	if rem := len(payload) % 2880; rem != 0 {
		payload = append(payload, make([]byte, 2880-rem)...)
	}
	if err := osWriteFile(path, hdr, payload); err != nil {
		t.Fatalf("write: %v", err)
	}
	img, _, err := LoadFITS(path)
	if err != nil {
		t.Fatalf("LoadFITS: %v", err)
	}
	if img.FloatAt(0) != 12 || img.FloatAt(1) != 14 {
		t.Fatalf("scaled pixels (%g, %g), want (12, 14)", img.FloatAt(0), img.FloatAt(1))
	}
}

// TestFITSContinueLongString pins the long-string convention: a value
// ending in '&' continues on the CONTINUE cards instead of truncating.
func TestFITSContinueLongString(t *testing.T) {
	hdr := cardBlock(
		card("SIMPLE  =                    T"),
		card("BITPIX  =                    8"),
		card("NAXIS   =                    0"),
		card(fmt.Sprintf("%-8s= '%s&'", "LONGKEY", "the first part of a long value ")),
		card(fmt.Sprintf("CONTINUE  '%s'", "and the second part")),
		card("END"),
	)
	path := filepath.Join(t.TempDir(), "long.fits")
	if err := osWriteFile(path, hdr); err != nil {
		t.Fatalf("write: %v", err)
	}
	_, headers, err := LoadFITS(path)
	if err != nil {
		t.Fatalf("LoadFITS: %v", err)
	}
	want := "the first part of a long value and the second part"
	if headers["LONGKEY"] != want {
		t.Fatalf("LONGKEY = %q, want %q", headers["LONGKEY"], want)
	}
}