'''
UFget.py - A hook into the University of Florida's Sparse Matrix Library
for use in Python.  For more detailed documentation, feel free to visit
the website at www.cs.hmc.edu/~koelze/

Written by Kevin Oelze, July 2008

'''


import urllib
import indexGet
import tarfile
import os
import pickle

#The folder you want to download files into
UFfolder = 'UFget/'

try:
#The dictionary that stores our index files
    fil = open('UFPickle.pk','rb')
    dictionary = pickle.load(fil)
except:
    try:
        print 'Files not found.  Refreshing index.'
        refresh()
        fil = open('UFPickle.pk','rb')
        dictioanry = pickle.load(fil)
    except:
        print 'Error refreshing.  Aborting for now.'
        quit()
            

#A dictionary-like class I use to hold the matrix's edges.  It extends
#dictionary and simply returns 0 if it's not in the dictionary
class ZeroDict(dict):
    #Only thing we need to overwrite.  Returns 0 instead of an error.
    def __getitem__(self,item):
        try:
            return dict.__getitem__(self,item)
        except:
            return 0


#The class where we store matrices and all of the valuable
#meta information from the UF Sparse Matrix Collection
class Matrix(object):
    #Reads in the appropriate data from the filename and
    #stores it in the fields
    def __init__(self,filename):
        #Let's parse that data!
        f = open(filename,'r')
        #The first line will tell us if the matrix is symmetric, skew-symmetric
        #or hermitian
        first = f.readline()
        if ('skew-symmetric' in first):
            self.type = 'skew-symmetric'
        elif('symmetric' in first):
            self.type = 'symmetric'
        elif('hermitian' in first or 'Hermitian' in first):
            self.type = 'hermitian'
        else:
            self.type = 'general'
        #Get rid of header stuff
        for i in range(3):
            f.readline()
        #Get the various parameters of the file
        self.name = " ".join(f.readline().strip().split(" ")[2:])
        self.title = " ".join(f.readline().strip().split(" ")[1:])
        self.id = int(" ".join(f.readline().strip().split(" ")[2:]))
        self.date = " ".join(f.readline().strip().split(" ")[2:])
        self.author = " ".join(f.readline().strip().split(" ")[2:])
        self.ed = " ".join(f.readline().strip().split(" ")[2:])
        #Throw away the line with fields..
        f.readline()
        self.kind = " ".join(f.readline().strip().split(" ")[2:])
        #Throw away the blank line
        f.readline()
        #Start parsing the data if there aren't comments, ignore any comments
        while(True):
            try:
                next = f.readline().split()
                self.nrows = int(next[0])
                self.ncols = int(next[1])
                entries = int(next[2])
                break
            except:
                pass
        self.dims = next[0] + ' by ' + next[1]
        self.name = filename
        self.A = ZeroDict()
        #Now the actual values
        for i in range(entries):
            next = f.readline().split()
            r = int(next[0])-1
            c = int(next[1])-1
            #Binary
            val = 1
            #Real/Integer
            if (len(next) == 3):
                val = float(next[2])
            #Complex-valued
            elif(len(next) == 4):
                val = float(next[2])+float(next[3])*1j
            self.A[r,c] = val
            if (r is not c):
                if (self.type is 'symmetric'):
                    self.A[c,r] = val
                elif(self.type is 'skew-symmetric'):
                    self.A[c,r] = -1*val
                elif(self.type is 'hermitian'):
                    self.A[c,r] = val.conjugate()
        #Correct the number of entries
        self.nnz = len(self.A)
                                        
                                                 
    #The string representation of this will show the metadata the
    #matlab stuff does too
    def __str__(self):
        res = ''
        res += 'Title: ' + self.title + '\n'
        res += 'Dimensions: ' + self.dims + '\n'
        res += 'Entries: ' + str(self.nnz) + '\n'
        res += 'Id: ' + str(self.id) + '\n'
        res += 'Name: ' + self.name + '\n'
        res += 'Date: ' + self.date + '\n'
        res += 'Author: ' + self.author + '\n'
        res += 'Ed: ' + self.ed + '\n'
        res += 'Kind: ' + self.kind + '\n'
        return res

    #Exact same as the string representation
    def __repr__(self):
        return str(self)
    
#Refreshes the index file.
def refresh():
    indexGet.getInfo()

#Download and extract the ith matrix from the sparse matrix library
def UFdownload(num):
    global UFfolder
    global dictionary
    baseURL = 'http://www.cise.ufl.edu/research/sparse/MM/'
    postfix = '.tar.gz'
    mtxfix = '.mtx'
    try: 
        data = dictionary[num]
    #They gave a string instead..
    except:
        data = [0, num, 0, 0]
    tarDir = baseURL + data[1]
    dlLoc = UFfolder + data[1][:data[1].find('/')]
    mtxFile = UFfolder + data[1] + data[1][data[1].find('/'):] + mtxfix
    tarFile = UFfolder + data[1] + postfix
    #Download and untar the file and make the directories (if necessary)
    if(not os.path.exists(mtxFile)):
        try:
            os.makedirs(os.path.dirname(tarFile))
        except:
            pass
        urllib.urlretrieve(tarDir +  postfix, tarFile)
        extract = tarfile.open(tarFile)
        extract.extractall(dlLoc)
        #Now get rid of that pesky tar file..
        os.remove(tarFile)
    else:
        print 'Matrix index', num, 'already exists, doing nothing.'
    #Return the file path of the .mtx file...
    return mtxFile

#Downloads all of the matrices in the UF Sparse Matrix Library
def UFgetAll():
    global dictionary
    for key in dictionary:
        print 'Downloading index number', key
        UFdownload(key)

#Download the matrix at index i from the
#sparse matrix database and read it into a Matrix object
def UFget(num):
    return Matrix(UFdownload(num))