from typing import Callable, NoReturn, Union
# https://stackoverflow.com/questions/46092104/subclass-in-type-hinting

import ciw
import numpy as np
from sklearn.base import BaseEstimator

# TODO: Write a function which returns new SKRouter classes in order to get around the
# fact that certain params will not be available at runtime.
class SKRouter(ciw.Node):
    '''
    Routes individuals depending on the system's state and time using a Scikit-Learn model.

    Args:
        get_pred_data (Callable): A function that takes an instance of Self and an individual as input
                                and returns the data used for prediction of the next node.
        skmodel (BaseEstimator): An instance of a scikit-learn compatible model.
        method (str, optional): Label indicating how the next node is obtained from skmodel.
                                   Defaults to `'predict_proba'`.

    Raises:
        ValueError: If method is predict_proba the model does not have the `predict_proba` method.

    Attributes:
        get_pred_data (Callable): The provided function for obtaining prediction data.
        skmodel (BaseEstimator): The scikit-learn model used for predicting the next node.

    Methods:
        next_node(ind: ciw.Individual) -> ciw.Node:
            Predicts the next node based on the provided individual's data using skmodel.

    '''

    def __init__(self, get_pred_data: Callable, skmodel: BaseEstimator, method: Union[str,Callable] = 'predict_proba') -> NoReturn:
        '''
        Initializes an instance of SKRouter.

        Args:
            get_pred_data (Callable): A function that takes an instance of Self and an individual as input
                                    and returns the data used for prediction of the next node.
            skmodel (BaseEstimator): An instance of a scikit-learn compatible model.
            method (str, optional): Label indicating how the next node is obtained from skmodel.
                                   Defaults to `'predict_proba'`.

        Raises:
            ValueError: If sampling is True and the skmodel does not have the `predict_proba` method.

        '''
        self.get_data = get_pred_data
        self.clf = skmodel

        if method not in ('predict_proba', 'predict') and not callable(method):
            raise NotImplementedError(f'{method} is not supported.')
        self.method: = method


    def next_node(self, ind: ciw.Individual) -> ciw.Node:
        '''
        Predicts the next node based on the provided individual's data using the scikit-learn model.

        Args:
            ind (ciw.Individual): The individual for which the next node needs to be predicted.

        Returns:
            ciw.Node: The predicted next node.

        '''
        
        pred_data = self.get_pred_data(self, ind)

        if self.method == 'predict_proba':
            if hasattr(self.skmodel, 'predict_proba'):
                probs = self.skmodel.predict_proba(pred_data)[0]
                classes = self.skmodel.classes_
            else:
                probs = self.skmodel.predict(pred_data)[0]
                chosen_node = range(len(self.simulation.nodes))
            chosen_node = classes @ np.random.multinomial(1, probs)
        elif self.method == 'predict':
            chosen_node = self.skmodel.predict(pred_data)[0]
        else:
            chosen_node = self.method(self.skmodel.predict(pred_data)[0])

        return self.simulation.nodes[chosen_node]
