# -*- coding: utf-8 -*-
# Copyright 2018 Google LLC
#
# 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.
# ==============================================================================

import os
import shutil

import numpy as np
import tensorflow as tf

from tensorflowjs import write_weights

TMP_DIR = '/tmp/write_weights_test/'


class TestWriteWeights(tf.test.TestCase):
  def setUp(self):
    if not os.path.isdir(TMP_DIR):
      os.makedirs(TMP_DIR)

  def tearDown(self):
    if os.path.isdir(TMP_DI
R):
      shutil.rmtree(TMP_DIR)

  def test_1_group_1_weight(self):
    groups = [
        [{
          
  'name': 'weight1',
            'data': np.array([1, 2, 3], 'float32')
        }]
    ]

    manifest = write_weights.write_weights(
        groups, TMP_DIR, shard_size_bytes=4 * 4)

    self.assertTrue(
        os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')),
        'weights_manifest.json does not exist')

    self.assertEqual(
        manifest,
        [{
            'paths': ['group1-shard1of1.bin'],
            'weights': [{
                'name': 'weight1',
                'shape': [3],
                'dtype': 'float32'
            }]
        }])

    weights_path = os.path.join(TMP_DIR, 'group1-shard1of1.bin')
    weight1 = np.fromfile(weights_path, 'float32')
    np.testing.assert_array_equal(weight1, np.array([1, 2, 3], 'float32'))

  def test_1_group_1_weight_bool(self):
    groups = [
        [{
            'name': 'weight1',
            'data': np.array([True, False, True], 'bool')
        }]
    ]

    manifest = write_weights.write_weights(
        groups, TMP_DIR, shard_size_bytes=4 * 4)

    self.assertTrue(
        os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')),
        'weights_manifest.json does not exist')

    self.assertEqual(
        manifest,
        [{
            'paths': ['group1-shard1of1.bin'],
            'weights': [{
                'name': 'weight1',
                'shape': [3],
                'dtype': 'bool'
            }]
        }])

    weights_path = os.path.join(TMP_DIR, 'group1-shard1of1.bin')
    weight1 = np.fromfile(weights_path, 'bool')
    np.testing.assert_array_equal(
        weight1, np.array([True, False, True], 'bool'))

  def test_1_group_1_weight_string(self):
    groups = [
        [{
            'name': 'weight1',
            'data': np.array([['здраво', 'end'], ['test', 'a']], 'object')
        }]
    ]

    manifest = write_weights.write_weights(
        groups, TMP_DIR, shard_size_bytes=4 * 1024 * 1024)

    self.assertTrue(
        os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')),
        'weights_manifest.json does not exist')

    self.assertEqual(
        manifest,
        [{
            'paths': ['group1-shard1of1.bin'],
            'weights': [{
                'name': 'weight1',
                'shape': [2, 2],
                'dtype': 'string'
            }]
        }])

    weights_path = os.path.join(TMP_DIR, 'group1-shard1of1.bin')
    with open(weights_path, 'rb') as f:
      weight_bytes = f.read()

      self.assertEqual(len(weight_bytes), 36)
      # 'здраво'
      size = np.frombuffer(weight_bytes[:4], 'uint32')[0]
      self.assertEqual(size, 12)  # 6 cyrillic chars (2 bytes each).
      string = weight_bytes[4:16].decode('utf-8')
      self.assertEqual(string, u'здраво')
      # 'end'
      size = np.frombuffer(weight_bytes[16:20], 'uint32')[0]
      self.assertEqual(size, 3)  # 3 ascii chars.
      string = weight_bytes[20:23].decode('utf-8')
      self.assertEqual(string, u'end')
      # 'test'
      size = np.frombuffer(weight_bytes[23:27], 'uint32')[0]
      self.assertEqual(size, 4)  # 4 ascii chars.
      string = weight_bytes[27:31].decode('utf-8')
      self.assertEqual(string, u'test')
      # 'a'
      size = np.frombuffer(weight_bytes[31:35], 'uint32')[0]
      self.assertEqual(size, 1)  # 4 ascii chars.
      string = weight_bytes[35:36].decode('utf-8')
      self.assertEqual(string, u'a')


  def test_1_group_1_weight_string_empty(self):
    groups = [
        [{
            'name': 'weight1',
            'data': np.array([''], 'object')
        }]
    ]

    manifest = write_weights.write_weights(
        groups, TMP_DIR, shard_size_bytes=4 * 1024 * 1024)

    self.assertTrue(
        os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')),
        'weights_manifest.json does not exist')

    self.assertEqual(
        manifest,
        [{
            'paths': ['group1-shard1of1.bin'],
            'weights': [{
                'name': 'weight1',
                'shape': [1],
                'dtype': 'string'
            }]
        }])

    weights_path = os.path.join(TMP_DIR, 'group1-shard1of1.bin')
    with open(weights_path, 'rb') as f:
      weight_bytes = f.read()
      self.assertEqual(len(weight_bytes), 4)
      size = np.frombuffer(weight_bytes[:4], 'uint32')[0]
      self.assertEqual(size, 0)  # Empty string.

  def test_1_group_1_weight_string_unicode(self):
    groups = [
        [{
            'name': 'weight1',
            'data': np.array([[u'здраво', u'end'], [u'test', u'a']], 'object')
        }]
    ]

    manifest = write_weights.write_weights(
        groups, TMP_DIR, shard_size_bytes=4 * 1024 * 1024)

    self.assertTrue(
        os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')),
        'weights_manifest.json does not exist')

    self.assertEqual(
        manifest,
        [{
            'paths': ['group1-shard1of1.bin'],
            'weights': [{
                'name': 'weight1',
                'shape': [2, 2],
                'dtype': 'string'
            }]
        }])

    weights_path = os.path.join(TMP_DIR, 'group1-shard1of1.bin')
    with open(weights_path, 'rb') as f:
      weight_bytes = f.read()

      self.assertEqual(len(weight_bytes), 36)
      # 'здраво'
      size = np.frombuffer(weight_bytes[:4], 'uint32')[0]
      self.assertEqual(size, 12)  # 6 cyrillic chars (2 bytes each).
      string = weight_bytes[4:16].decode('utf-8')
      self.assertEqual(string, u'здраво')
      # 'end'
      size = np.frombuffer(weight_bytes[16:20], 'uint32')[0]
      self.assertEqual(size, 3)  # 3 ascii chars.
      string = weight_bytes[20:23].decode('utf-8')
      self.assertEqual(string, u'end')
      # 'test'
      size = np.frombuffer(weight_bytes[23:27], 'uint32')[0]
      self.assertEqual(size, 4)  # 4 ascii chars.
      string = weight_bytes[27:31].decode('utf-8')
      self.assertEqual(string, u'test')
      # 'a'
      size = np.frombuffer(weight_bytes[31:35], 'uint32')[0]
      self.assertEqual(size, 1)  # 4 ascii chars.
      string = weight_bytes[35:36].decode('utf-8')
      self.assertEqual(string, u'a')

  def test_1_group_1_weight_string_sharded(self):
    groups = [
        [{
            'name': 'weight1',
            'data': np.array(['helloworld'], 'object')
        }]
    ]

    # The array takes up 14 bytes across 3 shards when shard size is 5 bytes.
    manifest = write_weights.write_weights(
        groups, TMP_DIR, shard_size_bytes=5)

    self.assertTrue(
        os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')),
        'weights_manifest.json does not exist')

    self.assertEqual(
        manifest,
        [{
            'paths': [
                'group1-shard1of3.bin',
                'group1-shard2of3.bin',
                'group1-shard3of3.bin'
            ],
            'weights': [{
                'name': 'weight1',
                'shape': [1],
                'dtype': 'string'
            }]
        }])

    weight_bytes = bytes()
    with open(os.path.join(TMP_DIR, 'group1-shard1of3.bin'), 'rb') as f:
      weight_byt