"""
Implements optimisation-based methods for solving linear inverse problems.

This module provides classical, deterministic approaches to inversion that seek
a single "best-fit" model. These methods are typically formulated as finding
the model `u` that minimizes a cost functional.

The primary goal is to find a stable solution to an ill-posed problem by
incorporating regularization, which balances fitting the data with controlling
the complexity or norm of the solution.

Key Classes
-----------
- `LinearLeastSquaresInversion`: Solves the inverse problem by minimizing a
  Tikhonov-regularized least-squares functional.
- `LinearMinimumNormInversion`: Finds the model with the smallest norm that
  fits the data to a statistically acceptable degree using the discrepancy
  principle.
- `ConstrainedLinearLeastSquaresInversion`: Solves a linear inverse problem
  subject to an affine subspace constraint.
"""

from __future__ import annotations
from typing import Optional, Union


from .nonlinear_operators import NonLinearOperator
from .inversion import LinearInversion


from .forward_problem import LinearForwardProblem
from .linear_operators import LinearOperator
from .linear_solvers import LinearSolver, IterativeLinearSolver
from .hilbert_space import Vector
from .subspaces import AffineSubspace


class LinearLeastSquaresInversion(LinearInversion):
    """
    Solves a linear inverse problem using Tikhonov-regularized least-squares.

    This method finds the model `u` that minimizes the functional:
    `J(u) = ||A(u) - d||² + α² * ||u||²`
    where `α` is the damping parameter. If a data error covariance is provided,
    the data misfit norm is appropriately weighted by the inverse covariance.
    """

    def __init__(self, forward_problem: "LinearForwardProblem", /) -> None:
        """
        Args:
            forward_problem: The forward problem. If it includes a data error
                measure, the measure's inverse covariance must be defined.
        """
        super().__init__(forward_problem)
        if self.forward_problem.data_error_measure_set:
            self.assert_inverse_data_covariance()

    def normal_operator(self, damping: float) -> LinearOperator:
        """
        Returns the Tikhonov-regularized normal operator.

        This operator, often written as `(A* @ W @ A + α*I)`, forms the left-hand
        side of the normal equations that must be solved to find the least-squares
        solution. `W` is the inverse data covariance (or identity).

        Args:
            damping: The Tikhonov damping parameter, `α`. Must be non-negative.

        Returns:
            The normal operator as a `LinearOperator`.
        """
        if damping < 0:
            raise ValueError("Damping parameter must be non-negative.")

        forward_operator = self.forward_problem.forward_operator
        identity = self.forward_problem.model_space.identity_operator()

        if self.forward_problem.data_error_measure_set:
            inverse_data_covariance = (
                self.forward_problem.data_error_measure.inverse_covariance
            )
            return (
                forward_operator.adjoint @ inverse_data_covariance @ forward_operator
                + damping * identity
            )
        else:
            return forward_operator.adjoint @ forward_operator + damping * identity

    def normal_rhs(self, data: Vector) -> Vector:
        """
        Returns the right hand side of the normal equations for given data.
        """

        forward_operator = self.forward_problem.forward_operator

        if self.forward_problem.data_error_measure_set:
            inverse_data_covariance = (
                self.forward_problem.data_error_measure.inverse_covariance
            )

            shifted_data = self.forward_problem.data_space.subtract(
                data, self.forward_problem.data_error_measure.expectation
            )

            return (forward_operator.adjoint @ inverse_data_covariance)(shifted_data)

        else:
            return forward_operator.adjoint(data)

    def least_squares_operator(
        self,
        damping: float,
        solver: "LinearSolver",
        /,
        *,
        preconditioner: Optional[LinearOperator] = None,
    ) -> Union[NonLinearOperator, LinearOperator]:
        """
        Returns an operator that maps data to the least-squares solution.

        The returned operator `L` gives the solution `u = L(d)`. If the data has
        errors with a non-zero mean, `L` is a general non-linear `Operator`.
        Otherwise, it is a `LinearOperator`.

        Args:
            damping: The Tikhonov damping parameter, `alpha`.
            solver: The linear solver for inverting the normal operator.
            preconditioner: An optional preconditioner for iterative solvers.

        Returns:
            An operator that maps from the data space to the model space.
        """

        forward_operator = self.forward_problem.forward_operator
        normal_operator = self.normal_operator(damping)

        if isinstance(solver, IterativeLinearSolver):
            inverse_normal_operator = solver(
                normal_operator, preconditioner=preconditioner
            )
        else:
            inverse_normal_operator = solver(normal_operator)

        if self.forward_problem.data_error_measure_set:
            inverse_data_covariance = (
                self.forward_problem.data_error_measure.inverse_covariance
            )

            # This mapping is affine, not linear, if the error measure has a non-zero mean.
            def mapping(data: Vector) -> Vector:
                shifted_data = self.forward_problem.data_space.subtract(
                    data, self.forward_problem.data_error_measure.expectation
                )
                return (
                    inverse_normal_operator
                    @ forward_operator.adjoint
                    @ inverse_data_covariance
                )(shifted_data)

            return NonLinearOperator(self.data_space, self.model_space, mapping)

        else:
            return inverse_normal_operator @ forward_operator.adjoint


class ConstrainedLinearLeastSquaresInversion(LinearInversion):
    """
    Solves a linear inverse problem subject to an affine subspace constraint.

    Problem:
        Minimize J(u) = || A(u) - d ||_D^2 + alpha * || u ||_M^2
        Subject to u in A (Affine Subspace)

    Method:
        The problem is reduced to an unconstrained minimization in the subspace.
        We decompose the model as u = u_base + w, where u_base is the element
        of the affine subspace closest to the origin (orthogonal to the tangent space),
        and w is a perturbation in the tangent space.

        The cost function separates (due to orthogonality) into:
        J(w) = || A(w) - (d - A(u_base)) ||^2 + alpha * || w ||^2 + (alpha * ||u_base||^2)

        This is solved using the standard LinearLeastSquaresInversion on a
        reduced forward problem.
    """

    def __init__(
        self, forward_problem: LinearForwardProblem, constraint: AffineSubspace
    ) -> None:
        """
        Args:
            forward_problem: The original unconstrained forward problem.
            constraint: The affine subspace A where the solution must lie.
        """
        super().__init__(forward_problem)
        self._constraint = constraint

        # 1. Compute the Orthogonal Base Vector (u_base)
        # u_base = (I - P) * translation
        # This is the unique vector in the affine space that is orthogonal to the tangent space.
        # It ensures ||u||^2 = ||u_base||^2 + ||w||^2, decoupling the regularization.
        self._u_base = constraint.domain.subtract(
            constraint.translation, constraint.projector(constraint.translation)
        )

        # 2. Construct Reduced Forward Problem
        # Operator: A_tilde = A @ P
        reduced_operator = forward_problem.forward_operator @ constraint.projector

        # The error measure on the data remains valid for the reduced problem
        # because the noise model is additive and independent of the model parameters.
        self._reduced_forward_problem = LinearForwardProblem(
            reduced_operator,
            data_error_measure=(
                forward_problem.data_error_measure
                if forward_problem.data_error_measure_set
                else None
            ),
        )

        # 3. Initialize the internal unconstrained solver
        self._unconstrained_inversion = LinearLeastSquaresInversion(
            self._reduced_forward_problem
        )

    def least_squares_operator(
        self,
        damping: float,
        solver: LinearSolver,
        /,
        **kwargs,
    ) -> NonLinearOperator:
        """
        Returns an operator that maps data to the constrained least-squares solution.

        Args:
            damping: The Tikhonov damping parameter.
            solver: The linear solver for the reduced normal equations.
            **kwargs: Additional arguments passed to the solver (e.g., preconditioner).

        Returns:
            A NonLinearOperator mapping d -> u_constrained.
        """

        # Get the operator L_tilde such that w = L_tilde(d_tilde)
        reduced_op = self._unconstrained_inversion.least_squares_operator(
            damping, solver, **kwargs
        )

        # Precompute A(u_base) to shift the data efficiently
        # This represents the data predicted by the "base" model.
        data_offset = self.forward_problem.forward_operator(self._u_base)

        domain = self.data_space
        codomain = self.model_space

        def mapping(d: Vector) -> Vector:
            # 1. Shift Data: d_tilde = d - A(u_base)
            d_tilde = domain.subtract(d, data_offset)

            # 2. Solve for perturbation w in the tangent space
            # w = (P A* A P + alpha I)^-1 P A* d_tilde
            w = reduced_op(d_tilde)

            # 3. Reconstruct full model: u = u_base + w
            # Note: w is guaranteed to be in the tangent space (Range of P)
            # because of the structure of the reduced normal equations.
            return codomain.add(self._u_base, w)

        return NonLinearOperator(domain, codomain, mapping)


class LinearMinimumNormInversion(LinearInversion):
    """
    Finds a regularized solution using the discrepancy principle.

    This method automatically selects a Tikhonov damping parameter `α` such that
    the resulting solution `u_α` fits the data to a statistically acceptable
    level. It finds the model with the smallest norm `||u||` that satisfies
    the target misfit, as determined by a chi-squared test.
    """

    def __init__(self, forward_problem: "LinearForwardProblem", /) -> None:
        """
        Args:
            forward_problem: The forward problem. Its data error measure and
                inverse covariance must be defined.
        """
        super().__init__(forward_problem)
        if self.forward_problem.data_error_measure_set:
            self.assert_inverse_data_covariance()

    def minimum_norm_operator(
        self,
        solver: "LinearSolver",
        /,
        *,
        preconditioner: Optional[LinearOperator] = None,
        significance_level: float = 0.95,
        minimum_damping: float = 0.0,
        maxiter: int = 100,
        rtol: float = 1.0e-6,
        atol: float = 0.0,
    ) -> Union[NonLinearOperator, LinearOperator]:
        """
        Returns an operator that maps data to the minimum-norm solution.

        The method uses a bracketing search to finds the damping parameter `alpha`
        such that `chi_squared(u_alpha, d)` matches a critical value. The mapping
        is non-linear if data errors are present.

        Args:
            solver: A solver for the linear systems.
            preconditioner: An optional preconditioner for iterative solvers.
            significance_level: The target significance level for the
                chi-squared test (e.g., 0.95).
            minimum_damping: A floor for the damping parameter search.
            maxiter: Maximum iterations for the bracketing search.
            rtol: Relative tolerance for the damping parameter.
            atol: Absolute tolerance for the damping parameter.

        Returns:
            An operator that maps data to the minimum-norm model.
        """
        if self.forward_problem.data_error_measure_set:
            critical_value = self.forward_problem.critical_chi_squared(
                significance_level
            )
            lsq_inversion = LinearLeastSquaresInversion(self.forward_problem)

            def get_model_for_damping(
                damping: float, data: Vector, model0: Optional[Vector] = None
            ) -> tuple[Vector, float]:
                """
                Computes the LS model and its chi-squared for a given damping.

                When an iterative solver is used, an initial guess can be provided.
                """

                normal_operator = lsq_inversion.normal_operator(damping)
                normal_rhs = lsq_inversion.normal_rhs(data)

                if isinstance(solver, IterativeLinearSolver):
                    model = solver.solve_linear_system(
                        normal_operator, preconditioner, normal_rhs, model0
                    )
                else:
                    inverse_normal_operator = solver(normal_operator)
                    model = inverse_normal_operator(normal_rhs)

                chi_squared = self.forward_problem.chi_squared(model, data)
                return model, chi_squared

            def mapping(data: Vector) -> Vector:
                """The non-linear mapping from data to the minimum-norm model."""

                # Check to see if the zero model fits the data.
                chi_squared = self.forward_problem.chi_squared_from_residual(data)
                if chi_squared <= critical_value:
                    return self.model_space.zero

                # Find upper and lower bounds for the optimal damping parameter
                damping = 1.0
                _, chi_squared = get_model_for_damping(damping, data)

                damping_lower = damping if chi_squared <= critical_value else None
                damping_upper = damping if chi_squared > critical_value else None

                it = 0
                if damping_lower is None:
                    while chi_squared > critical_value and it < maxiter:
                        it += 1
                        damping /= 2.0
                        _, chi_squared = get_model_for_damping(damping, data)
                        if damping < minimum_damping:
                            raise RuntimeError(
                                "Discrepancy principle has failed; critical value cannot be reached."
                            )
                    damping_lower = damping

                it = 0
                if damping_upper is None:
                    while chi_squared < critical_value and it < maxiter:
                        it += 1
                        damping *= 2.0
                        _, chi_squared = get_model_for_damping(damping, data)
                    damping_upper = damping

                if damping_lower is None or damping_upper is None:
                    raise RuntimeError(
                        "Failed to bracket the optimal damping parameter."
                    )

                # Bracket search for the optimal damping
                model0 = None
                for _ in range(maxiter):
                    damping = 0.5 * (damping_lower + damping_upper)
                    model, chi_squared = get_model_for_damping(damping, data, model0)

                    if chi_squared < critical_value:
                        damping_lower = damping
                    else:
                        damping_upper = damping

                    if damping_upper - damping_lower < atol + rtol * (
                        damping_lower + damping_upper
                    ):
                        return model

                    model0 = model

                raise RuntimeError("Bracketing search failed to converge.")

            return NonLinearOperator(self.data_space, self.model_space, mapping)

        else:
            # For error-free data, compute the minimum-norm solution via A*(A*A)^-1
            forward_operator = self.forward_problem.forward_operator
            normal_operator = forward_operator @ forward_operator.adjoint
            inverse_normal_operator = solver(normal_operator)
            return forward_operator.adjoint @ inverse_normal_operator


class ConstrainedLinearMinimumNormInversion(LinearInversion):
    """
    Finds the minimum-norm solution subject to an affine subspace constraint
    using the discrepancy principle.

    Problem:
        Minimize ||u||
        Subject to u in A (Affine Subspace)
        And chi_squared(u, d) <= critical_value

    Method:
        We decompose the model as u = u_base + w, where u_base is the element
        of the affine subspace with the smallest norm (orthogonal to the tangent
        space), and w is a perturbation in the tangent space.

        Because u_base and w are orthogonal, ||u||^2 = ||u_base||^2 + ||w||^2.
        Minimizing ||u|| is therefore equivalent to minimizing ||w||.

        The problem reduces to finding the minimum norm w such that:
        || A(w) - (d - A(u_base)) ||_D^2 <= critical_value

        This is solved using the standard LinearMinimumNormInversion on a
        reduced forward problem.
    """

    def __init__(
        self,
        forward_problem: LinearForwardProblem,
        constraint: AffineSubspace,
    ) -> None:
        """
        Args:
            forward_problem: The original unconstrained forward problem.
            constraint: The affine subspace A where the solution must lie.
        """
        super().__init__(forward_problem)
        if self.forward_problem.data_error_measure_set:
            self.assert_inverse_data_covariance()

        self._constraint = constraint

        # 1. Compute the Orthogonal Base Vector (u_base)
        # u_base = (I - P) * translation
        # This is the vector in the affine space closest to the origin.
        self._u_base = constraint.domain.subtract(
            constraint.translation, constraint.projector(constraint.translation)
        )

        # 2. Construct Reduced Forward Problem
        # Operator: A_tilde = A @ P
        reduced_operator = forward_problem.forward_operator @ constraint.projector

        self._reduced_forward_problem = LinearForwardProblem(
            reduced_operator,
            data_error_measure=(
                forward_problem.data_error_measure
                if forward_problem.data_error_measure_set
                else None
            ),
        )

        # 3. Initialize the internal unconstrained solver
        self._unconstrained_inversion = LinearMinimumNormInversion(
            self._reduced_forward_problem
        )

    def minimum_norm_operator(
        self,
        solver: LinearSolver,
        /,
        **kwargs,
    ) -> NonLinearOperator:
        """
        Returns an operator that maps data to the constrained minimum-norm solution.

        Args:
            solver: The linear solver for the reduced normal equations.
            **kwargs: Arguments passed to LinearMinimumNormInversion (e.g.,
                      significance_level, rtol, maxiter).

        Returns:
            A NonLinearOperator mapping d -> u_constrained.
        """

        # Get the operator L_tilde such that w = L_tilde(d_tilde)
        reduced_op = self._unconstrained_inversion.minimum_norm_operator(
            solver, **kwargs
        )

        # Precompute A(u_base) to shift the data
        data_offset = self.forward_problem.forward_operator(self._u_base)

        domain = self.data_space
        codomain = self.model_space

        def mapping(d: Vector) -> Vector:
            # 1. Shift Data: d_tilde = d - A(u_base)
            d_tilde = domain.subtract(d, data_offset)

            # 2. Solve for perturbation w in the tangent space
            w = reduced_op(d_tilde)

            # 3. Reconstruct full model: u = u_base + w
            return codomain.add(self._u_base, w)

        return NonLinearOperator(domain, codomain, mapping)
