import numpy as np
import pandas as pd

from CommonServerPython import *

ALL_LABELS = "*"
PREDICTIONS_OUT_FILE_NAME = 'predictions.csv'


def canonize_label(label):
    return label.replace(" ", "_")


def get_phishing_map_labels(comma_values):
    if comma_values == ALL_LABELS:
        return comma_values
    values = [x.strip() for x in comma_values.split(",")]
    labels_dict = {}
    for v in values:
        v = v.strip()
        if ":" in v:
            splited = v.rsplit(":", maxsplit=1)
            labels_dict[splited[0].strip()] = splited[1].strip()
        else:
            labels_dict[v] = v
    if len(set(labels_dict.values())) == 1:
        mapped_value = list(labels_dict.values())[0]
        error = ['Label mapping error: you need to map to at least two labels: {}.'.format(mapped_value)]
        return_error('\n'.join(error))
    return {k: canonize_label(v) for k, v in labels_dict.items()}


def get_data_with_mapped_label(y_true_list, labels_mapping):
    mapped_y_true = []
    relevant_indices = []
    for i, y_true in enumerate(y_true_list):
        if labels_mapping == ALL_LABELS:
            mapped_y_true.append(canonize_label(y_true))
            relevant_indices.append(i)
        elif y_true in labels_mapping:
            mapped_y_true.append(canonize_label(labels_mapping[y_true]))
            relevant_indices.append(i)
        else:
            continue
    return mapped_y_true, relevant_indices


def get_ml_model_evaluation(y_test, y_pred, target_accuracy, target_recall, detailed=False):
    res = demisto.executeCommand('GetMLModelEvaluation', {'yTrue': json.dumps(y_test),
                                                          'yPred': json.dumps(y_pred),
                                                          'targetPrecision': str(target_accuracy),
                                                          'targetRecall': str(target_recall),
                                                          'detailedOutput': 'true' if detailed else 'false'
                                                          })
    if is_error(res):
        return_error(get_error(res))
    return res


def output_model_evaluation(y_test, y_pred, res, context_field, human_readable_title=None):
    threshold = float(res[0]['Contents']['threshold'])
    confusion_matrix = json.loads(res[0]['Contents']['csr_matrix_at_threshold'])
    metrics_df = json.loads(res[0]['Contents']['metrics_df'])
    human_readable = res[0]['HumanReadable']
    if human_readable_title is not None:
        human_readable = '\n'.join([human_readable_title, human_readable])
    result_entry = {
        'Type': entryTypes['note'],
        'Contents': {'Threshold': threshold, 'ConfusionMatrixAtThreshold': confusion_matrix,
                     'Metrics': metrics_df, 'YTrue': y_test, 'YPred': y_pred},
        'ContentsFormat': formats['json'],
        'HumanReadable': human_readable,
        'HumanReadableFormat': formats['markdown'],
        'EntryContext': {
            context_field: {
                'EvaluationScores': metrics_df,
                'ConfusionMatrix': confusion_matrix,
            }
        }
    }
    demisto.results(result_entry)
    return confusion_matrix


def return_file_result_with_predictions_on_test_set(data, y_true, y_pred, y_pred_prob, additional_fields):
    predictions_data = {}
    for field in additional_fields:
        predictions_data[field] = [i.get(field, '') for i in data]
    predictions_data['y_true'] = y_true
    predictions_data['y_pred'] = y_pred
    predictions_data['y_pred_prob'] = y_pred_prob
    df = pd.DataFrame(predictions_data)
    non_empty_columns = [field for field in additional_fields if df[field].astype(bool).any()]
    csv_df = df.to_csv(columns=['y_true', 'y_pred', 'y_pred_prob'] + non_empty_columns, encoding='utf-8')
    demisto.results(fileResult(PREDICTIONS_OUT_FILE_NAME, csv_df))


def main(incident_types, incident_query, y_true_field, y_pred_field, y_pred_prob_field, model_target_accuracy,
         labels_mapping, additional_fields):
    non_empty_fields = '{},{}'.format(y_true_field.strip(), y_pred_field.strip())
    incidents_query_args = {'incidentTypes': incident_types,
                            'NonEmptyFields': non_empty_fields,
                            }
    if incident_query is not None:
        incidents_query_args['query'] = incident_query
    incidents_query_res = demisto.executeCommand('GetIncidentsByQuery', incidents_query_args)
    if is_error(incidents_query_res):
        return_error(get_error(incidents_query_res))
    incidents = json.loads(incidents_query_res[-1]['Contents'])
    demisto.results('Found {} incidents'.format(len(incidents)))
    y_true = []
    y_pred = []
    y_pred_prob = []
    incidents_with_missing_pred_prob = 0
    for i in incidents:
        y_true.append(i[y_true_field])
        y_pred.append(i[y_pred_field])
        if y_pred_prob_field not in i:
            incidents_with_missing_pred_prob += 1
        y_pred_prob.append(i.get(y_pred_prob_field, None))
    y_true, relevant_indices = get_data_with_mapped_label(y_true, labels_mapping)
    y_pred = [y_pred[i] for i in relevant_indices]
    y_pred_prob = [y_pred_prob[i] for i in relevant_indices]
    incidents = [incidents[i] for i in relevant_indices]
    y_pred_prob_is_given = incidents_with_missing_pred_prob == 0
    if y_pred_prob_is_given:
        y_pred_dict = [{label: prob} for label, prob in zip(y_pred, y_pred_prob)]
    else:
        y_pred_dict = [{label: 1.0} for label in y_pred]
    if y_pred_prob_is_given:
        res_threshold = get_ml_model_evaluation(y_true, y_pred_dict, model_target_accuracy, target_recall=0,
                                                detailed=True)
        # show results for the threshold found - last result so it will appear first
        output_model_evaluation(y_test=y_true, y_pred=y_pred_dict, res=res_threshold,
                                context_field='EvaluateMLModllAtProduction')
    # show results if no threshold (threhsold=0) was used. Following code is reached only if a legal thresh was found:
    if not y_pred_prob_is_given or not np.isclose(float(res_threshold[0]['Contents']['threshold']), 0):
        res = get_ml_model_evaluation(y_true, y_pred_dict, target_accuracy=0, target_recall=0)
        human_readable = '\n'.join(['## Results for No Threshold',
                                    'The following results were achieved by using no threshold (threshold equals 0)'])
        output_model_evaluation(y_test=y_true, y_pred=y_pred_dict, res=res,
                                context_field='EvaluateMLModllAtProductionNoThresh',
                                human_readable_title=human_readable)
    return_file_result_with_predictions_on_test_set(incidents, y_true, y_pred, y_pred_prob, additional_fields)


model_target_accuracy = demisto.args().get('modelTargetAccuracy', 0)
incident_types = demisto.args()['incidentTypes']
incident_query = demisto.args().get('incidentsQuery', None)
y_true_field = demisto.args()['emailTagKey']
y_pred_field = demisto.args()['emailPredictionKey']
y_pred_prob_field = demisto.args()['emailPredictionProbabilityKey']

labels_mapping = get_phishing_map_labels(demisto.args()['phishingLabels'])
additional_fields = demisto.args().get('additionalFields', '')
additional_fields = additional_fields.split(',')
additional_fields = [x.strip() for x in additional_fields]
main(incident_types, incident_query, y_true_field, y_pred_field, y_pred_prob_field, model_target_accuracy,
     labels_mapping, additional_fields)
