# Copyright (c) 2015, Disney Research
# All rights reserved.
#
# Author(s): Sehoon Ha <sehoon.ha@disneyresearch.com>
# Disney Research Robotics Group


import numpy as np


class QuadProb(object):
    def __init__(self, x=None):
        self.dim = 4
        self.x = np.array([0.5, 0.7, -0.5, 1.0])
        if x is not None:
            self.dim = len(x)
            self.x = x

    def f(self, x):
        diff = x - self.x
        return 0.5 * np.linalg.norm(diff) ** 2


class Rosen(object):
    def __init__(self):
        self.dim = 10

    def f(self, x):
        import scipy.optimize
        return scipy.optimize.rosen(x)

    def bounds(self):
        return [(-0.5, 0.5)] * self.dim

    def on_eval_f(self, solver):
        counter = solver.eval_counter
        value = solver.last_f
        # if counter % 10 == 0:
        #     solver.plot_convergence()


class ConstrainedQuadProb(object):
    def __init__(self, x=None):
        self.dim = 4
        self.x = x if x is not None else np.array([0.5, 0.7, -0.5, 1.0])

    def f(self, x):
        diff = x - self.x
        return 0.5 * np.linalg.norm(diff) ** 2

    def num_eq_constraints(self):
        return 2

    def c_eq(self, x, index):
        if index == 0:
            return x[0] + x[1] - 1.0
        elif index == 1:
            return x[1] + x[2] - 1.0

    def c_eq_jac(self, x, index):
        """ (optional) """
        if index == 0:
            return np.array([1.0, 1.0, 0.0, 0.0])
        elif index == 1:
            return np.array([0.0, 1.0, 1.0, 0.0])

    def num_ineq_constraints(self):
        return 1

    def c_ineq(self, x, index):
        if index == 0:
            return x[3] - 1.2


if __name__ == '__main__':
    print('test optimization problems')
    prob = QuadProb(np.array([0.5, -0.3, 0.9]))
    print('prob.dim = %d' % (prob.dim))
    from optimization.solver_cma import CMASolver as Solver
    solver = Solver(prob)
    res = solver.solve()
    print(res)