// Float32 tensors with independently selected native parameter encodings.
package main

import (
	"flag"
	"fmt"
	"github.com/openfluke/welvet/core"
	"github.com/openfluke/welvet/layers/mamba"
	"github.com/openfluke/welvet/nativemetal"
	"github.com/openfluke/welvet/quant"
	"log"
	"runtime"
	"strings"
)

func main() {
	if e := run(); e != nil {
		log.Fatal(e)
	}
}
func run() error {
	selected := flag.String("dtype", "all", "native parameter dtype name, or all 34")
	flag.Parse()
	if runtime.GOOS != "darwin" || runtime.GOARCH != "arm64" {
		fmt.Println("SKIP: Apple Metal needs Apple silicon macOS and cgo")
		return nil
	}
	count := 0
	for _, dt := range core.AllDTypes {
		if *selected != "all" && !strings.EqualFold(*selected, dt.String()) {
			continue
		}
		if e := forward(dt); e != nil {
			return fmt.Errorf("%s: %w", dt, e)
		}
		count++
	}
	if count == 0 {
		return fmt.Errorf("unknown dtype %q", *selected)
	}
	fmt.Printf("Mamba dtypes=%d; activations=float32; parameters=FormatNone native dtype\n", count)
	return nil
}
func forward(dt core.DType) error {
	const d, inner = 8, 16
	in, out := make([]float32, 2*inner*d), make([]float32, d*inner)
	for i := range in {
		in[i] = float32(i%17-8) * .01
	}
	for i := range out {
		out[i] = float32(i%13-6) * .01
	}
	model, e := mamba.NewConfigured(mamba.Config{DModel: d, DState: 4, Expand: 2, SeqLen: 3}, dt, quant.FormatNone, in, out, nil, nil)
	if e != nil {
		return e
	}
	s, e := nativemetal.Mount(model)
	if e != nil {
		return e
	}
	defer s.Close()
	x := core.NewTensor[float32](1, 3, d)
	for i := range x.Data {
		x.Data[i] = float32(i%7-3) * .05
	}
	post, e := s.Forward(x)
	if e != nil {
		return e
	}
	var checksum float64
	for _, v := range post.Data {
		checksum += float64(v)
	}
	fmt.Printf("%-10s shape=%v checksum=%.8g\n", dt, post.Shape, checksum)
	return nil
}