# -*- coding: utf-8 -*-
"""
Created on Tue Jul 12 09:07:40 2022

@author: janpe
"""

from .fancy_classes import Graph
from .saver import saver
from .lossfunctions import loss_dic
import numpy as np
from scipy import optimize

import logging
log = logging.getLogger(__name__)

class topological_opti:

    def __init__(self, start_graph: Graph, saver: saver, ent_dic=None, target_state=None,
                 config=None, safe_history=True):

        self.config = config
        self.imaginary = self.config['imaginary']
        if self.config['loss_func'] == 'ent':
            self.ent_dic = ent_dic
        else:
            self.target = target_state  # object of State class

        # do preoptimization on complete starting graph, this might already take some time
        self.graph = self.pre_optimize_start_graph(start_graph)
        self.saver = saver
        self.save_hist = safe_history
        self.history = []

    def check(self, result: object, lossfunctions: object):
        """
        check if all loss functions fulfill conditions for success. mostly defined through thresholds.

        Parameters
        ----------
        result : object
             from scipy.mimnimizer class
        lossfunctions : object
            list of all loss functions

        Returns
        -------
        bool
            False if we keep Graph or True if we can delete edge

        """

        if self.config['loss_func'] == 'ent':
            if abs(result.fun) - abs(self.loss_val[0]) > self.config['thresholds'][0]:
                return False
        else:
            # uncomment to see where checks fail
            # print(result.fun, self.config['thresholds'][0])
            if result.fun > self.config['thresholds'][0]:
                #if check fails return false
                return False
            # check if all loss functions are under the corresponding threshold
            for ii in range(1, len(lossfunctions)):
                if lossfunctions[ii](result.x) > self.config['thresholds'][ii]:
                    # if check fails return false
                    return False
        # when no check fails return True  = success
        return True

    def weights_to_valid_input(self, weights: list) -> list:
        """
        need to change weights from scip optimizer to proper input 
        if one optimize with complexe values

        Parameters
        ----------
        weights : list
            list of weights

        Raises
        ------
        ValueError
            if imaginary is not defined correctly in config file

        Returns
        -------
        list
            ordered list according to imaginary

        """
        if self.imaginary ==  'cartesian':

            if len(weights) % 2 == 0:
                ll2 = int(len(weights)/2)
            else:
                raise ValueError(
                    'odd number of weights for complex optimization')

            return [complex(real, imag) for real, imag in
                    zip(weights[:ll2], weights[ll2:])]

        elif self.imaginary == 'polar':

            if len(weights) % 2 == 0:
                ll2 = int(len(weights)/2)
            else:
                raise ValueError(
                    'odd number of weights for complex optimization')
            return [(radius, phase) for radius, phase in
                    zip(weights[:ll2], weights[ll2:])]
        elif self.imaginary is False:
            return weights
        else:
            raise ValueError('imaginary: only: polar,cartesian,false')

    def get_loss_functions(self, current_graph: Graph):
        """
        get a list of all loss functions mentioned in config

        Parameters
        ----------
        current_graph

        Returns
        -------
        callable_loss
            list of callable functions

        """
        # get loss function acc. to config
        lossfunctions = loss_dic[self.config['loss_func']]

        # entanglement loss function
        if self.config['loss_func'] == 'ent':  # we optimize for entanglement
            loss_specs = {'sys_dict': self.ent_dic,
                          'imaginary': self.imaginary,
                          'var_factor': self.config['var_factor']}

        # CR and FID
        elif self.config['loss_func'] in ['cr', 'fid']:
            loss_specs = {'target_state': self.target,
                          'cnfg': self.config}

        # fock basis
        elif self.config['loss_func'] in ['fockcr','fockfid']:
            #loss_specs = {'target_state': self.target,
            #              'cnfg': self.config}
            loss_specs = {'target_state': self.target,
                        'num_anc':self.config['num_anc'],
                        'amplitudes':self.config['amplitudes'],
                        'imaginary': self.imaginary}

        # custom loss functions
        elif self.config['loss_func'] == 'lff':
            loss_specs = {'cnfg': self.config}
        callable_loss = [func(current_graph, **loss_specs)
                         for func in lossfunctions]
        return callable_loss

    def update_losses(self, result, losses):
        """
        updates the losses for next steps

        Parameters
        ----------
        result : scipy.minimizer object
            minimize object from optimizaiton step
        losses : list
            list of loss functions

        Returns
        -------
        list of losses for corrosponding weights stored in result.x

        """
        loss_values = [result.fun]
        for ii in range(1, len(losses)):
            loss_values.append(losses[ii](result.x))
        return loss_values

    def pre_optimize_start_graph(self, graph) -> Graph:
        """
        first optimization of complete starting graph

        Parameters
        ----------
        graph

        Returns
        -------
        preopt_graph

        """
        # losses is a list of callable lossfunctions, e.g. [countrate(x), fidelity(x)], where x is a vector of edge weights
        # that can be given to scipy.optimize
        losses = self.get_loss_functions(graph)
        valid = False
        counter = 0
        # repeat optimization of complete graph until a good solution is found (which satifies self.check())
        while not valid:
            # prepare optimizer
            initial_values, bounds = self.prepOptimizer(len(graph))
            # optimization with scipy
            best_result = optimize.minimize(losses[0], x0=initial_values,
                                            bounds=bounds,
                                            method=self.config['optimizer'],
                                            options={'ftol': self.config['ftol']})
            self.loss_val = self.update_losses(best_result, losses)
            #check if solution is valid
            valid = self.check(best_result, losses)
            counter += 1
            #print a warning if preoptimization is stuck in a loop
            if counter % 10 == 0:
                print('10 invalid preoptimization, consider changing parameters.')
                log.info('10 invalid preoptimization, consider changing parameters.')

        # if num_pre is set to larger than 1 in config, do num_pre preoptimization and choose the best one.
        # for optimizations with concrete target state, num_pre = 1 is enough
        for __ in range(self.config['num_pre'] - 1):
            initial_values, bounds = self.prepOptimizer(len(graph))
            result = optimize.minimize(losses[0], x0=initial_values,
                                       bounds=bounds,
                                       method=self.config['optimizer'],
                                       options={'ftol': self.config['ftol']})

            if result.fun < best_result.fun:
                best_result = result
        self.loss_val = self.update_losses(best_result, losses)
        print(f'best result from pre-opt: {abs(best_result.fun)}')
        log.info(f'best result from pre-opt: {abs(best_result.fun)}')

        preopt_graph = Graph(graph.edges,
                             weights=self.weights_to_valid_input(
                                 best_result.x),
                             imaginary=self.imaginary)

        try:
            bulk_thr = self.config['bulk_thr']
        except:
            bulk_thr = 0
        if bulk_thr > 0:
            # cut all edges smaller than bulk_thr and optimize again
            # this can save a lot of time
            cont = True
            num_deleted = 0
            while cont:
                # delete smallest edges one by one
                idx_of_edge = preopt_graph.minimum()
                amplitude = preopt_graph[idx_of_edge]
                if self.imaginary == 'polar':
                    amplitude = amplitude[0]
                if abs(amplitude) < bulk_thr:
                    preopt_graph.remove(idx_of_edge)
                    num_deleted += 1
                else:
                    cont = False
            print(f'{num_deleted} edges deleted')
            log.info(f'{num_deleted} edges deleted')
            valid = False
            while not valid:
                # it is necessary that the truncated graph passes the checks
                initial_values, bounds = self.prepOptimizer(len(preopt_graph))
                losses = self.get_loss_functions(preopt_graph)
                trunc_result = optimize.minimize(losses[0], x0=initial_values,
                                                 bounds=bounds,
                                                 method=self.config['optimizer'],
                                                 options={'ftol': self.config['ftol']})
                self.loss_val = self.update_losses(trunc_result, losses)
                print(f'result after truncation: {abs(trunc_result.fun)}')
                log.info(f'result after truncation: {abs(trunc_result.fun)}')
                valid = self.check(trunc_result, losses)
            preopt_graph = Graph(preopt_graph.edges,
                                 weights=self.weights_to_valid_input(
                                     trunc_result.x),
                                 imaginary=self.imaginary)

        return preopt_graph

    def prepOptimizer(self, numweights, x=[]):
        """
        returns initial values and bounds for use in optimization.

        Parameters
        ----------
        numweights
        x

        Returns
        -------
        initial_values
        bounds
        """

        if not self.imaginary:
            bounds = numweights * [(-1, 1)]
            if len(x) == 0:
                initial_values = 2 * np.random.random(numweights) - 1
            else:
                initial_values = x
        else:
            bounds = numweights * [(-1, 1)] + numweights * [(-np.pi, np.pi)]
            if len(x) == 0:
                rands_r = 2 * np.random.random(numweights) - 1
                rands_th = 2 * np.pi * np.random.random(numweights) - np.pi
                initial_values = np.concatenate([rands_r, rands_th])
            else:
                initial_values = x

        return initial_values, bounds

    def continuationCondition(self, num_edge) -> bool:
        """
        conditions that stop optimization

        Parameters
        ----------
        num_edge

        Returns
        -------
        cont: bool
            if True, topological optimization is continued
        """
        if self.config['loss_func'] == 'ent':
            cont = len(self.graph) > self.config['min_edge'] and num_edge < len(
                self.graph)
        else:
            # if num_edge is higher than total number of edges in the graph
            # or higher than edges_tried, return False
            cont = num_edge < min(len(self.graph), self.config['edges_tried'])
        return cont

    def optimize_one_edge(self, num_edge: int,
                          num_tries_one_edge: int) -> (Graph, bool):
        """
        delete the num_edge-th smallest edge and optimize num_tries_one_edge times
        and check if corresponding loss function fulfills checks

        Parameters
        ----------
        num_edge
        num_tries_one_edge

        Returns
        -------
        new_graph
           if edge is successfully deleted, reduced graph is returned. else graph is not modified.
        success
        """
        # copy current Graph and delete num_edge´s smallest weight
        # set up reduced graph
        reduced_graph = self.graph.copy()
        #find index of num_edge smallest edge
        idx_of_edge = reduced_graph.minimum(num_edge)
        #store amplitude in case edge fails and needs to be put back in
        amplitude = reduced_graph[idx_of_edge]
        #remove smallest edge
        reduced_graph.remove(idx_of_edge, update=False)
        #update state catalog of reduced graph
        reduced_graph.getStateCatalog()
        #try a given number of times to delete this edge
        for ii in range(num_tries_one_edge):
            # if edge is tried for the first time, update loss function and other optimization parameters
            # also the weights of the old graph are used as initial values (this is much faster)
            if ii == 0:
                try:
                    # redefine loss functions for reduced graph
                    losses = self.get_loss_functions(reduced_graph)
                    # use initial values x0 from previous Graph
                    x0 = reduced_graph.weights
                    initial_values, bounds = self.prepOptimizer(len(reduced_graph),
                                                                x=x0)
                    # optimize with scipy
                    result = optimize.minimize(losses[0], x0=initial_values,
                                               bounds=bounds,
                                               method=self.config['optimizer'],
                                               options={'ftol': self.config['ftol']})
                except KeyError:
                    # if the target kets can not be produced with the given graph we can give up on this edge
                    # it wont work
                    reduced_graph[idx_of_edge] = amplitude
                    print('edge necessary for producing all kets')
                    log.info('edge necessary for producing all kets')
                    return reduced_graph, False  # no success keep current Graph
            # if edge has been tried before, reuse loss function etc. and use random initial values
            else:
                # random initial values
                initial_values, bounds = self.prepOptimizer(len(reduced_graph))
                # optimize with scipy
                result = optimize.minimize(losses[0], x0=initial_values,
                                           bounds=bounds,
                                           method=self.config['optimizer'],
                                           options={'ftol': self.config['ftol']})
            # check if solution is valid
            valid = self.check(result, losses)

            if valid:  # if criterion is reached then save reduced graph as graph, else leave graph as is
                # compute values for all losses
                self.loss_val = self.update_losses(result, losses)
                # if clean solution is encountered before optimization finishes, save that too
                # all weights are +-1 and solution is pretty good, it is interesting enough to be saved
                # to a file even if topological optimization can be continued
                if all(np.array(abs(self.graph)) > 0.95):
                    self.saver.save_graph(self)
                if self.save_hist:
                    self.history.append(self.loss_val)
                # return updated result graph
                return Graph(reduced_graph.edges,
                             weights=self.weights_to_valid_input(result.x),
                             imaginary=self.imaginary), True
        # all tries failed keep current Graph
        reduced_graph[idx_of_edge] = amplitude
        return reduced_graph, False

    def topologicalOptimization(self, save_hist=True) -> Graph:
        """
        does the topological main loop. deletes edges until continuation condition fails.
        Returns
        -------
        solution_graph
            result of optimization
        """
        #start with smallest edge
        num_edge = 0
        graph_history = []
        # if num_edge becomes too large, optimization is stopped
        while self.continuationCondition(num_edge):
            # try if num_edge smallest edge can be removed
            # if successful, return optimized graph after deletion
            # if not successful, return old graph
            self.graph, success = self.optimize_one_edge(
                num_edge, self.config['tries_per_edge'])
            #iterate num_edge to try next smallest edge
            num_edge += 1
            log.info(f'deleting edge {num_edge}')
            print(f'deleting edge {num_edge}')
            if success:
                print(
                    f"deleted: {len(self.graph)}  edges left with loss {self.loss_val[0]:.3f}")
                log.info(f"deleted: {len(self.graph)}  edges left with loss {self.loss_val[0]:.3f}")
                #reset to try smallest edge again for next iteration
                num_edge = 0
                graph_history.append(self.graph)

        return self.graph
