import itertools
from pathlib import Path

import msgpack
import numpy as np

from .cpp.refraction_correction import apply_correction_pipeline

LOAD_DIR = Path(__file__).parent / "refraction_models"
LOAD_VERSION = 1


class ModelDeserializationError(Exception):
    pass


class Refractionizer:
    def __init__(self, degree=3, type_="default", custom_load_dir=None):
        self.pipeline_radius_as_list = self.load_config_from_msgpack(
            "radius", type_, degree, custom_load_dir
        )

        self.pipeline_gaze_vector_as_list = self.load_config_from_msgpack(
            "gaze_vector", type_, degree, custom_load_dir
        )

        self.pipeline_sphere_center_as_list = self.load_config_from_msgpack(
            "sphere_center", type_, degree, custom_load_dir
        )

        self.pipeline_pupil_circle_as_list = self.load_config_from_msgpack(
            "pupil_circle", type_, degree, custom_load_dir
        )

    @staticmethod
    def load_config_from_msgpack(feature, type_, degree, custom_load_dir=None):
        load_dir = Path(custom_load_dir or LOAD_DIR).resolve()
        name = f"{type_}_refraction_model_{feature}_degree_{degree}.msgpack"
        path = load_dir / name
        with path.open("rb") as file:
            config_model = msgpack.unpack(file)
            Refractionizer._validate_loaded_model_config(config_model)
            try:
                return list(
                    itertools.chain(
                        Refractionizer._polynomial_features_from_config(config_model),
                        Refractionizer._standard_scaler_from_config(config_model),
                        Refractionizer._linear_regression_from_config(config_model),
                    )
                )
            except KeyError as err:
                raise ModelDeserializationError from err

    @staticmethod
    def _validate_loaded_model_config(config_model):
        if not isinstance(config_model, dict) or "version" not in config_model:
            raise ModelDeserializationError("Unrecognized format")
        if config_model["version"] != LOAD_VERSION:
            raise ModelDeserializationError(
                f"Unexpected version `{config_model['version']}` "
                f"(expected `{LOAD_VERSION}``)"
            )

    @staticmethod
    def _polynomial_features_from_config(config_model):
        yield np.array(config_model["steps"]["PolynomialFeatures"]["powers"])

    @staticmethod
    def _standard_scaler_from_config(config_model):
        config_scaler = config_model["steps"]["StandardScaler"]
        yield np.array(config_scaler["mean"])
        yield np.array(config_scaler["var"])

    @staticmethod
    def _linear_regression_from_config(config_model):
        config_lin_reg = config_model["steps"]["LinearRegression"]
        yield np.array(config_lin_reg["coef"])
        yield np.array(config_lin_reg["intercept"])

    @staticmethod
    def _apply_correction_pipeline(X, pipeline_arrays):
        return apply_correction_pipeline(np.asarray(X).T, *pipeline_arrays)

    def correct_radius(self, X):
        return self._apply_correction_pipeline(X, self.pipeline_radius_as_list)

    def correct_gaze_vector(self, X):
        return self._apply_correction_pipeline(X, self.pipeline_gaze_vector_as_list)

    def correct_sphere_center(self, X):
        return self._apply_correction_pipeline(X, self.pipeline_sphere_center_as_list)

    def correct_pupil_circle(self, X):
        return self._apply_correction_pipeline(X, self.pipeline_pupil_circle_as_list)


class SklearnRefractionizer(Refractionizer):
    def __init__(self, degree=3, type_="default", custom_load_dir=None):
        self.correct_radius = self.load_predict_fn_from_joblib_pickle(
            "radius", type_, degree, custom_load_dir
        )

        self.correct_gaze_vector = self.load_predict_fn_from_joblib_pickle(
            "gaze_vector", type_, degree, custom_load_dir
        )

        self.correct_sphere_center = self.load_predict_fn_from_joblib_pickle(
            "sphere_center", type_, degree, custom_load_dir
        )

        self.correct_pupil_circle = self.load_predict_fn_from_joblib_pickle(
            "pupil_circle", type_, degree, custom_load_dir
        )

    @staticmethod
    def load_predict_fn_from_joblib_pickle(
        feature, type_, degree, custom_load_dir=None
    ):
        import joblib

        load_dir = Path(custom_load_dir or LOAD_DIR).resolve()
        name = f"{type_}_refraction_model_{feature}_degree_{degree}.save"
        path = load_dir / name
        try:
            pipeline = joblib.load(path)
        except FileNotFoundError:
            raise
        except Exception as exc:
            raise ModelDeserializationError(
                f"Failed to load pickled model from {path}"
            ) from exc
        return pipeline.predict


if __name__ == "__main__":

    refractionizer = Refractionizer()

    print(refractionizer.correct_sphere_center([[0.0, 0.0, 35.0]]))
    print(refractionizer.correct_radius([[0.0, 0.0, 35.0, 0.0, 0.0, -1.0, 2.0]]))
    print(refractionizer.correct_gaze_vector([[0.0, 0.0, 35.0, 0.0, 0.0, -1.0, 2.0]]))
    print(refractionizer.correct_pupil_circle([[0.0, 0.0, 35.0, 0.0, 0.0, -1.0, 2.0]]))
