# Copyright (c) 2015, Disney Research
# All rights reserved.
#
# Author(s): Sehoon Ha <sehoon.ha@disneyresearch.com>
# Disney Research Robotics Group
#
# adapted by Seungmoon Song <seungmoon.song@gmail.com>

from __future__ import division  # '/' always means non-truncating division
from cmaes.solver import Solver
import numpy as np
import cma
import scipy.optimize
import time
from datetime import datetime
import sys


class CMASolver(Solver):
    def __init__(self, prob):
        Solver.__init__(self, prob)
        opts = cma.CMAOptions()
        # for k, v in opts.iteritems():
        #     print k, v
        # exit(0)
        self.p_dir = 'optim_data/cma/'
        opts.set('verb_disp', 1)
        opts.set('popsize', 8)
        opts.set('verb_filenameprefix', self.p_dir)
        opts.set('maxiter', 2000)
        self.options = opts
        self.cen = None
        self.rng = None

    def set_verbose(self, verbose):
        self.verbose = verbose
        if verbose:
            self.options['verb_disp'] = 1
        else:
            self.options['verb_disp'] = 0

    def create_directory(self):
        verbose = (self.options['verb_disp'] > 0)
        import os
        path = self.p_dir
        if verbose:
            print('cma path = ', path)

        if not os.path.exists(path):
            if verbose:
                print('CMA-ES: create directory [%s]' % path)
            os.makedirs(path)

    def eval_f(self, y):
        x = self.unnormalize(y)
        ret = super(CMASolver, self).eval_f(x)

        # for i in range(self.prob.num_eq_constraints()):
        #     ret_eq_i = self.prob.c_eq(x, i)
        #     # ret += 100.0 * (ret_eq_i ** 2)
        #     ret += 10.0 * (ret_eq_i)  # Assume the quadratic form
        # for i in range(self.prob.num_ineq_constraints()):
        #     ret_ineq_i = self.prob.c_ineq(x, i)
        #     if ret_ineq_i < 0:
        #         ret += 100.0 * (ret_ineq_i ** 2)
        return ret

    def clip(self, x):
        if self.rng is None:
            return x
        return np.clip(x, self.cen-self.rng, self.cen+self.rng)

    # normalize between [-1, 1]
    def normalize(self, x):
        if self.rng is None:
            return x
        return (x - self.cen) / self.rng

    def unnormalize(self, y):
        if self.rng is None:
            return y
        x = self.cen + y * self.rng
        return x

    def solve(self, x0=None, sigma=1.0):
        verbose = (self.options['verb_disp'] > 0)
        begin = time.time()
        if verbose:
            print('Optimization method = CMA-ES')
        if x0 is None:
            if verbose:
                print('Optimization: set x0 as zeros')
            if self.cen is not None:
                x0 = self.cen
            else:
                x0 = np.zeros(self.prob.dim)
        self.create_directory()
        if verbose:
            print('CMA-ES: cen = ', self.cen)
            print('CMA-ES: rng = ', self.rng)
            print('Optimization begins at ', str(datetime.now()))
            #print('normalized_center = ', self.normalize(x0))
            # for k, v in self.options.iteritems():
            #     print(k, '\t', v)

        res = cma.fmin(self.eval_f,
                       self.normalize(x0),
                       sigma,
                       options=self.options)
        if verbose:
            print('Optimization ends at ', str(datetime.now()))
            print('Total times = %.2fs' % (time.time() - begin))

        ret = scipy.optimize.OptimizeResult()
        ret['y'] = res[0]
        ret['x'] = self.unnormalize(res[0])
        ret['fun'] = res[1]
        # assert(np.allclose(res[1], self.prob.f(ret['x'])))
        ret['nfev'] = self.eval_counter
        # ret['jac'] = self.eval_g(ret['x'])
        ret['message'] = 'Optimization terminated successfully.'
        ret['status'] = 0
        ret['success'] = True
        return ret


class CMASolverPar(CMASolver):
    def solve(self, x0=None, sigma=1.0):
        verbose = (self.options['verb_disp'] > 0)
        begin = time.time()
        if verbose:
            print('Optimization method = CMA-ES')
        if x0 is None:
            if verbose:
                print('Optimization: set x0 as zeros')
            if self.cen is not None:
                x0 = self.cen
            else:
                x0 = np.zeros(self.prob.dim)
        self.create_directory()
        if verbose:
            print('CMA-ES: cen = ', self.cen)
            print('CMA-ES: rng = ', self.rng)
            print('Optimization begins at ', str(datetime.now()))
            #print('normalized_center = ', self.normalize(x0))
            # for k, v in self.options.iteritems():
            #     print(k, '\t', v)

        res = cma.fmin(None,
                       self.normalize(x0),
                       sigma,
                       parallel_objective=self.eval_f,
                       options=self.options)
        if verbose:
            print('Optimization ends at ', str(datetime.now()))
            print('Total times = %.2fs' % (time.time() - begin))

        ret = scipy.optimize.OptimizeResult()
        ret['y'] = res[0]
        ret['x'] = self.unnormalize(res[0])
        ret['fun'] = res[1]
        # assert(np.allclose(res[1], self.prob.f(ret['x'])))
        ret['nfev'] = self.eval_counter
        # ret['jac'] = self.eval_g(ret['x'])
        ret['message'] = 'Optimization terminated successfully.'
        ret['status'] = 0
        ret['success'] = True
        return ret


if __name__ == '__main__':
    import optimization.test_problems
    import numpy as np
    # prob = test_problems.QuadProb()
    prob = optimization.test_problems.Rosen()
    x0 = np.random.rand(prob.dim) - 0.5

    solver = CMASolver(prob)
    res = solver.solve(x0)
    print(res)