import os, itertools

from typing import Any, List, Mapping, Optional, Sequence, Tuple, Union

import numpy as np

from mlir.ir import *
from mlir.dialects import arith, builtin, linalg, scf, func, tensor

from ..core.compilation import attach_inplaceable_attributes, attach_passthrough
from ..core.problem_definition import *
from ..core.utils import *

from . import ops

# TODO: Orthogonal configuration object.
avx512 = True

RANK_RELATED_DIMS = "DHW"


def find_contiguous_rank_dims(lst: str) -> Tuple[int, int]:
  """Returns the positions of the rank-related dimensions in the list.

  Return a pair of values where the first value is the index of the first
  rank-related dimension and the second value is the index of the dimension
  immediately after the last rank-related dimension in the input format string.
  Expects rank-related dimensions to be contiguous in the format.

  Arguments:
  lst: convolution format string containing 'D', 'H' and 'W' to indicate
    rank-related dimensions.
  """
  start, end = None, None
  for i, char in enumerate(lst):
    if char in RANK_RELATED_DIMS:
      if start is None:
        start = i
      end = i
  return start, end + 1


class DepthwiseConvolutionProblem(ProblemDefinition):
  """Benchmarking problem definition for 1/2/3D depthwise convolutions.

  Supports various convolution formats, to be specified at construction time,
  as long as the "image" dimensions (depth, height, width) are contiguous and
  provided in the same order. For example, NCHW is supported but NHCW or NCWH
  are not. Note that in practice, "image" dimension names are interchangeable.
  """

  def __init__(self, input_format: str, kernel_format: str,
               strides: Optional[List[int]], dilations: Optional[List[int]]):
    """Creates a new Depthwise ConvolutionProblem definition with the given