from typing import Any, Optional

from vicentin.utils import Dispatcher
from vicentin.utils.logging import error, warn

disp_tychonov = Dispatcher()
disp_tv = Dispatcher()

try:
    from .regularization_np import (
        tychonov as tychonov_np,
        total_variation as total_variation_np,
    )

    disp_tychonov.register("numpy", tychonov_np)
    disp_tv.register("numpy", total_variation_np)
except ModuleNotFoundError:
    warn("NumPy backend not found.")
except Exception:
    error("NumPy backend failed to load.")


try:
    from .regularization_torch import (
        tychonov as tychonov_torch,
        total_variation as total_variation_torch,
    )

    disp_tychonov.register("torch", tychonov_torch)
    disp_tv.register("torch", total_variation_torch)
except ModuleNotFoundError:
    warn("Torch backend not found.")
except Exception:
    error("Torch backend failed to load.")


def tychonov(img: Any, backend: Optional[str] = None):
    ctx = disp_tychonov.detect_backend(img, backend)
    img = disp_tychonov.cast_values(ctx, img)

    return disp_tychonov(ctx, img)


def total_variation(
    img: Any, epsilon: float = 1e-2, backend: Optional[str] = None
):
    ctx = disp_tv.detect_backend(img, backend)
    img = disp_tv.cast_values(ctx, img)

    return disp_tv(ctx, img, epsilon)
