# This is the VTA performance modeling part of the homework.
# The goal is to build a statistical performance model that takes in workload
# characteristics (e.g., shapes) along with simulator metrics to predict the
# real wall-clock running time of a workload.
# We provide a dataset collected by running variants of ResNet-18 operators on
# VTA to train and evaluate the model.a
# The file has two TODOs: first, choose a better set of features that improves
# the performance of the statistical model.
# Second, explain why the chosen features improved the predictive model.
# Finally, you can optionally use a different model to further improve
# prediction performance.

# NOTE: this homework depends on the numpy and scikit-learn packages; you can
# install these through a python package manager like pip.

import csv
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import GradientBoostingRegressor


# Mapping of feature names to data columns
feature_map = {'n': 0, # the convolution batch size
               'ic': 1, # the number of input channels
               'img_size': 2, # the input size (height/width)
               'oc': 3, # the number of output channels
               'kern_size': 4, # the kernel size
               'strides': 5, # the stride of the convolution
               'padding': 6, # the size of the padding to the convolution
               'dilation': 7, # # the dilation of the convolution
               'inp_load_bytes': 8, # simulator: the number of input bytes loaded
               'wgt_load_bytes': 9, # simulator: the number of weight bytes loaded
               'acc_load_bytes': 10, # simulator: the number of accum bytes loaded
               'uop_load_bytes': 11, # simulator: the number of microop bytes loaded
               'out_store_bytes': 12, # simulator: the number of output bytes stored
               'gemm_counter': 13, # simulator: number of gemm ops executed
               'alu_counter': 14} # simulator: number of alu ops executed
                 
# Load data from csv file
def load_data(filename, features):
    data = list()
    labels = list()
    full = list()
    with open(filename, 'r') as f:
        csvreader = csv.reader(f)
        for row in csvreader:
            selected = list()
            for feature in features:
                selected.append(row[feature_map[feature]])
            data.append(selected)
            full.append(row)
            labels.append(row[-1])
    return data, labels, full

def main():
    # TODO: 1. choose a better set of features to improve the performance
    # (accuracy) of the statistical model
    features = ['img_size']
    # TODO: 2. explain why the features you chose would improve the accuracy of
    # the statistical model (e.g., how they are relevant to or would impact
    # performance)

    # Load data
    data, labels, full  = load_data('vta_data.csv', features=features)
    # Split data into training and test sets
    tr_idx = int(len(data)*0.8)
    tr_data = np.array(data[:tr_idx], dtype='float32')
    tr_labels = np.array(labels[:tr_idx], dtype='float32')
    tr_full = np.array(full[:tr_idx], dtype='float32')
    t_data = np.array(data[tr_idx:], dtype='float32')
    t_labels = np.array(labels[tr_idx:], dtype='float32')
    t_full = np.array(full[tr_idx:], dtype='float32')
    # TODO 3. (optional) try using a different model to further improve
    # prediction accuracy
    reg = LinearRegression().fit(tr_data, tr_labels)
    prds = reg.predict(t_data)
    mse = np.mean((prds - t_labels)**2)
    rel_error = np.abs(prds - t_labels)/t_labels
    mrel_error = np.mean(rel_error)
    
    rel_error_res = list()
    for i in range(0, len(t_data)):
        rel_error_res.append((t_full[i][:7], rel_error[i]))
    sorted_rel_error = sorted(rel_error_res, key=lambda i: str(i[0]))
    # Report the mean squared error and relative error
    print("mse: ", mse)
    print("mean rel_error: ", "{0:.2f}%".format(mrel_error * 100))

    

if __name__ == '__main__':
    main()