//go:build ignore

package main

import (
	"encoding/base64"
	"fmt"
	"image"
	"image/color"
	"math"
	
	"github.com/disintegration/imaging"
	"anticaptcha/pkg/onnx"
)

func main() {
	// 加载模型
	err := onnx.LoadModel("rotate", "models/[AntiCAP]-Rotation-RotNetR.onnx")
	if err != nil {
		fmt.Printf("Failed to load model: %v\n", err)
		return
	}
	fmt.Println("Model loaded")
	
	// 创建测试图片 (224x224)
	img := image.NewRGBA(image.Rect(0, 0, 224, 224))
	for y := 0; y < 224; y++ {
		for x := 0; x < 224; x++ {
			img.Set(x, y, color.RGBA{100, 150, 200, 255})
		}
	}
	for y := 50; y < 150; y++ {
		for x := 50; x < 100; x++ {
			img.Set(x, y, color.RGBA{255, 0, 0, 255})
		}
	}
	for y := 80; y < 140; y++ {
		for x := 120; x < 180; x++ {
			cx, cy := 150.0, 110.0
			rx, ry := 30.0, 30.0
			dx, dy := float64(x)-cx, float64(y)-cy
			if (dx*dx)/(rx*rx)+(dy*dy)/(ry*ry) <= 1 {
				img.Set(x, y, color.RGBA{0, 255, 0, 255})
			}
		}
	}
	
	// 预处理
	input := preprocessRotation(img)
	fmt.Printf("Input length: %d\n", len(input))
	
	// 推理
	sess, ok := onnx.GetSession("rotate")
	if !ok {
		fmt.Println("Session not found")
		return
	}
	
	dims := []int64{1, 3, 224, 224}
	output, err := sess.Run(input, dims)
	if err != nil {
		fmt.Printf("Run failed: %v\n", err)
		return
	}
	
	fmt.Printf("Output length: %d\n", len(output))
	fmt.Printf("Output first 10: %v\n", output[:10])
	
	// Argmax
	maxIdx := 0
	maxProb := float32(-math.MaxFloat32)
	for i := 0; i < len(output); i++ {
		if output[i] > maxProb {
			maxProb = output[i]
			maxIdx = i
		}
	}
	
	fmt.Printf("\nPredicted angle: %d degrees\n", maxIdx)
}

func preprocessRotation(img image.Image) []float32 {
	bounds := img.Bounds()
	w, h := bounds.Dx(), bounds.Dy()

	size := w
	if h < w {
		size = h
	}

	cropX := (w - size) / 2
	cropY := (h - size) / 2
	cropped := imaging.Crop(img, image.Rect(cropX, cropY, cropX+size, cropY+size))

	sqrt2 := math.Sqrt(2.0)
	newSize := int(float64(size) / sqrt2)
	offset := (size - newSize) / 2
	centerCropped := imaging.Crop(cropped, image.Rect(offset, offset, offset+newSize, offset+newSize))

	resized := imaging.Resize(centerCropped, 224, 224, imaging.Lanczos)
	rgb := imaging.Clone(resized)

	mean := [3]float32{0.485, 0.456, 0.406}
	std := [3]float32{0.229, 0.224, 0.225}

	pixels := make([]float32, 3*224*224)
	for y := 0; y < 224; y++ {
		for x := 0; x < 224; x++ {
			c := rgb.At(x, y)
			r, g, b, _ := c.RGBA()
			pixels[0*224*224+y*224+x] = (float32(r)/65535.0 - mean[0]) / std[0]
			pixels[1*224*224+y*224+x] = (float32(g)/65535.0 - mean[1]) / std[1]
			pixels[2*224*224+y*224+x] = (float32(b)/65535.0 - mean[2]) / std[2]
		}
	}

	return pixels
}