#!/usr/bin/env python3
# Copyright 2010-2022 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.
"""Solve a random Weighted Latency problem with the CP-SAT solver."""

import random
from typing import Sequence
from absl import app
from absl import flags

from google.protobuf import text_format
from ortools.sat.python import cp_model

_NUM_NODES = flags.DEFINE_integer('num_nodes', 12, 'Number of nodes to visit.')
_GRID_SIZE = flags.DEFINE_integer('grid_size', 20,
                                  'Size of the grid where nodes are.')
_PROFIT_RANGE = flags.DEFINE_integer('profit_range', 50, 'Range of profit.')
_SEED = flags.DEFINE_integer('seed', 0, 'Random seed.')
_PARAMS = flags.DEFINE_string('params',
                              'num_search_workers:16, max_time_in_seconds:5',
                              'Sat solver parameters.')
_PROTO_FILE = flags.DEFINE_string(
    'proto_file', '', 'If not empty, output the proto to this file.')


def build_model():
    """Create the nodes and the profit."""
    random.seed(_SEED.value)
    x = []
    y = []
    x.append(random.randint(0, _GRID_SIZE.value))
    y.append(random.randint(0, _GRID_SIZE.value))
    for _ in range(_NUM_NODES.value):
        x.append(random.randint(0, _GRID_SIZE.value))
        y.append(random.randint(0, _GRID_SIZE.value))

    profits = []
    profits.append(0)
    for _ in range(_NUM_NODES.value):
        profits.append(random.randint(1, _PROFIT_RANGE.value))
    sum_of_profits = sum(profits)
    profits = [p / sum_of_profits for p in profits]

    return x, y, profits


def solve_with_cp_sat(x, y, profits):
    """Solves the problem with the CP-SAT solver."""
    model = cp_model.CpModel()

    # because of the manhattan distance, the sum of distances is bounded by this.
    horizon = _GRID_SIZE.value * 2 * _NUM_NODES.value
    times = [
        model.NewIntVar(0, horizon, f'x_{i}')
        for i in range(_NUM_NODES.value + 1)
    ]

    # Node 0 is the start node.
    model.Add(times[0] == 0)

    # Create the circuit constraint.
    arcs = []
    for i in range(_NUM_NODES.value + 1):
        for j in range(_NUM_NODES.value + 1):
            if i == j:
                continue
            # We use a manhattan distance between nodes.
            distance = abs(x[i] - x[j]) + abs(y[i] - y[j])
            lit = model.NewBoolVar(f'{i}_to_{j}')
            arcs.append((i, j, lit))

            # Add transitions between nodes.
            if i == 0:
                # Initial transition
                model.Add(times[j] == distance).OnlyEnforceIf(lit)
            elif j != 0:
                # We do not care for the last transition.
                model.Add(times[j] == times[i] + distance).OnlyEnforceIf(lit)
    model.AddCircuit(arcs)

    model.Minimize(cp_model.LinearExpr.WeightedSum(times, profits))

    if _PROTO_FILE.value:
        model.ExportToFile(_PROTO_FILE.value)

    # Solve model.
    solver = cp_model.CpSolver()
    if _PARAMS.value:
        text_format.Parse(_PARAMS.value, solver.parameters)
    solver.parameters.log_search_progress = True
    solver.Solve(model)


def main(argv: Sequence[str]) -> None:
    if len(argv) > 1:
        raise app.UsageError('Too many command-line arguments.')

    x, y, profits = build_model()
    solve_with_cp_sat(x, y, profits)
    # TODO(user): Implement routing model.


if __name__ == '__main__':
    app.run(main)