# Copyright (C) Prizmi, LLC - All Rights Reserved
# Unauthorized copying or use of this file is strictly prohibited and subject to prosecution under applicable laws
# Proprietary and confidential

"""The main module for validating the output files generated by Learner."""

import sys
import logging

from learner.data_worker.data_loader import get_file_num_rows


class OutputValidator:
    def __init__(self, conf, models_dict):
        """Validates the outputs of prediction process.

         :param conf: configuration object
         :param models_dict: dict containing the model tag and number of rows scored
         """
        self._conf = conf
        self._models_dict = models_dict

    @property
    def conf(self):
        return self._conf

    @property
    def models_dict(self):
        return self._models_dict

    def validate(self):
        """The main function that runs all the instance methods to validate the outputs generated by Learner if the
        validation flag is set to true.

        :return: None
        """
        logging.info("Validating the output...")
        for tag, mdls in self._models_dict.items():
            mdl = mdls.get(0, mdls)
            self.validate_num_predictions(tag, mdl)
        logging.info("Successfully validated the output")

    def validate_num_predictions(self, tag, mdl):
        """Check if the output file generated by Learner contain the expected number of rows or not.

        :param tag: The tag for the model
        :return: None
        :raises SystemExit: if the number of rows in the file doesn't match the expected value
        """
        logging.info("Validating number of predictions")
        filename = mdl["path"] + self._conf.workspace.name + "_test_" + str(tag) + str(self._conf.sep_timetag) + ".csv"
        file_num_rows = get_file_num_rows(filename)

        diff = self.conf.model.nrows_score - file_num_rows
        if diff != 0:
            logging.error("Prediction input and output shapes do not match.\n "
                          "Expected: %i Found: %i", self.conf.model.nrows_score, file_num_rows)
            sys.exit(1)
