# -*- coding: iso-8859-1 -*-
from types import IntType,SliceType
import operator
from pysparse.sparse import spmatrix
from functools import reduce

class sparray:
    """
    d-dimensionnal sparse array emulation by a long sparse vector.
    supports syntax like:
    
    a[2,1,6,5]=2
    a[:,5,6,6]=list of appropriate length (one slice at a time only)
    b=a[:,5,6,6] (b is a numeric array)
    b=a[n], a[n]=6  if n in range
    
    a=sparray((2,6,9,8),dicto,shifts)
        where, optionnally, dicto is a dictionnary whose keys are tuple
        in range of (2,6,9,8), and shifts is a tuple
        to shift origins in case the smallest coordinate in dicto
        is not (0,...,0)
    
    """
    def __init__(self,dim,dicto=None,shifts=None):
        """
        attributes: shifts, dims, data, is1D, length
        methods : dump
        """

        self.shifts=shifts
        
        if type(dim)==type(()) or type(dim)==type([]):
            self.data = spmatrix.ll_mat(reduce(operator.mul, dim), 1)
            self.dims = dim        
            if dicto:
                for k, v in dicto.items():
                    shk = list(map(operator.__sub__, k, shifts))
                    self.data[self.comp(shk), 0]=v
        
        elif type(dim)==IntType:
            self.data = spmatrix.ll_mat(dim,1)
            self.dims = dim        
            if dicto:
                for k, v in dicto.items():
                    shk = k - shifts
                    self.data[shk,0] = v
        
        self.is1D = type(self.dims)==IntType
        
        

    def __get_shape0(self):return self.data.shape[0]
    length = property(__get_shape0, doc="sparray linear length")
    
    def decomp(self,ind):
        "from linear to multi indice"
        a = ind
        l = len(self.dims)
        res = [0]*l
        for i in range(l - 1, -1, -1):
            a, b = divmod(a,self.dims[i])
            res[i] = b
        return tuple(res)
    
    def comp(self,indice):
        "from multi indice to linear"
        l = len(self.dims)
        a = 0
        for i in range(l-1):
            a += reduce(operator.mul, self.dims[i+1:]) * indice[i]
        a += indice[l-1]
        return a
        
    def __setitem__(self, k, value):
        if type(k) is IntType:
            self.data[k, 0] = value
            return
        
        vec = [type(x) is SliceType for x in k]
        
        if True in vec: # suppose only one slice
            ii = vec.index(True)
            indices = []
            k = list(k)
            comp = self.comp
            for i in range(self.dims[ii]):
                k[ii] = i
                self.data[comp(k), 0] = value[i]
        else:
            self.data[self.comp(k),0]=value
    
    
    def __getitem__(self,k):
        """
        output a Numeric vector if slice in coordinates
        (one slice present only)
        """
        if type(k) is IntType: return self.data[k, 0]
        
        vec = [type(x) is SliceType for x in k]
        
        if True in vec: #suppose only one slice
            ii=vec.index(True)
            indices=[]
            k = list(k)
            import numpy
            rep = numpy.zeros((self.dims[ii],), 'd')
            for i in range(self.dims[ii]):
                k[ii] = i
                rep[i] = self.data[self.comp(k), 0]
            return rep
        else:
            return self.data[self.comp(k), 0]
            
    def dump(self):
        print(self.data)


