# Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import enum, shutil, pickle
from library import *
from matmul import *
from batch_matmul import *
from split_k_matmul import *
from pathlib import Path


###############################################################################
class EmitSourceMLIR:
  """Emitters for the operation MLIR source files."""

  def __init__(self, operation_path, dispatch_collection):
    self.operation_path = operation_path
    self.dispatch_collection = dispatch_collection
    self.operation = dispatch_collection.operation
    self.operation_kind = self.operation.operation_kind
    self.configuration_list = dispatch_collection.configuration_list
    self.operation_filepath = self.operation_path.joinpath(
        self.operation.name()).with_suffix(".mlir")

    mlir_configuration_emitter = {
        OperationKind.Matmul: EmitMatmulCompilationInfo,
        OperationKind.SplitkMatmul: EmitMatmulCompilationInfo,
        OperationKind.BatchMatmul: EmitMatmulCompilationInfo,
    }
    self.configuration_emitter = mlir_configuration_emitter[
        self.operation_kind]()

    mlir_dispatch_emitter = {
        OperationKind.Matmul: EmitLinalgMatmulDispatch,
        OperationKind.SplitkMatmul: EmitLinalgMatmulDispatch,
        OperationKind.BatchMatmul: EmitLinalgBatchMatmulDispatch,
    }
    self.dispatch_emitter = mlir_dispatch_emitter[self.operation_kind]()

  def __enter__(self):
    self.operation_file = open(self.operation_filepath, "w")
    self.operation_file.write(f'// Finename: {self.operation_filepath}')

    # Emit all the configuration attribute tags.
    for configuration in self.configuration_list:
      self.operation_file.write(self.configuration_emitter.emit(configuration))
    return self

  def emit(self):
    """Emit the op func.func for each dispatch (operation + configuration)"""
    for dispatch in self.dispatch_collection.get_dispatches():
      print(
          f"    Emitting tuning configuration : {dispatch.configuration.name()}"
      )
      self.operation_file.write(self.dispatch_emitter.emit(dispatch))

  def __exit__(self, exc_type, exc_value, traceback):
    self.operation_file.close()


###############################################################################
class Manifest:
  """Manifest collects, filters, and stores dispatches in a data structure.
     Manifest organizes the dispatches in a dictionary. 
     Usage:
      1. Create a manifest object with the command line arguments.
      2(a). Generate dispatches, append them in the manifest, and 
            serialize them into a file.
      2(b). Load dispatches from a serialized file.
      
      ```python
      # generator.py usage:
      manifest = Manifest(args)
      manifest.initialize()

      # compile.py or profile.py usage:
      manifest = Manifest(args)
      manifest.load()
      ```
  """

  def __init__(self, args):
    self.args = args

    # Dictionary of operation kind to a list of dispatch collections. We
    # initialize the dictionary during the generation of dispatches and
    # serialize it to a file. The serialized file is used to load the
    # dispatches for compilation and profiling.
    # Datatype: OperationKind -> [DispatchCollection]
    self.dispatch_collection_map = {}

    # For operation kind-based filtering of dispatches.
    self.operation_kind_enabled = []

    # For name-based filtering of dispatches.
    self.dispatch_names = []
    self.ignore_dispatch_names = []

    if args.operation_kind == 'all':
      self.operation_kind_enabled = []
    else:
      operations_kind_list = [
          OperationKind.Matmul,
          OperationKind.SplitkMatmul,
          OperationKind.BatchMatmul,
      ]
      self.operation_kind_enabled = [
          x for x in operations_kind_list
          if OperationKindNames[x] in args.operation_kind.split(',')
      ]

    if args.dispatches == 'all':
      self.dispatch_names = []
    else:
      self.dispatch_names = [x for x in args.dispatches.split(',') if x != '']

    # Paths to the generated directory (e.g. `./generated/linalg`).
    self.generated_path = Path(self.args.generated_dir, 'generated',
                               self.args.mlir_dialect)

    # Create the directories in self.generated_path, if it does not exist.
    if not self.generated_path.exists():
      self.generated_path.mkdir(parents=True, exist_ok=True)

    # Path to the serialized file.
    self.serialized_file_path = self.generated_path.joinpath(
        'serialized_file.pkl')

  def _filter_string_matches(self, filter_string, haystack):
    """Returns true if all substrings appear in the haystack in order"""
    substrings = filter_string.split('*')
    for sub in substrings:
      idx = haystack.find(sub)
      if idx < 0:
        return False
      haystack = haystack[idx + len(sub):]
    return True

  def is_enabled(self, dispatch):
    """Rerturns true if pass through filters based various criteria."""

    # Get the operation and configuration from the dispatch.
    operation = dispatch.operation
    configuration = dispatch.configuration

    # If the operation is not in the enabled list, return False.
    enabled = True

    # If operation_kind filter is enabled and the \
    # operation_kind in not in the enabled list, return False.
    if len(self.operation_kind_enabled) and \
      operation.operation_kind not in self.operation_kind_enabled:
      enabled = False

    # If dispatch name-based filter regex is enabled match the \
    # dispatch name (operation+configuration) against all regexs \
    # in self.dispatch_names.
    if len(self.dispatch_names):
      name = dispatch.name()
      enabled = False

      # compare against each regex included in self.dispatch_names.
      for substr_to_match in self.dispatch_names:
        if self._filter_string_matches(substr_to_match, name):
          enabled = True
          break

    # Return the result of the filter.
    return enabled

  def append_dispatch_collection(self, dispatch_collection):
    """Appends one instance of DispatchCollection to the manifest."""
    operation_kind = dispatch_collection.operation.operation_kind
    if operation_kind not in self.dispatch_collection_map.keys():
      self.dispatch_collection_map[operation_kind] = []

    # Get all the dispatches from the dispatch_collection.
    dispatches = dispatch_collection.get_dispatches()

    # Filter dispatches based on the filter criteria.
    filtered_dispatch_collection = DispatchCollection(
        dispatch_collection.operation, [])
    for dispatch in dispatches:
      if self.is_enabled(dispatch):
        filtered_dispatch_collection.append(dispatch)

    # Only append the filtered_dispatch_collection if it has an unfiltered configuration.
    if len(filtered_dispatch_collection.configuration_list):
      self.dispatch_collection_map[operation_kind].append(
          filtered_dispatch_collection)

  def append(self, dispatch_collection_list):
    """Appends one instance of DispatchCollection to the manifest."""
    for dispatch_collection in dispatch_collection_list:
      self.append_dispatch_collection(dispatch_collection)

  def initialize(self):
    """Initialize the mainfest object by generating dispatches for supported operations."""
    self.append(CudaMatmulGenerator(self.args).generate())
    self.append(CudaSplitKMatmulGenerator(self.args).generate())
    self.append(CudaBatchMatmulGenerator(self.args).generate())

    # Serialize the initialized mainfest state.
    self.dump()

  def dump(self):
    """Serialize (dump) the self.dispatch_collection_map to a pickle file."""
    with open(self.serialized_file_path, 'wb') as f:
      pickle.dump(self.dispatch_collection_map, f)

  def load(self):
    """Deserialize (load) the self.dispatch_collection_map from a pickle file."""
    if not self.serialized_file_path.exists():
      raise ValueError(f"Could not find : {self.serialized_file_path}")

    with open(self.serialized_file_path, 'rb') as load_file:
      self.dispatch_collection_map = pickle.load(load_file)

  def emit(self):
    """Emits the operations in the Manifest to the build directory as MLIR source files.
        The operations are emitted in the dialect specified by the `mlir_dialect` flag.
    """

    # For each operation_kind create a directory and emit the operations with
    # all the configurations in the configuration_list into their seperate directories.
    for operation_kind, dispatch_collection_list\
      in self.dispatch_collection_map.items():

      operation_kind_path = self.generated_path.joinpath(
          OperationKindNames[operation_kind])

      # If the operation_kind_path does not exists, create it.
      if not operation_kind_path.exists():
        operation_kind_path.mkdir(parents=True, exist_ok=True)

      for dispatch_collection in dispatch_collection_list:

        operation_path = operation_kind_path.joinpath(
            dispatch_collection.operation.name())

        # If the operation_path does not exists, create it.
        if not operation_path.exists():
          operation_path.mkdir()

        with EmitSourceMLIR(operation_path,
                            dispatch_collection) as emit_mlir_source:
          mlir_file_path = operation_path.joinpath(
              dispatch_collection.operation.name()).with_suffix('.mlir')
          print(f"[Generating]: {mlir_file_path}")

          # Emit mlir source file for the dispatch_collection.operation with all the configurations
          emit_mlir_source.emit()