"""
Photon Mosaic Main Workflow

This is the main Snakefile that orchestrates the entire photon mosaic processing pipeline.
It handles dataset discovery, target generation, and coordinates the preprocessing and
suite2p analysis workflows.

The workflow:
1. Discovers datasets and their TIFF files from the raw data directory
2. Generates preprocessing targets for each dataset/session combination
3. Generates suite2p analysis targets for processed data
4. Includes preprocessing.smk and suite2p.smk modules to execute the actual processing
"""

from pathlib import Path
from photon_mosaic.dataset_discovery import DatasetDiscoverer
from photon_mosaic import log_cuda_availability
import logging

# Configure logging based on config settings
log_level = logging.DEBUG if config.get("logging", {}).get("snakemake_verbose", False) else logging.INFO
logging.basicConfig(level=log_level)
logger = logging.getLogger("snakemake.workflow")

# CUDA availability check on workflow start
onstart:
    log_cuda_availability()

raw_data_base = Path(config["raw_data_base"]).resolve()
processed_data_base = Path(config["processed_data_base"]).resolve()
slurm_config = config.get("slurm", {})
output_pattern = config["preprocessing"]["output_pattern"]

logger.info(f"Raw data base: {raw_data_base}")
logger.info(f"Processed data base: {processed_data_base}")

# Log SLURM configuration
use_slurm = config.get("use_slurm", False)
if use_slurm:
    logger.info("SLURM execution enabled")
    logger.info(f"SLURM configuration: {slurm_config}")
else:
    logger.info("SLURM execution disabled - running locally")

# Discover datasets and their TIFF files using the new class-based approach
discoverer = DatasetDiscoverer(
    base_path=raw_data_base,
    pattern=config["dataset_discovery"].get("pattern", ".*"),
    exclude_datasets=config["dataset_discovery"].get("exclude_datasets"),
    exclude_sessions=config["dataset_discovery"].get("exclude_sessions"),
    tiff_patterns=config["dataset_discovery"].get("tiff_patterns"),
    neuroblueprint_format=config["dataset_discovery"].get("neuroblueprint_format", False),
)

discoverer.discover()

logger.info(f"Discovered datasets: {discoverer.transformed_datasets}")
logger.info(f"TIFF files: {discoverer.tiff_files}")

preproc_targets = [
    str(
        Path(processed_data_base)
        / dataset_name
        / discoverer.get_session_name(i, session_idx)
        / "funcimg"
        / f"{output_pattern}{tiff_name}"
    )
    for i, dataset_name in enumerate(discoverer.transformed_datasets)
    for session_idx, tiff_list in discoverer.tiff_files[discoverer.original_datasets[i]].items()
    for tiff_name in tiff_list
]

logger.info(f"Preprocessing targets: {preproc_targets}")

suite2p_targets = [
    str(
        Path(processed_data_base)
        / dataset_name
        / discoverer.get_session_name(i, session_idx)
        / "funcimg"
        / "suite2p"
        / "plane0"
        / fname
    )
    for i, dataset_name in enumerate(discoverer.transformed_datasets)
    for session_idx, tiff_list in discoverer.tiff_files[discoverer.original_datasets[i]].items()
    for fname in ["F.npy", "data.bin"]
    if tiff_list  # Only create targets for sessions that have files
]

logger.info(f"Suite2p targets: {suite2p_targets}")

include: "preprocessing.smk"
include: "suite2p.smk"

rule all:
    input:
        preproc_targets,
        suite2p_targets
