from torch.utils.data import Dataset, DataLoader, random_split
import pandas as pd
import torch
import h5py
import ast
import numpy as np
import pickle
import os

def ko_embedding_padding(embedding, max_len):
    pad_rows = max_len - embedding.shape[0]
    padded_embeddings = np.pad(
        embedding,
        pad_width=((0, pad_rows), (0, 0)),
        mode="constant",
        constant_values=0,
    )
    # mask = np.full(max_len, -np.inf, dtype=np.float32)
    # mask[:embedding.shape[0]] = 0.0
    mask = np.ones(max_len, dtype=bool)
    mask[:embedding.shape[0]] = False
    
    return padded_embeddings, mask

def get_input_token(inputs, vo_cab, max_len):
    inputs = sorted(inputs)
    input_token = [vo_cab['</s>']]
    for j in inputs:
        input_token.append(vo_cab[j])
    input_token.append(vo_cab['<s>'])
    padding_len = max_len - len(input_token)
    if padding_len > 0:
        input_token += [vo_cab['blank']] * padding_len
    return input_token

def get_dict(data, colum):
    all_list = list(data[colum])
    all_list = (pd.DataFrame(all_list, columns=[colum])
                  .drop_duplicates([colum])
                  .sort_values(by=colum)
                  .reset_index(drop=True))
    return dict(zip(all_list[colum], all_list.index))

class DataProcessor:
    def __init__(self, ko_count_path, data_path, embedding_h5_path, vo_cab_pkl_path, model_type, ko_max_len=4500, compound_max_len=98):
        self.ko_count_path = ko_count_path
        self.data_path = data_path
        self.embedding_h5_path = embedding_h5_path
        self.vo_cab_pkl_path = vo_cab_pkl_path
        self.model_type = model_type
        self.ko_max_len = ko_max_len
        self.compound_max_len = compound_max_len

        self.KO_count = None
        self.data = None
        self.vo_cab = None
        self.genome_dict = None
        self.medium_dict = None

    def load_data(self, count_min=800, count_max=4500):
        self.KO_count = pd.read_csv(self.ko_count_path)
        self.KO_count = self.KO_count[(self.KO_count['count'] >= count_min) & (self.KO_count['count'] <= count_max)]

        self.data = pd.read_csv(self.data_path)
        self.data = pd.merge(self.KO_count, self.data, how='inner', on='genome').drop_duplicates(['genome', 'compounds_list'])
        self.data['compounds_list'] = self.data['compounds_list'].apply(lambda x: list(ast.literal_eval(x)))
        
        
        # Load vo_cab: use pkl file
        if self.vo_cab_pkl_path and os.path.exists(self.vo_cab_pkl_path):
            print(f"Loading predefined compound_cab from {self.vo_cab_pkl_path}")
            with open(self.vo_cab_pkl_path, 'rb') as f:
                data_pkl = pickle.load(f)
                if isinstance(data_pkl, dict) and 'compound_cab' in data_pkl:
                    self.vo_cab = data_pkl['compound_cab']
                    print(f"Successfully loaded compound_cab, vocabulary size: {len(self.vo_cab)}")
                    
        # genome and medium dictionaries
        self.genome_dict = get_dict(self.data, 'genome')
        self.medium_dict = get_dict(self.data, 'media_name')

        # 保存 genome_dict 和 medium_dict 到文件
        genome_filename = f"{self.model_type}_genome_dict.pkl"
        medium_filename = f"{self.model_type}_medium_dict.pkl"
        # 保存 genome_dict
        with open(genome_filename, 'wb') as f:
            pickle.dump(self.genome_dict, f)
        # 保存 medium_dict
        with open(medium_filename, 'wb') as f:
            pickle.dump(self.medium_dict, f)
            
    def get_sample_by_index(self, index):
        """根据索引获取单个样本的数据"""
        if self.data is None:
            raise ValueError("请先调用load_data()方法加载数据")
        
        if index >= len(self.data):
            raise IndexError("索引超出数据范围")
        
        row = self.data.iloc[index]
        genome = row['genome']
        compounds = row['compounds_list']
        medium = row['media_name']
        
        # 从HDF5文件读取embedding
        with h5py.File(self.embedding_h5_path, "r") as hf:
            try:
                embeddings = hf[genome][:]
            except KeyError:
                print(f"⚠️ Genome {genome} not found in HDF5")
                return None
        
        # 处理token和mask
        genome_token = self.genome_dict[genome]
        medium_token = self.medium_dict[medium]
        compound_token = get_input_token(inputs=compounds, vo_cab=self.vo_cab, max_len=self.compound_max_len)
        
        embedding_pad, mask = ko_embedding_padding(embeddings, self.ko_max_len)
        
        return {
            'ko_token': torch.from_numpy(embedding_pad).float(),
            'compound_token': torch.tensor(compound_token),
            'ko_mask': torch.tensor(mask),
            'genome_token': torch.tensor(genome_token),
            'medium_token': torch.tensor(medium_token)
        }
    
    def get_data_length(self):
        """获取数据集长度"""
        if self.data is None:
            raise ValueError("请先调用load_data()方法加载数据")
        return len(self.data)

    def process_test_data(self, genome_list=None):
        # 如果 genome_list 为空，就使用整个数据集
        if genome_list is None:
            test_data = self.data
        else:
            test_data = self.data[self.data['genome'].isin(genome_list)]
    
        ko_token = []
        compound_token = []
        medium_token = []
        genome_token = []
        ko_mask = []
    
        with h5py.File(self.embedding_h5_path, "r") as hf:
            for genome, compounds, medium in zip(test_data['genome'], test_data['compounds_list'], test_data['media_name']):
                try:
                    embeddings = hf[genome][:]
                except KeyError:
                    print(f"⚠️ Genome {genome} not found in HDF5, skipped")
                    continue  # 跳过这个 genome
                
                # 处理 token 和 mask
                genome_token.append(self.genome_dict[genome])
                medium_token.append(self.medium_dict[medium])
                compound_token.append(get_input_token(inputs=compounds, vo_cab=self.vo_cab, max_len=self.compound_max_len))
                
                embedding_pad, mask = ko_embedding_padding(embeddings, self.ko_max_len)
                ko_token.append(embedding_pad)
                ko_mask.append(mask)
    
        # 转成 tensor
        medium_token = torch.tensor(medium_token)
        genome_token = torch.tensor(genome_token)
        ko_mask = torch.tensor(ko_mask)
        ko_token = torch.stack([torch.from_numpy(x).float() for x in ko_token])
        compound_token = torch.tensor(compound_token)
    
        return ko_token, compound_token, ko_mask, genome_token, medium_token


class MyDataset(Dataset):
    def __init__(self, src, trg, src_mask, genome, medium):
        self.src = src
        self.trg = trg
        self.src_mask = src_mask
        self.genome = genome
        self.medium = medium

    def __getitem__(self, index):
        return self.src[index], self.trg[index], self.src_mask[index], self.genome[index], self.medium[index]

    def __len__(self):
        return self.trg.size(0)


class LazyDataset(Dataset):
    """支持动态加载数据的Dataset类"""
    def __init__(self, data_processor, indices=None):
        self.data_processor = data_processor
        if indices is None:
            self.indices = list(range(data_processor.get_data_length()))
        else:
            self.indices = indices

    def __getitem__(self, index):
        # 根据索引动态从DataProcessor获取数据
        actual_index = self.indices[index]
        sample = self.data_processor.get_sample_by_index(actual_index)
        
        if sample is None:
            # 如果无法获取数据，返回下一个有效样本
            for i in range(index + 1, len(self.indices)):
                actual_index = self.indices[i]
                sample = self.data_processor.get_sample_by_index(actual_index)
                if sample is not None:
                    break
            
            if sample is None:
                raise RuntimeError(f"无法获取索引 {index} 的数据")
        
        return (sample['ko_token'], sample['compound_token'], sample['ko_mask'], 
                sample['genome_token'], sample['medium_token'])

    def __len__(self):
        return len(self.indices)

# def build_dataloaders(ko_token, compound_token, ko_mask, genome_token, medium_token, batch_size):
#     dataset = MyDataset(ko_token, compound_token, ko_mask, genome_token, medium_token)
#     n = len(dataset)
#     train_size = int(n * 10 / 10)
#     # valid_size = int((n - train_size) / 2)
#     # test_size = n - train_size - valid_size

#     # train_ds, valid_ds, test_ds = random_split(dataset, [train_size, valid_size, test_size])
#     train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
#     # train_loader_no_shuffle = DataLoader(train_ds, batch_size=batch_size, shuffle=False)
#     # valid_loader = DataLoader(valid_ds, batch_size=batch_size, shuffle=False)
#     # test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False)

#     # return train_loader, train_loader_no_shuffle, valid_loader, test_loader
#     return train_loader

def build_dataloaders(ko_token, compound_token, ko_mask, genome_token, medium_token, batch_size, shuffle):
    dataset = MyDataset(ko_token, compound_token, ko_mask, genome_token, medium_token)
    train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
    return train_loader  # 返回单个 DataLoader


def build_single_lazy_dataloader(data_processor, batch_size=1, shuffle=False):
    """为单个数据集创建懒加载数据加载器"""
    dataset = LazyDataset(data_processor)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
    return dataloader


def build_lazy_dataloaders(data_processor, batch_size=1, split_ratio=8, shuffle=True):
    """使用LazyDataset构建数据加载器，支持动态加载（保持向后兼容）"""
    total_length = data_processor.get_data_length()
    indices = list(range(total_length))
    
    # 计算分割点
    train_size = int(total_length * split_ratio / 10)
    valid_size = int((total_length - train_size) / 2)
    
    # 分割索引
    if shuffle:
        import random
        random.shuffle(indices)
    
    train_indices = indices[:train_size]
    valid_indices = indices[train_size:train_size + valid_size]
    test_indices = indices[train_size + valid_size:]
    
    # 创建数据集和数据加载器
    train_dataset = LazyDataset(data_processor, train_indices)
    valid_dataset = LazyDataset(data_processor, valid_indices)
    test_dataset = LazyDataset(data_processor, test_indices)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle)
    train_loader_no_shuffle = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)
    valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    
    return train_loader, train_loader_no_shuffle, valid_loader, test_loader

# Dataset class specifically for model inference (using unseen data)
class InferenceDataset(Dataset):
    """
    Dataset class for loading data during inference phase.
    Designed to handle unseen data after model training.
    """
    def __init__(self, src, src_mask, genome):
        """
        Initialize the inference dataset.
        
        Args:
            src: Input features (e.g., embeddings) with shape [num_samples, seq_len, d_model]
            src_mask: Corresponding mask tensor with shape [num_samples, seq_len]
        """
        self.src = src  # Input features (e.g., ko_token embeddings)
        self.src_mask = src_mask  # Mask for input features
        self.genome = genome

    def __getitem__(self, index):
        """
        Get a single sample from the dataset by index.
        
        Args:
            index: Index of the sample to retrieve
            
        Returns:
            Tuple containing (src_feature, src_mask) for the specified index
        """
        return self.src[index], self.src_mask[index], self.genome[index]

    def __len__(self):
        """
        Return the total number of samples in the dataset.
        
        Returns:
            int: Number of samples (matches length of source features)
        """
        return len(self.src)


def build_inference_dataloader(ko_token, ko_mask, genome, batch_size=1, shuffle=False):
    """
    Build a DataLoader for model inference on unseen data.
    
    Args:
        ko_token: Input features (embeddings) for inference, shape [num_samples, seq_len, d_model]
        ko_mask: Corresponding mask tensor, shape [num_samples, seq_len]
        batch_size: Batch size for inference (default: 1)
        shuffle: Whether to shuffle data (default: False for inference to preserve order)
        
    Returns:
        DataLoader: Ready-to-use DataLoader for inference
    """
    # Initialize dataset with inference data
    dataset = InferenceDataset(
        src=ko_token,
        src_mask=ko_mask,
        genome=genome
    )
    
    # Create DataLoader with inference-optimized settings
    return DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle
    )