#!python
"""
Makes temporary seed correlation maps using 3 rois of interest
then makes pretty picture of these maps.

Usage:
    cifti_vis_RSN cifti-snaps [options] <func.dtseries.nii> <subject>
    cifti_vis_RSN cifti-subject [options] <func.dtseries.nii> <subject>
    cifti_vis_RSN nifti-snaps [options] <func.nii.gz> <subject>
    cifti_vis_RSN nifti-subject [options] <func.nii.gz> <subject>
    cifti_vis_RSN index [options]

Arguments:
    <func.dtseries.nii>  A 3D cifti dtseries file to view RSN maps from
    <func.nii.gz>        A 4D nifti file to view RSN maps from
    <subject>            Subject ID for HCP surfaces

Options:
  --qcdir PATH             Full path to location of QC directory
  --ciftify-work-dir PATH  The directory for HCP subjects (overrides CIFTIFY_WORKDIR/ HCP_DATA
                           enivironment variables)
  --subjects-filter STR    A string that can be used to filter out subject directories for index
  --subjects-list TXT      Input a text file of subject ids to include in index
  --colour-palette STR     Specify the colour palette for the seed correlation maps
  --hcp-data-dir PATH      DEPRECATED, use --ciftify-work-dir instead
  -v,--verbose             Verbose logging
  --debug                  Debug logging in Erin's very verbose style
  --help                   Print help

DETAILS
This makes pretty pictures of your hcp views using connectome workbenches "show scene" commands
It pastes the pretty pictures together into some .html QC pages

By default, all subjects in the qc directory will be included in the index.
The --subject-list argument to specify a text file containing the subject list.

You can changed the color palette for all pics using the --colour-palette flag.
The default colour palette is videen_style. Some people like 'PSYCH-NO-NONE' better.
For more info on palettes see wb_command help.

Requires connectome workbench (i.e. wb_command and imagemagick)

Written by Erin W Dickie, Feb 2016
"""
import os
import subprocess
import glob
import sys
import tempfile
import shutil
import logging
import logging.config
from PIL import Image

import pandas as pd
import numpy as np
import nibabel as nib
from docopt import docopt

import ciftify
from ciftify.utils import run
from ciftify.qc_config import replace_all_references, replace_path_references

config_path = os.path.join(os.path.dirname(ciftify.config.find_ciftify_global()), 'bin', "logging.conf")
logging.config.fileConfig(config_path, disable_existing_loggers=False)
logger = logging.getLogger(os.path.basename(__file__))

arguments       = docopt(__doc__)
snaps_cifti      = arguments['cifti-subject'] or arguments['cifti-snaps']
snaps_nifti      = arguments['nifti-subject'] or arguments['nifti-snaps']
index_only      = arguments['index']
subject           = arguments['<subject>']
qcdir           = arguments['--qcdir']
func_nifti      = arguments['<func.nii.gz>']
func_cifti      = arguments['<func.dtseries.nii>']
if arguments['--ciftify-work-dir']:
    work_dir = arguments['--ciftify-work-dir']
else:
    work_dir = arguments['--hcp-data-dir']
seedcorr_palette = arguments['--colour-palette']
subjects_filter = arguments['--subjects-filter']
VERBOSE         = arguments['--verbose']
DEBUG           = arguments['--debug']

if DEBUG:
    logger.setLevel(logging.DEBUG)
    logging.getLogger('ciftify').setLevel(logging.DEBUG)

ciftify.utils.log_arguments(arguments)

## define the location of the template scene_file
template_dir = ciftify.config.find_scene_templates()
## define the settings for the qcpages
# Note: order for the list is the order in the scene file
# Name: the name that will apear as filenames and in title of qc page
# MakeIndex : if True, this will create a html page of everyparticipant for this views
# SplitHorizontal: Wether or not to split the image in half and display it as a line.
QCmodes = {
    'RSN': {
      'TemplateFile' : 'mapvis.scene',
      'scene_list': [
        {'Idx': 1, 'Name': 'VolAx', 'MakeIndex': True, 'SplitHorizontal': True, 'Keep': True, 'Order': 3},
        {'Idx': 2, 'Name': 'VolSag', 'MakeIndex': True, 'SplitHorizontal': True, 'Keep': True, 'Order': 4},
        {'Idx': 3, 'Name': 'VolCor', 'MakeIndex': True, 'SplitHorizontal': True, 'Keep': True, 'Order': 5},
        {'Idx': 4, 'Name': 'LM', 'MakeIndex': True, 'SplitHorizontal': True, 'Keep': True, 'Order': 1},
        {'Idx': 5, 'Name': 'APDV', 'MakeIndex': True, 'SplitHorizontal': True, 'Keep': True, 'Order': 2}],
    }
}

seed_vertices = [ {'Name':'L-precuneus', "vertex" : 13363, 'hemi': 'L', 'best_view': "LM"},
                  {'Name':'L-superior-parietal', "vertex" : 14785, 'hemi': 'L', 'best_view': "APDV"},
                  {'Name':'L-anterior-insula', "vertex" : 10939, 'hemi': 'L', 'best_view': "LM"},
                  {'Name':'R-mpfc', "vertex" : 28901, 'hemi': 'R', 'best_view': "LM"}]

## set radius of geodesic ROIS
SEED_RADIUS = 8

## to run qc
## copy qc_template.scene into your hcp directory as qc_views.scene
def png_SplitHorizontal(input_png,output_png):
    '''
    uses pillow to split the image top and bottom halves to one line
    '''
    with Image.open(input_png) as img:
        height = img.size[1]
        width = img.size[0]
        half_the_height = height // 2
        img_top = img.crop((0, 0, width, half_the_height))
        img_btm = img.crop((0, half_the_height, width, height))
        im2 = Image.new('RGBA', (int(width*2), half_the_height))
        im2.paste(img_top, (0, 0))
        im2.paste(img_btm, (width, 0))
        im2.save(output_png)

def write_htmlheader(htmlhandle):
    ''' writes some style elements into the html header '''
    htmlhandle.write('''
    <head>
    <link rel="stylesheet" href="https://maxcdn.bootstrapcdn.com/bootstrap/3.3.6/css/bootstrap.min.css" integrity="sha384-1q8mTJOASx8j1Au+a5WDVnPi2lkFfwwEAa8hDDdjZlpLegxhjVME1fgjWPGmkzs7" crossorigin="anonymous">
    <style>
    body{
      background-color: rgba(0, 0, 0, 0.1);
      margin-top: 100px;
    }
    h1 {
      text-indent: 15px
    }
    </style>
    </head>
    ''')

def html_add_image(htmlhandle, colwidth, href, src, label):
    ''' writes an image to the html page with hyperlink'''
    htmlhandle.write('''
    <div class="theme-table-image col-sm-{colwidth}">
      <a href="{href}"><img src="{src}" class="img-responsive img-rounded">{label}</a><br>
    </div>
    '''.format(colwidth = colwidth,
               href = href,
               src = src,
               label = label))

def write_htmlnavbar(htmlhandle, brandname, nav_list, activelink = None):
    '''
    uses information from the nav_dict to build a fixed navigation bar
    nav dict contains a list of item to go in the Navigation bar
    '''
    htmlhandle.write('''
  <nav class="navbar navbar-inverse navbar-fixed-top">
    <div class="container-fluid">
      <div class="navbar-header">
        <a class="navbar-brand">{}</a>
      </div>
    <ul class="nav navbar-nav navbar-right">
    '''.format(brandname))
    for nav_dict in nav_list:
        activeclass = ' class="active"' if nav_dict['href'] == activelink else ''
        htmlhandle.write('<li{}><a href="{}">{}</a></li>\n'.format(
            activeclass,
            nav_dict['href'],
            nav_dict['label']
            ))
    htmlhandle.write('   </ul>\n  </div>\n</nav>\n')

def write_index(qcdir, subjects, pic_name, colwidth, indexname, title):
    '''
    Writes html file with all subjects for one pic shown together
    '''
    # open the file
    htmlindex = open(os.path.join(qcdir,indexname),'w')

    # write the header and title line
    htmlindex.write('<!DOCTYPE html>\n<HTML><TITLE>{}</TITLE>\n'.format(title))
    write_htmlheader(htmlindex)
    htmlindex.write('<body>\n')
    ## writing the navigation bar
    nav_list = [{'href': "", 'label':'Seed:'}]
    for seed_dict in seed_vertices:
        nav_list.append({ 'href': "{}.html".format(seed_dict['Name']),
                           'label': seed_dict['Name'] })
    nav_list.append({'href': "index.html", 'label':'Index'})
    write_htmlnavbar(htmlindex,
                    "RSN Seed Correlations",
                    nav_list,
                    activelink  = indexname)

    ## add the main title
    htmlindex.write('<h1>{}</h1>\n'.format(title))

    # add the pic and a link to the subjects' page
    for subject in subjects:
        subjectpage = os.path.join(qcdir,subject,'qc_sub.html')
        htmlindex.write('<div class="container" style="width: 100%;">')
        pic = os.path.join(qcdir,subject,pic_name)
        picrelpath = os.path.relpath(pic,os.path.dirname(htmlindex.name))
        subrelpath = os.path.relpath(subjectpage,os.path.dirname(htmlindex.name))
        html_add_image(htmlindex, colwidth,
                       subrelpath, picrelpath,
                       "{}".format(subject))
        htmlindex.write('</div>\n</br>')
    ## close the html index stream
    htmlindex.write('</body>\n')
    htmlindex.close()



if arguments['cifti-snaps'] or arguments['nifti-snaps']:
    logger.warning("The 'snaps' argument has be deprecated. Please use 'subject' in the future.")

## pic a QC dict
qc_dict = QCmodes['RSN']
scene_list = qc_dict['scene_list']
template_scene = os.path.join(template_dir,qc_dict['TemplateFile'])

## open the template scene for reading
template_txt = open(template_scene, 'r').read()

## get the work_dir - cause we use it a lot
if not work_dir:
    work_dir = ciftify.config.find_work_dir()

#mkdir a tmpdir for the
tmpdirbase = tempfile.mkdtemp()

#tmpdirbase = os.path.join(work_dir,'temp')

## set the qcdir Structure
if not qcdir:
    qcdir = os.path.join(work_dir,'qc_{}'.format('RSN_seedcorr'))

if snaps_cifti:
    funcbase = os.path.basename(func_cifti).replace('.dtseries.nii','')
    func = func_cifti
    if not os.path.isfile(func):
        sys.exit("{} not found...exiting".format(func))

if snaps_nifti:
    funcbase = os.path.basename(func_nifti).replace('.nii','').replace('.gz','')
    if not os.path.isfile(func_nifti):
        sys.exit("{} not found...exiting".format(func_nifti))

## make pics and qcpage for each subject
if snaps_cifti or snaps_nifti:

    qc_subdir = os.path.join(qcdir,'{}_{}'.format(subject,funcbase))

    run(['mkdir','-p', qc_subdir])

    ## make a tempdir for this subject
    tmpdir = os.path.join(tmpdirbase,subject)
    run(['mkdir','-p', tmpdir])

    scene_dir=tmpdir

    ## do the conversion if this input is nifti...
    if snaps_nifti:
        func = os.path.join(tmpdir,'{}.dtseries.nii'.format(funcbase))
        func_vx_size = ciftify.niio.voxel_spacing(func_nifti)
        expected_resolution = (2.0, 2.0, 2.0)
        if func_vx_size != expected_resolution :
            run(['ciftify_vol_result', '--resample-nifti',
                    '--ciftify-work-dir', work_dir,
                    subject,
                    func_nifti,
                    func])
        else:
            run(['ciftify_vol_result',
                    '--ciftify-work-dir',work_dir,
                    subject,
                    func_nifti,
                    func])

    func_fnifti = os.path.join(tmpdir,'func.nii.gz')
    ## convert funcfile to fake nifti
    run(['wb_command','-cifti-convert','-to-nifti',func, func_fnifti])
    ## load functional file as nibabel object using ciftify-utilities
    f_data, affine, header, dims = ciftify.niio.load_nifti(func_fnifti)
    ## mean the fuctional cifti image to make cifti template for the output
    run(['wb_command', '-cifti-reduce', func, 'MIN', os.path.join(tmpdir, 'template.dscalar.nii')])

    ## write the header for the subjects qc page
    qcsubhtml = os.path.join(qc_subdir, 'qc_sub.html')
    qcsubhtml = open(qcsubhtml,'w')
    qcsubhtml.write('<!DOCTYPE html>\n<HTML><TITLE> {} RSN seed correlations</TITLE>\n'.format(subject))
    write_htmlheader(qcsubhtml)
    qcsubhtml.write('<body>\n')
    ## writing the navigation bar
    nav_list = [{'href': "", 'label':'Seed:'}]
    for seed_dict in seed_vertices:
        nav_list.append({ 'href': "../{}.html".format(seed_dict['Name']),
                           'label': seed_dict['Name'] })
    nav_list.append({'href': "../index.html", 'label':'Index'})
    write_htmlnavbar(qcsubhtml,
                    "RSN Seed Correlations",
                    nav_list)
    ## writing the title
    qcsubhtml.write('<h1>{} {} Seed Correlations</h1>\n'.format(subject, funcbase))

    for seednum in range(len(seed_vertices)):
        '''
        for each seed vertex make and roi and generate a seed map
        '''
        ## get info from the seed_dict
        seed_dict = seed_vertices[seednum]
        seed_txt = os.path.join(tmpdir,'seed{}vertex.txt'.format(seednum))
        subject_surf = os.path.join(work_dir, subject, 'MNINonLinear','fsaverage_LR32k',
            '{}.{}.midthickness.32k_fs_LR.surf.gii'.format(subject, seed_dict['hemi']))
        seed_roi_gii = os.path.join(tmpdir,'roi{}.{}.shape.gii'.format(seednum,seed_dict['hemi']))
        seed_roi_cifti = os.path.join(tmpdir,'roi{}.dscalar.nii'.format(seednum))
        seed_meants_txt = os.path.join(tmpdir,'roi{}.meants.txt'.format(seednum))
        seedcorr_dscalar = os.path.join(tmpdir,'seedcorr{}.dscalar.nii'.format(seednum))

        ## write the vertex value to a tmp txt file
        with open(seed_txt, 'w') as f:
            f.write(str(seed_dict['vertex']))
        ## feed that txt file to wb_command surface-geodesic-rois
        run(['wb_command', '-surface-geodesic-rois', subject_surf,
                str(SEED_RADIUS), seed_txt, seed_roi_gii])
        ## use wb_command -cifti-roi-average to get the mean timeseries
        if seed_dict['hemi'] == 'L':
            metric_structure = 'CORTEX_LEFT'
            roi_arg = '-left-roi'
        if seed_dict['hemi'] == 'R':
            metric_structure = 'CORTEX_RIGHT'
            roi_arg = '-right-roi'
        run(['wb_command', '-cifti-roi-average', func, seed_meants_txt,
            roi_arg, seed_roi_gii])
        run(['wb_command', '-cifti-create-dense-from-template',
                os.path.join(tmpdir, 'template.dscalar.nii'),
                seed_roi_cifti,
                '-metric', metric_structure, seed_roi_gii])

        ## change the default color palette so that it shows up grey in our picks
        run(['wb_command', '-cifti-palette', seed_roi_cifti,
             'MODE_AUTO_SCALE', seed_roi_cifti,
             '-palette-name', 'power_surf'])

        ## read the mean timeseries into numpy
        meants = np.loadtxt(seed_meants_txt)
        ## correlated the mean timeseries with the func data
        out = np.zeros([dims[0]*dims[1]*dims[2], 1])

        ## determine brainmask bits..
        std_array = np.std(f_data, axis=1)
        mask_indices = np.where(std_array > 0)[0]

        for i in mask_indices:
            out[i] = np.corrcoef(meants, f_data[i, :])[0][1]
        ## reshape data and write it out to a fake nifti file
        out = out.reshape([dims[0], dims[1], dims[2], 1])
        out = nib.nifti1.Nifti1Image(out, affine)
        out.to_filename(os.path.join(tmpdir,'seedcorr{}.nii.gz'.format(seednum)))
        ## convert back
        run(['wb_command','-cifti-convert','-from-nifti',
            os.path.join(tmpdir,'seedcorr{}.nii.gz'.format(seednum)),
            os.path.join(tmpdir, 'template.dscalar.nii'),
            seedcorr_dscalar])

        if seedcorr_palette:
            run(['wb_command', '-cifti-palette', seedcorr_dscalar,
                'MODE_AUTO_SCALE', seedcorr_dscalar,
                '-palette-name', seedcorr_palette])

        ## copy out template .scene file and find and replace the subid and hcppath
        hcp_data_realpath = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(subject_surf)))))
        scene_file = os.path.join(scene_dir,'seedcorr_{}_{}.scene'.format(subject, seednum))
        scenestream = open(scene_file,'w')
        txt = template_txt.replace('SURFS_SUBJECT', subject)
        txt = txt.replace('SURFS_MESHNAME', '.32k_fs_LR')
        txt = replace_path_references(txt, 'SURFSDIR',
                os.path.join(work_dir, subject, 'MNINonLinear','fsaverage_LR32k'),
                scene_file)
        txt = replace_all_references(txt, 'T1W',
                os.path.join(work_dir, subject, 'MNINonLinear', 'T1w.nii.gz'),
                scene_file)
        txt = replace_all_references(txt, 'TOPSCALAR', seed_roi_cifti, scene_file)
        txt = replace_all_references(txt, 'MIDSCALAR', seedcorr_dscalar, scene_file)
        scenestream.write(txt)
        scenestream.close()

        ## write the header for the subjects qc page
        qchtml = os.path.join(qc_subdir, 'qc_{}.html'.format(seed_dict['Name']))
        qchtml = open(qchtml,'w')
        qchtml.write('<!DOCTYPE html>\n<HTML><TITLE> qc {} {} seed correlation</TITLE>\n'.format(subject,seed_dict['Name']))
        write_htmlheader(qchtml)
        qchtml.write('<body>\n')
        write_htmlnavbar(qchtml,
                        "{} {} Seed {}".format(subject, funcbase, seed_dict['Name']),
                        [{ 'href': "qc_sub.html", 'label': "Return to Subject Page"}])

        qchtml.write('\n<h1>{} {} {} seed correlation</h1>\n'.format(subject,funcbase,seed_dict['Name']))

        ## for each scene, generate the pics and add to qc page
        for scene_dict in scene_list:
            if scene_dict['Keep']==True:
                pic = os.path.join(qc_subdir,'{}_{}.png'.format(seed_dict['Name'],scene_dict['Name']))
                picrelpath = os.path.relpath(pic,os.path.dirname(qchtml.name))
                html_add_image(qchtml, 12, picrelpath, picrelpath, "")
            else:
                pic = os.path.join(tmpdir,'{}_{}.png'.format(seed_dict['Name'],scene_dict['Name']))
            if scene_dict['SplitHorizontal']==True:
                tmppic = os.path.join(tmpdir, "scene{}.png".format(scene_dict['Idx']))
                run(['wb_command', '-logging', 'WARNING', '-show-scene',
                        scene_file, str(scene_dict['Idx']),
                        tmppic,'600', '400'])
                png_SplitHorizontal(tmppic,pic)
            else:
                run(['wb_command', '-logging', 'WARNING', '-show-scene',
                        scene_file, str(scene_dict['Idx']),
                        pic,'600', '400'])

        ## close the htmlpage stream
        qchtml.close()

        ## add the fav pic to the qc_sub page
        favpic = os.path.join(qc_subdir,'{}_{}.png'.format(seed_dict['Name'],seed_dict['best_view']))
        picrelpath = os.path.relpath(favpic,os.path.dirname(qcsubhtml.name))
        subrelpath = os.path.relpath(qchtml.name,os.path.dirname(qcsubhtml.name))
        qcsubhtml.write('<div>\n')
        html_add_image(qcsubhtml, 12, subrelpath, picrelpath,
                        "{} Seed".format(seed_dict['Name']))
        qcsubhtml.write('</div>\n')

    ## close the subject's htmlpage stream
    qcsubhtml.close()

## Start the index html file
if index_only:
    subjects = list(ciftify.utils.get_subj(qcdir))

    if subjects_filter:
        subjects = list(filter(lambda x: subjects_filter in x, subjects))
    ## remove common directories that are not subjects from subjects list

    ## write the major index
    indexhtml = os.path.join(qcdir, 'index.html')
    indexhtml = open(indexhtml,'w')
    indexhtml.write('<!DOCTYPE html>\n<HTML><TITLE>RSN Seed Correlations</TITLE>\n')
    write_htmlheader(indexhtml)
    indexhtml.write('<body>\n')

    ## writing the navigation bar
    nav_list = [{'href': "", 'label':'Seed:'}]
    for seed_dict in seed_vertices:
        nav_list.append({ 'href': "{}.html".format(seed_dict['Name']),
                           'label': seed_dict['Name'] })
    nav_list.append({'href': "index.html", 'label':'Index'})
    write_htmlnavbar(indexhtml,
                    "RSN Seed Correlations",
                    nav_list,
                    activelink  = 'index.html')

    ## writing the lists to the main index page
    indexhtml.write('<h1>Resting State Network Seed Correlations Index</h1>\n')
    indexhtml.write('<h2>All subjects together</h2>\n')
    indexhtml.write('<ul>\n  ')
    for seed_dict in seed_vertices:
        indexhtml.write('<li><a href="{}.html">{} Seed Correlations</a></li>\n'.format(seed_dict['Name'], seed_dict['Name']))
    indexhtml.write('</ul>\n')
    indexhtml.write('<h2>Subject Pages</h2>\n')
    indexhtml.write('<ul>\n  ')
    for subject in subjects:
        indexhtml.write('<li><a href="{}/qc_sub.html">{}</a></li>\n'.format(subject,subject))
    indexhtml.write('</ul>\n')
    indexhtml.write('</body>')
    indexhtml.close()

    for seed_dict in seed_vertices:
        write_index(qcdir,
                    subjects,
                    '{}_{}.png'.format(seed_dict['Name'],seed_dict['best_view']),
                    "theme-table-image col-sm-12",
                    '{}.html'.format(seed_dict['Name']),
                    "{} Seed Correlations".format(seed_dict['Name']))


#get rid of the tmpdir
shutil.rmtree(tmpdirbase)
