package mnist

import (
	"archive/zip"
	"encoding/binary"
	"fmt"
	"io"
	"math"
	"os"
	"testing"
	"time"

	"tests/Examples/lattigo/ckks/mnist/mnist_utils"
)

const dataBase = "../../../common/mnist/data/"
const (
	modelPath  = dataBase + "traced_model.pt"
	imagesPath = dataBase + "t10k-images-idx3-ubyte"
	labelsPath = dataBase + "t10k-labels-idx1-ubyte"
)

func loadWeights(path string) ([][]float32, error) {
	r, err := zip.OpenReader(path)
	if err != nil {
		return nil, err
	}
	defer r.Close()

	weights := make([][]float32, 4)
	for i := 0; i < 4; i++ {
		f, err := r.Open(fmt.Sprintf("traced_model/data/%d", i))
		if err != nil {
			return nil, err
		}
		defer f.Close()

		data, err := io.ReadAll(f)
		if err != nil {
			return nil, err
		}

		numFloats := len(data) / 4
		weights[i] = make([]float32, numFloats)
		for j := 0; j < numFloats; j++ {
			bits := binary.LittleEndian.Uint32(data[j*4 : (j+1)*4])
			weights[i][j] = math.Float32frombits(bits)
		}
	}
	return weights, nil
}

func loadMNISTImages(path string) ([][]float64, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	header := make([]byte, 16)
	if _, err := f.Read(header); err != nil {
		return nil, err
	}

	numImages := int(binary.BigEndian.Uint32(header[4:8]))
	rows := int(binary.BigEndian.Uint32(header[8:12]))
	cols := int(binary.BigEndian.Uint32(header[12:16]))

	pixelsPerImage := rows * cols
	images := make([][]float64, numImages)
	for i := 0; i < numImages; i++ {
		imgData := make([]byte, pixelsPerImage)
		if _, err := f.Read(imgData); err != nil {
			return nil, err
		}
		images[i] = make([]float64, pixelsPerImage)
		for j := 0; j < pixelsPerImage; j++ {
			// Normalize: (X/255.0 - 0.1307) / 0.3081
			val := float64(imgData[j]) / 255.0
			images[i][j] = (val - 0.1307) / 0.3081
		}
	}
	return images, nil
}

func loadMNISTLabels(path string) ([]int, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	header := make([]byte, 8)
	if _, err := f.Read(header); err != nil {
		return nil, err
	}

	numLabels := int(binary.BigEndian.Uint32(header[4:8]))
	labels := make([]int, numLabels)
	labelData := make([]byte, numLabels)
	if _, err := f.Read(labelData); err != nil {
		return nil, err
	}
	for i := 0; i < numLabels; i++ {
		labels[i] = int(labelData[i])
	}
	return labels, nil
}

func TestMNIST(t *testing.T) {
	if wd, err := os.Getwd(); err == nil {
		t.Logf("PWD: %s", wd)
	}
	t.Logf("modelPath: %s", modelPath)
	weights, err := loadWeights(modelPath)
	if err != nil {
		t.Fatalf("Failed to load weights: %v", err)
	}

	images, err := loadMNISTImages(imagesPath)
	if err != nil {
		t.Fatalf("Failed to load images: %v", err)
	}

	labels, err := loadMNISTLabels(labelsPath)
	if err != nil {
		t.Fatalf("Failed to load labels: %v", err)
	}

	// Use the autogenerated configuration which has the correct parameters
	// determined by the HEIR compiler
	startConfig := time.Now()
	evaluator, params, encoder, encryptor, decryptor := Mnist__configure()
	t.Logf("Mnist__configure took %v", time.Since(startConfig))

	// Preprocess weights
	startPreprocess := time.Now()
	preprocessedWeights := mnist_utils.Mnist__preprocessing(params, encoder, weights[3], weights[2], weights[1], weights[0])
	t.Logf("mnist_utils.Mnist__preprocessing took %v", time.Since(startPreprocess))

	// Test with 1 sample since the computation is slow with the
	// large security parameters (LogN=15)
	total := 1
	correct := 0

	for i := 0; i < total; i++ {
		input := images[i]
		label := labels[i]

		// Convert float64 input to float32 for encryption helper
		inputFloat32 := make([]float32, len(input))
		for j := 0; j < len(input); j++ {
			inputFloat32[j] = float32(input[j])
		}

		// Use generated encryption helper
		startEncrypt := time.Now()
		ctInput := Mnist__encrypt__arg4(evaluator, params, encoder, encryptor, inputFloat32)
		t.Logf("Mnist__encrypt__arg4 took %v", time.Since(startEncrypt))

		startZero1 := time.Now()
		ctZero1 := Mnist__encrypt__zero__0(evaluator, params, encoder, encryptor)
		t.Logf("mnist__encrypt__zero__0 took %v", time.Since(startZero1))

		startZero2 := time.Now()
		ctZero2 := Mnist__encrypt__zero__1(evaluator, params, encoder, encryptor)
		t.Logf("mnist__encrypt__zero__1 took %v", time.Since(startZero2))

		startTime := time.Now()
		// Calling the generated preprocessed mnist function
		resCt := Mnist__preprocessed(evaluator, params, encoder, weights[0], weights[1], weights[2], weights[3], ctInput, ctZero1, ctZero2, preprocessedWeights)
		duration := time.Since(startTime)
		t.Logf("Sample %d (Mnist__preprocessed) took %v", i, duration)

		startDecrypt := time.Now()
		// Use generated decryption helper
		resValues := Mnist__decrypt__result0(evaluator, params, encoder, decryptor, resCt)
		t.Logf("Mnist__decrypt__result0 took %v", time.Since(startDecrypt))

		maxVal := float32(-math.MaxFloat32)
		maxIdx := -1
		for j := 0; j < 10; j++ {
			if resValues[j] > maxVal {
				maxVal = resValues[j]
				maxIdx = j
			}
		}

		if maxIdx == label {
			correct++
		}
		t.Logf("Sample %d: predicted %d, actual %d", i, maxIdx, label)
		t.Logf("logits:\n")
		for j := 0; j < 10; j++ {
			t.Logf("%d, %.6f\n", j, resValues[j])
		}
	}

	accuracy := float64(correct) / float64(total)
	t.Logf("Accuracy: %.2f (%d/%d correct)", accuracy, correct, total)
	// With only 3 samples, we need at least 2/3 correct (0.67)
	if accuracy < 0.6 {
		t.Errorf("Accuracy too low: %.2f (%d/%d correct)", accuracy, correct, total)
	}
}
