from numba import types, config
from numba.cuda.stubs import _vector_type_stubs


class SimulatedVectorType:
    attributes = ['x', 'y', 'z', 'w']

    def __init__(self, *args):
        args_flattened = []
        for arg in args:
            if isinstance(arg, SimulatedVectorType):
                args_flattened += arg.as_list()
            else:
                args_flattened.append(arg)
        self._attrs = self.attributes[:len(args_flattened)]
        if not self.num_elements == len(args_flattened):
            raise TypeError(
                f"{self.name} expects {self.num_elements}"
                f" elements, got {len(args_flattened)}"
            )

        for arg, attr in zip(args_flattened, self._attrs):
            setattr(self, attr, arg)

    @property
    def name(self):
        raise NotImplementedError()

    @property
    def num_elements(self):
        raise NotImplementedError()

    def as_list(self):
        return [getattr(self, attr) for attr in self._attrs]


def make_simulated_vector_type(num_elements, name):
    if config.USE_LEGACY_TYPE_SYSTEM:
        base_type = types.float32
    else:
        base_type = types.np_float32

    obj = type(name, (SimulatedVectorType,), {
        "num_elements": num_elements,
        "base_type": base_type,
        "name": name
    })
    obj.user_facing_object = obj
    return obj


def _initialize():
    _simulated_vector_types = {}
    for stub in _vector_type_stubs:
        num_elements = int(stub.__name__[-1])
        _simulated_vector_types[stub.__name__] = (
            make_simulated_vector_type(num_elements, stub.__name__)
        )
        _simulated_vector_types[stub.__name__].aliases = stub.aliases
    return _simulated_vector_types


vector_types = _initialize()
