import inspect
import traceback
from io import StringIO
from typing import Dict, Any, List, Tuple

from sklearn.multioutput import MultiOutputRegressor

from enmapbox.typeguard import typechecked
from enmapboxprocessing.enmapalgorithm import EnMAPProcessingAlgorithm, Group
from enmapboxprocessing.typing import RegressorDump
from enmapboxprocessing.utils import Utils
from qgis.core import QgsProcessingContext, QgsProcessingFeedback


@typechecked
class FitRegressorAlgorithmBase(EnMAPProcessingAlgorithm):
    P_DATASET, _DATASET = 'dataset', 'Training dataset'
    P_REGRESSOR, _REGRESSOR = 'regressor', 'Regressor'
    P_OUTPUT_REGRESSOR, _OUTPUT_REGRESSOR = 'outputRegressor', 'Output regressor'

    def helpParameters(self) -> List[Tuple[str, str]]:
        return [
            (self._DATASET, 'Training dataset skops file used for fitting the classifier. '
                            'If not specified, an unfitted classifier is created.'),
            (
                self._REGRESSOR, self.helpParameterCode()
                + '\nNote: The Python code provided here is executed locally with the permissions of the current '
                  'user during algorithm execution.'
            ),
            (self._OUTPUT_REGRESSOR, self.SkopsFileDestination)
        ]

    def displayName(self) -> str:
        raise NotImplementedError()

    def shortDescription(self) -> str:
        raise NotImplementedError()

    def code(self):
        raise NotImplementedError()

    def helpParameterCode(self) -> str:
        raise NotImplementedError()

    def group(self):
        return Group.Regression.value

    def initAlgorithm(self, configuration: Dict[str, Any] = None):
        self.addParameterCode(self.P_REGRESSOR, self._REGRESSOR, self.defaultCodeAsString())
        self.addParameterRegressionDataset(self.P_DATASET, self._DATASET, None, True)
        self.addParameterFileDestination(self.P_OUTPUT_REGRESSOR, self._OUTPUT_REGRESSOR, self.SkopsFileFilter)

    def defaultCodeAsString(self):
        try:
            lines = [line[8:] for line in inspect.getsource(self.code).split('\n')][1:-2]
        except OSError:
            lines = ['']
        lines = '\n'.join(lines)
        return lines

    def parameterAsRegressor(self, parameters: Dict[str, Any], name, context: QgsProcessingContext):
        namespace = dict()
        code = self.parameterAsString(parameters, name, context)

        # nosec B102 # User-defined scikit-learn model code execution by design; equivalent to the QGIS Python Console.
        # The code execution is transparently documented for users (e.g., via the Processing algorithm help).
        exec(code, namespace)  # nosec B102 # User-defined scikit-learn model code execution by design;
        return namespace['regressor']

    def checkParameterValues(self, parameters: Dict[str, Any], context: QgsProcessingContext) -> Tuple[bool, str]:
        valid, message = super().checkParameterValues(parameters, context)
        if not valid:
            return valid, message
        # check code
        try:
            self.parameterAsRegressor(parameters, self.P_REGRESSOR, context)
        except Exception:
            return False, traceback.format_exc()
        return True, ''

    def processAlgorithm(
        self, parameters: Dict[str, Any], context: QgsProcessingContext, feedback: QgsProcessingFeedback
    ) -> Dict[str, Any]:
        dump = self.parameterAsRegressorDump(parameters, self.P_DATASET, context)
        filename = self.parameterAsFileOutput(parameters, self.P_OUTPUT_REGRESSOR, context)
        regressor = self.parameterAsRegressor(parameters, self.P_REGRESSOR, context)

        with open(filename + '.log', 'w') as logfile:
            feedback, feedback2 = self.createLoggingFeedback(feedback, logfile)
            self.tic(feedback, parameters, context)

            if dump is not None:
                feedback.pushInfo(
                    f'Load training dataset: X=array{list(dump.X.shape)} y=array{list(dump.y.shape)} '
                    f'targets={[c.name for c in dump.targets]}')
                feedback.pushInfo('Fit regressor')

                if isinstance(regressor, MultiOutputRegressor):
                    dumpY = dump.y
                elif dump.y.shape[1] == 1:
                    dumpY = dump.y.ravel()
                else:
                    dumpY = dump.y

                try:
                    regressor.fit(dump.X, dumpY, log_cout=StringIO(), log_cerr=StringIO())  # fixes issue #790
                except Exception:
                    regressor.fit(dump.X, dumpY)

            else:
                feedback.pushInfo('Store unfitted classifier')
                dump = RegressorDump(None, None, None, None, regressor)

            dump.regressor = regressor
            Utils.modelDump(dump.__dict__, filename)

            result = {self.P_OUTPUT_REGRESSOR: filename}
            self.toc(feedback, result)

        return result
