DATASET_DIR = "/home/aidrive/zhengxj/projects/eeg_enhanced_driving/brain_like_driving_model/data/real_car/processed_8s"

WHOLE_VERSION_IDENTIFIER = "v1.5_freeze_resnet_linear_linear_single_gpu+novice"
VERSION_IDENTIFIER = WHOLE_VERSION_IDENTIFIER.split("_")[0]
LARGE_VERSION_IDENTIFIER, SMALL_VERSION_IDENTIFIER = VERSION_IDENTIFIER.split(".")

# sample rate of different data
SAMPLE_RATE_DICT = {
    "video": 30,
    "eeg": 1000,
    "canbus": 1000,
    "origin_human": 1000, # the origin sample rate of human data
}
assert SAMPLE_RATE_DICT["eeg"] == SAMPLE_RATE_DICT["canbus"], "the sample rate of eeg and canbus are not equal"

# if we change the sample rate in the next code block, we should interpolate data of the origin dataset
INTERPOLATION_NEEDED = False

# get n_sample_per_frame
# let their sample rate could be divided
# so that each frame can stand for a integer number of sample
N_SAMPLE_PER_FRAME = SAMPLE_RATE_DICT["eeg"] // SAMPLE_RATE_DICT["video"]
if SAMPLE_RATE_DICT["video"] * N_SAMPLE_PER_FRAME != SAMPLE_RATE_DICT["eeg"]:
    INTERPOLATION_NEEDED = True
    N_SAMPLE_PER_FRAME += 1
    SAMPLE_RATE_DICT["eeg"] = SAMPLE_RATE_DICT["video"] * N_SAMPLE_PER_FRAME
    SAMPLE_RATE_DICT["canbus"] = SAMPLE_RATE_DICT["video"] * N_SAMPLE_PER_FRAME

# driving_conditions
# Baseline removed
DRIVING_CONDITIONS = ["A1", "A2", "D1", "S1", "A3", "A4", "L1", "L2", "R1", "R2", "R3", "R4", "L3"]
# subject_dict
SUBJECT_DICT = {
    # 'expert': ["P06", "P09", "P11", "P13", "P14", "P15", "P16", "P17", "P18", "P19"],
    'novice': ["P05", "P07", "P08", "P20", "P21", "P22", "P23", "P25", "P26", "P27"]
}

# YOLO feature
YOLO_FEATURE_LABELS = ["person_count", "bicycle_count", "car_count", "motorcycle_count", "truck_count", "traffic light_count", "bus_count", "stop sign_count"]

# EEG channel
EEG_CHANNEL_LABELS = [
    # Visual Area
    "PO7_voltage", "PO5_voltage", "PO3_voltage", "POZ_voltage", "PO4_voltage", "PO6_voltage", "PO8_voltage", "O1_voltage", "OZ_voltage", "O2_voltage",
    # Decision Area
    "FP1_voltage", "FPZ_voltage", "FP2_voltage", "AF3_voltage", "AF4_voltage",
    # Motor Seneory Area
    "FC1_voltage", "FCZ_voltage", "FC2_voltage", "C1_voltage", "CZ_voltage", "C2_voltage", "CP1_voltage", "CPZ_voltage", "CP2_voltage"
]

# Action channel
CANBUS_CHANNEL_LABELS = ["throttle_percentage", "brake_percentage", "steering_percentage"]

# Visual Network
# for different visual network, we need to use different preprocessed video
VISUAL_NETWORK = "resnet"

# Log dataloading process
LOG_DATA_LOADING_PROCESS = False

DATASET_CONFIG_DICT = {
    "dataset_dir": DATASET_DIR,
    "interpolation_needed": INTERPOLATION_NEEDED,
    "sample_rate_dict": SAMPLE_RATE_DICT,
    "n_sample_per_frame": N_SAMPLE_PER_FRAME,
    "driving_conditions": DRIVING_CONDITIONS,
    "subject_dict": SUBJECT_DICT,
    "visual_network": VISUAL_NETWORK,
    "eeg_channel_labels": EEG_CHANNEL_LABELS,
    "canbus_channel_labels": CANBUS_CHANNEL_LABELS,
    "log_data_loading_process": LOG_DATA_LOADING_PROCESS
}


VISUAL_NETWORK_OUTCHANNEL_CONFIG_DICT = {
    # the origin dimension of feature output of ResNet50
    "v1": 2048,
    "v2": {
        "0": 2048, # origin (lstm1024+lstm512)
        "1": 1024, # origin/2 (lstm512+lstm256)
        "2": 512, # origin/4 (lstm256+lstm128)
    }
}

VISUAL_NETWORK_OUTCHANNEL = 0
if LARGE_VERSION_IDENTIFIER == "v1":
    VISUAL_NETWORK_OUTCHANNEL = VISUAL_NETWORK_OUTCHANNEL_CONFIG_DICT[LARGE_VERSION_IDENTIFIER]
elif LARGE_VERSION_IDENTIFIER == "v2":
    VISUAL_NETWORK_OUTCHANNEL = VISUAL_NETWORK_OUTCHANNEL_CONFIG_DICT[LARGE_VERSION_IDENTIFIER][SMALL_VERSION_IDENTIFIER]
else:
    raise Exception("not implemented yet!")

# Neural network config
VISUAL_NETWORK_CONFIG_DICT = {
    "freeze_parameters": True,
    "out_channel": VISUAL_NETWORK_OUTCHANNEL
} 

DECISION_MAKING_NETWORK_CONFIG_DICT = {
    # the dimension of feature output of ResNet50
    "in_channel": VISUAL_NETWORK_CONFIG_DICT["out_channel"],
    "d_hidden": VISUAL_NETWORK_CONFIG_DICT["out_channel"] / 2,
    "out_patch_sample": N_SAMPLE_PER_FRAME,
    "out_channel": len(EEG_CHANNEL_LABELS)
}

ACTION_PLANNING_NETWORK_CONFIG_DICT = {
    "in_channel": {
        "v1": len(EEG_CHANNEL_LABELS), # DECISION_MAKING_NETWORK_CONFIG_DICT["out_channel"]
        "v2": DECISION_MAKING_NETWORK_CONFIG_DICT["d_hidden"]
    },
    # As the network goes deeper, the feature of network extract is more compressed, so the number of neurons decreses here
    "d_hidden":{
      "v1": 512,
      "v2": DECISION_MAKING_NETWORK_CONFIG_DICT["d_hidden"] / 2,
    },
    "out_patch_sample": N_SAMPLE_PER_FRAME,
    "out_channel": len(CANBUS_CHANNEL_LABELS)
}

BRAIN_LIKE_MODEL_CONFIG_DICT = {
    "visual_network": VISUAL_NETWORK_CONFIG_DICT,
    "decision_making_network": DECISION_MAKING_NETWORK_CONFIG_DICT,
    "action_planning_network": ACTION_PLANNING_NETWORK_CONFIG_DICT,
    "eeg_loss_weight": 1,
    "canbus_loss_weight": 1
}

OPTIMIZER = "Adam"
LEARNING_RATE = 0.001

EPOCH = 15
BATCH_SIZE = 8
