package io.synor.compute;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.DisplayName;
import static org.junit.jupiter.api.Assertions.*;

import java.util.Arrays;

@DisplayName("Tensor Tests")
class TensorTest {

    @Test
    @DisplayName("Creates tensor with shape and data")
    void testCreateTensor() {
        double[] data = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
        int[] shape = {2, 3};
        Tensor tensor = new Tensor(shape, data);

        assertArrayEquals(shape, tensor.getShape());
        assertEquals(6, tensor.getSize());
        assertEquals(2, tensor.getNdim());
    }

    @Test
    @DisplayName("Creates zeros tensor")
    void testZeros() {
        Tensor tensor = Tensor.zeros(new int[]{3, 4});

        assertArrayEquals(new int[]{3, 4}, tensor.getShape());
        assertEquals(12, tensor.getSize());
        for (double v : tensor.getData()) {
            assertEquals(0.0, v);
        }
    }

    @Test
    @DisplayName("Creates ones tensor")
    void testOnes() {
        Tensor tensor = Tensor.ones(new int[]{2, 2});

        assertArrayEquals(new int[]{2, 2}, tensor.getShape());
        for (double v : tensor.getData()) {
            assertEquals(1.0, v);
        }
    }

    @Test
    @DisplayName("Creates random tensor")
    void testRandom() {
        Tensor tensor = Tensor.random(new int[]{10, 10});

        assertArrayEquals(new int[]{10, 10}, tensor.getShape());
        assertEquals(100, tensor.getSize());
        for (double v : tensor.getData()) {
            assertTrue(v >= 0 && v < 1, "Random values should be in [0, 1)");
        }
    }

    @Test
    @DisplayName("Creates randn tensor with normal distribution")
    void testRandn() {
        Tensor tensor = Tensor.randn(new int[]{1000});

        double mean = tensor.mean();
        double std = tensor.std();

        assertTrue(Math.abs(mean) < 0.2, "Mean should be close to 0");
        assertTrue(std > 0.8 && std < 1.2, "Std should be close to 1");
    }

    @Test
    @DisplayName("Creates identity matrix")
    void testEye() {
        Tensor tensor = Tensor.eye(3);

        assertArrayEquals(new int[]{3, 3}, tensor.getShape());
        assertEquals(1.0, tensor.get(0, 0));
        assertEquals(1.0, tensor.get(1, 1));
        assertEquals(1.0, tensor.get(2, 2));
        assertEquals(0.0, tensor.get(0, 1));
        assertEquals(0.0, tensor.get(1, 0));
    }

    @Test
    @DisplayName("Creates arange tensor")
    void testArange() {
        Tensor tensor = Tensor.arange(0, 5, 1);

        assertArrayEquals(new int[]{5}, tensor.getShape());
        assertEquals(0.0, tensor.get(0));
        assertEquals(4.0, tensor.get(4));
    }

    @Test
    @DisplayName("Creates linspace tensor")
    void testLinspace() {
        Tensor tensor = Tensor.linspace(0, 10, 11);

        assertArrayEquals(new int[]{11}, tensor.getShape());
        assertEquals(0.0, tensor.get(0), 0.0001);
        assertEquals(10.0, tensor.get(10), 0.0001);
        assertEquals(5.0, tensor.get(5), 0.0001);
    }

    @Test
    @DisplayName("Reshapes tensor")
    void testReshape() {
        Tensor tensor = new Tensor(new int[]{6}, new double[]{1, 2, 3, 4, 5, 6});
        Tensor reshaped = tensor.reshape(new int[]{2, 3});

        assertArrayEquals(new int[]{2, 3}, reshaped.getShape());
        assertEquals(6, reshaped.getSize());
    }

    @Test
    @DisplayName("Throws on invalid reshape")
    void testInvalidReshape() {
        Tensor tensor = Tensor.zeros(new int[]{4});

        assertThrows(IllegalArgumentException.class, () -> {
            tensor.reshape(new int[]{2, 3});
        });
    }

    @Test
    @DisplayName("Transposes 2D tensor")
    void testTranspose() {
        Tensor tensor = new Tensor(new int[]{2, 3}, new double[]{1, 2, 3, 4, 5, 6});
        Tensor transposed = tensor.transpose();

        assertArrayEquals(new int[]{3, 2}, transposed.getShape());
        assertEquals(1.0, transposed.get(0, 0));
        assertEquals(4.0, transposed.get(0, 1));
    }

    @Test
    @DisplayName("Gets element at index")
    void testGet() {
        Tensor tensor = new Tensor(new int[]{2, 3}, new double[]{1, 2, 3, 4, 5, 6});

        assertEquals(1.0, tensor.get(0, 0));
        assertEquals(3.0, tensor.get(0, 2));
        assertEquals(5.0, tensor.get(1, 1));
    }

    @Test
    @DisplayName("Calculates sum")
    void testSum() {
        Tensor tensor = new Tensor(new int[]{4}, new double[]{1, 2, 3, 4});

        assertEquals(10.0, tensor.sum());
    }

    @Test
    @DisplayName("Calculates mean")
    void testMean() {
        Tensor tensor = new Tensor(new int[]{4}, new double[]{1, 2, 3, 4});

        assertEquals(2.5, tensor.mean());
    }

    @Test
    @DisplayName("Calculates std")
    void testStd() {
        Tensor tensor = new Tensor(new int[]{4}, new double[]{1, 2, 3, 4});

        assertEquals(1.118, tensor.std(), 0.001);
    }

    @Test
    @DisplayName("Finds min")
    void testMin() {
        Tensor tensor = new Tensor(new int[]{4}, new double[]{3, 1, 4, 2});

        assertEquals(1.0, tensor.min());
    }

    @Test
    @DisplayName("Finds max")
    void testMax() {
        Tensor tensor = new Tensor(new int[]{4}, new double[]{3, 1, 4, 2});

        assertEquals(4.0, tensor.max());
    }

    @Test
    @DisplayName("Applies ReLU activation")
    void testRelu() {
        Tensor tensor = new Tensor(new int[]{5}, new double[]{-2, -1, 0, 1, 2});
        Tensor result = tensor.relu();

        assertArrayEquals(new double[]{0, 0, 0, 1, 2}, result.getData(), 0.0001);
    }

    @Test
    @DisplayName("Applies sigmoid activation")
    void testSigmoid() {
        Tensor tensor = new Tensor(new int[]{1}, new double[]{0});
        Tensor result = tensor.sigmoid();

        assertEquals(0.5, result.get(0), 0.0001);
    }

    @Test
    @DisplayName("Applies softmax activation")
    void testSoftmax() {
        Tensor tensor = new Tensor(new int[]{3}, new double[]{1, 2, 3});
        Tensor result = tensor.softmax();

        assertEquals(1.0, result.sum(), 0.0001);
        assertTrue(result.get(2) > result.get(1));
        assertTrue(result.get(1) > result.get(0));
    }

    @Test
    @DisplayName("Serializes and deserializes tensor")
    void testSerialization() {
        Tensor original = new Tensor(new int[]{2, 3}, new double[]{1, 2, 3, 4, 5, 6});
        var json = original.toJson();

        assertNotNull(json.get("shape"));
        assertNotNull(json.get("data"));
        assertNotNull(json.get("dtype"));

        Tensor restored = Tensor.fromJson(json);
        assertArrayEquals(original.getShape(), restored.getShape());
        assertArrayEquals(original.getData(), restored.getData(), 0.0001);
    }

    @Test
    @DisplayName("Converts to string representation")
    void testToString() {
        Tensor tensor = Tensor.zeros(new int[]{2, 3});
        String str = tensor.toString();

        assertTrue(str.contains("Tensor"));
        assertTrue(str.contains("2, 3") || str.contains("[2,3]"));
    }
}
