from numbers import Real
from typing import *

import numpy as np

__seed__ = ["polygrad", "gradientdescent", "lossandgradient"]


def grade_polygrad(student_func: Callable[[Tuple[float, ...], float], float]):
    from hashlib import sha224

    from bwsi_grader import compare_functions
    from bwsi_grader.print import print_passed

    def soln(coeffs, x):
        assert all(isinstance(c, Real) for c in coeffs)
        assert isinstance(x, Real)
        coeffs = coeffs[1:]
        return float(
            sum(
                coef * ex * (x ** (ex - 1))
                for coef, ex in zip(coeffs, range(1, len(coeffs) + 1))
            )
        )

    # check some standard cases
    for coeffs, x in zip(((), (1.0,), (0, 2.0, 3.5)), (0.0, 12.4, 5.2)):
        compare_functions(student=student_func, soln=soln, fn_args=(coeffs, x))

    # check random cases

    for i in range(1000):
        num_coeffs = np.random.randint(0, 10)
        coeffs = tuple(float(i) for i in np.round(np.random.rand(num_coeffs) * 10, 2))
        x = np.round(np.random.rand(1) * 100, 2).item()
        compare_functions(student=student_func, soln=soln, fn_args=(coeffs, x))

    print_passed(sha224(str.encode(__seed__[0])).hexdigest())


def grade_gradient_descent(
    student_func: Callable[[List[float], float, int, float], List[float]]
):

    from hashlib import sha224

    from bwsi_grader import compare_functions
    from bwsi_grader.print import print_passed

    def soln(poly, step_size=0.1, iterations=10, x=100.0):
        """Returns a list of x-values visited when optimizing `poly` through gradient descent

        Parameters
        ----------
        poly : List[float]
            Polynomial coefficients in increasing order
        step_size : Optional[float], default: 0.1
            The magnitude of the step to take for each update of x
        iterations : Optional[int], default: 10
            After this number of iterations, the grad_descent function should return
        x : Optional[float], default: 100.
            The initial value of x

        Returns
        -------
        x_list : List[float]
            A list of the values of x that were visited, including
            the initial value of x.
        """
        assert isinstance(x, Real)
        assert isinstance(iterations, int)
        assert isinstance(step_size, Real)

        def poly_grad(coeffs, x):
            assert all(isinstance(c, Real) for c in coeffs)
            assert isinstance(x, Real)
            coeffs = coeffs[1:]
            return float(
                sum(
                    coef * ex * (x ** (ex - 1))
                    for coef, ex in zip(coeffs, range(1, len(coeffs) + 1))
                )
            )

        x_list = [x]
        for _ in range(iterations):
            x -= step_size * poly_grad(poly, x)
            x_list.append(x)
        return x_list

    # test static case
    x = 2.0
    coeffs = [1.0, 2.0, -4.0]
    iterations = 3
    step_size = 0.1

    compare_functions(
        student=student_func,
        soln=soln,
        fn_args=(coeffs,),
        fn_kwargs=dict(step_size=step_size, iterations=iterations, x=x),
    )

    for i in range(1000):
        num_coeffs = np.random.randint(0, 4)
        poly = [float(i) for i in np.random.randint(0, 10, size=(num_coeffs,))]
        x = float((np.random.rand() - 0.5))
        step_size = float(np.random.rand() * 0.01)
        iterations = int(np.random.randint(5, 15))
        _soln = soln(poly, step_size=step_size, iterations=iterations, x=x)

        if not np.isfinite(_soln).all():
            continue  # Only test happy paths
        else:
            compare_functions(
                student=student_func,
                soln=soln,
                fn_args=(poly,),
                fn_kwargs=dict(step_size=step_size, iterations=iterations, x=x),
            )

    print_passed(sha224(str.encode(__seed__[1])).hexdigest())


def grade_loss_and_gradient(
    student_func: Callable[
        [np.ndarray, np.ndarray, np.ndarray], Tuple[float, np.ndarray]
    ]
):
    from hashlib import sha224

    from bwsi_grader import compare_functions
    from bwsi_grader.print import print_passed

    def soln(X, A, y):
        """Given the input data (X), the collection of model parameters (A), and the
        truth data (y). Compute the loss, and the gradient of the loss (evaluated at A).

        Parameters
        ----------
        X : numpy.ndarray (shape = (N, 4))
            The data for all movies, with the first column of the array
            always equal to 1
        A : np.ndarray (shape = (4,))
            The model parameters: A_bias, A_prod, A_prom, A_book
        y : np.ndarray (shape = (N,))
            The true box office sales for each movie

        Returns
        -------
        out : tuple
            The first entry is the calculated loss and the second entry is the gradient
            evaluated at A.

            The loss should be a single floating-point number
            The gradient should be a numpy array of shape-(4,)

        Fun challenge (optional): For both the loss and the gradient try using matrix
        multiplication and broadcasting to avoid using any for-loops.
        """
        import numpy as np

        assert X.ndim == 2
        assert A.ndim == 1
        assert y.ndim == 1

        N = X.shape[0]
        assert X.shape == (N, 4)
        assert A.shape == (4,)
        assert y.shape == (N,)

        assert np.all(X[:, 0] == 1.0)

        y_hat = X.dot(A)
        loss = np.mean((y - y_hat) ** 2)
        grad = -2 * np.mean((y - y_hat) * X.T, axis=1)
        return loss, grad

    X = np.array([[1.0, 2.0, 0.1, 8.0], [1.0, 5.0, 2.0, 6.0]])

    A = np.array([0.5, 1.0, 2.0, 0.0])

    y = np.array((20.0, 30.0))

    compare_functions(student=student_func, soln=soln, fn_args=(X, A, y))

    for _ in range(50):
        p = np.random.randint(1, 10)
        X, A, y = (
            np.random.rand(p, 4).round(2),
            np.random.rand(4).round(2),
            np.random.rand(p).round(2),
        )
        X[:, 0] = 1.0

        compare_functions(student=student_func, soln=soln, fn_args=(X, A, y))
    print_passed(sha224(str.encode(__seed__[2])).hexdigest())
