import argparse
import pickle
import pandas as pd
from pathlib import Path
import os
import json 

BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'datasets')

def get_shhs1_demographics(data_path, output_path):
    """Extract demographics from SHHS1 dataset."""
    # Assuming SHHS1 has a CSV file with demographics

    demo_file = Path(data_path) / 'shhs1-dataset-0.20.0.csv'
    output_file = Path(output_path) / 'shhs1_demographics.pkl'
    df = pd.read_csv(demo_file)
    
    demographics = {}
    for _, row in df.iterrows():
        patient_id = row['nsrrid']  # Adjust column name as needed
        demographics[patient_id] = {
            'Age': row['age_s1'],  # Adjust column names as needed
            'Sex': row['gender'], # 1 or 2
            'BMI': row['bmi_s1'],
            'AHI': row['ahi_a0h3']
        }
    
    
    with open(output_file, 'wb') as f:
        pickle.dump(demographics, f)
    print(f"SHHS1 demographics saved to {output_file}")

def get_shhs2_demographics(data_path, output_path):
    """Extract demographics from SHHS2 dataset."""
    # Assuming SHHS2 has a CSV file with demographics
    
    demo_file = Path(data_path) / 'shhs2-dataset-0.20.0.csv'
    output_file = Path(output_path) / 'shhs2_demographics.pkl'
    df = pd.read_csv(demo_file, encoding='cp1252')

    demographics = {}
    for _, row in df.iterrows():
        patient_id = row['nsrrid']  # Adjust column name as needed
        demographics[patient_id] = {
            'Age': row['age_s2'],  # Adjust column names as needed
            'Sex': row['gender'], # 1 or 2
            'BMI': row['bmi_s2'],
            'AHI': row['ahi_a0h3']
        }
    
    
    with open(output_file, 'wb') as f:
        pickle.dump(demographics, f)
    print(f"SHHS2 demographics saved to {output_file}")

def get_cfs_demographics(data_path, output_path):
    """Extract demographics from CFS dataset."""
    # Assuming CFS has a CSV file with demographics

    demo_file = Path(data_path) / 'cfs-visit5-harmonized-dataset-0.7.0.csv'
    output_file = Path(output_path) / 'cfs_demographics.pkl'
    df = pd.read_csv(demo_file)
    
    demographics = {}
    for _, row in df.iterrows():
        patient_id = row['nsrrid']  # Adjust column name as needed
        demographics[patient_id] = {
            'Age': row['nsrr_age'],  # Adjust column names as needed
            'Sex': 1 if row['nsrr_sex'] == 'male' else 2, # 1 or 2
            'BMI': row['nsrr_bmi']            
        }
    
    demo_file2 = Path(data_path) / 'cfs-visit5-dataset-0.7.0.csv'
    df2 = pd.read_csv(demo_file2)
    for _, row in df2.iterrows():
        patient_id = row['nsrrid']  # Adjust column name as needed
        demographics[patient_id]['AHI'] = row['ahi_a0h3']  # Update AHI if available in the second file

    print(demographics)


    with open(output_file, 'wb') as f:
        pickle.dump(demographics, f)
    print(f"CFS demographics saved to {output_file}")

def get_kiss_demographics(data_path, output_path):
    """Extract demographics from KISS dataset."""
    # KISS has a single JSON file that includes all subjects info
    demo_file = Path(data_path) / '1_PSG_210426.json'
    # df = pd.read_csv(demo_file)
    with open (demo_file, 'r') as f:
        ann = json.load(f)
    
    demographics = {}
    for patient_info in ann['Patient']:
        patient_id = patient_info['Report']['Patient Serial Number']
        demographics[patient_id] = {
            'Age': patient_info['Report']['Age'],
            'Sex': 1 if patient_info['Report']['Sex'] == 'Male' else 2,
            'BMI': patient_info['Report']['BMI'],
            'AHI': patient_info['Report']['AHI'],

        }
    # for _, row in df.iterrows():
    #     patient_id = row['patient_id']  # Adjust column name as needed
    #     demographics[patient_id] = {
    #         'Age': row['age'],  # Adjust column names as needed
    #         'Sex': row['sex'],
    #         'BMI': row['bmi']
    #     }
    
    output_file = Path(output_path) / 'kiss_demographics.pkl'
    with open(output_file, 'wb') as f:
        pickle.dump(demographics, f)
    print(f"KISS demographics saved to {output_file}")
    
def get_kvss_demographics(data_path, output_path):
    """Extract demographics from KVSS dataset."""
    
    data_dirs = os.listdir(data_path)
    demographics = {}
    for patient_snum in data_dirs:
        json_path = os.path.join(data_path, patient_snum, patient_snum + "_annotation.json")
        with open(json_path, "r") as f:
            annotations = json.load(f)    

        demographics[patient_snum] = {
            'Age': annotations['Report']['Age'],
            'Sex': 1 if annotations['Report']['Sex'] == 'Male' else 2,
            'BMI': annotations['Report']['BMI'],
            'AHI': annotations['Report']['AHI'],

        }    
    
    output_file = Path(output_path) / 'kvss_demographics.pkl'
    with open(output_file, 'wb') as f:
        pickle.dump(demographics, f)
    print(f"KVSS demographics saved to {output_file}")


def get_phy_demographics(data_path, output_path):
    """Extract demographics from Physionet dataset."""
    # Assuming Physionet has a CSV file with demographics
    demo_file = Path(data_path) / 'age-sex.csv'
    df = pd.read_csv(demo_file)
    
    demographics = {}
    for _, row in df.iterrows():
        patient_id = row['Record']  # Adjust column name as needed
        demographics[patient_id] = {
            'Age': row['Age'],  # Adjust column names as needed
            'Sex': 1 if row['Sex'] == 'M' else 2, # 1 or 2
        }
    
    output_file = Path(output_path) / 'phy_demographics.pkl'
    with open(output_file, 'wb') as f:
        pickle.dump(demographics, f)
    print(f"PHY demographics saved to {output_file}")

def main():
    parser = argparse.ArgumentParser(description='Collect demographics from sleep study datasets')
    parser.add_argument('--dataset', type=str, required=True, 
                        choices=['shhs1', 'shhs2', 'kiss', 'kvss', 'phy', 'cfs', 'all'],
                        help='Dataset to process (shhs1, shhs2, kiss, phy, or all)')
    # parser.add_argument('--data_path', type=str, required=True,
    #                     help='Path to the dataset directory')
    parser.add_argument('--output_path', type=str, default='../datasets/demographics/',
                        help='Path to save demographics pickle files')
    
    args = parser.parse_args()
    
    # Create output directory if it doesn't exist
    Path(args.output_path).mkdir(parents=True, exist_ok=True)
    
    if args.dataset == 'shhs1' or args.dataset == 'all':
        data_path = os.path.join(BASE_PATH, "#_SHHS/datasets")
        get_shhs1_demographics(data_path, args.output_path)
    
    if args.dataset == 'shhs2' or args.dataset == 'all':
        data_path = os.path.join(BASE_PATH, "#_SHHS/datasets")
        get_shhs2_demographics(data_path, args.output_path)
    
    if args.dataset == 'kiss' or args.dataset == 'all':
        data_path = os.path.join(BASE_PATH, "#_2020_Sleep_Quality")
        get_kiss_demographics(data_path, args.output_path)
        
    if args.dataset == 'kvss' or args.dataset == 'all':
        data_path = os.path.join(BASE_PATH, "#_2021_Sleep_Video")
        get_kvss_demographics(data_path, args.output_path)

    if args.dataset == 'phy' or args.dataset == 'all':
        data_path = os.path.join(BASE_PATH, "physionet.org/files/challenge-2018/1.0.0")
        get_phy_demographics(data_path, args.output_path)
    
    if args.dataset == 'cfs' or args.dataset == 'all':
        data_path = os.path.join(BASE_PATH, "CFS_raw/datasets")
        get_cfs_demographics(data_path, args.output_path)


if __name__ == '__main__':
    main()