#!/usr/bin/env python
# Created by: Robert Cimrman, 05.12.2005

"""Benchamrks for umfpack module"""

from __future__ import division, print_function, absolute_import

from optparse import OptionParser
import time
import urllib.request, urllib.parse, urllib.error
import gzip

import numpy as np

import sys

import scipy.sparse as sp
import scipy.sparse.linalg.dsolve.umfpack as um
import scipy.linalg as nla

defaultURL = 'http://www.cise.ufl.edu/research/sparse/HBformat/'

usage = """%%prog [options] <matrix file name> [<matrix file name>, ...]

<matrix file name> can be a local or distant (gzipped) file

default url is:
        %s

supported formats are:
        triplet .. [nRow, nCol, nItem] followed by 'nItem' * [ir, ic, value]
        hb      .. Harwell-Boeing format N/A
""" % defaultURL


##
# 05.12.2005, c
def read_triplet( fd ):
    nRow, nCol = map( int, fd.readline().split() )
    nItem = int( fd.readline() )

    ij = np.zeros( (nItem,2), np.int32 )
    val = np.zeros( (nItem,), np.float64 )
    for ii, row in enumerate( fd.readlines() ):
        aux = row.split()
        ij[ii] = int( aux[0] ), int( aux[1] )
        val[ii] = float( aux[2] )

    mtx = sp.csc_matrix( (val, ij), dims = (nRow, nCol), nzmax = nItem )

    return mtx

##
# 06.12.2005, c
def read_triplet2( fd ):
    nRow, nCol = map( int, fd.readline().split() )
    nItem = int( fd.readline() )

    ij, val = io.read_array( fd,
                             columns = [(0,1), (2,)],
                             atype = (np.int32, np.float64),
                             rowsize = nItem )

    mtx = sp.csc_matrix( (val, ij), dims = (nRow, nCol), nzmax = nItem )

    return mtx


formatMap = {'triplet' : read_triplet}
##
# 05.12.2005, c
def readMatrix( matrixName, options ):

    if options.default_url:
        matrixName = defaultURL + matrixName

    print('url:', matrixName)

    if matrixName[:7] == 'http://':
        if sys.version_info[0] >= 3:
            fileName, status = urllib.request.urlretrieve( matrixName )
        else:
            fileName, status = urllib.urlretrieve( matrixName )
##        print status
    else:
        fileName = matrixName

    print('file:', fileName)

    try:
        readMatrix = formatMap[options.format]
    except:
        raise ValueError('unsupported format: %s' % options.format)

    print('format:', options.format)

    print('reading...')
    if fileName.endswith('.gz'):
        fd = gzip.open( fileName )
    else:
        fd = open( fileName )

    mtx = readMatrix( fd )

    fd.close()

    print('ok')

    return mtx

##
# 05.12.2005, c
def main():
    parser = OptionParser( usage = usage )
    parser.add_option( "-c", "--compare",
                       action = "store_true", dest = "compare",
                       default = False,
                       help = "compare with default scipy.sparse solver [default: %default]" )
    parser.add_option( "-p", "--plot",
                       action = "store_true", dest = "plot",
                       default = False,
                       help = "plot time statistics [default: %default]" )
    parser.add_option( "-d", "--default-url",
                       action = "store_true", dest = "default_url",
                       default = False,
                       help = "use default url [default: %default]" )
    parser.add_option( "-f", "--format", type = type( '' ),
                       dest = "format", default = 'triplet',
                       help = "matrix format [default: %default]" )
    (options, args) = parser.parse_args()

    if (len( args ) >= 1):
        matrixNames = args;
    else:
        parser.print_help(),
        return

    sizes, nnzs, times, errors = [], [], [], []
    legends = ['umfpack', 'sparse.solve']
    for ii, matrixName in enumerate( matrixNames ):

        print('*' * 50)
        mtx = readMatrix( matrixName, options )

        sizes.append( mtx.shape )
        nnzs.append( mtx.nnz )
        tts = np.zeros( (2,), dtype = np.double )
        times.append( tts )
        err = np.zeros( (2,2), dtype = np.double )
        errors.append( err )

        print('size              : %s (%d nnz)' % (mtx.shape, mtx.nnz))

        sol0 = np.ones( (mtx.shape[0],), dtype = np.double )
        rhs = mtx * sol0

        umfpack = um.UmfpackContext()

        tt = time.clock()
        sol = umfpack( um.UMFPACK_A, mtx, rhs, autoTranspose = True )
        tts[0] = time.clock() - tt
        print("umfpack           : %.2f s" % tts[0])

        error = mtx * sol - rhs
        err[0,0] = nla.norm( error )
        print('||Ax-b||          :', err[0,0])

        error = sol0 - sol
        err[0,1] = nla.norm( error )
        print('||x - x_{exact}|| :', err[0,1])

        if options.compare:
            tt = time.clock()
            sol = sp.solve( mtx, rhs )
            tts[1] = time.clock() - tt
            print("sparse.solve      : %.2f s" % tts[1])

            error = mtx * sol - rhs
            err[1,0] = nla.norm( error )
            print('||Ax-b||          :', err[1,0])

            error = sol0 - sol
            err[1,1] = nla.norm( error )
            print('||x - x_{exact}|| :', err[1,1])

    if options.plot:
        try:
            import pylab
        except ImportError:
            raise ImportError("could not import pylab")
        times = np.array( times )
        print(times)
        pylab.plot( times[:,0], 'b-o' )
        if options.compare:
            pylab.plot( times[:,1], 'r-s' )
        else:
            del legends[1]

        print(legends)

        ax = pylab.axis()
        y2 = 0.5 * (ax[3] - ax[2])
        xrng = list(range( len( nnzs )))
        for ii in xrng:
            yy = y2 + 0.4 * (ax[3] - ax[2])\
                 * np.sin( ii * 2 * np.pi / (len( xrng ) - 1) )

            if options.compare:
                pylab.text( ii+0.02, yy,
                            '%s\n%.2e err_umf\n%.2e err_sp'
                            % (sizes[ii], np.sum( errors[ii][0,:] ),
                               np.sum( errors[ii][1,:] )) )
            else:
                pylab.text( ii+0.02, yy,
                            '%s\n%.2e err_umf'
                            % (sizes[ii], np.sum( errors[ii][0,:] )) )
            pylab.plot( [ii, ii], [ax[2], ax[3]], 'k:' )

        pylab.xticks( xrng, ['%d' % (nnzs[ii] ) for ii in xrng] )
        pylab.xlabel( 'nnz' )
        pylab.ylabel( 'time [s]' )
        pylab.legend( legends )
        pylab.axis( [ax[0] - 0.05, ax[1] + 1, ax[2], ax[3]] )
        pylab.show()

if __name__ == '__main__':
    main()
