import numpy as np
from typing import Union, Callable, Any
from scipy.interpolate import LinearNDInterpolator


def interpolate_solution(solution: Any, points: np.ndarray) -> np.ndarray:
    """
    Interpolate a finite element solution at arbitrary points.
    
    Parameters
    ----------
    solution : Any
        Solution object with field, value, and mesh attributes
    points : np.ndarray
        Array of shape (n_points, 2) containing the coordinates of points to evaluate
        
    Returns
    -------
    np.ndarray
        Array of shape (n_points,) containing the interpolated values
    """
    # Extract mesh and solution values
    mesh = solution.mesh
    u_values = solution.value
    
    # Prepare mesh points for interpolation
    mesh_points = np.vstack((mesh.p[0, :], mesh.p[1, :])).T
    
    # Create interpolator
    interpolator = LinearNDInterpolator(mesh_points, u_values)
    
    # Interpolate at given points
    return interpolator(points)


def compute_mse(
    solution: Any, 
    points: np.ndarray, 
    ground_truth: Union[np.ndarray, Callable]
) -> float:
    """
    Compute the Mean Squared Error between a solution and ground truth at given points.
    
    Parameters
    ----------
    solution : Any
        Solution object with field, value, and mesh attributes
    points : np.ndarray
        Array of shape (n_points, 2) containing the coordinates of points to evaluate
    ground_truth : Union[np.ndarray, Callable]
        Either an array of shape (n_points,) containing ground truth values,
        or a callable that takes points as input and returns values
        
    Returns
    -------
    float
        Mean Squared Error
    """
    # Interpolate solution at points
    predicted = interpolate_solution(solution, points)
    
    # Get ground truth values
    if callable(ground_truth):
        true_values = ground_truth(points)
    else:
        true_values = ground_truth
    
    # Compute MSE
    mse = np.mean((predicted - true_values) ** 2)
    
    return mse 