# -*- coding: utf-8 -*-
# ____________________________________________________________________________________
#
# 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.
# ____________________________________________________________________________________
#
#

import pyomo.common.unittest as unittest
from pyomo.environ import (
    ConcreteModel,
    Var,
    Param,
    Set,
    Constraint,
    Objective,
    Expression,
    Suffix,
    RangeSet,
    ExternalFunction,
    units,
    maximize,
    sin,
    cos,
    sqrt,
)
from pyomo.network import Port, Arc
from pyomo.dae import ContinuousSet, DerivativeVar
from pyomo.gdp import Disjunct, Disjunction
from pyomo.core.base.units_container import pint_available, UnitsError
from pyomo.util.check_units import (
    assert_units_consistent,
    assert_units_equivalent,
    check_units_equivalent,
    identify_inconsistent_units,
)


def python_callback_function(arg1, arg2):
    return 42.0


@unittest.skipIf(not pint_available, "Testing units requires pint")
class TestUnitsChecking(unittest.TestCase):
    def _create_model_and_vars(self):
        u = units
        m = ConcreteModel()
        m.dx = Var(units=u.m, initialize=0.10188943773836046)
        m.dy = Var(units=u.m, initialize=0.0)
        m.vx = Var(units=u.m / u.s, initialize=0.7071067769802851)
        m.vy = Var(units=u.m / u.s, initialize=0.7071067769802851)
        m.t = Var(units=u.s, bounds=(1e-5, 10.0), initialize=0.0024015570927624456)
        m.theta = Var(
            bounds=(0, 0.49 * 3.14), initialize=0.7853981693583533, units=u.radians
        )
        m.a = Param(initialize=-32.2, units=u.ft / u.s**2)
        m.x_unitless = Var()
        return m

    def test_assert_units_consistent_equivalent(self):
        u = units
        m = ConcreteModel()
        m.dx = Var(units=u.m, initialize=0.10188943773836046)
        m.dy = Var(units=u.m, initialize=0.0)
        m.vx = Var(units=u.m / u.s, initialize=0.7071067769802851)
        m.vy = Var(units=u.m / u.s, initialize=0.7071067769802851)
        m.t = Var(units=u.min, bounds=(1e-5, 10.0), initialize=0.0024015570927624456)
        m.theta = Var(
            bounds=(0, 0.49 * 3.14), initialize=0.7853981693583533, units=u.radians
        )
        m.a = Param(initialize=-32.2, units=u.ft / u.s**2)
        m.x_unitless = Var()

        m.obj = Objective(expr=m.dx, sense=maximize)
        m.vx_con = Constraint(expr=m.vx == 1.0 * u.m / u.s * cos(m.theta))
        m.vy_con = Constraint(expr=m.vy == 1.0 * u.m / u.s * sin(m.theta))
        m.dx_con = Constraint(expr=m.dx == m.vx * u.convert(m.t, to_units=u.s))
        m.dy_con = Constraint(
            expr=m.dy
            == m.vy * u.convert(m.t, to_units=u.s)
            + 0.5
            * (u.convert(m.a, to_units=u.m / u.s**2))
            * (u.convert(m.t, to_units=u.s)) ** 2
        )
        m.ground = Constraint(expr=m.dy == 0)
        m.unitless_con = Constraint(expr=m.x_unitless == 5.0)

        assert_units_consistent(m)  # check model
        assert_units_consistent(m.dx)  # check var - this should never fail
        assert_units_consistent(
            m.x_unitless
        )  # check unitless var - this should never fail
        assert_units_consistent(m.vx_con)  # check constraint
        assert_units_consistent(m.unitless_con)  # check unitless constraint

        assert_units_equivalent(m.dx, m.dy)  # check var
        assert_units_equivalent(m.x_unitless, u.dimensionless)  # check unitless var
        assert_units_equivalent(m.x_unitless, None)  # check unitless var
        assert_units_equivalent(m.vx_con.body, u.m / u.s)  # check constraint
        assert_units_equivalent(
            m.unitless_con.body, u.dimensionless
        )  # check unitless constraint
        assert_units_equivalent(m.dx, m.dy)  # check var
        assert_units_equivalent(m.x_unitless, u.dimensionless)  # check unitless var
        assert_units_equivalent(m.x_unitless, None)  # check unitless var
        assert_units_equivalent(m.vx_con.body, u.m / u.s)  # check constraint

        m.broken = Constraint(expr=m.dy == 42.0 * u.kg)
        with self.assertRaises(UnitsError):
            assert_units_consistent(m)
        assert_units_consistent(m.dx)
        assert_units_consistent(m.vx_con)
        with self.assertRaises(UnitsError):
            assert_units_consistent(m.broken)

        self.assertTrue(check_units_equivalent(m.dx, m.dy))
        self.assertFalse(check_units_equivalent(m.dx, m.vx))

    def test_assert_units_consistent_on_datas(self):
        u = units
        m = ConcreteModel()
        m.S = Set(initialize=[1, 2, 3])
        m.x = Var(m.S, units=u.m)
        m.t = Var(m.S, units=u.s)
        m.v = Var(m.S, units=u.m / u.s)
        m.unitless = Var(m.S)

        @m.Constraint(m.S)
        def vel_con(m, i):
            return m.v[i] == m.x[i] / m.t[i]

        @m.Constraint(m.S)
        def unitless_con(m, i):
            return m.unitless[i] == 42.0

        @m.Constraint(m.S)
        def sqrt_con(m, i):
            return sqrt(m.v[i]) == sqrt(m.x[i] / m.t[i])

        assert_units_consistent(m)  # check model
        assert_units_consistent(m.x)  # check var
        assert_units_consistent(m.t)  # check var
        assert_units_consistent(m.v)  # check var
        assert_units_consistent(m.unitless)  # check var
        assert_units_consistent(m.vel_con)  # check constraint
        assert_units_consistent(m.unitless_con)  # check unitless constraint

        assert_units_consistent(m.x[2])  # check var data
        assert_units_consistent(m.t[2])  # check var data
        assert_units_consistent(m.v[2])  # check var data
        assert_units_consistent(m.unitless[2])  # check var
        assert_units_consistent(m.vel_con[2])  # check constraint data
        assert_units_consistent(m.unitless_con[2])  # check unitless constraint data

        assert_units_equivalent(m.x[2], m.x[1])  # check var data
        assert_units_equivalent(m.t[2], u.s)  # check var data
        assert_units_equivalent(m.v[2], u.m / u.s)  # check var data
        assert_units_equivalent(
            m.unitless[2], u.dimensionless
        )  # check var data unitless
        assert_units_equivalent(m.unitless[2], None)  # check var
        assert_units_equivalent(m.vel_con[2].body, u.m / u.s)  # check constraint data
        assert_units_equivalent(
            m.unitless_con[2].body, u.dimensionless
        )  # check unitless constraint data

        @m.Constraint(m.S)
        def broken(m, i):
            return m.x[i] == 42.0 * m.v[i]

        with self.assertRaises(UnitsError):
            assert_units_consistent(m)
        with self.assertRaises(UnitsError):
            assert_units_consistent(m.broken)
        with self.assertRaises(UnitsError):
            assert_units_consistent(m.broken[1])

        # all of these should still work
        assert_units_consistent(m.x)  # check var
        assert_units_consistent(m.t)  # check var
        assert_units_consistent(m.v)  # check var
        assert_units_consistent(m.unitless)  # check var
        assert_units_consistent(m.vel_con)  # check constraint
        assert_units_consistent(m.unitless_con)  # check unitless constraint

        assert_units_consistent(m.x[2])  # check var data
        assert_units_consistent(m.t[2])  # check var data
        assert_units_consistent(m.v[2])  # check var data
        assert_units_consistent(m.unitless[2])  # check var
        assert_units_consistent(m.vel_con[2])  # check constraint data
        assert_units_consistent(m.unitless_con[2])  # check unitless constraint data

    def test_assert_units_consistent_all_components(self):
        # test all scalar components consistent
        u = units
        m = self._create_model_and_vars()
        m.obj = Objective(expr=m.dx / m.t - m.vx)
        m.con = Constraint(expr=m.dx / m.t == m.vx)
        # vars already added
        m.exp = Expression(expr=m.dx / m.t - m.vx)
        m.suff = Suffix(direction=Suffix.LOCAL)
        # params already added
        # sets already added
        m.rs = RangeSet(5)
        m.disj1 = Disjunct()
        m.disj1.constraint = Constraint(expr=m.dx / m.t <= m.vx)
        m.disj2 = Disjunct()
        m.disj2.constraint = Constraint(expr=m.dx / m.t <= m.vx)
        m.disjn = Disjunction(expr=[m.disj1, m.disj2])
        # block tested as part of model
        m.extfn = ExternalFunction(
            python_callback_function, units=u.m / u.s, arg_units=[u.m, u.s]
        )
        m.conext = Constraint(expr=m.extfn(m.dx, m.t) - m.vx == 0)
        m.cset = ContinuousSet(bounds=(0, 1))
        m.svar = Var(m.cset, units=u.m)
        m.dvar = DerivativeVar(sVar=m.svar, units=u.m / u.s)

        def prt1_rule(m):
            return {"avar": m.dx}

        def prt2_rule(m):
            return {"avar": m.dy}

        m.prt1 = Port(rule=prt1_rule)
        m.prt2 = Port(rule=prt2_rule)

        def arcrule(m):
            return dict(source=m.prt1, destination=m.prt2)

        m.arc = Arc(rule=arcrule)

        # complementarities do not work yet
        # The expression system removes the u.m since it is multiplied by zero.
        # We need to change the units_container to allow 0 when comparing units
        # m.compl = Complementarity(expr=complements(m.dx/m.t >= m.vx, m.dx == 0*u.m))

        assert_units_consistent(m)

    def test_units_roundoff_error(self):
        # Issue 2393: this example resulted in roundoff error where the
        # computed units for var_1 were
        # (0.9999999999999986,
        # <Unit('kilogram ** 1 / kelvin / meter ** 1.66533e-16 / second ** 3')>
        # )
        m = ConcreteModel()
        m.var_1 = Var(
            initialize=400,
            units=(
                (units.J**0.4)
                * (units.kg**0.2)
                * (units.W**0.6)
                / units.K
                / (units.m**2.2)
                / (units.Pa**0.2)
                / (units.s**0.8)
            ),
        )
        m.var_1.fix()
        m.var_2 = Var(initialize=400, units=units.kg / units.s**3 / units.K)
        assert_units_equivalent(m.var_1, m.var_2)

    def test_identify_inconsistent_units(self):
        u = units
        m = ConcreteModel()
        m.S = Set(initialize=[1, 2, 3])
        m.x = Var(units=u.m)
        m.t = Var(units=u.s)

        # Scalar constraints
        m.c1 = Constraint(expr=m.x == 10 * u.m)  # Units consistent
        m.c2 = Constraint(expr=m.x == m.t)  # Units inconsistent

        # Indexed Constraint
        @m.Constraint(m.S)
        def c3(blk, i):
            if i == 1:
                return m.t == 10 * u.m  # Inconsistent units
            return m.t == 10 * u.s  # Consistent units

        # Scalar Expressions
        m.e1 = Expression(expr=m.x + 10 * u.m)  # Units consistent
        m.e2 = Expression(expr=m.x + m.t)  # Units inconsistent

        # Indexed Expression
        @m.Expression(m.S)
        def e3(blk, i):
            if i == 1:
                return m.t + 10 * u.m  # Inconsistent units
            return m.t + 10 * u.s  # Consistent units

        # Scalar Objectives
        m.o1 = Objective(expr=m.x + 10 * u.m)  # Units consistent
        m.o2 = Objective(expr=m.x + m.t)  # Units inconsistent

        # Indexed Objective
        @m.Objective(m.S)
        def o3(blk, i):
            if i == 1:
                return m.t + 10 * u.m  # Inconsistent units
            return m.t + 10 * u.s  # Consistent units

        failures = identify_inconsistent_units(m)

        assert len(failures) == 6
        assert m.c2 in failures
        assert m.c3[1] in failures
        assert m.e2 in failures
        assert m.e3[1] in failures
        assert m.o2 in failures
        assert m.o3[1] in failures


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