import openpyxl

from prettytable import PrettyTable
from sklearn import linear_model

# Workday Compensation Excel Sheet
WD_DATA_FILE_PATH = '/home/pradeep/Research/avsal.xlsx'

# Data columns based on Workday format
DATE_COLUMN = 'A'
PAY_COLUMN = 'D'
REASON_COLUMN = 'C'
START_ROW = 3
END_ROW = 8
REASON_MAP = {
    'adjustment': 1,
    'profile': 2,
    'performance': 3,
}


# YEAR MONTH REASON
PREDICT_FOR = [
    [2017, 5, 'profile'],
    [2018, 2, 'performance'],
    [2019, 2, 'performance'],
    [2020, 2, 'performance'],
]


class Regressor(object):
    def __init__(self, train_data, train_labels, model='linear'):
        if model == 'linear':
            self.regr = linear_model.LinearRegression()
        self.train(train_data, train_labels)

    def train(self, train_data, train_labels):
        self.regr.fit(train_data, train_labels)

    def predict(self, input):
        return self.regr.predict(input)

    def score(self, train_data, train_labels):
        return self.regr.score(train_data, train_labels)


class ExcelParser(object):
    def __init__(self, file_path):
        self.wb = openpyxl.load_workbook(file_path)

    def get_sheet(self, sheet_name=None):
        if not sheet_name:
            sheet_name = self.wb.get_sheet_names()[0]
        return self.wb.get_sheet_by_name(sheet_name)

    def get_data(self, column, start_row, end_row, sheet=None, only_val=True):
        if not sheet:
            sheet = self.get_sheet()
        cells = sheet['{column}{start_row}:{column}{end_row}'.format(
            column=column, start_row=start_row, end_row=end_row)]
        if only_val:
            return [entry[0].value for entry in cells]
        return cells


def main():
    ex = ExcelParser(WD_DATA_FILE_PATH)
    date_cells = ex.get_data(DATE_COLUMN, START_ROW, END_ROW)
    pay_cells = ex.get_data(PAY_COLUMN, START_ROW, END_ROW)
    reason_cells = ex.get_data(REASON_COLUMN, START_ROW, END_ROW)

    # Training vector
    X = [[entry[0].year, entry[0].month, REASON_MAP[reason]]
         for entry in zip(date_cells, reason_cells)
         for reason in REASON_MAP
         if reason in entry[1].lower()]
    y = [float(entry) for entry in pay_cells]

    regr = Regressor(X, y)

    print "\nYour, salary prediction (Accuracy %s%%):\n" \
          % round(regr.score(X, y)*100, 2)

    X_test = [[entry[0], entry[1], REASON_MAP[entry[2]]]
              for entry in PREDICT_FOR]
    y_test = list(regr.predict(X_test))

    tb = PrettyTable(["Year", "Month", "Raise Type",
                      "Revised Salary", "Percent"])
    tb.add_row([
        PREDICT_FOR[0][0],
        PREDICT_FOR[0][1],
        PREDICT_FOR[0][2],
        round(y_test[0], 2),
        '-'])
    for i in range(1, len(y_test)):
        perc = round((y_test[i] - y_test[i-1])*100.0 / y_test[i-1], 2)
        tb.add_row([
            PREDICT_FOR[i][0],
            PREDICT_FOR[i][1],
            PREDICT_FOR[i][2],
            round(y_test[i], 2),
            perc])
    print tb


if __name__ == '__main__':
    main()
