from .tsn import tsn_to_str_list, tsn_to_tuple


def int_shape(*s):
    if len(s) == 1: 
        assert isinstance(s, (tuple,list))
        s = s[0]
    else: s = tuple(s)
    return tuple([int(d) for d in s])

def select(x, dv_dict, squeeze=False):
    '''
    Index using dimension shorthands
    
    x: (t, 'bcd') -- tensor, shape tuple (can be indexed in numpy notation : x[:,0,:])
    dv_dict: {'b': 0, 'c': 5} 
    squeeze: [True, False] or a tsn list ('b,c') of dims to be squeezed
    '''
    assert not squeeze, 'not implemented'

    assert isinstance(tuple), 'The first argument should be a tuple of (vector, shape)'
    xv, xs = x
    shape, is_seq = tsn_to_str_list(xs)
    if not is_seq:
        raise NotImplementedError(f"get from shape {xs} not implemented")

    colon = slice(None)
    slice_tuple = [colon] * len(shape)
    for pos, sh in shape:
        if sh in dv_dict:
            slice_tuple[pos] = dv_dict[sh]

    y = x[slice_tuple]
    return y

def size_assert(x_size, sa, dims=None):
    '''
    x_size: integer tuple
    sa: TSA
    dims: None or Sequence[int], e.g., [0,1]
    Check if size of tensor x matches TSA `sa` along `dims` axes
    '''
    x_size, sa = tuple(x_size), tuple(sa)
    if dims is not None:
        assert isinstance(dims, (list, tuple))
        x_size = [x_size[d] for d in dims]
        sa = [sa[d] for d in dims]

    if x_size != sa:
        print(f'Size mismatch: size = {x_size}, expected: {sa}')
        assert False


def reduce_dims (tfm):
    '''
    tfm: str, 'btd->b'
    '''
    src, to = tfm.split('->')
    src = tsn_to_tuple(src.strip())
    to = tsn_to_tuple(to.strip())

    assert isinstance(src, tuple)
    assert isinstance(to, tuple)

    drops = []
    #check src includes all dims in to
    for i, d in enumerate(src):
        if d not in to:
            drops.append(i)

    return tuple(drops)

