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

package io_test

// The godoc examples: one runnable, checked snippet per format the
// package speaks. `go test` executes them, so the documentation cannot
// rot. Each one writes into a fresh temporary directory and removes it
// again.

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"

	tensor "sourcedock.dev/petrbalvin/tensor"
	"sourcedock.dev/petrbalvin/tensor/io"
)

// tempDir makes a fresh temporary directory for one example; the
// example removes it with a deferred os.RemoveAll.
func tempDir() string {
	dir, err := os.MkdirTemp("", "tensor-io-example-")
	if err != nil {
		log.Fatal(err)
	}
	return dir
}

// A 2-D array goes out as comma-separated text and reads back as the
// same numbers.
func ExampleSaveCSV() {
	dir := tempDir()
	defer os.RemoveAll(dir)
	path := filepath.Join(dir, "readings.csv")

	a, err := tensor.FromFloats([]float64{18.5, 21.25, 19.75, 23, 20.5, 17.25}, 2, 3)
	if err != nil {
		log.Fatal(err)
	}
	if err := io.SaveCSV(path, a); err != nil {
		log.Fatal(err)
	}
	back, err := io.LoadCSV(path, false)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("shape:", back.Shape())
	fmt.Println("first:", back.FloatAt(0), "last:", back.FloatAt(back.Len()-1))
	// Output:
	// shape: [2 3]
	// first: 18.5 last: 17.25
}

// The stream form writes CSV to any writer and reads it from any
// reader, skipping a header row on the way back.
func ExampleSaveCSVWriter() {
	a, err := tensor.FromFloats([]float64{1, 2, 3, 4}, 2, 2)
	if err != nil {
		log.Fatal(err)
	}
	var buf strings.Builder
	if err := io.SaveCSVWriter(&buf, a); err != nil {
		log.Fatal(err)
	}
	fmt.Print(buf.String())

	back, err := io.LoadCSVReader(strings.NewReader("row,c1,c2\n"+buf.String()), true)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("shape:", back.Shape(), "last:", back.FloatAt(3))
	// Output:
	// 1,2
	// 3,4
	// shape: [2 2] last: 4
}

// A float64 image goes out as a FITS primary image with header cards
// and reads back with the cards beside the values.
func ExampleSaveFITS() {
	dir := tempDir()
	defer os.RemoveAll(dir)
	path := filepath.Join(dir, "image.fits")

	a, err := tensor.FromFloats([]float64{1, 2.5, 3, 4, 5, 6}, 2, 3)
	if err != nil {
		log.Fatal(err)
	}
	headers := map[string]string{"OBJECT": "M31", "EXPTIME": "600"}
	if err := io.SaveFITS(path, a, headers); err != nil {
		log.Fatal(err)
	}
	img, header, err := io.LoadFITS(path)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("shape:", img.Shape(), "dtype:", img.Dtype())
	fmt.Println("OBJECT:", header["OBJECT"], "EXPTIME:", header["EXPTIME"])
	fmt.Println("last:", img.FloatAt(img.Len()-1))
	// Output:
	// shape: [2 3] dtype: float
	// OBJECT: M31 EXPTIME: 600
	// last: 6
}

// A table extension holds a character column and a numeric column; the
// reader returns both, parallel to the file's column list.
func ExampleSaveFITSTable() {
	dir := tempDir()
	defer os.RemoveAll(dir)
	path := filepath.Join(dir, "catalogue.fits")

	flux, err := tensor.FromFloats([]float64{1.5, 2.25, 3.75}, 3)
	if err != nil {
		log.Fatal(err)
	}
	cols := []io.FITSTableColumn{
		{Name: "source", Form: "8A", Text: []string{"alpha", "beta", "gamma"}},
		{Name: "flux", Unit: "Jy", Form: "D", Data: flux},
	}
	if err := io.SaveFITSTable(path, false, cols, nil); err != nil {
		log.Fatal(err)
	}
	table, err := io.LoadFITSTable(path)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("kind:", table.Kind, "rows:", table.Rows)
	fmt.Println("names:", table.Names, "unit:", table.Units[1])
	fmt.Println("source[1]:", table.Text[0][1], "flux[1]:", table.Columns[1].FloatAt(1))
	// Output:
	// kind: BINTABLE rows: 3
	// names: [source flux] unit: Jy
	// source[1]: beta flux[1]: 2.25
}

// Two datasets, one in a group, with attributes on the datasets and on
// the root and the group: the file reads back with the same paths,
// shapes and dtypes, and each dataset carries the attributes of the
// enclosing groups.
func ExampleSaveHDF5() {
	dir := tempDir()
	defer os.RemoveAll(dir)
	path := filepath.Join(dir, "scan.h5")

	temp, err := tensor.FromFloat32s([]float32{1.5, 2.5, 3.5, 4.5}, 2, 2)
	if err != nil {
		log.Fatal(err)
	}
	counts, err := tensor.FromInts([]int64{100, 200, 300}, 3)
	if err != nil {
		log.Fatal(err)
	}
	sets := []io.HDF5Dataset{
		{Path: "/scan/temperature", Values: temp, Attrs: map[string]string{"units": "degC"}},
		{Path: "/scan/background", Values: counts, Attrs: map[string]string{"units": "counts"}},
	}
	groupAttrs := map[string]map[string]string{
		"/":     {"title": "cruise"},
		"/scan": {"instrument": "thermistor"},
	}
	if err := io.SaveHDF5(path, sets, groupAttrs); err != nil {
		log.Fatal(err)
	}
	got, err := io.LoadHDF5(path)
	if err != nil {
		log.Fatal(err)
	}
	for _, d := range got {
		fmt.Printf("%s %v %s %s %s\n", d.Path, d.Shape, d.Values.Dtype(), d.Attrs["instrument"], d.Attrs["units"])
	}
	// Output:
	// /scan/background [3] int thermistor counts
	// /scan/temperature [2 2] float32 thermistor degC
}

// Dimensions, a variable with an attribute and a global attribute go
// through NetCDF classic and back; the writer stores float64 as
// NC_DOUBLE, which the reader lands as float64 again, and the integer
// type codes land their own dtypes the same way.
func ExampleSaveNetCDF() {
	dir := tempDir()
	defer os.RemoveAll(dir)
	path := filepath.Join(dir, "field.nc")

	dims := []io.NetCDFDim{{Name: "lat", Length: 2}, {Name: "lon", Length: 3}}
	temp, err := tensor.FromFloats([]float64{1, 2, 3, 4, 5, 6}, 2, 3)
	if err != nil {
		log.Fatal(err)
	}
	vars := []io.NetCDFVar{{
		Name:   "temp",
		Dims:   []string{"lat", "lon"},
		Values: temp,
		Attrs:  map[string]string{"units": "degC"},
	}}
	if err := io.SaveNetCDF(path, dims, vars, map[string]string{"title": "cruise"}); err != nil {
		log.Fatal(err)
	}
	gotDims, gotVars, gotAttrs, err := io.LoadNetCDF(path)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("dims:", gotDims)
	fmt.Println("var:", gotVars[0].Name, gotVars[0].Dims, gotVars[0].Attrs["units"])
	fmt.Println("title:", gotAttrs["title"], "last:", gotVars[0].Values.FloatAt(5))
	// Output:
	// dims: [{lat 2} {lon 3}]
	// var: temp [lat lon] degC
	// title: cruise last: 6
}

// A native-endian file of float64 values maps into a read-only array
// without being read; release unmaps it and comes strictly last.
func ExampleSaveNativeFloats() {
	dir := tempDir()
	defer os.RemoveAll(dir)
	path := filepath.Join(dir, "values.bin")

	values := []float64{1.5, 2.5, 3.5, 4.5, 5.5}
	if err := io.SaveNativeFloats(path, values); err != nil {
		log.Fatal(err)
	}
	a, release, err := io.MapFloats(path, 0, len(values))
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("mapped:", a.Len(), a.Dtype(), a.FloatAt(0), a.FloatAt(4))
	if err := release(); err != nil {
		log.Fatal(err)
	}
	// Output:
	// mapped: 5 float 1.5 5.5
}