# SPDX-FileCopyrightText: 2025 RTE (https://www.rte-france.com)
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
# SPDX-License-Identifier: MPL-2.0

import numpy as np
import pytest

from thermohl.utils import bisect_v, quasi_newton_2d

_nprs = 3141592654


def test_bisect():
    np.random.seed(_nprs)
    size = 99
    tol = 1.0e-09

    c = np.linspace(1, size, size) + np.random.randn(size)
    c = np.abs(c)

    def fun(x):
        return x**2 - c

    r0 = np.sqrt(c)

    x0, err = bisect_v(
        fun,
        0.0,
        np.sqrt(np.max(c)) * 1.1,
        (size,),
        print_error=False,
        tolerance=tol,
        max_iterations=99,
    )

    assert np.max(np.abs(x0 - r0) <= np.minimum(tol, err))


def test_bisect_scalar():
    def f(x):
        return x**2 - 2

    tol = 1e-6
    x0, err = bisect_v(f, lower_bound=0, upper_bound=2, output_shape=1, tolerance=tol)
    np.testing.assert_allclose(x0, np.sqrt(2), atol=tol)


def test_bisect_vector():
    def f(x):
        return np.array([x[0] ** 2 - 2, x[1] ** 3 - 2])
        # Not the best way to use bisect_v, but just for testing.

    tol = 1e-6
    x0, err = bisect_v(
        f, lower_bound=0, upper_bound=2, output_shape=(2,), tolerance=tol
    )
    np.testing.assert_allclose(x0, np.array([np.sqrt(2), np.cbrt(2)]), atol=tol)


def test_bisect_array():
    c = np.arange(1, 28).reshape(3, 3, 3)

    def f(x):
        return x**2 - c

    tol = 1e-6
    x0, err = bisect_v(
        f, lower_bound=0, upper_bound=30, output_shape=c.shape, tolerance=tol
    )

    np.testing.assert_allclose(x0, np.sqrt(c), atol=tol)


def test_bisect_no_convergence():
    def f(x):
        return x**2 + 1  # No root

    with pytest.raises(ValueError):
        bisect_v(
            f,
            lower_bound=-50,
            upper_bound=50,
            output_shape=(1,),
        )


def test_bisect_no_convergence_array():
    c = np.arange(-1, 1)

    def f(x):
        return x**2 + c  # No root for c > 0

    with pytest.raises(ValueError):
        bisect_v(
            f,
            lower_bound=-50,
            upper_bound=50,
            output_shape=(1,),
        )


#
def test_quasi_newton_2d_convergence():
    np.random.seed(_nprs)
    size = 10
    tol = 1.0e-12

    a = np.abs(1 + np.random.randn(size))
    b = np.abs(1 + np.random.randn(size))

    def f1(x, y):
        return y - a * x**2

    def f2(x, y):
        return y - b * x**3

    xg = np.ones((size,))
    yg = np.ones((size,))
    x, y, count, err = quasi_newton_2d(
        func1=f1,
        func2=f2,
        x_init=xg,
        y_init=yg,
        relative_tolerance=tol,
        max_iterations=999,
        delta_x=1.0e-09,
        delta_y=1.0e-09,
    )

    assert np.logical_and(
        np.max(np.abs(f1(x, y))) < tol, np.max(np.abs(f2(x, y))) < tol
    )
    assert count == 999


def test_quasi_newton_2d_no_convergence():
    np.random.seed(_nprs)
    size = 10
    tol = 1.0e-12

    a = np.abs(1 + np.random.randn(size))
    b = np.abs(1 + np.random.randn(size))

    def f1(x, y):
        return y - a * x**2

    def f2(x, y):
        return y - b * x**3

    xg = np.ones((size,))
    yg = np.ones((size,))
    x, y, count, err = quasi_newton_2d(
        func1=f1,
        func2=f2,
        x_init=xg,
        y_init=yg,
        relative_tolerance=tol,
        max_iterations=1,  # Set max iterations to 1 to force no convergence
        delta_x=1.0e-09,
        delta_y=1.0e-09,
    )

    assert count == 1
    assert np.logical_or(
        np.max(np.abs(f1(x, y))) >= tol, np.max(np.abs(f2(x, y))) >= tol
    )


def test_quasi_newton_2d_large_system():
    np.random.seed(_nprs)
    size = 1000
    tol = 1.0e-12

    a = np.abs(1 + np.random.randn(size))
    b = np.abs(1 + np.random.randn(size))

    def f1(x, y):
        return y - a * x**2

    def f2(x, y):
        return y - b * x**3

    xg = np.ones((size,))
    yg = np.ones((size,))
    x, y, count, err = quasi_newton_2d(
        func1=f1,
        func2=f2,
        x_init=xg,
        y_init=yg,
        relative_tolerance=tol,
        max_iterations=999,
        delta_x=1.0e-09,
        delta_y=1.0e-09,
    )

    assert np.logical_and(
        np.max(np.abs(f1(x, y))) < tol, np.max(np.abs(f2(x, y))) < tol
    )
