#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright 2017 National Technology and Engineering Solutions of Sandia, LLC
#  Under the terms of Contract DE-NA0003525 with National Technology and 
#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain 
#  rights in this software.
#  This software is distributed under the 3-clause BSD License.
#  ___________________________________________________________________________

"""
Unit Tests for pyomo.dae.set_utils
"""
import os
from os.path import abspath, dirname

from six import StringIO

import pyutilib.th as unittest

from pyomo.core.base import (Block, Constraint, ConcreteModel, Var, Set,
        TransformationFactory)
from pyomo.common.collections import ComponentMap
from pyomo.common.log import LoggingIntercept
from pyomo.dae import *
from pyomo.dae.set_utils import *

currdir = dirname(abspath(__file__)) + os.sep


def make_model():
    m = ConcreteModel()
    m.time = ContinuousSet(bounds=(0, 10))
    m.space = ContinuousSet(bounds=(0, 5))
    m.set1 = Set(initialize=['a', 'b', 'c'])
    m.set2 = Set(initialize=['d', 'e', 'f'])
    m.fs = Block()

    m.fs.v0 = Var(m.space, initialize=1)

    @m.fs.Block()
    def b1(b):
        b.v = Var(m.time, m.space, initialize=1)
        b.dv = DerivativeVar(b.v, wrt=m.time, initialize=0)

        b.con = Constraint(m.time, m.space,
                rule=lambda b, t, x: b.dv[t, x] == 7 - b.v[t, x])
        # Inconsistent

        @b.Block(m.time)
        def b2(b, t):
            b.v = Var(initialize=2)

    @m.fs.Block(m.time, m.space)
    def b2(b, t, x):
        b.v = Var(m.set1, initialize=2)

        @b.Block(m.set1)
        def b3(b, c):
            b.v = Var(m.set2, initialize=3)

            @b.Constraint(m.set2)
            def con(b, s):
                return (5*b.v[s] ==
                        m.fs.b2[m.time.first(), m.space.first()].v[c])
                # inconsistent

    @m.fs.Constraint(m.time)
    def con1(fs, t):
        return fs.b1.v[t, m.space.last()] == 5
    # Will be inconsistent

    @m.fs.Constraint(m.space)
    def con2(fs, x):
        return fs.b1.v[m.time.first(), x] == fs.v0[x]
    # will be consistent

    disc = TransformationFactory('dae.collocation')
    disc.apply_to(m, wrt=m.time, nfe=5, ncp=2, scheme='LAGRANGE-RADAU')
    disc.apply_to(m, wrt=m.space, nfe=5, ncp=2, scheme='LAGRANGE-RADAU')

    return m


class TestDaeSetUtils(unittest.TestCase):
    
    # Test explicit/implicit index detection functions
    def test_indexed_by(self):
        m = ConcreteModel()
        m.time = ContinuousSet(bounds=(0, 10))
        m.space = ContinuousSet(bounds=(0, 10))
        m.set = Set(initialize=['a', 'b', 'c'])
        m.set2 = Set(initialize=[('a', 1), ('b', 2)])
        m.v = Var()
        m.v1 = Var(m.time)
        m.v2 = Var(m.time, m.space)
        m.v3 = Var(m.set, m.space, m.time)
        m.v4 = Var(m.time, m.set2)
        m.v5 = Var(m.set2, m.time, m.space)

        @m.Block()
        def b(b):
            b.v = Var()
            b.v1 = Var(m.time)
            b.v2 = Var(m.time, m.space)
            b.v3 = Var(m.set, m.space, m.time)

        @m.Block(m.time)
        def b1(b):
            b.v = Var()
            b.v1 = Var(m.space)
            b.v2 = Var(m.space, m.set)

        @m.Block(m.time, m.space)
        def b2(b):
            b.v = Var()
            b.v1 = Var(m.set)

            @b.Block()
            def b(bl):
                bl.v = Var()
                bl.v1 = Var(m.set)
                bl.v2 = Var(m.time)

        @m.Block(m.set2, m.time)
        def b3(b):
            b.v = Var()
            b.v1 = Var(m.space)

            @b.Block(m.space)
            def b(bb):
                bb.v = Var(m.set)

        disc = TransformationFactory('dae.collocation')
        disc.apply_to(m, wrt=m.time, nfe=5, ncp=2, scheme='LAGRANGE-RADAU')
        disc.apply_to(m, wrt=m.space, nfe=5, ncp=2, scheme='LAGRANGE-RADAU')

        self.assertFalse(is_explicitly_indexed_by(m.v, m.time))
        self.assertTrue(is_explicitly_indexed_by(m.b.v2, m.space))
        self.assertTrue(is_explicitly_indexed_by(m.b.v3, m.time, m.space))

        self.assertFalse(is_in_block_indexed_by(m.v1, m.time))
        self.assertFalse(is_in_block_indexed_by(m.v2, m.set))
        self.assertTrue(is_in_block_indexed_by(m.b1[m.time[1]].v2, m.time))

        self.assertTrue(is_in_block_indexed_by(
            m.b2[m.time[1], m.space[1]].b.v1, m.time))
        self.assertTrue(is_in_block_indexed_by(
            m.b2[m.time[1], m.space[1]].b.v2, m.time))
        self.assertTrue(is_explicitly_indexed_by(
            m.b2[m.time[1], m.space[1]].b.v2, m.time))
        self.assertFalse(is_in_block_indexed_by(
            m.b2[m.time[1], m.space[1]].b.v1, m.set))

        self.assertFalse(is_in_block_indexed_by(
            m.b2[m.time[1], m.space[1]].b.v1, 
            m.space, stop_at=m.b2[m.time[1], m.space[1]]))

        # Explicit indexing with multi-dimensional set:
        self.assertTrue(is_explicitly_indexed_by(m.v4, m.time, m.set2))
        self.assertTrue(is_explicitly_indexed_by(m.v5, m.time, m.set2, m.space))

        # Implicit indexing with multi-dimensional set:
        self.assertTrue(is_in_block_indexed_by(
            m.b3['a', 1, m.time[1]].v, m.set2))
        self.assertTrue(is_in_block_indexed_by(
            m.b3['a', 1, m.time[1]].v, m.time))
        self.assertTrue(is_in_block_indexed_by(
            m.b3['a', 1, m.time[1]].v1[m.space[1]], m.set2))
        self.assertFalse(is_in_block_indexed_by(
            m.b3['a', 1, m.time[1]].v1[m.space[1]], m.space))
        self.assertTrue(is_in_block_indexed_by(
            m.b3['b', 2, m.time[2]].b[m.space[2]].v['b'], m.set2))
        self.assertTrue(is_in_block_indexed_by(
            m.b3['b', 2, m.time[2]].b[m.space[2]].v['b'], m.time))
        self.assertTrue(is_in_block_indexed_by(
            m.b3['b', 2, m.time[2]].b[m.space[2]].v['b'], m.space))
        self.assertFalse(is_in_block_indexed_by(
            m.b3['b', 2, m.time[2]].b[m.space[2]].v['b'], m.set))
        self.assertFalse(is_in_block_indexed_by(
            m.b3['b', 2, m.time[2]].b[m.space[2]].v['b'], m.time,
            stop_at=m.b3['b', 2, m.time[2]]))
        self.assertFalse(is_in_block_indexed_by(
            m.b3['b', 2, m.time[2]].b[m.space[2]].v['b'], m.time,
            stop_at=m.b3))


    # Test get_index_set_except and _complete_index
    def test_get_index_set_except(self):
        '''
        Tests:
          For components indexed by 0, 1, 2, 3, 4 sets:
            get_index_set_except one, then two (if any) of those sets
            check two items that should be in set_except
            insert item(s) back into these sets via index_getter
        '''
        m = ConcreteModel()
        m.time = ContinuousSet(bounds=(0, 10))
        m.space = ContinuousSet(bounds=(0, 10))
        m.set1 = Set(initialize=['a', 'b', 'c'])
        m.set2 = Set(initialize=['d', 'e', 'f'])
        m.v = Var()
        m.v1 = Var(m.time)
        m.v2 = Var(m.time, m.space)
        m.v3 = Var(m.time, m.space, m.set1)
        m.v4 = Var(m.time, m.space, m.set1, m.set2)

        # Multi-dimensional set:
        m.set3 = Set(initialize=[('a', 1), ('b', 2)])
        m.v5 = Var(m.set3)
        m.v6 = Var(m.time, m.space, m.set3)
        m.v7 = Var(m.set3, m.space, m.time)

        disc = TransformationFactory('dae.collocation')
        disc.apply_to(m, wrt=m.time, nfe=5, ncp=2, scheme='LAGRANGE-RADAU')
        disc.apply_to(m, wrt=m.space, nfe=5, ncp=2, scheme='LAGRANGE-RADAU')

        # Want this to give a TypeError
        # info = get_index_set_except(m.v, m.time)

        # Indexed by one set
        info = get_index_set_except(m.v1, m.time)
        set_except = info['set_except']
        index_getter = info['index_getter']
        self.assertTrue(set_except == [None])
        # Variable is not indexed by anything except time
        # Test that index_getter returns only the new value given,
        # regardless of whether it was part of the set excluded (time):
        self.assertEqual(index_getter((), -1), -1)

        # Indexed by two sets
        info = get_index_set_except(m.v2, m.time)
        set_except = info['set_except']
        index_getter = info['index_getter']
        self.assertTrue(m.space[1] in set_except
                        and m.space.last() in set_except)
	# Here (2,) is the partial index, corresponding to space.
        # Can be provided as a scalar or tuple. 4, the time index,
        # should be inserted before (2,)
        self.assertEqual(index_getter((2,), 4), (4, 2))
        self.assertEqual(index_getter(2, 4), (4, 2))

        # Case where every set is "omitted," now for multiple sets
        info = get_index_set_except(m.v2, m.space, m.time)
        set_except = info['set_except']
        index_getter = info['index_getter']
        self.assertTrue(set_except == [None])
        # 5, 7 are the desired index values for space, time 
        # index_getter should put them in the right order for m.v2,
        # even if they are not valid indices for m.v2
        self.assertEqual(index_getter((), 5, 7), (7, 5))

        # Indexed by three sets
        info = get_index_set_except(m.v3, m.time)
        # In this case set_except is a product of the two non-time sets
        # indexing v3
        set_except = info['set_except']
        index_getter = info['index_getter']
        self.assertTrue((m.space[1], 'b') in set_except
                        and (m.space.last(), 'a') in set_except)
        # index_getter inserts a scalar index into an index of length 2
        self.assertEqual(index_getter((2, 'b'), 7), (7, 2, 'b'))

        info = get_index_set_except(m.v3, m.space, m.time)
        # Two sets omitted. Now set_except is just set1
        set_except = info['set_except']
        index_getter = info['index_getter']
        self.assertTrue('a' in set_except)
        # index_getter inserts the two new indices in the right order
        self.assertEqual(index_getter('b', 1.2, 1.1), (1.1, 1.2, 'b'))

        # Indexed by four sets
        info = get_index_set_except(m.v4, m.set1, m.space)
        # set_except is a product, and there are two indices to insert
        set_except = info['set_except']
        index_getter = info['index_getter']
        self.assertTrue((m.time[1], 'd') in set_except)
        self.assertEqual(index_getter((4, 'f'), 'b', 8), (4, 8, 'b', 'f'))
        
        # The intended usage of this function looks something like:
        index_set = m.v4.index_set()
        for partial_index in set_except:
            complete_index = index_getter(partial_index, 'a', m.space[2])
            self.assertTrue(complete_index in index_set)
            # Do something for every index of v4 at 'a' and space[2]

        # Indexed by a multi-dimensional set
        info = get_index_set_except(m.v5, m.set3)
        set_except = info['set_except']
        index_getter = info['index_getter']
        self.assertEqual(set_except, [None])
        self.assertEqual(index_getter((), ('a', 1)), ('a', 1))

        info = get_index_set_except(m.v6, m.set3, m.time)
        set_except = info['set_except']
        index_getter = info['index_getter']
        self.assertTrue(m.space[1] in set_except)
        self.assertEqual(index_getter(m.space[1], ('b', 2), m.time[1]),
                (m.time[1], m.space[1], 'b', 2))

        info = get_index_set_except(m.v7, m.time)
        set_except = info['set_except']
        index_getter = info['index_getter']
        self.assertIn(('a', 1, m.space[1]), set_except)
        self.assertEqual(index_getter(('a', 1, m.space[1]), m.time[1]),
                         ('a', 1, m.space[1], m.time[1]))

        m.v8 = Var(m.time, m.set3, m.time)
        with self.assertRaises(ValueError):
            info = get_index_set_except(m.v8, m.time)
        with self.assertRaises(ValueError):
            info = get_index_set_except(m.v8, m.space)

    def test_deactivate_model_at(self):
        m = make_model()

        deactivate_model_at(m, m.time, m.time[2])
        self.assertTrue(m.fs.con1[m.time[1]].active)
        self.assertFalse(m.fs.con1[m.time[2]].active)
        self.assertTrue(m.fs.con2[m.space[1]].active)
        self.assertFalse(m.fs.b1.con[m.time[2], m.space[1]].active)
        self.assertFalse(m.fs.b2[m.time[2], m.space.last()].active)
        self.assertTrue(m.fs.b2[m.time[2], m.space.last()].b3['a'].con['e'].active)

        deactivate_model_at(m, m.time, [m.time[1], m.time[3]])
        # disc equations at time.first()
        self.assertFalse(m.fs.con1[m.time[1]].active)
        self.assertFalse(m.fs.con1[m.time[3]].active)
        self.assertFalse(m.fs.b1.con[m.time[1], m.space[1]].active)
        self.assertFalse(m.fs.b1.con[m.time[3], m.space[1]].active)

        with self.assertRaises(KeyError):
            deactivate_model_at(m, m.time, m.time[1], allow_skip=False,
                    suppress_warnings=True)


if __name__ == "__main__":
    unittest.main()