# Copyright 2022 DeepMind Technologies Limited
#
# 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.
# ==============================================================================

"""Best known ranks for different matrix multiplication sizes."""

import math


# Best known ranks of matrix multiplication tensors including improvements found
# by AlphaTensor, but excluding improvements found by recombining solutions.
_SOTA = {
    (3, 3, 3): 23,
    (3, 3, 4): 29,
    (3, 3, 5): 36,
    (3, 3, 6): 40,
    (3, 3, 7): 49,
    (3, 3, 8): 55,
    (3, 3, 9): 63,
    (3, 3, 10): 69,
    (3, 3, 11): 76,
    (3, 3, 12): 80,
    (3, 4, 4): 38,
    (3, 4, 5): 47,  # Discovered by AlphaTensor.
    (3, 4, 6): 56,
    (3, 4, 7): 66,
    (3, 4, 8): 75,
    (3, 4, 9): 85,
    (3, 4, 10): 94,
    (3, 4, 11): 104,
    (3, 4, 12): 112,
    (3, 5, 5): 58,
    (3, 5, 6): 70,
    (3, 5, 7): 83,
    (3, 5, 8): 94,
    (3, 5, 9): 106,
    (3, 5, 10): 116,
    (3, 5, 11): 128,
    (3, 5, 12): 140,
    (3, 6, 6): 80,
    (3, 6, 7): 96,
    (3, 6, 8): 110,
    (3, 6, 9): 120,
    (3, 6, 10): 136,
    (3, 6, 11): 150,
    (3, 6, 12): 160,
    (3, 7, 7): 115,
    (3, 7, 8): 130,
    (3, 7, 9): 145,
    (3, 7, 10): 162,
    (3, 7, 11): 179,
    (3, 7, 12): 192,
    (3, 8, 8): 149,
    (3, 8, 9): 165,
    (3, 8, 10): 185,
    (3, 8, 11): 204,
    (3, 8, 12): 220,
    (3, 9, 9): 183,
    (3, 9, 10): 205,
    (3, 9, 11): 226,
    (3, 9, 12): 240,
    (3, 10, 10): 230,
    (3, 10, 11): 252,
    (3, 10, 12): 272,
    (3, 11, 11): 276,
    (3, 11, 12): 300,
    (3, 12, 12): 320,
    (4, 4, 4): 49,
    (4, 4, 5): 63,  # Discovered by AlphaTensor.
    (4, 4, 6): 75,
    (4, 4, 7): 87,
    (4, 4, 8): 98,
    (4, 4, 9): 104,
    (4, 4, 10): 120,
    (4, 4, 11): 130,
    (4, 4, 12): 142,
    (4, 5, 5): 76,  # Discovered by AlphaTensor.
    (4, 5, 6): 93,
    (4, 5, 7): 109,
    (4, 5, 8): 122,
    (4, 5, 9): 140,
    (4, 5, 10): 154,
    (4, 5, 11): 170,
    (4, 5, 12): 183,
    (4, 6, 6): 105,
    (4, 6, 7): 125,
    (4, 6, 8): 140,
    (4, 6, 9): 160,
    (4, 6, 10): 175,
    (4, 6, 11): 195,
    (4, 6, 12): 210,
    (4, 7, 7): 147,
    (4, 7, 8): 164,
    (4, 7, 9): 188,
    (4, 7, 10): 207,
    (4, 7, 11): 229,
    (4, 7, 12): 246,
    (4, 8, 8): 182,
    (4, 8, 9): 208,
    (4, 8, 10): 231,
    (4, 8, 11): 255,
    (4, 8, 12): 272,
    (4, 9, 9): 225,
    (4, 9, 10): 259,
    (4, 9, 11): 284,
    (4, 9, 12): 300,
    (4, 10, 10): 280,
    (4, 10, 11): 312,
    (4, 10, 12): 336,
    (4, 11, 11): 346,
    (4, 11, 12): 372,
    (4, 12, 12): 390,
    (5, 5, 5): 98,
    (5, 5, 6): 116,
    (5, 5, 7): 136,
    (5, 5, 8): 152,
    (5, 5, 9): 173,
    (5, 5, 10): 190,
    (5, 5, 11): 206,
    (5, 5, 12): 221,
    (5, 6, 6): 137,
    (5, 6, 7): 159,
    (5, 6, 8): 176,
    (5, 6, 9): 200,
    (5, 6, 10): 218,
    (5, 6, 11): 236,
    (5, 6, 12): 250,
    (5, 7, 7): 185,
    (5, 7, 8): 206,
    (5, 7, 9): 235,
    (5, 7, 10): 258,
    (5, 7, 11): 281,
    (5, 7, 12): 300,
    (5, 8, 8): 230,
    (5, 8, 9): 264,
    (5, 8, 10): 291,
    (5, 8, 11): 319,
    (5, 8, 12): 341,
    (5, 9, 9): 297,
    (5, 9, 10): 325,
    (5, 9, 11): 359,
    (5, 9, 12): 384,
    (5, 10, 10): 352,
    (5, 10, 11): 392,
    (5, 10, 12): 424,
    (5, 11, 11): 429,
    (5, 11, 12): 462,
    (5, 12, 12): 491,
    (6, 6, 6): 160,
    (6, 6, 7): 185,
    (6, 6, 8): 203,
    (6, 6, 9): 225,
    (6, 6, 10): 247,
    (6, 6, 11): 268,
    (6, 6, 12): 280,
    (6, 7, 7): 215,
    (6, 7, 8): 239,
    (6, 7, 9): 273,
    (6, 7, 10): 297,
    (6, 7, 11): 324,
    (6, 7, 12): 344,
    (6, 8, 8): 266,
    (6, 8, 9): 296,
    (6, 8, 10): 336,
    (6, 8, 11): 368,
    (6, 8, 12): 390,
    (6, 9, 9): 343,
    (6, 9, 10): 375,
    (6, 9, 11): 416,
    (6, 9, 12): 435,
    (6, 10, 10): 406,
    (6, 10, 11): 454,
    (6, 10, 12): 488,
    (6, 11, 11): 496,
    (6, 11, 12): 530,
    (6, 12, 12): 560,
    (7, 7, 7): 250,
    (7, 7, 8): 279,
    (7, 7, 9): 321,
    (7, 7, 10): 352,
    (7, 7, 11): 387,
    (7, 7, 12): 414,
    (7, 8, 8): 310,
    (7, 8, 9): 360,
    (7, 8, 10): 398,
    (7, 8, 11): 438,
    (7, 8, 12): 468,
    (7, 9, 9): 406,
    (7, 9, 10): 450,
    (7, 9, 11): 492,
    (7, 9, 12): 522,
    (7, 10, 10): 494,
    (7, 10, 11): 543,
    (7, 10, 12): 576,
    (7, 11, 11): 589,
    (7, 11, 12): 626,
    (7, 12, 12): 660,
    (8, 8, 8): 343,
    (8, 8, 9): 392,
    (8, 8, 10): 443,
    (8, 8, 11): 492,
    (8, 8, 12): 508,
    (8, 9, 9): 430,
    (8, 9, 10): 492,
    (8, 9, 11): 543,
    (8, 9, 12): 570,
    (8, 10, 10): 559,
    (8, 10, 11): 610,
    (8, 10, 12): 645,
    (8, 11, 11): 660,
    (8, 11, 12): 699,
    (8, 12, 12): 735,
    (9, 9, 9): 511,
    (9, 9, 10): 540,
    (9, 9, 11): 600,
    (9, 9, 12): 600,
    (9, 10, 10): 625,
    (9, 10, 11): 680,
    (9, 10, 12): 708,
    (9, 11, 11): 754,
    (9, 11, 12): 768,
    (9, 12, 12): 800,
    (10, 10, 10): 686,
    (10, 10, 11): 758,
    (10, 10, 12): 812,
    (10, 11, 11): 836,
    (10, 11, 12): 894,
    (10, 12, 12): 936,
    (11, 11, 11): 919,
    (11, 11, 12): 972,
    (11, 12, 12): 1022,
    (12, 12, 12): 1040
}


def get_sota_rank(a: int, b: int, c: int) -> int:
  """Returns best known rank of T_{a, b, c} (without recombination results)."""
  a, b, c = sorted([a, b, c])
  if a == 0:
    return 0
  if a == 1:
    return b * c
  if a == 2:
    # Hopcroft & Kerr.
    return int(math.ceil((3 * b * c + max(b, c)) / 2))

  return _SOTA[(a, b, c)]
