# -*- coding: utf-8 -*-
"""
Created on Thu Jul  9 18:03:39 2020

@author: akanksha
"""

import argparse
import warnings
from functools import reduce

import joblib
import pandas as pd
import sklearn
from pandas.core.common import SettingWithCopyWarning

warnings.simplefilter(action="ignore", category=SettingWithCopyWarning)

parser = argparse.ArgumentParser(description='Prediction from combined models for the reads.')

parser.add_argument('--methodsfile', '-i', type=str, required=True,
                    help='TSV file containing name and path of the method output tsv file. The output tsv file from the method should be in the format [ID,Pos,Strand,Score]. Can be compressed in gz.')

parser.add_argument('--model', '-m', choices=["default", "optimized"], required=True, type=str,
                    help='which model to select from default RF or optimized RF with max_depth 3 and n_estimator 10')

parser.add_argument('--output', '-o', type=str, required=True,
                    help='Where to store the outputs')

parser.add_argument('--model-base-dir', '-b', type=str, required=True,
                    help='Where is the model dir', default='/projects/li-lab/yang/test/METEORE')

options = parser.parse_args()


def mod_file(data_file_path):
    data_file = pd.read_csv(data_file_path, header=0, sep="\t")
    name = data_file_path.split("/")[-1].split(".")[0]  # HL60_DeepSignal-METEORE-perRead-score
    data_file.drop_duplicates(subset=['Chr', "ID", "Pos", "Strand"], inplace=True)  # add chr
    data_file.reset_index(inplace=True, drop=True)

    df_id_strand = data_file[["ID", "Strand"]]
    df_id_strand.drop_duplicates(subset=["ID", "Strand"], inplace=True)

    # mask = data_file.index[data_file.Strand == "-"].tolist() # select negative strand
    # data_file["Pos"][mask] = data_file["Pos"][mask] - 1 # merge to positive strand
    data_file.drop(["Strand"], axis=1, inplace=True)
    data_file.rename(columns={"Score": name}, inplace=True)
    data_file.reset_index(inplace=True, drop=True)

    # print(data_file)
    return (data_file, df_id_strand)


def main(mp, combine_file, combine_id_strand):
    loaded_model = joblib.load(open(mp, 'rb'))
    print(loaded_model)

    combine_file.dropna(inplace=True)
    X = combine_file[combine_file.columns[3:]]  # 2:
    X = sklearn.preprocessing.MinMaxScaler().fit_transform(X)
    prediction = pd.DataFrame(loaded_model.predict(X))  ##
    prediction_prob = pd.DataFrame(loaded_model.predict_proba(X))
    prediction.rename(columns={0: "Prediction"}, inplace=True)
    prediction_prob = prediction_prob[[1]]
    prediction_prob.rename(columns={1: "Prob_methylation"}, inplace=True)
    final_output = pd.concat([combine_file[combine_file.columns[:3]], prediction, prediction_prob], axis=1)  #:2

    final_output.dropna(inplace=True)
    import numpy as np
    final_output['Pos'] = final_output['Pos'].astype(np.int64)
    final_output['Prediction'] = final_output['Prediction'].astype(int)
    outdf = final_output.merge(combine_id_strand, how='left', on='ID')

    outdf.to_csv(options.output, header=True, index=None, sep='\t')
    print(f"save to {options.output}")


if __name__ == '__main__':
    # Used to find the randome forest model in Github
    METEORE_Dir = options.model_base_dir
    df_file = pd.read_csv(options.methodsfile, header=None, sep='\t')
    if options.model == "default":
        fillval = "default"
    else:
        fillval = "max_depth_3_n_estimator_10"
    modelname = '_'.join(df_file[0])
    mp = f'{METEORE_Dir}/saved_models/rf_model_' + fillval + '_' + modelname + '.model'
    dfs = []
    dfs_id_strand = []
    for i in df_file[1]:
        df, df_id_strand = mod_file(i)
        dfs.append(df)
        dfs_id_strand.append(df_id_strand)
    combine_file = reduce(lambda left, right: pd.merge(left, right, how='inner', on=["ID", "Chr", "Pos"]),
                          dfs)  # add chr
    combine_file.drop_duplicates(subset=["ID", "Chr", "Pos"], inplace=True)  # add chr
    combine_file.reset_index(inplace=True, drop=True)
    # print(combine_file)

    # outfn = options.output.replace('-perRead.combine.tsv.gz', '-intermediate.combinedf.tsv.gz')
    # combine_file.to_csv(outfn, sep='\t')
    #
    # print(f"save to {outfn}")

    # Used for add strand later
    combine_id_strand = pd.concat(dfs_id_strand)
    combine_id_strand.drop_duplicates(inplace=True)
    # print(combine_id_strand)
    if len(combine_file) > 0:
        main(mp, combine_file, combine_id_strand)  ##
