from dataclasses import asdict

from sklearn.exceptions import NotFittedError
from sklearn.utils.validation import check_is_fitted

from logic import *


class UserError(Exception):
    """A mistake the student can fix; its message is shown in the UI."""


class Manager:
    def __init__(self):
        self._dataset: Dataset | None = None
        self._model: BaseEstimator | None = None
        self._model_config: ModelSettings | None = None
        # Maps plot coordinates back to inputs; None when the inputs are the plot coordinates
        self._projection: Projection | None = None
        self._input_names: list[str] = []

    def handle_set_dataset(self, dataset: dict) -> None:
        self._reset_model()

        try:
            parsed_dataset = Dataset(
                inputs=np.array(dataset["inputs"]),
                outputs=np.array(dataset["outputs"]),
            )
            input_names = [str(name) for name in dataset.get("input_names", [])]
            projection = dataset.get("projection")
            parsed_projection = None if projection is None else Projection(
                origin=np.array(projection["origin"]),
                components=np.array(projection["components"]),
            )
        except Exception as e:
            raise ValueError(f"Invalid dataset format: {e}")

        self._dataset = parsed_dataset
        self._projection = parsed_projection
        n_inputs = parsed_dataset.inputs.shape[1] if parsed_dataset.inputs.ndim == 2 else 0
        self._input_names = input_names or [f"x{i + 1}" for i in range(n_inputs)]

    def handle_set_model_config(self, settings: dict) -> None:
        self._reset_model()

        try:
            parsed_settings = ModelSettings(
                type=settings["type"],
                params=dict(settings.get("params", {})),
            )
        except Exception as e:
            raise ValueError(f"Invalid model settings format: {e}")

        self._model_config = parsed_settings

    def handle_build_model(self) -> dict | None:
        """Fit the model; returns its metrics and learned parameters (None with no data)."""
        self._reset_model()

        if self._dataset is None:
            raise ValueError("Dataset not set")
        if self._model_config is None:
            raise ValueError("Model settings not set")

        try:
            self._model = init_model(self._model_config)
        except Exception as e:
            # user error
            raise UserError(f"Error initializing model: {e}")

        if self._dataset.inputs.size > 0 and self._dataset.outputs.size > 0:
            try:
                train_model(self._model, self._dataset)
            except Exception as e:
                # user error
                self._reset_model()
                raise UserError(f"Error training model: {e}")
            return describe_model(self._model, self._dataset, self._input_names)
        return None

    def handle_get_decision_boundary(self, settings: dict) -> dict:
        if self._model is None:
            return dict(x=[], y=[], labels=[])

        try:
            check_is_fitted(self._model)
        except NotFittedError:
            return dict(x=[], y=[], labels=[])

        try:
            parsed_settings = DecisionBoundarySettings(
                xmin=settings["xmin"],
                xmax=settings["xmax"],
                ymin=settings["ymin"],
                ymax=settings["ymax"],
                resolution=settings.get("resolution", 500),
            )
        except Exception as e:
            raise ValueError(f"Invalid decision boundary settings format: {e}")

        try:
            result = get_decision_boundary_values(
                self._model, 
                parsed_settings, 
                self._projection.lift if self._projection else None,
            )
            
        except Exception as e:
            raise UserError(f"Error computing decision boundary: {e}")

        return asdict(result)

    def _reset_model(self) -> None:
        self._model = None
        print("Model reset")