package gotensor

import (
	"testing"

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

// MockLayer 是一个用于测试的模拟层
type MockLayer struct {
	Weight *Tensor
}

// NewMockLayer 创建一个新的模拟层
func NewMockLayer() *MockLayer {
	weight, _ := gomatrix.NewMatrix([]float64{0.5, 0.3, 0.4, 0.7}, []int{2, 2})
	grad, _ := gomatrix.NewZeros([]int{2, 2})

	return &MockLayer{
		Weight: &Tensor{
			Data: weight,
			Grad: grad,
		},
	}
}

func (m *MockLayer) Forward(inputs *Tensor) (*Tensor, error) {
	// 简单的矩阵乘法
	output, err := m.Weight.MatMul(inputs)
	if err != nil {
		return nil, err
	}
	return output, nil
}

func (m *MockLayer) Parameters() []*Tensor {
	return []*Tensor{m.Weight}
}

func (m *MockLayer) ZeroGrad() {
	m.Weight.ZeroGrad()
}

// MockModel 是一个用于测试的模拟模型
type MockModel struct {
	Layer *MockLayer
}

func (m *MockModel) Forward(inputs *Tensor) (*Tensor, error) {
	return m.Layer.Forward(inputs)
}

func (m *MockModel) Parameters() []*Tensor {
	return m.Layer.Parameters()
}

func (m *MockModel) ZeroGrad() {
	m.Layer.ZeroGrad()
}

func NewVector(data []float64) (*Tensor, error) {
	return NewTensor(data, []int{len(data), 1})

}

func Must[T any](t *T, err error) *T {
	if err != nil {
		panic(err)
	}
	return t
}

// TestTrainer 测试训练器的基本功能
func TestTrainer(t *testing.T) {
	// 创建模型
	mockLayer := NewMockLayer()
	model := &MockModel{
		Layer: mockLayer,
	}

	// 创建优化器
	optimizer := NewSGD(model.Parameters(), 0.01)

	// 创建训练器
	trainer := NewTrainer(model, optimizer)

	// 创建训练数据
	inputs := []*Tensor{
		Must(NewVector([]float64{1, 0})),
		Must(NewVector([]float64{0, 1})),
	}

	targets := []*Tensor{
		Must(NewVector([]float64{1, 0})),
		Must(NewVector([]float64{0, 1})),
	}

	// 定义损失函数
	lossFn := func(output, target *Tensor) *Tensor {
		// MSE 损失函数
		diff, _ := output.Data.Subtract(target.Data)
		squared, _ := diff.Multiply(diff)

		// 计算矩阵元素的总和
		var total float64
		shape := squared.Shape()
		if len(shape) == 1 {
			for i := 0; i < squared.Size(); i++ {
				val, _ := squared.Get(i)
				total += val
			}
		} else if len(shape) == 2 {
			rows, cols := shape[0], shape[1]
			for i := 0; i < rows; i++ {
				for j := 0; j < cols; j++ {
					val, _ := squared.Get(i, j)
					total += val
				}
			}
		} else {
			// 对于其他维度，遍历所有元素
			for i := 0; i < squared.Size(); i++ {
				var val float64
				var err error

				if len(shape) == 1 {
					val, err = squared.Get(i)
				} else if len(shape) == 2 {
					cols := shape[1]
					val, err = squared.Get(i/cols, i%cols)
				} else {
					// 对于高维张量，按顺序访问元素
					val, err = squared.Get(i)
				}

				if err != nil {
					continue
				}

				total += val
			}
		}

		size := float64(output.Size())
		resultVal := total / size

		// 创建结果张量
		result, _ := NewTensor([]float64{resultVal}, []int{1})
		return result
	}

	// 测试TrainEpoch
	avgLoss, err := trainer.TrainEpoch(inputs, targets, lossFn)
	if err != nil {
		t.Errorf("TrainEpoch failed: %v", err)
	}

	if avgLoss < 0 {
		t.Errorf("Expected non-negative loss, got %v", avgLoss)
	}

	// 测试Evaluate
	evalLoss, err := trainer.Evaluate(inputs, targets, lossFn)
	if err != nil {
		t.Errorf("Evaluate failed: %v", err)
	}

	if evalLoss < 0 {
		t.Errorf("Expected non-negative evaluation loss, got %v", evalLoss)
	}
}

// TestTrainerFullTrain 测试完整的训练过程
func TestTrainerFullTrain(t *testing.T) {
	// 创建一个简单的线性回归模型
	mockLayer := NewMockLayer()
	model := &MockModel{
		Layer: mockLayer,
	}

	// 创建优化器
	optimizer := NewSGD(model.Parameters(), 0.1)

	// 创建训练器
	trainer := NewTrainer(model, optimizer)

	// 创建训练数据 (简单的 y = x 示例)
	trainInputs := []*Tensor{
		Must(NewVector([]float64{1, 0})),
		Must(NewVector([]float64{0, 1})),
		Must(NewVector([]float64{1, 1})),
	}

	trainTargets := []*Tensor{
		Must(NewVector([]float64{1, 0})),
		Must(NewVector([]float64{0, 1})),
		Must(NewVector([]float64{1, 1})),
	}

	// 定义损失函数
	lossFn := func(output, target *Tensor) *Tensor {
		// MSE 损失函数
		diff, _ := output.Data.Subtract(target.Data)
		squared, _ := diff.Multiply(diff)
		sum := squared.Sum()
		size := float64(output.Size())
		result := sum / size
		return Must(NewTensor([]float64{result}, []int{1}))
		// return result
	}

	// 训练模型
	epochs := 5
	err := trainer.Train(trainInputs, trainTargets, epochs, lossFn, false)
	if err != nil {
		t.Errorf("Full training failed: %v", err)
	}

	// 验证训练后的损失应该比训练前低（理想情况下）
	evalLoss, err := trainer.Evaluate(trainInputs, trainTargets, lossFn)
	if err != nil {
		t.Errorf("Post-training evaluation failed: %v", err)
	}

	if evalLoss < 0 {
		t.Errorf("Expected non-negative evaluation loss after training, got %v", evalLoss)
	}
}

// TestTrainerWithDifferentOptimizers 测试使用不同优化器的训练器
func TestTrainerWithDifferentOptimizers(t *testing.T) {
	// 创建模型
	mockLayer := NewMockLayer()
	model := &MockModel{
		Layer: mockLayer,
	}

	// 测试Adam优化器
	adamOpt := NewAdam(model.Parameters(), 0.001, 0.9, 0.999, 1e-8)
	trainer := NewTrainer(model, adamOpt)

	// 创建训练数据
	inputs := []*Tensor{
		Must(NewVector([]float64{1, 0})),
		Must(NewVector([]float64{0, 1})),
	}

	targets := []*Tensor{
		Must(NewVector([]float64{1, 0})),
		Must(NewVector([]float64{0, 1})),
	}

	// 损失函数
	lossFn := func(output, target *Tensor) *Tensor {
		diff, _ := output.Data.Subtract(target.Data)
		squared, _ := diff.Multiply(diff)

		// 计算矩阵元素的总和
		var total float64
		shape := squared.Shape()
		if len(shape) == 1 {
			for i := 0; i < squared.Size(); i++ {
				val, _ := squared.Get(i)
				total += val
			}
		} else if len(shape) == 2 {
			rows, cols := shape[0], shape[1]
			for i := 0; i < rows; i++ {
				for j := 0; j < cols; j++ {
					val, _ := squared.Get(i, j)
					total += val
				}
			}
		} else {
			// 对于其他维度，遍历所有元素
			for i := 0; i < squared.Size(); i++ {
				var val float64
				var err error

				if len(shape) == 1 {
					val, err = squared.Get(i)
				} else if len(shape) == 2 {
					cols := shape[1]
					val, err = squared.Get(i/cols, i%cols)
				} else {
					// 对于高维张量，按顺序访问元素
					val, err = squared.Get(i)
				}

				if err != nil {
					continue
				}

				total += val
			}
		}

		size := float64(output.Size())
		resultVal := total / size

		// 创建结果张量
		result, _ := NewTensor([]float64{resultVal}, []int{1})
		return result
	}

	// 训练
	err := trainer.Train(inputs, targets, 3, lossFn, false)
	if err != nil {
		t.Errorf("Training with Adam optimizer failed: %v", err)
	}

	evalLoss, err := trainer.Evaluate(inputs, targets, lossFn)
	if err != nil {
		t.Errorf("Evaluation with Adam optimizer failed: %v", err)
	}

	if evalLoss < 0 {
		t.Errorf("Expected non-negative evaluation loss with Adam, got %v", evalLoss)
	}
}