#  ___________________________________________________________________________
#
#  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.init_cond
"""
import os
from os.path import abspath, dirname

from six import StringIO

import pyutilib.th as unittest

from pyomo.core.base import *
from pyomo.environ import SolverFactory
from pyomo.common.collections import ComponentMap
from pyomo.common.log import LoggingIntercept
from pyomo.dae import *
from pyomo.dae.initialization import *

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

ipopt_available = SolverFactory('ipopt').available()


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

    @m.fs.Constraint(m.time, m.space)
    def con3(fs, t, x):
        if x == m.space.first():
            return Constraint.Skip
        return fs.b2[t, x].v['a'] == 7.

    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 TestDaeInitCond(unittest.TestCase):
    
    def test_get_inconsistent_initial_conditions(self):
        m = make_model()
        inconsistent = get_inconsistent_initial_conditions(m, m.time)

        self.assertIn(m.fs.b1.con[m.time[1], m.space[1]], inconsistent)
        self.assertIn(m.fs.b2[m.time[1], m.space[1]].b3['a'].con['d'],
                inconsistent)
        self.assertIn(m.fs.con1[m.time[1]], inconsistent)
        self.assertNotIn(m.fs.con2[m.space[1]], inconsistent)


    @unittest.skipIf(not ipopt_available, 'ipopt is not available')
    def test_solve_consistent_initial_conditions(self):
        m = make_model()
        solver = SolverFactory('ipopt')
        solve_consistent_initial_conditions(m, m.time, solver, allow_skip=True)
        inconsistent = get_inconsistent_initial_conditions(m, m.time)
        self.assertFalse(inconsistent)

        self.assertTrue(m.fs.con1[m.time[1]].active)
        self.assertTrue(m.fs.con1[m.time[3]].active)
        self.assertTrue(m.fs.b1.con[m.time[1], m.space[1]].active)
        self.assertTrue(m.fs.b1.con[m.time[3], m.space[1]].active)

        with self.assertRaises(KeyError):
            solve_consistent_initial_conditions(
                    m,
                    m.time,
                    solver,
                    allow_skip=False,
                    )


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