"""
CBMPy: fluxmodules decomposiiton test module

This module only exists for testing purposes
=====================
PySCeS Constraint Based Modelling (http://cbmpy.sourceforge.net)

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>

Author: Arne C. Reimers
Contact email: arne.c.reimers@gmail.com
"""

# preparing for Python 3 port
from __future__ import division, print_function
from __future__ import absolute_import
#from __future__ import unicode_literals

import unittest
import os

from cbmpy import CBRead

from . import matroid
from .decomposition import Decomposition, Vertex
from .sparserationals import Matrix


class TestDecomposition(unittest.TestCase):

    def setUp(self):
        """ create test matroid """
        mat = Matrix([[-2, -1, 1, 0, 0, 0, 0, 0, 0],
                      [1, 0, 0, -1, 0, 0, 0, 0, 0],
                      [0, 1, -1, 1, -1, -1, 0, 0, 0],
                      [1, 0, 0, 0, 1, 0, -1, 0, 0],
                      [0, 0, 0, 0, 0, 1, 1, -1, 1],
                      [0, 0, 0, 0, 0, 0, 0, 1, -1]])

        elems = [1, 2, 3, 4, 5, 6, 7, 8, 9]
        self.matroid = matroid.fromMatrix(mat, elems)

        self.decomp = Decomposition(self.matroid)

    def tearDown(self):
        pass

    def testZeroModules(self):
        self.decomp.splitZeroModules(self.decomp.root)

        self.assertEqual(len(self.decomp.root.edges), 2)
        for e in self.decomp.root.edges.values():
            self.assertEqual(len(e), 0)

        for e in self.decomp.root.edges.keys():
            self.assertIsInstance(e, Vertex)
            self.assertEqual(len(e.edges), len(e.matroid.elems) + 1)
            if len(e.matroid.elems) == 2:
                self.assertSetEqual(set(e.matroid.elems), {8, 9})
            else:
                self.assertSetEqual(set(e.matroid.elems), {1, 2, 3, 4, 5, 6, 7})

    def testOneModules(self):
        for v in self.decomp.listNonLeaves():
            self.decomp.splitOneModules(v)

        for v in self.decomp.listNonLeaves():
            print(v)
            print(v.edges)
            print("")

        self.assertTrue(self.decomp.isParallelSeries())

        for v in self.decomp.listNonLeaves():
            for (w, interface) in v.edges.items():
                sep = w.getLeaves(v)
                print(sep)
                conn = self.matroid.connectivity(sep)
                self.assertLessEqual(conn, 1)
                self.assertEqual(len(interface), conn)


class TestDecomposition2(unittest.TestCase):
    def setUp(self):
        cDir = os.path.dirname(__file__)
        model_dir = os.path.join(cDir, '../models')
        model_file = "toy_model_maarleveld2014.l3.xml"

        self.cmod = CBRead.readSBML3FBC(model_file, model_dir)
        matrix = Matrix()
        labels, _ = matrix.addMetabolicNetwork(self.cmod)
        self.matroid = matroid.fromMatrix(matrix, labels)

        print(self.matroid.basis)
        print(self.matroid.nbasis)
        print(self.matroid.rep.toNumpy())

    def tearDown(self):
        pass

    def testPoolmanDirectRows(self):
        decomp = Decomposition(self.matroid)
        sim = self.matroid.similarityMatrix("rows")
#        print("rows:")
#        print(sim.matrix)
        # we expect coparallel elements to have similarity 1
        coparallel = self.matroid.coparallel()
        for group in coparallel:
            for e in group:
                for f in group:
                    if e != f:
                        ei = sim.elems.index(e)
                        fi = sim.elems.index(f)
                        self.assertAlmostEqual(sim[ei, fi], 1,
                                               msg="at %s(%d) %s(%d)" % (e, ei, f, fi))

        decomp.poolmanMethod(decomp.root, sim)
#        for v in decomp.listNonLeaves():
#            print(v)
#            print(v.edges)
#            print("")
        write(decomp.makeGraphViz(), "poolman_direct_rows.gv")

    def testPoolmanDirectCols(self):
        decomp = Decomposition(self.matroid)
        sim = self.matroid.similarityMatrix("cols")
#        print("cols:")
#        print(sim.matrix)
        # we expect parallel elements to have similarity 1
        parallel = self.matroid.parallel()
        for group in parallel:
            for e in group:
                for f in group:
                    if e != f:
                        ei = sim.elems.index(e)
                        fi = sim.elems.index(f)
                        self.assertAlmostEqual(sim[ei, fi], 1,
                                               msg="at %s(%d) %s(%d)" % (e, ei, f, fi))

        decomp.poolmanMethod(decomp.root, sim)
#        for v in decomp.listNonLeaves():
#            print(v)
#            print(v.edges)
#            print("")
        write(decomp.makeGraphViz(), "poolman_direct_cols.gv")

    def testPoolmanDirectDirect(self):
        decomp = Decomposition(self.matroid)
        sim = self.matroid.similarityMatrix("direct")
#        print("cols:")
#        print(sim.matrix)
        # we expect parallel elements to have similarity 1
        parallel = self.matroid.parallel()
        for group in parallel:
            for e in group:
                for f in group:
                    if e != f:
                        ei = sim.elems.index(e)
                        fi = sim.elems.index(f)
                        self.assertAlmostEqual(sim[ei, fi], 1,
                                               msg="at %s(%d) %s(%d)" % (e, ei, f, fi))

        decomp.poolmanMethod(decomp.root, sim)
#        for v in decomp.listNonLeaves():
#            print(v)
#            print(v.edges)
#            print("")
        write(decomp.makeGraphViz(), "poolman_direct_direct.gv")

    def testPoolmanDirectCombinded(self):
        decomp = Decomposition(self.matroid)
        sim = self.matroid.similarityMatrix("combined")
#        print("combined:")
#        print(sim.matrix)
        # we expect coparallel elements to have similarity 1
        coparallel = self.matroid.coparallel()
        for group in coparallel:
            for e in group:
                for f in group:
                    if e != f:
                        ei = sim.elems.index(e)
                        fi = sim.elems.index(f)
                        self.assertAlmostEqual(sim[ei, fi], 1,
                                               msg="at %s(%d) %s(%d)" % (e, ei, f, fi))
        # we expect parallel elements to have similarity 1
        parallel = self.matroid.parallel()
        for group in parallel:
            for e in group:
                for f in group:
                    if e != f:
                        ei = sim.elems.index(e)
                        fi = sim.elems.index(f)
                        self.assertAlmostEqual(sim[ei, fi], 1,
                                               msg="at %s(%d) %s(%d)" % (e, ei, f, fi))

        decomp.poolmanMethod(decomp.root, sim)
#        for v in decomp.listNonLeaves():
#            print(v)
#            print(v.edges)
#            print("")
        write(decomp.makeGraphViz(), "poolman_direct_combined.gv")

    def testOneModules(self):
        decomp = Decomposition(self.matroid)
        decomp.splitOneModules(decomp.root)
        write(decomp.makeGraphViz(), "one_modules.gv")

    def testPoolmanRows(self):
        decomp = Decomposition(self.matroid)
        decomp.splitOneModules(decomp.root)
        sim = decomp.root.matroid.similarityMatrix("rows")
        decomp.poolmanMethod(decomp.root, sim)
        write(decomp.makeGraphViz(), "poolman_rows.gv")


# @unittest.skip("takes quite long")
class TestDecompositionITM686(unittest.TestCase):
    def setUp(self):
        cDir = os.path.dirname(__file__)
        model_dir = os.path.join(cDir, '../models')
        model_file = "iTM686_sbml3.xml"

        self.cmod = CBRead.readSBML3FBC(model_file, model_dir)
        matrix = Matrix()
        labels, _ = matrix.addMetabolicNetwork(self.cmod)
        self.matroid = matroid.fromMatrix(matrix, labels)

    def tearDown(self):
        pass

    @unittest.skip("takes quite long")
    def testPoolmanRows(self):
        decomp = Decomposition(self.matroid)
        decomp.splitSimple()
        for v in decomp.listNonLeaves():
            if not v.isSimple():
                sim = v.matroid.similarityMatrix("rows")
                decomp.poolmanMethod(v, sim)

        write(decomp.makeGraphViz(), "iTM686_poolman_rows.gv")
        self.assertTrue(decomp.verifyFullyBranched())
        self.assertTrue(decomp.verifyEdgeWidth())
        print("width rows = %d" % decomp.getWidth())

    @unittest.skip("takes quite long")
    def testPoolmanCols(self):
        decomp = Decomposition(self.matroid)
        decomp.splitSimple()
        for v in decomp.listNonLeaves():
            if not v.isSimple():
                sim = v.matroid.similarityMatrix("cols")
                decomp.poolmanMethod(v, sim)

        write(decomp.makeGraphViz(), "iTM686_poolman_cols.gv")
        # self.assertTrue(decomp.verifyFullyBranched())
        # self.assertTrue(decomp.verifyEdgeWidth())
        print("width cols = %d" % decomp.getWidth())

    # @unittest.skip("takes quite long")
    def testPoolmanCombined(self):
        decomp = Decomposition(self.matroid)
        decomp.splitSimple()
        for v in decomp.listNonLeaves():
            if not v.isSimple():
                sim = v.matroid.similarityMatrix("combined")
                decomp.poolmanMethod(v, sim)

        write(decomp.makeGraphViz(), "iTM686_poolman_combined.gv")
        # self.assertTrue(decomp.verifyFullyBranched())
        # self.assertTrue(decomp.verifyEdgeWidth())
        print("width combined = %d" % decomp.getWidth())

    @unittest.skip("takes quite long")
    def testPoolmanDirect(self):
        decomp = Decomposition(self.matroid)
        decomp.splitSimple()
        sim = decomp.matroid.similarityMatrix("direct")
        for v in decomp.listNonLeaves():
            if not v.isSimple():
                decomp.poolmanMethod(v, sim)

        write(decomp.makeGraphViz(), "iTM686_poolman_direct.gv")
        # self.assertTrue(decomp.verifyFullyBranched())
        # self.assertTrue(decomp.verifyEdgeWidth())
        print("width direct = %d" % decomp.getWidth())


class TestDecompositionEcore(unittest.TestCase):
    def setUp(self):
        cDir = os.path.dirname(__file__)
        model_dir = os.path.join(cDir, '../models')
        model_file = "ecoli_core_COBRA.xml"

        self.cmod = CBRead.readCOBRASBML(model_file, model_dir)
        matrix = Matrix()
        labels, _ = matrix.addMetabolicNetwork(self.cmod)
        self.matroid = matroid.fromMatrix(matrix, labels)

    def tearDown(self):
        pass

    def testPoolmanRows(self):
        decomp = Decomposition(self.matroid)
        decomp.splitSimple()
        for v in decomp.listNonLeaves():
            if not v.isSimple():
                sim = v.matroid.similarityMatrix("rows")
                decomp.poolmanMethod(v, sim)

        write(decomp.makeGraphViz(cmod=self.cmod), "ecoli_core_poolman_rows.gv")
        self.assertTrue(decomp.verifyFullyBranched())
        self.assertTrue(decomp.verifyEdgeWidth())
        print("width rows = %d" % decomp.getWidth())

    def testPoolmanCols(self):
        decomp = Decomposition(self.matroid)
        decomp.splitSimple()
        for v in decomp.listNonLeaves():
            if not v.isSimple():
                sim = v.matroid.similarityMatrix("cols")
                decomp.poolmanMethod(v, sim)

        write(decomp.makeGraphViz(cmod=self.cmod), "ecoli_core_poolman_cols.gv")
        self.assertTrue(decomp.verifyFullyBranched())
        self.assertTrue(decomp.verifyEdgeWidth())
        print("width cols = %d" % decomp.getWidth())

    def testPoolmanCombined(self):
        decomp = Decomposition(self.matroid)
        decomp.splitSimple()
        for v in decomp.listNonLeaves():
            if not v.isSimple():
                sim = v.matroid.similarityMatrix("combined")
                decomp.poolmanMethod(v, sim)

        write(decomp.makeGraphViz(cmod=self.cmod), "ecoli_core_poolman_combined.gv")
        self.assertTrue(decomp.verifyFullyBranched())
        self.assertTrue(decomp.verifyEdgeWidth())
        print("width combined = %d" % decomp.getWidth())

    def testPoolmanDirect(self):
        decomp = Decomposition(self.matroid)
        decomp.splitSimple()
        sim = decomp.matroid.similarityMatrix("direct")
        for v in decomp.listNonLeaves():
            if not v.isSimple():
                decomp.poolmanMethod(v, sim)

        write(decomp.makeGraphViz(cmod=self.cmod), "ecoli_core_poolman_direct.gv")
        self.assertTrue(decomp.verifyFullyBranched())
        self.assertTrue(decomp.verifyEdgeWidth())
        print("width direct = %d" % decomp.getWidth())


def write(result, filename):
    cDir = os.path.dirname(__file__)
    result_dir = os.path.join(cDir, '../../results')
    result_file = os.path.join(result_dir, filename)
    fo = open(result_file, mode='w')
    fo.write(result)
    fo.close()


if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testName']
    unittest.main()
