import os
import catboost
from sklearn.externals import joblib


class BaseModel(object):
    def __init__(self, model_name, binary_path, is_local,
                 score_name='undefined-score-name',
                 date_str='undefined-date',
                 save_info=False):
        self.model_name = model_name
        self.binary_path = binary_path
        self.is_local = is_local
        self.is_binary_ready = False
        self.score_name = score_name
        self.date_str = date_str
        self.save_info = save_info

        if self.is_local:
            self.load_binary()

    def load_binary(self):
        if self.is_binary_ready:
            return
        if self.is_local:
            self.binary = joblib.load(self.binary_path)
        else:
            self.binary = joblib.load(os.path.basename(self.binary_path))
        self.is_binary_ready = True

    def apply(self, data):
        raise 'Not implemented. Inherit class and implement.'


class PredictProbaModel(BaseModel):
    """This should return probability of class 1"""
    def __init__(self, model_name, binary_path, is_local, date_str,
                 score_name, save_info=False):
        super(PredictProbaModel, self).__init__(
            model_name=model_name, binary_path=binary_path, is_local=is_local,
            date_str=date_str, score_name=score_name, save_info=save_info)

    def apply(self, data):
        if not self.is_local:
            self.load_binary()
        return float(self.binary.predict_proba([data])[0, 1])


class InvertedPredictProbaModel(BaseModel):
    """This should return probability of class 0"""
    def __init__(self, model_name, binary_path, is_local, date_str,
                 score_name, save_info=False):
        super(InvertedPredictProbaModel, self).__init__(
            model_name=model_name, binary_path=binary_path, is_local=is_local,
            date_str=date_str, score_name=score_name, save_info=save_info)

    def apply(self, data):
        if not self.is_local:
            self.load_binary()
        return float(self.binary.predict_proba([data])[0, 0])


class PredictModel(BaseModel):
    def __init__(self, model_name, binary_path, is_local, date_str,
                 score_name, save_info=False):
        super(PredictModel, self).__init__(
            model_name=model_name, binary_path=binary_path, is_local=is_local,
            date_str=date_str, score_name=score_name, save_info=save_info)

    def apply(self, data):
        self.load_binary()
        return float(self.binary.predict([data]))


class CatboostPredictProbaModel(BaseModel):
    def __init__(self, model_name, binary_path, is_local, date_str,
                 score_name, save_info=False):
        super(CatboostPredictProbaModel, self).__init__(
            model_name=model_name, binary_path=binary_path, is_local=is_local,
            date_str=date_str, score_name=score_name, save_info=save_info)

    def apply(self, data):
        self.load_binary()
        return float(self.binary.predict_proba([data])[0, 1])

    def load_binary(self):
        if self.is_binary_ready:
            return

        self.binary = catboost.CatBoostClassifier()
        if self.is_local:
            self.binary.load_model(self.binary_path)
        else:
            self.binary.load_model(os.path.basename(self.binary_path))

        self.is_binary_ready = True
