/**
 * @file  relabel_fsseg.cc
 * @brief Changes the labels of the subcortical segmentation of FreeSurfer
 *        such that it contains only labels for WM, GM, and CSF.
 *
 * The relabeling is based on Table 41 of the master thesis of Kowkabzadeh
 * with some additional mappings not mentioned in this table.
 *
 * K. Kowkabzadeh, Evaluations of Tissue Segmentation of brain MR Images,
 * Master of Science Thesis in Biomedical Engineering,
 * Department of Signals and Systems, Division of Biomedical Engineering,
 * Chalmers University of Technology, Göteborg, Sweden, 2010
 * Report No. EX048/2010
 *
 * Copyright (c) 2011 University of Pennsylvania. All rights reserved.
 * See https://www.rad.upenn.edu/sbia/software/license.html or COPYING file.
 *
 * Contact: SBIA Group <sbia-software at uphs.upenn.edu>
 */

#include <iostream>
#include <map>
#include <nifti1_io.h>
#include <getopt.h>


// acceptable in .cc file
using namespace std;


// ===========================================================================
// help / version
// ===========================================================================

// ---------------------------------------------------------------------------
void print_usage(const char* exec_name) {
    cout << "Usage:" << endl;
    cout << "  " << exec_name << " [options] <aseg.nii>" << endl;
    cout << endl;
    cout << "  Attention: This program overwrites the input image!" << endl;
}

// ---------------------------------------------------------------------------
void print_options() {
    cout << "Required Options:" << endl;
    cout << "  <aseg.nii>             The subcortical segmentation volume aseg.mgz of FreeSurfer" << endl;
    cout << "                         in the NIfTI-1 image format." << endl;
    cout << endl;
    cout << "Options:" << endl;
    cout << "  --csf-label <0..255>   Output label of CSF (default: 10)." << endl;
    cout << "  --gm-label <0..255>    Output label of GM (default: 150)." << endl;
    cout << "  --wm-label <0..255>    Output label of WM (default: 250)." << endl;
    cout << "  --verbose, -v          Increases verbosity of output messages." << endl;
    cout << "                         Can be given multiple times." << endl;
    cout << "  --help, -h             Print help and exit." << endl;
    cout << "  --helpshort            Print short help and exit." << endl;
    cout << "  --version              Print version and exit." << endl;
}

// ---------------------------------------------------------------------------
void print_contact() {
    cout << "Contact:" << endl;
    cout << "  SBIA Group <sbia-software at uphs.upenn.edu>" << endl;
}

// ---------------------------------------------------------------------------
void print_help(const char* exec_name) {
    cout << endl;
    print_usage(exec_name);
    cout << endl;
    cout << "Description:" << endl;
    cout << "  This program reads the aseg.nii image which was created from the FreeSurfer" << endl;
    cout << "  output file aseg.mgz using the command mri_convert and changes the labels" << endl;
    cout << "  such that only CSF, GM, and WM are labeled." << endl;
    cout << endl;
    print_options();
    cout << endl;
    cout << "Example:" << endl;
    cout << "  mri_convert aseg.mgz fsseg.nii; " << exec_name << " fsseg.nii" << endl;
    cout << "    Converts the aseg.mgz volume to NIfTI-1 and then changes the labels to either" << endl;
    cout << "    background, CSF, GM, or WM." << endl;
    cout << endl;
    print_contact();
    cout << endl;
}

// ---------------------------------------------------------------------------
void print_helpshort(const char* exec_name) {
    cout << endl;
    print_usage(exec_name);
    cout << endl;
    print_options();
    cout << endl;
    print_contact();
    cout << endl;
}

// ---------------------------------------------------------------------------
void print_version() {
    cout << "relabel_fsseg (MICO) version 1.0.0" << endl;
    cout << "Copyright (c) 2011 University of Pennsylvania. All rights reserved." << endl;
    cout << "See https://www.rad.upenn.edu/sbia/software/license.html or COPYING file." << endl;
}

// ===========================================================================
// freesurfer labels
// ===========================================================================

// ---------------------------------------------------------------------------
// background
unsigned char bg_labels[] = {
    0,  // unknown
    16, // brain stem
    85, // optic chiasm
};

// ---------------------------------------------------------------------------
// white matter
unsigned char wm_labels[] = {
    2,   // left cerebral white matter
    7,   // left cerebellum white matter
    11,  // left cerebral white matter
    30,  // left vessel
    41,  // right cerebral white matter
    46,  // right cerebellum white matter
    62,  // right vessel
    77,  // white matter hypointensities
    78,  // left white matter hypointensities
    79,  // right white matter hypointensities
    251, // corpus callosum posterior
    252, // corpus callosum mid posterior
    253, // corpus callosum central
    254, // corpus callosum mid anterior
    255  // corpus callosum anterior
};

// ---------------------------------------------------------------------------
// gray matter
unsigned char gm_labels[] = {
    3,  // left cerebral cortex
    8,  // left cerebellum cortex
    9,  // left thalamus
    10, // left thalamus proper
    12, // left putamen
    13, // left pallidum
    17, // left hippocampus
    18, // left amygdala
    26, // left accumbens area
    28, // left ventral DC
    42, // right cerebral cortex
    47, // right cerebellum cortex
    48, // right thalamus
    49, // right thalamus proper
    50, // right caudate
    51, // right putamen
    52, // right pallidum
    53, // right hippocampus
    54, // right amygdala
    58, // right accumbens area
    60  // right ventral DC
};

// ---------------------------------------------------------------------------
// cerebrospinal fluid
unsigned char csf_labels[] = {
    1,  // left cerebral exterior
    4,  // left lateral ventricle
    5,  // left inferior lateral ventricle
    14, // third ventricle
    15, // fourth ventricle
    24, // CSF
    31, // left choroid plexus
    40, // right cerebral exterior
    43, // right lateral ventricle
    44, // right inferior lateral ventricle
    63, // right choroid plexus
    72, // fifth ventricle
    75, // left lateral ventricle
    76, // right lateral ventricle
    80, // non-white matter hypointensities
    81, // left non-white matter hypointensities
    82, // right non-white matter hypointensities
};

// ===========================================================================
// main
// ===========================================================================

// ---------------------------------------------------------------------------
int main(int argc, char *argv[])
{
    bool ok = true;

    const char* const exec_name = argv[0]; // used for help output

    const char* aseg_file = NULL; // name of input/output image file
    unsigned char bg = 0;         // output label for background
    unsigned char csf = 10;       // output label for CSF
    unsigned char gm  = 150;      // output label for GM
    unsigned char wm  = 250;      // output label for WM
    int verbose = 0;              // verbosity of output messages

    // parse arguments
    static struct option long_options [] =
    {
        {"csf-label", required_argument, NULL, 'c'},
        {"gm-label",  required_argument, NULL, 'g'},
        {"wm-label",  required_argument, NULL, 'w'},
        {"verbose",   no_argument,       NULL, 'v'},
        {"version",   no_argument,       NULL, 'V'},
        {"help",      no_argument,       NULL, 'h'},
        {"helpshort", no_argument,       NULL, 'u'}
    };

    int oc = -1;
    while ((oc = getopt_long(argc, argv, "vh", long_options, NULL)) != -1) {
        switch (oc) {
            case 'c':
                csf = static_cast<unsigned char>(atoi(optarg));
                break;

            case 'g':
                gm = static_cast<unsigned char>(atoi(optarg));
                break;

            case 'w':
                wm = static_cast<unsigned char>(atoi(optarg));
                break;

            case 'v':
                verbose++;
                break;

            case 'V':
                print_version();
                exit(EXIT_SUCCESS);

            case 'h':
                print_help(exec_name);
                exit(EXIT_SUCCESS);
                break;

            case 'u':
                print_helpshort(exec_name);
                exit(EXIT_SUCCESS);

            default:
                exit(EXIT_FAILURE);
        }
    }

    argc -= optind;
    argv += optind;

    if (argc != 1) {
        print_help(exec_name);
        exit(EXIT_FAILURE);
    }

    aseg_file = argv[0];

    // read segmentation
    nifti_image* aseg = nifti_image_read(aseg_file, 1);

    if (aseg->datatype != DT_UNSIGNED_CHAR) {
        cerr << "Label image must be of type DT_UNSIGNED_CHAR!" << endl;
        ok = false;
    }

    // create map of input to output labels for faster lookup
    typedef map<unsigned char, unsigned char> LabelsMap;
    LabelsMap labels_map;

    if (ok) {
        for (unsigned int i = 0; i < sizeof(bg_labels); i++) {
            labels_map[bg_labels[i]] = bg;
        }
        for (unsigned int i = 0; i < sizeof(wm_labels); i++) {
            labels_map[wm_labels[i]] = wm;
        }
        for (unsigned int i = 0; i < sizeof(gm_labels); i++) {
            labels_map[gm_labels[i]] = gm;
        }
        for (unsigned int i = 0; i < sizeof(csf_labels); i++) {
            labels_map[csf_labels[i]] = csf;
        }
    }

    // relabel segmentation
    if (ok) {
        map<unsigned char, int> skipped_labels;
        int number_of_voxels = aseg->nx * aseg->ny * aseg->nz;
        unsigned char* img = reinterpret_cast<unsigned char*>(aseg->data);
        for (int i = 0; i < number_of_voxels; i++) {
            LabelsMap::iterator it = labels_map.find(img[i]);
            if (it == labels_map.end()) {
                ++skipped_labels[img[i]];
            } else {
                img[i] = it->second;
            }
        }
        if (verbose > 0) {
            for (map<unsigned char, int>::const_iterator it = skipped_labels.begin(); it != skipped_labels.end(); ++it) {
                cout << "Skipped unknown label " << static_cast<unsigned int>(it->first);
                cout << " in total " << it->second << " times." << endl;
            }
        }
    }

    // write segmentation
    if (ok) {
        nifti_image_write(aseg);
    }

    exit(ok ? EXIT_SUCCESS : EXIT_FAILURE);
}