import sys
sys.path.append("./src/")

from hw_generator.generator_conv import CONVGenerator 
from hw_generator.generator_gemm import GEMMGenerator 

from benchmark.benchmark import Workload, Benchmark
from codesign.hw_evaluation import evaluation_function_model, gen_software, reset

from benchmark.computations import conv2d_compute

if __name__ == '__main__':


    dtype = "int8"
    
    print("Testing accelerators with GEMM intrinsic ...")
    generator = GEMMGenerator() 
    parameterization = {'x':16, 'sp_capacity':512, 'local_capacity':1024, 
                        'sp_banks':4, 'dma_buswidth':256, 'dma_maxbytes':128, 
                        'dataflow': "OS", 'dtype':"int8"}
    benchmark = Benchmark("conv_test")
    conv1 = Workload("conv_1", "CONV", conv2d_compute, (1, 128, 32, 256, 32, 3, 3, 1, dtype, generator.type))  
    conv2 = Workload("conv_2", "CONV", conv2d_compute, (1, 1, 64, 64, 256, 1, 1, 1, dtype, generator.type))  
    benchmark.add(conv1) 
    benchmark.add(conv2) 
    evaluation_function_model(parameterization, benchmark, generator)
    arg_str = [str(i) for i in parameterization.values()]
    tag = '_'.join(arg_str)
    software = gen_software(tag, is_print=True)
    print("Passed.")

    print("Testing accelerators with CONV intrinsic ...")
    reset()
    generator = CONVGenerator() 
    parameterization = {'x':4, 'y':4, 'sp_capacity':2048, 'local_capacity':1024, 
                            'sp_banks':4, 'dma_buswidth':256, 'dma_maxbytes':128, 
                            'dataflow': "FIXED", 'dtype':"int8"}
    benchmark = Benchmark("conv_test")
    conv1 = Workload("conv_1", "CONV", conv2d_compute, (1, 128, 32, 256, 32, 3, 3, 1, dtype, generator.type))  
    conv2 = Workload("conv_2", "CONV", conv2d_compute, (1, 256, 64, 64, 256, 1, 1, 1, dtype, generator.type))  
    benchmark.add(conv1) 
    benchmark.add(conv2) 
    evaluation_function_model(parameterization, benchmark, generator)
    arg_str = [str(i) for i in parameterization.values()]
    tag = '_'.join(arg_str)
    software = gen_software(tag, is_print=True)
    print("Passed.")
  
    print("Testing accelerators with GEMV intrinsic ...")
    reset()
    generator = GEMVGenerator() 
    parameterization = {'x':16, 'sp_capacity':512, 'local_capacity':1024, 
                            'sp_banks':4, 'dma_buswidth':256, 'dma_maxbytes':128, 
                            'dataflow': "FIXED", 'dtype':"int8"}
    benchmark = Benchmark("conv_test")
    conv1 = Workload("conv_1", "CONV", conv2d_compute, (1, 128, 32, 256, 32, 3, 3, 1, dtype, generator.type))  
    conv2 = Workload("conv_2", "CONV", conv2d_compute, (1, 256, 64, 64, 256, 1, 1, 1, dtype, generator.type))  
    benchmark.add(conv1) 
    benchmark.add(conv2) 
    evaluation_function_model(parameterization, benchmark, generator)
    arg_str = [str(i) for i in parameterization.values()]
    tag = '_'.join(arg_str)
    software = gen_software(tag, is_print=True)
    
    print("Passed.")

    print("Testing accelerators with DOT intrinsic ...")
    reset()
    generator = DOTGenerator() 
    parameterization = {'x':16, 'y':8, 'sp_capacity':512, 'local_capacity':1024, 
                            'sp_banks':4, 'dma_buswidth':256, 'dma_maxbytes':128, 
                            'dataflow': "FIXED", 'dtype':"int8"}     
    benchmark = Benchmark("conv_test")
    conv1 = Workload("conv_1", "CONV", conv2d_compute, (1, 128, 32, 256, 32, 3, 3, 1, dtype, generator.type))  
    conv2 = Workload("conv_2", "CONV", conv2d_compute, (1, 256, 64, 64, 256, 1, 1, 1, dtype, generator.type))  
    benchmark.add(conv1) 
    benchmark.add(conv2) 
    evaluation_function_model(parameterization, benchmark, generator)
    arg_str = [str(i) for i in parameterization.values()]
    tag = '_'.join(arg_str)
    software = gen_software(tag, is_print=True)
    print("Passed.")


