# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# =============================================================================
"""Utils to lstm_test_data_generator.py that helps to generate the test data for lstm kernel (lstm_test_data.cc)"""

import numpy as np
from copy import deepcopy


def clip_range(vals, bit_width):
  """Mimic integer calculation.
  Clip the range of vals based on bit width.
  e.g., clip_range([300], 8) = [127] since int8 have range [-128, 127]
  Args:
      vals (np.array): float representation of the integer values
      bit_width (int): number of desired bits for vals
  Returns:
      np.array : clipped vals
  """
  # Numpy integer calculation does not do saturation. Implement here
  min_val = -2**(bit_width - 1)
  max_val = 2**(bit_width - 1) - 1
  if vals.max() > max_val or vals.min() < min_val:
    print(f"WARNING: integer overflow!")
  return np.clip(vals, min_val, max_val)


def quantize_data(data, scale, zero_point=0, bit_width=8):
  """Quantize the data to integer type with desired bit width.
  The quantized data is represented using float since integer calculation in
  numpy may differ from other implementations (e.g., no integer saturation
  protection in numpy)
  Args:
      data (np.array): float data
      scale (float): quantization scale of the data
      zero_point (integer): quantization zero point of the data
      bit_width (int): number of representative bits for vals
  Returns:
      np.array : quantized data in float but clipped range
  """
  vals = np.round(data / scale) + zero_point
  return clip_range(vals, bit_width)


def dequantize_data(quantized_data, scale, zero_point=0):
  """Dequantize the data to integer type with desired bit width.
  Args:
      quantized_data (np.array): quantized data
      scale (float): quantization scale of the data
      zero_point (integer): quantization zero point of the data
  Returns:
      np.array : dequantized data
  """
  return scale * (quantized_data - zero_point)


def rescale(data, effective_scale, zero_point, num_bits):
  """Rescale the data to the effective scale """
  # q = r/s + z
  rescaled = np.round(data * effective_scale) + zero_point
  return clip_range(rescaled, num_bits)


def calculate_scale(min_val, max_val, num_bits=8, symmetry=False):
  """Calculate quantization scale from the range and bit width"""
  num_bins = np.power(2, num_bits) - 1
  if symmetry:
    return max(abs(min_val), abs(max_val)) / int(num_bins / 2)
  return np.array((max_val - min_val) / num_bins, dtype=np.float32)


def calculate_zp(min_val, scale, num_bits=8):
  """Calculate the zero point from the minimal value"""
  quantized_floor = -np.power(2, num_bits) / 2
  return int(quantized_floor - min_val / scale)


def sigmoid(x):
  """Sigmoid (floating point)"""
  return 1 / (1 + np.exp(-x))


def quantized_sigmoid(input, input_scale, output_scale, num_bits=16):
  """Sigmoid (interger)"""
  float_input = input * input_scale
  float_result = sigmoid(float_input)
  return quantize_data(float_result, output_scale, bit_width=num_bits)


def quantized_tanh(input, input_scale, output_scale, num_bits=16):
  """Tanh (interger)"""
  float_input = input * input_scale
  float_result = np.tanh(float_input)
  return quantize_data(float_result, output_scale, bit_width=num_bits)


class QuantizedTensor:
  """Data structure for a quantized tensor"""

  def __init__(self, float_data, scale, zero_point, symmetry, num_bits=8):
    """Tensor is initialized using the floating point data"""
    self.float_data = float_data
    self.scale = scale
    self.zero_point = int(zero_point)
    self.symmetry = symmetry
    self.num_bits = num_bits
    self.quantized_data = quantize_data(float_data, scale, zero_point,
                                        num_bits)

  @property
  def dequantized_data(self):
    """Dequantize the quantized tensor data back to floating point"""
    return dequantize_data(self.quantized_data, self.scale,
                           self.zero_point).flatten()


class QuantizedGateParams:
  """Hold the quantization data and corresponding information for a LSTM gate (forget/input/cell/output gate) """

  def __init__(
      self,
      quantized_activation_weight,
      quantized_recurrent_weight,
      bias_data_float,
      shape_info,
      bias_num_bits=32,
      cell_num_bits=16,
      modulation=False,
  ):
    self.shape_info = shape_info
    self.activation_weight = quantized_activation_weight
    self.recurrent_weight = quantized_recurrent_weight
    self.bias_data_float = bias_data_float
    self.modulation = modulation
    self.bias_num_bits = bias_num_bits
    self.cell_num_bits = cell_num_bits
    # For INT16 cell state, the input scale is Q3.12
    self.nonlinear_input_scale = np.power(2.0, -(cell_num_bits - 4))
    # For INT16 cell state, the output scale is Q0.15
    self.nonlinear_output_scale = np.power(2.0, -(cell_num_bits - 1))

  def quantize_bias_data(self, input_scale):
    bias_scale = self.activation_weight.scale * input_scale
    return quantize_data(self.bias_data_float, bias_scale, 0,
                         self.bias_num_bits)

  def fold_zeropoint(self, weight, zero_point):
    # W*real = W*(quant-zero_pt) = Wquant - Wzero_pt
    # Wzero_pt is precomputed here as a constant (implemented in TFLM)
    zp_vector = zero_point * np.ones(shape=(self.shape_info['input_dim'], 1))
    zero_folded_vector = np.dot(weight, zp_vector)
    return -1 * clip_range(zero_folded_vector, self.bias_num_bits)

  def compute_activation_bias(self, input_scale, input_zp):
    # Wz is precomputed here and added it to the original bias (same scale)
    zero_folded_vector = self.fold_zeropoint(
        self.activation_weight.quantized_data, input_zp)
    quantized_bias = self.quantize_bias_data(input_scale)
    return zero_folded_vector + quantized_bias

  def compute_recurrent_bias(self, recurrent_zp):
    # Wz is precomputed here
    return self.fold_zeropoint(self.recurrent_weight.quantized_data,
                               recurrent_zp)

  def effective_activation_scale(self, input_scale):
    # Combine input scale with output scale. Used for fc calculation
    return (self.activation_weight.scale * input_scale /
            self.nonlinear_input_scale)

  def effective_recurrence_scale(self, recurrent_scale):
    # Combine input scale with output scale. Used for fc calculation
    return (self.recurrent_weight.scale * recurrent_scale /
            self.nonlinear_input_scale)


def assemble_quantized_tensor(float_data,
                              min_val,
                              max_val,
                              symmetry,
                              num_bits=8):
  """Create a QuantizedTensor using floating point data, range information, and bit width"""
  scale = calculate_scale(min_val, max_val, num_bits, symmetry)
  zp = 0
  if not symmetry:
    zp = calculate_zp(min_val, scale, num_bits)
  return QuantizedTensor(float_data,
                         scale,
                         zp,
                         symmetry=symmetry,
                         num_bits=num_bits)


def create_gate_params(gate_parameters, model_config, modulation=False):
  """Create a QuantizedGateParams using the gate paramater information and the model configuration"""
  shape_info = model_config['shape_info']
  quantization_settings = model_config['quantization_settings']

  activation_weight_data = np.array(
      gate_parameters['activation_weight_data']).reshape(
          (shape_info['input_dim'], shape_info['state_dim']))
  activation_weight = assemble_quantized_tensor(
      activation_weight_data,
      activation_weight_data.min(),
      activation_weight_data.max(),
      True,
      quantization_settings['weight_bits'],
  )

  recurrent_weight_data = np.array(
      gate_parameters['recurrent_weight_data']).reshape(
          (shape_info['input_dim'], shape_info['state_dim']))

  recurrent_weight = assemble_quantized_tensor(
      recurrent_weight_data,
      recurrent_weight_data.min(),
      recurrent_weight_data.max(),
      True,
      quantization_settings['weight_bits'],
  )

  bias_data_float = np.array(gate_parameters['bias_data']).reshape(
      (shape_info['input_dim'], 1))
  gate_params = QuantizedGateParams(
      activation_weight,
      recurrent_weight,
      bias_data_float,
      shape_info,
      bias_num_bits=quantization_settings['bias_bits'],
      cell_num_bits=quantization_settings['cell_bits'],
      modulation=modulation,
  )
  return gate_params


def gate_calculation(input, hidden_state, gate_params, debug=False):
  """
  A gate calculation is tanh(FC(activation, activation weight) + FC(recurrent, recurrent weight)). 
  For modulation gate, sigmoid is used instead of tanh.

  Note: for debugging purpose, floating point calculation is conducted in parallel with the integer calculation
  """
  # Quantized Version
  input_fc = np.dot(gate_params.activation_weight.quantized_data,
                    input.quantized_data)
  input_fc += gate_params.compute_activation_bias(input.scale,
                                                  input.zero_point)
  input_fc = rescale(input_fc,
                     gate_params.effective_activation_scale(input.scale), 0,
                     gate_params.cell_num_bits)
  recurrent_fc = np.dot(gate_params.recurrent_weight.quantized_data,
                        hidden_state.quantized_data)
  recurrent_fc += gate_params.compute_recurrent_bias(hidden_state.zero_point)
  recurrent_fc = rescale(
      recurrent_fc, gate_params.effective_recurrence_scale(hidden_state.scale),
      0, gate_params.cell_num_bits)

  before_activation = clip_range(input_fc + recurrent_fc,
                                 gate_params.cell_num_bits)

  # Float Version
  float_result = np.dot(gate_params.activation_weight.float_data,
                        input.float_data)
  float_result += np.dot(gate_params.recurrent_weight.float_data,
                         hidden_state.float_data)
  float_result += gate_params.bias_data_float

  if debug:
    print(f'input fc: {input_fc.flatten()}')
    print(f'recurrent fc: {recurrent_fc.flatten()}')

    dequantized_res = dequantize_data(before_activation,
                                      gate_params.nonlinear_input_scale)
    print(f'Intermediate before activation: {before_activation.flatten()}')
    print(f'dequantized :{dequantized_res.flatten()} ')
    print(f'float computation result: {float_result.flatten()} ')

    diff = dequantized_res - float_result
    print(f'diff percentage (%): {abs(diff/float_result).flatten()*100}')

  if gate_params.modulation:
    activated = quantized_tanh(before_activation,
                               gate_params.nonlinear_input_scale,
                               gate_params.nonlinear_output_scale,
                               gate_params.cell_num_bits)
    float_result = np.tanh(float_result)
  else:
    activated = quantized_sigmoid(before_activation,
                                  gate_params.nonlinear_input_scale,
                                  gate_params.nonlinear_output_scale,
                                  gate_params.cell_num_bits)
    float_result = sigmoid(float_result)

  if debug:
    dequantized_res = dequantize_data(activated,
                                      gate_params.nonlinear_output_scale)
    print(f'Gate result: {activated.flatten()} ')
    print(f'Dequantized: {dequantized_res.flatten()} ')
    print(f'float computation result: {float_result.flatten()} ')
    diff = dequantized_res - float_result
    print(f'diff percentage (%): {abs(diff/float_result).flatten()*100}')

  return activated, float_result


# The LSTM class
class QuantizedLSTMDebugger(object):
  """Help the debugging process of the LSTM kernel implementation by
  1. Exposing the kernel internal computation 
  2. Run floating point calculation in parallel with the integer version
  """

  def __init__(
      self,
      kernel_config,
      kernel_params,
      init_hidden_state_vals,
      hiddens_state_range,
      init_cell_state_vals,
      cell_state_range,
      cell_clip=8,
  ):
    self.kernel_config = kernel_config
    self.forget_gate_params = create_gate_params(
        kernel_params['forget_gate_data'], kernel_config)
    self.input_gate_params = create_gate_params(
        kernel_params['input_gate_data'], kernel_config)
    self.modulation_gate_params = create_gate_params(
        kernel_params['cell_gate_data'], kernel_config, modulation=True)
    self.output_gate_params = create_gate_params(
        kernel_params['output_gate_data'], kernel_config)
    self.quantization_settings = kernel_config['quantization_settings']

    self.hidden_state_tensor = assemble_quantized_tensor(
        np.array(init_hidden_state_vals).reshape((-1, 1)),
        hiddens_state_range[0],
        hiddens_state_range[1],
        False,
        self.quantization_settings['activation_bits'],
    )
    self.cell_state_tensor = assemble_quantized_tensor(
        np.array(init_cell_state_vals).reshape((-1, 1)),
        cell_state_range[0],
        cell_state_range[1],
        True,
        self.quantization_settings['cell_bits'],
    )

    self.quantized_cell_clip = quantize_data(
        cell_clip,
        self.cell_state_tensor.scale,
        self.cell_state_tensor.zero_point,
        self.quantization_settings['cell_bits'],
    )

  def invoke(self, input_tensor, debug=False):
    assert (
        input_tensor.num_bits == self.quantization_settings['activation_bits'])

    prev_hidden_state_tensor = deepcopy(self.hidden_state_tensor)
    prev_cell_state_tensor = deepcopy(self.cell_state_tensor)

    prev_hidden_state_float = prev_hidden_state_tensor.float_data
    prev_cell_state_float = prev_cell_state_tensor.float_data

    # forget gate
    forget_gate_quant, forget_gate_float = gate_calculation(
        input_tensor, prev_hidden_state_tensor, self.forget_gate_params)

    self.cell_state_tensor.quantized_data = rescale(
        prev_cell_state_tensor.quantized_data * forget_gate_quant,
        self.forget_gate_params.nonlinear_output_scale,
        0,
        self.quantization_settings['cell_bits'],
    )
    self.cell_state_tensor.float_data = (prev_cell_state_float *
                                         forget_gate_float)

    # input gate
    input_gate_quant, input_gate_float = gate_calculation(
        input_tensor, prev_hidden_state_tensor, self.input_gate_params)

    modulation_gate_quant, modulation_gate_float = gate_calculation(
        input_tensor, prev_hidden_state_tensor, self.modulation_gate_params)

    gated_input_quant = rescale(
        input_gate_quant * modulation_gate_quant,
        self._calculate_effective_cell_scale(),
        0,
        self.quantization_settings['cell_bits'],
    )
    gated_input_float = input_gate_float * modulation_gate_float

    if (
        debug
    ):  # Hidden/cell state will be updated, break up the debug to record the intermediate state
      print('======================One Step LSTM======================')
      print('###### Forget Gate Output: ######')
      print(f'Quantized: {forget_gate_quant.flatten()}')
      dequantized_val = dequantize_data(
          forget_gate_quant, self.forget_gate_params.nonlinear_output_scale, 0)
      print(f'Dequantized : {dequantized_val.flatten()}')
      print(f'Float : {forget_gate_float.flatten()}')

      print('###### Cell state after forgetting: ######')
      print(f'Quantized: {self.cell_state_tensor.quantized_data.flatten()}')
      print(
          f'Dequantized: {self.cell_state_tensor.dequantized_data.flatten()}')
      print(f'Float : {self.cell_state_tensor.float_data.flatten()}')

      print('###### Input gate output: ######')
      print(f'Quantized: {input_gate_quant.flatten()}')
      dequantized_val = dequantize_data(
          input_gate_quant, self.input_gate_params.nonlinear_output_scale, 0)
      print(f'Dequantized: {dequantized_val.flatten()}')
      print(f'Float : {input_gate_float.flatten()}')

      print('###### cell gate output: ######')
      print(f'Quantized: {modulation_gate_quant.flatten()}')
      dequantized_val = dequantize_data(
          modulation_gate_quant,
          self.modulation_gate_params.nonlinear_output_scale,
          0,
      )
      print(f'Dequantized: {dequantized_val.flatten()}')
      print(f'Float : {modulation_gate_float.flatten()}')

      print('###### Gated input (input_gate * cell_gate): ######')
      print(f'Quantized: {gated_input_quant.flatten()}')
      dequantized_val = dequantize_data(gated_input_quant,
                                        self.cell_state_tensor.scale, 0)
      print(f'Dequantized: {dequantized_val.flatten()}')
      print(f'Float : {gated_input_float.flatten()}')

    # Update the cell state
    self.cell_state_tensor.quantized_data += gated_input_quant
    self._apply_cell_clip()
    self.cell_state_tensor.float_data += gated_input_float

    # output gate
    output_gate_quant, output_gate_float = gate_calculation(
        input_tensor, prev_hidden_state_tensor, self.output_gate_params)

    # Update the hidden state
    transformed_cell_quant = quantized_tanh(
        self.cell_state_tensor.quantized_data,
        self.output_gate_params.nonlinear_input_scale,
        self.output_gate_params.nonlinear_output_scale,
        self.cell_state_tensor.num_bits,
    )

    transformed_cell_float = np.tanh(self.cell_state_tensor.float_data)

    gated_output_quant = rescale(
        output_gate_quant * transformed_cell_quant,
        self._calculate_effective_output_scale(),
        self.hidden_state_tensor.zero_point,
        self.hidden_state_tensor.num_bits,
    )
    gated_output_float = output_gate_float * transformed_cell_float

    self.hidden_state_tensor.quantized_data = gated_output_quant
    self.hidden_state_tensor.float_data = gated_output_float

    if debug:
      print('###### Updated cell state): ######')
      print(f'Quantized: {self.cell_state_tensor.quantized_data.flatten()}')
      print(
          f'Dequantized: {self.cell_state_tensor.dequantized_data.flatten()}')
      print(f'Float : {self.cell_state_tensor.float_data.flatten()}')

      print('###### Output gate: ######')
      print(f'Quantized : {output_gate_quant.flatten()}')
      dequantized_val = dequantize_data(
          output_gate_quant, self.output_gate_params.nonlinear_output_scale, 0)
      print(f'Dequantized: {dequantized_val.flatten()}')
      print(f'Float : {output_gate_float.flatten()}')

      print('###### Tanh transformed cell: ######')
      print(f'Quantized: {transformed_cell_quant.flatten()}')
      dequantized_val = dequantize_data(
          transformed_cell_quant,
          self.output_gate_params.nonlinear_output_scale,
          0,
      )
      print(f'Dequantized: {dequantized_val.flatten()}')
      print(f'Float : {transformed_cell_float.flatten()}')

      print('###### Updated hidden state: ######')
      print(f'Quantized: {gated_output_quant.flatten()}')
      print(
          f'Dequantized: {self.hidden_state_tensor.dequantized_data.flatten()}'
      )
      print(f'Float : {gated_output_float.flatten()}')

      diff = abs(self.hidden_state_tensor.dequantized_data -
                 gated_output_float.flatten())
      max_diff_perc = diff / gated_output_float.flatten() * 100
      print(f'Max diff perc (%): {max_diff_perc}')
    return gated_output_quant, gated_output_float

  def _calculate_effective_output_scale(self):
    return (self.output_gate_params.nonlinear_output_scale *
            self.modulation_gate_params.nonlinear_output_scale /
            self.hidden_state_tensor.scale)

  def _calculate_effective_cell_scale(self):
    return (self.input_gate_params.nonlinear_output_scale *
            self.modulation_gate_params.nonlinear_output_scale /
            self.cell_state_tensor.scale)

  def _apply_cell_clip(self):
    cell_vals = self.cell_state_tensor.quantized_data
    if (cell_vals.max() > self.quantized_cell_clip
        or cell_vals.min() < -self.quantized_cell_clip):
      print(f'WARNING: cell values clip to {self.quantized_cell_clip}!')

    self.cell_state_tensor.quantized_data = np.round(
        np.clip(cell_vals, -self.quantized_cell_clip,
                self.quantized_cell_clip))
