#!/usr/bin/env python3
"""
BrainAnytime Inference Engine for Hugging Face Space

提供模型加载、推理和结果处理功能。
支持4个任务：CN vs AD, CN vs MCI, MMSE, AGE
支持5种模态组合：T, TF, TMF, TFP, TMFP
"""

import os
import sys
import warnings
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Union
import json

import numpy as np
import torch
import torch.nn as nn
import nibabel as nib

warnings.filterwarnings("ignore")

# 添加 BrainAnytime 代码路径
BASE_DIR = Path(__file__).parent
BRAINANYTIME_DIR = BASE_DIR / "BrainAnytime"
if BRAINANYTIME_DIR.exists():
    sys.path.insert(0, str(BRAINANYTIME_DIR))

# 导入 BrainAnytime 模型
from models.multimae3d import create_multimae3d


# =============================================================================
# 配置常量
# =============================================================================

# 模态配置
MODALITY_ORDER = ['T1', 'T2', 'Flair', 'PET']
MODALITY_SHORT = {'T1': 'T', 'T2': 'M', 'Flair': 'F', 'PET': 'P'}
SHORT_TO_FULL = {'T': 'T1', 'M': 'T2', 'F': 'Flair', 'P': 'PET'}

# 任务配置
TASKS = {
    'CN_vs_AD': {
        'type': 'classification',
        'display_name': 'CN vs AD',
        'num_classes': 2,
        'classes': ['CN', 'AD'],
    },
    'CN_vs_MCI': {
        'type': 'classification',
        'display_name': 'CN vs MCI',
        'num_classes': 2,
        'classes': ['CN', 'MCI'],
    },
    'MMSE': {
        'type': 'regression',
        'display_name': 'MMSE Score',
        'min': 10.0,
        'max': 30.0,
    },
    'AGE': {
        'type': 'regression',
        'display_name': 'Age',
        'min': 50.0,
        'max': 100.0,
    },
}

# 模型默认参数（与训练时一致）
# IMPORTANT: decoder_embed_dim must be divisible by 6 for 3D sincos position embedding
DEFAULT_MODEL_ARGS = {
    'img_size': 128,
    'patch_size': 16,
    'embed_dim': 768,
    'depth': 12,
    'num_heads': 12,
    'decoder_embed_dim': 384,  # Must be divisible by 6 (384/6=64) ✓
    'decoder_depth': 2,        # Pretrain default
    'decoder_num_heads': 12,
    'pool': 'mean',
    'dropout': 0.5,
}

# 归一化参数
MMSE_MIN, MMSE_MAX = 10.0, 30.0


# =============================================================================
# 下游任务模型头
# =============================================================================

class MultiMAE3DForDownstream(nn.Module):
    """
    MultiMAE3D + 下游任务头 (分类/回归)
    从 BrainAnytime/finetune_main.py 复制
    """
    def __init__(
        self,
        encoder,
        embed_dim: int = 768,
        num_outputs: int = 1,
        pool: str = 'mean',
        dropout: float = 0.5,
    ):
        super().__init__()
        self.encoder = encoder
        self.pool = pool
        self.num_patches_per_modality = encoder.num_patches
        self.num_global_tokens = encoder.num_global_tokens
        
        # LayerNorm before head (important for checkpoint compatibility)
        self.norm = nn.LayerNorm(embed_dim)
        
        # 预测头
        self.head = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(embed_dim, num_outputs)
        )
        
    def forward(self, images: torch.Tensor, observed: torch.Tensor) -> torch.Tensor:
        """
        Args:
            images: [B, 4, D, H, W] - 4 modalities
            observed: [B, 4] - 0/1 mask for available modalities
        Returns:
            logits: [B, num_outputs]
        """
        # encode() returns [B, 1 + 4*num_patches, embed_dim]
        encoder_out = self.encoder.encode(images, observed)

        if self.pool == 'cls':
            features = encoder_out[:, 0]  # CLS token -> [B, D]
        elif self.pool == 'mean':
            # Mean pool over modality tokens with masking for missing modalities
            tokens = encoder_out[:, self.num_global_tokens:]  # [B, 4*N_p, D]
            B, _, D = tokens.shape
            N = self.num_patches_per_modality
            # Build per-token mask: repeat each modality's observed flag N times
            mask = observed.unsqueeze(-1).expand(-1, -1, N)  # [B, 4, N]
            mask = mask.reshape(B, 4 * N).unsqueeze(-1)  # [B, 4*N, 1]
            features = (tokens * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
        else:
            raise ValueError(f"Unknown pool type: {self.pool}")

        features = self.norm(features)
        logits = self.head(features)
        return logits


# =============================================================================
# 推理引擎
# =============================================================================

class BrainAnytimeInference:
    """
    BrainAnytime 推理引擎
    
    用法:
        engine = BrainAnytimeInference(checkpoints_dir)
        engine.load_model('CN_vs_AD')
        result = engine.predict(nifti_data, modalities)
    """
    
    def __init__(
        self,
        checkpoints_dir: Optional[str] = None,
        device: Optional[str] = None,
    ):
        """
        Args:
            checkpoints_dir: 模型检查点目录，默认从 HF Hub 下载
            device: 'cuda' 或 'cpu'，默认自动选择
        """
        self.checkpoints_dir = checkpoints_dir
        self.device = torch.device(
            device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        )
        print(f"Inference device: {self.device}")
        
        # 缓存加载的模型
        self.models: Dict[str, MultiMAE3DForDownstream] = {}
        
    def _get_checkpoint_path(self, task: str) -> str:
        """获取任务对应的 checkpoint 路径"""
        checkpoint_files = {
            'CN_vs_AD': 'CN_vs_AD_seed_0_best.pth',
            'CN_vs_MCI': 'CN_vs_MCI_seed_0_best.pth',
            'MMSE': 'MMSE_seed_0_best.pth',
            'AGE': 'AGE_seed_0_best.pth',
        }
        
        if task not in checkpoint_files:
            raise ValueError(f"Unknown task: {task}. Available: {list(checkpoint_files.keys())}")
        
        if self.checkpoints_dir:
            path = os.path.join(self.checkpoints_dir, checkpoint_files[task])
            if os.path.exists(path):
                return path
        
        # 从 Hugging Face Hub 下载
        try:
            from huggingface_hub import hf_hub_download
            path = hf_hub_download(
                repo_id="Simmonstt/BrainAnytime",
                filename=checkpoint_files[task],
            )
            return path
        except Exception as e:
            raise RuntimeError(
                f"Cannot load checkpoint for {task}. "
                f"Please provide checkpoints_dir or ensure HF Hub access. Error: {e}"
            )
    
    def load_model(self, task: str) -> MultiMAE3DForDownstream:
        """
        加载指定任务的模型（懒加载）
        
        Args:
            task: 任务名称 ('CN_vs_AD', 'CN_vs_MCI', 'MMSE', 'AGE')
        
        Returns:
            加载好的模型
        """
        if task in self.models:
            return self.models[task]
        
        print(f"Loading model for task: {task}")
        
        # 创建编码器
        encoder = create_multimae3d(
            img_size=DEFAULT_MODEL_ARGS['img_size'],
            patch_size=DEFAULT_MODEL_ARGS['patch_size'],
            embed_dim=DEFAULT_MODEL_ARGS['embed_dim'],
            depth=DEFAULT_MODEL_ARGS['depth'],
            num_heads=DEFAULT_MODEL_ARGS['num_heads'],
            decoder_embed_dim=DEFAULT_MODEL_ARGS['decoder_embed_dim'],
            decoder_depth=DEFAULT_MODEL_ARGS['decoder_depth'],
            decoder_num_heads=DEFAULT_MODEL_ARGS['decoder_num_heads'],
        )
        
        # 创建下游任务模型
        model = MultiMAE3DForDownstream(
            encoder=encoder,
            embed_dim=DEFAULT_MODEL_ARGS['embed_dim'],
            num_outputs=1,
            pool=DEFAULT_MODEL_ARGS['pool'],
            dropout=DEFAULT_MODEL_ARGS['dropout'],
        ).to(self.device)
        
        # 加载 checkpoint
        checkpoint_path = self._get_checkpoint_path(task)
        print(f"  Loading checkpoint: {checkpoint_path}")
        
        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        model.load_state_dict(ckpt['model_state_dict'])
        
        print(f"  Loaded (epoch={ckpt.get('epoch', '?')}, "
              f"best_metric={ckpt.get('best_metric', '?')})")
        
        model.eval()
        self.models[task] = model
        
        return model
    
    def preprocess_nifti(
        self,
        nifti_path: str,
        target_size: Tuple[int, int, int] = (128, 128, 128),
    ) -> Optional[np.ndarray]:
        """
        加载并预处理 NIfTI 文件
        
        Args:
            nifti_path: NIfTI 文件路径
            target_size: 目标尺寸 (D, H, W)
        
        Returns:
            预处理后的数据 [D, H, W]，失败返回 None
        """
        try:
            # 加载
            nii = nib.load(nifti_path)
            data = nii.get_fdata().astype(np.float32)
            
            # 确保 3D
            if data.ndim == 4:
                data = data[..., 0]
            
            # 检查尺寸
            if data.shape != target_size:
                print(f"  Warning: Size mismatch {data.shape} != {target_size}")
                # 如果需要，可以在这里添加 resize 逻辑
                # 但目前假设输入已经是 128x128x128
                return None
            
            # Min-Max 归一化
            data_min, data_max = data.min(), data.max()
            if data_max > data_min:
                data = (data - data_min) / (data_max - data_min)
            else:
                print(f"  Warning: Empty image (max == min)")
                return None
            
            return data
            
        except Exception as e:
            print(f"  Error loading {nifti_path}: {e}")
            return None
    
    def prepare_input(
        self,
        nifti_files: Dict[str, str],
        modality_combo: str,
    ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
        """
        准备模型输入
        
        Args:
            nifti_files: {模态名: 文件路径}，如 {'T1': 'path/to/T1.nii.gz'}
            modality_combo: 模态组合字符串，如 'T', 'TF'
        
        Returns:
            (images, observed) 或 None
            - images: [1, 4, 128, 128, 128]
            - observed: [1, 4]
        """
        expected_modalities = [SHORT_TO_FULL[c] for c in modality_combo]
        
        # 加载所有模态
        images_list = []
        observed_list = []
        
        for mod in MODALITY_ORDER:
            if mod in expected_modalities and mod in nifti_files:
                data = self.preprocess_nifti(nifti_files[mod])
                if data is not None:
                    images_list.append(data)
                    observed_list.append(1.0)
                else:
                    return None
            else:
                # 缺失模态用零填充
                images_list.append(np.zeros((128, 128, 128), dtype=np.float32))
                observed_list.append(0.0)
        
        # 转换为张量
        images = np.stack(images_list, axis=0)  # [4, 128, 128, 128]
        images = torch.from_numpy(images).unsqueeze(0).to(self.device)  # [1, 4, 128, 128, 128]
        
        observed = torch.tensor(observed_list, dtype=torch.float32).unsqueeze(0).to(self.device)
        
        return images, observed
    
    @torch.no_grad()
    def predict(
        self,
        nifti_files: Dict[str, str],
        task: str,
        modality_combo: str,
    ) -> Optional[Dict]:
        """
        执行推理
        
        Args:
            nifti_files: {模态名: 文件路径}
            task: 任务名称
            modality_combo: 模态组合字符串
        
        Returns:
            推理结果字典
        """
        # 加载模型
        model = self.load_model(task)
        task_config = TASKS[task]
        
        # 准备输入
        input_data = self.prepare_input(nifti_files, modality_combo)
        if input_data is None:
            return None
        
        images, observed = input_data
        
        # 推理
        logits = model(images, observed)
        
        # 后处理
        if task_config['type'] == 'classification':
            # 分类：sigmoid → probability
            prob = torch.sigmoid(logits).item()
            pred_class = 1 if prob > 0.5 else 0
            pred_label = task_config['classes'][pred_class]
            confidence = prob if pred_class == 1 else 1 - prob
            
            result = {
                'task': task,
                'task_type': 'classification',
                'prediction': pred_label,
                'probability': prob,
                'confidence': confidence,
                'classes': task_config['classes'],
                'logits': logits.item(),
            }
        else:
            # 回归：反归一化
            normalized_pred = logits.item()
            
            if task == 'MMSE':
                # 反归一化到 [10, 30]
                pred = normalized_pred * (MMSE_MAX - MMSE_MIN) + MMSE_MIN
                pred = max(MMSE_MIN, min(MMSE_MAX, pred))
                unit = 'points'
                reference_range = [MMSE_MIN, MMSE_MAX]
            else:  # AGE
                # 假设 AGE 是 z-score，这里简化处理
                pred = normalized_pred * 20 + 75  # 近似反归一化
                unit = 'years'
                reference_range = [50, 100]
            
            result = {
                'task': task,
                'task_type': 'regression',
                'prediction': pred,
                'unit': unit,
                'reference_range': reference_range,
                'normalized_value': normalized_pred,
                'logits': logits.item(),
            }
        
        # 添加输入信息
        result['input'] = {
            'modality_combo': modality_combo,
            'modalities': [SHORT_TO_FULL[c] for c in modality_combo],
            'files': nifti_files,
        }
        
        return result
    
    def predict_from_sample(
        self,
        sample_dir: str,
        task: str,
        modality_combo: str,
    ) -> Optional[Dict]:
        """
        从样本目录执行推理
        
        Args:
            sample_dir: 样本目录路径 (包含 .nii.gz 文件)
            task: 任务名称
            modality_combo: 模态组合
        
        Returns:
            推理结果
        """
        sample_dir = Path(sample_dir)
        
        # 查找 NIfTI 文件
        expected_modalities = [SHORT_TO_FULL[c] for c in modality_combo]
        nifti_files = {}
        
        for mod in expected_modalities:
            # 查找匹配的文件
            pattern = f"*{mod}*.nii.gz"
            matches = list(sample_dir.glob(pattern))
            
            if not matches:
                # 尝试更宽松的匹配
                pattern = f"*.nii.gz"
                for f in sample_dir.glob(pattern):
                    if mod.lower() in f.name.lower():
                        matches.append(f)
                        break
            
            if matches:
                nifti_files[mod] = str(matches[0])
            else:
                print(f"  Error: Cannot find {mod} in {sample_dir}")
                return None
        
        return self.predict(nifti_files, task, modality_combo)


# =============================================================================
# 辅助函数
# =============================================================================

def create_inference_engine(checkpoints_dir: Optional[str] = None) -> BrainAnytimeInference:
    """
    创建推理引擎实例
    
    Args:
        checkpoints_dir: 检查点目录，默认从 HuggingFace Hub 下载
    
    Returns:
        BrainAnytimeInference 实例
    """
    return BrainAnytimeInference(checkpoints_dir=checkpoints_dir)


def format_result(result: Dict) -> str:
    """格式化推理结果为可读字符串"""
    if result is None:
        return "Inference failed"
    
    task = result['task']
    task_type = result['task_type']
    
    lines = [
        f"Task: {task}",
        f"Type: {task_type}",
    ]
    
    if task_type == 'classification':
        lines.extend([
            f"Prediction: {result['prediction']}",
            f"Probability: {result['probability']:.4f}",
            f"Confidence: {result['confidence']:.2%}",
        ])
    else:
        lines.extend([
            f"Prediction: {result['prediction']:.2f} {result['unit']}",
            f"Reference Range: {result['reference_range']}",
        ])
    
    lines.append(f"Modalities: {', '.join(result['input']['modalities'])}")
    
    return '\n'.join(lines)


# =============================================================================
# 测试
# =============================================================================

if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='Test inference engine')
    parser.add_argument('--checkpoints_dir', type=str, 
                        default='/home/23037125r/code/random/multimae_freeze_then_finetune/',
                        help='Checkpoints directory')
    parser.add_argument('--sample_dir', type=str,
                        default='/home/23037125r/code/Downstream_tasks/hf_space/demo_samples/CN_vs_AD/TMFP/sample_005',
                        help='Sample directory for testing')
    parser.add_argument('--task', type=str, default='CN_vs_AD',
                        choices=list(TASKS.keys()),
                        help='Task to test')
    parser.add_argument('--combo', type=str, default='TMFP',
                        choices=['T', 'TF', 'TMF', 'TFP', 'TMFP'],
                        help='Modality combination')
    
    args = parser.parse_args()
    
    # 创建引擎
    engine = create_inference_engine(args.checkpoints_dir)
    
    # 执行推理
    print(f"\nTesting inference:")
    print(f"  Task: {args.task}")
    print(f"  Combo: {args.combo}")
    print(f"  Sample: {args.sample_dir}")
    
    result = engine.predict_from_sample(args.sample_dir, args.task, args.combo)
    
    if result:
        print("\nResult:")
        print(format_result(result))
    else:
        print("\nInference failed!")