#!/usr/bin/env python
"""
run_post_mcdc
"""
from dhcp.func import registration as reg, surface
from fsl.utils.filetree import FileTree
from dhcp.util.io import Path
from dhcp.util import util, fslpy as fsl
from dhcp.util.enums import SegType
from tempfile import TemporaryDirectory
import shutil

import numpy as np
import pandas as pd
import logging
import os
import os.path as op
import sys
import argparse
import logging


def _run(subid, sesid, workdir):
    
    # environment variables
    os.environ['DHCP_ANTS_PATH'] = "/vol/dhcp-results/bin/ANTs/bin/bin"
    os.environ['DHCP_C3D_PATH'] = "/vol/dhcp-results/bin/c3d/bin"

    sinfo = pd.read_csv('subject-info/fmri_all_10072020.csv', dtype={'subid': str, 'sesid': str})
    sinfo = sinfo[(sinfo.subid==subid) & (sinfo.sesid==sesid)]

    scan_ga = sinfo.iloc[0]['scan_ga']
    birth_ga = sinfo.iloc[0]['birth_ga']

    # setup logging
    logging.basicConfig(
        format='[%(asctime)s - %(name)s.%(funcName)s ] %(levelname)s : %(message)s',
        level=logging.INFO
    )

    # setup defaults
    defaults = FileTree.read(util.get_resource('dhcp-defaults.tree'), partial_fill=False).update(
        subid=subid,
        sesid=sesid,
        workdir=workdir,
    )
    
    # extract PVE (toblerone) if not already on disk
    
    if not defaults.on_disk('T2w_probseg', error=False):
        
        print(f'Calculate probseg: {subid} {sesid}')
    
        srf_defaults = defaults.update(mesh='native', space='T2w')
        surface.toblerone_pve(
            white_surf_right=srf_defaults.update(hemi='right').get('surf_white'),
            white_surf_left=srf_defaults.update(hemi='left').get('surf_white'),
            pial_surf_right=srf_defaults.update(hemi='right').get('surf_pial'),
            pial_surf_left=srf_defaults.update(hemi='left').get('surf_pial'),
            ref=defaults.get('T2w'),
            struct=defaults.get('T2w'),
            struct_to_ref_xfm='I',
            pve_name=defaults.get('T2w_probseg', make_dir=True),
        )
        util.update_sidecar(
            Path(defaults.get('T2w_probseg', make_dir=True)),
            seg_type=SegType.toblerone.name
        )
    else:
        print(f'Probseg already exists: {subid} {sesid}')


    ############
    # STANDARD
    ############
    
    # setup atlas

    age = int(np.round(scan_ga))

    atlas_tree = Path(util.get_setting('dhcp_volumetric_atlas_tree'))
    atlas = FileTree.read(atlas_tree).update(path=atlas_tree.dirname)


    # template (scan-age) -> struct
    
    if not defaults.update(src_space=f'template-{age}', ref_space='struct').on_disk('warp', error=False):
        
        reg.DEFAULTS = defaults.update(src_space=f'template-{age}', ref_space='struct')
    
        fsl.fslroi(input=defaults.get('T2w_probseg'), output=reg.DEFAULTS.get('ref_gmprob', make_dir=True), tmin=0, tsize=1)
        fsl.fslroi(input=defaults.get('T2w_probseg'), output=reg.DEFAULTS.get('ref_wmprob', make_dir=True), tmin=1, tsize=1)

        reg.template_to_struct(
            age=age,
            struct_brainmask=defaults.get('T2w_brainmask'),
            struct_T2w=defaults.get('T2w'),
            struct_gmprob=reg.DEFAULTS.get('ref_gmprob'),
            struct_wmprob=reg.DEFAULTS.get('ref_wmprob'),
            atlas=atlas,
            quick=False,
        )
    
    else:
        print(f'template-to-struct already exists: {subid} {sesid}')


    # struct -> age-matched template -> standard template (composite)

    standard_age = 40
    reg.DEFAULTS = defaults.update(src_space='struct', ref_space='standard')

    if not defaults.update(src_space='struct', ref_space='standard').on_disk('inv_warp', error=False):
    
        reg.struct_to_template_composite(
            struct=defaults.get('T2w'),
            struct2template_warp=defaults.update(src_space=f'template-{age}', ref_space='struct').get('inv_warp'),
            age=age,
            standard_age=standard_age,
            atlas=atlas,
        )
        
    else:
        print(f'struct-to-standard already exists: {subid} {sesid}')
    
    # func (undistorted) -> struct -> age-matched template -> standard template (composite)
    
    defaults = defaults.update(dc='dc')

    reg.DEFAULTS = defaults.update(src_space='func-mcdc', ref_space='standard')

    if defaults.update(src_space='func-mcdc', ref_space='struct').on_disk('affine', error=False):
    
        reg.func_to_template_composite(
            func=defaults.get('func_mcdc_mean'),
            func2struct_affine=defaults.update(src_space='func-mcdc', ref_space='struct').get('affine'),
            struct2template_warp=defaults.update(src_space='struct', ref_space='standard').get('warp'),
            standard_age=standard_age,
            atlas=atlas,
        )
        
    else:
        print(f'missing func-mcdc_to_struct: {subid} {sesid}')



if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('subid', help='subject ID')
    parser.add_argument('sesid', help='session ID')
    parser.add_argument('workdir', help='base workdir')

    args = parser.parse_args()
    _run(**vars(args))