# Copyright (c) 2015, Disney Research

# All rights reserved.

#

# Author(s): Sehoon Ha <sehoon.ha@disneyresearch.com>

# Disney Research Robotics Group

"""

Example Usage:

# prob has f() and g() as member functions

solver = Solver(prob)

res = solver.solve()

x = res['x']

f = res['f']

"""

import cmaes.utils

import numpy as np

class Solver(object):

    def __init__(self, prob):

        self.prob = prob

        self.eval_counter = 0

        self.numerical_diff_step = 1e-4

        self.iter_values = list()

        self.verbose = True

        self.check_gradient = False

    def get_check_gradient(self, ):

        return self.check_gradient

    def set_check_gradient(self, check_gradient=True):

        self.check_gradient = check_gradient

    def eval_f(self, x):

        ret = self.prob.f(x)

        self.eval_counter += 1

        self.last_x = x

        self.last_f = ret

        self.iter_values.append(ret)

        if hasattr(self.prob, 'on_eval_f'):

            self.prob.on_eval_f(self)

        return ret

    def eval_g(self, x):

        if hasattr(self.prob, 'g'):

            ret = self.prob.g(x)

            if self.get_check_gradient():

                h = self.numerical_diff_step

                ret2 = utils.grad(self.prob.f, x, h)

                isgood = np.allclose(ret, ret2, atol=1e-05)

                if not isgood:

                    print(ret)

                    print(ret2)

                    print("diff = %.12f" % np.linalg.norm(ret - ret2))

                print('check_gradient... %s' % isgood)

        else:

            h = self.numerical_diff_step

            ret = utils.grad(self.prob.f, x, h)

        return ret

    def eval_c_eq_jac(self, x, i):

        if hasattr(self.prob, 'c_eq_jac'):

            ret = self.prob.c_eq_jac(x, i)

        else:

            h = self.numerical_diff_step

            def c_eq_f(x):

                return self.prob.c_eq(x, i)

            ret = utils.grad(c_eq_f, x, h)

        return ret

    def eval_c_ineq_jac(self, x, i):

        if hasattr(self.prob, 'c_ineq_jac'):

            ret1 = self.prob.c_ineq_jac(x, i)

            return ret1

        else:

            h = self.numerical_diff_step

            def c_ineq_f(x):

                return self.prob.c_ineq(x, i)

            ret2 = utils.grad(c_ineq_f, x, h)

            return ret2

    def collect_constraints(self):

        constraints = list()

        if hasattr(self.prob, 'num_eq_constraints'):

            num = self.prob.num_eq_constraints()

            if self.verbose:

                print('  [Solver]: num_eq_constraints = %d' % num)

            for i in range(num):

                c = dict()

                c['type'] = 'eq'

                assert(hasattr(self.prob, 'c_eq'))

                c['fun'] = self.prob.c_eq

                c['jac'] = self.eval_c_eq_jac

                c['args'] = [i]

                constraints.append(c)

        if hasattr(self.prob, 'num_ineq_constraints'):

            num = self.prob.num_ineq_constraints()

            if self.verbose:

                print('  [Solver]: num_ineq_constraints = %d' % num)

            for i in range(num):

                c = dict()

                c['type'] = 'ineq'

                assert(hasattr(self.prob, 'c_ineq'))

                c['fun'] = self.prob.c_ineq

                c['jac'] = self.eval_c_ineq_jac

                c['args'] = [i]

                constraints.append(c)

        if self.verbose:

            print('  [Solver]: num_constraints = %d' % len(constraints))

        return constraints

    def bounds(self):

        if hasattr(self.prob, 'bounds'):

            return self.prob.bounds()

        else:

            return None

    def solve(self, x0=None):

        pass

    def set_verbose(self, verbose):

        self.verbose = verbose

    def plot_convergence(self, filename=None):

        yy = self.iter_values

        xx = range(len(yy))

        import matplotlib.pyplot as plt

        # Plot

        plt.ioff()

        fig = plt.figure()

        fig.set_size_inches(18.5, 10.5)

        font = {'size': 28}

        plt.title('Value over # evaluations')

        plt.xlabel('X', fontdict=font)

        plt.ylabel('Y', fontdict=font)

        plt.plot(xx, yy)

        plt.axes().set_yscale('log')

        if filename is None:

            filename = 'plots/iter.png'

        plt.savefig(filename, bbox_inches='tight')

        plt.close(fig)

        print('plotting convergence OK.. ' + filename)

    def save_result(self, res, filename):

        with open(filename, 'w+') as fin:

            fin.write(str(res))

        print('writing result OK.. ' + filename)