import numpy as np
import scipy as sp

class BaseSystemGet:
    """ Get methods for the base system class"""
    def get_sym(self, tensor: np.ndarray, i: int, j: int) -> np.ndarray:
        """Gets the i,j component of a symmetric tensor saved in an array structure.

        Args: 
            tensor (numpy.ndarray): The symmetric tensor.
            i: The row index. (int)
            j: The column index. (int)
        
        Returns:
            The i,j component of the tensor. (numpy.ndarray) 
        """

        if self.dim == 2:
            if i == 0:
                return tensor[0] if j == 0 else tensor[1]
            elif i == 1:
                return tensor[1] if j == 0 else tensor[2]

        elif self.dim == 3:
            if i == 0:
                return tensor[0] if j == 0 else tensor[1] if j == 1 else tensor[2]
            elif i == 1:
                return tensor[1] if j == 0 else tensor[3] if j == 1 else tensor[4]
            elif i == 2:
                return tensor[2] if j == 0 else tensor[4] if j == 1 else tensor[5]

    def get_sym_tl(self, tensor: np.ndarray, i: int, j: int) -> np.ndarray:
        """Gets the i,j component of a symmetric traceless tensor saved in an array structure.

        Args:
            tensor: The symmetric traceless tensor. (numpy.ndarray)
            i: The row index. (int)
            j: The column index. (int)
        
        Returns:
            The i,j component of the tensor. (numpy.ndarray) 
        """
        if self.dim == 2:
            if i == 0:
                return tensor[0] if j == 0 else tensor[1]
            elif i == 1:
                return tensor[1] if j == 0 else -tensor[0]
            
        elif self.dim == 3:
            if i == 0:
                return tensor[0] if j == 0 else tensor[1] if j == 1 else tensor[2]
            elif i == 1:
                return tensor[1] if j == 0 else tensor[3] if j == 1 else tensor[4]
            elif i == 2:
                return tensor[2] if j == 0 else tensor[4] if j == 1 else -(tensor[0] + tensor[3])

    def get_anti_sym(self, omega: np.ndarray, i: int, j: int) -> np.ndarray:
        """Gets the i,j component of an anti-symmetric tensor saved in an array structure.

        Args:
            omega: The anti-symmetric tensor. (numpy.ndarray)
            i: The row index. (int)
            j: The column index. (int)
        
        Returns:
            The i,j component of the tensor. (numpy.ndarray) 
        """
        # TODO: I don't like that the input vector is a scalar field in 2 dimensions. (Vidar 11.03.24)
        if self.dim == 2:
            if i==j:
                return 0
            return (-1)**i *omega
        elif self.dim ==3:
            if i ==j:
                return 0
            else:
                return np.sign(j-i)*omega[i+j-1]
