import tempfile
import os
import platform
import numpy as np

class ScalarProblemWrapper:
    # pylint: disable=too-few-public-methods
    '''Helper class that evalues a DaisyOptimizationProblem and extract the value from the returned
    dict'''

    def __init__(self, problem):
        '''
        Parameters
        ----------
        problem : DaisyOptimizationProblem
        '''
        self.problem = problem

    def __call__(self, parameter_values):
        '''
        Parameters
        ----------
        parameter_values : sequence
          Parameter values. Lenght MUST match length of `self.parameters`

        Returns
        -------
        float
        '''
        result = self.problem(parameter_values)
        err_msg = 'Expected a dict with exactly one scalar valued objective mapping'
        try:
            value = result.popitem()[1]
            if len(result) != 0 or not isinstance(value, (int, float)):
                raise RuntimeError(err_msg)
            return value
        except (KeyError, AttributeError, TypeError) as e:
            raise RuntimeError(err_msg) from e


class DaisyOptimizationProblem:
    # pylint: disable=too-many-arguments,too-many-positional-arguments,too-few-public-methods
    '''Class that knows how to run simulation and compute objective for a parameter set'''
    def __init__(
            self, runner, file_generator, objective_fn, parameters, data_dir=None, debug=False
    ):
        """
        Parameters
        ----------
        runner : DaisyRunner

        file_generator : FileGenerator

        objective_fn : DaisyObjective

        parameters : [DaisyParameter] OR dict of (str, [DaisyParameter])
          Parameters to optimize. If a list it is assumed that all parameters are for the 'dai' file

        data_dir : str
          If not None then temporary directories will be created in this directory. Otherwise, they
          will be created in a default location depending on platform.

        debug: bool
          If True do not delete the temporary directory where Daisy output is stored
        """
        self.runner = runner
        self.file_generator = file_generator
        self.objective_fn = objective_fn
        self.parameter_kind = {}
        if not isinstance(parameters, dict):
            parameters = { 'dai' : parameters }

        # Convert dict of parameters to a list of parameters and verify that there are no name
        # clashes.
        self.parameters = []
        for kind, params in parameters.items():
            for param in params:
                if param.name in self.parameter_kind:
                    raise ValueError(
                        'Parameters must have unique names across all templates. '
                        f'{param.name} from {kind} is also in {self.parameter_kind[param.name]}'
                    )
                self.parameter_kind[param.name] = kind
                self.parameters.append(param)

        self.data_dir = data_dir
        if data_dir is None and platform.system().lower() == 'linux':
            # There is a good chance that we are using flatpak, in which case we need to use a tmp
            # location that flatpak Daisy can read and write. Anything inside the user home is good.
            self.data_dir = os.path.expanduser('~/.tmp/daisy')
        if self.data_dir is not None:
            os.makedirs(self.data_dir, exist_ok=True)
        self.debug = debug

    def __call__(self, parameter_values):
        # TODO: Rewrite to accept a dict of parameters. This is too brittle
        """Run Daisy with the given parameters and evaluate the objective. The return value depends
        on

        Parameters
        ----------
        parameter_values : sequence
          Parameter values. Lenght MUST match length of `self.parameters`

        Returns
        -------
        objective_map : dict of [str, float]
          Mapping from objective names to objective values
        """
        named_parameters = { 'dai' : {} }
        for p, value in zip(self.parameters, parameter_values):
            kind = self.parameter_kind[p.name]
            if kind not in named_parameters:
                named_parameters[kind] = { p.name : value }
            else:
                named_parameters[kind][p.name] = value

        # If we debug then we dont want the directory to be deleted after use
        # From python 3.12 we can pass delete=False to TemporaryDirectory, but prior to that we need
        # to use mkdtemp.
        if self.debug:
            output_directory = tempfile.mkdtemp(dir=self.data_dir)
            return self._run(output_directory, named_parameters)

        with tempfile.TemporaryDirectory(dir=self.data_dir) as output_directory:
            return self._run(output_directory, named_parameters)

    def _run(self, output_directory, named_parameters):
        dai_file = self.file_generator(output_directory, named_parameters, tagged=True)['dai']
        sim_result = self.runner(dai_file, output_directory)
        if sim_result.returncode != 0:
            print(sim_result)
            return { self.objective_fn.name : np.nan }
        return self.objective_fn(output_directory)
