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

// Command fits builds a synthetic star field, saves it as a FITS
// primary image, reads it back and recovers the brightest star's
// position by a centre-of-mass centroid, the first step of any
// aperture photometry pipeline.
//
// Usage: go run ./examples/fits
package main

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

	"sourcedock.dev/petrbalvin/tensor"
)

func main() {
	const (
		size  = 128
		sigma = 2.0 // pixels, the seeing disk
	)
	// Three stars of different brightness on a flat sky background.
	type star struct {
		x, y, flux float64
	}
	stars := []star{
		{40.5, 60.5, 900},
		{80.5, 30.5, 300},
		{95.5, 95.5, 120},
	}
	field := make([]float64, size*size)
	for i := range size {
		for j := range size {
			v := 100.0 // sky
			for _, s := range stars {
				d2 := (float64(i)-s.y)*(float64(i)-s.y) + (float64(j)-s.x)*(float64(j)-s.x)
				v += s.flux * math.Exp(-d2/(2*sigma*sigma))
			}
			field[i*size+j] = v
		}
	}
	img, err := tensor.FromFloats(field, size, size)
	if err != nil {
		log.Fatal(err)
	}

	path := filepath.Join(os.TempDir(), "tensor-example-stars.fits")
	defer os.Remove(path)
	headers := map[string]string{
		"OBJECT":  "synthetic field",
		"EXPTIME": "30",
		"FILTER":  "V",
	}
	if err := tensor.SaveFITS(path, img, headers); err != nil {
		log.Fatal(err)
	}
	back, hdr, err := tensor.LoadFITS(path)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("wrote and read %s\n", filepath.Base(path))
	for _, k := range []string{"OBJECT", "EXPTIME", "FILTER"} {
		fmt.Printf("  %s = %s\n", k, hdr[k])
	}
	if back.Shape()[0] != size || back.Shape()[1] != size {
		log.Fatalf("round trip changed the shape: %v", back.Shape())
	}

	// Locate the brightest pixel, then centroid a 9x9 window around
	// it with the sky level subtracted.
	best, bestVal := 0, -1.0
	for i := range size * size {
		if v := back.FloatAt(i); v > bestVal {
			best, bestVal = i, v
		}
	}
	by, bx := best/size, best%size
	sum, sx, sy := 0.0, 0.0, 0.0
	for i := by - 4; i <= by+4; i++ {
		for j := bx - 4; j <= bx+4; j++ {
			w := back.FloatAt(i*size+j) - 100
			if w < 0 {
				w = 0
			}
			sum += w
			sx += w * float64(j)
			sy += w * float64(i)
		}
	}
	fmt.Printf("\nbrightest star: peak at (x=%d, y=%d), %.0f counts\n", bx, by, bestVal)
	fmt.Printf("centroid of the 9x9 window: (x=%.2f, y=%.2f)\n", sx/sum, sy/sum)
	fmt.Println("true position:              (x=40.50, y=60.50)")
}