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

// Command netcdf writes a synthetic climate field to a NetCDF classic
// file, reads it back, and computes the zonal statistics a climate
// workflow starts from. The point is the round trip: dimensions,
// attributes and values survive the file exactly.
//
// Usage: go run ./examples/netcdf
package main

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

	"sourcedock.dev/petrbalvin/tensor"
)

func main() {
	const (
		nLat = 36 // 5-degree grid
		nLon = 72
	)
	// A warm anomaly centred on 45 N, 15 E over a zonal gradient, the
	// shape of a heat island in a coarse climate model.
	field := make([]float64, nLat*nLon)
	for i := range nLat {
		lat := -90 + 5*(float64(i)+0.5)
		for j := range nLon {
			lon := -180 + 5*(float64(j)+0.5)
			base := 30*math.Cos(lat*math.Pi/180) - 5
			dLat := (lat - 45) / 15
			dLon := math.Sin((lon - 15) * math.Pi / 180)
			anomaly := 8 * math.Exp(-(dLat*dLat + dLon*dLon))
			field[i*nLon+j] = base + anomaly
		}
	}
	temp, err := tensor.FromFloats(field, nLat, nLon)
	if err != nil {
		log.Fatal(err)
	}

	path := filepath.Join(os.TempDir(), "tensor-example-climate.nc")
	defer os.Remove(path)
	dims := []tensor.NetCDFDim{
		{Name: "lat", Length: nLat},
		{Name: "lon", Length: nLon},
	}
	vars := []tensor.NetCDFVar{{
		Name:   "temperature",
		Dims:   []string{"lat", "lon"},
		Values: temp,
		Attrs: map[string]string{
			"units":       "degC",
			"long_name":   "synthetic air temperature",
			"anomaly_lon": "15",
		},
	}}
	attrs := map[string]string{
		"title":  "tensor NetCDF example",
		"source": "synthetic Gaussian anomaly",
	}
	if err := tensor.SaveNetCDF(path, dims, vars, attrs); err != nil {
		log.Fatal(err)
	}
	fmt.Printf("wrote %s: %d x %d grid, %d variables\n\n", path, nLat, nLon, len(vars))

	gotDims, gotVars, gotAttrs, err := tensor.LoadNetCDF(path)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("dimensions: ")
	for _, d := range gotDims {
		fmt.Printf("%s(%d) ", d.Name, d.Length)
	}
	fmt.Printf("\nglobal attributes: %v\n\n", gotAttrs)

	back := gotVars[0].Values
	if back.Shape()[0] != nLat || back.Shape()[1] != nLon {
		log.Fatalf("round trip changed the shape: %v", back.Shape())
	}
	maxDiff := 0.0
	for i := range nLat * nLon {
		d := math.Abs(back.FloatAt(i) - field[i])
		if d > maxDiff {
			maxDiff = d
		}
	}
	fmt.Printf("largest round-trip difference: %g (exact for float64)\n\n", maxDiff)

	// Zonal means: the latitude profile of the field, the first thing
	// a climate diagnostic asks for.
	fmt.Println("latitude    zonal mean temperature")
	for i := 0; i < nLat; i += 6 {
		row, err := tensor.Slice(back, 0, i, i+1)
		if err != nil {
			log.Fatal(err)
		}
		mean, err := tensor.Mean(row)
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("%6.1f°    %10.3f degC\n", -90+5*(float64(i)+0.5), mean)
	}
}