from pathlib import Path

import tensorflow as tf

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


class TFMatmulBench(TFBase):
    def __init__(self, output_path: Path):
        super().__init__(output_path, Op.MATMUL)

    def experiment(self, experiment_args: tuple[int, int], length: int, dtype: tf.DType, device: tf.device):
        shape_1, shape_2 = experiment_args
        with device:
            tensor_1 = tf.ones(shape_1, dtype=dtype)
            tensor_2 = tf.ones(shape_2, dtype=dtype)

            for _ in range(length):
                _ = tensor_1 @ tensor_2

    def name(self, experiment_args: tuple[int, int]) -> str:
        shape_1, shape_2 = experiment_args
        return f'{shape_1[0]}x{shape_1[1]} @ {shape_2[0]}x{shape_2[1]}'

    def mop(self, experiment_args: tuple[int, int]) -> float:
        shape_1, shape_2 = experiment_args
        return (shape_1[0] * shape_2[1] / 1000_000) * 2 * (shape_1[1] - 1)

    def run(self,
            experiment_args: list[tuple[tuple[int, int], tuple[int, int]]],
            experiment_count: int,
            data_type: DataType):
        super().run(experiment_args, experiment_count, data_type)