# pylint: disable=no-member
import traceback

import demisto_ml
from CommonServerPython import *

TARGET_PRECISION = 0.97
THRESHOLD = 0.9
OUT_OF_THE_BOX_MODEL_NAME = "demisto_out_of_the_box_model_v2"
OUT_OF_THE_BOX_MODEL_PATH = "/ml/encrypted_model.b"
EVALUATION_PATH = "/ml/oob_evaluation.txt"
OOB_VERSION_INFO_KEY = "oob_version"


def oob_model_exists_and_updated():
    res_model = demisto.executeCommand("getMLModel", {"modelName": OUT_OF_THE_BOX_MODEL_NAME})[0]
    if is_error(res_model):
        return False
    model_type = dict_safe_get(res_model, [0, "Contents", "model", "type", "type"], "UNKNOWN_MODEL_TYPE")
    return model_type == demisto_ml.ModelType.Torch.value


def load_oob_model():
    try:
        encoded_model = demisto_ml.load_oob(OUT_OF_THE_BOX_MODEL_PATH)
    except Exception:
        return_error(traceback.format_exc())
    res = demisto.executeCommand(
        "createMLModel",
        {
            "modelData": encoded_model,
            "modelName": OUT_OF_THE_BOX_MODEL_NAME,
            "modelLabels": ["Malicious", "Non-Malicious"],
            "modelOverride": "true",
            "modelType": demisto_ml.ModelType.Torch.value,
            "modelExtraInfo": {"threshold": THRESHOLD},
        },
    )
    if is_error(res):
        return_error(get_error(res))

    with open(EVALUATION_PATH) as json_file:
        data = json.load(json_file)
    y_test = data["YTrue"]
    y_pred = data["YPred"]
    y_pred_prob = data["YPredProb"]

    y_pred_evaluation = [{pred: prob} for pred, prob in zip(y_pred, y_pred_prob)]
    res = demisto.executeCommand(
        "GetMLModelEvaluation",
        {
            "yTrue": json.dumps(y_test),
            "yPred": json.dumps(y_pred_evaluation),
            "targetPrecision": str(0.85),
            "targetRecall": str(0),
            "detailedOutput": "true",
        },
    )
    if is_error(res):
        return_error(get_error(res))
    confusion_matrix = json.loads(res[0]["Contents"]["csr_matrix_at_threshold"])
    confusion_matrix_no_all = {k: v for k, v in confusion_matrix.items() if k != "All"}
    confusion_matrix_no_all = {
        k: {sub_k: sub_v for sub_k, sub_v in v.items() if sub_k != "All"} for k, v in confusion_matrix_no_all.items()
    }
    res = demisto.executeCommand(
        "evaluateMLModel",
        {
            "modelConfusionMatrix": confusion_matrix_no_all,
            "modelName": OUT_OF_THE_BOX_MODEL_NAME,
            "modelEvaluationVectors": {"Ypred": y_pred, "Ytrue": y_test, "YpredProb": y_pred_prob},
            "modelConfidenceThreshold": THRESHOLD,
            "modelTargetPrecision": TARGET_PRECISION,
        },
    )

    if is_error(res):
        return_error(get_error(res))


def predict_phishing_words():
    if not oob_model_exists_and_updated():
        load_oob_model()
    dargs = demisto.args()
    dargs["modelName"] = OUT_OF_THE_BOX_MODEL_NAME
    dargs["modelStoreType"] = "mlModel"
    res = demisto.executeCommand("DBotPredictPhishingWords", dargs)
    if is_error(res):
        return_error(get_error(res))
    return res


def main():
    res = predict_phishing_words()
    return res


if __name__ in ["__main__", "__builtin__", "builtins"]:
    demisto.results(main())