"""This holds code that deals with many file formats.

Functions:
open_training       gene, drug, has_relationship
open_scan           pmid, sentnum, lex phrase, text phrase, start, end, score
open_cooccurrences  pmid, <SCAN1>, <SCAN2>
open_observations   gene, drug, has_relationship, vector, index
open_results        gene, drug, has_rel, prediction, p0, p1, is_correct
open_score          ScoreData data structure

open_vocabulary  word, index
open_wordcounts  word, count
open_sentences   pmid, snum, GENE_INFO, DRUG_INFO, sentence, has_rel, shows_rel
open_sentence_observations  pmid, snum, gene_start, drug_start, vector, index

"""
from Extracto import filefns
lwrite = filefns.lwrite

# lexicon phrase     actual phrase found in the lexicon
# text phrase        instantiation of the phrase found in the text

class _HandleOpener:
    def __init__(self, reader_fn, writer_fn):
        self._reader_fn = reader_fn
        self._writer_fn = writer_fn
    def __call__(self, filename, mode='r'):
        if mode not in ['r', 'w']:
            raise ValueError, "Invalid mode %r" % mode
        if mode == 'r':
            if self._reader_fn is None:
                raise AssertionError, "Sorry, this file cannot be read"
            return self._reader_fn(filename)
        elif mode == 'w':
            if self._writer_fn is None:
                raise AssertionError, "Sorry, this file cannot be written"
            return self._writer_fn(filename)

class _training_reader:
    # Return gene, drug, has_relationship
    def __init__(self, filename):
        self.handle = filefns.openfh(filename)
    def next(self):
        line = self.handle.readline()
        if not line:
            raise StopIteration
        x = line.rstrip().split("|", 2)
        gene, drug, has_rel = x
        try:
            # now can be a string!
            has_rel = int(has_rel)
        except ValueError, x:
            pass
        return gene, drug, has_rel
    def __iter__(self):
        return self

class _training_writer:
    def __init__(self, filename):
        self.handle = filefns.LockedHandle(filefns.openfh(filename, 'w'))
    def close(self):
        self.handle.close()
    def flush(self):
        self.handle.flush()
    def write(self, gene, drug, has_relationship):
        self.handle.write("%s|%s|%s\n" % (gene, drug, has_relationship))

open_training = _HandleOpener(_training_reader, _training_writer)


class _scan_reader:
    # Return pmid, sentnum, lexicon_phrase, text_phrase, tp_start, tp_end, score
    def __init__(self, filename):
        self.handle = filefns.openfh(filename)
    def next(self):
        line = self.handle.readline()
        if not line:
            raise StopIteration
        x = line.rstrip().split("\t")
        pmid, snum, lex_phrase, text_phrase, start, end, score = x
        snum, start, end = map(int, (snum, start, end))
        if score == 'None':
            score = None
        else:
            score = float(score)
        return pmid, snum, lex_phrase, text_phrase, start, end, score
    def __iter__(self):
        return self

class _scan_writer:
    def __init__(self, filename):
        self.handle = filefns.LockedHandle(filefns.openfh(filename, 'w'))
    def close(self):
        self.handle.close()
    def flush(self):
        self.handle.flush()
    def write(self, pmid, snum, lex_phrase, text_phrase, start, end, score):
        self.handle.write("%s\t%d\t%s\t%s\t%d\t%d\t%s\n" %
                          (pmid, snum, lex_phrase,
                           text_phrase, start, end, score))

open_scan = _HandleOpener(_scan_reader, _scan_writer)


class _cooccurrences_reader:
    # Return pmid, sentnum1, lex_phrase1, text_phrase1, start1, end1,
    #              sentnum2, lex_phrase2, text_phrase2, start2, end2,
    # start and end are offsets from beginning of document
    def __init__(self, filename):
        self.handle = filefns.openfh(filename)
    def next(self):
        line = self.handle.readline()
        if not line:
            raise StopIteration
        x = line.rstrip().split("\t")
        pmid, snum1, lex_phrase1, text_phrase1, start1, end1, \
              snum2, lex_phrase2, text_phrase2, start2, end2 = x
        snum1, start1, end1 = map(int, (snum1, start1, end1))
        snum2, start2, end2 = map(int, (snum2, start2, end2))
        return pmid, snum1, lex_phrase1, text_phrase1, start1, end1, \
               snum2, lex_phrase2, text_phrase2, start2, end2
    def __iter__(self):
        return self

class _cooccurrences_writer:
    def __init__(self, filename):
        self.handle = filefns.LockedHandle(filefns.openfh(filename, 'w'))
    def close(self):
        self.handle.close()
    def flush(self):
        self.handle.flush()
    def write(self, pmid, snum1, lex_phrase1, text_phrase1, start1, end1,
              snum2, lex_phrase2, text_phrase2, start2, end2):
        self.handle.write("%s\t%d\t%s\t%s\t%d\t%d\t%d\t%s\t%s\t%d\t%d\n" %
                          (pmid, snum1, lex_phrase1,
                           text_phrase1, start1, end1,
                           snum2, lex_phrase2, text_phrase2, start2, end2))

open_cooccurrences = _HandleOpener(
    _cooccurrences_reader, _cooccurrences_writer)


class _observations_reader:
    # Return gene, drug, has_relationship, vector, index
    def __init__(self, filename):
        self.handle = filefns.openfh(filename)
    def next(self):
        from Extracto import sparsevec

        line = self.handle.readline()
        if not line:
            raise StopIteration
        x = line.rstrip().split("\t", 4)
        index, gene, drug, has_relationship, sparse_str = x
        index, has_relationship = int(index), int(has_relationship)

        sparse_vector = sparsevec.load_str(sparse_str)
        vector = sparsevec.to_array(sparse_vector)
        return gene, drug, has_relationship, vector, index
    def __iter__(self):
        return self

class _observations_writer:
    def __init__(self, filename):
        self.handle = filefns.LockedHandle(filefns.openfh(filename, 'w'))
    def close(self):
        self.handle.close()
    def flush(self):
        self.handle.flush()
    def write(self, gene, drug, has_relationship, vector, index):
        from Extracto import sparsevec
        sparse_vector = sparsevec.from_array(vector)
        sparse_str = sparsevec.save_str(sparse_vector)
        self.handle.write("%d\t%s\t%s\t%d\t%s\n" % (
            index, gene, drug, has_relationship, sparse_str))

open_observations = _HandleOpener(_observations_reader, _observations_writer)


class _results_reader:
    # Return gene, drug, has_relationship, prediction, p0, p1, is_correct
    def __init__(self, filename):
        self.handle = filefns.openfh(filename)
    def next(self):
        line = self.handle.readline()
        if not line:
            raise StopIteration
        x = line.rstrip().split("|")
        gene, drug, has_relationship, prediction, p0, p1, is_correct = x
        has_relationship, prediction, is_correct = map(
            int, (has_relationship, prediction, is_correct))
        p0, p1 = float(p0), float(p1)
        return gene, drug, has_relationship, prediction, p0, p1, is_correct
    def __iter__(self):
        return self

class _results_writer:
    def __init__(self, filename):
        self.handle = filefns.LockedHandle(filefns.openfh(filename, 'w'))
    def close(self):
        self.handle.close()
    def flush(self):
        self.handle.flush()
    def write(self, gene, drug, has_relationship, 
              prediction, p0, p1, is_correct):
        self.handle.write("%s|%s|%d|%d|%g|%g|%d\n" % (
            gene, drug, has_relationship, prediction, p0, p1, is_correct))

open_results = _HandleOpener(_results_reader, _results_writer)



class ScoreData:
    """
    
    Members:
    pairs_tested        Number of gene/drug pairs tested.
    num_relationships   Number of relationships.
    max_fscore          Maximum f-score.
    data            List of (gene, drug, has_relationship, score, rec, prec).

    """
    def __init__(self):
        self.pairs_tested = None
        self.num_relationships = None
        self.max_fscore = None
        self.data = []

class _score_reader:
    # Return ScoreData
    def __init__(self, filename):
        self.handle = filefns.openfh(filename)
        self.read = 0
    def next(self):
        from Extracto import tokenfns
        
        if self.read:
            raise StopIteration
        score_data = ScoreData()
        
        line = self.handle.readline()
        assert line.startswith("Tested")
        score_data.pairs_tested = int(line.split()[1])

        line = self.handle.readline()
        assert line.startswith("Total relationships")
        score_data.num_relationships = int(line.rstrip().split()[-1])
        
        line = self.handle.readline()
        assert line.startswith("Maximum F-score")
        score_data.max_fscore = float(line.rstrip().split()[-1])

        line = self.handle.readline()
        columns = line.strip().split()
        offsets = tokenfns.find_offsets(columns, line)
        
        assert line.lstrip().startswith("REL")
        for line in self.handle:
            cols = []
            for i in range(len(offsets)-1):
                cols.append(line[offsets[i]:offsets[i+1]])
            cols.append(line[offsets[-1]:])
            cols = [x.strip() for x in cols]
            assert len(cols) == 6
            has_relationship, score, rec, prec, gene, drug = cols
            has_relationship = int(has_relationship)
            score, rec, prec = float(score), float(rec), float(prec)
            x = gene, drug, has_relationship, score, rec, prec
            score_data.data.append(x)

        return score_data
            
    def __iter__(self):
        return self

class _score_writer:
    def __init__(self, filename):
        self.handle = filefns.LockedHandle(filefns.openfh(filename, 'w'))
    def close(self):
        self.handle.close()
    def flush(self):
        self.handle.flush()
    def write(self, score_data):
        from Extracto.genename.fileformats import _align_columns
        self.handle.write("Tested %d pairs\n" % score_data.pairs_tested)
        self.handle.write("Total relationships: %d\n" %
                          score_data.num_relationships)
        self.handle.write("Maximum F-score: %g\n" % score_data.max_fscore)

        columns = []
        headings =["REL", "SCORE", "REC", "PREC", "GENE\tDRUG"]
        formats = ["%d", "%.4f", "%.3f", "%.3f", "%s"]
        data_index = [2, 3, 4, 5, 0]  # index into ScoreData.data
        for i in range(len(headings)):
            index = data_index[i]
            if headings[i] == "GENE\tDRUG":
                col = [headings[i]] + \
                      ["%s\t%s" % (x[index], x[index+1])
                       for x in score_data.data]
            else:
                col = [headings[i]] + \
                      [formats[i] % x[index] for x in score_data.data]
            columns.append((col, "left"))
        _align_columns(self.handle, " ", *columns)

open_score = _HandleOpener(_score_reader, _score_writer)


class _sentences_reader:
    # Return pmid, sentnum, gene, gene start, gene end, drug, drug start, drug end, sentence, has_relationship, sentence_shows_relationship
    def __init__(self, filename):
        self.handle = filefns.openfh(filename)
    def next(self):
        line = self.handle.readline()
        if not line:
            raise StopIteration
        x = line.rstrip().split("\t", 10)
        pmid, snum, gene, gs, ge, drug, ds, de, \
              sentence, has_rel, shows_rel = x
        snum, gs, ge, ds, de, has_rel, shows_rel = map(
            int, (snum, gs, ge, ds, de, has_rel, shows_rel))
        return pmid, snum, gene, gs, ge, drug, ds, de, \
               sentence, has_rel, shows_rel
    def __iter__(self):
        return self

class _sentences_writer:
    def __init__(self, filename):
        self.handle = filefns.LockedHandle(filefns.openfh(filename, 'w'))
    def close(self):
        self.handle.close()
    def flush(self):
        self.handle.flush()
    def write(self, pmid, sentnum, gene, gene_start, gene_end,
              drug, drug_start, drug_end,
              sentence, has_relationship, sentence_shows_relationship):
        cols = pmid, sentnum, gene, gene_start, gene_end, \
               drug, drug_start, drug_end, \
               sentence, has_relationship, sentence_shows_relationship
        cols = map(str, cols)
        self.handle.write("%s\n" % "\t".join(cols))

open_sentences = _HandleOpener(_sentences_reader, _sentences_writer)


class _vocabulary_reader:
    # Return word, index
    def __init__(self, filename):
        self.handle = filefns.openfh(filename)
    def next(self):
        line = self.handle.readline()
        if not line:
            raise StopIteration
        x = line.rstrip().split("|")
        word, index = x
        index = int(index)
        return word, index
    def __iter__(self):
        return self

class _vocabulary_writer:
    def __init__(self, filename):
        self.handle = filefns.LockedHandle(filefns.openfh(filename, 'w'))
    def close(self):
        self.handle.close()
    def flush(self):
        self.handle.flush()
    def write(self, word, index):
        self.handle.write("%s|%d\n" % (word, index))

open_vocabulary = _HandleOpener(_vocabulary_reader, _vocabulary_writer)


class _sentence_observations_reader:
    # Return pmid, sentnum, gene_start, drug_start, vector, index
    def __init__(self, filename):
        self.handle = filefns.openfh(filename)
    def next(self):
        from Extracto import sparsevec

        line = self.handle.readline()
        if not line:
            raise StopIteration
        x = line.rstrip().split("\t", 9)
        index, pmid, sentnum, gs, ds, sparse_str = x
        index, sentnum, gs, ds = map(int, (index, sentnum, gs, ds))

        sparse_vector = sparsevec.load_str(sparse_str)
        vector = sparsevec.to_array(sparse_vector)
        return pmid, sentnum, gs, ds, vector, index
    def __iter__(self):
        return self

class _sentence_observations_writer:
    def __init__(self, filename):
        self.handle = filefns.LockedHandle(filefns.openfh(filename, 'w'))
    def close(self):
        self.handle.close()
    def flush(self):
        self.handle.flush()
    def write(self, pmid, sentnum, gene_start, drug_start, vector, index):
        from Extracto import sparsevec
        sparse_vector = sparsevec.from_array(vector)
        sparse_str = sparsevec.save_str(sparse_vector)
        self.handle.write("%d\t%s\t%d\t%d\t%d\t%s\n" % (
            index, pmid, sentnum, gene_start, drug_start, sparse_str))

open_sentence_observations = _HandleOpener(
    _sentence_observations_reader, _sentence_observations_writer)


class _wordcounts_reader:
    # Return word, count
    def __init__(self, filename):
        self.handle = filefns.openfh(filename)
    def next(self):
        line = self.handle.readline()
        if not line:
            raise StopIteration
        x = line.rstrip().split("\t")
        word, count = x
        count = int(count)
        return word, count
    def __iter__(self):
        return self

class _wordcounts_writer:
    def __init__(self, filename):
        self.handle = filefns.LockedHandle(filefns.openfh(filename, 'w'))
    def close(self):
        self.handle.close()
    def flush(self):
        self.handle.flush()
    def write(self, word, count):
        self.handle.write("%s\t%d\n" % (word, count))

open_wordcounts = _HandleOpener(_wordcounts_reader, _wordcounts_writer)