1# Copyright 2010-2021 Google LLC
2# Licensed under the Apache License, Version 2.0 (the "License");
3# you may not use this file except in compliance with the License.
4# You may obtain a copy of the License at
5#
6#     http://www.apache.org/licenses/LICENSE-2.0
7#
8# Unless required by applicable law or agreed to in writing, software
9# distributed under the License is distributed on an "AS IS" BASIS,
10# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11# See the License for the specific language governing permissions and
12# limitations under the License.
13"""Collection of helpers to visualize cp_model solutions in colab."""
14
15# pylint: disable=g-import-not-at-top
16import random
17try:
18    from IPython.display import display
19    from IPython.display import SVG
20    import plotly.figure_factory as ff
21    import svgwrite
22    correct_imports = True
23except ImportError:
24    correct_imports = False
25
26
27def RunFromIPython():
28    if not correct_imports:
29        return False
30    try:
31        return __IPYTHON__ is not None
32    except NameError:
33        return False
34
35
36def ToDate(v):
37    return '2016-01-01 6:%02i:%02i' % (v / 60, v % 60)
38
39
40class ColorManager(object):
41    """Utility to create colors to use in visualization."""
42
43    def ScaledColor(self, sr, sg, sb, er, eg, eb, num_steps, step):
44        """Creates an interpolated rgb color between two rgb colors."""
45        num_intervals = num_steps - 1
46        dr = (er - sr) / num_intervals
47        dg = (eg - sg) / num_intervals
48        db = (eb - sb) / num_intervals
49        r = sr + dr * step
50        g = sg + dg * step
51        b = sb + db * step
52        return 'rgb(%i, %i, %i)' % (r, g, b)
53
54    def SeedRandomColor(self, seed=0):
55        random.seed(seed)
56
57    def RandomColor(self):
58        return 'rgb(%i,%i,%i)' % (random.randint(0, 255), random.randint(
59            0, 255), random.randint(0, 255))
60
61
62def DisplayJobshop(starts, durations, machines, name):
63    """Simple function to display a jobshop solution using plotly."""
64
65    jobs_count = len(starts)
66    machines_count = len(starts[0])
67    all_machines = range(0, machines_count)
68    all_jobs = range(0, jobs_count)
69    df = []
70    for i in all_jobs:
71        for j in all_machines:
72            df.append(
73                dict(Task='Resource%i' % machines[i][j],
74                     Start=ToDate(starts[i][j]),
75                     Finish=ToDate(starts[i][j] + durations[i][j]),
76                     Resource='Job%i' % i))
77
78    sorted_df = sorted(df, key=lambda k: k['Task'])
79
80    colors = {}
81    cm = ColorManager()
82    cm.SeedRandomColor(0)
83    for i in all_jobs:
84        colors['Job%i' % i] = cm.RandomColor()
85
86    fig = ff.create_gantt(sorted_df,
87                          colors=colors,
88                          index_col='Resource',
89                          title=name,
90                          show_colorbar=False,
91                          showgrid_x=True,
92                          showgrid_y=True,
93                          group_tasks=True)
94    fig.show()
95
96
97class SvgWrapper(object):
98    """Simple SVG wrapper to use in colab."""
99
100    def __init__(self, sizex, sizey, scaling=20.0):
101        self.__sizex = sizex
102        self.__sizey = sizey
103        self.__scaling = scaling
104        self.__offset = scaling
105        self.__dwg = svgwrite.Drawing(
106            size=(self.__sizex * self.__scaling + self.__offset,
107                  self.__sizey * self.__scaling + self.__offset * 2))
108
109    def Display(self):
110        display(SVG(self.__dwg.tostring()))
111
112    def AddRectangle(self, x, y, dx, dy, fill, stroke='black', label=None):
113        """Draw a rectangle, dx and dy must be >= 0."""
114        s = self.__scaling
115        o = self.__offset
116        corner = (x * s + o, (self.__sizey - y - dy) * s + o)
117        size = (dx * s - 1, dy * s - 1)
118        self.__dwg.add(
119            self.__dwg.rect(insert=corner, size=size, fill=fill, stroke=stroke))
120        self.AddText(x + 0.5 * dx, y + 0.5 * dy, label)
121
122    def AddText(self, x, y, label):
123        text = self.__dwg.text(
124            label,
125            insert=(x * self.__scaling + self.__offset,
126                    (self.__sizey - y) * self.__scaling + self.__offset),
127            text_anchor='middle',
128            font_family='sans-serif',
129            font_size='%dpx' % (self.__scaling / 2))
130        self.__dwg.add(text)
131
132    def AddXScale(self, step=1):
133        """Add an scale on the x axis."""
134        o = self.__offset
135        s = self.__scaling
136        y = self.__sizey * s + o / 2.0 + o
137        dy = self.__offset / 4.0
138        self.__dwg.add(
139            self.__dwg.line((o, y), (self.__sizex * s + o, y), stroke='black'))
140        for i in range(0, int(self.__sizex) + 1, step):
141            self.__dwg.add(
142                self.__dwg.line((o + i * s, y - dy), (o + i * s, y + dy),
143                                stroke='black'))
144
145    def AddYScale(self, step=1):
146        """Add an scale on the y axis."""
147        o = self.__offset
148        s = self.__scaling
149        x = o / 2.0
150        dx = self.__offset / 4.0
151        self.__dwg.add(
152            self.__dwg.line((x, o), (x, self.__sizey * s + o), stroke='black'))
153        for i in range(0, int(self.__sizey) + 1, step):
154            self.__dwg.add(
155                self.__dwg.line((x - dx, i * s + o), (x + dx, i * s + o),
156                                stroke='black'))
157
158    def AddTitle(self, title):
159        """Add a title to the drawing."""
160        text = self.__dwg.text(
161            title,
162            insert=(self.__offset + self.__sizex * self.__scaling / 2.0,
163                    self.__offset / 2),
164            text_anchor='middle',
165            font_family='sans-serif',
166            font_size='%dpx' % (self.__scaling / 2))
167        self.__dwg.add(text)
168