package gotensor

import (
	"math"
	"testing"

	"git.kingecg.top/kingecg/gomatrix"
)

// Linear 是一个简单的线性层实现，用于测试
type Linear struct {
	Weight *Tensor
	Bias   *Tensor
}

// NewLinear 创建一个新的线性层
func NewLinear(weight, bias *Tensor) *Linear {
	return &Linear{
		Weight: weight,
		Bias:   bias,
	}
}

func (l *Linear) Forward(inputs *Tensor) (*Tensor, error) {
	// 执行线性变换: output = inputs * weight^T + bias
	weightTransposed, err := l.Weight.Data.Transpose()
	if err != nil {
		return nil, err
	}

	mulResult, err := inputs.MatMul(&Tensor{Data: weightTransposed}) // 需要包装为Tensor
	if err != nil {
		return nil, err
	}

	output, err := mulResult.Add(l.Bias)
	if err != nil {
		return nil, err
	}

	return output, nil
}

func (l *Linear) Parameters() []*Tensor {
	return []*Tensor{l.Weight, l.Bias}
}

func (l *Linear) ZeroGrad() {
	l.Weight.ZeroGrad()
	l.Bias.ZeroGrad()
}

// TestSequential 测试Sequential模型的基本功能
func TestSequential(t *testing.T) {
	// 创建一个简单的线性层实现用于测试
	weight, err := gomatrix.NewMatrix([]float64{1, 2, 3, 4}, []int{2, 2})
	if err != nil {
		t.Fatalf("Failed to create weight matrix: %v", err)
	}
	bias, err := gomatrix.NewMatrix([]float64{0.5, 0.5}, []int{1, 2}) // 改为 1x2 矩阵以匹配输出形状
	if err != nil {
		t.Fatalf("Failed to create bias vector: %v", err)
	}

	linearLayer := &Linear{
		Weight: &Tensor{
			Data: weight,
			Grad: Must(gomatrix.NewMatrix([]float64{0, 0, 0, 0}, []int{2, 2})),
		},
		Bias: &Tensor{
			Data: bias,
			Grad: Must(gomatrix.NewMatrix([]float64{0, 0}, []int{1, 2})),
		},
	}

	// 创建Sequential模型
	seq := &Sequential{
		Layers: []Layer{linearLayer},
	}

	// 测试前向传播
	inputData, err := gomatrix.NewMatrix([]float64{1, 1}, []int{1, 2})
	if err != nil {
		t.Fatalf("Failed to create input vector: %v", err)
	}
	input := &Tensor{
		Data: inputData,
	}
	output, err := seq.Forward(input)
	if err != nil {
		t.Errorf("Sequential forward failed: %v", err)
	}

	if output == nil {
		t.Error("Expected output tensor, got nil")
	}

	// 测试Parameters方法
	params := seq.Parameters()
	if len(params) == 0 {
		t.Error("Expected non-empty parameters list")
	}

	// 测试ZeroGrad方法
	seq.ZeroGrad()
	for _, param := range params {
		// 检查梯度是否被清零
		shape := param.Shape()
		for i := 0; i < param.Size(); i++ {
			var gradVal float64
			if len(shape) == 1 {
				gradVal, _ = param.Grad.Get(i)
			} else if len(shape) == 2 {
				cols := shape[1]
				gradVal, _ = param.Grad.Get(i/cols, i%cols)
			}
			if gradVal != 0 {
				t.Errorf("Expected gradient to be zero, got %v", gradVal)
			}
		}
	}
}

// TestSaveLoadModel 测试模型保存和加载功能
func TestSaveLoadModel(t *testing.T) {
	weight, err := gomatrix.NewMatrix([]float64{1, 2, 3, 4}, []int{2, 2})
	if err != nil {
		t.Fatalf("Failed to create weight matrix: %v", err)
	}
	bias, err := gomatrix.NewMatrix([]float64{0.5, 0.5}, []int{1, 2}) // 改为 1x2 矩阵以匹配输出形状
	if err != nil {
		t.Fatalf("Failed to create bias matrix: %v", err)
	}

	linearLayer := &Linear{
		Weight: &Tensor{
			Data: weight,
			Grad: Must(gomatrix.NewMatrix([]float64{0, 0, 0, 0}, []int{2, 2})),
		},
		Bias: &Tensor{
			Data: bias,
			Grad: Must(gomatrix.NewMatrix([]float64{0, 0}, []int{1, 2})),
		},
	}

	model := &Sequential{
		Layers: []Layer{linearLayer},
	}

	// 保存模型
	filepath := "/tmp/test_model.json"
	err = SaveModel(model, filepath)
	if err != nil {
		t.Errorf("SaveModel failed: %v", err)
	}

	// 修改原始模型参数
	weightVal, _ := linearLayer.Weight.Data.Get(0, 0)
	if math.Abs(weightVal-1.0) > 1e-9 {
		t.Errorf("Expected weight to be 1.0, got %v", weightVal)
	}

	// 加载模型
	err = LoadModel(model, filepath)
	if err != nil {
		t.Errorf("LoadModel failed: %v", err)
	}

	// 验证加载后的参数
	loadedWeight, _ := linearLayer.Weight.Data.Get(0, 0)
	if math.Abs(loadedWeight-1.0) > 1e-9 {
		t.Errorf("Expected loaded weight to be 1.0, got %v", loadedWeight)
	}
}

// TestLinearLayer 测试线性层功能
func TestLinearLayer(t *testing.T) {
	weightData, err := gomatrix.NewMatrix([]float64{2, 0, 0, 3}, []int{2, 2})
	if err != nil {
		t.Fatalf("Failed to create weight matrix: %v", err)
	}
	weight := &Tensor{
		Data: weightData,
		Grad: Must(gomatrix.NewMatrix([]float64{0, 0, 0, 0}, []int{2, 2})),
	}

	biasData, err := gomatrix.NewMatrix([]float64{0.5, 0.5}, []int{1, 2}) // 改为 1x2 矩阵
	if err != nil {
		t.Fatalf("Failed to create bias vector: %v", err)
	}
	bias := &Tensor{
		Data: biasData,
		Grad: Must(gomatrix.NewVector([]float64{0, 0})),
	}

	layer := NewLinear(weight, bias)

	// 测试前向传播
	inputData, err := gomatrix.NewMatrix([]float64{1, 1}, []int{1, 2})
	if err != nil {
		t.Fatalf("Failed to create input vector: %v", err)
	}
	input := &Tensor{
		Data: inputData,
	}
	output, err := layer.Forward(input)
	if err != nil {
		t.Fatalf("Linear layer forward failed: %v", err)
	}

	// 计算期望输出: input * weight^T + bias = [1,1] * [[2,0],[0,3]]^T + [0.5,0.5] = [1,1] * [[2,0],[0,3]] + [0.5,0.5] = [2,3] + [0.5,0.5] = [2.5,3.5]
	expected0, _ := output.Data.Get(0, 0)
	expected1, _ := output.Data.Get(0, 1)

	if math.Abs(expected0-2.5) > 1e-9 {
		t.Errorf("Expected output[0] to be 2.5, got %v", expected0)
	}
	if math.Abs(expected1-3.5) > 1e-9 {
		t.Errorf("Expected output[1] to be 3.5, got %v", expected1)
	}

	// 测试Parameters方法
	params := layer.Parameters()
	if len(params) != 2 { // 权重和偏置
		t.Errorf("Expected 2 parameters, got %d", len(params))
	}

	// 测试ZeroGrad方法
	layer.ZeroGrad()
	for _, param := range params {
		shape := param.Shape()
		for i := 0; i < param.Size(); i++ {
			var gradVal float64
			if len(shape) == 1 {
				gradVal, _ = param.Grad.Get(i)
			} else if len(shape) == 2 {
				cols := shape[1]
				gradVal, _ = param.Grad.Get(i/cols, i%cols)
			}
			if math.Abs(gradVal) > 1e-9 {
				t.Errorf("Expected gradient to be zero, got %v", gradVal)
			}
		}
	}
}