##############################################################################
# Institute for the Design of Advanced Energy Systems Process Systems
# Engineering Framework (IDAES PSE Framework) Copyright (c) 2018-2020, by the
# software owners: The Regents of the University of California, through
# Lawrence Berkeley National Laboratory,  National Technology & Engineering
# Solutions of Sandia, LLC, Carnegie Mellon University, West Virginia
# University Research Corporation, et al. All rights reserved.
#
# Please see the files COPYRIGHT.txt and LICENSE.txt for full copyright and
# license information, respectively. Both files are also available online
# at the URL "https://github.com/IDAES/idaes-pse".
##############################################################################
"""
Tests for dynamic utility methods.
"""

import pytest
from pyomo.environ import (ConcreteModel, Block, Constraint, Var, Set,
        TransformationFactory)
from pyomo.dae import ContinuousSet, DerivativeVar
from pyomo.dae.set_utils import (is_explicitly_indexed_by,
        is_in_block_indexed_by, get_index_set_except)
import idaes.logger as idaeslog
from idaes.core.util.dyn_utils import *

__author__ = "Robert Parker"


@pytest.mark.unit
def test_fix_and_deactivate():
    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)

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

        @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])

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

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


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

    for t in m.time:
        m.fs.b1.v[t, m.space.first()].fix()

    active_dict = get_activity_dict(m.fs)
    for comp in m.fs.component_data_objects(Constraint, Block):
        assert active_dict[id(comp)] == True

    deactivate_model_at(m, m.time, m.time[2])
    assert m.fs.con1[m.time[1]].active
    assert not m.fs.con1[m.time[2]].active
    assert m.fs.con2[m.space[1]].active
    assert not m.fs.b1.con[m.time[2], m.space[1]].active
    assert not m.fs.b2[m.time[2], m.space.last()].active
    assert 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]], 
            outlvl=idaeslog.ERROR)
    # Higher outlvl threshold as will encounter warning trying to deactivate
    # disc equations at time.first()
    assert not m.fs.con1[m.time[1]].active
    assert not m.fs.con1[m.time[3]].active
    assert not m.fs.b1.con[m.time[1], m.space[1]].active
    assert not m.fs.b1.con[m.time[3], m.space[1]].active

    init_derivs = get_derivatives_at(m, m.time, m.time.first())[m.time.first()]
    init_deriv_names = [var.name for var in init_derivs]

    assert m.fs.b1.dv[m.time.first(), m.space.first()] in init_derivs
    assert m.fs.b1.dv[m.time[1], m.space[1]].name in init_deriv_names

    deriv_dict = get_derivatives_at(m, m.time, [m.time.first(), m.time.last()])
    deriv_name_dict = {t: [d.name for d in deriv_dict[t]] 
                          for t in deriv_dict.keys()}
    assert m.time.first() in deriv_name_dict.keys()
    assert m.time.last() in deriv_name_dict.keys()
    assert (m.fs.b1.dv[m.time.last(), m.space[1]].name 
            in deriv_name_dict[m.time.last()])
    assert (m.fs.b1.dv[m.time.last(), m.space[1]].name 
            not in deriv_name_dict[m.time.first()])

    vars_unindexed = fix_vars_unindexed_by(m, m.time)
    cons_unindexed = deactivate_constraints_unindexed_by(m, m.time)
    assert m.fs.v0[m.space[1]] in vars_unindexed
    assert m.fs.b1.b2[m.time[1]].v not in vars_unindexed
    assert m.fs.con2[m.space[2]] in cons_unindexed
    assert m.fs.con1[m.time[1]] not in cons_unindexed
    assert not m.fs.con2[m.space[1]].active
    assert m.fs.v0[m.space[1]].fixed

    path = path_from_block(m.fs.b2[m.time[1], m.space[1]].b3['a'].v,
                           m, include_comp=False)
    assert path == [('fs', None), ('b2', (m.time[1], m.space[1])),
                    ('b3', 'a')]
    path = path_from_block(m.fs.b2[m.time[1], m.space[1]].b3['a'].v,
                           m, include_comp=True)
    assert path == [('fs', None), ('b2', (m.time[1], m.space[1])),
                    ('b3', 'a'), ('v', None)]
    path = path_from_block(m.fs.b2[m.time[1], m.space[1]].b3['a'].v['f'],
                           m, include_comp=True)
    assert path == [('fs', None), ('b2', (m.time[1], m.space[1])),
                    ('b3', 'a'), ('v', 'f')]
    path = path_from_block(m.fs.b2[m.time[1], m.space[1]].b3['a'].v['f'],
                           m.fs.b2[m.time[1], m.space[1]], include_comp=True)
    assert path == [('b3', 'a'), ('v', 'f')]
    path = path_from_block(m.fs.b1.con[m.time[1], m.space[1]], m.fs)
    assert path == [('b1', None)]

    m.fs.b1.b2[m.time[1]].v.set_value(-1)
    for x in m.space:
        m.fs.b1.v[m.time[1], x].set_value(-1)
        m.fs.b1.dv[m.time[1], x].set_value(-1)
        for c1 in m.set1:
            m.fs.b2[m.time[1], x].v[c1].set_value(-1)
            for c2 in m.set2:
                m.fs.b2[m.time[1], x].b3[c1].v[c2].set_value(-1)

    copy_values_at_time(m, m, m.time.last(), m.time[1], copy_fixed=False)
    assert m.fs.b1.b2[m.time.last()].v.value == -1
    for x in m.space:
        if x != m.space.first():
            assert m.fs.b1.v[m.time.last(), x].value == -1
        else:
            assert m.fs.b1.v[m.time.last(), x].value == 1
            assert m.fs.b1.v[m.time.last(), x].fixed
        assert m.fs.b1.dv[m.time.last(), x].value == -1
        for c1 in m.set1:
            assert m.fs.b2[m.time.last(), x].v[c1].value == -1
            for c2 in m.set2:
                assert m.fs.b2[m.time[1], x].b3[c1].v[c2].value == -1

@pytest.mark.unit
def test_copy_non_time_indexed_values():
    m1 = ConcreteModel()
    m1.time = Set(initialize=[1,2,3,4,5])
    m1.v1 = Var(m1.time, initialize=1)
    m1.v2 = Var(initialize=1)
    @m1.Block(['a','b'])
    def b1(b, i):
        b.v3 = Var(initialize=1)

        @b.Block(m1.time)
        def b2(b, t):
            b.v4 = Var(initialize=1)

        @b.Block()
        def b4(b):
            b.v6 = Var(initialize=1)
            
    @m1.Block(m1.time)
    def b3(b, t):
        b.v5 = Var(initialize=1)

    m2 = ConcreteModel()
    m2.time = Set(initialize=[1,2,3,4,5])
    m2.v1 = Var(m2.time, initialize=2)
    m2.v2 = Var(initialize=2)
    @m2.Block(['a','b'])
    def b1(b):
        b.v3 = Var(initialize=2)

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

        @b.Block()
        def b4(b):
            b.v6 = Var(initialize=2)

    @m2.Block(m1.time)
    def b3(b, t):
        b.v5 = Var(initialize=2)

    copy_non_time_indexed_values(m1, m2)
    assert m1.v1[1].value != m2.v1[1].value
    assert m1.v2.value == m2.v2.value == 2
    assert m1.b1['a'].v3.value == m2.b1['a'].v3.value == 2
    assert m1.b1['b'].b4.v6.value == m2.b1['b'].b4.v6.value == 2
    assert m1.b3[3].v5.value != m2.b3[3].v5.value


@pytest.mark.unit
def test_find_comp_in_block():
    m1 = ConcreteModel()

    @m1.Block([1,2,3])
    def b1(b):
        b.v = Var([1,2,3])

    m2 = ConcreteModel()

    @m2.Block([1,2,3])
    def b1(b):
        b.v = Var([1,2,3,4])

    @m2.Block([1,2,3])
    def b2(b):
        b.v = Var([1,2,3])

    v1 = m1.b1[1].v[1]

    assert find_comp_in_block(m2, m1, v1) is m2.b1[1].v[1]

    v2 = m2.b2[1].v[1]
    v3 = m2.b1[3].v[4]

    # These should result in Attribute/KeyErrors

    with pytest.raises(AttributeError, match=r'.*has no attribute.*'):
        find_comp_in_block(m1, m2, v2)
    with pytest.raises(KeyError, match=r'.*is not a valid index.*'):
        find_comp_in_block(m1, m2, v3)
    assert find_comp_in_block(m1, m2, v2, allow_miss=True) is None
    assert find_comp_in_block(m1, m2, v3, allow_miss=True) is None


@pytest.mark.unit
def test_find_comp_in_block_at_time():
    m1 = ConcreteModel()
    m1.time = Set(initialize=[1,2,3])

    @m1.Block(m1.time)
    def b1(b):
        b.v = Var(m1.time)

    @m1.Block([1,2,3])
    def b(bl):
        bl.v = Var(m1.time)
        bl.v2 = Var(m1.time, ['a','b','c'])

    m2 = ConcreteModel()
    m2.time = Set(initialize=[1,2,3,4,5,6])

    @m2.Block(m2.time)
    def b1(b):
        b.v = Var(m2.time)

    @m2.Block([1,2,3,4])
    def b(bl):
        bl.v = Var(m2.time)
        bl.v2 = Var(m2.time, ['a','b','c'])

    @m2.Block([1,2,3])
    def b2(b):
        b.v = Var(m2.time)

    v1 = m1.b1[1].v[1]
    v3 = m1.b[3].v[1]

    assert find_comp_in_block_at_time(m2, m1, v1, m2.time, 4) is m2.b1[4].v[4]
    assert find_comp_in_block_at_time(m2, m1, v3, m2.time, 5) is m2.b[3].v[5]

    v = m1.b[1].v2[1, 'a']
    assert find_comp_in_block_at_time(m2, m1, v, m2.time, 6) is m2.b[1].v2[6, 'a']

    v2 = m2.b2[1].v[1]
    v4 = m2.b[4].v[1]

    # Should result in exceptions:
    with pytest.raises(AttributeError, match=r'.*has no attribute.*'):
        find_comp_in_block_at_time(m1, m2, v2, m1.time, 3)
    with pytest.raises(KeyError, match=r'.*is not a valid index.*'):
        find_comp_in_block_at_time(m1, m2, v4, m1.time, 3)

    assert find_comp_in_block_at_time(m1, m2, v2, m1.time, 3, allow_miss=True) is None
    assert find_comp_in_block_at_time(m1, m2, v4, m1.time, 3, allow_miss=True) is None


@pytest.mark.unit
def test_get_location_of_coordinate_set():
    m = ConcreteModel()
    m.s1 = Set(initialize=[1,2,3])
    m.s2 = Set(initialize=[('a',1), ('b',2)])
    m.s3 = Set(initialize=[('a',1,0), ('b',2,1)])
    m.v1 = Var(m.s1)
    m.v2 = Var(m.s1, m.s2)
    m.v121 = Var(m.s1, m.s2, m.s1)
    m.v3 = Var(m.s3, m.s1, m.s2)

    assert get_location_of_coordinate_set(m.v1.index_set(), m.s1) == 0
    assert get_location_of_coordinate_set(m.v2.index_set(), m.s1) == 0
    assert get_location_of_coordinate_set(m.v3.index_set(), m.s1) == 3

    # These should each raise a ValueError
    with pytest.raises(ValueError):
        get_location_of_coordinate_set(m.v1.index_set(), m.s2)
    with pytest.raises(ValueError):
        get_location_of_coordinate_set(m.v121.index_set(), m.s1)


@pytest.mark.unit
def test_get_index_of_set():
    m = ConcreteModel()
    m.s1 = Set(initialize=[1,2,3])
    m.s2 = Set(initialize=[('a',1), ('b',2)])
    m.s3 = Set(initialize=[('a',1,0), ('b',2,1)])
    m.v0 = Var()
    m.v1 = Var(m.s1)
    m.v2 = Var(m.s1, m.s2)
    m.v121 = Var(m.s1, m.s2, m.s1)
    m.v3 = Var(m.s3, m.s1, m.s2)

    assert get_index_of_set(m.v1[2], m.s1) == 2
    assert get_index_of_set(m.v2[2,'a',1], m.s1) == 2
    assert get_index_of_set(m.v3['b',2,1,3,'b',2], m.s1) == 3

    with pytest.raises(ValueError) as exc_test:
        get_index_of_set(m.v0, m.s1)
    with pytest.raises(ValueError) as exc_test:
        get_index_of_set(m.v2[1,'a',1], m.s2)
    with pytest.raises(ValueError) as exc_test:
        get_index_of_set(m.v2[1,'b',2], m.s3)


@pytest.mark.unit
def test_get_implicit_index_of_set():
    m = ConcreteModel()
    m.s1 = Set(initialize=[1,2,3])
    m.s2 = Set(initialize=['a', 'b', 'c'])
    m.s3 = Set(initialize=[('d',4), ('e',5)])

    @m.Block()
    def b1(b):
        @b.Block(m.s3, m.s1)
        def b2(b):
            @b.Block()
            def b3(b):
                b.v1 = Var(m.s2)
                b.v2 = Var(m.s1)

    assert get_implicit_index_of_set(m.b1.b2['d',4,1].b3.v1['a'], m.s1) == 1
    assert get_implicit_index_of_set(m.b1.b2['d',4,1].b3.v1['a'], m.s2) == 'a'

    with pytest.raises(ValueError) as exc_test:
        get_implicit_index_of_set(m.b1.b2['e',5,2].b3.v2[1], m.s1)
