# frozen_string_literal: true

require 'minitest/autorun'
require_relative '../lib/synor_compute'

class TestTensor < Minitest::Test
  def test_creation
    tensor = SynorCompute::Tensor.new([2, 3], [1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
    assert_equal [2, 3], tensor.shape
    assert_equal 6, tensor.size
    assert_equal 2, tensor.ndim
  end

  def test_zeros
    tensor = SynorCompute::Tensor.zeros([3, 4])
    assert_equal [3, 4], tensor.shape
    assert_equal 12, tensor.size
    assert tensor.data.all? { |v| v == 0.0 }
  end

  def test_ones
    tensor = SynorCompute::Tensor.ones([2, 2])
    assert tensor.data.all? { |v| v == 1.0 }
  end

  def test_random
    tensor = SynorCompute::Tensor.rand([10, 10])
    assert_equal [10, 10], tensor.shape
    assert tensor.data.all? { |v| v >= 0 && v < 1 }
  end

  def test_randn
    tensor = SynorCompute::Tensor.randn([1000])
    mean = tensor.mean
    std = tensor.std
    assert mean.abs < 0.2, "Mean should be close to 0, got #{mean}"
    assert std > 0.8 && std < 1.2, "Std should be close to 1, got #{std}"
  end

  def test_eye
    tensor = SynorCompute::Tensor.eye(3)
    assert_equal [3, 3], tensor.shape
    assert_equal 1.0, tensor[[0, 0]]
    assert_equal 1.0, tensor[[1, 1]]
    assert_equal 0.0, tensor[[0, 1]]
  end

  def test_arange
    tensor = SynorCompute::Tensor.arange(0, 5, 1)
    assert_equal [5], tensor.shape
    assert_equal 0.0, tensor[[0]]
    assert_equal 4.0, tensor[[4]]
  end

  def test_linspace
    tensor = SynorCompute::Tensor.linspace(0, 10, 11)
    assert_equal [11], tensor.shape
    assert_in_delta 0.0, tensor[[0]], 0.0001
    assert_in_delta 10.0, tensor[[10]], 0.0001
  end

  def test_reshape
    tensor = SynorCompute::Tensor.new([6], [1, 2, 3, 4, 5, 6])
    reshaped = tensor.reshape([2, 3])
    assert_equal [2, 3], reshaped.shape
    assert_equal 6, reshaped.size
  end

  def test_transpose
    tensor = SynorCompute::Tensor.new([2, 3], [1, 2, 3, 4, 5, 6])
    transposed = tensor.transpose
    assert_equal [3, 2], transposed.shape
  end

  def test_sum
    tensor = SynorCompute::Tensor.new([4], [1, 2, 3, 4])
    assert_equal 10.0, tensor.sum
  end

  def test_mean
    tensor = SynorCompute::Tensor.new([4], [1, 2, 3, 4])
    assert_equal 2.5, tensor.mean
  end

  def test_min_max
    tensor = SynorCompute::Tensor.new([4], [3, 1, 4, 2])
    assert_equal 1.0, tensor.min
    assert_equal 4.0, tensor.max
  end

  def test_relu
    tensor = SynorCompute::Tensor.new([5], [-2, -1, 0, 1, 2])
    result = tensor.relu
    assert_equal [0, 0, 0, 1, 2], result.data
  end

  def test_sigmoid
    tensor = SynorCompute::Tensor.new([1], [0])
    result = tensor.sigmoid
    assert_in_delta 0.5, result[[0]], 0.0001
  end

  def test_softmax
    tensor = SynorCompute::Tensor.new([3], [1, 2, 3])
    result = tensor.softmax
    assert_in_delta 1.0, result.sum, 0.0001
    assert result[[2]] > result[[1]]
  end

  def test_serialization
    original = SynorCompute::Tensor.new([2, 3], [1, 2, 3, 4, 5, 6])
    json = original.to_json
    restored = SynorCompute::Tensor.from_json(json)
    assert_equal original.shape, restored.shape
    assert_equal original.data, restored.data
  end
end