control/control_case/hdf5.go

package control_case

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
)

type FieldMapResponse struct {
	RequestedT float64                `json:"requested_t"`
	StageT     float64                `json:"stage_t"`
	Rows       int                    `json:"rows"`
	Cols       int                    `json:"cols"`
	Fields     map[string][][]float64 `json:"fields"`
}

var (
	stagePathRe   = regexp.MustCompile(`^\s*group\s+/map/(stage_t=[^\s]+)\s*$`)
	dimsRe        = regexp.MustCompile(`\(\s*(\d+)\s*,\s*(\d+)\s*\)`)
	floatInDataRe = regexp.MustCompile(`[-+]?(?:\d*\.\d+|\d+\.?\d*)(?:[eE][-+]?\d+)?`)
)

func ReadNearestFieldMap(filePath string, requestedT float64) (*FieldMapResponse, error) {
	stageGroups, err := listStageGroups(filePath)
	if err != nil {
		return nil, err
	}
	if len(stageGroups) == 0 {
		return nil, fmt.Errorf("no stage_t groups found")
	}

	selected := stageGroups[0]
	for _, stage := range stageGroups[1:] {
		if absFloat64(stage.stageT-requestedT) < absFloat64(selected.stageT-requestedT) {
			selected = stage
		}
	}

	result := &FieldMapResponse{
		RequestedT: requestedT,
		StageT:     selected.stageT,
		Fields:     make(map[string][][]float64, 4),
	}

	for _, fieldName := range []string{"psi", "phi", "T", "C"} {
		dims, values, err := readDatasetMatrix(filePath, selected.path, fieldName)
		if err != nil {
			return nil, err
		}
		if len(dims) != 2 {
			return nil, fmt.Errorf("unexpected dims for %s: %v", fieldName, dims)
		}

		matrix, err := reshapeFloat64(values, dims[0], dims[1])
		if err != nil {
			return nil, fmt.Errorf("failed to reshape dataset %s: %w", fieldName, err)
		}
		matrix = transposeFloat64(matrix)
		if result.Rows == 0 && result.Cols == 0 {
			result.Rows = len(matrix)
			if result.Rows > 0 {
				result.Cols = len(matrix[0])
			}
		}
		result.Fields[fieldName] = matrix
	}

	return result, nil
}

func CopyLastStateStorageH5(sourcePath, outputPath string) error {
	if strings.TrimSpace(outputPath) == "" {
		return fmt.Errorf("output path is empty")
	}

	sourceAbs, err := filepath.Abs(sourcePath)
	if err != nil {
		return err
	}
	outputAbs, err := filepath.Abs(outputPath)
	if err != nil {
		return err
	}
	if sourceAbs == outputAbs {
		return fmt.Errorf("output path must be different from source path")
	}

	stageGroups, err := listStageGroups(sourcePath)
	if err != nil {
		return err
	}
	if len(stageGroups) == 0 {
		return fmt.Errorf("no stage_t groups found")
	}

	latest := stageGroups[0]
	for _, stage := range stageGroups[1:] {
		if stage.stageT > latest.stageT {
			latest = stage
		}
	}

	if err := os.MkdirAll(filepath.Dir(outputPath), os.ModePerm); err != nil {
		return err
	}

	tmpFile, err := os.CreateTemp(filepath.Dir(outputPath), ".storage-*.h5")
	if err != nil {
		return err
	}
	tmpPath := tmpFile.Name()
	if err := tmpFile.Close(); err != nil {
		_ = os.Remove(tmpPath)
		return err
	}
	if err := os.Remove(tmpPath); err != nil {
		return err
	}

	if err := runH5Copy(sourcePath, tmpPath, latest.path, "/map/stage_t=0", "-p"); err != nil {
		_ = os.Remove(tmpPath)
		return err
	}

	if err := os.Rename(tmpPath, outputPath); err != nil {
		_ = os.Remove(tmpPath)
		return err
	}

	return nil
}

type stageFieldMap struct {
	path   string
	stageT float64
}

func listStageGroups(filePath string) ([]stageFieldMap, error) {
	out, err := runH5Dump(filePath, "-n")
	if err != nil {
		return nil, err
	}

	var stages []stageFieldMap
	for _, line := range strings.Split(string(out), "\n") {
		matches := stagePathRe.FindStringSubmatch(line)
		if len(matches) != 2 {
			continue
		}
		stageT, err := strconv.ParseFloat(strings.TrimPrefix(matches[1], "stage_t="), 64)
		if err != nil {
			continue
		}
		stages = append(stages, stageFieldMap{
			path:   "/map/" + matches[1],
			stageT: stageT,
		})
	}

	return stages, nil
}

func readDatasetMatrix(filePath, stagePath, fieldName string) ([]int, []float64, error) {
	datasetPath := stagePath + "/" + fieldName

	headerOut, err := runH5Dump(filePath, "-H", "-d", datasetPath)
	if err != nil {
		return nil, nil, err
	}

	rows, cols, err := parseDimsFromHeader(headerOut)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to parse dims for %s: %w", datasetPath, err)
	}

	dataOut, err := runH5Dump(filePath, "-d", datasetPath)
	if err != nil {
		return nil, nil, err
	}

	values, err := parseFloatData(dataOut)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to parse data for %s: %w", datasetPath, err)
	}

	if rows*cols != len(values) {
		return nil, nil, fmt.Errorf("value count %d does not match %dx%d for %s", len(values), rows, cols, datasetPath)
	}

	return []int{rows, cols}, values, nil
}

func parseDimsFromHeader(out []byte) (int, int, error) {
	matches := dimsRe.FindSubmatch(out)
	if len(matches) != 3 {
		return 0, 0, fmt.Errorf("dimensions not found")
	}

	rows, err := strconv.Atoi(string(matches[1]))
	if err != nil {
		return 0, 0, err
	}
	cols, err := strconv.Atoi(string(matches[2]))
	if err != nil {
		return 0, 0, err
	}
	return rows, cols, nil
}

func parseFloatData(out []byte) ([]float64, error) {
	lines := strings.Split(string(out), "\n")
	values := make([]float64, 0, 2048)
	inData := false

	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		switch {
		case trimmed == "DATA {":
			inData = true
			continue
		case inData && trimmed == "}":
			return values, nil
		case !inData:
			continue
		}

		colon := strings.Index(trimmed, ":")
		if colon < 0 {
			continue
		}
		payload := strings.TrimSpace(trimmed[colon+1:])
		if payload == "" {
			continue
		}

		matches := floatInDataRe.FindAllString(payload, -1)
		for _, match := range matches {
			value, err := strconv.ParseFloat(match, 64)
			if err != nil {
				return nil, err
			}
			values = append(values, value)
		}
	}

	if !inData {
		return nil, fmt.Errorf("data block not found")
	}
	return nil, fmt.Errorf("data block not terminated")
}

func reshapeFloat64(values []float64, rows, cols int) ([][]float64, error) {
	if rows*cols != len(values) {
		return nil, fmt.Errorf("value count %d does not match %dx%d", len(values), rows, cols)
	}

	matrix := make([][]float64, rows)
	for r := 0; r < rows; r++ {
		start := r * cols
		row := make([]float64, cols)
		copy(row, values[start:start+cols])
		matrix[r] = row
	}
	return matrix, nil
}

func transposeFloat64(matrix [][]float64) [][]float64 {
	if len(matrix) == 0 || len(matrix[0]) == 0 {
		return matrix
	}

	rows := len(matrix)
	cols := len(matrix[0])
	transposed := make([][]float64, cols)
	for c := 0; c < cols; c++ {
		transposed[c] = make([]float64, rows)
		for r := 0; r < rows; r++ {
			transposed[c][r] = matrix[r][c]
		}
	}
	return transposed
}

func runH5Dump(filePath string, args ...string) ([]byte, error) {
	cmdArgs := append(args, filePath)
	cmd := exec.Command("h5dump", cmdArgs...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return nil, fmt.Errorf("h5dump %v failed: %w: %s", cmdArgs, err, strings.TrimSpace(string(out)))
	}
	return out, nil
}

func runH5Copy(sourceFile, outputFile, sourcePath, outputPath string, args ...string) error {
	cmdArgs := []string{"-i", sourceFile, "-o", outputFile, "-s", sourcePath, "-d", outputPath}
	cmdArgs = append(cmdArgs, args...)
	cmd := exec.Command("h5copy", cmdArgs...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("h5copy %v failed: %w: %s", cmdArgs, err, strings.TrimSpace(string(out)))
	}
	return nil
}

func absFloat64(value float64) float64 {
	if value < 0 {
		return -value
	}
	return value
}