run.py

from __future__ import annotations
import pdb

import logging
import pathlib
import time
from datetime import datetime
import warnings

warnings.filterwarnings(action="ignore")
from typing import Any, Dict, List
from tqdm import tqdm

import numpy as np
import pandas as pd
from scipy.stats import rankdata
from sklearn import metrics
import torch
from sklearn.ensemble import IsolationForest
import pickle

torch.set_printoptions(linewidth=250)
torch.cuda.reset_peak_memory_stats("cuda")
from torch.utils.data import TensorDataset, DataLoader

from data.dataloader import (
    DeltaXAIDataset,
    Mimic3o,
    Physionet19,
    SimulatedStateo,
    SimulatedSwitchFeatureo,
    SimulatedSpikeo,
    SimulatedDelayedSpikeo,
    SimulatedActivityo,
    Mimic,
    SimulatedData,
    SimulatedSpike,
    SimulatedSwitch,
    SimulatedState,
)
from explainer.base import BaseExplainer, ExplainerConfig, BaselineType
from explainer.fit import FIT
from explainer.winit import WinIT
from explainer.timex import TimeX
from explainer.timexpp import TimeXPP
from models.generator import GeneratorTrainingResults
from modeltrainer import ModelTrainerWithCv
from utils.masker import Masker
from utils.plot import (
    BoxPlotter,
    visualize_per_sample_attribution,
    visualize_temporal_evolution,
)
from utils.utils import aggregate_scores, append_df_to_csv, set_seed
from utils.config import get_args, FEATURE_MAP

from txai.models.bc_model import AblationParameters, transformer_default_args


class Params:
    def __init__(self, argdict: Dict[str, Any]):
        self.argdict = argdict

        self._all_explainer_dict: Dict[str, List[Dict[str, Any]]] | None = None
        self._generators_to_train: Dict[str, List[Dict[str, Any]]] | None = None

        self._outpath: pathlib.Path | None = None
        self._ckptpath: pathlib.Path | None = None
        self._model_args: Dict[str, Any] | None = None
        self._model_train_args: Dict[str, Any] | None = None

        self._datasets = self._resolve_datasets()
        self._resolve_model_args()
        self._resolve_explainers()
        self._init_logging()

    @property
    def datasets(self) -> DeltaXAIDataset:
        return self._datasets

    @property
    def model_args(self) -> Dict[str, Any]:
        return {} if self._model_args is None else self._model_args

    @property
    def model_train_args(self) -> Dict[str, Any]:
        return {} if self._model_train_args is None else self._model_train_args

    @property
    def all_explainer_dict(self) -> Dict[str, List[Dict[str, Any]]]:
        return {} if self._all_explainer_dict is None else self._all_explainer_dict

    @property
    def generators_to_train(self) -> Dict[str, List[Dict, str, Any]]:
        return {} if self._generators_to_train is None else self._generators_to_train

    @property
    def outpath(self) -> pathlib.Path | None:
        return None if self._outpath is None else self._outpath

    @property
    def ckptpath(self) -> pathlib.Path | None:
        return None if self._ckptpath is None else self._ckptpath

    def _resolve_datasets(self) -> DeltaXAIDataset:
        data = self.argdict["data"]
        testbs = self.argdict["testbs"]
        data_path = self.argdict["datapath"]
        data_seed = self.argdict["explainerseed"]
        cv_to_use = self.argdict["cv"]
        print(f"{self.argdict=}")
        kwargs = {
            "batch_size": testbs,
            "seed": data_seed,
            "cv_to_use": cv_to_use,
            "time_difference": self.argdict["time_difference"],
            "num_samples_per_instance": self.argdict["num_samples_per_instance"]
        }
        if data_path is not None:
            kwargs["data_path"] = data_path

        if data == "spike":
            kwargs["testbs"] = 300 if testbs == -1 else testbs
            delay = self.argdict["delay"]
            return SimulatedSpike(delay=delay, **kwargs)
        elif data == "mimic":
            kwargs["testbs"] = 1000 if testbs == -1 else testbs
            return Mimic(**kwargs)
        elif data == "switch":
            kwargs["testbs"] = 300 if testbs == -1 else testbs
            return SimulatedSwitch(**kwargs)
        elif data == "state":
            kwargs["testbs"] = 300 if testbs == -1 else testbs
            return SimulatedState(**kwargs)
        elif data == "mimic3o":
            kwargs["testbs"] = 300 if testbs == -1 else testbs
            return Mimic3o(self.argdict, **kwargs)
        elif data == "physionet19":
            kwargs["testbs"] = 300 if testbs == -1 else testbs
            return Physionet19(self.argdict, **kwargs)
        elif data == "stateo":
            kwargs["testbs"] = 300 if testbs == -1 else testbs
            return SimulatedStateo(self.argdict, **kwargs)
        elif data == "switch_featureo":
            kwargs["testbs"] = 300 if testbs == -1 else testbs
            return SimulatedSwitchFeatureo(self.argdict, **kwargs)
        elif data == "spikeo":
            kwargs["testbs"] = 300 if testbs == -1 else testbs
            return SimulatedSpikeo(self.argdict, **kwargs)
        elif data == "delayed_spikeo":
            kwargs["testbs"] = 300 if testbs == -1 else testbs
            return SimulatedDelayedSpikeo(self.argdict, **kwargs)
        elif data == "activityo":
            kwargs["testbs"] = 300 if testbs == -1 else testbs
            return SimulatedActivityo(self.argdict, **kwargs)

        raise ValueError(f"Unknown data {data}")

    def _resolve_explainers(self) -> None:
        explainers = self.argdict["explainer"]
        nsamples = self.argdict["samples"]
        all_explainer_dict = {}
        generator_dict = {}
        for explainer in explainers:
            if explainer == "dynamask":
                explainer_dict = self._resolve_dynamask_explainer_dict()
                all_explainer_dict[explainer] = [explainer_dict]
            elif explainer == "winit":
                windows = self.argdict["window"]
                winit_metrics = self.argdict["winitmetric"]
                winit_explainer_dict_list = []
                generator_dict_list = []
                for window in windows:
                    explainer_dict_window = {
                        "window_size": window,
                        "joint": self.argdict["joint"],
                        "conditional": self.argdict["conditional"],
                        "usedatadist": self.argdict["usedatadist"],
                        "random_state": self.argdict["explainerseed"],
                    }
                    if nsamples != -1:
                        explainer_dict_window["n_samples"] = nsamples
                    for winit_metric in winit_metrics:
                        explainer_dict = explainer_dict_window.copy()
                        explainer_dict["metric"] = winit_metric
                        winit_explainer_dict_list.append(explainer_dict)

                    generator_dict_list.append(explainer_dict_window)
                all_explainer_dict[explainer] = winit_explainer_dict_list
                generator_dict["winit"] = generator_dict_list
            else:
                explainer_dict = {}
                if explainer in ["FIT", "FO", "AFO"] and nsamples != -1:
                    explainer_dict["n_samples"] = nsamples
                if explainer == "FIT":
                    generator_dict["FIT"] = [explainer_dict]
                if explainer in ["TimeX", "TimeXPP"] :
                    explainer_dict["pretrained_model_args"] = self._model_args
                    generator_dict[explainer] = [explainer_dict]
                if explainer == "TIMING":
                    explainer_dict["num_segments"] = self.argdict["num_segments"]
                    explainer_dict["max_seg_len"] = self.argdict["max_seg_len"]
                    explainer_dict["min_seg_len"] = self.argdict["min_seg_len"]
                    
                all_explainer_dict[explainer] = [explainer_dict]
        self._all_explainer_dict = all_explainer_dict
        self._generators_to_train = generator_dict

    def _resolve_dynamask_explainer_dict(self) -> Dict[str, Any]:
        data = self.argdict["data"]
        area = self.argdict["area"]
        loss = self.argdict["loss"]
        timereg = self.argdict["timereg"]
        sizereg = self.argdict["sizereg"]
        deletion_mode = self.argdict["deletion"]
        blur_type = self.argdict["blurtype"]
        explainer_dict = {"num_epoch": self.argdict["epoch"]}
        if loss is not None:
            explainer_dict["loss"] = loss
        if area is not None:
            explainer_dict["area_list"] = area
        elif data == "mimic":
            explainer_dict["area_list"] = [0.05]
        if timereg is not None:
            explainer_dict["time_reg_factor"] = timereg
        elif data == "mimic":
            explainer_dict["time_reg_factor"] = 0
        if sizereg is not None:
            explainer_dict["size_reg_factor_dilation"] = sizereg
        elif data == "mimic":
            explainer_dict["size_reg_factor_dilation"] = 10000
        if deletion_mode is not None:
            explainer_dict["deletion_mode"] = deletion_mode
        elif data == "mimic":
            explainer_dict["deletion_mode"] = True
        if blur_type is not None:
            explainer_dict["blur_type"] = blur_type
        elif data == "mimic":
            explainer_dict["blur_type"] = "fadema"
        return explainer_dict

    def _resolve_model_args(self) -> None:
        model_type = self.argdict["modeltype"].upper()
        hidden_size = 200
        dropout = 0.5
        num_layers = 1
        num_ensemble = 1
        lr = 1e-3
        self._model_args = {
            "hidden_size": hidden_size,
            "dropout": dropout,
            "num_layers": num_layers,
            "model_type": model_type,
            "num_ensemble": num_ensemble,
        }

        if lr is None:
            if isinstance(self._datasets, Mimic):
                if model_type == "MTAND":
                    lr = 1e-4
                elif model_type in ["SEFT", "GRU"]:
                    lr = 1e-3
            else:
                lr = 1e-3

        if isinstance(self._datasets, Mimic):
            if model_type == "GRU":
                num_epochs = 100
            elif model_type == "CONV":
                num_epochs = 10
            elif model_type == "LSTM":
                num_epochs = 30
            elif model_type in ["MTAND", "SEFT"]:
                num_epochs = 300
        else:
            num_epochs = 30
        self._model_train_args = {"num_epochs": num_epochs, "lr": lr}

        base_out_path = pathlib.Path(self.argdict["outpath"])
        base_ckpt_path = pathlib.Path(self.argdict["ckptpath"])
        self._outpath = base_out_path
        self._ckptpath = base_ckpt_path

    def _init_logging(self) -> logging.Logger:
        format = "%(asctime)s %(levelname)8s %(name)25s: %(message)s"
        log_formatter = logging.Formatter(format)

        if self.argdict["logfile"] is None:
            time_str = datetime.now().strftime("%Y%m%d-%H%M")
            log_file_name = f"log_{time_str}.log"
        else:
            log_file_name = self.argdict["logfile"]

        log_path = pathlib.Path(self.argdict["outpath"])
        log_path.mkdir(parents=True, exist_ok=True)
        log_file_name = log_path / log_file_name
        logging.basicConfig(format=format, level="INFO")

        root_logger = logging.getLogger()
        file_handler = logging.FileHandler(str(log_file_name))
        file_handler.setFormatter(log_formatter)
        root_logger.addHandler(file_handler)
        return root_logger

    def _resolve_path(self, base_path: pathlib.Path, model_type: str, num_layers: int):
        if model_type == "GRU":
            return base_path / f"gru{num_layers}layer"
        elif model_type == "LSTM":
            return base_path / "lstm"
        elif model_type == "CONV":
            return base_path / "conv"
        else:
            return base_path / f"{model_type}"

    def get_maskers(self, explainer: BaseExplainer) -> List[Masker]:
        maskers = []
        absolutize = not isinstance(explainer, (WinIT, FIT, TimeX, TimeXPP))
        #absolutize = not isinstance(explainer, (WinIT, FIT))
        substitution = self.argdict["substitution"]

        top_points = self.argdict["top_points"]
        top_ratio = self.argdict["top_ratio"]
        use_top_points = self.argdict["use_top_points"]
        masker = Masker(top_points, top_ratio, use_top_points, absolutize, substitution)

        maskers.append(masker)
        return maskers



class ExplanationRunner:
    """
    Our main class for training the model, training the generator, running explanations and
    evaluating explanations for various datasets.
    """

    def __init__(
        self,
        args,
        dataset: DeltaXAIDataset,
        device,
        out_path: pathlib.Path,
        ckpt_path: pathlib.Path,
    ):
        self.args = args
        self.dataset = dataset
        self.device = device
        self.out_path = out_path
        self.ckpt_path = ckpt_path

        self.model_trainers: ModelTrainerWithCv | None = None
        self.explainers: Dict[int, BaseExplainer] | None = None
        self.importances: Dict[int, np.ndarray] | None = None
        self.elapsed_times: Dict[int, float] | None = None

        self.log = logging.getLogger(ExplanationRunner.__name__)

    def init_model(
        self,
        hidden_size: int,
        dropout: float,
        num_layers: int,
        model_type: str = "GRU",
        verbose_eval: int = 10,
        early_stopping: bool = True,
        num_ensemble: int = 10,
        multi_gpu: bool = True,
    ) -> None:
        self.model_trainers = ModelTrainerWithCv(
            self.args,
            self.dataset,
            self.ckpt_path,
            hidden_size,
            dropout,
            num_layers,
            model_type,
            self.device,
            verbose_eval,
            early_stopping,
            num_ensemble,
            multi_gpu=multi_gpu,
        )

    def train_model(
        self,
        num_epochs: int,
        lr: float = 0.001,
        weight_decay: float = 0.001,
        use_all_times: bool = True,
    ) -> None:
        if self.model_trainers is None:
            raise RuntimeError("Initialize the model first.")
        self.model_trainers.train_models(
            num_epochs, lr, weight_decay, use_all_times=use_all_times
        )
        self.model_trainers.load_model()
        self._get_test_results(use_all_times)

    def load_model(self, use_all_times: bool = True) -> None:
        if self.model_trainers is None:
            raise RuntimeError("Initialize the model first.")
        self.model_trainers.load_model()
        self._get_test_results(use_all_times)

    def _get_test_results(self, use_all_times: bool) -> None:
        test_results = self.model_trainers.get_test_results(use_all_times)
        test_accs = [round(v.accuracy, 6) for v in test_results.values()]
        test_aucs = [round(v.auc, 6) for v in test_results.values()]
        test_auprcs = [round(v.auprc, 6) for v in test_results.values()]
        test_f1s = [round(v.f1, 6) for v in test_results.values()]
        self.log.info(
            f"Average Accuracy = {np.mean(test_accs):.4f}\u00b1{np.std(test_accs):.4f}"
        )
        self.log.info(
            f"Average AUROC = {np.mean(test_aucs):.4f}\u00b1{np.std(test_aucs):.4f}"
        )
        self.log.info(
            f"Average AUPRC = {np.mean(test_auprcs):.4f}\u00b1{np.std(test_auprcs):.4f}"
        )
        self.log.info(
            f"Average F1 = {np.mean(test_f1s):.4f}\u00b1{np.std(test_f1s):.4f}"
        )

    def _get_optimal_cutoff(self, use_all_times: bool, split, metric) -> None:
        optimal_cutoff_dict = self.model_trainers.get_optimal_cutoff(
            use_all_times, split, metric
        )
        for cv, cutoff in optimal_cutoff_dict.items():
            self.log.info(f"Optimal cutoff for CV {cv} = {cutoff:.4f}")
        return optimal_cutoff_dict

    def run_inference(
        self,
        dataloader: DataLoader,
    ) -> Dict[int, np.ndarray]:
        return self.model_trainers.run_inference(dataloader)

    def clean_up(self, clean_importance=True, clean_explainer=True, clean_model=False):
        if clean_model and self.model_trainers is not None:
            del self.model_trainers
            self.model_trainers = None
        if clean_explainer and self.explainers is not None:
            del self.explainers
            self.explainers = None
        if clean_importance and self.importances is not None:
            del self.importances
            self.importances = None
        torch.cuda.empty_cache()

    def get_explainers(
        self,
        args,
        explainer_name: str,
        explainer_dict: Dict[str, Any],
    ) -> None:
        config = ExplainerConfig(
            device=self.device,
            baseline_type=BaselineType.ZERO,
            window_size=self.dataset.num_timesteps,
            time_difference=args["time_difference"],
            additional_args=explainer_dict,
        )
        self.explainers = {}
        try:
            module_name = f"explainer.{explainer_name.lower()}"
            module = __import__(module_name, fromlist=[explainer_name])
            ExplainerClass = getattr(module, explainer_name)
        except (ImportError, AttributeError) as e:
            raise ValueError(f"Explainer '{explainer_name}' not found: {e}")

        for cv in self.dataset.cv_to_use():
            if explainer_name == "AFO":
                explainer = ExplainerClass(config, self.dataset.train_loaders[cv])
            elif explainer_name == "FIT":
                # Check if generator exists first
                generator_path = self._get_generator_path(cv) / "joint_generator" / f"len_{self.dataset.feature_size}.pt"
                generator_exists = generator_path.exists()

                explainer = FIT(
                    config,
                    device=self.device,
                    feature_size=self.dataset.feature_size,
                    data_name=self.dataset.get_name(),
                    path=self._get_generator_path(cv),
                    num_samples=explainer_dict.get("n_samples", 10),
                    train_loader=None if generator_exists else self.dataset.train_loaders[cv],
                    valid_loader=None if generator_exists else self.dataset.valid_loaders[cv]
                )
                
                try:
                    explainer.load_generators()
                except FileNotFoundError:
                    self.log.info(f"Generator not found for FIT (CV={cv}). Training a new one...")
                    explainer.train_generators(self.dataset.train_loaders[cv], self.dataset.valid_loaders[cv], num_epochs=300)
                except Exception as e:
                    self.log.error(f"Error with FIT generator for CV={cv}: {str(e)}")
                    raise
            elif explainer_name == "WinIT":
                # Check if generator exists first - could be joint or feature generator
                feature_generator_path_list = [self._get_generator_path(cv) / "feature_generator" / f"feature_{i}_len_10_cond_False.pt" for i in range(self.dataset.feature_size)]
                generator_exists = all(feature_generator_path.exists() for feature_generator_path in feature_generator_path_list)

                explainer = ExplainerClass(
                    config,
                    num_features=self.dataset.feature_size,
                    data_name=self.dataset.get_name(),
                    path=self._get_generator_path(cv),
                    train_loader=None if generator_exists else self.dataset.train_loaders[cv],
                    valid_loader=None if generator_exists else self.dataset.valid_loaders[cv],
                    args=args,
                )
                
                # If we initially thought generator existed but it doesn't, we'll need to provide loaders
                try:
                    explainer.load_generators()
                except (FileNotFoundError, RuntimeError):
                    self.log.info(f"Generator not found for WinIT (CV={cv}). Training a new one...")
                    explainer.load_generators(train_loader=self.dataset.train_loaders[cv], valid_loader=self.dataset.valid_loaders[cv])
            elif explainer_name == "GT":
                config.top = args["top"]
                explainer = ExplainerClass(config)
            elif explainer_name in ['TimeX', 'TimeXPP']:
                default_transformer_args = transformer_default_args
                default_abl_params = AblationParameters(archtype=args['archtype'].lower())
                loss_weight_dict = {
                    "gsat": args['gsat_weight'],
                    "connect": args['connect_weight'],
                }
                print(f"{self.dataset=}")
                explainer = ExplainerClass(
                    config=config,
                    device=config.device,
                    pret_ckpt_path=args['pret_ckpt_path'],
                    ckpt_save_path=args['ckpt_save_path'],
                    d_inp=self.dataset.train_mu.shape[-1],
                    max_len=self.dataset.train_mu.shape[0],
                    n_classes=self.dataset.num_classes,
                    n_prototypes=args['n_prototypes'],
                    gsat_r=args['gsat_r'],
                    transformer_args=default_transformer_args,
                    ablation_parameters=default_abl_params,
                    loss_weight_dict=loss_weight_dict,
                    archtype=args['archtype'],
                    mu=self.dataset.train_mu,
                    std=self.dataset.train_std,
                )
            elif "SWING" in explainer_name:
                config.additional_args = {
                    "swing_overall_path": args['swing_overall_path'],
                    "swing_baseline": args['swing_baseline'],
                    "swing_path": args['swing_path'],
                    "swing_n_samples": args['swing_n_samples'],
                    "swing_td": args['swing_td'],
                }
                explainer = ExplainerClass(config)
            else:
                explainer = ExplainerClass(config)

            self.explainers[cv] = explainer

    def train_generators(
        self, num_epochs: int
    ) -> Dict[int, GeneratorTrainingResults] | None:
        """
        Train the generator if applicable. Test the generator and save the generator
        training results.

        Args:
            num_epochs:
                Train the generator for number of epochs.

        Returns:
            The generator training results. None if the explainer has no generator to train.
        """
        if self.explainers is None:
            raise RuntimeError(
                "explainer is not initialized. Call get_explainer to initialize."
            )
        results = {}
        generator_array_path = self._get_generator_array_path()
        generator_array_path.mkdir(parents=True, exist_ok=True)
        for cv in self.dataset.cv_to_use():
            self.log.info(f"Training generator for cv={cv}")
            gen_result = self.explainers[cv].train_generators(
                self.dataset.train_loaders,
                self.dataset.valid_loaders,
                num_epochs,
            )
            self.explainers[cv].test_generators(self.dataset.test_loader)
            if gen_result is not None:
                results[cv] = gen_result
                np.save(
                    generator_array_path / f"{gen_result.name}_train_loss_cv_{cv}.npy",
                    gen_result.train_loss_trends,
                )
                np.save(
                    generator_array_path / f"{gen_result.name}_valid_loss_cv_{cv}.npy",
                    gen_result.valid_loss_trends,
                )
                np.save(
                    generator_array_path / f"{gen_result.name}_best_epoch_cv_{cv}.npy",
                    gen_result.best_epochs,
                )
        if len(results) > 0:
            return results
        return None

    def load_generators(self):
        """
        Load generators for explainers that use them, or train if they don't exist.
        """
        if self.explainers is None:
            raise RuntimeError(
                "explainer is not initialized. Call get_explainer to initialize."
            )
        if not any(explainer in self.args["explainer"] for explainer in ["AFO", "FIT", "WinIT", "TimeX", "TimeXPP"]):
            return

        for cv in self.dataset.cv_to_use():
            # Pass train and validation loaders so generator can be trained if needed
            self.explainers[cv].load_generators()

    def _get_generator_path(self, cv: int) -> pathlib.Path:
        # return pathlib.Path(self.ckpt_path).parent / self.dataset.get_name() / str(cv)
        return pathlib.Path(self.ckpt_path).parent.parent

    def set_model_for_explainer(self, set_eval: bool = True):
        if self.explainers is None:
            raise RuntimeError(
                "explainer is not initialized. Call get_explainer to initialize."
            )

        for cv in self.dataset.cv_to_use():
            model = self.model_trainers.model_trainers[cv].model
            self.explainers[cv].set_model(model, set_eval=set_eval)

    def run_attributes(self) -> None:
        if self.explainers is None:
            raise RuntimeError(
                "explainer is not initialized. Call get_explainer to initialize."
            )

        self.importances, self.elapsed_times = self._run_attributes_recursive(
            self.dataset.test_loader
        )

    # # TODO: remove below
    # def run_attributes(self) -> None:
    #     """
    #     Run attribution method for the explainer on the test set.
    #     """
    #     if self.explainers is None:
    #         raise RuntimeError(
    #             "explainer is not initialized. Call get_explainer to initialize."
    #         )

    #     # import pdb; pdb.set_trace()

    #     # Create a new test loader with batch size 1
    #     original_test_dataset = self.dataset.test_loader.dataset
    #     single_batch_loader = torch.utils.data.DataLoader(
    #         original_test_dataset,
    #         batch_size=1,
    #         shuffle=False,  # Maintain original order
    #         num_workers=self.dataset.test_loader.num_workers,
    #         pin_memory=getattr(self.dataset.test_loader, 'pin_memory', False)
    #     )

    #     self.importances, self.elapsed_times = self._run_attributes_recursive(
    #         single_batch_loader
    #     )

    def _run_attributes_recursive(
        self, dataloader: DataLoader
    ) -> Dict[int, np.ndarray]:
        try:
            return self._run_attributes(dataloader)
        except RuntimeError as e:
            if "CUDA out of memory" in str(e):
                # reduce batch size
                new_batch_size = dataloader.batch_size // 2
                if 0 < new_batch_size < dataloader.batch_size:
                    self.log.warning(
                        f"CUDA out of memory! Reducing batch size from "
                        f"{dataloader.batch_size} to {new_batch_size}"
                    )
                    new_loader = DataLoader(dataloader.dataset, new_batch_size)
                    # self.test_loader.batch_size = new_batch_size
                    return self._run_attributes_recursive(new_loader)
            raise e

    def _run_attributes(self, dataloader: DataLoader) -> Dict[int, np.ndarray]:
        all_importance_scores = {}
        elapsed_times = {}  # New dictionary to store elapsed times

        for cv in self.dataset.cv_to_use():
            importance_scores = []
            total_time = 0
            num_batches = 0

            for i, batch in tqdm(enumerate(dataloader)):
                batch = [x.to(self.device) for x in batch]
                start_time = time.time()
                #cur_batch = [batch[0][:, -48:], batch[1][:, -48:], batch[2], batch[3], batch[4]]
                #prev_batch = [batch[0][:, -49:-1], batch[1][:, -49:-1], batch[2], batch[3], batch[4]]

                score = self.explainers[cv].attribute(batch).detach().cpu().numpy()
                #cur_score = self.explainers[cv].attribute(cur_batch).detach().cpu().numpy()
                #prev_score = self.explainers[cv].attribute(prev_batch).detach().cpu().numpy()

                #score = np.zeros(batch[0].shape)
                #score[:, -48:] += cur_score
                #score[:, -49:-1] -= prev_score

                end_time = time.time()
                total_time += end_time - start_time
                num_batches += 1

                importance_scores.append(score)
            importance_scores = np.concatenate(importance_scores, 0)
            all_importance_scores[cv] = importance_scores
            elapsed_times[cv] = total_time / num_batches

            self.log.info(
                f"CV {cv}: Average attribution time per batch: {elapsed_times[cv]:.4f} seconds"
            )

        return all_importance_scores, elapsed_times

    def evaluate_simulated_importance(self, aggregate_methods) -> pd.DataFrame:
        if not isinstance(self.dataset, SimulatedData):
            raise ValueError(
                "non simulated dataset does not have simulated importances."
            )

        if self.importances is None:
            raise ValueError(
                "No importances is loaded. Call load_importance or run_attribute first."
            )

        ground_truth_importance = self.dataset.load_ground_truth_importance()
        absolutize = not isinstance(next(iter(self.explainers.values())), (FIT, WinIT))
        df = self._evaluate_importance_with_gt(
            ground_truth_importance, absolutize, aggregate_methods
        )
        self._plot_boxes(num_to_plot=20, aggregate_methods=aggregate_methods)
        return df

    def evaluate_performance_drop(
        self,
        maskers: List[Masker],
    ) -> pd.DataFrame:
        # 1) Load test arrays once
        testset = list(self.dataset.test_loader.dataset)
        x_test = torch.stack([it[0] for it in testset]).cpu().numpy()
        mask_test = torch.stack([it[1] for it in testset]).cpu().numpy()
        pid_test = torch.stack([it[2] for it in testset]).cpu().numpy()
        interval_test = torch.stack([it[3] for it in testset]).cpu().numpy()
        y_test = torch.stack([it[4] for it in testset]).cpu().numpy()

        # 2) Original predictions
        orig_pred = self.run_inference(self.dataset.test_loader)

        dfs = {}
        self.masked_data = {}
        # bs = self.dataset.test_loader.batch_size
        bs = 2048

        metric_cols = [
            "Prediction Difference",
            "Area Under Prediction Difference",
            "Macro Prediction Difference",
            "Area Under Macro Prediction Difference",

            "Prediction Preservation",
            "Area Under Prediction Preservation",
            "Macro Prediction Preservation",
            "Area Under Macro Prediction Preservation",

            "Correlation Coefficient in All Importances",
            "Correlation Coefficient in Top Importances",
            "Correlation Coefficient in Bottom Importances",

            "Avg. Mask Count",
        ]

        B, T, F = x_test.shape

        mask_sum = mask_test.sum(axis=(1, 2))
        print(f"Mask statistics:")
        print(f"  Shape: {mask_sum.shape}")
        print(f"  Min: {mask_sum.min()}")
        print(f"  Max: {mask_sum.max()}")
        print(f"  Mean: {mask_sum.mean():.2f}")
        print(f"  Std: {mask_sum.std():.2f}")
        print(f"  Median: {np.median(mask_sum):.2f}")
        print(f"  25th percentile: {np.percentile(mask_sum, 25):.2f}")
        print(f"  75th percentile: {np.percentile(mask_sum, 75):.2f}")
        print(f"  Non-zero count: {np.count_nonzero(mask_sum)}/{mask_sum.size}")
        print(f"{(T * F)=}")

        for masker in maskers:
            self.log.info(f"Evaluating mask={masker.get_name()}")
            total = masker.top_points if masker.use_top_points else int(np.ceil(masker.top_ratio * T * F))

            # helper to run inference on masked arrays
            def infer(xx, mm):
                ds = TensorDataset(
                    torch.from_numpy(xx),
                    torch.from_numpy(mm),
                    torch.from_numpy(pid_test),
                    torch.from_numpy(interval_test),
                    torch.from_numpy(y_test),
                )
                loader = DataLoader(ds, batch_size=bs, shuffle=False)
                preds = self.run_inference(loader)
                del loader, ds
                return preds

            # Precompute for correlation
            cv0 = self.dataset.cv_to_use()[0]
            imp0 = self.importances[cv0]
            B, T, F = imp0.shape
            
            # For top importances (drop metrics)
            flat0_abs = np.abs(imp0).reshape(B, -1)
            sorted_desc = np.sort(flat0_abs, axis=1)[:, ::-1]
            
            # For bottom importances (preservation metrics)
            # Use raw importances without abs() to ensure we have variation
            flat0_raw = imp0.reshape(B, -1)
            # Only consider observed points (mask > 0)
            observed_mask = mask_test.reshape(B, -1) > 0
            flat0_observed = np.where(observed_mask, flat0_raw, np.nan)
            
            # Sort by absolute value but keep original sign
            abs_indices = np.argsort(np.abs(flat0_observed), axis=1)
            sorted_asc = np.take_along_axis(flat0_observed, abs_indices, axis=1)
            
            # Replace NaNs with zeros for the correlation calculation
            sorted_asc = np.nan_to_num(sorted_asc)

            # Precompute mask rankings for both top and bottom
            masker.precompute_rankings(
                self.importances,
                x_test,
                mask_test,
                pid_test,
                interval_test,
                select_top=True,
            )
            masker.precompute_rankings(
                self.importances,
                x_test,
                mask_test,
                pid_test,
                interval_test,
                select_top=False,
            )
            
            # Initialize fully masked data for preservation metrics
            masker.select_top = False
            original_top_points = masker.top_points
            
            # For macro-level preservation, we need to ensure we're starting with all features masked
            # Save original state
            original_select_top = masker.select_top
            
            # Set to mask all features (for preservation metrics)
            masker.select_top = False  # Select bottom features (which means mask all when top_points=0)
            masker.top_points = 0      # Start with all masked
            
            # Get fully masked predictions for instance-level
            xs_fully_masked, ms_fully_masked, _ = masker.mask(x_test, mask_test, self.importances)
            fully_masked_pred = infer(xs_fully_masked[0], ms_fully_masked[0])
            
            # Get fully masked predictions for macro-level
            xs_macro_fully_masked, ms_macro_fully_masked, _ = masker.mask(x_test, mask_test, self.importances, macro=True)
            macro_fully_masked_pred = infer(xs_macro_fully_masked[0], ms_macro_fully_masked[0])
            
            # Restore original settings
            masker.top_points = original_top_points
            masker.select_top = original_select_top

            # Log the initial state
            self.log.info(f"Initial fully masked prediction mean: {fully_masked_pred.mean()}")
            self.log.info(f"Initial macro fully masked prediction mean: {macro_fully_masked_pred.mean()}")
            self.log.info(f"Original prediction mean: {orig_pred.mean()}")

            # Initialize tracking variables
            prev_drop, prev_mdrop = orig_pred.copy(), orig_pred.copy()
            prev_pres, prev_mpres = fully_masked_pred.copy(), macro_fully_masked_pred.copy()
            
            # Make sure the macro masking is working correctly for preservation
            self.log.info(f"Fully masked vs original diff: {np.abs(fully_masked_pred - orig_pred).mean()}")
            self.log.info(f"Macro fully masked vs original diff: {np.abs(macro_fully_masked_pred - orig_pred).mean()}")

            drop_means, pres_means = [], []
            mdrop_means, mpres_means = [], []
            per_drop, per_pres = [], []
            per_mdrop, per_mpres = [], []
            removed_vals, preserved_vals = [], []

            last_info = {
                "drop": {},
                "preserve": {},
                "macro_drop": {},
                "macro_preserve": {},
            }

            # Use consistent top_points for all iterations by saving the value
            masker.top_points = total
            if not masker.use_top_points:
                masker.use_top_points = True

            for i in range(total):
                masker.top_points = i + 1

                # instance-level drop
                masker.select_top = True
                xs_d, ms_d, im_d = masker.mask(x_test, mask_test, self.importances)
                pd_d = infer(xs_d[0], ms_d[0])
                if self.dataset.num_classes < 2:
                    dd = np.abs(prev_drop.flatten() - pd_d.flatten())
                else:
                    dd = np.mean(np.abs(prev_drop - pd_d), axis=1) # L1 Loss
                    # dd = np.sum(np.abs(prev_drop - pd_d), axis=1) # Sum of Loss
                drop_means.append(dd.mean())
                per_drop.append(dd)
                removed_vals.append(sorted_desc[:, i])
                prev_drop = pd_d

                # macro-level drop
                xs_md, ms_md, im_md = masker.mask(x_test, mask_test, self.importances, macro=True)
                pd_md = infer(xs_md[0], ms_md[0])
                if self.dataset.num_classes < 2:
                    dm = np.abs(prev_mdrop.flatten() - pd_md.flatten())
                else:
                    dm = np.mean(np.abs(prev_mdrop - pd_md), axis=1)
                    # dm = np.sum(np.abs(prev_mdrop - pd_md), axis=1)
                mdrop_means.append(dm.mean())
                per_mdrop.append(dm)
                prev_mdrop = pd_md

                if i == total - 1:
                    last_info["drop"] = {
                        "new_xs": xs_d,
                        "new_masks": ms_d,
                        "importance_masks": im_d,
                        "new_preds": pd_d,
                    }
                    last_info["macro_drop"] = {
                        "new_xs": xs_md,
                        "new_masks": ms_md,
                        "importance_masks": im_md,
                        "new_preds": pd_md,
                    }

                # instance-level preserve
                masker.select_top = False
                xs_p, ms_p, im_p = masker.mask(x_test, mask_test, self.importances)
                pd_p = infer(xs_p[0], ms_p[0])
                if self.dataset.num_classes < 2:
                    dp = np.abs(prev_pres.flatten() - pd_p.flatten())
                else:
                    dp = np.mean(np.abs(prev_pres - pd_p), axis=1)
                    # dp = np.sum(np.abs(prev_pres - pd_p), axis=1)
                pres_means.append(dp.mean())
                per_pres.append(dp)
                # print(f"{sorted_asc[:, i]=}")
                preserved_vals.append(sorted_asc[:, i])
                prev_pres = pd_p

                # macro-level preserve
                xs_mp, ms_mp, im_mp = masker.mask(x_test, mask_test, self.importances, macro=True)
                pd_mp = infer(xs_mp[0], ms_mp[0])
                if self.dataset.num_classes < 2:
                    dmp = np.abs(prev_mpres.flatten() - pd_mp.flatten())
                else:
                    dmp = np.mean(np.abs(prev_mpres - pd_mp), axis=1)
                    # dmp = np.sum(np.abs(prev_mpres - pd_mp), axis=1)
                mpres_means.append(dmp.mean())
                per_mpres.append(dmp)
                prev_mpres = pd_mp

                if i == total - 1:
                    last_info["preserve"] = {
                        "new_xs": xs_p,
                        "new_masks": ms_p,
                        "importance_masks": im_p,
                        "new_preds": pd_p,
                    }
                    last_info["macro_preserve"] = {
                        "new_xs": xs_mp,
                        "new_masks": ms_mp,
                        "importance_masks": im_mp,
                        "new_preds": pd_mp,
                    }

            # flatten for correlation and cumulative
            drop_vec = np.concatenate(per_drop)
            pres_vec = np.concatenate(per_pres)
            all_vec = np.concatenate([drop_vec, pres_vec])

            remove_imp = np.concatenate(removed_vals)
            preserve_imp = np.concatenate(preserved_vals)
            all_imp = np.concatenate([remove_imp, preserve_imp])

            cc = np.corrcoef(all_imp, all_vec)[0, 1]
            cc_drop = np.corrcoef(remove_imp, drop_vec)[0, 1]
            cc_pres = np.corrcoef(preserve_imp, pres_vec)[0, 1]

            print(f"{drop_means=}\n{pres_means=}")

            cum_d = np.cumsum(drop_means)
            cum_p = np.cumsum(pres_means)
            cum_md = np.cumsum(mdrop_means)
            cum_mp = np.cumsum(mpres_means)
            area_d = np.trapz(cum_d) / total
            area_p = np.trapz(cum_p) / total
            area_md = np.trapz(cum_md) / total
            area_mp = np.trapz(cum_mp) / total

            avg_mask = masker.all_masked_count[cv0].sum() / len(x_test)

            # Build and populate DataFrame
            df = pd.DataFrame(index=self.dataset.cv_to_use(), columns=metric_cols, dtype=float)
            
            for cv in df.index:
                df.at[cv, "Prediction Difference"] = cum_d[-1]
                df.at[cv, "Area Under Prediction Difference"] = area_d
                df.at[cv, "Macro Prediction Difference"] = cum_md[-1]
                df.at[cv, "Area Under Macro Prediction Difference"] = area_md

                df.at[cv, "Prediction Preservation"] = cum_p[-1]
                df.at[cv, "Area Under Prediction Preservation"] = area_p
                df.at[cv, "Macro Prediction Preservation"] = cum_mp[-1]
                df.at[cv, "Area Under Macro Prediction Preservation"] = area_mp

                df.at[cv, "Correlation Coefficient in All Importances"] = cc
                df.at[cv, "Correlation Coefficient in Top Importances"] = cc_drop
                df.at[cv, "Correlation Coefficient in Bottom Importances"] = cc_pres

                df.at[cv, "Avg. Mask Count"] = avg_mask

            dfs[masker.get_name()] = df
            self.masked_data[masker.get_name()] = last_info

        result = pd.concat(dfs, axis=0)
        result.index.name = "mask method"
        return result
    
    def evaluate_performance_drop_ood(
        self,
        maskers: List[Masker],
    ) -> pd.DataFrame:
        # 1) Load test arrays once
        testset = list(self.dataset.test_loader.dataset)
        x_test = torch.stack([it[0] for it in testset]).cpu().numpy()
        mask_test = torch.stack([it[1] for it in testset]).cpu().numpy()
        pid_test = torch.stack([it[2] for it in testset]).cpu().numpy()
        interval_test = torch.stack([it[3] for it in testset]).cpu().numpy()
        y_test = torch.stack([it[4] for it in testset]).cpu().numpy()

        # 2) Original predictions
        orig_pred = self.run_inference(self.dataset.test_loader)

        dfs = {}
        self.masked_data = {}
        bs = self.dataset.test_loader.batch_size

        metric_cols = [
            "OOD score",
            "Prediction Difference",
            "Correlation Coefficient in All Importances",
            "Avg. Mask Count",
        ]

        B, T, F = x_test.shape
        
        from models.ae_lstm import TimeSeriesAutoencoder
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        seq_len    = self.dataset.num_timesteps
        in_dim     = self.dataset.feature_size
        patch_size = 50
        embed_dim  = 64
        lr         = 1e-3
        n_epochs   = 10
        
        ae = TimeSeriesAutoencoder(
            seq_len=seq_len,
            in_dim=in_dim,
            patch_size=patch_size,
            embed_dim=embed_dim,
            encoder_depth=3,
            decoder_depth=2,
            num_heads=1,
        ).to(device)
        
        ae.load_state_dict(torch.load(self.args['ckpt_save_path'], map_location=device))
        ae.eval()

        for masker in maskers:
            self.log.info(f"Evaluating mask={masker.get_name()}")
            total = masker.top_points if masker.use_top_points else int(np.ceil(masker.top_ratio * T * F))

            # helper to run inference on masked arrays
            def infer(xx, mm):
                ds = TensorDataset(
                    torch.from_numpy(xx),
                    torch.from_numpy(mm),
                    torch.from_numpy(pid_test),
                    torch.from_numpy(interval_test),
                    torch.from_numpy(y_test),
                )
                loader = DataLoader(ds, batch_size=bs, shuffle=False)
                preds = self.run_inference(loader)
                del loader, ds
                return preds

            # Precompute for correlation
            cv0 = self.dataset.cv_to_use()[0]
            imp0 = self.importances[cv0]
            B, T, F = imp0.shape
            
            # For top importances (drop metrics)
            flat0_abs = np.abs(imp0).reshape(B, -1)
            sorted_desc = np.sort(flat0_abs, axis=1)[:, ::-1]
            
            # For bottom importances (preservation metrics)
            # Use raw importances without abs() to ensure we have variation
            flat0_raw = imp0.reshape(B, -1)
            # Only consider observed points (mask > 0)
            observed_mask = mask_test.reshape(B, -1) > 0
            flat0_observed = np.where(observed_mask, flat0_raw, np.nan)
            
            # Sort by absolute value but keep original sign
            abs_indices = np.argsort(np.abs(flat0_observed), axis=1)
            sorted_asc = np.take_along_axis(flat0_observed, abs_indices, axis=1)
            
            # Replace NaNs with zeros for the correlation calculation
            sorted_asc = np.nan_to_num(sorted_asc)

            # Precompute mask rankings for both top and bottom
            masker.precompute_rankings(
                self.importances,
                x_test,
                mask_test,
                pid_test,
                interval_test,
                select_top=True,
            )
            masker.precompute_rankings(
                self.importances,
                x_test,
                mask_test,
                pid_test,
                interval_test,
                select_top=False,
            )
            
            # Initialize fully masked data for preservation metrics
            masker.select_top = False
            original_top_points = masker.top_points
            
            # For macro-level preservation, we need to ensure we're starting with all features masked
            # Save original state
            original_select_top = masker.select_top
            
            # Set to mask all features (for preservation metrics)
            masker.select_top = False  # Select bottom features (which means mask all when top_points=0)
            masker.top_points = 0      # Start with all masked
            
            # Get fully masked predictions for instance-level
            xs_fully_masked, ms_fully_masked, _ = masker.mask(x_test, mask_test, self.importances)
            fully_masked_pred = infer(xs_fully_masked[0], ms_fully_masked[0])
            
            # Get fully masked predictions for macro-level
            xs_macro_fully_masked, ms_macro_fully_masked, _ = masker.mask(x_test, mask_test, self.importances, macro=True)
            macro_fully_masked_pred = infer(xs_macro_fully_masked[0], ms_macro_fully_masked[0])
            
            # Restore original settings
            masker.top_points = original_top_points
            masker.select_top = original_select_top

            # Log the initial state
            self.log.info(f"Initial fully masked prediction mean: {fully_masked_pred.mean()}")
            self.log.info(f"Initial macro fully masked prediction mean: {macro_fully_masked_pred.mean()}")
            self.log.info(f"Original prediction mean: {orig_pred.mean()}")

            # Initialize tracking variables
            prev_drop, prev_mdrop = orig_pred.copy(), orig_pred.copy()
            prev_pres, prev_mpres = fully_masked_pred.copy(), macro_fully_masked_pred.copy()
            
            # Make sure the macro masking is working correctly for preservation
            self.log.info(f"Fully masked vs original diff: {np.abs(fully_masked_pred - orig_pred).mean()}")
            self.log.info(f"Macro fully masked vs original diff: {np.abs(macro_fully_masked_pred - orig_pred).mean()}")

            drop_means, pres_means = [], []
            mdrop_means, mpres_means = [], []
            per_drop, per_pres = [], []
            per_mdrop, per_mpres = [], []
            removed_vals, preserved_vals = [], []

            last_info = {
                "drop": {},
                "preserve": {},
                "macro_drop": {},
                "macro_preserve": {},
            }

            # Use consistent top_points for all iterations by saving the value
            masker.top_points = total
            if not masker.use_top_points:
                masker.use_top_points = True

            for i in range(total):
                masker.top_points = i + 1

                # instance-level drop
                masker.select_top = True
                xs_d, ms_d, im_d = masker.mask(x_test, mask_test, self.importances)
                x_all = xs_d[0][:, -self.dataset.num_timesteps:]
                mask_all = im_d[0][:, -self.dataset.num_timesteps:]
                x_all = torch.tensor(x_all, device=device)
                mask_all = torch.tensor(mask_all, device=device)
                batch_size = self.dataset.test_loader.batch_size
                
                total_loss = 0.0
                total_samples = 0

                with torch.no_grad():
                    for start in range(0, x_all.size(0), batch_size):
                        end = start + batch_size
                        batch = x_all[start:end]
                        # print("mask")
                        # print(mask_all[start:end][1])
                        # print("orig")
                        # print(batch[1])
                        loss = ae(batch, mask_all[start:end])
                        bs = batch.size(0)
                        total_loss += loss.item() * bs
                        total_samples += bs
                        
                overall_mean_loss = total_loss / total_samples
                self.log.info(f"Overall mean loss ({i + 1}): {overall_mean_loss:.6f}")
                
                pd_d = infer(xs_d[0], ms_d[0])
                
                if self.dataset.num_classes < 2:
                    dd = np.abs(prev_drop.flatten() - pd_d.flatten())
                else:
                    dd = np.mean(np.abs(prev_drop - pd_d), axis=1) # L1 Loss
                    # dd = np.sum(np.abs(prev_drop - pd_d), axis=1) # Sum of Loss
                drop_means.append(dd.mean())
                per_drop.append(dd)
                removed_vals.append(sorted_desc[:, i])
                prev_drop = pd_d

                # macro-level drop
                xs_md, ms_md, im_md = masker.mask(x_test, mask_test, self.importances, macro=True)
                pd_md = infer(xs_md[0], ms_md[0])
                if self.dataset.num_classes < 2:
                    dm = np.abs(prev_mdrop.flatten() - pd_md.flatten())
                else:
                    dm = np.mean(np.abs(prev_mdrop - pd_md), axis=1)
                    # dm = np.sum(np.abs(prev_mdrop - pd_md), axis=1)
                mdrop_means.append(dm.mean())
                per_mdrop.append(dm)
                prev_mdrop = pd_md

                if i == total - 1:
                    last_info["drop"] = {
                        "new_xs": xs_d,
                        "new_masks": ms_d,
                        "importance_masks": im_d,
                        "new_preds": pd_d,
                    }
                    last_info["macro_drop"] = {
                        "new_xs": xs_md,
                        "new_masks": ms_md,
                        "importance_masks": im_md,
                        "new_preds": pd_md,
                    }

                # instance-level preserve
                masker.select_top = False
                xs_p, ms_p, im_p = masker.mask(x_test, mask_test, self.importances)
                pd_p = infer(xs_p[0], ms_p[0])
                if self.dataset.num_classes < 2:
                    dp = np.abs(prev_pres.flatten() - pd_p.flatten())
                else:
                    dp = np.mean(np.abs(prev_pres - pd_p), axis=1)
                    # dp = np.sum(np.abs(prev_pres - pd_p), axis=1)
                pres_means.append(dp.mean())
                per_pres.append(dp)
                # print(f"{sorted_asc[:, i]=}")
                preserved_vals.append(sorted_asc[:, i])
                prev_pres = pd_p

                # macro-level preserve
                xs_mp, ms_mp, im_mp = masker.mask(x_test, mask_test, self.importances, macro=True)
                pd_mp = infer(xs_mp[0], ms_mp[0])
                if self.dataset.num_classes < 2:
                    dmp = np.abs(prev_mpres.flatten() - pd_mp.flatten())
                else:
                    dmp = np.mean(np.abs(prev_mpres - pd_mp), axis=1)
                    # dmp = np.sum(np.abs(prev_mpres - pd_mp), axis=1)
                mpres_means.append(dmp.mean())
                per_mpres.append(dmp)
                prev_mpres = pd_mp

                if i == total - 1:
                    last_info["preserve"] = {
                        "new_xs": xs_p,
                        "new_masks": ms_p,
                        "importance_masks": im_p,
                        "new_preds": pd_p,
                    }
                    last_info["macro_preserve"] = {
                        "new_xs": xs_mp,
                        "new_masks": ms_mp,
                        "importance_masks": im_mp,
                        "new_preds": pd_mp,
                    }

            # flatten for correlation and cumulative
            drop_vec = np.concatenate(per_drop)
            pres_vec = np.concatenate(per_pres)
            all_vec = np.concatenate([drop_vec, pres_vec])

            remove_imp = np.concatenate(removed_vals)
            preserve_imp = np.concatenate(preserved_vals)
            all_imp = np.concatenate([remove_imp, preserve_imp])

            cc = np.corrcoef(all_imp, all_vec)[0, 1]
            cc_drop = np.corrcoef(remove_imp, drop_vec)[0, 1]
            cc_pres = np.corrcoef(preserve_imp, pres_vec)[0, 1]

            cum_d = np.cumsum(drop_means)
            cum_p = np.cumsum(pres_means)
            cum_md = np.cumsum(mdrop_means)
            cum_mp = np.cumsum(mpres_means)
            area_d = np.trapz(cum_d) / total
            area_p = np.trapz(cum_p) / total
            area_md = np.trapz(cum_md) / total
            area_mp = np.trapz(cum_mp) / total

            avg_mask = masker.all_masked_count[cv0].sum() / len(x_test)
            
            # Build and populate DataFrame
            df = pd.DataFrame(index=self.dataset.cv_to_use(), columns=metric_cols, dtype=float)
            
            for cv in df.index:
                df.at[cv, "OOD score"] = overall_mean_loss
                df.at[cv, "Prediction Difference"] = cum_d[-1]
                df.at[cv, "Correlation Coefficient in All Importances"] = cc
                df.at[cv, "Avg. Mask Count"] = avg_mask

            dfs[masker.get_name()] = df
            self.masked_data[masker.get_name()] = last_info

        result = pd.concat(dfs, axis=0)
        result.index.name = "mask method"
        return result

    def _plot_boxes(
        self,
        num_to_plot,
        aggregate_methods: List[str],
        x_other: List[Dict[int, torch.Tensor]] | None = None,
        mask_other: List[Dict[int, torch.Tensor]] | None = None,
        importance_mask_other: List[Dict[int, torch.Tensor]] | None = None,
        mask_name: str = "",
    ) -> None:
        explainer_name = self.get_explainer_name()
        plotter = BoxPlotter(self.dataset, self.out_path, num_to_plot, explainer_name, self.args["time_difference"])

        if self.importances is not None and x_other is not None:
            for aggregate_method in aggregate_methods:
                testset = list(self.dataset.test_loader.dataset)
                x_test = torch.stack([x[0] for x in testset]).cpu().numpy()
                mask_test = torch.stack([x[1] for x in testset]).cpu().numpy()
                y_test = torch.stack([x[-1] for x in testset]).cpu().numpy()

                # Get original predictions
                orig_preds = self.run_inference(self.dataset.test_loader)

                # Create masked testset for inference
                masked_testset = [
                    (x_other[0][i], mask_other[0][i]) + item[2:]
                    for i, item in enumerate(testset)
                ]
                batch_size = self.dataset.test_loader.batch_size
                masked_test_loader = DataLoader(
                    masked_testset, batch_size=batch_size, shuffle=False
                )

                new_preds = self.run_inference(masked_test_loader)

                plotter.plot_combined_visualization(
                    self.importances,
                    aggregate_method,
                    x_test,
                    mask_test,
                    {k: v.cpu().numpy() for k, v in x_other[0].items()},
                    {k: v.cpu().numpy() for k, v in mask_other[0].items()},
                    {k: v.cpu().numpy() for k, v in importance_mask_other[0].items()},
                    orig_preds,
                    new_preds,
                    y_test,
                    mask_name,
                )

    def _evaluate_importance_with_gt(
        self,
        ground_truth_importance: np.ndarray,
        absolutize: bool,
        aggregate_methods: List[str],
    ):
        ground_truth_importance = ground_truth_importance[:, :, 1:].reshape(
            len(ground_truth_importance), -1
        )
        dfs = {}
        for aggregate_method in aggregate_methods:
            df = pd.DataFrame()
            for cv, importance_unaggregated in self.importances.items():
                importance_scores = aggregate_scores(
                    importance_unaggregated, aggregate_method
                )
                importance_scores = importance_scores[:, :, 1:].reshape(
                    len(importance_scores), -1
                )

                # compute mean ranks
                ranks = rankdata(-importance_scores, axis=1)
                ranks_min = rankdata(-importance_scores, axis=1, method="min")
                gt_positions = np.where(ground_truth_importance)
                gt_ranks = ranks[gt_positions]
                gt_ranks_min = ranks_min[gt_positions]
                mean_rank = np.mean(gt_ranks)
                mean_rank_min = np.mean(gt_ranks_min)

                gt_score = ground_truth_importance.flatten()
                explainer_score = importance_scores.flatten()

                if absolutize:
                    explainer_score = np.abs(explainer_score)

                if np.any(np.isnan(explainer_score)):
                    self.log.warning("NaNs appear in explainer scores!")

                explainer_score = np.nan_to_num(explainer_score)
                auc_score = metrics.roc_auc_score(gt_score, explainer_score)
                aupr_score = metrics.average_precision_score(gt_score, explainer_score)
                prec_score, rec_score, thresholds = metrics.precision_recall_curve(
                    gt_score, explainer_score
                )
                auprc_score = (
                    metrics.auc(rec_score, prec_score) if rec_score.shape[0] > 1 else -1
                )

                pos_ratio = ground_truth_importance.sum() / len(ground_truth_importance)
                result = {
                    "Auroc": auc_score,
                    "Avpr": aupr_score,
                    "Auprc": auprc_score,
                    "Mean rank": mean_rank,
                    "Mean rank (min)": mean_rank_min,
                    "Pos ratio": pos_ratio,
                }
                self.log.info(f"cv={cv}")
                for k, v in result.items():
                    self.log.info(f"{k:20}: {v:.4f}")
                df[cv] = pd.Series(result)
            df = df.transpose()
            df.index.name = "cv"
            dfs[aggregate_method] = df
        df_all = pd.concat(dfs, axis=0)
        df_all.index.name = "aggregate method"
        return df_all

    def _get_mask_array_path(self) -> pathlib.Path:
        return self.out_path / self.dataset.get_name() / "array"

    def _get_cum_array_path(self) -> pathlib.Path:
        return self.out_path / self.dataset.get_name() / "cum_array"

    def _get_generator_array_path(self) -> pathlib.Path:
        return self.out_path / self.dataset.get_name() / "generator_array"

    def get_explainer_name(self) -> str:
        if self.explainers is None:
            return ""
        return next(iter(self.explainers.values())).get_name()

    def visualize(self):
        self.log.info("Visualizing masked data and feature attributions...")
        
        # Create plot directory if it doesn't exist
        plot_path = pathlib.Path(self.out_path)
        plot_path.mkdir(parents=True, exist_ok=True)
        
        # Get test data
        testset = list(self.dataset.test_loader.dataset)
        x_test = torch.stack([it[0] for it in testset]).cpu().numpy()
        mask_test = torch.stack([it[1] for it in testset]).cpu().numpy()
        y_test = torch.stack([it[4] for it in testset]).cpu().numpy()
        
        # Create BoxPlotter instance
        plotter = BoxPlotter(
            dataset=self.dataset,
            plot_path=plot_path,
            num_to_plot=min(300, len(x_test)),  # Limit to 20 samples or dataset size
            explainer_name=self.get_explainer_name(),
            time_difference=self.args["time_difference"]
        )

        plotter._plot_single_visualization(
            _inputs={0: x_test},
            y_test=y_test,
            input_name="Original_Input",
        )

        for mask_name, mask_data in self.masked_data.items():
            self.log.info(f"Visualizing for mask method: {mask_name}")
            for mode in ["drop"]:
                if mode not in mask_data:
                    self.log.warning(f"Mode {mode} not found in mask_data for {mask_name}")
                    continue
                
                mode_data = mask_data[mode]
                # Check if we have the necessary data
                if not all(k in mode_data for k in ["new_xs", "new_masks", "importance_masks", "new_preds"]):
                    self.log.warning(f"Missing required data for {mode} visualization in {mask_name}")
                    continue

                # Plot combined visualization
                plotter._plot_single_visualization(
                    _inputs={0: self.importances[0]},  # Use first CV's importances
                    y_test=y_test,
                    input_name=self.get_explainer_name(),
                )
        
        # For each masker, visualize the results
        """
        for mask_name, mask_data in self.masked_data.items():
            self.log.info(f"Visualizing for mask method: {mask_name}")
            
            # Process drop and preserve data separately
            for mode in ["drop", "preserve"]:
                if mode not in mask_data:
                    self.log.warning(f"Mode {mode} not found in mask_data for {mask_name}")
                    continue
                
                mode_data = mask_data[mode]
                
                # Check if we have the necessary data
                if not all(k in mode_data for k in ["new_xs", "new_masks", "importance_masks", "new_preds"]):
                    self.log.warning(f"Missing required data for {mode} visualization in {mask_name}")
                    continue
                
                # Prepare data for visualization
                cv_data = {
                    0: {  # Use 0 as the CV key since we're only visualizing one set
                        "new_xs": mode_data["new_xs"][0],
                        "new_masks": mode_data["new_masks"][0],
                        "importance_masks": mode_data["importance_masks"][0],
                        "new_preds": mode_data["new_preds"]
                    }
                }
                
                # Plot combined visualization
                plotter.plot_combined_visualization(
                    importances={0: self.importances[0]},  # Use first CV's importances
                    x_test=x_test,
                    mask_test=mask_test,
                    new_xs={0: cv_data[0]["new_xs"]},
                    new_masks={0: cv_data[0]["new_masks"]},
                    importance_masks={0: cv_data[0]["importance_masks"]},
                    orig_preds=orig_preds,
                    new_preds={0: cv_data[0]["new_preds"]},
                    y_test=y_test,
                    mask_name=f"{mask_name}_{mode}"
                )
                
                self.log.info(f"Completed visualization for {mask_name} - {mode}")
        
        # Also visualize the macro-level results if available
        for mask_name, mask_data in self.masked_data.items():
            for mode in ["macro_drop", "macro_preserve"]:
                if mode not in mask_data:
                    continue
                
                mode_data = mask_data[mode]
                
                # Check if we have the necessary data
                if not all(k in mode_data for k in ["new_xs", "new_masks", "importance_masks", "new_preds"]):
                    continue
                
                # Prepare data for visualization
                cv_data = {
                    0: {  # Use 0 as the CV key
                        "new_xs": mode_data["new_xs"][0],
                        "new_masks": mode_data["new_masks"][0],
                        "importance_masks": mode_data["importance_masks"][0],
                        "new_preds": mode_data["new_preds"]
                    }
                }
                
                # Plot combined visualization
                plotter.plot_combined_visualization(
                    importances={0: self.importances[0]},
                    x_test=x_test,
                    mask_test=mask_test,
                    new_xs={0: cv_data[0]["new_xs"]},
                    new_masks={0: cv_data[0]["new_masks"]},
                    importance_masks={0: cv_data[0]["importance_masks"]},
                    orig_preds=orig_preds,
                    new_preds={0: cv_data[0]["new_preds"]},
                    y_test=y_test,
                    mask_name=f"{mask_name}_{mode}"
                )
                
                self.log.info(f"Completed visualization for {mask_name} - {mode}")
        """
        
        self.log.info("Visualization completed.")
        
    def backward_fill(self, x, nan=0):        
        mask_all_zero = (x == 0).all(dim=(1, 2))  # (B,)

        zero_indices = torch.nonzero(mask_all_zero, as_tuple=True)[0]

        zero_samples = x[mask_all_zero]
        print(len(zero_samples))
        print(x.shape)
        
            
        B, T, F = x.shape
        print(x[0])
        
        nz = x != 0

        nz_t = nz.permute(0, 2, 1)              # (B, F, T)
        first_idx = nz_t.float().argmax(dim=2)  # (B, F)

        x_t = x.permute(0, 2, 1)                # (B, F, T)
        first_val = torch.gather(
            x_t, 2, first_idx.unsqueeze(2)
        ).squeeze(2)                           # (B, F)

        cum_nz = nz.cumsum(dim=1)              # (B, T, F)
        lead_mask = cum_nz == 0                # (B, T, F)
        print(first_val[0])
        fill_vals = first_val.unsqueeze(1).expand(B, T, F)
        out = x.clone()
        out = torch.where(lead_mask, fill_vals, out)
        
        print(out[0])
        raise RuntimeError

        return out
    
    def train_ood(self):
        from tqdm import tqdm
        for cv in self.dataset.cv_to_use():
            train_loader = self.dataset.train_loaders[cv]
            valid_loader = self.dataset.valid_loaders[cv]
            
            # from models.mae import MaskedAutoencoderTimeSeries
            from models.ae_lstm import TimeSeriesAutoencoder
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

            seq_len    = self.dataset.num_timesteps
            in_dim     = self.dataset.feature_size
            patch_size = 50
            embed_dim  = 64
            lr         = 1e-3
            n_epochs   = 50
            
            model = TimeSeriesAutoencoder(
                seq_len=seq_len,
                in_dim=in_dim,
                patch_size=patch_size,
                embed_dim=embed_dim,
                encoder_depth=3,
                decoder_depth=2,
                num_heads=1,
                # mask_ratio=0.6
            ).to(device)
            optimizer = torch.optim.Adam(model.parameters(), lr=lr)

            best_val_loss = float('inf')
            ckpt_path     = self.args['ckpt_save_path']

            for epoch in range(1, n_epochs+1):
                # --- Training ---
                model.train()
                total_train_loss = 0.0

                for batch in tqdm(train_loader):
                    x = batch[0]
                    x = x.to(device)

                    optimizer.zero_grad()
                    loss = model(x)
                    loss.backward()
                    optimizer.step()

                    total_train_loss += loss.item() * x.size(0)

                avg_train_loss = total_train_loss / len(train_loader.dataset)

                # --- Validation ---
                model.eval()
                total_val_loss = 0.0
                with torch.no_grad():
                    for batch in tqdm(valid_loader):
                        x = batch[0]
                        x = x.to(device)

                        loss = model(x)
                        total_val_loss += loss.item() * x.size(0)

                avg_val_loss = total_val_loss / len(valid_loader.dataset)

                print(f"[Epoch {epoch}/{n_epochs}] "
                    f"Train Loss: {avg_train_loss:.6f}  "
                    f"Val Loss:   {avg_val_loss:.6f}")

                # --- Save best model ---
                if avg_val_loss < best_val_loss:
                    best_val_loss = avg_val_loss
                    torch.save(model.state_dict(), ckpt_path)
                    print(f"Best model saved (val_loss: {best_val_loss:.6f})")
                    
        raise RuntimeError


def display_compact_dataframe(df, max_rows=20, max_colwidth=50, col_truncate=50):
    log = logging.getLogger("Base")

    original_settings = {
        "max_columns": pd.get_option("display.max_columns"),
        "width": pd.get_option("display.width"),
        "max_rows": pd.get_option("display.max_rows"),
        "float_format": pd.get_option("display.float_format"),
        "max_colwidth": pd.get_option("display.max_colwidth"),
    }

    try:
        # Configure display settings for compact output
        pd.set_option("display.max_columns", None)
        pd.set_option("display.width", None)
        pd.set_option("display.max_rows", max_rows)
        pd.set_option("display.float_format", "{:.2g}".format)
        pd.set_option("display.max_colwidth", max_colwidth)

        # Create copy of DataFrame to avoid modifying original
        df_display = df.copy()
        df_display.columns = [col[:col_truncate] for col in df_display.columns]

        output = ["\nResults:"]
        for col in df_display.columns:
            value = df_display[col].values[0] if len(df_display) > 0 else None
            output.append(" " * 4 + f"{col}: {value}")

        log.info("\n".join(output))

    finally:
        # Restore original settings
        for setting, value in original_settings.items():
            pd.set_option(f"display.{setting}", value)


if __name__ == "__main__":
    parser = get_args()
    data = parser.data
    explainers = parser.explainer
    skip_explain = parser.skipexplain
    eval_explain = parser.eval
    train_models = parser.train
    train_gen = parser.traingen
    train_ood = parser.trainood
    ood = parser.ood
    result_file = parser.resultfile
    explainer_seed = parser.explainerseed
    device = parser.device

    params = Params(parser.__dict__)
    dataset = params.datasets
    model_args = params.model_args
    model_train_args = params.model_train_args
    all_explainer_dict = params.all_explainer_dict
    out_path = params.outpath
    ckpt_path = params.ckptpath

    start_time = time.time()
    log = logging.getLogger("Base")
    for k, v in parser.__dict__.items():
        log.info(f"{k:20}: {v}")
    first = True
    save_failed = False

    parser.window_size = dataset.num_timesteps
    parser.feature_size = dataset.feature_size

    all_df = []
    print("Start to load data...")
    dataset.load_data()
    print("Loading successfully")
    runner = ExplanationRunner(parser.__dict__, dataset, device, out_path, ckpt_path)
    runner.init_model(**model_args)
    use_all_times = not isinstance(dataset, Mimic)
    if train_models:
        runner.train_model(**model_train_args, use_all_times=use_all_times)
    else:
        runner.load_model(use_all_times)
        
    if train_ood:
        runner.train_ood()

    if train_gen:
        generators_to_train = params.generators_to_train
        for explainer_name, explainer_dict_list in generators_to_train.items():
            for explainer_dict in explainer_dict_list:
                runner.get_explainers(parser.__dict__, explainer_name, explainer_dict=explainer_dict)
                log.info(f"Training Generator...Data={dataset.get_name()}, Explainer={explainer_name}")
                runner.train_generators(num_epochs=parser.epoch)
        log.info("Training Generator Done.")

    for explainer_name, explainer_dict_list in all_explainer_dict.items():
        for explainer_dict in explainer_dict_list:
            explainer_start_time = time.time()

            set_seed(explainer_seed)
            runner.clean_up(
                clean_importance=True, clean_explainer=True, clean_model=False
            )
            runner.get_explainers(
                parser.__dict__, explainer_name, explainer_dict=explainer_dict
            )
            runner.set_model_for_explainer(set_eval=explainer_name != "fit")

            if not skip_explain:
                log.info(
                    f"Running Explanations...Data={dataset.get_name()}, Explainer={explainer_name}, Dict={explainer_dict}"
                )

                runner.load_generators()
                runner.run_attributes()

                importances = runner.importances
                elapsed_times = runner.elapsed_times

                log.info("Explanations done.")

            if not eval_explain:
                continue

            log.info(
                f"Evaluating importance...Data={dataset.get_name()}, Explainer={explainer_name}"
            )

            if runner.importances is None:
                runner.load_importance()
            if isinstance(dataset, SimulatedData):
                df = runner.evaluate_simulated_importance(parser["aggregate"])
            elif ood:
                maskers = params.get_maskers(next(iter(runner.explainers.values())))
                df = runner.evaluate_performance_drop_ood(maskers)
            else:
                maskers = params.get_maskers(next(iter(runner.explainers.values())))
                df = runner.evaluate_performance_drop(maskers)
            log.info("Evaluating importance done.")

            df = df.reset_index()
            original_columns = df.columns
            df["Dataset"] = dataset.get_name()
            df["Explainer"] = next(iter(runner.explainers.values())).get_name()
            df["Elapsed Real Time"] = np.mean(list(elapsed_times.values()))
            df["Peak Memory"] = (
                f"{torch.cuda.max_memory_allocated(device) / 1024**2:.2f}"
            )

            df = df[
                ["Dataset", "Explainer"]
                + list(original_columns)
                + ["Elapsed Real Time", "Peak Memory"]
            ]

            all_df.append(df)
            result_path = out_path / dataset.get_name()
            result_path.mkdir(parents=True, exist_ok=True)
            error_code = append_df_to_csv(df, result_path / result_file)
            if first:
                first = False
            else:
                if error_code != 0:
                    save_failed = True
            display_compact_dataframe(df)

            if parser.vis:
                runner.visualize()
    log.info(f"All done! Time elapsed: {time.time() - start_time}")