from FiniteFieldMatrix import FiniteFieldMatrix
from FiniteField import FiniteFieldGF257

def _create_transform_matrix(k, f):
    """ Used by both the encoder and decoder to generate a transform
        matrix used to go from k pieces to f pieces.
        For the encoder it is used directly for the decoder a subset is used
    """
    a = []
    for i in range(0, f-k):
        a.append([])
        for j in range(0, k):
            a[i].append((FiniteFieldGF257(2) ** (FiniteFieldGF257(f)-FiniteFieldGF257(k)+i)) ** j)
    A = FiniteFieldMatrix(a, class_type=FiniteFieldGF257)
    
    b = []        
    for i in range(0, k):
        b.append([])
        for j in range(0, k):
            if (i == 0 and j == 0):
                b[i].append(FiniteFieldGF257(1))
            elif (i == 0):
                b[i].append(FiniteFieldGF257(0))
            else:
                b[i].append(FiniteFieldGF257((FiniteFieldGF257(2) ** (FiniteFieldGF257(i)-1)) ** j))
    B = FiniteFieldMatrix(b, class_type=FiniteFieldGF257)
    
    B_inv = B.inverse()
    C = A * B_inv
    E = []
    for i in range(0, k):
        E.append([])
        for j in range(0, k):
            if i == j:
                E[i].append(FiniteFieldGF257(1))
            else:
                E[i].append(FiniteFieldGF257(0))
    for i in range(k, f):
        E.append([])
        for j in range(0, k):
            E[i].append(C.m[i-k][j])
    return FiniteFieldMatrix(E, class_type=FiniteFieldGF257)

class Decoder(object):
    """ Decode an IDA encoded matrix data
        encoded_data - an array of storage bin objects
                       Storage bin objects must have two items
                        * matrix_data() - returns an array of bytes
                        * series_id - property representing the sequence number
                                      of the piece of data (between 0 and f-1)
    """
    def __init__(self, encoded_data, f):
        self.encoded_data = encoded_data
        self.f = f
        
    def _encoded_matrix(self):
        """ Build a GF(2^8) Matrix out of the matrix data """
        em = []
        for storage_bin in self.encoded_data:
            # NOTE: the storage_bin has to have a matrix_data function
            #       that returns the data actually in the storage bin
            em.append(storage_bin.matrix_data())
        return FiniteFieldMatrix(em, class_type=FiniteFieldGF257)
    
    def _translate_matrix(self):
        """ builds the subset of the transform matrix that will be
            used to decode the encoded message
            
            Note that the encoded data is expected to have an attribute
            called series_id that is the sequence it is stored in
        """
        trans = _create_transform_matrix(len(self.encoded_data),
                                         self.f)
        new_trans = []
        for storage_bin in self.encoded_data:
            # pick out the transform matrix row corresponding to the
            # to the series_id found in the storage_bin
            new_trans.append(trans.m[storage_bin.series_id])
        return FiniteFieldMatrix(new_trans,
                                 class_type=FiniteFieldGF257)
    
    def decode_matrix(self):
        """ create the encoding matrix, take the inverse then multiply by
            the encoded data.
        """
        # subset of the transform matrix
        trans = self._translate_matrix()
        # the encoded data as a matrix
        encoded = self._encoded_matrix()
        # find the inverse matrix of the transform matrix
        trans_inv = trans.inverse();
        # mutliple the inverse of the transform matrix by the encoded data
        decoded = trans_inv * encoded
        return decoded
    
class Encoder(object):
    """ Returns the encoded data given a data matrix """
    def __init__(self, data_matrix, number_of_fragments):
        self.data_matrix = data_matrix
        self.number_of_fragments = number_of_fragments

    def encoded_matrix(self):
        trans = _create_transform_matrix(self.data_matrix.rows(),
                                         self.number_of_fragments)
        encoded = trans * self.data_matrix
        return encoded

import unittest
class TestIda(unittest.TestCase):
    def setUp(self):
        self.test_file = "/etc/hosts"

    def testIt(self):
        from Splitter import FileSplitter
        k = 7
        f = 10
        fs = FileSplitter(self.test_file, k)
        M = fs.matrix()
        ida = Encoder(M, f)
        E = ida.encoded_matrix()
        self.failUnless(E.rows() == f, "Number of matrix rows do not match f (%d vs %d)" % (E.rows(),f))
        
    def testDecoder(self):
        from Splitter import FileSplitter
        from Splitter import Joiner
        k = 7
        f = 10
        fs = FileSplitter(self.test_file, k)
        M = fs.matrix()
        enc = Encoder(M, f)
        E = enc.encoded_matrix()
        self.failUnless(E.rows() == f, "Number of matrix rows do not match f (%d vs %d)" % (E.rows(),f))
        class samplesb:
            def __init__(self,data,series_id):
                self.data = data
                self.series_id = series_id
            def matrix_data(self):
                return self.data
        recovered_data = []
        for i in [7,3,2,8,1,5,9]:
            recovered_data.append(samplesb(E.m[i],i))
        dec = Decoder(recovered_data, f)
        D = dec.decode_matrix()
        j = Joiner(D)
        orig_file = file(self.test_file)
        orig_data = orig_file.read()
        self.failUnless(str(j) == orig_data, "Failed to properly decode the original file")
        
if __name__ == "__main__":
    unittest.main()