from joblib import dump, load
import numpy as np


class FacilityFatalityModel:
    """Predict annual bird fatalities per MW for a wind-facility design."""

    def __init__(self, estimator, numeric_features, categorical_features,
                 prediction_interval_half_width, metadata=None):
        self.estimator = estimator
        self.numeric_features = list(numeric_features)
        self.categorical_features = list(categorical_features)
        self.prediction_interval_half_width = float(prediction_interval_half_width)
        self.metadata = metadata or {}

    @property
    def feature_columns(self):
        return self.numeric_features + self.categorical_features

    def predict(self, rows):
        predictions = np.asarray(
            self.estimator.predict(rows[self.feature_columns]), dtype=float
        )
        return np.maximum(predictions, 0.0)

    def predict_interval(self, rows):
        predictions = self.predict(rows)
        half_width = self.prediction_interval_half_width
        return np.column_stack([
            np.maximum(predictions - half_width, 0.0),
            predictions + half_width,
        ])

    def save(self, filename):
        dump(self, filename)

    @staticmethod
    def load(filename):
        return load(filename)