# src/autoclean/tasks/assr_default.py
"""Task implementation for assr EEG preprocessing."""

from typing import Any, Dict

from autoclean.core.task import Task
from autoclean.io.export import save_raw_to_set
from autoclean.step_functions.continuous import step_create_bids_path


class AssrDefault(Task):
    """Task implementation for assr EEG preprocessing."""

    def __init__(self, config: Dict[str, Any]):
        self.raw = None
        self.cleaned_raw = None
        self.epochs = None
        super().__init__(config)

    def run(self) -> None:
        """Run the complete assr processing pipeline."""
        # Import and save raw EEG data
        self.import_raw()

        # Continue with other preprocessing steps
        self.run_basic_steps()

        # Store a copy of the pre-cleaned raw data for comparison in reports
        self.original_raw = self.raw.copy()

        # Create BIDS-compliant paths and filenames
        self.raw, self.config = step_create_bids_path(self.raw, self.config)

        self.clean_bad_channels(cleaning_method="interpolate", reset_bads=True)

        self.rereference_data()

        self.annotate_noisy_epochs()

        self.annotate_uncorrelated_epochs()

        # #Segment rejection
        self.detect_dense_oscillatory_artifacts()

        # #ICA
        self.run_ica()

        self.run_ICLabel()

        save_raw_to_set(
            raw=self.raw,
            autoclean_dict=self.config,
            stage="post_clean_raw",
            flagged=self.flagged,
        )

        # --- EPOCHING BLOCK START ---
        self.create_eventid_epochs()  # Using fixed-length epochs

        # Prepare epochs for ICA
        self.detect_outlier_epochs()

        # Clean epochs using GFP
        self.gfp_clean_epochs()
        # --- EPOCHING BLOCK END ---

        # Generate visualization reports
        # self.generate_reports()

    def _validate_task_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
        """Validate assr specific configuration.

        # Args:
        #     config: Configuration dictionary that has passed common validation

        # Returns:
        #     Validated configuration dictionary

        # Raises:
        #     ValueError: If required fields are missing or invalid
        #"""
        # # Validate assr specific fields
        required_fields = {
            "task": str,
            "eeg_system": str,
            "tasks": dict,
        }

        for field, field_type in required_fields.items():
            if field not in config:
                raise ValueError(f"Missing required field: {field}")
            if not isinstance(config[field], field_type):
                raise ValueError(f"Field {field} must be of type {field_type}")

        # Stage files are now auto-generated by export functions when needed
        # No need to validate stage_files structure as it's handled dynamically

        return config
