"""Generates the specialized gemm functions."""

import mul_Nx8_Mx8_neon
import qnt_Nx8_neon
import zip_Nx8_neon

_QUANTIZED_8BIT = 'quantized_8bit'
_FULL_32BIT = 'full_32bit'
_FULL_FLOAT = 'full_float'


class Error(Exception):
  """Module level error."""


class ConfigurationError(Error):
  """Runtime configuration error."""


def GenerateCommonTempsCountersAndConsts(emitter, rows):
  emitter.EmitDeclare('const std::int32_t', 'row_chunks', 'm / 3')
  emitter.EmitDeclare('const std::int32_t', 'col_chunks', 'n / 3')
  emitter.EmitDeclare('const std::int32_t', 'padded_k', '((k + 7) / 8) * 8')
  emitter.EmitDeclare('const std::int32_t', 'chunk_size', 'k * 3')
  emitter.EmitDeclare('const std::int32_t', 'zipped_chunk_size',
                      '(padded_k + 16) * 3')
  emitter.EmitDeclare('const std::int32_t', 'zipped_rhs_size',
                      '(padded_k + 16) * n')
  emitter.EmitDeclare('const std::uint8_t*', 'lhs_chunk', 'lhs')
  emitter.EmitDeclare('const std::uint8_t*', 'rhs_chunk', 'rhs')
  emitter.EmitDeclare('std::uint8_t*', 'zipped_lhs', 'scratch')
  emitter.EmitDeclare(
      'std::int32_t*', 'zipped_lhs_3_offsets',
      'reinterpret_cast<std::int32_t*>(zipped_lhs + padded_k * 3)')
  if rows is not 0:
    emitter.EmitDeclare(
        'std::int32_t*', 'zipped_lhs_%d_offsets' % rows,
        'reinterpret_cast<std::int32_t*>(zipped_lhs + padded_k * %d)' % rows)
  emitter.EmitDeclare('std::uint8_t*', 'zipped_rhs',
                      'scratch + zipped_chunk_size')
  emitter.EmitDeclare('std::uint8_t*', 'zipped_rhs_chunk', 'zipped_rhs')
  emitter.EmitDeclare('const std::int32_t', 'result_chunk_stride',
                      'result_stride * 3')
  emitter.EmitNewline()


def GenerateQuantized8BitTempsCountersAndConsts(emitter, rows):
  """Generates all the boilerplate variables for the q8 gemm function."""
  GenerateCommonTempsCountersAndConsts(emitter, rows)
  emitter.EmitDeclare('const std::int32_t', 'const_offset',
                      'lhs_offset * rhs_offset * k + result_offset')
  emitter.EmitDeclare('const std::int32_t', 'rounding_offset',
                      '(1 << (shift - 1))')
  emitter.EmitDeclare('std::int32_t*', 'temp_result',
                      'reinterpret_cast<std::int32_t*>('
                      'scratch + zipped_chunk_size + zipped_rhs_size)')
  emitter.EmitDeclare('std::uint8_t*', 'result_chunk', 'result')
  emitter.EmitDeclare('std::int32_t*', 'mul_result_chunk', 'temp_result')
  emitter.EmitDeclare('const std::int32_t', 'mul_result_chunk_stride_bytes',
                      '((n * 4 + 7) / 8) * 8')
  emitter.EmitNewline()


def GenerateFullTempsCountersAndConsts(emitter, result_type, rows):
  """Generates all the boilerplate variables for the int32 and float gemms."""
  GenerateCommonTempsCountersAndConsts(emitter, rows)
  emitter.EmitDeclare('const std::int32_t', 'const_offset',
                      'lhs_offset * rhs_offset * k')
  emitter.EmitDeclare(result_type, 'result_chunk', 'result')
  emitter.EmitDeclare(result_type, 'mul_result_chunk', 'result')
  emitter.EmitDeclare('const std::int32_t', 'mul_result_chunk_stride_bytes',
                      'result_stride * 4')
  emitter.EmitNewline()


def ZipName(rows, leftovers, aligned):
  return zip_Nx8_neon.BuildName(rows, leftovers, aligned)


def GenerateZipRhs(emitter, aligned, cols, leftovers):
  """Emits the code responsible for zipping the rhs matrix."""
  emitter.EmitOpenBracket('for (int i = 0; i < col_chunks; ++i)')
  emitter.EmitCall(
      ZipName(3, leftovers, aligned),
      ['rhs_chunk', 'k', 'k', 'zipped_rhs_chunk', 'lhs_offset', 0])
  emitter.EmitAssignIncrement('rhs_chunk', 'chunk_size')
  emitter.EmitAssignIncrement('zipped_rhs_chunk', 'zipped_chunk_size')
  emitter.EmitCloseBracket()

  if cols is not 0:
    emitter.EmitCall(
        ZipName(cols, leftovers, aligned),
        ['rhs_chunk', 'k', 'k', 'zipped_rhs_chunk', 'lhs_offset', 0])
  emitter.EmitNewline()


def MulName(result_type, lhs_add, rhs_add, rows, cols):
  return mul_Nx8_Mx8_neon.BuildName(result_type, lhs_add, rhs_add, rows, cols)


def GetMulParams(result_type):
  params = ['zipped_lhs', 'zipped_rhs_chunk', 'padded_k', 'mul_result_chunk',
            'mul_result_chunk_stride_bytes']
  if result_type is 'float':
    params.append('result_scale')
  return params


def GenerateMulRows(emitter, result, result_type, lhs_add, rhs_add, aligned,
                    rows, cols, leftovers):
  """Emits code responsible for multiplication of one horizontal lhs strip."""
  emitter.EmitCall(
      ZipName(rows, leftovers, aligned),
      ['lhs_chunk', 'k', 'k', 'zipped_lhs', 'rhs_offset', 'const_offset'])
  emitter.EmitAssign('zipped_rhs_chunk', 'zipped_rhs')
  emitter.EmitAssign('mul_result_chunk', result)

  emitter.EmitOpenBracket('for (int j = 0; j < col_chunks; ++j)')

  emitter.EmitCall(
      MulName(result_type, lhs_add, rhs_add, rows, 3),
      GetMulParams(result_type))
  emitter.EmitAssignIncrement('zipped_rhs_chunk', 'zipped_chunk_size')
  emitter.EmitAssignIncrement('mul_result_chunk', 3)

  emitter.EmitCloseBracket()

  if cols is not 0:
    emitter.EmitCall(
        MulName(result_type, lhs_add, rhs_add, rows, cols),
        GetMulParams(result_type))


def GenerateQuantized8BitMul(emitter, aligned, rows, cols, leftovers):
  """Emits code for all lhs strips & leftover rows. Quantize after mul code."""
  emitter.EmitOpenBracket('for (int i = 0; i < row_chunks; ++i)')
  GenerateMulRows(emitter, 'temp_result', 'int32', False, True, aligned, 3,
                  cols, leftovers)
  emitter.EmitCall(
      qnt_Nx8_neon.BuildMultiQuantizeName(aligned, 3),
      ['temp_result', 'n', 'mul_result_chunk_stride_bytes',
       'zipped_lhs_3_offsets', 'result_chunk', 'result_stride',
       'multiplicative_offset', 'rounding_offset', '-shift'])
  emitter.EmitAssignIncrement('lhs_chunk', 'chunk_size')
  emitter.EmitAssignIncrement('result_chunk', 'result_chunk_stride')
  emitter.EmitCloseBracket()
  emitter.EmitNewline()

  if rows is not 0:
    GenerateMulRows(emitter, 'temp_result', 'int32', False, True, aligned, rows,
                    cols, leftovers)
    emitter.EmitCall(
        qnt_Nx8_neon.BuildMultiQuantizeName(aligned, rows),
        ['temp_result', 'n', 'mul_result_chunk_stride_bytes',
         'zipped_lhs_%d_offsets' % rows, 'result_chunk', 'result_stride',
         'multiplicative_offset', 'rounding_offset', '-shift'])


def GenerateFullMul(emitter, result_type, aligned, rows, cols, leftovers):
  emitter.EmitOpenBracket('for (int i = 0; i < row_chunks; ++i)')
  GenerateMulRows(emitter, 'result_chunk', result_type, True, True, aligned, 3,
                  cols, leftovers)
  emitter.EmitAssignIncrement('lhs_chunk', 'chunk_size')
  emitter.EmitAssignIncrement('result_chunk', 'result_chunk_stride')
  emitter.EmitCloseBracket()
  emitter.EmitNewline()

  if rows is not 0:
    GenerateMulRows(emitter, 'result_chunk', result_type, True, True, aligned,
                    rows, cols, leftovers)


def BuildName(output_type, aligned, rows, cols, leftover):
  name = BuildMainGemmName(output_type) + '_%d_%d_%d' % (rows, cols, leftover)
  if aligned:
    name += '_aligned'
  return name


def GetCommonGemmParameters():
  return [['std::uint8_t*', 'scratch'], ['const std::uint8_t*', 'lhs'],
          ['const std::uint8_t*', 'rhs'], ['std::int32_t', 'm'],
          ['std::int32_t', 'n'], ['std::int32_t', 'k'],
          ['std::int32_t', 'lhs_offset'], ['std::int32_t', 'rhs_offset']]


def GetGemmParameters(output_type, extra_params=None):
  """Prepares a (type, parameter) array for the gemm functions."""
  if extra_params is None:
    extra_params = []
  params = GetCommonGemmParameters()
  if output_type is _QUANTIZED_8BIT:
    params += [['std::int32_t', 'result_offset'],
               ['std::int32_t', 'multiplicative_offset'],
               ['std::int32_t', 'shift'], ['std::uint8_t*', 'result']]
  elif output_type is _FULL_32BIT:
    params += [['std::int32_t*', 'result']]
  elif output_type is _FULL_FLOAT:
    params += [['float', 'result_scale'], ['float*', 'result']]
  else:
    raise ConfigurationError('Unsupported output type: %s' % output_type)
  return params + extra_params


def GetStridedGemmParameters(output_type):
  return GetGemmParameters(output_type, [['std::int32_t', 'result_stride']])


def GenerateGemm(emitter, output_type, aligned, rows, cols, leftovers):
  """Build one gemm function for given row, col, and depth leftovers."""
  emitter.EmitFunctionBeginA(
      BuildName(output_type, aligned, rows, cols, leftovers),
      GetStridedGemmParameters(output_type), 'void')

  emitter.EmitAssert('m %% 3 == %d' % rows)
  emitter.EmitAssert('n %% 3 == %d' % cols)
  emitter.EmitAssert('k %% 8 == %d' % leftovers)

  if output_type is _QUANTIZED_8BIT:
    GenerateQuantized8BitTempsCountersAndConsts(emitter, rows)
    GenerateZipRhs(emitter, aligned, cols, leftovers)
    GenerateQuantized8BitMul(emitter, aligned, rows, cols, leftovers)
  elif output_type is _FULL_32BIT:
    GenerateFullTempsCountersAndConsts(emitter, 'std::int32_t*', rows)
    GenerateZipRhs(emitter, aligned, cols, leftovers)
    GenerateFullMul(emitter, 'int32', aligned, rows, cols, leftovers)
  elif output_type is _FULL_FLOAT:
    GenerateFullTempsCountersAndConsts(emitter, 'float*', rows)
    GenerateZipRhs(emitter, aligned, cols, leftovers)
    GenerateFullMul(emitter, 'float', aligned, rows, cols, leftovers)
  else:
    raise ConfigurationError('Unknown output type: %s' % output_type)

  emitter.EmitFunctionEnd()


def GenerateGemmCall(emitter, output_type, aligned, m_mod, n_mod, leftovers):
  emitter.EmitCall(
      emitter.Scope('internal',
                    BuildName(output_type, aligned, m_mod, n_mod, leftovers)),
      [p for (unused_t, p) in GetStridedGemmParameters(output_type)])


def GenerateGemmSwitch3(emitter, output_type, aligned, m_mod, n_mod):
  """Third level of main switch, choose optimized version on depth leftover."""
  emitter.EmitSwitch('k % 8')

  for leftovers in range(0, 8):
    emitter.EmitCase(leftovers)
    emitter.PushIndent()
    GenerateGemmCall(emitter, output_type, aligned, m_mod, n_mod, leftovers)
    emitter.EmitBreak()
    emitter.PopIndent()

  emitter.EmitSwitchEnd()


def GenerateGemmSwitch2(emitter, output_type, aligned, m_mod):
  """Second level of main switch, choose optimized version on cols leftover."""
  emitter.EmitSwitch('n % 3')

  for n_mod in range(0, 3):
    emitter.EmitCase(n_mod)
    emitter.PushIndent()
    GenerateGemmSwitch3(emitter, output_type, aligned, m_mod, n_mod)
    emitter.EmitBreak()
    emitter.PopIndent()

  emitter.EmitSwitchEnd()


def GenerateGemmSwitch1(emitter, output_type, aligned):
  """First level of main switch, choose optimized version on rows leftover."""
  emitter.EmitSwitch('m % 3')

  for m_mod in range(0, 3):
    emitter.EmitCase(m_mod)
    emitter.PushIndent()
    GenerateGemmSwitch2(emitter, output_type, aligned, m_mod)
    emitter.EmitBreak()
    emitter.PopIndent()

  emitter.EmitSwitchEnd()


def BuildMainGemmName(output_type):
  if output_type is _QUANTIZED_8BIT:
    return 'gemm_q8'
  elif output_type is _FULL_32BIT:
    return 'gemm_i32'
  elif output_type is _FULL_FLOAT:
    return 'gemm_f'
  else:
    raise ConfigurationError('Unsupported output type: %s' % output_type)


def BuildStridedMainGemmName(output_type):
  return BuildMainGemmName(output_type) + '_strided'


def GenerateMainGemmFunction(emitter, output_type):
  """Emit high level gemm function that switches between optimized versions."""
  emitter.EmitFunctionBeginA(
      BuildStridedMainGemmName(output_type),
      GetStridedGemmParameters(output_type), 'void')

  emitter.EmitDeclare('const bool', 'lhs_aligned',
                      '((reinterpret_cast<std::uintptr_t>(lhs) % 8) == 0)')
  emitter.EmitDeclare('const bool', 'rhs_aligned',
                      '((reinterpret_cast<std::uintptr_t>(rhs) % 8) == 0)')
  emitter.EmitDeclare('const bool', 'k_aligned', '((k % 8) == 0)')

  if output_type is _QUANTIZED_8BIT:
    emitter.EmitDeclare('const bool', 'result_aligned',
                        '((reinterpret_cast<std::uintptr_t>(result) % 8) == 0)')
    emitter.EmitDeclare('const bool', 'result_stride_aligned',
                        '((result_stride % 8) == 0)')
    emitter.EmitDeclare('const bool', 'aligned',
                        'lhs_aligned && rhs_aligned && result_aligned '
                        '&& k_aligned && result_stride_aligned')
  else:
    emitter.EmitDeclare('const bool', 'aligned',
                        'lhs_aligned && rhs_aligned && k_aligned')

  emitter.EmitIf('aligned')
  GenerateGemmSwitch1(emitter, output_type, True)
  emitter.EmitElse()
  GenerateGemmSwitch1(emitter, output_type, False)
  emitter.EmitEndif()
  emitter.EmitFunctionEnd()


def GenerateWrapperGemmFunction(emitter, output_type):
  emitter.EmitFunctionBeginA(
      BuildMainGemmName(output_type), GetGemmParameters(output_type), 'void')
  emitter.EmitCall(
      BuildStridedMainGemmName(output_type),
      [p for (unused_t, p) in GetGemmParameters(output_type)] + ['n'])
  emitter.EmitFunctionEnd()


def GenerateInternalFunctions(emitter):
  """Generate all the functions hidden in the internal namespace."""
  for output_type in [_QUANTIZED_8BIT, _FULL_32BIT, _FULL_FLOAT]:
    for aligned in [True, False]:
      for rows in range(0, 3):
        for cols in range(0, 3):
          for leftover in range(0, 8):
            GenerateGemm(emitter, output_type, aligned, rows, cols, leftover)
            emitter.EmitNewline()


def GeneratePublicFunctions(emitter):
  for output_type in [_QUANTIZED_8BIT, _FULL_32BIT, _FULL_FLOAT]:
    GenerateMainGemmFunction(emitter, output_type)
    emitter.EmitNewline()

    GenerateWrapperGemmFunction(emitter, output_type)
    emitter.EmitNewline()