from __future__ import annotations
import torch as t

class SparseAct:
    """
    A SparseAct is a helper class which represents a vector in the sparse feature basis provided by an SAE, jointly with the SAE error term.
    A SparseAct may have three fields:
    act (... seq d_sae): the feature activations in the sparse basis
    res (... seq d_model): the SAE error term
    resc (... seq): a contracted SAE error term, useful for when we want one number per feature and error 
            (instead of having d_model numbers per error)
    """

    def __init__(
        self, 
        act: t.Tensor,
        res: t.Tensor | None = None,
        resc: t.Tensor | None = None, # contracted residual
        ) -> None:

        self.act = act
        self.res = res
        self.resc = resc

    def _map(self, f, aux=None) -> 'SparseAct':
        kwargs = {}
        
        # if isinstance(aux, SparseAct):
        if aux.__class__.__name__ == 'SparseAct': # NOTE: not recommended but this is for fixing relative import
            for attr in ['act', 'res', 'resc']:
                if getattr(self, attr) is not None and getattr(aux, attr) is not None:
                    kwargs[attr] = f(getattr(self, attr), getattr(aux, attr))
        else:
            for attr in ['act', 'res', 'resc']:
                if getattr(self, attr) is not None:
                    kwargs[attr] = f(getattr(self, attr), aux)
        return SparseAct(**kwargs)
        
    def __mul__(self, other) -> SparseAct:
        return self._map(lambda x, y: x * y, other)

    def __rmul__(self, other) -> SparseAct:
        # This will handle float/int * SparseAct by reusing the __mul__ logic
        return self.__mul__(other)
    
    def __matmul__(self, other: SparseAct) -> SparseAct:
        if other.res is not None and self.res is not None:
            # Normal mode, where act is features, and res is error term of SAE
            return SparseAct(act = self.act * other.act, resc=(self.res * other.res).sum(dim=-1))
        else:
            return SparseAct(act = self.act * other.act)
    
    def __add__(self, other) -> SparseAct:
        return self._map(lambda x, y: x + y, other)
    
    def __radd__(self, other: SparseAct) -> SparseAct:
        return self.__add__(other)
    
    def __sub__(self, other) -> SparseAct:
        return self._map(lambda x, y: x - y, other)
    
    def __rsub__(self, other) -> SparseAct:
        return self._map(lambda x, y: y - x, other)
    
    def __truediv__(self, other) -> SparseAct:
        if isinstance(other, SparseAct):
            kwargs = {}
            for attr in ['act', 'res', 'resc']:
                if getattr(self, attr) is not None:
                    kwargs[attr] = getattr(self, attr) / getattr(other, attr)
        else:
            kwargs = {}
            for attr in ['act', 'res', 'resc']:
                if getattr(self, attr) is not None:
                    kwargs[attr] = getattr(self, attr) / other
        return SparseAct(**kwargs)

    def __rtruediv__(self, other) -> SparseAct:
        if isinstance(other, SparseAct):
            kwargs = {}
            for attr in ['act', 'res', 'resc']:
                if getattr(self, attr) is not None:
                    kwargs[attr] = other / getattr(self, attr)
        else:
            kwargs = {}
            for attr in ['act', 'res', 'resc']:
                if getattr(self, attr) is not None:
                    kwargs[attr] = other / getattr(self, attr)
        return SparseAct(**kwargs)

    def __neg__(self) -> SparseAct:
        return self._map(lambda x, _: -x)
    
    def __invert__(self) -> SparseAct:
            return self._map(lambda x, _: ~x)
    
    def __getitem__(self, index: int):
        return self.act[index]
    
    def __repr__(self):
        return f"SparseAct(act={self.act}, res={self.res}), resc={self.resc}"
    
    def sum(self, dim=None):
        kwargs = {}
        for attr in ['act', 'res', 'resc']:
            if getattr(self, attr) is not None:
                kwargs[attr] = getattr(self, attr).sum(dim)
        return SparseAct(**kwargs)
    
    def mean(self, dim: int):
        kwargs = {}
        for attr in ['act', 'res', 'resc']:
            if getattr(self, attr) is not None:
                kwargs[attr] = getattr(self, attr).mean(dim)
        return SparseAct(**kwargs)

    @property
    def grad(self):
        kwargs = {}
        for attribute in ['act', 'res', 'resc']:
            if getattr(self, attribute) is not None:
                kwargs[attribute] = getattr(self, attribute).grad
        return SparseAct(**kwargs)
    
    def clone(self):
        kwargs = {}
        for attribute in ['act', 'res', 'resc']:
            if getattr(self, attribute) is not None:
                kwargs[attribute] = getattr(self, attribute).clone()
        return SparseAct(**kwargs)
    
    @property
    def value(self):
        kwargs = {}
        for attribute in ['act', 'res', 'resc']:
            if getattr(self, attribute) is not None:
                kwargs[attribute] = getattr(self, attribute).value
        return SparseAct(**kwargs)

    def save(self):
        return self._map(lambda x, _: x.save())
    
    def detach(self):
        return self._map(lambda x, _: x.detach())
    
    def to_tensor(self, contracted: bool = True):
        list_tens = [self.act]
        if self.res is not None:
            list_tens.append(self.res)
        if self.resc is not None:
            if contracted:
                assert len(self.resc.shape)+1 == len(self.act.shape), "Unvalid resc shape."
                list_tens.append(t.unsqueeze(self.resc, -1))
            else:
                assert len(self.resc.shape) == len(self.act.shape), "Unvalid resc shape."
                list_tens.append(self.resc)
        return t.cat(list_tens, dim=-1)
    
    def to_sparse_like_self(self, tens: t.Tensor, contracted: bool = True):
        assert tens.shape == self.to_tensor(contracted).shape
        if self.res is None and self.resc is None:
            return SparseAct(act=tens)
        elif self.resc is None and self.res is not None:
            return SparseAct(
                act = tens[..., :self.act.shape[-1]],
                res = tens[..., self.act.shape[-1]:]
            )
        elif self.res is None and self.resc is not None:
            return SparseAct(
                act = tens[..., :-1],
                resc = tens[..., -1]
            )
        else:
            return SparseAct(
                act = tens[..., :self.act.shape[-1]],
                res = tens[..., self.act.shape[-1]:(self.act.shape[-1]+self.res.shape[-1])], # type: ignore
                resc = tens[..., -1]
            )

    def to(self, device):
        for attr in ['act', 'res', 'resc']:
            if getattr(self, attr) is not None:
                setattr(self, attr, getattr(self, attr).to(device))
        return self

    def __eq__(self, other): # type: ignore
        return self._map(lambda x, y: x == y, other)
    
    def __gt__(self, other):
        return self._map(lambda x, y: x > y, other)
    
    def __lt__(self, other):
        return self._map(lambda x, y: x < y, other)
    
    def __le__(self, other):
        return self._map(lambda x, y: x <= y, other)
    
    def __ge__(self, other):
        return self._map(lambda x, y: x >= y, other)
    
    def nonzero(self):
        return self._map(lambda x, _: x.nonzero())
    
    def squeeze(self, dim):
        return self._map(lambda x, _: x.squeeze(dim=dim))
    
    def expand_as(self, other):
        return self._map(lambda x, y: x.expand_as(y), other)
    
    def zeros_like(self):
        return self._map(lambda x, _: t.zeros_like(x))
    
    def ones_like(self):
        return self._map(lambda x, _: t.ones_like(x))
    
    def abs(self):
        return self._map(lambda x, _: x.abs())
    
    def numel(self):
        numel = 0
        for attr in ['act', 'res', 'resc']:
            if getattr(self, attr) is not None:
                numel += getattr(self, attr).numel()
        return numel
    
    def item(self):
        if self.numel() == 1:
            return self.act.item()
        
    def contract(self):
        """
        Contract the residuals along the given dimension.
        """
        if self.resc is not None:
            self.resc = self.resc.sum(dim=-1)
        return self