from math import nan
from typing import Dict, Any, List, Tuple

import numpy as np

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


@typechecked
class SpectralIndexOptimizerAlgorithm(EnMAPProcessingAlgorithm):
    P_DATASET, _DATASET = 'dataset', 'Training dataset'
    P_FORMULA, _FORMULA = 'formula', 'Formula'
    P_MAX_FEATURES, _MAX_FEATURES = 'maxFeatures', 'Max. features'
    P_F1, _F1 = 'f1', 'Fixed feature F1'
    P_F2, _F2 = 'f2', 'Fixed feature F2'
    P_F3, _F3 = 'f3', 'Fixed feature F3'
    P_OUTPUT_MATRIX, _OUTPUT_MATRIX = 'outScoreMatrix', 'Output score matrix'

    @classmethod
    def displayName(cls) -> str:
        return 'Spectral Index Optimizer'

    def shortDescription(self) -> str:
        return 'This algorithm finds the optimal two-feature index by modelling a target variable via linear ' \
               'regression.'

    def helpParameters(self) -> List[Tuple[str, str]]:
        return [
            (self._DATASET, 'The regression dataset.'),
            (self._FORMULA, 'The formula with variable features A and B to be optimized, '
                            'and up to three fixed features F1, F2 and F3.\n'
                            'Note: The Python code provided here is executed locally with the permissions of the '
                            'current user during algorithm execution.'),
            (self._MAX_FEATURES, 'Limit the number of features to be evaluated. Default is to use all features.'),
            (self._F1, 'Specify to use a fixed feature F1 in the formula.'),
            (self._F2, 'Specify to use a fixed feature F2 in the formula.'),
            (self._F3, 'Specify to use a fixed feature F3 in the formula.'),
            (self._OUTPUT_MATRIX, self.RasterFileDestination)
        ]

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

    def initAlgorithm(self, configuration: Dict[str, Any] = None):
        self.addParameterRegressionDataset(self.P_DATASET, self._DATASET)
        self.addParameterString(self.P_FORMULA, self._FORMULA, '(A-B) / (A+B)')
        self.addParameterInt(self.P_MAX_FEATURES, self._MAX_FEATURES, None, True, 2, None, True)
        self.addParameterInt(self.P_F1, self._F1, None, True, 1, None, True)
        self.addParameterInt(self.P_F2, self._F2, None, True, 1, None, True)
        self.addParameterInt(self.P_F3, self._F3, None, True, 1, None, True)
        self.addParameterRasterDestination(self.P_OUTPUT_MATRIX, self._OUTPUT_MATRIX)

    def processAlgorithm(
        self, parameters: Dict[str, Any], context: QgsProcessingContext, feedback: QgsProcessingFeedback
    ) -> Dict[str, Any]:
        from sklearn.linear_model import LinearRegression
        from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error

        filenameDataset = self.parameterAsFile(parameters, self.P_DATASET, context)
        formula = self.parameterAsString(parameters, self.P_FORMULA, context)
        maxFeatures = self.parameterAsInt(parameters, self.P_MAX_FEATURES, context)
        f1No = self.parameterAsInt(parameters, self.P_F1, context)
        f2No = self.parameterAsInt(parameters, self.P_F2, context)
        f3No = self.parameterAsInt(parameters, self.P_F3, context)
        filename = self.parameterAsOutputLayer(parameters, self.P_OUTPUT_MATRIX, context)

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

            dump = RegressorDump.fromDict(Utils.modelLoad(filenameDataset))
            features = np.array(dump.features)
            targets = dump.targets
            X = np.array(dump.X, np.float32)
            y = np.array(dump.y, np.float32)

            F1 = F2 = F3 = None
            if f1No is not None:
                F1 = X[:, f1No - 1]
            if f2No is not None:
                F2 = X[:, f2No - 1]
            if f3No is not None:
                F3 = X[:, f3No - 1]

            if maxFeatures is not None:
                maxFeatures = min(maxFeatures, len(features))
                select = np.round(np.linspace(0, len(features) - 1, maxFeatures)).astype(int)
                features = features[select]
                X = X[:, select]

            feedback.pushInfo('Used features:')
            for featureNo, feature in enumerate(features, 1):
                feedback.pushInfo(f'{featureNo}: {feature}')
            nfeatures = len(features)
            ntargets = len(targets)
            olsr = LinearRegression()
            nbands = ntargets * 3
            scores = np.full((nbands, nfeatures, nfeatures), nan)
            for ai in range(nfeatures):
                feedback.setProgress(ai / nfeatures * 100)
                A = X[:, ai]
                for bi in range(ai + 1, nfeatures):
                    B = X[:, bi]
                    for yi in range(ntargets):
                        # nosec B307 # User-defined spectral index code evaluation by design;
                        # equivalent to the QGIS Python Console.
                        # The code execution is transparently documented for users
                        # (e.g. via the Processing algorithm help).
                        S = eval(formula,
                                 {'A': A, 'B': B, 'F1': F1, 'F2': F2, 'F3': F3}
                                 )  # nosec B307 # User-defined spectral index evaluation like in Python Console.

                        if not isinstance(S, np.ndarray):
                            raise ValueError(
                                f'formula must evaluate to a numpy.ndarray, got {type(S).__name__}'
                            )
                        Y = y[:, yi].flatten()

                        # formula may eval to not finite values
                        valid = np.isfinite(S)
                        S = S[valid]
                        Y = Y[valid]

                        S = S.reshape((-1, 1))
                        olsr.fit(S, Y)
                        yP = olsr.predict(S)

                        rmse = mean_squared_error(Y, yP) ** 0.5
                        mae = mean_absolute_error(Y, yP)
                        r2 = r2_score(Y, yP)
                        scores[yi * 3 + 0, ai, bi] = scores[yi * 3 + 0, bi, ai] = rmse
                        scores[yi * 3 + 1, ai, bi] = scores[yi * 3 + 1, bi, ai] = mae
                        scores[yi * 3 + 2, ai, bi] = scores[yi * 3 + 2, bi, ai] = r2

            bandNames = list()
            scoreNames = ['RMSE', 'MAE', 'R^2']
            for target in targets:
                for scoreName in scoreNames:
                    bandNames.append(f'{target.name} - {scoreName}')

            writer = Driver(filename).createFromArray(scores)

            # report best scores
            feedback.pushInfo('Best feature combinations:')
            for arr, bandName in zip(scores, bandNames):
                if 'R^2' in bandName:
                    a, b = np.unravel_index(np.nanargmax(arr), arr.shape)
                else:
                    a, b = np.unravel_index(np.nanargmin(arr), arr.shape)
                feedback.pushInfo(f'{bandName}: {arr[a, b]}; {features[a]}, {features[b]}')
                writer.setMetadataItem(bandName, f'{arr[a, b]}; {features[a]}, {features[b]}')

            for bandNo, bandName in enumerate(bandNames, 1):
                writer.setBandName(bandName, bandNo)
            writer.setNoDataValue(nan)
            writer.setMetadataItem('features', list(features))
            writer.close()
            result = {self.P_OUTPUT_MATRIX: filename}

            self.toc(feedback, result)

        return result
