# Copyright (c) 2025 Salvador E. Tropea
# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
# License: GPLv3
# Project: SeCoNoHe
#
# Implements: Fast Multi-Level Foreground Estimation
# by Thomas Germer, Tobias Uelwer, Stefan Conrad, Stefan Harmeling
# arXiv:2006.14970 26 Jun 2020
#
# Most code comes from the PyMatting project: https://github.com/pymatting/pymatting
# by: Thomas Germer, Tobias Uelwer, Joseph Adam, Christian Clauss
# License of PyMatting: MIT
#
# This is the wrapper to select one of the implementations.
# Mostly generated by Gemini 2.5 Pro
#
import torch
import logging
import time
import traceback
import importlib
from typing import List, Tuple

# --- Global cache for failed imports ---
# We use a set for efficient checking.
_FAILED_IMPORTS = set()

# --- Backend Configuration ---
# Defines the priority and necessary information for each implementation.
# The 'module' is the file name (without .py) (the function must be the same)
BACKENDS = [
    {'name': 'cupy', 'module': 'fmlfe_cupy'},
    {'name': 'opencl', 'module': 'fmlfe_pyopencl'},
    {'name': 'numba', 'module': 'fmlfe_numba'},
    {'name': 'torch', 'module': 'fmlfe_torch'},
]
IMPL_PRIORITY = [b['name'] for b in BACKENDS]


def fmlfe(
    images: torch.Tensor,
    masks: torch.Tensor,
    logger: logging.Logger,
    implementation: str = 'auto',
    regularization: float = 1e-5,
    n_small_iterations: int = 10,
    n_big_iterations: int = 2,
    small_size: int = 32,
    gradient_weight: float = 1.0
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    A high-level wrapper to run the Fast Multi-Level Foreground Estimation algorithm
    using the best available backend implementation.

    Accepts and returns ComfyUI-style tensors (BHWC/BHW).
    """
    # --- 1. Determine which implementation(s) to try ---
    if implementation == 'auto':
        implementations_to_try = IMPL_PRIORITY
    elif implementation in IMPL_PRIORITY:
        implementations_to_try = [implementation]
    else:
        raise ValueError(f"Invalid implementation '{implementation}'. Must be one of {IMPL_PRIORITY + ['auto']}")

    # --- 2. Iterate through implementations until one succeeds ---
    last_exception = None
    for impl_name in implementations_to_try:
        # If in auto mode, skip backends that have failed to import before.
        if implementation == 'auto' and impl_name in _FAILED_IMPORTS:
            continue

        backend_func = None
        try:
            # --- 3. Dynamically import the chosen backend ---
            backend_info = next(b for b in BACKENDS if b['name'] == impl_name)
            module_name = backend_info['module']
            module = importlib.import_module(f".{module_name}", package=__package__)
            backend_func = getattr(module, module_name)
            logger.debug(f"Successfully loaded '{impl_name}' implementation.")

            # --- 4. Prepare data and run the batch ---
            start_time = time.perf_counter()

            # All backends operate on single images, so we loop through the batch.
            batch_foregrounds: List[torch.Tensor] = []
            batch_backgrounds: List[torch.Tensor] = []

            for i in range(images.shape[0]):
                image_hwc = images[i]
                mask_hw = masks[i]

                # --- Data Conversion for each backend ---
                if impl_name == 'cupy':
                    import cupy
                    image_in = cupy.asarray(image_hwc)
                    alpha_in = cupy.asarray(mask_hw)
                elif impl_name in ['numba', 'opencl']:
                    image_in = image_hwc.cpu().numpy()
                    alpha_in = mask_hw.cpu().numpy()
                else:  # 'torch'
                    image_in = image_hwc
                    alpha_in = mask_hw

                # --- Execute the backend function ---
                F_out, B_out = backend_func(
                    image_in, alpha_in,
                    regularization=regularization,
                    n_small_iterations=n_small_iterations,
                    n_big_iterations=n_big_iterations,
                    small_size=small_size,
                    gradient_weight=gradient_weight,
                    return_background=True
                )

                # --- Convert output back to a Torch tensor on the correct device ---
                if impl_name == 'cupy':
                    batch_foregrounds.append(torch.as_tensor(F_out, device=images.device))
                    batch_backgrounds.append(torch.as_tensor(B_out, device=images.device))
                elif impl_name in ['numba', 'opencl']:
                    batch_foregrounds.append(torch.from_numpy(F_out).to(images.device))
                    batch_backgrounds.append(torch.from_numpy(B_out).to(images.device))
                else:  # 'torch'
                    batch_foregrounds.append(F_out)
                    batch_backgrounds.append(B_out)

            # --- 5. Finalize results and log success ---
            final_F = torch.stack(batch_foregrounds, dim=0)
            final_B = torch.stack(batch_backgrounds, dim=0)

            if images.device.type == 'cuda':
                torch.cuda.synchronize()
            end_time = time.perf_counter()

            logger.debug(f"'{impl_name}' implementation successful.")
            logger.debug(f"Processed images of shape {images.shape} in {end_time - start_time:.4f} seconds.")

            return final_F, final_B

        except ImportError as e:
            if implementation == 'auto':
                logger.debug(f"Could not import '{impl_name}' backend: {e}. Caching failure and trying next.")
                _FAILED_IMPORTS.add(impl_name)
                last_exception = e
            else:
                logger.error(f"Failed to import required backend '{impl_name}'. Please ensure it is installed.")
                raise e

        except Exception as e:
            if implementation == 'auto':
                logger.error(f"Execution failed for '{impl_name}' backend.")
                logger.info("Traceback for the error was:\n" + traceback.format_exc())
                last_exception = e
                # Continue to the next implementation in the priority list
            else:
                logger.error(f"Execution failed for the selected '{impl_name}' backend.")
                raise e

    # --- 9. If all implementations fail ---
    final_error_msg = f"All available foreground estimation implementations failed. Last error: {last_exception}"
    logger.error(final_error_msg)
    raise RuntimeError(final_error_msg)


# Example usage (would be in your ComfyUI node)
if __name__ == '__main__':
    # This block is for standalone testing.
    # It requires mock files for the backends to exist.
    print("This file is a library and should be imported by a ComfyUI node.")
    # Example of how a node would call it:
    # logger = logging.getLogger("MyNode")
    # images = torch.randn(2, 256, 256, 3, device='cuda') # BHWC
    # masks = torch.randn(2, 256, 256, device='cuda')      # BHW
    # F, B = fmlfe(images, masks, logger, implementation='auto')
    # print(f"Output F shape: {F.shape}, Output B shape: {B.shape}")
