models/admet_ft/_modules/dataset.py

"""
dataset.py - 데이터셋 클래스와 데이터 처리 로직
"""
import os
import sys
import numpy as np
import pandas as pd
from tqdm.auto import tqdm
from typing import Dict, List, Optional, Tuple, Union

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer

# 프로젝트 루트 경로 추가
_current_dir = os.path.dirname(os.path.abspath(__file__))
_project_root = os.path.dirname(os.path.dirname(_current_dir))
if _project_root not in sys.path:
    sys.path.insert(0, _project_root)
    
from admet_ft._modules.utils import (
    get_task_list, load_vocabs,
    DIDB_FILTER_COLS, LOAD_DATA_PATH, LOAD_PREPARED_DATA_PATH
)
from admet_ft._modules.scaler import (
    save_scaler, apply_scaler, save_scaler_power, apply_scaler_power,
    save_scaler_minmax, apply_scaler_minmax,
    save_scaler_adaptive, apply_scaler_adaptive
)


class ChemMultiTaskDataset(Dataset):
    """ChemBERTa 기반 Multi-Task Learning을 위한 데이터셋"""
    def __init__(
        self,
        df_dataset: pd.DataFrame,
        all_vocabs: Dict[str, Dict[str, int]],
        model_name: Optional[str] = None,
        task_type: str = 'cls',
    ):
        """
        Args:
            df_dataset: 입력 데이터프레임
            all_vocabs: 분류 태스크를 위한 레이블 매핑 사전
            model_name: 사용할 transformer 모델 이름 (ChemBERTa)
            task_type: 'cls' 또는 'reg' (분류 또는 회귀)
        """
        self.model_name = model_name
        self.task_type = task_type
        
        # 태스크 목록 설정
        self.y_cols = get_task_list(task_type)
        self.all_vocabs = all_vocabs

        # 데이터 저장
        self.df = df_dataset.copy()

        # Tokenizer는 __init__에서 준비만 하고, 실제 토큰화는 __getitem__에서 수행
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)

    def __len__(self):
        """데이터셋 길이 반환"""
        return len(self.df)

    def __getitem__(self, idx):
        """데이터셋에서 idx번째 항목 반환"""
        row = self.df.iloc[idx]

        # Tokenization은 여기서 수행
        objs = self.tokenizer(row['SMILES'], padding='max_length', max_length=510,
                              truncation=True, return_tensors='pt')

        # 출력 레이블 처리
        if self.task_type == 'cls':
            labels = []
            for target in self.y_cols:
                cls_value = str(row[target]).lower()
                if target in self.all_vocabs and cls_value in self.all_vocabs[target]:
                    label_id = self.all_vocabs[target][cls_value]
                else:
                    # 어휘에 없는 경우 기본값 사용
                    label_id = 0  # 기본값으로 0 사용
                labels.append(label_id)
        else:  # reg
            labels = [row[target] for target in self.y_cols]

        labels = torch.tensor(labels, dtype=torch.float32)

        # squeeze(0)로 배치 차원 제거
        return {'input_ids': objs['input_ids'].squeeze(0),
                'attention_mask': objs['attention_mask'].squeeze(0),
                'labels': labels,
                'smiles': row['SMILES']}  # SMILES 추가


class ChemMultiTaskDataModule:
    def __init__(
        self, 
        batch_size: int = 32, 
        scaling: bool = True, 
        scaler_path: str = "./results/scaler",
        task_type: str = 'cls', 
        model_name: Optional[str] = None, 
        missing_label_strategy: str = 'any',
        data_type: str = 'didb',
        load_type: str = 'default',
        num_workers: int = 2,
        output_dir: str = "./results",
        scaler_type: str = 'power'
    ):
        """
        Args:
            data_folder: 데이터 폴더 경로
            batch_size: 배치 크기
            scaling: 스케일링 적용 여부
            task_type: 'cls' 또는 'reg' (분류 또는 회귀)
            model_name: 사용할 transformer 모델 이름 (ChemBERTa)
            missing_label_strategy: 
                'any' - (default) 한 개 이상의 클래스에 라벨이 있으면 포함
                'all' - 모든 클래스 라벨이 있어야 포함
        """
        super().__init__()
        self.batch_size = batch_size
        self.model_name = model_name
        self.task_type = task_type
        self.missing_label_strategy = missing_label_strategy  # <--- 추가
        self.data_type = data_type
        self.load_type = load_type

        self.scaling = scaling
        self.scaler_path = scaler_path
        self.output_dir = output_dir
        self.scaler_type = scaler_type
        self.scaler_file: Optional[str] = None

        # multi-gpu training setting
        self.num_workers = max(0, int(num_workers))
        world_size_env = os.environ.get("WORLD_SIZE", "1")
        try:
            world_size = int(world_size_env)
        except ValueError:
            world_size = 1
        world_size = max(1, world_size)
        self.workers_per_rank = max(1, self.num_workers // world_size) if self.num_workers > 0 else 0

        # 데이터 유형에 따른 필터 컬럼 설정
        self.filter_cols = DIDB_FILTER_COLS

        # 태스크 목록 설정
        if task_type == 'cls':
            self.task_list = [f'{x}.cls' for x in self.filter_cols]
        else:
            self.task_list = self.filter_cols

        # 어휘 및 데이터셋 준비
        self.all_vocabs = None
        self.train_dataset, self.valid_dataset, self.test_dataset = None, None, None

        self.all_df = None
        self.train_df, self.valid_df, self.test_df = None, None, None

        # Default Data Loading
        if load_type == 'default':
            self.all_df = self._load_and_merge_data()
            self.all_df = self._nan_column_filter(self.all_df, self.filter_cols, self.missing_label_strategy)

        elif load_type == 'prepared':
            self.train_df, self.valid_df, self.test_df = self._load_and_merge_prepared_data()

            self.train_df = self._nan_column_filter(self.train_df, self.filter_cols, self.missing_label_strategy)
            self.valid_df = self._nan_column_filter(self.valid_df, self.filter_cols, self.missing_label_strategy)
            self.test_df = self._nan_column_filter(self.test_df, self.filter_cols, self.missing_label_strategy)

        else:
            raise ValueError(f"Unknown load_type: {load_type}")

    
    def _nan_column_filter(self, df_input: pd.DataFrame, filter_cols: List[str], filter_strategy: str = 'any') -> pd.DataFrame:
        df_input = df_input[['SMILES'] + filter_cols].reset_index(drop=True) # [1] 결측 행 처리 (여기서 옵션 적용)

        if filter_strategy == 'all':
            valid_rows = df_input[filter_cols].notna().all(axis=1) # [2] 모든 컬럼이 결측 아닌 행만 유지
        elif filter_strategy == 'any':
            valid_rows = df_input[filter_cols].notna().any(axis=1)
        else:
            raise ValueError(f"알 수 없는 missing_label_strategy: {filter_strategy}")
        
        df_input = df_input[valid_rows].reset_index(drop=True)
        return df_input


    def _check_numeric_columns(self, df: pd.DataFrame) -> pd.DataFrame:
        # SMILES 컬럼 보존 및 결측 처리
        if "SMILES" not in df.columns:
            df["SMILES"] = ""
        df["SMILES"] = df["SMILES"].fillna("").astype(str)

        for col in DIDB_FILTER_COLS:
            if col not in df.columns:
                df[col] = np.nan
            df[col] = pd.to_numeric(df[col], errors='coerce')

        return df

    def _load_and_merge_data(self):
        if self.data_type == "didb":
            path = LOAD_DATA_PATH["didb"]
        elif self.data_type == "portal":
            path = LOAD_DATA_PATH["portal"]
        elif self.data_type == "half":
            path = LOAD_DATA_PATH["all"]
        elif self.data_type == "all":
            path = LOAD_DATA_PATH["all"]
        else:
            raise ValueError(f"Unknown data_type: {self.data_type}")

        # data load and convert string to numeric
        df = pd.read_csv(path)
        df = self._check_numeric_columns(df)

        if self.data_type == "half":
            return df.sample(n=int(len(df) * 0.5), random_state=42)
        else:
            return df

    def _load_and_merge_prepared_data(self):
        if self.data_type == "didb":
            path = LOAD_PREPARED_DATA_PATH["didb"]
        elif self.data_type == "portal":
            path = LOAD_PREPARED_DATA_PATH["portal"]
        elif self.data_type == "half" or self.data_type == "all":
            path = LOAD_PREPARED_DATA_PATH["all"]
        elif self.data_type == "all_rating":
            path = LOAD_PREPARED_DATA_PATH["all_rating"]
        else:
            raise ValueError(f"Unknown data_type: {self.data_type}")

        # Load datasets
        train_df, valid_df, test_df = pd.read_csv(os.path.join(path, "train.csv")), pd.read_csv(os.path.join(path, "valid.csv")), pd.read_csv(os.path.join(path, "test.csv"))
        train_df, valid_df, test_df = self._check_numeric_columns(train_df), self._check_numeric_columns(valid_df), self._check_numeric_columns(test_df)

        if self.data_type == "half":
            train_df = train_df.sample(n=int(len(train_df) * 0.5), random_state=42)
            valid_df = valid_df.sample(n=int(len(valid_df) * 0.5), random_state=42)

        return train_df, valid_df, test_df

    def setup(self):
        """데이터셋 준비 및 분할"""
        self.scaler_file = None

        # 어휘 로딩 (분류 작업용)
        if self.task_type == 'cls':
            source_df = self.all_df if self.load_type == 'default' else self.train_df
            cls_targets = [x for x in source_df.columns.tolist() if x.endswith('.cls')]
            self.all_vocabs = load_vocabs(self.output_dir, cls_targets)
        else:
            self.all_vocabs = {}

        # Scaler 함수 선택
        if self.scaler_type == 'power':
            save_fn = save_scaler_power
            apply_fn = apply_scaler_power
            scaler_filename = 'scaler_power_config.csv'
        elif self.scaler_type == 'minmax':
            save_fn = save_scaler_minmax
            apply_fn = apply_scaler_minmax
            scaler_filename = 'scaler_minmax_config.csv'
        elif self.scaler_type == 'adapt':
            save_fn = save_scaler_adaptive
            apply_fn = apply_scaler_adaptive
            scaler_filename = 'scaler_adapt.csv'
        else:  # zscore (default)
            save_fn = save_scaler
            apply_fn = apply_scaler
            scaler_filename = 'scaler_config.csv'

        if self.load_type == 'default':
            # 학습/검증/테스트 데이터셋 분할
            data_save_path = os.path.join(self.output_dir, "valid_test_split/")
            os.makedirs(data_save_path, exist_ok=True)

            if self.scaling:
                if not self.scaler_path:
                    raise ValueError("Scaler path must be provided when scaling is enabled.")
                scaler_file_path = os.path.join(self.scaler_path, scaler_filename)

                # Adaptive scaler는 feature_cols 파라미터가 필요
                if self.scaler_type == 'adapt':
                    self.all_df = save_fn(self.all_df, scaler_path=scaler_file_path, feature_cols=self.filter_cols)
                else:
                    self.all_df = save_fn(self.all_df, scaler_path=scaler_file_path)

            T = int(len(self.all_df) * 0.7)  # 70% 훈련 데이터
            train_df = self.all_df.sample(n=T, random_state=42)
            t_indexs = self.all_df.index.isin(train_df.index)
            other_df = self.all_df[~t_indexs]
            
            V = int(len(other_df) * 0.5)  # 남은 데이터의 50%를 검증 데이터로
            valid_df = other_df.sample(n=V, random_state=42)
            v_indexs = other_df.index.isin(valid_df.index)
            test_df = other_df[~v_indexs]

            self.train_df = train_df
            self.valid_df = valid_df
            self.test_df = test_df

            self.valid_df.to_csv(os.path.join(data_save_path, "valid.csv"), index=False)
            self.test_df.to_csv(os.path.join(data_save_path, "test.csv"), index=False)

            # 데이터셋 생성
            self.train_dataset = ChemMultiTaskDataset(
                self.train_df, self.all_vocabs, self.model_name, self.task_type)
            
            self.valid_dataset = ChemMultiTaskDataset(
                self.valid_df, self.all_vocabs, self.model_name, self.task_type)
            
            self.test_dataset = ChemMultiTaskDataset(
                self.test_df, self.all_vocabs, self.model_name, self.task_type)

        else:
            if self.scaling: # [3] 스케일링 적용 (기존 방식): 전체 데이터프레임에 대해 표준화 수행 및 저장
                if not self.scaler_path:
                    raise ValueError("Scaler path must be provided when scaling is enabled.")
                scaler_file_path = os.path.join(self.scaler_path, scaler_filename)

                # Index 백업 (재분할을 위해)
                train_indices = self.train_df.index
                valid_indices = self.valid_df.index

                # Train/Valid 병합하여 Scaler fit+transform
                merged_df = pd.concat([self.train_df, self.valid_df], axis=0)

                # Adaptive scaler는 feature_cols 파라미터가 필요
                if self.scaler_type == 'adapt':
                    scaled_merged = save_fn(merged_df, scaler_path=scaler_file_path, feature_cols=self.filter_cols)
                else:
                    scaled_merged = save_fn(merged_df, scaler_path=scaler_file_path)

                # Index 기반 재분할 (스케일링된 데이터)
                self.train_df = scaled_merged.loc[train_indices].copy()
                self.valid_df = scaled_merged.loc[valid_indices].copy()

                # Test 데이터는 원본 스케일 유지 (일반화 성능 평가 목적)
                # self.test_df는 그대로 유지

                self.scaler_file = scaler_file_path
            else:
                self.scaler_file = None
                
            self.train_dataset = ChemMultiTaskDataset(
                self.train_df, self.all_vocabs, self.model_name, self.task_type)
            self.valid_dataset = ChemMultiTaskDataset(
                self.valid_df, self.all_vocabs, self.model_name, self.task_type)
            self.test_dataset = ChemMultiTaskDataset(
                self.test_df, self.all_vocabs, self.model_name, self.task_type)


    def get_dataloaders(self) -> Tuple[DataLoader, DataLoader, DataLoader]:
        """훈련, 검증, 테스트 데이터로더 반환"""
        if self.train_dataset is None:
            self.setup()
        
        train_size = len(self.train_dataset) if self.train_dataset else 0
        valid_size = len(self.valid_dataset) if self.valid_dataset else 0
        test_size = len(self.test_dataset) if self.test_dataset else 0
        print(f"[DataModule] Train samples: {train_size}, Val samples: {valid_size}, Test samples: {test_size}")
        
        train_loader = DataLoader(
            self.train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            collate_fn=self._collate_fn,
            num_workers=self.workers_per_rank,
            pin_memory=True,
            drop_last=False,
            persistent_workers=self.workers_per_rank > 0
        )
        
        valid_loader = DataLoader(
            self.valid_dataset,
            batch_size=self.batch_size,
            collate_fn=self._collate_fn,
            num_workers=self.workers_per_rank,
            pin_memory=True,
            drop_last=False,
            persistent_workers=self.workers_per_rank > 0
        )
        
        test_loader = DataLoader(
            self.test_dataset,
            batch_size=self.batch_size,
            collate_fn=self._collate_fn,
            num_workers=self.workers_per_rank,
            pin_memory=True,
            drop_last=False,
            persistent_workers=self.workers_per_rank > 0
        )
        
        return train_loader, valid_loader, test_loader
    
    def get_prediction_dataloader(self, df_predict: pd.DataFrame) -> DataLoader:
        """예측용 데이터로더 반환"""
        predict_dataset = ChemMultiTaskDataset(
            df_predict, self.all_vocabs, self.model_name, self.task_type)
        
        predict_loader = DataLoader(
            predict_dataset,
            batch_size=self.batch_size,
            collate_fn=self._collate_fn,
            num_workers=self.workers_per_rank,
            pin_memory=True,
            drop_last=False,
            persistent_workers=self.workers_per_rank > 0
        )
        
        return predict_loader
    
    def _collate_fn(self, batch):
        """배치 데이터 조합 함수"""
        if not batch:
            raise ValueError("Empty batch encountered")

        # 레이블 처리
        labels = torch.stack([item['labels'] for item in batch])

        # 입력 ID 및 어텐션 마스크 처리
        input_ids = torch.stack([item['input_ids'] for item in batch])
        attention_mask = torch.stack([item['attention_mask'] for item in batch])

        # SMILES 수집
        smiles = [item['smiles'] for item in batch]

        return {
            'input_ids': input_ids,
            'attention_mask': attention_mask,
            'labels': labels,
            'smiles': smiles  # SMILES 추가
        }