#!/usr/bin/env python
# Fri Nov  6 17:28:57 EST 2009
# convert a matrix from one format to another

import sys
import optparse

from rasmus import util
from rasmus import matrixlib


o = optparse.OptionParser()
o.add_option("-o", "--outmatrix", dest="outmatrix", metavar="FILENAME")
o.add_option("-r", "--rowlabels", dest="rowlabels", metavar="FILENAME")
o.add_option("-c", "--collabels", dest="collabels", metavar="FILENAME")


conf, args = o.parse_args()
inputformat, outputformat = args


#=============================================================================
# inputs

rowlabels = None
collabels = None

if inputformat == "lmat":
    dat = util.read_delim(sys.stdin)
    for row in dat:
        row[2] = float(row[2])
    rowlabels = util.unique(row[0] for row in dat)
    collabels = util.unique(row[1] for row in dat)
    imat = list(util.unique(matrixlib.ilmat2imat(dat, rowlabels, collabels)))
    
else:
    raise Exception("unknown input format '%s'" % inputformat)


#=============================================================================
# setup output

if conf.outmatrix:
    out = open(conf.outmatrix, "w")
else:
    out = sys.stdout

if rowlabels and conf.rowlabels:
    util.write_list(conf.rowlabels, rowlabels)
if collabels and conf.collabels:
    util.write_list(conf.collabels, collabels)



#=============================================================================
# output matrix

if outputformat == "rmat":
    nrows = max(row[0] for row in imat) + 1
    ncols = max(row[1] for row in imat) + 1
    nnz = len(imat)
    rmat = matrixlib.imat2rmat(nrows, ncols, nnz, imat)    
    matrixlib.write_rmat(out, nrows, ncols, nnz, rmat)

elif outputformat == "srmat":

    nrows = max(row[0] for row in imat) + 1
    ncols = max(row[1] for row in imat) + 1

    assert nrows == ncols

    # permute columns
    lookup = util.list2lookup(rowlabels)
    cperm = util.mget(lookup, collabels)

    imat = map(list, imat)
    for row in imat:
        row[1] = lookup[collabels[row[1]]]

    nnz = None
    rmat = matrixlib.imat2rmat(nrows, ncols, nnz, imat)
    nnz = sum(map(len, rmat))

    matrixlib.write_rmat(out, nrows, ncols, nnz, rmat, square=True)
    
else:
    raise Exception("unknown output format '%s'" % outputformat)
    




    
