from collections import OrderedDict
from archx.utils import get_prod

def input_gemm(architecture_dict: OrderedDict, workload_dict: OrderedDict=None)->OrderedDict:
    performance_dict = OrderedDict()
    router_dim = get_prod(architecture_dict['irouter']['instance']) if 'irouter' in architecture_dict else 1

    frequency = architecture_dict['ififo']['query']['frequency']
    ififo_dim = get_prod(architecture_dict['ififo']['instance'][-2:])
    input_register_dim = get_prod(architecture_dict['input_register']['instance'][-3:])

    cycle_count = 1 / router_dim
    performance_dict['cycle_count'] = OrderedDict({'value': cycle_count, 'unit': 'cycle'})
    performance_dict['runtime'] = OrderedDict({'value': cycle_count / 1000 / frequency, 'unit': 'ms'})
    
    ififo_dict = OrderedDict({'count': ififo_dim})
    input_register_dict = OrderedDict({'count': input_register_dim})

    performance_dict['subevent'] = OrderedDict({'ififo': ififo_dict,
                                                'input_register': input_register_dict})
    return performance_dict

def weight_gemm(architecture_dict: OrderedDict, workload_dict: OrderedDict=None)->OrderedDict:
    performance_dict = OrderedDict()
    router_dim = get_prod(architecture_dict['irouter']['instance']) if 'irouter' in architecture_dict else 1

    frequency = architecture_dict['wfifo']['query']['frequency']
    wfifo_dim = get_prod(architecture_dict['wfifo']['instance'][-2:])
    weight_register_dim = get_prod(architecture_dict['weight_register']['instance'][-3:])

    if 'int_to_fp' in architecture_dict:
        int_to_fp_dim = get_prod(architecture_dict['int_to_fp']['instance'][-2:])

    cycle_count = 1 / router_dim
    performance_dict['cycle_count'] = OrderedDict({'value': cycle_count, 'unit': 'cycle'})
    performance_dict['runtime'] = OrderedDict({'value': cycle_count / 1000 / frequency, 'unit': 'ms'})
    
    wfifo_dict = OrderedDict({'count': wfifo_dim})
    weight_register_dict = OrderedDict({'count': weight_register_dim})
    if 'int_to_fp' in architecture_dict:
        int_to_fp_dict = OrderedDict({'count': int_to_fp_dim})

    performance_dict['subevent'] = OrderedDict({'wfifo': wfifo_dict,
                                                'weight_register': weight_register_dict})
    if 'int_to_fp' in architecture_dict:
        performance_dict['subevent']['int_to_fp'] = int_to_fp_dict
    return performance_dict

def array_gemm(architecture_dict: OrderedDict, workload_dict: OrderedDict=None)->OrderedDict:
    performance_dict = OrderedDict()
    router_dim = get_prod(architecture_dict['irouter']['instance']) if 'irouter' in architecture_dict else 1

    frequency = architecture_dict['multiplier']['query']['frequency']

    multiplier_dim = get_prod(architecture_dict['multiplier']['instance'][-3:])
    pe_register_dim = get_prod(architecture_dict['pe_register']['instance'][-3:])
    accumulator_dim = get_prod(architecture_dict['accumulator']['instance'][-3:])
    accumulator_register_dim = get_prod(architecture_dict['accumulator_register']['instance'][-2:])
    adder_dim = get_prod(architecture_dict['adder']['instance'][-2:])
    ofifo_dim = get_prod(architecture_dict['ofifo']['instance'][-2:])

    if 'ch_aloc' in architecture_dict:
        ch_aloc_dim = get_prod(architecture_dict['ch_aloc']['instance'][-2:])
    if 'ch_dealoc' in architecture_dict:
        ch_dealoc_dim = get_prod(architecture_dict['ch_dealoc']['instance'][-2:])
    if 'int_to_fp_figna' in architecture_dict:
        int_to_fp_figna_dim = get_prod(architecture_dict['int_to_fp_figna']['instance'][-2:])
    if 'prealigner' in architecture_dict:
        prealigner_dim = get_prod(architecture_dict['prealigner']['instance'][-2:])

    cycle_count = 1 / router_dim
    performance_dict['cycle_count'] = OrderedDict({'value': cycle_count, 'unit': 'cycle'})
    performance_dict['runtime'] = OrderedDict({'value': cycle_count / 1000 / frequency, 'unit': 'ms'})
    
    multiplier_dict = OrderedDict({'count': multiplier_dim})
    pe_register_dict = OrderedDict({'count': pe_register_dim})
    accumulator_dict = OrderedDict({'count': accumulator_dim})
    accumulator_register_dict = OrderedDict({'count': accumulator_register_dim * cycle_count})
    adder_dict = OrderedDict({'count': adder_dim * cycle_count})
    ofifo_dict = OrderedDict({'count': ofifo_dim * cycle_count})

    if 'ch_aloc' in architecture_dict:
        ch_aloc_dict = OrderedDict({'count': ch_aloc_dim})
    if 'ch_dealoc' in architecture_dict:
        ch_dealoc_dict = OrderedDict({'count': ch_dealoc_dim})
    if 'int_to_fp_figna' in architecture_dict:
        int_to_fp_figna_dict = OrderedDict({'count': int_to_fp_figna_dim})
    if 'prealigner' in architecture_dict:
        prealigner_dict = OrderedDict({'count': prealigner_dim})

    performance_dict['subevent'] = OrderedDict({'multiplier': multiplier_dict,
                                                'pe_register': pe_register_dict,
                                                'accumulator': accumulator_dict,
                                                'accumulator_register': accumulator_register_dict,
                                                'adder': adder_dict,
                                                'ofifo': ofifo_dict})
    if 'ch_aloc' in architecture_dict:
        performance_dict['subevent']['ch_aloc'] = ch_aloc_dict
    if 'ch_dealoc' in architecture_dict:
        performance_dict['subevent']['ch_dealoc'] = ch_dealoc_dict
    if 'int_to_fp_figna' in architecture_dict:
        performance_dict['subevent']['int_to_fp_figna'] = int_to_fp_figna_dict
    if 'prealigner' in architecture_dict:
        performance_dict['subevent']['prealigner'] = prealigner_dict

    return performance_dict

def vector_gemm(architecture_dict: OrderedDict, workload_dict: OrderedDict=None)->OrderedDict:
    performance_dict = OrderedDict()
    router_dim = get_prod(architecture_dict['irouter']['instance']) if 'irouter' in architecture_dict else 1

    frequency = architecture_dict['multiplier_vector']['query']['frequency']
    multiplier_vector_dim = get_prod(architecture_dict['multiplier_vector']['instance'][-2:])
    register_vector_dim = get_prod(architecture_dict['register_vector']['instance'][-2:])

    cycle_count = 1 / router_dim
    performance_dict['cycle_count'] = OrderedDict({'value': cycle_count, 'unit': 'cycle'})
    performance_dict['runtime'] = OrderedDict({'value': cycle_count / 1000 / frequency, 'unit': 'ms'})
    
    multiplier_vector_dict = OrderedDict({'count': multiplier_vector_dim})
    accumulator_vector_dict = OrderedDict({'count': 0})
    register_vector_dict = OrderedDict({'count': register_vector_dim})
    mac_register_vector_dict = OrderedDict({'count': 0})

    performance_dict['subevent'] = OrderedDict({'multiplier_vector': multiplier_vector_dict,
                                                'accumulator_vector': accumulator_vector_dict,
                                                'register_vector': register_vector_dict,
                                                'mac_register_vector': mac_register_vector_dict})
    return performance_dict