"""This module implementation functionality to read and write files in
PyML format.

class cannot contain ; or :
class should be string
class cannot end with ','


Functions:
save_vector
load_vector
save_sparse
load_sparse

"""

def save_vector(handle, data):
    """

    """
    # data is a list of (id, klasses, vector)
    # id and klasses can be None
    
    # Convert the vectors to dictionaries.
    data = [(id, klasses, _vec2dict(vec)) for (id, klasses, vec) in data]
    save_sparse(handle, data)

def load_vector(handle):
    """Return a list of (id, klasses, vector)."""
    assert hasattr(handle, 'readline'), "I need a handle"
    data = load_sparse(handle)

    # The keys of the data are the dimensions of the vector.
    data = [(id, klasses, _str2int_keys(dict)) for (id, klasses, dict) in data]

    # The number of dimensions of the vector is the maximum dimension
    # in the dictionary.  Add 1, since the vector dimension starts at
    # 0.
    ndims = 0
    for x, x, dict in data:
        if dict:
            ndims = max(ndims, max(dict.keys())+1)

    # Convert the dictionaries to vectors.
    data = [(id, klasses, _dict2vec(dict, ndims))
            for (id, klasses, dict) in data]
    return data

def save_sparse(handle, data):
    # data is a list of (id, klasses, dict)
    # id and klasses can be None

    for id, klasses, dict in data:
        cols = []
        if id is not None:
            cols.append("%s," % id)
        if klasses:
            klasses = map(str, klasses)
            cols.append(";".join(klasses))
        fids = dict.keys()
        fids.sort()
        for fid in fids:
            fval = dict[fid]
            cols.append("%s:%s" % (fid, fval))
        handle.write("%s\n" % (" ".join(cols)))
    handle.flush()

def load_sparse(handle):
    """Return a list of (id, klasses, dict)."""
    assert hasattr(handle, 'readline'), "I need a handle"
    data = []
    for line in handle:
        cols = line.rstrip().split(" ")
        id, klasses, dict = None, None, {}

        id = klasses = None
        dict = {}
        if cols[0].endswith(","):
            id = cols[0][:-1]
            del cols[0]
        if cols[0].find(":") < 0:
            klasses = cols[0].split(";")
            del cols[0]
        # cols might be empty if there is an empty vector
        #assert cols, "No data found.  Error in file?"
        for c in cols:
            i = c.rindex(":")
            fid, fval = c[:i], c[i+1:]
            dict[fid] = float(fval)
        data.append((id, klasses, dict))
    return data

def _str2int_keys(dict):
    """Return dict, where the keys have been converted to integers."""
    intdict = {}
    for k, v in dict.items():
        intdict[int(k)] = v
    return intdict

def _dict2vec(dict, ndims):
    vector = [0]*ndims
    for i, v in dict.items():
        vector[int(i)] = v
    return vector

def _vec2dict(vector):
    dict = {}
    for i in range(len(vector)):
        if vector[i]:
            dict[str(i)] = vector[i]
    return dict