#!/usr/bin/env python3
# Copyright 2010-2021 Google LLC
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for model_builder."""

import math

from ortools.model_builder.python import model_builder
import unittest
import os


class ModelBuilderTest(unittest.TestCase):
    # Number of decimal places to use for numerical tolerance for
    # checking primal, dual, objective values and other values.
    NUM_PLACES = 5

    # pylint: disable=too-many-statements
    def run_minimal_linear_example(self, solver_name):
        """Minimal Linear Example."""
        model = model_builder.ModelBuilder()
        model.name = 'minimal_linear_example'
        x1 = model.new_num_var(0.0, math.inf, 'x1')
        x2 = model.new_num_var(0.0, math.inf, 'x2')
        x3 = model.new_num_var(0.0, math.inf, 'x3')
        self.assertEqual(3, model.num_variables)
        self.assertFalse(x1.is_integral)
        self.assertEqual(0.0, x1.lower_bound)
        self.assertEqual(math.inf, x2.upper_bound)
        x1.lower_bound = 1.0
        self.assertEqual(1.0, x1.lower_bound)

        model.maximize(10.0 * x1 + 6 * x2 + 4.0 * x3 - 3.5)
        self.assertEqual(4.0, x3.objective_coefficient)
        self.assertEqual(-3.5, model.objective_offset)
        model.objective_offset = -5.5
        self.assertEqual(-5.5, model.objective_offset)

        c0 = model.add(x1 + x2 + x3 <= 100.0)
        self.assertEqual(100, c0.upper_bound)
        c1 = model.add(10 * x1 + 4.0 * x2 + 5.0 * x3 <= 600.0, 'c1')
        self.assertEqual('c1', c1.name)
        c2 = model.add(2.0 * x1 + 2.0 * x2 + 6.0 * x3 <= 300.0)
        self.assertEqual(-math.inf, c2.lower_bound)

        solver = model_builder.ModelSolver(solver_name)
        self.assertEqual(model_builder.SolveStatus.OPTIMAL, solver.solve(model))

        # The problem has an optimal solution.
        self.assertAlmostEqual(733.333333 + model.objective_offset,
                               solver.objective_value,
                               places=self.NUM_PLACES)
        self.assertAlmostEqual(33.333333,
                               solver.value(x1),
                               places=self.NUM_PLACES)
        self.assertAlmostEqual(66.666667,
                               solver.value(x2),
                               places=self.NUM_PLACES)
        self.assertAlmostEqual(0.0, solver.value(x3), places=self.NUM_PLACES)

        dual_objective_value = (solver.dual_value(c0) * c0.upper_bound +
                                solver.dual_value(c1) * c1.upper_bound +
                                solver.dual_value(c2) * c2.upper_bound +
                                model.objective_offset)
        self.assertAlmostEqual(solver.objective_value,
                               dual_objective_value,
                               places=self.NUM_PLACES)

        # x1 and x2 are basic
        self.assertAlmostEqual(0.0,
                               solver.reduced_cost(x1),
                               places=self.NUM_PLACES)
        self.assertAlmostEqual(0.0,
                               solver.reduced_cost(x2),
                               places=self.NUM_PLACES)
        # x3 is non-basic
        x3_expected_reduced_cost = (4.0 - 1.0 * solver.dual_value(c0) -
                                    5.0 * solver.dual_value(c1))
        self.assertAlmostEqual(x3_expected_reduced_cost,
                               solver.reduced_cost(x3),
                               places=self.NUM_PLACES)

        self.assertIn('minimal_linear_example',
                      model.export_to_lp_string(False))
        self.assertIn('minimal_linear_example',
                      model.export_to_mps_string(False))

    def test_minimal_linear_example(self):
        self.run_minimal_linear_example('glop')

    def test_import_from_mps_string(self):
        mps_data = """
* Generated by MPModelProtoExporter
*   Name             : SupportedMaximizationProblem
*   Format           : Free
*   Constraints      : 0
*   Variables        : 1
*     Binary         : 0
*     Integer        : 0
*     Continuous     : 1
NAME          SupportedMaximizationProblem
OBJSENSE
  MAX
ROWS
 N  COST
COLUMNS
    X_ONE   COST         1
BOUNDS
 UP BOUND   X_ONE        4
ENDATA
"""
        model = model_builder.ModelBuilder()
        self.assertTrue(model.import_from_mps_string(mps_data))
        self.assertEqual(model.name, 'SupportedMaximizationProblem')

    def test_import_from_mps_file(self):
        path = os.path.dirname(__file__)
        mps_path = f'{path}/../data/tests/maximization.mps'
        model = model_builder.ModelBuilder()
        self.assertTrue(model.import_from_mps_file(mps_path))
        self.assertEqual(model.name, 'SupportedMaximizationProblem')

    def test_import_from_lp_string(self):
        lp_data = """
      min: x + y;
      bin: b1, b2, b3;
      1 <= x <= 42;
      constraint_num1: 5 b1 + 3b2 + x <= 7;
      4 y + b2 - 3 b3 <= 2;
      constraint_num2: -4 b1 + b2 - 3 z <= -2;
"""
        model = model_builder.ModelBuilder()
        self.assertTrue(model.import_from_lp_string(lp_data))
        self.assertEqual(6, model.num_variables)
        self.assertEqual(3, model.num_constraints)
        self.assertEqual(1, model.var_from_index(0).lower_bound)
        self.assertEqual(42, model.var_from_index(0).upper_bound)
        self.assertEqual('x', model.var_from_index(0).name)

    def test_import_from_lp_file(self):
        path = os.path.dirname(__file__)
        lp_path = f'{path}/../data/tests/small_model.lp'
        model = model_builder.ModelBuilder()
        self.assertTrue(model.import_from_lp_file(lp_path))
        self.assertEqual(6, model.num_variables)
        self.assertEqual(3, model.num_constraints)
        self.assertEqual(1, model.var_from_index(0).lower_bound)
        self.assertEqual(42, model.var_from_index(0).upper_bound)
        self.assertEqual('x', model.var_from_index(0).name)

    def test_variables(self):
        model = model_builder.ModelBuilder()
        x = model.new_int_var(0.0, 4.0, 'x')
        self.assertEqual(0, x.index)
        self.assertEqual(0.0, x.lower_bound)
        self.assertEqual(4.0, x.upper_bound)
        self.assertEqual('x', x.name)
        x.lower_bound = 1.0
        x.upper_bound = 3.0
        self.assertEqual(1.0, x.lower_bound)
        self.assertEqual(3.0, x.upper_bound)
        self.assertTrue(x.is_integral)

        # Tests the equality operator.
        y = model.new_int_var(0.0, 4.0, 'y')
        x_copy = model.var_from_index(0)
        self.assertEqual(x, x)
        self.assertEqual(x, x_copy)
        self.assertNotEqual(x, y)

        # Tests the hash method.
        var_set = set()
        var_set.add(x)
        self.assertIn(x, var_set)
        self.assertIn(x_copy, var_set)
        self.assertNotIn(y, var_set)


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