"""
Creating a registration template for FreeSurfer:
following http://surfer.nmr.mgh.harvard.edu/fswiki/SurfaceRegAndTemplates
(c) @rno klein, 2009

The process of creating a template involves calculating a "mean" (across subjects) 
pattern of curvature-related values, along with the variance of these variables. 
This task is performed by program mris_make_template. However, mris_make_template 
expects the input subjects to be already aligned (ie: already have ?h.sphere.reg files). 
This of course will not be the case if you are about to create a template!

The process starts ("Prep") by choosing one reference subject, and producing a template 
(Here, mytemp0.tif). No ?h.sphere is required; this step uses the subject's own pattern 
as the reference.

In Round 1, all reference subjects data are registered to the initial template, 
producing ?h.sphere.myreg0.  Using that registration, a new template (mytemp1.tif) 
is made using all reference subjects as input.

Round 2 is essentially a repeat of Round 1, but produces an improved template (mytemp2.tif) 
because Round 2's initial registration step is based on a better template than in Round 1.  
The steps of Round 2 could be iterated again if needed, but from what I've learned from Bruce 
Fischl, this is generally not necessary.

Notes

1. The website's diagram assumes that the subjects have all been processed once by 
FreeSurfer's recon-all pipeline or similar to produce all the required files needed for the 
template-creation process. Recon-all will also have produced ?h.sphere.reg aligned to the 
standard reference template (?h.average.curvature.filled.buckner40.tif), and corresponding 
annotation files -- for purposes here, these can be ignored. Indeed, once a satisfactory new 
template has been created, the recon-all script can be modified to use the new template, and 
the subjects can be re-run to create new improved annotations.

2. GW comment: Since the process uses a particular subject as the starting point, the final 
position of the template pattern is dependent on that first subject. The iterations of 
registration and template creation work to improve the variance aspect of the template.
"""

import os, sys

make_initial_template = 0
make_template = 1
iterations = 2
out_dir  = 'freesurfer_templates/'

subject_list = ['S01','S02','S03','S04','S05','S06','S07','S08','S09','S10',
                'S11','S12','S13','S14','S15','S16','S17','S18','S19','S20']
tempname = '1thru20'
#subject_list = ['S21','S22','S23','S24','S25','S26','S27','S28','S29','S30',
#                'S31','S32','S33','S34','S35','S36','S37','S38','S39','S40']
#tempname = '21thru40'
subjects_string = ' '.join(subject_list)

# Make initial template:                                        
# mris_make_template [options] <hemi> <surface name> <subject> <subject> ... <output name>
if make_initial_template:
    for hemi in ['lh','rh']:
        template = out_dir+hemi+'.'+tempname+'_0.tif'
        cmd = 'mris_make_template ' + hemi + ' sphere ' + subject_list[0] + ' ' + template
        print(cmd); os.system(cmd)

# Iteratively refine template:
if make_template:
    for i in range(0,iterations):
        for hemi in ['lh','rh']:
            for subject in subject_list:
                src_file = '$SUBJECTS_DIR/'+subject+'/surf/'+hemi+'.sphere'
                reg_file = '$SUBJECTS_DIR/'+subject+'/surf/'+hemi+'.sphere.'+tempname+'reg'+str(i)
                template = out_dir+hemi+'.'+tempname+'_'+str(i)+'.tif'
                cmd = 'mris_register -curv ' + src_file + ' ' + template + ' ' + reg_file
                print(cmd); os.system(cmd)

            template_new = out_dir+hemi+'.'+tempname+'_'+str(i+1)+'.tif'
            cmd = 'mris_make_template '+hemi+' sphere.'+tempname+'reg'+str(i)+' '+subjects_string+' '+template_new
            print(cmd); os.system(cmd)