import matplotlib.pyplot as plt
import numpy as np
from numpy import typing as npt
from typing import Callable

from . import ensure_iterable, SearchSpace
from .constraints import constrain_likelihood, snap_to_equality_constraints, snap_to_inequality_constraints
from .covariance import CovarianceMatrix


def ls_solver(y: Callable[[npt.NDArray[np.float64]], npt.NDArray[np.float64]],
              jacobian: Callable[[npt.NDArray[np.float64]], npt.NDArray[np.float64]],
              cov: CovarianceMatrix,
              x_init: npt.NDArray[np.float64],
              epsilon:float=1e-3,
              max_num_iterations:int=int(10e3),
              force_full_calc:bool=False,
              plot_progress:bool=False,
              eq_constraints:list=None,
              ineq_constraints:list=None,
              constraint_tolerance:float=1e-6)-> tuple[npt.NDArray, npt.NDArray]:
    """
    Computes the least square solution for geolocation processing.
    
    Ported from MATLAB Code.
    
    Nicholas O'Donoughue
    14 January 2021
    
    :param y: Measurement error vector function handle (accepts n_dim vector of source position estimate, responds with
                 error between received and modeled data vector)
    :param jacobian: Jacobian matrix function handle (accepts n_dim vector of source position estimate, and responds 
                     with n_dim x n_sensor Jacobian matrix)
    :param cov: Measurement error covariance matrix
    :param x_init: Initial estimate of source position
    :param epsilon: Desired position error tolerance (stopping condition)
    :param max_num_iterations: Maximum number of LS iterations to perform
    :param force_full_calc: Forces all max_num_iterations to be calculated
    :param plot_progress: Binary flag indicating whether to plot error/pos est over time
    :param eq_constraints: List of equality constraint functions (see constraints)
    :param ineq_constraints: List of inequality constraint functions (see constraints)
    :param constraint_tolerance: Tolerance to apply to equality constraints (default = 1e-6); any deviations with a
                                 Euclidean norm less than this tolerance are considered to satisfy the constraint.
    :return x: Estimated source position.
    :return x_full: Iteration-by-iteration estimated source positions
    """

    # Parse inputs
    n_dims = np.size(x_init)

    # Make certain that eq_constraints and ineq_constraints are iterable
    if ineq_constraints is not None:
        ensure_iterable(ineq_constraints, flatten=True)
    if eq_constraints is not None:
        ensure_iterable(eq_constraints, flatten=True)

    # Initialize loop
    current_iteration = 0
    error = np.inf
    x_full = np.zeros(shape=(n_dims, max_num_iterations))
    x_prev = x_init
    x_full[:, current_iteration] = x_prev

    # Initialize Plotting
    if plot_progress:
        _, ax = plt.subplots()
        plt.xlabel('Iteration Number')
        plt.ylabel('Change in Position Estimate')
        plt.yscale('log')

    # Divergence Detection
    num_expanding_iterations = 0
    max_num_expanding_iterations = 10
    prev_error = np.inf

    # Loop until either the desired tolerance is achieved or the maximum number of iterations have occurred
    while current_iteration < (max_num_iterations-1) and (force_full_calc or error >= epsilon):
        current_iteration += 1

        # Evaluate Residual and Jacobian Matrix
        y_i = y(x_prev)
        jacobian_i = np.squeeze(jacobian(x_prev))  # Use the squeeze command to drop the third dim (n_source = 1)

        # Compute delta_x^(i), according to 10.20
        delta_x = cov.solve_lstsq(y_i, jacobian_i)
        if np.ndim(delta_x) > 1:
            # There's a second dimension; take the average step across them
            delta_x = np.mean(delta_x, axis=1)

        # Update predicted location
        x_update = x_prev + delta_x

        # Apply Equality Constraints
        if eq_constraints is not None:
            x_update = snap_to_equality_constraints(x_update, eq_constraints=eq_constraints,
                                                    tol=constraint_tolerance)

        # Apply Inequality Constraints
        if ineq_constraints is not None:
            x_update = snap_to_inequality_constraints(x_update, ineq_constraints=ineq_constraints)

        # TODO: What to do if both ineq and eq constraints are in use?

        # Update variables
        x_full[:, current_iteration] = x_update
        x_prev = x_update
        error = np.linalg.norm(delta_x)

        if plot_progress:
            plt.plot(current_iteration, error, '.')

        # Check for divergence
        if error <= prev_error:
            num_expanding_iterations = 0
        else:
            num_expanding_iterations += 1
            if num_expanding_iterations >= max_num_expanding_iterations:
                # Divergence detected
                x_full[:, current_iteration:] = np.nan
                break
                
        prev_error = error

    x = x_prev

    # Bookkeeping
    if not force_full_calc:
        x_full[:, current_iteration+1:] = x[:, np.newaxis]

    return x, x_full


def gd_solver(y,
              jacobian,
              cov: CovarianceMatrix,
              x_init:npt.ArrayLike,
              alpha:float=0.3,
              beta:float=0.8,
              epsilon:float=1.e-3,
              max_num_iterations:int=int(10e3),
              force_full_calc:bool=False,
              plot_progress:bool=False,
              eq_constraints:list=None,
              ineq_constraints:list=None,
              constraint_tolerance:float=1e-6)-> tuple[npt.NDArray, npt.NDArray]:
    """
    Computes the gradient descent solution for localization given the provided measurement and Jacobian function 
    handles, and measurement error covariance.

    Ported from MATLAB code.
    
    Nicholas O'Donoughue
    29 April 2025
        
    :param y: Measurement vector function handle (accepts n_dim vector of source position estimate, responds with error 
              between received and modeled data vector)
    :param jacobian: Jacobian matrix function handle (accepts n_dim vector of source position estimate, and responds 
                     with n_dim x n_sensor Jacobian matrix)
    :param cov: Measurement error covariance matrix
    :param x_init: Initial estimate of source position
    :param alpha: Backtracking line search parameter
    :param beta: Backtracking line search parameter
    :param epsilon: Desired position error tolerance (stopping condition)
    :param max_num_iterations: Maximum number of LS iterations to perform
    :param force_full_calc: Forces all max_num_iterations to be executed
    :param plot_progress: Binary flag indicating whether to plot error/pos est over time
    :param eq_constraints: List of equality constraint functions (see constraints)
    :param ineq_constraints: List of inequality constraint functions (see constraints)
    :param constraint_tolerance: Tolerance to apply to equality constraints (default = 1e-6); any deviations with a
                                 Euclidean norm less than this tolerance are considered to satisfy the constraint.
    :return x: Estimated source position
    :return x_full: Iteration-by-iteration estimated source positions
    """
    # Parse inputs
    n_dims = np.size(x_init)

    # Make certain that eq_constraints and ineq_constraints are iterable
    if ineq_constraints is not None:
        ensure_iterable(ineq_constraints, flatten=True)
    if eq_constraints is not None:
        ensure_iterable(eq_constraints, flatten=True)

    # Initialize loop
    current_iteration = 0
    error = np.inf
    x_full = np.zeros(shape=(n_dims, max_num_iterations))
    x_prev = x_init
    x_full[:, current_iteration] = x_prev
    
    # Cost Function for Gradient Descent
    def cost_fxn(z):
        this_y = y(z)
        return np.mean(cov.solve_aca(this_y.T))

    # Initialize Plotting
    if plot_progress:
        fig_plot, (ax1, ax2) = plt.subplots(ncols=1, nrows=2)
        ax1.set_xlabel('Iteration Number')
        ax1.set_ylabel('Change in Position Estimate')
        ax1.set_yscale('log')
        
        ax2.set_xlabel('x')
        ax2.set_ylabel('y')
    else:
        ax1 = ax2 = None  # This removes a 'variable possibly unset' warning

    # Divergence Detection
    num_expanding_iterations = 0
    max_num_expanding_iterations = 5
    prev_error = np.inf

    # Loop until either the desired tolerance is achieved or the maximum
    # number of iterations have occurred
    while current_iteration < (max_num_iterations-1) and (force_full_calc or error >= epsilon):
        current_iteration += 1
    
        # Evaluate Residual and Jacobian Matrix
        y_i = y(x_prev)
        jacobian_i = np.squeeze(jacobian(x_prev))  # Use squeeze to remove third dimension (n_source=1)

        # If y_i is zero, gradient descent will fail because we're at the bottom. Skip to the next iteration without
        # changing anything
        if np.sum(np.abs(y_i)) < 1e-20:
            x_update = x_prev
            t = 0.
        else:
            # Compute Gradient and Cost function
            grad = -2 * cov.solve_acb(jacobian_i, y_i)
            if np.ndim(grad) > 1:
                # There's a second dimension; take the average gradient across them
                grad = np.mean(grad, axis=1)

            # Descent direction is the negative of the gradient
            del_x = -grad #/np.linalg.norm(grad))

            # Compute the step size
            t = backtracking_line_search(cost_fxn, x_prev, grad, del_x, alpha, beta)

            # Update x position
            x_update = x_prev + t*del_x

        # Apply Equality Constraints
        if eq_constraints is not None:
            x_update = snap_to_equality_constraints(x_update, eq_constraints=eq_constraints,
                                                    tol=constraint_tolerance)

        # Apply Inequality Constraints
        if ineq_constraints is not None:
            x_update = snap_to_inequality_constraints(x_update, ineq_constraints=ineq_constraints)

        # TODO: What to do if both ineq and eq constraints are in use?

        # Update variables
        error = np.linalg.norm(x_update - x_prev)  # how much did we adjust our position
        x_full[:, current_iteration] = x_update
        x_prev = x_update

        if plot_progress:
            ax1.plot(current_iteration, error, '.')
            ax2.plot(x_full[0, np.arange(current_iteration)], x_full[1, np.arange(current_iteration)], '-+')
        
        # Check for divergence
        if error <= prev_error:
            num_expanding_iterations = 0
        else:
            num_expanding_iterations += 1
            if num_expanding_iterations >= max_num_expanding_iterations:
                # Divergence detected
                x_full[:, current_iteration:] = np.nan
                break

        prev_error = error

    x = x_prev

    # Bookkeeping
    if not force_full_calc:
        x_full[:, current_iteration + 1:] = x[:, np.newaxis]

    return x, x_full
        

def backtracking_line_search(f, x, grad, del_x, alpha=0.5, beta=0.8, do_plot: bool=False):
    """
    # Performs backtracking line search according to algorithm 9.2 of
    # Stephen Boyd's, Convex Optimization

    Ported from MATLAB Code

    Nicholas O'Donoughue
    14 January 2021

    :param f: Function handle to minimize
    :param x: Current estimate of x
    :param grad: Gradient of f at x
    :param del_x: Descent direction
    :param alpha: Constant between 0 and 0.5 (default=0.3)
    :param beta: Constant between 0 and 1 (default=0.8)
    :return t: Optimal step size for the current estimate x.
    """

    # Initialize the search parameters and direction
    t = 1
    starting_val = np.squeeze(f(x))
    slope = np.squeeze(grad.T @ del_x)
    if slope > 0:
        raise ValueError(f"Error conducting backtracking line search; slope should be <0, but {slope} was encountered.")

    # Make sure that x, del_x are arrays (not matrices)
    del_x = np.squeeze(del_x)
    x = np.squeeze(x)
    # Make sure the starting value is large enough
    while f(x+t*del_x) <= starting_val+alpha*t*slope:
        t = 2*t
    
    # Conduct the backtracking line search
    while f(x+t*del_x) > starting_val+alpha*t*slope:
        t = beta*t

    if do_plot:
        start_val = np.max((2*t, 10))
        t_vec = np.linspace(start=-1*start_val, stop=start_val, num=100)
        plt.figure()
        plt.plot(t_vec, [f(x+t*del_x) for t in t_vec])
        plt.plot([0, t], [starting_val, f(x+t*del_x)],'-.')
        plt.grid(True)
        plt.show()

    # Ensure that the cost function did not increase
    if f(x+t*del_x) > starting_val:
        raise ValueError(f"Error conducting backtracking line search; unable to reduce the cost function.")

    return t


def ml_solver(ell, search_space: SearchSpace, eq_constraints=None, ineq_constraints=None, constraint_tolerance=None,
              prior=None, prior_wt: float = 0., num_levels: int=1, zoom_per_level: float=2,
              **kwargs)-> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], tuple[npt.NDArray[np.float64], ...]]:
    """
    Execute ML estimation through brute force computational methods.

    Any unexpected keyword arguments will be pass on as keyword arguments to the function ell, when it is called.

    Ported from MATLAB code.

    Nicholas O'Donoughue
    14 January 2021

    :param ell: Function handle for the likelihood of a given position must accept x_ctr (and similar sized vectors)
                as the sole input.
    :param search_space: SearchSpace object defining the space over which to search
    :param eq_constraints: List of equality constraint functions (see constraints)
    :param ineq_constraints: List of inequality constraint functions (see constraints)
    :param constraint_tolerance: Tolerance to apply to equality constraints (default = 1e-6); any deviations with a
                                 Euclidean norm less than this tolerance are considered to satisfy the constraint.
    :param prior: Function handle that accepts one or more positions and returns the probability of that position being
                  the true source location, according to some prior distribution. Will be multiplied by log10 when
                  combined with the likelihood distribution (which is assumed to be a log likelihood).
    :param prior_wt: Weight to apply to the prior distribution; (1-prior_wt) will be applied to the likelihood function.
    :param num_levels: Number of repeated searches to conduct; each with a finer mesh around the previously estimated
                       solution (int)
    :param zoom_per_level: Decrease in the grid search area per search level, as a multiplier > 1  (float)
    :return x_est: Estimated minimum
    :return A: Likelihood computed at each x position in the search space
    :return x_grid: Set of x positions for the entire search space (M x N) for N=1, 2, or 3.
    """

    if num_levels > 1:
        if zoom_per_level <= 1: raise ValueError(f"zoom_per_level must be greater than one; received {zoom_per_level}")

        # Call the solver with one less level to get the coarse estimate
        x_est, _, _ = ml_solver(ell, search_space, eq_constraints, ineq_constraints, constraint_tolerance,
                                prior, prior_wt, num_levels-1, zoom_per_level, **kwargs)
        # Zoom in and run as normal
        search_space = search_space.zoom_in(new_ctr=x_est, zoom=zoom_per_level**(num_levels-1), overwrite=False)

    # Set up the search space
    x_set, x_grid = search_space.x_set, search_space.x_grid

    # Constrain the likelihood, if needed
    if ineq_constraints is not None or eq_constraints is not None:
        ell = constrain_likelihood(ell=ell,
                                   eq_constraints=eq_constraints,
                                   ineq_constraints=ineq_constraints,
                                   tol=constraint_tolerance)

    # Evaluate the likelihood function at each coordinate in the search space
    likelihood = ell(x_set, **kwargs)

    if prior is not None and prior_wt > 0:
        pdf_prior = np.reshape(prior(x_set), shape=np.shape(likelihood))

        likelihood = (1 - prior_wt) * np.array(likelihood) + prior_wt * np.log10(pdf_prior)

    # Find the peak
    idx_pk = np.argmax(likelihood)
    x_est = x_set[:, idx_pk]

    return x_est, likelihood, x_grid


def bestfix_solver(pdfs, search_space: SearchSpace):
    """
    Based on the BESTFIX algorithm, invented in 1990 by Eric Hodson (R&D Associates, now with Naval Postgraduate
    School).  Patent is believed to be in the public domain.

    Ref:
     Eric Hodson, "Method and arrangement for probabilistic determination of a target location,"
     U.S. Patent US5045860A, 1990, https://patents.google.com/patent/US5045860A

    Ported from MATLAB Code.

    Nicholas O'Donoughue
    14 January 2021

    :param pdfs: Lx1 cell list of PDF functions, each of which represents the probability that an input position
                 (Ndim x 3 array) is the true source position for one of the measurements.
    :param search_space: SearchSpace object defining the space over which to search
    :return x_est: Estimated position
    :return result: Likelihood computed at each x position in the search space
    :return x_grid: Set of x positions for the entire search space (M x N) for N=1, 2, or 3.
    """

    # Set up the search space
    x_set, x_grid = search_space.x_set, search_space.x_grid

    # Apply each PDF to all input coordinates and then multiply across PDFs
    result = np.asarray([np.prod(np.asarray([this_pdf(this_x) for this_pdf in pdfs])) for this_x in x_set.T])

    # Reshape
    result = np.reshape(result, search_space.grid_shape)

    # Find the highest scoring position
    idx_pk = result.argmax()
    x_est = x_set[:, idx_pk]

    return x_est, result, x_grid
