tensor

from typing import Self
import unittest

# O(k)
def product(dimensions: list[int]):
    p = 1
    for d in dimensions:
        p *= d
    return p

class Tensor:
    index = []
    def __init__(self, dimensions: list[int]):
        self.dimensions = dimensions
        p = product(dimensions)
        self.values: list = [None for n in range(p)]
        def assign(index, value):
            self.values[index] = value
        self.assign = assign

    @property
    def order(self):
        return len(self.dimensions)

    # O(k)
    def compute_tensor_index(self, index: int) -> list[int]:
        final = [0 for _ in self.dimensions]
        p = product(self.dimensions) # this is computed too many times for the same array 'dimensions'
        r = 0
        for i, dimension in enumerate(self.dimensions):
            p //= dimension
            final[i] = (index - r) // p
            r += final[i]*p
        return final

    # O(k)
    def compute_linear_index(self, index: list[int]) -> int:
        if len(index) != len(self.dimensions):
            raise Exception("Index's length and dimensions's length MUST be equal.")
        p, final = 1, 0
        for i, dimension in enumerate(self.dimensions[::-1]):
            final += index[-i-1]*p
            p *= dimension
        return final

    def __getitem__(self, index):
        if type(index) == list:
            if len(index) == len(self.dimensions):
                pos = self.compute_linear_index(index)
                return self.values[pos]
            else:
                raise IndexError("Index and dimensions MUST have the same length.")
        elif len(self.index) == len(self.dimensions)-1:
            index = [*self.index, index]
            pos = self.compute_linear_index(index)
            return self.values[pos]
        else:
            newt = Tensor(self.dimensions[:])
            newt.values = self.values[:]
            newt.index = [*self.index, index]
            newt.assign = self.assign
            return newt

    def __setitem__(self, index, value):
        index = [*self.index, index]
        if len(index) == len(self.dimensions):
            pos = self.compute_linear_index(index)
            self.assign(pos, value)
        elif isinstance(value, Tensor):
            if len(self.dimensions) == len(index) + len(value.dimensions):
                index_extended = [*index, *[0 for _ in range(len(value.dimensions))]]
                pos = self.compute_linear_index(index_extended)
                for i, x in enumerate(value.values):
                    self.assign(pos + i, x)
            else:
                raise IndexError("Tensor provided is not compatible with current index.")
        else:
            raise IndexError("Not enough indexes provided.")

    def __mul__(self, other: Self):
        newt = Tensor([*self.dimensions, *other.dimensions])
        p2 = product(other.dimensions)
        for pos1, v1 in enumerate(self.values):
            for pos2, v2 in enumerate(other.values):
                # first approach - bad!
                index1 = self.compute_tensor_index(pos1) # O(k1)
                index2 = other.compute_tensor_index(pos2) # O(k2)
                newpos = newt.compute_linear_index([*index1,*index2]) # O(k1+k2)

                # second approach - good!
                newpos2 = pos1 * p2 + pos2
                assert newpos == newpos2
                newt.assign(newpos, v1 * v2)

        return newt

    def contraction(self, i, j):
        if self.dimensions[i] == self.dimensions[j]:
            dimensions = [d for k, d in enumerate(self.dimensions) if k not in [i,j]]
            newt = Tensor(dimensions)
            newt.values = [0 for _ in range(product(dimensions))]

            for pos, v in enumerate(self.values):
                index = self.compute_tensor_index(pos)
                if index[i] == index[j]:
                    new_index = [d for k, d in enumerate(index) if k not in [i,j]]
                    new_pos = newt.compute_linear_index(new_index)
                    newt.values[new_pos] += v
            return newt
        else:
            raise IndexError("Bad indices, they MUST have the same dimension.")

class TestTensorMethods(unittest.TestCase):
    def test_computes(self):
        linear_index = 2*4*5+3*5+4
        tensor_index = [2,3,4]
        dimensions = [3,4,5]
        tensor = Tensor(dimensions)
        self.assertEqual(
            tensor.compute_linear_index(tensor_index),
            linear_index
        )
        self.assertListEqual(
            tensor.compute_tensor_index(linear_index),
            tensor_index
        )

    def test_element_assign(self):
        dimensions = [3,4,5,6]
        t1 = Tensor(dimensions)
        t1.values = list(range(product(dimensions)))
        self.assertEqual(t1.order, len(dimensions))
        for i in range(3):
            for j in range(4):
                for k in range(5):
                    for l in range(6):
                        v = i*4*5*6 + j*5*6 + k*6 + l
                        self.assertEqual(t1[i][j][k][l], v)
                        t1[i][j][k][l] = v+1
                        self.assertEqual(t1[i][j][k][l], v+1)

    def test_tensor_assign(self):
        d1, d2 = [3,4,5,6], [5,6]
        t1, t2 = Tensor(d1), Tensor(d2)
        t1.values = list(range(product(d1)))
        for k in range(5):
            for l in range(6):
                t2[k][l] = k*6+l
        for i in range(3):
            for j in range(4):
                t1[i][j] = t2
                for k in range(5):
                    for l in range(6):
                        self.assertEqual(t1[i][j][k][l], k*6 + l)

    def test_tensor_mult(self):
        d1, d2 = [3,4], [5,6,7]
        t1, t2 = Tensor(d1), Tensor(d2)
        t1.values = list(range(product(d1)))
        t2.values = list(range(product(d2)))
        t3 = t1 * t2
        self.assertEqual(t1.order, len(d1))
        self.assertEqual(t2.order, len(d2))
        for i in range(3):
            for j in range(4):
                for k in range(5):
                    for l in range(6):
                        for m in range(7):
                            self.assertEqual(
                                t1[i][j] * t2[k][l][m],
                                t3[i][j][k][l][m]
                            )

    def test_matrix_mult(self):
        d1 = d2 = [2,2]
        t1 = Tensor(d1)
        t2 = Tensor(d2)
        t1.values = [1,2,3,4]
        t2.values = [5,6,7,8]
        newt = (t1*t2).contraction(1,2)
        self.assertListEqual(newt.values, [19,22,43,50])



if __name__ == '__main__':
    unittest.main()