#coding:utf-8
'''
write by ygq
create on 2025-09-04

OASIS（Open Access Series of Imaging Studies） 是一个旨在向科研界免费提供脑部MRI数据的项目。本横断面（Cross-Sectional）数据集是其第一个版本，发布于2007年。
OASIS-1 是横断面的，意味着它无法捕捉个体随时间的动态变化。对于研究疾病进展，后续的 OASIS-2 和 OASIS-3（纵向数据集）是更好的选择。

OASIS-2，全称为 Longitudinal Multimodal Neuroimaging: Principal 150 Subjects，是 OASIS 项目发布的第二个核心数据集。顾名思义，它的核心特点是 纵向（Longitudinal）。

核心目标：
    研究正常衰老和阿尔茨海默病（AD）中的大脑结构随时间变化的模式。
研究设计：
    纵向研究。同一批受试者被多次扫描和评估，持续数年。
样本量：
    150 名年龄在 60 到 96 岁之间的受试者。
人群组成：
    所有 150 名受试者在首次扫描时都被诊断为认知正常（CDR = 0）。
    在研究期间，部分受试者仍然保持认知正常，而另一部分则发展为痴呆（被临床诊断为可能患有阿尔茨海默病）。
数据采集：
    每名受试者进行了 至少 2 次 的访视会话（session），最多达到了 5 次。
    每次访视之间的平均间隔时间约为 2.2 年，整个研究跨度最长超过 7 年。
    每次访视都包括：3-4 次 T1 加权 MRI 扫描（在单次会话中完成，用于平均以提高信噪比）和详细的临床神经心理评估。
数据内容：
    与 OASIS-1 类似，包括原始 DICOM 图像、预处理后的 Analyze 格式图像，以及全面的临床认知评估数据。

    
关键区别的详细解释
横断面 vs. 纵向 (Cross-Sectional vs. Longitudinal)：
    OASIS-1 像是在给一个城市的所有人在同一天拍一张照片。你可以比较年轻人和老年人、健康人和病人的区别，但看不到任何一个人是如何变老或生病的。
    OASIS-2 像是挑选了150位健康的老年人，然后每年都给他们拍一张照片，持续好几年。这样你就能亲眼看到有些人如何慢慢地出现变化，最终生病。这对于理解疾病的过程至关重要。
受试者群体的区别：
    OASIS-1 包含了已经确诊的AD患者，非常适合训练一个模型来学习“AD大脑看起来是什么样”。
    OASIS-2 的受试者起点都是健康的，这使得它成为研究疾病前驱期（即临床症状出现之前）的宝贵资源。你可以分析那些最终患病的人，在多年前其大脑是否就已经存在细微的、可检测的差异。
数据分析方法的差异：
    分析 OASIS-1 通常使用跨主体（cross-sectional）比较，例如比较AD组和正常对照组的平均海马体积。
    分析 OASIS-2 则侧重于个体内部随时间的变化（within-subject change）。例如，为每个受试者计算其年化脑萎缩率，然后比较保持正常组和转化组之间的萎缩速率差异。这需要更复杂的纵向统计模型。


    1. 人口统计学信息
        性别（M/F）
        用手习惯（Hand）（均为右利手）
        年龄（Age）
        教育程度（Educ）（1-5级）
        社会经济地位（SES）

    2. 临床评估
        MMSE（简易精神状态检查）
        CDR（临床痴呆评级：0=正常，0.5=非常轻微，1=轻度，2=中度）

    3. 衍生解剖指标
        eTIV：估计颅内容积
        ASF：图谱缩放因子
        nWBV：标准化全脑体积
'''
import os
import glob,re
import pandas as pd
import SimpleITK as sitk
import argparse
import json
from tqdm import tqdm
from util import meta_data
import util
import numpy as np
# from bert_helper import *

import shutil
##dataset_meta
import warnings
warnings.filterwarnings("ignore")
meta_id_name='MRI ID'
##性别（M/F）,用手习惯（Hand）（均为右利手）,年龄（Age）,教育程度（Educ）（1-5级）,社会经济地位（SES）,MMSE（简易精神状态检查）,CDR（临床痴呆评级：0=正常，0.5=非常轻微，1=轻度，2=中度）,eTIV：估计颅内容积,ASF：图谱缩放因子,nWBV：标准化全脑体积
#META_COLUMN=['ID', 'M/F', 'Hand', 'Age', 'Educ', 'SES', 'MMSE', 'CDR', 'eTIV','nWBV', 'ASF', 'Delay']
META_COLUMN=['Subject ID', 'Group', 'Visit', 'MR Delay', 'M/F', 'Hand','Age', 'EDUC', 'SES', 'MMSE', 'CDR', 'eTIV', 'nWBV', 'ASF']

TASK_VALUE="segmentation"
CLAMP_RANGE_CT = [-300,300]
CLAMP_RANGE_MRI = None # MRI images threshold placeholder TBC...
TARGET_VOXEL_SPACING=None


# def find_metadata_files(path):
#     # for Cancer Image Archive (TCIA) dataset
#     search_pattern = os.path.join(path, '**', 'metadata.csv')
#     return glob.glob(search_pattern, recursive=True)

def find_metadata_files(path):
    # for Cancer Image Archive (TCIA) dataset
    search_pattern = os.path.join(path, '*.csv')
    return glob.glob(search_pattern, recursive=True)
##added by yanguoqing on 20250527
def find_image_dirs(path):
    return os.listdir(path)

##modify by yanguoqing on 20250527
def load_dicom_images(folder_path):
    reader = sitk.ImageSeriesReader()
    dicom_names = reader.GetGDCMSeriesFileNames(folder_path)
    reader.SetFileNames(dicom_names)
    image = reader.Execute()
    return dicom_names,image

##added by yanguoqing on 20250527
def load_dicom_tag(imgs):
    reader = sitk.ImageFileReader()
    # dicom_names = reader.GetGDCMSeriesFileNames(folder_path)
    reader.SetFileName(imgs)
    reader.ReadImageInformation()  # 仅读取元信息，不加载像素数据
    # metadata_keys = reader.GetMetaDataKeys()
    tag=reader.Execute()
    return tag

def load_nrrd(fp):
    return sitk.ReadImage(fp)

##modify by yanguoqing on 20250904
def load_raw_images(series_files):
    '''
    每个病例包含3到4种RAW的单次平扫MR
    将多个分开的模态合并，构建第四个维度的数组，分别按照MPR-1,MPR-2...顺序存放
    '''
    reader = sitk.ImageSeriesReader()
    reader.SetFileNames(series_files)
    image = reader.Execute()
    return image

def save_nifti(image, output_path, folder_path):
    # Set metadata in the NIfTI file's header
    output_dirpath = os.path.dirname(output_path)
    if not os.path.exists(output_dirpath):
        print(f"Creating directory {output_dirpath}")
        os.makedirs(output_dirpath)
    # Set metadata in the NIfTI file's header
    image.SetMetaData("FolderPath", folder_path)
    sitk.WriteImage(image, output_path)

##modify by yanguoqing on 20250527
def convert_windows_to_linux_path(windows_path):
    # Replace backslashes with forward slashes and remove the drive letter
    # Some meta files have windows paths, but the data is stored on a linux server
    linux_path = windows_path.replace('\\', '/')
    if ':' in linux_path:
        linux_path = linux_path.split(':', 1)[1]
    return linux_path

def main(target_path, output_dir):
    pid_dirs=find_image_dirs(target_path)
    failed_files = []
    if not os.path.isdir(output_dir):
        os.makedirs(output_dir)
    json_output_path = os.path.join(output_dir, 'nifti_mappings.json')
    failed_files_path = os.path.join(output_dir, 'failed_files.json')
    meta = meta_data()
    
    # Initialize the JSON file
    if not os.path.exists(json_output_path):
        with open(json_output_path, 'w') as json_file:
            json.dump({}, json_file)
    ##方便处理解析信息，转成csv文件
    meta_file=os.path.join(os.path.dirname(os.path.realpath(__file__)),'oasis2_longitudinal_demographics.csv')
    meta_file_ori=os.path.join(target_path,'oasis_longitudinal_demographics-8d83e569fa2e2d30 (1).xlsx')
    if os.path.isfile(meta_file):
        mf_flag=True
        df_meta=pd.read_csv(meta_file,sep=',')
    else:
        mf_flag=False


    if pid_dirs:
        for pid_dir in tqdm(pid_dirs, desc="Processing pid dirs"):
            if not os.path.isdir(os.path.join(target_path,pid_dir)):
                continue
            
            ##遍历所有目录下的病例数据
            image_dirs=find_image_dirs(os.path.join(target_path,pid_dir))
            
            for data_dir in tqdm(image_dirs, desc="Processing images files"):
                ##data_dir即id
                full_path=os.path.join(target_path,pid_dir,data_dir)
                
                modality="MRI"
                study='OASIS_2'##Dataset_name
                CIA_other_info = {'metadata_file':''}
                CIA_other_info['split'] = "train"
                CIA_other_info['metadata_file']=meta_file_ori
                data_info_row=df_meta[df_meta[meta_id_name]==data_dir]
                
                if data_info_row.shape[0]>0:
                    data_info_row=data_info_row.reset_index()
                    #print(data_info_row[meta_id_name])
                    for keyname in META_COLUMN[:]:
                        CIA_other_info[keyname]=str(data_info_row[keyname][0])
                    
                    CIA_other_info['Image_id']=data_dir
                  

                else:
                    meta_image_id=data_dir
                    for keyname in META_COLUMN[1:]:
                        CIA_other_info[keyname]=''
                
                

                try:
                    ##读取原始的RAW目录下多次单扫img
                    #\RAW\OAS1_0001_MR1_mpr-1_anon.img
                    series_files=glob.glob("%s/RAW/mpr-*.img"%(full_path))
                    series_files.sort()
                    
                    if len(series_files)>0:
                        ##存在有效的MRI影像数据进行后续处理
                        sitk_img_original=load_raw_images(series_files)
                        submodality=[re.search(r"mpr-\d{1}",os.path.basename(fp)).group(0) for fp in series_files]
                        sub_modality_dict={}
                        for idx,value in enumerate(submodality):
                            sub_modality_dict[idx]=value

                        meta.add_keyvalue('Sub_modality',sub_modality_dict)
                        
                    else:
                        print("病例数据%s为空"%data_dir)
                        continue
                    
                    
                    original_spacing = list(sitk_img_original.GetSpacing())
                    original_size = list(sitk_img_original.GetSize())

                    

                    meta.add_keyvalue('Spacing_mm',min(original_spacing))
                    meta.add_keyvalue('OriImg_path',",".join(series_files))
                    meta.add_keyvalue('Size',original_size)  # 这里用处理后的size -- YH Jachin
                    meta.add_keyvalue('Modality',modality)
                    meta.add_keyvalue('Dataset_name',study)
                    meta.add_keyvalue('ROI','head')

                    

                    output_image_file = os.path.join(output_dir,data_dir, f"{data_dir}.nii.gz")
                    # output_path=convert_windows_to_linux_path(output_path)
                    ##
                    save_nifti(sitk_img_original, output_image_file, full_path)
                    print(f"Saved NIfTI file to {output_image_file}")
                    ##Label processing

                    

                except Exception as e:
                    print(e)
                    failed_files.append(data_dir)
                    print(f"Failed to load OASIS images from {data_dir}")
                    continue

                
                
                meta.add_extra_keyvalue('Metadata',CIA_other_info)


                # Write the mapping to the JSON file on the fly
                with open(json_output_path, 'r+') as json_file:
                    existing_mappings = json.load(json_file)
                    existing_mappings[output_image_file] = meta.get_meta_data()
                    json_file.seek(0)
                    # print(existing_mappings)
                    json.dump(existing_mappings, json_file, indent=4)
                    json_file.truncate()
    # else:
    #     print("No metadata.csv files found.")
    
    with open(failed_files_path, "w") as json_file:
        json.dump(failed_files, json_file)
        
    print(f"The list has been written to {failed_files_path}")
    print(f"Saved NIfTI mappings to {json_output_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process DICOM files and save as NIfTI.")
    parser.add_argument("--target_path", type=str, help="Path to the target directory containing metadata files.", default="/home/data/Github/data/data_gen_def/DATASETS/OASIS/OASIS_2/OAS2_RAW//")
    parser.add_argument("--output_dir", type=str, help="Directory to save the NIfTI files.", default="/home/data/Github/data/data_gen_def/DATASETS_processed/OASIS/OASIS_2/RAW")
    args = parser.parse_args()
    print(args.target_path, args.output_dir)
    main(args.target_path, args.output_dir)