"""
Copyright, the CVXPY authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""


import numpy as np
import pytest

import cvxpy as cp
from cvxpy.reductions.solvers.defines import INSTALLED_SOLVERS
from cvxpy.tests.nlp_tests.derivative_checker import DerivativeChecker
from cvxpy.tests.test_conic_solvers import is_knitro_available

# Always parametrize all solvers, skip at runtime if not available
NLP_SOLVERS = [
    pytest.param('IPOPT', marks=pytest.mark.skipif(
        'IPOPT' not in INSTALLED_SOLVERS, reason='IPOPT is not installed.')),
    pytest.param('KNITRO', marks=pytest.mark.skipif(
        not is_knitro_available(), reason='KNITRO is not installed or license not available.')),
    pytest.param('UNO', marks=pytest.mark.skipif(
        'UNO' not in INSTALLED_SOLVERS, reason='UNO is not installed.')),
    pytest.param('COPT', marks=pytest.mark.skipif(
        'COPT' not in INSTALLED_SOLVERS, reason='COPT is not installed.')),
]


@pytest.mark.parametrize("solver", NLP_SOLVERS)
class TestNLPExamples:
    """
    Nonlinear test problems taken from the IPOPT documentation and
    the Julia documentation: https://jump.dev/JuMP.jl/stable/tutorials/nonlinear/simple_examples/.
    """

    def test_hs071(self, solver):
        x = cp.Variable(4, bounds=[0, 6])
        x.value = np.array([1.0, 5.0, 5.0, 1.0])
        objective = cp.Minimize(x[0]*x[3]*(x[0] + x[1] + x[2]) + x[2])

        constraints = [
            x[0]*x[1]*x[2]*x[3] >= 25,
            cp.sum(cp.square(x)) == 40,
        ]
        problem = cp.Problem(objective, constraints)
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        assert np.allclose(x.value, np.array([0.75450865, 4.63936861, 3.78856881, 1.88513184]))

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    def test_mle(self, solver):
        n = 1000
        np.random.seed(1234)
        data = np.random.randn(n)

        mu = cp.Variable((1,), name="mu")
        mu.value = np.array([0.0])
        sigma = cp.Variable((1,), name="sigma")
        sigma.value = np.array([1.0])

        constraints = [mu == sigma**2]
        log_likelihood = (
            (n / 2) * cp.log(1 / (2 * np.pi * (sigma)**2))
            - cp.sum(cp.square(data-mu)) / (2 * (sigma)**2)
        )

        objective = cp.Maximize(log_likelihood)
        problem = cp.Problem(objective, constraints)
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        assert np.allclose(sigma.value, 0.77079388)
        assert np.allclose(mu.value, 0.59412321)

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    def test_portfolio_opt(self, solver):
        # data taken from https://jump.dev/JuMP.jl/stable/tutorials/nonlinear/portfolio/
        # r and Q are pre-computed from historical data of 3 assets
        r = np.array([0.026002150277777, 0.008101316405671, 0.073715909491990])
        Q = np.array([
            [0.018641039983891, 0.003598532927677, 0.001309759253660],
            [0.003598532927677, 0.006436938322676, 0.004887265158407],
            [0.001309759253660, 0.004887265158407, 0.068682765454814],
        ])
        x = cp.Variable(3)
        x.value = np.array([10.0, 10.0, 10.0])
        variance = cp.quad_form(x, Q)
        expected_return = r @ x
        problem = cp.Problem(
            cp.Minimize(variance),
            [
                cp.sum(x) <= 1000,
                expected_return >= 50,
                x >= 0
            ]
        )
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        # Second element can be slightly negative due to numerical tolerance
        assert np.allclose(x.value, np.array([4.97045504e+02, 0.0, 5.02954496e+02]), atol=1e-4)

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    def test_portfolio_opt_sum_multiply(self, solver):
        # data taken from https://jump.dev/JuMP.jl/stable/tutorials/nonlinear/portfolio/
        # r and Q are pre-computed from historical data of 3 assets
        r = np.array([0.026002150277777, 0.008101316405671, 0.073715909491990])
        Q = np.array([
            [0.018641039983891, 0.003598532927677, 0.001309759253660],
            [0.003598532927677, 0.006436938322676, 0.004887265158407],
            [0.001309759253660, 0.004887265158407, 0.068682765454814],
        ])
        x = cp.Variable(3)
        x.value = np.array([10.0, 10.0, 10.0])
        variance = cp.quad_form(x, Q)
        expected_return = cp.sum(cp.multiply(r, x))
        problem = cp.Problem(
            cp.Minimize(variance),
            [
                cp.sum(x) <= 1000,
                expected_return >= 50,
                x >= 0
            ]
        )
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        # Second element can be slightly negative due to numerical tolerance
        assert np.allclose(x.value, np.array([4.97045504e+02, 0.0, 5.02954496e+02]), atol=1e-4)

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    @pytest.mark.skipif('UNO' in INSTALLED_SOLVERS, reason='UNO has an algorithmic error')
    def test_rosenbrock(self, solver):
        x = cp.Variable(2, name='x')
        objective = cp.Minimize((1 - x[0])**2 + 100 * (x[1] - x[0]**2)**2)
        problem = cp.Problem(objective, [])
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        assert np.allclose(x.value, np.array([1.0, 1.0]))
        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    def test_qcp(self, solver):
        x = cp.Variable(1)
        y = cp.Variable(1, bounds=[0, np.inf])
        z = cp.Variable(1, bounds=[0, np.inf])

        objective = cp.Maximize(x)

        constraints = [
            x + y + z == 1,
            x**2 + y**2 - z**2 <= 0,
            x**2 - cp.multiply(y, z) <= 0
        ]
        problem = cp.Problem(objective, constraints)
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        assert np.allclose(x.value, np.array([0.32699284]))
        assert np.allclose(y.value, np.array([0.25706586]))
        assert np.allclose(z.value, np.array([0.4159413]))

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    def test_analytic_polytope_center(self, solver):
        # Generate random data
        np.random.seed(0)
        m, n = 50, 4
        b = np.ones(m)
        rand = np.random.randn(m - 2*n, n)
        A = np.vstack((rand, np.eye(n), np.eye(n) * -1))

        # Define the variable
        x = cp.Variable(n)
        # set initial value for x
        objective = cp.Minimize(-cp.sum(cp.log(b - A @ x)))
        problem = cp.Problem(objective, [])
        # Solve the problem
        problem.solve(solver=solver, nlp=True)

        assert problem.status == cp.OPTIMAL

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    def test_analytic_polytope_center_x_column_vector(self, solver):
        # Generate random data
        np.random.seed(0)
        m, n = 50, 4
        b = np.ones((m, 1))
        rand = np.random.randn(m - 2*n, n)
        A = np.vstack((rand, np.eye(n), np.eye(n) * -1))

        # Define the variable
        x = cp.Variable((n, 1))
        # set initial value for x
        objective = cp.Minimize(-cp.sum(cp.log(b - A @ x)))
        problem = cp.Problem(objective, [])
        # Solve the problem
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL

        checker = DerivativeChecker(problem)
        checker.run_and_assert()


    def test_socp(self, solver):
        # Define variables
        x = cp.Variable(3)
        y = cp.Variable()

        # Define objective function
        objective = cp.Minimize(3 * x[0] + 2 * x[1] + x[2])

        # Define constraints
        constraints = [
            cp.norm(x, 2) <= y,
            x[0] + x[1] + 3*x[2] >= 1.0,
            y <= 5
        ]

        # Create and solve the problem
        problem = cp.Problem(objective, constraints)
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        assert np.allclose(objective.value, -13.548638814247532)
        assert np.allclose(x.value, [-3.87462191, -2.12978826, 2.33480343])
        assert np.allclose(y.value, 5)

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    def test_portfolio_socp(self, solver):
        np.random.seed(858)
        n = 100
        x = cp.Variable(n, name='x')
        mu = np.random.randn(n)
        Sigma = np.random.randn(n, n)
        Sigma = Sigma.T @ Sigma
        gamma = 0.1
        t = cp.Variable(name='t', bounds=[0, None])
        L = np.linalg.cholesky(Sigma, upper=False)

        objective = cp.Minimize(- mu.T @ x + gamma * t)
        constraints = [cp.norm(L.T @ x, 2) <= t,
                       cp.sum(x) == 1,
                       x >= 0]
        problem = cp.Problem(objective, constraints)
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        assert np.allclose(problem.value, -1.93414338e+00)

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    def test_portfolio_socp_x_column_vector(self, solver):
        np.random.seed(858)
        n = 100
        x = cp.Variable((n, 1), name='x')
        mu = np.random.randn(n, 1)
        Sigma = np.random.randn(n, n)
        Sigma = Sigma.T @ Sigma
        gamma = 0.1
        t = cp.Variable(name='t', bounds=[0, None])
        L = np.linalg.cholesky(Sigma, upper=False)

        objective = cp.Minimize(-cp.sum(cp.multiply(mu, x)) + gamma * t)
        constraints = [cp.norm(L.T @ x, 2) <= t,
                       cp.sum(x) == 1,
                       x >= 0]
        problem = cp.Problem(objective, constraints)
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        assert np.allclose(problem.value, -1.93414338e+00)

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    def test_localization(self, solver):
        np.random.seed(42)
        m = 10
        dim = 2
        x_true = np.array([2.0, -1.5])
        a = np.random.uniform(-5, 5, (m, dim))
        rho = np.linalg.norm(a - x_true, axis=1)  # no noise
        x = cp.Variable(2, name='x')
        t = cp.Variable(m, name='t')
        constraints = [t == cp.sqrt(cp.sum(cp.square(x - a), axis=1))]
        objective = cp.Minimize(cp.sum_squares(t - rho))
        problem = cp.Problem(objective, constraints)
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        assert np.allclose(x.value, x_true)

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    def test_localization2(self, solver):
        np.random.seed(42)
        m = 10
        dim = 2
        x_true = np.array([2.0, -1.5])
        a = np.random.uniform(-5, 5, (m, dim))
        rho = np.linalg.norm(a - x_true, axis=1)  # no noise
        x = cp.Variable((1, 2), name='x')
        t = cp.Variable(m, name='t')
        constraints = [t == cp.sqrt(cp.sum(cp.square(x - a), axis=1))]
        objective = cp.Minimize(cp.sum_squares(t - rho))
        problem = cp.Problem(objective, constraints)
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        assert np.allclose(x.value, x_true)

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    def test_circle_packing_formulation_one(self, solver):
        """Epigraph formulation."""
        rng = np.random.default_rng(5)
        n = 3
        radius = rng.uniform(1.0, 3.0, n)

        centers = cp.Variable((2, n), name='c')
        constraints = []
        for i in range(n - 1):
            for j in range(i + 1, n):
                constraints += [cp.sum(cp.square(centers[:, i] - centers[:, j])) >=
                                (radius[i] + radius[j]) ** 2]

        centers.value = rng.uniform(-5.0, 5.0, (2, n))
        t = cp.Variable()
        obj = cp.Minimize(t)
        constraints += [cp.max(cp.norm_inf(centers, axis=0) + radius) <= t]
        problem = cp.Problem(obj, constraints)
        problem.solve(solver=solver, nlp=True)

        true_sol = np.array([[1.73655994, -1.98685738, 2.57208783],
                             [1.99273311, -1.67415425, -2.57208783]])
        assert np.allclose(centers.value, true_sol)

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    @pytest.mark.skipif('UNO' in INSTALLED_SOLVERS, reason='UNO finds a KKT point with worse obj')
    def test_circle_packing_formulation_two(self, solver):
        """Using norm_inf. This test revealed a very subtle bug in the unpacking of
        the ipopt solution. Some variables were mistakenly reordered. It was fixed
        in https://github.com/cvxgrp/cvxpy-ipopt/pull/82"""
        rng = np.random.default_rng(5)
        n = 3
        radius = rng.uniform(1.0, 3.0, n)

        centers = cp.Variable((2, n), name='c')
        constraints = []
        for i in range(n - 1):
            for j in range(i + 1, n):
                constraints += [cp.sum(cp.square(centers[:, i] - centers[:, j])) >=
                                (radius[i] + radius[j]) ** 2]

        centers.value = rng.uniform(-5.0, 5.0, (2, n))
        obj = cp.Minimize(cp.max(cp.norm_inf(centers, axis=0) + radius))
        prob = cp.Problem(obj, constraints)
        prob.solve(solver=solver, nlp=True)

        assert np.allclose(obj.value, 4.602738956101437)

        residuals = []
        for i in range(n - 1):
            for j in range(i + 1, n):
                dist_sq = np.linalg.norm(centers.value[:, i] - centers.value[:, j]) ** 2
                min_dist_sq = (radius[i] + radius[j]) ** 2
                residuals.append(dist_sq - min_dist_sq)

        assert(np.all(np.array(residuals) <= 1e-6))

        # Ipopt finds these centers, but Knitro rotates them (but finds the same
        # objective value)
        #true_sol = np.array([[1.73655994, -1.98685738, 2.57208783],
        #                     [1.99273311, -1.67415425, -2.57208783]])
        #assert np.allclose(centers.value, true_sol)

        checker = DerivativeChecker(prob)
        checker.run_and_assert()

    def test_circle_packing_formulation_three(self, solver):
        """Using max max abs."""
        rng = np.random.default_rng(5)
        n = 3
        radius = rng.uniform(1.0, 3.0, n)

        centers = cp.Variable((2, n), name='c')
        constraints = []
        for i in range(n - 1):
            for j in range(i + 1, n):
                constraints += [cp.sum(cp.square(centers[:, i] - centers[:, j])) >=
                                (radius[i] + radius[j]) ** 2]

        centers.value = rng.uniform(-5.0, 5.0, (2, n))
        obj = cp.Minimize(cp.max(cp.max(cp.abs(centers), axis=0) + radius))
        prob = cp.Problem(obj, constraints)
        prob.solve(solver=solver, nlp=True)

        # after the chain rule implementation, we find another configuration
        # with a minus sign
        true_sol = -np.array([[1.73655994, -1.98685738, 2.57208783],
                             [1.99273311, -1.67415425, -2.57208783]])
        assert np.allclose(centers.value, true_sol)

        checker = DerivativeChecker(prob)
        checker.run_and_assert()

    @pytest.mark.skipif('UNO' in INSTALLED_SOLVERS, reason='UNO reaches iteration limit')
    def test_geo_mean(self, solver):
        x = cp.Variable(3, nonneg=True)
        geo_mean = cp.geo_mean(x)
        objective = cp.Maximize(geo_mean)
        constraints = [cp.sum(x) == 1]
        problem = cp.Problem(objective, constraints)
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        assert np.allclose(x.value, np.array([1/3, 1/3, 1/3]))

        checker = DerivativeChecker(problem)
        checker.run_and_assert()

    @pytest.mark.skipif('UNO' in INSTALLED_SOLVERS, reason='UNO reaches iteration limit')
    def test_geo_mean2(self, solver):
        p = np.array([.07, .12, .23, .19, .39])
        x = cp.Variable(5, nonneg=True)
        prob = cp.Problem(cp.Maximize(cp.geo_mean(x, p)), [cp.sum(x) <= 1])
        prob.solve(solver=solver, nlp=True)
        x_true = p/sum(p)
        assert prob.status == cp.OPTIMAL
        assert np.allclose(x.value, x_true)
        checker = DerivativeChecker(prob)
        checker.run_and_assert()

    def test_div_composition(self, solver):
        x = cp.Variable(nonneg=True, bounds=[1, 5])
        prob = cp.Problem(cp.Maximize(cp.exp(1 / x)))
        prob.solve(solver=solver, nlp=True)
        x_true = 1.0
        assert prob.status == cp.OPTIMAL
        assert np.allclose(x.value, x_true)
        checker = DerivativeChecker(prob)
        checker.run_and_assert()

    def test_clnlbeam(self, solver):
        N = 1000
        h = 1 / N
        alpha = 350

        t = cp.Variable(N+1, bounds=[-1, 1])
        x = cp.Variable(N+1, bounds=[-0.05, 0.05])
        u = cp.Variable(N+1)
        u.value = np.zeros(N+1)
        control_terms = cp.multiply(0.5 * h, cp.power(u[1:], 2) + cp.power(u[:-1], 2))
        trigonometric_terms = cp.multiply(0.5 * alpha * h, cp.nlp.cos(t[1:]) + cp.nlp.cos(t[:-1]))
        objective_terms = cp.sum(control_terms + trigonometric_terms)

        objective = cp.Minimize(objective_terms)
        constraints = []
        position_constraints = (x[1:] - x[:-1] -
                                cp.multiply(0.5 * h, cp.nlp.sin(t[1:]) + cp.nlp.sin(t[:-1])) == 0)
        constraints.append(position_constraints)
        angle_constraint = (t[1:] - t[:-1] - 0.5 * h * (u[1:] + u[:-1]) == 0)
        constraints.append(angle_constraint)

        problem = cp.Problem(objective, constraints)
        problem.solve(solver=solver, nlp=True)
        assert problem.status == cp.OPTIMAL
        assert np.allclose(problem.value, 3.500e+02)

        # the derivative checker takes more than 10 seconds on this problem
        #checker = DerivativeChecker(problem)
        #checker.run_and_assert()
