import os
from typing import Dict, List
import numpy as np
import absl.testing.absltest
import tests.Examples.openfhe.ckks.rotom.chained_matmul.chained_matmul_lib as chained_matmul


def read_from_directory(dirpath: str) -> Dict[str, List[float]]:
  """Reads all .npz files from a directory and maps filename to list of floats.

  File format: compressed numpy .npz files with a 'data' key containing the
  array.
  Returns a dict from filename (with extension) to list of floats.
  """
  result = {}

  if not os.path.isdir(dirpath):
    print(f"Error: Could not open directory {dirpath}")
    return result

  for filename in os.listdir(dirpath):
    if filename == "." or filename == "..":
      continue

    if not filename.endswith(".npz"):
      continue

    fullpath = os.path.join(dirpath, filename)
    try:
      npz_data = np.load(fullpath)
      if "data" not in npz_data:
        print(f"Warning: 'data' key not found in {fullpath}")
        continue

      # Extract data array and flatten to 1D list
      data_array = npz_data["data"]
      values = data_array.flatten().tolist()
      result[filename] = values
      npz_data.close()
    except Exception as e:
      print(f"Warning: Could not load file {fullpath}: {e}")
      continue

  return result


class ChainedMatmulTest(absl.testing.absltest.TestCase):
  """End-to-end test for chained matrix multiplication (a @ b @ c).

  Adds an end-to-end test for chained matrix multiplication (a @ b @ c) where a
  is an encrypted input matrix and b, c are plaintext matrices. All matrices
  have shape 64x64 with n=4096.

  Rotom finds the layout assignment:
  a: [1:64:1][0:64:1] (column-major)
  b: roll(0,1) [0:64:1];[1:64:1][R:64:1] (repeated diagonal-major)
  c: roll(0,1) [0:64:1];[1:64:1][R:64:1] (repeated diagonal-major)

  Rotom optimizes the matrix multiplication kernel using baby-step giant-step.

  The layout assignment was automatically generated by Rotom.
  """

  def test_chained_matmul(self):
    # Initialize crypto context
    crypto_context = chained_matmul.chained_matmul__generate_crypto_context()
    key_pair = crypto_context.KeyGen()
    crypto_context = chained_matmul.chained_matmul__configure_crypto_context(
        crypto_context, key_pair.secretKey
    )

    # Read inputs from inputs/ directory and map filename to vector
    inputs_map = read_from_directory(
        "tests/Examples/openfhe/ckks/rotom/chained_matmul/inputs"
    )
    # Read result from results/ directory
    expected_result = read_from_directory(
        "tests/Examples/openfhe/ckks/rotom/chained_matmul/results"
    )

    # Encrypt the specified input as the secret argument
    print(
        f"\n=== Encrypting input 2.npz (length={len(inputs_map['2.npz'])}) ==="
    )
    ct_encrypted = chained_matmul.chained_matmul__encrypt__arg0(
        crypto_context, inputs_map["2.npz"], key_pair.publicKey
    )

    # Call chained_matmul with 128 plaintext arguments
    result = chained_matmul.chained_matmul(
        crypto_context, ct_encrypted, inputs_map["4.npz"], inputs_map["224.npz"]
    )

    actual = chained_matmul.chained_matmul__decrypt__result0(
        crypto_context, result, key_pair.secretKey
    )

    self.assertGreater(len(actual), 0, "Actual result should not be empty")
    self.assertGreater(
        len(expected_result["result.npz"]),
        0,
        "Expected result should not be empty",
    )

    # Check that the actual result is close to the expected result
    for i in range(len(actual)):
      self.assertAlmostEqual(
          expected_result["result.npz"][i], actual[i], delta=1e-3
      )


if __name__ == "__main__":
  absl.testing.absltest.main()
