from pathlib import Path
from typing import List, Tuple

import tensorflow.compat.v1 as tf

from src.common import DataType, Op
from src.tf_1.base import TFBase


class TFNNDenseX5Bench(TFBase):
    def __init__(self, output_path: Path, data_type: DataType):
        super().__init__(output_path, Op.NN_DENSE_X5, data_type)
        self.dense_op = None

    def pre_experiment(self, experiment_args: Tuple[int, int]):
        super().pre_experiment(experiment_args)
        batch_size, dimension = experiment_args
        input_tensor = tf.get_variable('input_tensor', shape=(batch_size, dimension), dtype=self.dtype,
                                       initializer=tf.initializers.ones, trainable=False)
        output_tensor = input_tensor
        for layer in range(5):
            weights = tf.get_variable(f'Weights_{layer}', shape=(dimension, dimension), dtype=self.dtype,
                                      initializer=tf.initializers.ones, trainable=False)
            biases = tf.get_variable(f'Biases_{layer}', shape=dimension, dtype=self.dtype,
                                     initializer=tf.initializers.ones, trainable=False)
            output_tensor = tf.matmul(output_tensor, weights) + biases
        self.dense_op = output_tensor

        self.session.run(tf.initializers.global_variables())

    def experiment(self):
        self.session.run(self.dense_op)

    def run(self, experiment_args: List[Tuple[int, int]], experiment_count: int):
        super().run(experiment_args, experiment_count)