#  Part of the LLVM Project, 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

"""Experimental MLIR-PyTACO with sparse tensor support.

See http://tensor-compiler.org/ for TACO tensor compiler.

This module implements the PyTACO API for writing a tensor to a file or reading
a tensor from a file.

See the following links for Matrix Market Exchange (.mtx) format and FROSTT
(.tns) format:
  https://math.nist.gov/MatrixMarket/formats.html
  http://frostt.io/tensors/file-formats.html
"""

from typing import List, TextIO

from . import mlir_pytaco

# Define the type aliases so that we can write the implementation here as if
# it were part of mlir_pytaco.py.
Tensor = mlir_pytaco.Tensor
Format = mlir_pytaco.Format
DType = mlir_pytaco.DType
Type = mlir_pytaco.Type

# Constants used in the implementation.
_MTX_FILENAME_SUFFIX = ".mtx"
_TNS_FILENAME_SUFFIX = ".tns"


def read(filename: str, fmt: Format,
         dtype: DType = DType(Type.FLOAT32)) -> Tensor:
  """Inputs a tensor from a given file.

  The name suffix of the file specifies the format of the input tensor. We
  currently only support .mtx format for support sparse tensors.

  Args:
    filename: A string input filename.
    fmt: The storage format of the tensor.
    dtype: The data type, default to float32.

  Raises:
    ValueError: If filename doesn't end with .mtx or .tns, or fmt is not an
    instance of Format or fmt is not a sparse tensor.
  """
  if (not isinstance(filename, str) or
      (not filename.endswith(_MTX_FILENAME_SUFFIX) and
       not filename.endswith(_TNS_FILENAME_SUFFIX))):
    raise ValueError("Expected string filename ends with "
                     f"{_MTX_FILENAME_SUFFIX} or {_TNS_FILENAME_SUFFIX}: "
                     f"{filename}.")

  return Tensor.from_file(filename, fmt, dtype)


def write(filename: str, tensor: Tensor) -> None:
  """Outputs a tensor to a given file.

  The name suffix of the file specifies the format of the output. We currently
  only support .tns format.

  Args:
    filename: A string output filename.
    tensor: The tensor to output.

  Raises:
    ValueError: If filename doesn't end with .tns or tensor is not a Tensor.
  """
  if (not isinstance(filename, str) or
      not filename.endswith(_TNS_FILENAME_SUFFIX)):
    raise ValueError("Expected string filename ends with"
                     f" {_TNS_FILENAME_SUFFIX}: {filename}.")
  if not isinstance(tensor, Tensor):
    raise ValueError(f"Expected a Tensor object: {tensor}.")

  tensor.to_file(filename)