// cmd/eval — CLI driver for the jepa-fx-risk evaluation harness.
// Reads embeddings from a parquet/npy-style JSON export (embeddings.json)
// and targets from eurusd_daily.parquet, then runs the requested metric.
//
//	./bin/eval -metric probe|silhouette|erank [-emb embeddings.json]
//
// embeddings.json format: {"embeddings": [[...], ...], "dates": ["2022-01-03", ...]}
// Generated by train.py when run with EXPORT_EMBEDDINGS=1.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log/slog"
	"math"
	"os"

	"gitea.d-ma.be/mathias/jepa-fx-risk/internal/eval"
)

func main() {
	metric := flag.String("metric", "probe", "probe | silhouette | erank")
	embFile := flag.String("emb", "embeddings.json", "path to embeddings JSON")
	flag.Parse()

	log := slog.New(slog.NewJSONHandler(os.Stdout, nil))

	emb, labels, y, err := loadEmbeddings(*embFile)
	if err != nil {
		log.Error("load embeddings", "err", err)
		os.Exit(1)
	}
	log.Info("loaded", "n", len(emb), "dim", len(emb[0]), "metric", *metric)

	switch *metric {
	case "probe":
		r2 := eval.LinearProbe(emb, y, 1e-3)
		fmt.Printf(`{"metric":"val_vol_r2","value":%.6f}`+"\n", r2)
		log.Info("linear probe", "val_vol_r2", fmt.Sprintf("%.4f", r2))
	case "silhouette":
		if labels == nil {
			log.Error("silhouette requires HV labels in embeddings.json")
			os.Exit(1)
		}
		sil, err := eval.Silhouette(emb, labels)
		if err != nil {
			log.Error("silhouette", "err", err)
			os.Exit(1)
		}
		fmt.Printf(`{"metric":"silhouette","value":%.6f}`+"\n", sil)
		log.Info("silhouette", "score", fmt.Sprintf("%.4f", sil))
	case "erank":
		er := eval.EffectiveRank(emb)
		fmt.Printf(`{"metric":"effective_rank","value":%.6f}`+"\n", er)
		log.Info("effective rank", "erank", fmt.Sprintf("%.2f", er))
	default:
		log.Error("unknown metric", "metric", *metric)
		os.Exit(1)
	}
}

type embJSON struct {
	Embeddings [][]float64 `json:"embeddings"`
	Dates      []string    `json:"dates"`
	RealizedVol []float64  `json:"realized_vol"`
	HVLabel     []int      `json:"hv_label"`
}

func loadEmbeddings(path string) (emb [][]float64, labels []int, y []float64, err error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("open %s: %w", path, err)
	}
	defer func() { _ = f.Close() }()

	var d embJSON
	if err := json.NewDecoder(f).Decode(&d); err != nil {
		return nil, nil, nil, fmt.Errorf("decode: %w", err)
	}
	if len(d.Embeddings) == 0 {
		return nil, nil, nil, fmt.Errorf("empty embeddings in %s", path)
	}

	// standardise embeddings (zero mean, unit std) per dimension
	n, dim := len(d.Embeddings), len(d.Embeddings[0])
	mu := make([]float64, dim)
	for _, row := range d.Embeddings {
		for j, v := range row {
			mu[j] += v
		}
	}
	for j := range mu {
		mu[j] /= float64(n)
	}
	sd := make([]float64, dim)
	for _, row := range d.Embeddings {
		for j, v := range row {
			diff := v - mu[j]
			sd[j] += diff * diff
		}
	}
	for j := range sd {
		sd[j] = math.Sqrt(sd[j]/float64(n)) + 1e-8
	}
	norm := make([][]float64, n)
	for i, row := range d.Embeddings {
		norm[i] = make([]float64, dim)
		for j, v := range row {
			norm[i][j] = (v - mu[j]) / sd[j]
		}
	}

	if len(d.HVLabel) > 0 {
		labels = d.HVLabel
	}
	return norm, labels, d.RealizedVol, nil
}