#!/usr/bin/env python

# Part of a library of functions for dealing with obliquity in
# datasets.
#
# Here, we store an object that contains functionality for working
# with child_dsets when processing obliquity.
#
# auth : PA Taylor (SSCC, NIMH, NIH, USA)
# ============================================================================

import os, glob
import shutil

from   afnipy import afni_base    as ab
from   afnipy import lib_obl_name as LON

# ============================================================================

class ChildObj:
    """Object for setting up the processing of an child_dset.

After running this object to initialize and check things:
  
    cobj = ChildObj(...)

execute the proc_child() method to actually generate datasets:

    cobj.proc_child()

Then hopefully check your well-created results.

Parameters
----------
child_dset : str 
    (req to run) name of child_dset being input for processing
child_prefix : str 
    (req to run) name of child_prefix for outputting results
mat_obl : str
    (req to run) name of matrix with obliquity to apply to child_dset,
    which must be in ONELINE format; likely created earlier in
    obliquity_tool.py's processing
do_qc : bool
    do we do before/after QC images of overlaps? requires both inset_*
    args
inset_raw : str  
    name of the original inset for the obliquity tool program (just
    needed if do_qc=True)
inset_proc : str  
    name of the processed inset from the obliquity tool program (just
    needed if do_qc=True)

    """

    def __init__(self, 
                 child_dset=None, child_prefix=None, mat_obl=None, 
                 do_qc=True, inset_raw=None, inset_proc=None,
                 overwrite='',
                 do_clean=True, verb=1):

        # ----- set up attributes

        # main input variables
        self.child_dset         = child_dset    # input child_dset
        self.child_prefix       = child_prefix  # output prefix of child
        self.child_avsp         = ''            # 3dinfo -av_space CHILD_DSET
        self.child_nobj         = None          # name obj for child prefix
        self.mat_obl            = mat_obl       # obliquity matr to use

        # for QC images
        self.do_qc              = do_qc         # make olap images?
        self.qc_img_00          = "img_olap_00" # initial olap
        self.qc_img_01          = "img_olap_01" # post-proc olap
        self.inset_raw          = inset_raw     # inset to obliquity_tool.py
        self.inset_proc         = inset_proc    # inset after deobl

        # intermediate: generated by methods
        self.ok_to_proc         = False         # True after init is done
        self.wdir               = ''            # basename of working dir
        self.wdir_full          = ''            # full path of working dir
        self.child_00_dset      = ''            # name of dset_child in wdir
        self.mat_00             = 'mat_obl.aff12.1D'        # mat_obl in wdir
        self.child_00_aform     = 'child_00_aform.aff12.1D'  # child aform matr
        self.child_01_aform     = 'child_01_aform.aff12.1D'
        self.child_01_ijk2dicom = 'child_01_ijk2dicom.aff12.1D' # calc+create
        self.child_01_origin    = 'child_01_origin.1D'          # calc+create

        self.verb               = verb          # verbosity level
        self.overwrite          = overwrite     # string to overwrite or not
        self.do_clean           = do_clean      # rm wdir when done?

        # ----- take action(s)

        if not(self.child_dset is None) and \
           not(self.child_prefix is None) and \
           not(self.mat_obl is None) :
            tmp1 = self.basic_setup()
            tmp2 = self.make_wdir()

            # make sure we got to this point to be able to run proc
            self.ok_to_proc = True

    # ----- methods

    def proc_child(self):
        """Do main processing of applying obliquity matrix to child_dset"""

        if not(self.ok_to_proc) :
            txt = "Not enough inputs provided to run setup and prepare "
            txt+= "for main processing.  Please check ChildObj help and "
            txt+= "try again"
            ab.EP(txt)

        if self.verb:
            ab.IP("Process child dset: " + self.child_dset)

        # present dir
        cwd = os.getcwd()

        # copy dset to wdir as BRIK/HEAD
        cmd  = '3dTcat -overwrite -prefix {}/child_00 '.format(self.wdir_full)
        cmd += '"{}"'.format(self.child_dset)
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        if stat : 
            ab.EP("Failed in proc_child: 3dTcat")

        # fine the full name to the output (probably has +orig, but not sure)
        sss = '{}/child_00{}.HEAD'.format(self.wdir_full, self.child_avsp)
        self.child_00_dset = glob.glob(sss)
        if len(self.child_00_dset) != 1 :
            txt = "Should have 1 child_00_dset "
            txt+= "but actually have {}:\n".format(len(self.child_00_dset))
            txt+= "{}".format('\n'.join(self.child_00_dset))
            ab.EP(txt)
        # ... and store just the basename
        self.child_00_dset = os.path.basename(glob.glob(sss)[0])
            
        # get obl mat: cp into wdir
        cmd  = '\\cp {} '.format(self.mat_obl)
        cmd += '{}/{}'.format(self.wdir_full, self.mat_00)
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        if stat : 
            ab.EP("Failed in proc_child: cp")
        
        if self.do_qc :
            # inset_raw under child_raw (all obl applied)
            ulay = self.inset_raw
            olay = self.child_dset
            oimg = '{}/{}'.format(self.wdir_full, self.qc_img_00)

            # make the image
            tmp1 = self.make_qc_image(ulay, olay, oimg, 'in')

        # jump to wdir, simpler proc
        os.chdir(self.wdir_full)
        
        if self.verb :
            ab.IP("Applying obliquity")

        # get child aform (= IJK_TO_DICOM_REAL)
        cmd  = '3dinfo -aform_real_oneline {} '.format(self.child_00_dset)
        cmd += '> {}'.format(self.child_00_aform)
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        if stat : 
            ab.EP("Failed in proc_child: 3dinfo")

        # main calc for new matrix: apply obl mat to child aform
        cmd  = 'cat_matvec -ONELINE '
        cmd += '{} {} '.format(self.child_00_aform, self.mat_00)
        cmd += '> {}'.format(self.child_01_aform)
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        if stat : 
            ab.EP("Failed in proc_child: cat_matvec")

        # ----- read in extra info for attribute resetting

        # 1) read in mat: new aform (= IJK_TO_DICOM_REAL)
        cmd  = 'cat {}'.format(self.child_01_aform)
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        if stat : 
            ab.EP("Failed in proc_child: cat new aform")
        child_mat_aform = [float(x) for x in com.so[0].strip().split()]

        # 2) read in mat: current IJK_TO_DICOM, which we will edit
        cmd  = 'cat_matvec -ONELINE {}::IJK_TO_DICOM'.format(self.child_00_dset)
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        if stat : 
            ab.EP("Failed in proc_child: cat_matvec IJK_TO_DICOM")
        child_mat_ijk2dicom = [float(x) for x in com.so[0].strip().split()]
        # ... and update origin/offset part from new aform
        child_mat_ijk2dicom[3]  = child_mat_aform[3]
        child_mat_ijk2dicom[7]  = child_mat_aform[7]
        child_mat_ijk2dicom[11] = child_mat_aform[11]

        # ----- calc+write out new info 

        # write out mat: new IJK_TO_DICOM
        fff = open(self.child_01_ijk2dicom, 'w')
        ttt = '  '.join([str(x) for x in child_mat_ijk2dicom])
        fff.write(ttt + '\n')
        fff.close()

        # write out file: new ORIGIN
        fff = open(self.child_01_origin, 'w')
        ttt = '  '.join([str(x) for x in child_mat_aform[3::4]])
        fff.write(ttt + '\n')
        fff.close()

        # ----- attach new info to dset

        # attach new aform to dset
        cmd  = '3drefit -atrfloat IJK_TO_DICOM_REAL '
        cmd += '{} '.format(self.child_01_aform)
        cmd += '{}'.format(self.child_00_dset)
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        if stat : 
            ab.EP("Failed in proc_child: 3drefit IJK_TO_DICOM_REAL")

        # attach new IJK_TO_DICOM to dset
        cmd  = '3drefit -atrfloat IJK_TO_DICOM '
        cmd += '{} '.format(self.child_01_ijk2dicom)
        cmd += '{}'.format(self.child_00_dset)
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        if stat : 
            ab.EP("Failed in proc_child: 3drefit IJK_TO_DICOM")

        # attach new origin to dset
        cmd  = '3drefit -atrfloat ORIGIN '
        cmd += '{} '.format(self.child_01_origin)
        cmd += '{}'.format(self.child_00_dset)
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        if stat : 
            ab.EP("Failed in proc_child: 3drefit ORIGIN")

        # ----- finalize

        # put dset in output location (just up)
        cmd  = '3dcopy {} '.format(self.overwrite)
        cmd += '{} ../{}'.format(self.child_00_dset, self.child_nobj.bname)
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        if stat : 
            ab.EP("Failed in proc_child: 3dcopy")

        # ... and jump back to prior location
        os.chdir(cwd)

        if self.do_qc :
            # inset_proc under child_proc (all obl applied)
            ulay = self.inset_proc
            if not(os.path.isfile(ulay)):
                txt  = "For QC, cannot find ulay dset: " + ulay
                tmp1 = ab.EP1(txt)
                return -1
            olay = '{}/{}'.format(self.wdir_full, self.child_00_dset)
            oimg = '{}/{}'.format(self.wdir_full, self.qc_img_01)

            # make the image
            tmp2 = self.make_qc_image(ulay, olay, oimg, 'out')

        if self.do_clean :
            if self.verb > 1 :
                ab.IP("Cleaning work dir: " + self.wdir_full)
            shutil.rmtree(self.wdir_full)

        return 0

    def make_qc_image(self, ulay, olay, oimg, label=''):
        """Make a QC image and put it into the final spot"""

        msg = "child olap '{}'".format(label)
        if self.verb:
            ab.IP("Make QC image, " + msg)

        # for speed up (just copy 0th vol, if subbrick selectors
        # aren't being used already)
        if not("[" in ulay) :
            ulay+= "[0]"
        if not("[" in olay) :
            olay+= "[0]"

        cmd  = '@djunct_overlap_check '
        cmd += '-ulay "{}" '.format(ulay)
        cmd += '-olay "{}" '.format(olay)
        cmd += '-prefix {}'.format(oimg)
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        if stat : 
            ab.EP1("Failed to make QC image, " + msg)

        gimg = glob.glob(oimg + '_DEOB.jpg')
        if len(gimg) == 0 :
            gimg = glob.glob(oimg + '.jpg')
        if len(gimg) == 0 :
            txt = 'Cannot find any QC image: {}*'.format(oimg)
            ab.WP(txt)
        else:
            fimg = self.child_nobj.name_full_noext
            fimg+= '_olap_' + label + '.jpg'
            cmd  = '\\cp {} {}'.format(gimg[0], fimg)
            com  = ab.shell_com(cmd, capture=1)
            stat = com.run()
            if stat : 
                ab.EP("Failed to copy QC image, " + msg)
        
        return 0

    def basic_setup(self):
        """Make sure all inputs have been provided for proceeding, and
        generate some basic pieces.
        """

        # required inputs to run

        if self.child_dset is None :
            ab.EP("No child_dset entered")

        # get AFNI view of inset
        cmd  = '3dinfo -av_space ' + self.child_dset
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        avsp = com.so[0].strip()
        self.child_avsp = avsp

        if self.child_prefix is None :
            ab.EP("No child_prefix entered")

        if self.mat_obl is None :
            ab.EP("No mat_obl entered")

        # optional inputs

        if self.do_qc :
            if self.inset_raw is None :
                ab.EP("For the child dset QC, must provide inset_raw")
            if self.inset_proc is None :
                ab.EP("For the child dset QC, must provide inset_proc")

        # create object with pieces of names
        self.child_nobj = LON.NameObj(self.child_prefix)
        if self.verb > 5 :
            (self.child_nobj).print_dict()

        return 0

    def make_wdir(self):
        """Make working dir name and full path"""

        # child output basename, for wdir name
        bname_ne = self.child_nobj.bname_noext

        # child output path dirname, for wdir_full
        dname = self.child_nobj.dname

        # random str for wdir name
        cmd  = '3dnewid -fun11'
        com  = ab.shell_com(cmd, capture=1)
        stat = com.run()
        rstr = com.so[0].strip()

        # build wdir, name and full path
        self.wdir      = '__wdir_child_' + bname_ne + '_' + rstr
        self.wdir_full = dname + '/' + self.wdir

        # make dir (if nec)
        if not(os.path.isdir(self.wdir_full)) :
            try:
                os.makedirs(self.wdir_full, exist_ok=True)
            except:
                ab.EP("Could not make child's wdir_full: " + self.wdir_full)

        return 0

# ============================================================================

if __name__ == "__main__" :

    # an example use case
    print("++ No example")

