#!/usr/bin/env python
#
# Copyright 2013 by Leighton Pritchard.  All rights reserved.
# This code is part of the Biopython distribution and governed by its
# license.  Please see the LICENSE file that should have been included
# as part of this package.

""" Tests for general functionality of the KGML parser and pathway model
"""

# Builtins
from __future__ import with_statement
import os
import unittest
import tempfile

# Biopython Bio.KEGG.KGML (?)
from Bio.KEGG.KGML.KGML_parser import read


class PathwayData(object):
    """ Convenience structure for testing pathway data
    """
    def __init__(self, infilename, outfilename, element_counts,
                 pathway_image, show_pathway_image=False):
        self.infilename = infilename
        self.outfilename = outfilename
        self.element_counts = element_counts
        self.pathway_image = pathway_image
        self.show_pathway_image = show_pathway_image


class KGMLPathwayTest(unittest.TestCase):
    """ Import the ko01100 metabolic map from a local .xml KGML file, and from
        the KEGG site, and write valid KGML output for each
    """
    def setUp(self):
        # Does our output director exist?  If not, create it
        if not os.path.isdir('KEGG'):
            os.mkdir('KEGG')
        # Define some data to work with as a list of tuples:
        # (infilename, outfilename, (entry_count, ortholog_count,
        # compound_count, map_counts), pathway_image,
        # show_image_map)
        self.data = [PathwayData(os.path.join("KEGG", "ko01100.xml"),
                                 tempfile.gettempprefix() + ".ko01100.kgml",
                                 (3628, 1726, 1746, 149),
                                 os.path.join("KEGG", "map01100.png")),
                     PathwayData(os.path.join("KEGG", "ko03070.xml"),
                                 tempfile.gettempprefix() + ".ko03070.kgml",
                                 (81, 72, 8, 1),
                                 os.path.join("KEGG", "map03070.png"),
                                 True)]
        # A list of KO IDs that we're going to use to modify pathway
        # appearance. These are KO IDs for reactions that take part in ko00020,
        # the TCA cycle
        self.ko_ids = \
            set(['ko:K00239', 'ko:K00240', 'ko:K00241', 'ko:K00242', 'ko:K00244',
                 'ko:K00245', 'ko:K00246', 'ko:K00247', 'ko:K00174', 'ko:K00175',
                 'ko:K00177', 'ko:K00176', 'ko:K00382', 'ko:K00164', 'ko:K00164',
                 'ko:K00658', 'ko:K01902', 'ko:K01903', 'ko:K01899', 'ko:K01900',
                 'ko:K01899', 'ko:K01900', 'ko:K00031', 'ko:K00030', 'ko:K00031',
                 'ko:K01648', 'ko:K00234', 'ko:K00235', 'ko:K00236', 'ko:K00237',
                 'ko:K01676', 'ko:K01677', 'ko:K01678', 'ko:K01679', 'ko:K01681',
                 'ko:K01682', 'ko:K01681', 'ko:K01682', 'ko:K01647', 'ko:K00025',
                 'ko:K00026', 'ko:K00024', 'ko:K01958', 'ko:K01959', 'ko:K01960',
                 'ko:K00163', 'ko:K00161', 'ko:K00162', 'ko:K00163', 'ko:K00161',
                 'ko:K00162', 'ko:K00382', 'ko:K00627', 'ko:K00169', 'ko:K00170',
                 'ko:K00172', 'ko:K00171', 'ko:K01643', 'ko:K01644', 'ko:K01646',
                 'ko:K01610', 'ko:K01596'])

    def tearDown(self):
        for p in self.data:
            if os.path.isfile(p.outfilename):
                os.remove(p.outfilename)

    def test_read_and_write_KGML_files(self):
        """ Read KGML from, and write KGML to, local files.
            Check we read/write the correct number of elements.
        """
        for p in self.data:
            # Test opening file
            with open(p.infilename, 'rU') as f:
                pathway = read(f)
                # Do we have the correct number of elements of each type
                self.assertEqual((len(pathway.entries),
                                  len(pathway.orthologs),
                                  len(pathway.compounds),
                                  len(pathway.maps)),
                                 p.element_counts)
            # Test writing file
            with open(p.outfilename, 'w') as f:
                f.write(pathway.get_KGML())
            # Can we read the file we wrote?
            with open(p.outfilename, 'rU') as f:
                pathway = read(f)
                # Do we have the correct number of elements of each type
                self.assertEqual((len(pathway.entries),
                                  len(pathway.orthologs),
                                  len(pathway.compounds),
                                  len(pathway.maps)),
                                 p.element_counts)


if __name__ == '__main__':
    runner = unittest.TextTestRunner(verbosity=2)
    unittest.main(testRunner=runner)