# ____________________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright (c) 2008-2026 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


import pyomo.common.unittest as unittest

from pyomo.environ import (
    SolverFactory,
    ConcreteModel,
    Set,
    Block,
    Var,
    Constraint,
    TransformationFactory,
)
from pyomo.dae import ContinuousSet, DerivativeVar
from pyomo.dae.initialization import (
    solve_consistent_initial_conditions,
    get_inconsistent_initial_conditions,
)

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.0

    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()
