import pickle
import numpy as np
import pandas as pd
import argparse


def list_of_columns(pdata, rdata):
    """
    Using R and Python files, make a list of columns that will be used
    for comparison.
    """
    r_cols = set(rdata.columns)
    print('Number of columns in R file...{0}'.format(len(r_cols)))

    p_cols = set(pdata.columns)
    print('Number of columns in Python file...{0}'.format(len(p_cols)))

    if (len(p_cols) == len(r_cols)):
        selected_cols = []
        for col in p_cols:
            if (isinstance(col, int)):
                if ('pp.' + str(col) in r_cols):
                    selected_cols.append(col)
            elif (col in r_cols):
                selected_cols.append(col)
            else:
                print('Unmatched column...{0}'.format(col))
        print('Total common column...{0}'.format(len(selected_cols)))
    else:
        print("Python and R file do not have the same number of features")
        return 0

    # if selected columns is less than total columns, something is wrong.
    if (len(selected_cols) == len(p_cols)):
        return selected_cols
    else:
        print("Python and R files have different columns")
        return 0


def check_feature_difference(cols, pdata, rdata, dp):
    """
    Find if R and Python files have different features
    """
    matched_cols = set()
    unmatched_cols = set()
    for col in cols:
        if isinstance(col, int):
            r_col_data = set([round(float(v), dp) for v in rdata['pp.' + str(col)]])
            p_col_data = set([round(float(v), dp) for v in pdata[col]])
        else:
            r_col_data = set([round(float(v), dp) for v in rdata[col]])
            p_col_data = set([round(float(v), dp) for v in pdata[col]])

        unmatched = len(r_col_data.difference(p_col_data))
        if (unmatched == 0):
            matched_cols.add(col)
        else:
            unmatched_cols.add(col)

    return matched_cols, unmatched_cols


if __name__ == "__main__":
    """
    start of the code
    """
    dp = 2  # decimal places
    parser = argparse.ArgumentParser(description='Compare features generated by R and Python')
    parser.add_argument('--pythonfile', required=True, help='Features pickle file generated by Python')
    parser.add_argument('--rfile', required=True, help='Features pickle file generated from RDS file')
    parser.add_argument('--decimalplace', default=dp, type=int,
                        help='Decimal places of features to select for comparison')
    parser.add_argument("-v", "--verbose", action="count", default=0, help="verbosity")

    args = parser.parse_args()

    # load R and Python data
    pdata = pickle.load(open(args.pythonfile, 'rb'))
    rdata = pickle.load(open(args.rfile, 'rb'))

    # Processing starts here for training set
    selected_cols = list_of_columns(pdata, rdata)
    if selected_cols == 0:
        print("Program failed!!!")
    else:
        matched_cols, unmatched_cols = check_feature_difference(selected_cols, pdata, rdata, args.decimalplace)
        print("Number of features with the same values in R and Python files...{0}".format(len(matched_cols)))
        print("Number of features with the different values in R and Python files...{0}".format(len(unmatched_cols)))
        print("Features with different values in R and Python files:")
        print(unmatched_cols)
