# Copyright 2024 The JAX SC Authors.
#
# 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.
"""An FDO client implementation that uses CSV files as storage."""

import collections
from collections.abc import Mapping
import glob
import itertools
import os
import re
import time

from absl import logging
import jax
from jax_tpu_embedding.sparsecore.lib.fdo import fdo_client
from jax_tpu_embedding.sparsecore.lib.nn import embedding
import numpy as np


_FILE_NAME = 'fdo_stats'
_FILE_EXTENSION = 'npz'
_MAX_ID_STATS_SUFFIX = '_max_ids'
_MAX_UNIQUE_ID_STATS_SUFFIX = '_max_unique_ids'


class NPZFileFDOClient(fdo_client.FDOClient):
  """FDO client that writes stats to a file in .npz format.

  Usage:
    # Create a FDO client.
    client = NPZFileFDOClient(base_dir='/path/to/base/dir')

    # Record observed stats from sparse input processing
    max_ids_per_process, max_uniques_per_process =
      embedding.preprocess_sparse_dense_matmul_input(...)
    client.record(max_ids_per_process, max_uniques_per_process)

    # Publish process local stats to a file.
    client.publish()

    # Load stats from all files in the base_dir.
    max_ids_per_process, max_uniques_per_process = client.load()
  """

  def __init__(self, base_dir: str):
    self._base_dir = base_dir
    self._max_ids_per_partition = collections.defaultdict(np.ndarray)
    self._max_unique_ids_per_partition = collections.defaultdict(np.ndarray)
    self._required_coo_buffer_size_per_sc = collections.defaultdict(np.ndarray)

  def record(self, data: embedding.SparseDenseMatmulInputStats) -> None:
    """Records stats per process.

    Accumulates the max ids observed per process per sparsecore per device for
    each embedding table.
    Accumulates the max unique ids observed per process per sparsecore per
    device for each embedding table.
    Args:
      data: A mapping representing data to be recorded.
    """
    max_ids_per_process = data.max_ids_per_partition
    for table_name, stats in max_ids_per_process.items():
      logging.vlog(
          2, 'Recording observed max ids for table: %s -> %s', table_name, stats
      )
      if table_name not in self._max_ids_per_partition:
        self._max_ids_per_partition[table_name] = stats
      else:
        self._max_ids_per_partition[table_name] = np.vstack(
            (self._max_ids_per_partition[table_name], stats)
        )
    max_uniques_per_process = data.max_unique_ids_per_partition
    for table_name, stats in max_uniques_per_process.items():
      logging.vlog(
          2,
          'Recording observed max unique ids for table: %s -> %s',
          table_name,
          stats,
      )
      if table_name not in self._max_unique_ids_per_partition:
        self._max_unique_ids_per_partition[table_name] = stats
      else:
        self._max_unique_ids_per_partition[table_name] = np.vstack(
            (self._max_unique_ids_per_partition[table_name], stats)
        )
    used_coo_buffer_size = data.required_buffer_size_per_sc
    for table_name, stats in used_coo_buffer_size.items():
      if table_name not in self._required_coo_buffer_size_per_sc:
        self._required_coo_buffer_size_per_sc[table_name] = stats
      else:
        self._required_coo_buffer_size_per_sc[table_name] = np.vstack(
            (self._required_coo_buffer_size_per_sc[table_name], stats)
        )

  # LINT.IfChange(generate_file_name)
  def _generate_file_name(self) -> str:
    """Generates a file name for the stats."""
    # File Format: `fdo_stats_<process_id>_<timestamp>.npz`
    filename = '{}_{}_{}.{}'.format(
        _FILE_NAME, jax.process_index(), time.time_ns(), _FILE_EXTENSION
    )
    return os.path.join(self._base_dir, filename)
  # LINT.ThenChange(:_get_latest_files_by_process)

  def _get_latest_files_by_process(self, files: list[str]) -> list[str]:
    """Returns the latest file for each process."""
    # Note: This function depends on the file name format generated by
    # _generate_file_name.
    # Returns a list of latest file in each `fdo_stats_<process_id>` group.
    if not files:
      return []
    # Read files that match the file creation pattern.
    dir_path = os.path.split(files[0])[0]
    dir_prefix = len(dir_path) + 1
    pattern = fr'{_FILE_NAME}_(\d+)_(\d+)\.{_FILE_EXTENSION}'
    # Group files as list of (process_idx, timestamp, files)
    file_groups = []
    for file in files:
      match = re.search(pattern, file[dir_prefix:])
      if match:
        file_groups.append((match.group(1), int(match.group(2)), file))
    if not file_groups:
      return []
    # Sort the files in descending order to get latest files first.
    file_groups = sorted(file_groups, reverse=True)
    latest_files = []
    for _, file_group in itertools.groupby(file_groups, key=lambda x: x[0]):
      # Get the first item in the group since already sorted by timestamp
      _, _, file_name = next(file_group)
      latest_files.append(file_name)
    return latest_files

  def _write_to_file(self, stats: Mapping[str, np.ndarray]) -> None:
    """Writes stats to a npz file."""
    file_name = self._generate_file_name()
    logging.info('Write stats to %s', file_name)
    jax.numpy.savez(file_name, **stats)

  def publish(self) -> None:
    """Publishes locally accmulatedstats to a file in the base_dir.

    Publish is called by each process there by collecting stats from all
    processes.
    """
    merged_stats = {
        f'{table_name}{_MAX_ID_STATS_SUFFIX}': stats
        for table_name, stats in self._max_ids_per_partition.items()
    }
    merged_stats.update({
        f'{table_name}{_MAX_UNIQUE_ID_STATS_SUFFIX}': stats
        for table_name, stats in self._max_unique_ids_per_partition.items()
    })
    self._write_to_file(merged_stats)

  def _read_from_file(self, files_glob: str) -> Mapping[str, np.ndarray]:
    """Reads stats from a npz file."""
    files = self._get_latest_files_by_process(glob.glob(files_glob))
    if not files:
      raise FileNotFoundError('No stats files found in %s' % files_glob)
    stats = collections.defaultdict(np.ndarray)
    for file_name in files:
      logging.info('Reading stats from %s', file_name)
      loaded = np.load(file_name)
      loaded_d = {key: loaded[key] for key in loaded.files}
      for key, value in loaded_d.items():
        if stats.get(key) is None:
          stats[key] = value
        else:
          stats[key] = np.max(np.vstack((stats[key], value)), axis=0)
    return stats

  def load(
      self,
  ) -> tuple[Mapping[str, np.ndarray], Mapping[str, np.ndarray]]:
    """Loads state of local FDO client from disk.

    Reads all files in the base_dir and aggregates stats.
    Returns:
      A tuple of (max_ids_per_partition, max_unique_ids_per_partition)
    Raises:
      FileNotFoundError: If no stats files are found in the base_dir.
      ValueError: If the stats files do not have expected keys fro max ids and
      max unique ids.
    """
    # TODO(b/393435682): Fallback to default limits if no stats on disk.
    # read files from base_dir and aggregate stats.
    files_glob = os.path.join(
        self._base_dir, '{}*.{}'.format(_FILE_NAME, _FILE_EXTENSION)
    )
    stats = self._read_from_file(files_glob)
    max_id_stats, max_unique_id_stats = {}, {}
    for table_name, stats in stats.items():
      if table_name.endswith(f'{_MAX_ID_STATS_SUFFIX}'):
        max_id_stats[table_name[: -len(_MAX_ID_STATS_SUFFIX)]] = stats
      elif table_name.endswith(f'{_MAX_UNIQUE_ID_STATS_SUFFIX}'):
        max_unique_id_stats[table_name[: -len(_MAX_UNIQUE_ID_STATS_SUFFIX)]] = (
            stats
        )
      else:
        raise ValueError(
            f'Unexpected table name and stats key: {table_name}, expected to'
            f' end with {_MAX_ID_STATS_SUFFIX} or {_MAX_UNIQUE_ID_STATS_SUFFIX}'
        )
    self._max_ids_per_partition = max_id_stats
    self._max_unique_ids_per_partition = max_unique_id_stats
    return (max_id_stats, max_unique_id_stats)

  def get_required_buffer_size_per_sc(self) -> Mapping[str, np.ndarray]:
    return dict(self._required_coo_buffer_size_per_sc)
