# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import time
import multiprocessing
import string


class Optimizer:
    """Base class for all optimization methods.

    Parameters
    ----------
    dimensions: int
        Number of dimensions of the search space i.e. number of optimization variables.
    evaluation_function: callable
        A function (or a callable class) which takes a Numpy:ndarray as argument and returns a Numpy:ndarray containing
        objectives' and constraint's values. Optionally, it can take a keyword argument :code:`s` in order to handle
        unique string generated by the optimizer.
    monitoring: str or None
        Adjust the monitoring style and level of output produced during the optimization. :code:`monitoring='None'` or
        :code:`monitoring=None` suppress all output. For displaying a per iteration output to the terminal use
        :code:`monitoring='basic'` which can display a colored output if supported by your terminal. For displaying a
         progress bar and dynamically updated best candidate state use :code:`monitoring='dashboard'`.
    best: CandidateState
        The best candidate state encountered in the optimization process.
    """

    def __init__(self):
        """Initialization for Optimizer class defines member variables and set initial values.

        Returns
        -------
        Optimizer
            Optimizer instance
        """

        self.dimensions = None
        self.evaluation_function = None
        self.number_of_processes = 1
        
        self.objectives = 1
        self.objective_weights = None
        self.objective_labels = None
        self.constraints = 0
        self.constraint_labels = None
        
        self.iterations = None
        self.evaluations_count = 0
        self.maximum_evaluations = None
        self.target_fitness = None
        
        self.lb = None
        self.ub = None
        self.best = None
        
        self.monitoring = None
        self._progress_log = None
        self._time1 = time.time()

        self.forward_unique_str = False
        self._unique_str_list = []
        self.it = 0
        
    def _init_optimizer(self):
        """A private method for Optimizer initialization which checks the types and values of Optimizer member
        parameters. Automatically sets optional values if they are not provided.

        Returns
        -------
        None
            Nothing
        """
        
        assert isinstance(self.dimensions, int) and self.dimensions > 0, \
            "optimizer.dimensions should be positive integer"
            
        if np.size(self.lb) == 1 and np.size(self.ub) == 1:
            self.lb = np.full(self.dimensions, self.lb)
            self.ub = np.full(self.dimensions, self.ub)
            
        assert np.size(self.lb) == self.dimensions and \
            np.size(self.ub) == self.dimensions, \
            "optimizer.lb and optimizer.ub should be of size optimizer.dimensions"

        assert callable(self.evaluation_function), "optimizer.evaluation_function needs to be callable"

        assert isinstance(self.objectives, int) and self.objectives > 0, \
            "optimizer.objectives should be positive integer"
        
        assert isinstance(self.constraints, int) and self.constraints >= 0, \
            "optimizer.constraints should be non-negative integer"
        
        if self.objective_weights is None:
            self.objective_weights = np.ones(self.objectives) / self.objectives   
        else:
            assert len(self.objective_weights) == self.objectives, \
                "optimizer.objective_weights list should contain number of elements equal to optimizer.objectives"
        
        if self.objective_labels is None:
            self.objective_labels = [f'obj_{o}' for o in range(self.objectives)]
        else:
            assert len(self.objective_labels) == self.objectives, \
                "optimizer.objective_labels list should contain number of strings equal to optimizer.objectives"
            
        if self.constraint_labels is None:
            self.constraint_labels = [f'cnstr_{c}' for c in range(self.constraints)]
        else:
            assert len(self.constraint_labels) == self.constraints, \
                "optimizer.constraint_labels list should contain number of strings equal to optimizer.constraints"

        assert (isinstance(self.number_of_processes, int) and self.number_of_processes > 0)\
               or self.number_of_processes == 'maximum', \
               "optimizer.number_of_processes should be positive integer or \'maximum\'"
        if self.number_of_processes == 'maximum':
            self.number_of_processes = multiprocessing.cpu_count()
        if self.number_of_processes > 1: 
            if self.monitoring == 'basic':
                self.log(f'Preparing a multiprocess pool for {self.number_of_processes} processes.')
            self.pool = multiprocessing.Pool(self.number_of_processes)

        if self.maximum_evaluations:
            assert isinstance(self.maximum_evaluations, int) and self.maximum_evaluations > 0, \
                "optimizer.maximum_evaluations should be positive integer"

        if self.iterations is None:
            self.iterations = 100 * self.dimensions

        assert isinstance(self.monitoring, str) or self.monitoring is None, \
            'optimizer.monitoring should be a string or None'

        if isinstance(self.monitoring, str):
            # Convert to lowercase string
            self.monitoring = self.monitoring.lower()
        if self.monitoring is None:
            self.monitoring = 'none'

        if self.monitoring == 'none':
            self._progress_log = lambda remark=None: None

        elif self.monitoring == 'basic':
            # self.progress_log = self._progress_log2
            from rich.console import Console
            self.rich_console = Console(highlight=False)
            self._progress_log = self.__progress_log_basic

        elif self.monitoring == 'dashboard':
            from rich.progress import Progress, BarColumn, TimeRemainingColumn, TimeElapsedColumn
            from rich.live import Live
            from rich.console import RenderGroup, Text
            # from rich.panel import Panel
            # from datetime import timedelta
            import indago

            self._progress_log = lambda remark=None: None
            self.__progress_bar = Progress(
                "[progress.description]{task.description}",
                # "{task.completed:} of {task.total}",
                BarColumn(),
                "[progress.percentage]{task.percentage:>3.0f}%",
                # "Elapsed: {task.time_remaining}",
                "Elapsed:",
                TimeElapsedColumn(),
                "Remaining:",
                TimeRemainingColumn(),
                # expand=True,
            )

            self.__progress_bar_tasks = [self.__progress_bar.add_task("Max. iterations",
                                                                      total=self.iterations)]
            if self.maximum_evaluations:
                self.__progress_bar_tasks.append(self.__progress_bar.add_task("Max. evaluations",
                                                                              total=self.maximum_evaluations))
            if self.target_fitness:
                self.__progress_bar_tasks.append(self.__progress_bar.add_task("Target fitness", total=1))

            def update_progress_bar():

                progress = [self.it]

                if self.maximum_evaluations:
                    progress.append(self.evaluations_count)

                if self.target_fitness and self.best:
                    f0 = None
                    for h in self.results.cHistory:
                        if np.all(h[2].C <= 0):
                            f0 = h[2].f
                            break

                    if f0:
                        if self.best.f > self.target_fitness:
                            p = (self.best.f - self.target_fitness) / (f0 - self.target_fitness)
                        else:
                            p = 0
                        p = 1 - np.log10(1 + 9 * p)
                    else:
                        p = 0
                    progress.append(p)

                for task, prog in zip(self.__progress_bar_tasks, progress):
                    self.__progress_bar.update(task, completed=prog)

                best = 'None'
                if self.best is not None:
                    best = str(self.best.X)
                    all_labels = self.objective_labels + self.constraint_labels
                    label_len = np.max([len(lbl) for lbl in all_labels]) + 3
                    for o, o_label in zip(self.best.O, self.objective_labels):
                        best += '\n[magenta]' + f'{o_label}: '.rjust(label_len) + f'{o:12.5e} [/magenta]'
                    for c, c_label in zip(self.best.C, self.constraint_labels):
                        if c <= 0:
                            best += '\n[green]' + f'{c_label}: '.rjust(label_len) + f'{c:12.5e}   => (satisfied)[/green]'
                        else:
                            best += '\n[red]' + f'{c_label}: '.rjust(label_len) + f'{c:12.5e}   => (violated)[/red]'
                    best += '\n' + 'Fitness: '.rjust(label_len) + f'{self.best.f:12.5e}'

                return RenderGroup(f'[b]INDAGO[/b] {indago.__version__}',
                                   f'{self.__class__.__name__} optimization running\n',
                                   self.__progress_bar,
                                   f'\nCurrent iteration:{self.it}\nTotal evaluations:{self.evaluations_count}\nBest:\n{best}')
            self.__update_progress_bar = update_progress_bar
            self.live = Live(self.__update_progress_bar())

        else:
            assert False, f'Unknown value for optimizer.monitoring parameter ({self.monitoring})'
        self.evaluations_count = 0
        self.it = 0

        # Initialize history
        self.results = OptimizationResults(self)
        self.results.cHistory = []
        self.best = None

    def _check_params(self, mandatory_params, optional_params, defined_params):
        """Private method which checks if optimizer parameters are defined in Optimizer.params dictionary.
        Should be called in initializations of derived Optimizer classes. It asserts if any of mandatory parameters is
        missing and prints a warning if unknown/excessive parameter is provided.

        Parameters
        ----------
        mandatory_params : list of str
            A list of mandatory parameter names/keys for given method.
        optional_params: list of str
            A list of optional parameter names/keys for given method.
        defined_params: list of str
            A list of method's parameters names/keys defined in Optimizer.params dictionary.

        Returns
        -------
        None
            Nothing
        """
        for param in mandatory_params:
            assert param in defined_params, f'Error: Missing parameter {param}'

        for param in defined_params:
            if param not in mandatory_params and param not in optional_params:
                print(f'Warning: Excessive parameter {param}')

    @staticmethod
    def title(msg, line_length=60):
        """Utility function for consistently displaying titles in the output. The text of the title is decorated and
        printed on the center of the screen.

        Parameters
        ----------
        msg: str
            A text of the title.
        line_length: int
            A length (number of characters) of the line used for drawing the decoration border and
        centering the title. Default is 60.

        Returns
        -------
        None
            Nothing
        """
        t = '-' * line_length
        t += '\n' + msg.center(line_length) + '\n'
        t += '-' * line_length
        print(t)

    @staticmethod
    def log(msg, indent=0):
        """Utility function for consistently displaying messages in the output.

        Parameters
        ----------
        msg: str
            Message to be displayed in the output.
        indent: int
            Number of characters used for indentation of the message. Default is 0.
        Returns
        -------
        None
            Nothing
        """
        print(' ' * indent + msg)

    def __progress_log_basic(self, remark=None):
        line = f'iter: {self.it:{int(np.ceil(np.log10(self.iterations))) + 1}d}/{self.iterations}'
        line += f', eval: {self.evaluations_count:10d}'
        if self.maximum_evaluations:
            line += f'/{self.maximum_evaluations}'
        for io, ol in enumerate(self.objective_labels):
            line += f', {ol}: [blue]{self.best.O[io]}[/blue]'
        for cl, cv in zip(self.constraint_labels, self.best.C):
            if cv > 0:
                line += f', {cl}: [red]{cv}[/red]'
            else:
                line += f', {cl}: [green]{cv}[/green]'
        line += f', fit: [magenta bold underline]{self.best.f}[/magenta bold underline]'
        if remark:
            line += f' {remark}'
        self.rich_console.print(line)
        #self.log(line)

    def tic(self):
        """Utility function for measuring the computational time. Optimizer.tic() is used for starting the stopwatch.
        It is used in pair with Optimizer.toc() method which stops the stopwatch.

        Returns
        -------
        Float
            Return the recorded current time in seconds since the Epoch.
        """
        self._time1 = time.time()
        return self._time1

    def toc(self, msg='', silent=False):
        """Utility function for measuring the computational time. Optimizer.toc() stops the stopwatch and returns
        time elapsed since Optimizer.tic() was triggered.

        Parameters
        ----------
        msg: str, default=''
            A message which is displayed with the measured time. Not used if silent=True.
        silent: bool, default=True
            An option for disabling the output. Default is False. If True, the value of elapsed time is returned but no
            message is shown.

        Returns
        -------
        elapsed_time: float
            Time elapsed since calling the last Optimizer.toc() in seconds.
        """

        s = time.time() - self._time1
        
        if not silent:
            if msg == '':
                msg = 'Elapsed time:'

            if s < 1:
                dur = f'{s*1e3:.3f} milliseconds.'
            else:
                dur = f'{s:.3f} seconds.'
            self.log(f'{msg} {dur}')
            
        return s

    def _gen_unique_str(self):
        """A private method for generating a unique string of 16 characters used when forwarding a unique string to the
        evaluation function is enabled (Optimizer.forward_unique_str = True). The generated string is guaranteed to be
        unique within an optimizer.

        Returns
        -------
        s: str
            A unique string.
        """
        while True:
            _s = ''.join(np.random.choice([c for c in string.ascii_lowercase]) for _ in range(16))
            if _s not in self._unique_str_list:
                self._unique_str_list.append(_s)
                return _s

    def optimize(self, seed=None):
        """A method which starts the optimization. The method wraps :code:`_run` method of the optimizer's subclass.

        Parameters
        ----------
        seed: int or None
            A random seed. Use the same value for reproducing identical stochastic procedures.

        Returns
        -------
        optimum: OptimizationResults
            The best solution found in the optimization.
        """
        self._init_optimizer()
        np.random.seed(seed)

        if self.monitoring == 'dashboard':
            with self.live:
                r = self._run()
        else:
            r = self._run()

        return r

    def _stopping_criteria(self):
        """A private method used to evaluate all specified stopping conditions.

        Returns
        -------
        stop: bool
            A :code:`True` or :code:`False`, whether the optimization should stop.
        """

        # Stop if fitness threshold achieved
        if self.target_fitness:
            if self.best.f < self.target_fitness and np.all(self.best.C <= 0):
                return True

        # Stop if maximum number of evaluations surpassed
        if self.maximum_evaluations:
            if self.evaluations_count >= self.maximum_evaluations:
                return True

        return False

    def __getstate__(self):
        self_dict = self.__dict__.copy()
        del self_dict['pool']
        return self_dict

    def __setstate__(self, state):
        self.__dict__.update(state)

    def _multiprocess_evaluate(self, candidates):
        """A private method used for calling parallel evaluation of multiple candidates. It relies on multiprocessing
        pool using map and starmap methods. The objectives, constraints and fitness of candidates in the list is updated
        after the evaluation.

        Parameters
        ----------
        candidates: list of CandidateState
            A list of candidates to be evaluated.

        Returns
        -------
        candidates: list of CandidateState
            A list of evaluated candidates.
        """

        if self.forward_unique_str:
            result = self.pool.starmap(self.evaluation_function,
                                       [tuple([c.X, self._gen_unique_str()]) for c in candidates])
        else:
            result = self.pool.map(self.evaluation_function, [c.X for c in candidates])

        if self.objectives == 1 and self.constraints == 0:
            # Fast evaluation
            for p in range(len(candidates)):
                candidates[p].f = result[p]
                candidates[p].O[0] = result[p]
        else:
            # Full evaluation
            for p in range(len(candidates)):
                for io in range(self.objectives):
                    candidates[p].O[io] = result[p][io]
                candidates[p].f = np.dot(candidates[p].O, self.objective_weights)
                for ic in range(self.constraints):
                    candidates[p].C[ic] = result[p][self.objectives + ic]

        return candidates  # TODO do we need to return candidates or they can be updated in-place?

    def collective_evaluation(self, candidates):
        """Function used for evaluation of multiple candidates which automatically conducts parallel or serial
        evaluation and forwards a unique string to the evaluation function. Evaluation is performed in-place and the
        candidates provided via argument are updated.

        Parameters
        ----------
        candidates: list of CandidateState
            A list of candidates to be evaluated.

        Returns
        -------
        None
            Nothing
        """

        n = len(candidates)
        self.evaluations_count += n

        if self.number_of_processes > 1:
            self._multiprocess_evaluate(candidates)
        else:
            if self.forward_unique_str:
                for p in range(n):
                    candidates[p].evaluate(self._gen_unique_str())
            else:
                for p in range(n):
                    candidates[p].evaluate()

        # Determine new best candidate state
        if self.best is None:
            self.best = np.min(candidates).copy()
            self.results.cHistory.append((self.it, self.evaluations_count, self.best))
        else:
            # self.best = np.min(np.append(candidates, [self.best])).copy()
            candidates_best = np.min(candidates)
            if candidates_best <= self.best:
                self.best = candidates_best.copy()
                self.results.cHistory.append((self.it, self.evaluations_count, self.best))

        if self.monitoring == 'dashboard':
            self.live.update(self.__update_progress_bar())

    def multiprocess_evaluate_test(self, number_of_candidates, seed=None):
        # Function for testing purpose only - remove it and use it in the test sources!

        self._init_optimizer()
        if None:
            np.random.seed() # Possible fix for some seed issues in parallel execution
        else:
            np.random.seed(seed)

        candidates = np.array([CandidateState(self) for _ in range(number_of_candidates)], dtype=CandidateState)

        for p in range(number_of_candidates):
            candidates[p].X = np.random.uniform(self.lb, self.ub)

        self.collective_evaluation(candidates)

        return candidates


class CandidateState:
    """Candidate solution for optimization problem"""
    
    def __init__(self, optimizer: Optimizer):
        self._optimizer = optimizer
        self.X = np.zeros(optimizer.dimensions)
        self.O = np.zeros(optimizer.objectives) * np.nan
        self.C = np.zeros(optimizer.constraints) * np.nan
        self.f = np.nan
                    
        # Evaluation
        if self._optimizer.objectives == 1 and self._optimizer.constraints == 0:
            self._evaluate = self._eval_fast
        else:
            self._evaluate = self._eval_full
        
        # Comparison operators
        if self._optimizer.objectives == 1 and self._optimizer.constraints == 0:
            self._eq_fn = CandidateState._eq_fast
            self._lt_fn = CandidateState._lt_fast 
            #self.__gt__ = self._gt_fast
        else:
            self._eq_fn = CandidateState._eq_full
            self._lt_fn = CandidateState._lt_full 
            #self.__gt__ = self._gt_full
        
        
    def copy(self):
        
        cP = CandidateState(self._optimizer)        
        # Real copy for ndarrays
        cP.X = np.copy(self.X)
        cP.O = np.copy(self.O)
        cP.C = np.copy(self.C)
        cP.f = self.f
        return cP
    
    def __str__(self): 
        nch = np.max([np.size(self.X), np.size(self.O), np.size(self.C)]) 
        nch = np.min([8, nch])
        #print(nch)
        s = '-' * (12 * nch + nch + 20)
        
        for i in range(int(np.ceil(np.size(self.X)/nch))):
            if i == 0:
                s += '\n' + 'X: '.rjust(20)
            else:
                s += '\n' + ' ' * 20
            s += ' '.join([f'{x:12.5e}' for x in self.X[i*nch:(i+1)*nch]])

        all_labels = self._optimizer.objective_labels + self._optimizer.constraint_labels
        label_len = np.max([len(lbl) for lbl in all_labels]) + 3
        for o, o_label in zip(self.O, self._optimizer.objective_labels):
            s += '\n' + f'{o_label}: '.rjust(label_len) + f'{o:12.5e}'
        for c, c_label in zip(self.C, self._optimizer.constraint_labels):
            s += '\n' + f'{c_label}: '.rjust(label_len) + f'{c:12.5e}'
        s += '\n' + 'Fitness: '.rjust(label_len) + f'{self.f:12.5e}' + \
             '\n' + '-' * (12 * nch + nch + 20)
        return s
    
    # Equality operator
    def __eq__(self, other): 
        return self._eq_fn(self, other) 
    
    @staticmethod
    def _eq_fast(a, b): 
        return a.f == b.f

    @staticmethod
    def _eq_full(a, b):
        return np.sum(np.abs(a.X - b.X)) + np.sum(np.abs(a.O - b.O)) + np.sum(np.abs(a.C - b.C)) == 0.0

    # Inequality operator
    def __ne__(self, other):
        return self.f != other.f
    
    # Less-then operator
    def __lt__(self, other):
        return self._lt_fn(self, other)    

    def evaluate(self, s=None):
        return self._evaluate(s)

    @staticmethod
    def _lt_fast(a, b): 
        #print('_lt_fast')
        return a.f < b.f   
    
    @staticmethod     
    def _lt_full(a, b):  
        #print("_lt_full")               
        if np.sum(a.C > 0) == 0 and np.sum(b.C > 0) == 0: 
            # Both are feasible
            # Better candidate is the one with smaller fitness
            return a.f < b.f
        
        elif np.sum(a.C > 0) == np.sum(b.C > 0): 
            # Both are unfeasible and break same number of constraints
            # Better candidate is the one with smaller sum of constraint values
            return np.sum(a.C) < np.sum(b.C)
        
        else:
            # The number of unsatisfied constraints is not the same
            # Better candidate is the one which breaks fewer constraints
            return np.sum(a.C > 0) < np.sum(b.C > 0)
       
    
    def __gt__(self, other):
        #print('__gt__')
        #return self.f > other.f
        return not (self._lt_fn(self, other) or self._eq_fn(self, other))
    """
    def _gt_fast(self, other):
        return self.f > other.f
    def _gt_full(self, other): 
        return not (self.__eq__(other) or self.__lt__(other))
    """    
    
    def __le__(self, other):
        return self._lt_fn(self, other) or self.__eq__(other)
    
    def __ge__(self, other):
        return self.__gt__(other) or self.__eq__(other)

    def set_X(self, X: np.ndarray):

        assert np.size(self.X) == np.size(X), 'Unexpected size of optimization vector X'
        self.X = X
        
    def _eval_full(self, s=None):

        if s is None:
            oc = self._optimizer.evaluation_function(self.X)
        else:
            oc = self._optimizer.evaluation_function(self.X, s)

        for io in range(self._optimizer.objectives):
            self.O[io] = oc[io]
        
        self.f = np.dot(self.O, self._optimizer.objective_weights)
        
        for ic in range(self._optimizer.constraints):
            self.C[ic] = oc[self._optimizer.objectives + ic]

        return self.O, self.C, self.f
        
    def _eval_fast(self, s=None):

        if s is None:
            f = self._optimizer.evaluation_function(self.X)
        else:
            f = self._optimizer.evaluation_function(self.X, s)

        self.f = f
        self.O[0] = f
        return self.O, self.C, self.f


class OptimizationResults:
    """Data holder for optimization results."""
    
    def __init__(self, optimizer: Optimizer):
        
        self.optimizer = optimizer
        self.cHistory = []
        
    def plot_convergence(self, axes=None):
        """Plot convergence

        Parameters
        ----------
        axes :

        Returns
        -------

        """
        
        if axes is None:
            fig = plt.figure(constrained_layout=True)
            spec = gridspec.GridSpec(ncols=1, nrows=2, figure=fig)
            axo = fig.add_subplot(spec[0])
            axc = fig.add_subplot(spec[1], sharex=axo)
        else:
            axo, axc = axes
        
        for io in range(self.optimizer.objectives):
            O = np.array([[ev, cB.O[io]] for (it, ev, cB) in self.cHistory])
            axo.plot(O[:, 0], O[:, 1], label=self.optimizer.objective_labels[io], lw=1)
            
        F = np.array([[ev, cB.f] for (it, ev, cB) in self.cHistory])
        axo.plot(F[:, 0], F[:, 1], label='f', lw=2, ls='-', c='b')
        axo.legend()
        
        for ic in range(self.optimizer.constraints):
            C = np.array([[ev, cB.C[ic]] for (it, ev, cB) in self.cHistory])
            axc.plot(C[:, 0], C[:, 1], label=self.optimizer.constraint_labels[ic], lw=1)

            IX = C[:, 1] > 0
            C = C[IX, :]

            # nc = np.sum(C>0)
            # axc.plot(I[:nc], C[:nc], c='r', ls='-', lw=2)
            # if nc == 1:
            #     axc.plot(I[:nc], C[:nc],'ro')
                            
        axc.legend()       
        axo.set_ylabel('Objectives')
        axc.set_xlabel('Evaluations')
        axc.set_ylabel('Constraints')
        #ax.set_yscale('log')
        