"""Image functions for PET data reconstruction and processing."""
__author__      = "Pawel J. Markiewicz"
__copyright__   = "Copyright 2018"
#-------------------------------------------------------------------------------

import sys
import os

import numpy as np
import math
from math import pi
import random
import scipy.ndimage as ndi

import nibabel as nib
import pydicom as dcm
import re
import glob
import logging

from subprocess import call
import time
import multiprocessing

from niftypet.nipet import mmraux
from niftypet import nimpa

import resources as rs

#===================================================================================
# IMAGE ROUTINES
#===================================================================================

def convert2e7(img, Cnt):
    '''Convert GPU optimised image to Siemens/E7 image shape (127,344,344).'''

    margin = (Cnt['SO_IMX']-Cnt['SZ_IMX'])/2

    #permute the dims first
    imo = np.transpose(img, (2,0,1))

    nvz = img.shape[2]

    #> get the x-axis filler and apply it
    filler = np.zeros((nvz, Cnt['SZ_IMY'], margin), dtype=np.float32)
    imo = np.concatenate((filler, imo, filler), axis=2)

    #> get the y-axis filler and apply it
    filler = np.zeros((nvz, margin, Cnt['SO_IMX']), dtype=np.float32)
    imo = np.concatenate((filler, imo, filler), axis=1)
    return imo

def convert2dev(im, Cnt):
    '''Reshape Siemens/E7 (default) image for optimal GPU execution.'''
    log = logging.getLogger(__name__)

    if im.shape[1]!=Cnt['SO_IMY'] or im.shape[2]!=Cnt['SO_IMX']:
        raise ValueError('e> input image array is not of the correct Siemens shape.')

    if 'rSZ_IMZ' in Cnt and im.shape[0]!=Cnt['rSZ_IMZ']:
        log.warning('the axial number of voxels does not match the reduced rings.')
    elif not 'rSZ_IMZ' in Cnt and im.shape[0]!=Cnt['SZ_IMZ']:
        log.warning('the axial number of voxels does not match the rings.')

    im_sqzd = np.zeros((im.shape[0], Cnt['SZ_IMY'], Cnt['SZ_IMX']), dtype=np.float32)
    margin = (Cnt['SO_IMX']-Cnt['SZ_IMX'])/2
    margin_=-margin
    if margin==0:
        margin = None
        margin_= None

    im_sqzd = im[:, margin:margin_, margin:margin_]
    im_sqzd = np.transpose(im_sqzd, (1, 2, 0))
    return im_sqzd

def cropxy(im, imsize, datain, Cnt, store_pth=''):
    '''Crop image transaxially to the size in tuple <imsize>.
    Return the image and the affine matrix.
    '''
    log = logging.getLogger(__name__)
    if not imsize[0]%2==0 and not imsize[1]%2==0:
        log.error('image size has to be an even number!')
        return None

    # cropping indexes
    i0 = (Cnt['SO_IMX']-imsize[0])/2
    i1 = (Cnt['SO_IMY']+imsize[1])/2

    B = image_affine(datain, Cnt, gantry_offset=False)
    B[0,3] -= 10*Cnt['SO_VXX']*i0
    B[1,3] += 10*Cnt['SO_VXY']*(Cnt['SO_IMY']-i1)

    cim = im[:, i0:i1, i0:i1]

    if store_pth!='':
        nimpa.array2nii( cim[::-1,::-1,:], B, store_pth, descrip='cropped')
        log.debug('saved cropped image to:' + store_pth)

    return cim, B
#-------------------------------------------------------------------------------------------

def image_affine(datain, Cnt, gantry_offset=False):
    '''Creates a blank reference image, to which another image will be resampled'''

    #------get necessary data for -----
    # gantry offset
    if gantry_offset:
        goff, tpo = mmraux.lm_pos(datain, Cnt)
    else:
        goff = np.zeros((3))
    vbed, hbed = mmraux.vh_bedpos(datain, Cnt)
    # create a reference empty mu-map image
    B = np.diag(np.array([-10*Cnt['SO_VXX'], 10*Cnt['SO_VXY'], 10*Cnt['SO_VXZ'], 1]))
    B[0,3] = 10*(.5*Cnt['SO_IMX']*Cnt['SO_VXX']      + goff[0])
    B[1,3] = 10*((-.5*Cnt['SO_IMY']+1)*Cnt['SO_VXY'] - goff[1])
    B[2,3] = 10*((-.5*Cnt['SO_IMZ']+1)*Cnt['SO_VXZ'] - goff[2] + hbed)
    # -------------------------------------------------------------------------------------
    return B

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

def getmu_off(mu, Cnt, Offst=np.array([0., 0., 0.])):
    #number of voxels
    nvx = mu.shape[0]
    #change the shape to 3D
    mu.shape = (Cnt['SO_IMZ'], Cnt['SO_IMY'], Cnt['SO_IMX'])

    #-------------------------------------------------------------------------
    # CORRECT THE MU-MAP for GANTRY OFFSET
    #-------------------------------------------------------------------------
    Cim = {
        'VXSOx':0.208626,
        'VXSOy':0.208626,
        'VXSOz':0.203125,
        'VXNOx':344,
        'VXNOy':344,
        'VXNOz':127,

        'VXSRx':0.208626,
        'VXSRy':0.208626,
        'VXSRz':0.203125,
        'VXNRx':344,
        'VXNRy':344,
        'VXNRz':127
    }
    #original image offset
    Cim['OFFOx'] = -0.5*Cim['VXNOx']*Cim['VXSOx']
    Cim['OFFOy'] = -0.5*Cim['VXNOy']*Cim['VXSOy']
    Cim['OFFOz'] = -0.5*Cim['VXNOz']*Cim['VXSOz']
    #resampled image offset
    Cim['OFFRx'] = -0.5*Cim['VXNRx']*Cim['VXSRx']
    Cim['OFFRy'] = -0.5*Cim['VXNRy']*Cim['VXSRy']
    Cim['OFFRz'] = -0.5*Cim['VXNRz']*Cim['VXSRz']
    #transformation matrix
    A = np.array(
        [[ 1., 0., 0.,  Offst[0] ],
        [  0., 1., 0.,  Offst[1] ],
        [  0., 0., 1.,  Offst[2] ],
        [  0., 0., 0.,  1. ]], dtype=np.float32
        )
    #apply the gantry offset to the mu-map
    mur = nimpa.prc.improc.resample(mu, A, Cim)
    return mur

def getinterfile_off(fmu, Cnt, Offst=np.array([0., 0., 0.])):
    '''
    Return the floating point mu-map in an array from Interfile, accounting for image offset (does slow interpolation).
    '''
    #read the image file
    f = open(fmu, 'rb')
    mu = np.fromfile(f, np.float32)
    f.close()

    # save_im(mur, Cnt, os.path.dirname(fmu) + '/mur.nii')
    #-------------------------------------------------------------------------
    mur = getmu_off(mu, Cnt)
    #create GPU version of the mu-map
    murs = convert2dev(mur, Cnt)
    #number of voxels
    nvx = im.shape[0]
    #get the basic stats
    mumax = np.max(mur)
    mumin = np.min(mur)
    #number of voxels greater than 10% of max image value
    n10mx = np.sum(mur>0.1*mumax)
    #return image dictionary with the image itself and some other stats
    mu_dct = {'im':mur,
              'ims':murs,
              'max':mumax,
              'min':mumin,
              'nvx':nvx,
              'n10mx':n10mx}
    return mu_dct

#================================================================================
def getinterfile(fim, Cnt):
    '''Return the floating point image file in an array from an Interfile file.'''

    #read the image file
    f = open(fim, 'rb')
    im = np.fromfile(f, np.float32)
    f.close()

    #number of voxels
    nvx = im.shape[0]
    #change the shape to 3D
    im.shape = (Cnt['SO_IMZ'], Cnt['SO_IMY'], Cnt['SO_IMX'])

    #get the basic stats
    immax = np.max(im)
    immin = np.min(im)

    #number of voxels greater than 10% of max image value
    n10mx = np.sum(im>0.1*immax)

    #reorganise the image for optimal gpu execution
    im_sqzd = convert2dev(im, Cnt)

    #return image dictionary with the image itself and some other stats
    im_dct = {'im':im,
              'ims':im_sqzd,
              'max':immax,
              'min':immin,
              'nvx':nvx,
              'n10mx':n10mx}

    return im_dct
#================================================================================

#=================================================================================
#-define uniform cylinder
def get_cylinder(Cnt, rad=25, xo=0, yo=0, unival=1, gpu_dim=False):
    '''Outputs image with a uniform cylinder of intensity = unival, radius = rad, and transaxial centre (xo, yo)'''
    imdsk = np.zeros((1, Cnt['SO_IMX'], Cnt['SO_IMY']), dtype=np.float32)
    for t in np.arange(0, math.pi, math.pi/(2*360)):
        x = xo+rad*math.cos(t)
        y = yo+rad*math.sin(t)
        yf = np.arange(-y+2*yo, y, Cnt['SO_VXY']/2)
        v = np.int32(.5*Cnt['SO_IMX'] - np.ceil(yf/Cnt['SO_VXY']))
        u = np.int32(.5*Cnt['SO_IMY'] + np.floor(x/Cnt['SO_VXY']))
        imdsk[0,v,u] = unival
    if 'rSO_IMZ' in Cnt:
        nvz = Cnt['rSO_IMZ']
    else:
        nvz = Cnt['SO_IMZ']
    imdsk = np.repeat(imdsk, nvz, axis=0)
    if gpu_dim: imdsk = convert2dev(imdsk, Cnt)
    return imdsk

#================================================================================
def hu2mu(im):
    '''HU units to 511keV PET mu-values'''

    # convert nans to -1024 for the HU values only
    im[np.isnan(im)] = -1024
    # constants
    muwater  = 0.096
    mubone   = 0.172
    rhowater = 0.158
    rhobone  = 0.326
    uim = np.zeros(im.shape, dtype=np.float32)
    uim[im<=0] = muwater * ( 1+im[im<=0]*1e-3 )
    uim[im> 0] = muwater * \
        ( 1+im[im>0]*1e-3 * rhowater/muwater*(mubone-muwater)/(rhobone-rhowater) )
    # remove negative values
    uim[uim<0] = 0
    return uim

# =====================================================================================
# object/patient mu-map resampling to NIfTI
# better use dcm2niix
def mudcm2nii(datain, Cnt):
    '''DICOM mu-map to NIfTI'''
    log = logging.getLogger(__name__)

    mu, pos, ornt = nimpa.dcm2im(datain['mumapDCM'])
    mu *= 0.0001
    A = pos['AFFINE']
    A[0,0] *= -1
    A[0,3] *= -1
    A[1,3] += A[1,1]
    nimpa.array2nii(mu[:,::-1,:], A, os.path.join(os.path.dirname(datain['mumapDCM']),'mu.nii.gz'))

    #------get necessary data for creating a blank reference image (to which resample)-----
    # gantry offset
    goff, tpo = mmraux.lm_pos(datain, Cnt)
    ihdr, csainfo = mmraux.hdr_lm(datain)
    #start horizontal bed position
    p = re.compile(r'start horizontal bed position.*\d{1,3}\.*\d*')
    m = p.search(ihdr)
    fi = ihdr[m.start():m.end()].find('=')
    hbedpos = 0.1*float(ihdr[m.start()+fi+1:m.end()])

    B = np.diag(np.array([-10*Cnt['SO_VXX'], 10*Cnt['SO_VXY'], 10*Cnt['SO_VXZ'], 1]))
    B[0,3] = 10*(.5*Cnt['SO_IMX']*Cnt['SO_VXX']      + goff[0])
    B[1,3] = 10*((-.5*Cnt['SO_IMY']+1)*Cnt['SO_VXY'] - goff[1])
    B[2,3] = 10*((-.5*Cnt['SO_IMZ']+1)*Cnt['SO_VXZ'] - goff[2] + hbedpos)
    im = np.zeros((Cnt['SO_IMZ'], Cnt['SO_IMY'], Cnt['SO_IMX']), dtype=np.float32)
    nimpa.array2nii(im, B, os.path.join(os.path.dirname(datain['mumapDCM']),'muref.nii.gz'))
    # -------------------------------------------------------------------------------------
    fmu = os.path.join(os.path.dirname(datain['mumapDCM']),'mu_r.nii.gz')
    if os.path.isfile( Cnt['RESPATH'] ):
        call( [ Cnt['RESPATH'],
                    '-ref', os.path.join(os.path.dirname(datain['mumapDCM']),'muref.nii.gz'),
                    '-flo', os.path.join(os.path.dirname(datain['mumapDCM']),'mu.nii.gz'),
                    '-res', fmu,
                    '-pad', '0'] )
    else:
        log.error('path to resampling executable is incorrect!')
        raise IOError('Error launching NiftyReg for image resampling.')

    return fmu

# =====================================================================================
def obj_mumap(
        datain,
        params=None,
        outpath='',
        comment='',
        store=False,
        store_npy=False,
        gantry_offset=True,
        del_auxilary=True,
        ):
    '''Get the object mu-map from DICOM images'''
    log = logging.getLogger(__name__)
    if params is None:
        params = {}

    # three ways of passing scanner constants <Cnt> are here decoded
    if 'Cnt' in params:
        Cnt = params['Cnt']
    elif 'SO_IMZ' in params:
        Cnt = params
    else:
        Cnt = rs.get_mmr_constants()

    # output folder
    if outpath=='':
        fmudir = os.path.join( datain['corepath'], 'mumap-obj' )
    else:
        fmudir = os.path.join( outpath, 'mumap-obj' )
    nimpa.create_dir(fmudir)
    # ref file name
    fmuref = os.path.join(fmudir, 'muref.nii.gz')
    # ref affine
    B = image_affine(datain, Cnt, gantry_offset=gantry_offset)
    # ref image (blank)
    im = np.zeros((Cnt['SO_IMZ'], Cnt['SO_IMY'], Cnt['SO_IMX']), dtype=np.float32)
    # store ref image
    nimpa.array2nii(im, B, fmuref)

    # check if the object dicom files for MR-based mu-map exists
    if not 'mumapDCM' in datain or not os.path.isdir(datain['mumapDCM']):
        log.error('DICOM forlder for the mu-map does not exist.')
        return None

    fnii = 'converted-from-object-DICOM_'
    tstmp = nimpa.time_stamp(simple_ascii=True)

    # find residual(s) from previous runs and delete them
    resdcm = glob.glob( os.path.join(fmudir, '*'+fnii+'*.nii*') )
    for d in resdcm:
        os.remove(d)

    # convert the DICOM mu-map images to nii
    call( [ Cnt['DCM2NIIX'], '-f', fnii+tstmp, '-o', fmudir, datain['mumapDCM'] ] )
    #files for the T1w, pick one:
    fmunii = glob.glob( os.path.join(fmudir, '*'+fnii+tstmp+'*.nii*') )[0]
    # fmunii = glob.glob( os.path.join(datain['mumapDCM'], '*converted*.nii*') )
    # fmunii = fmunii[0]

    # the converted nii image resample to the reference size
    fmu = os.path.join(fmudir, comment+'mumap_tmp.nii.gz')
    if os.path.isfile( Cnt['RESPATH'] ):
        cmd = [ Cnt['RESPATH'],
                    '-ref', fmuref,
                    '-flo', fmunii,
                    '-res', fmu,
                    '-pad', '0']
        if log.getEffectiveLevel() > logging.DEBUG:
            cmd.append('-voff')
        call(cmd)
    else:
        log.error('path to resampling executable is incorrect!')
        sys.exit()

    nim = nib.load(fmu)
    # get the affine transform
    A = nim.get_sform()
    mu = nim.get_data()
    mu = np.transpose(mu[:,::-1,::-1], (2, 1, 0))
    # convert to mu-values
    mu = np.float32(mu)/1e4
    mu[mu<0] = 0

    #> return image dictionary with the image itself and some other stats
    mu_dct = dict(im=mu, affine=A)
    if not del_auxilary:
        mu_dct['fmuref'] = fmuref

    # store the mu-map if requested
    if store_npy:
        # to numpy array
        fnp = os.path.join(fmudir, 'mumap-from-DICOM.npy' )
        np.save(fnp, (mu, A))

    if store:
        # with this file name
        fmumap = os.path.join(fmudir, 'mumap-from-DICOM_no-alignment'+comment+'.nii.gz')
        nimpa.array2nii(mu[::-1,::-1,:], A, fmumap)
        mu_dct['fim'] = fmumap

    if del_auxilary:
        os.remove(fmuref)
        os.remove(fmunii)
        os.remove(fmu)

    return mu_dct

#=================================================================================
# pCT/UTE MU-MAP ALIGNED
#---------------------------------------------------------------------------------
def align_mumap(
        datain,
        scanner_params=None,
        outpath='',
        use_stored=False,
        hst=None,
        t0=0, t1=0,
        itr=2,
        faff='',
        fpet='',
        fcomment='',
        store=False,
        store_npy=False,
        petopt='ac',
        musrc='ute', # another option is pct for mu-map source
        ute_name='UTE2',
        del_auxilary=True,
        verbose=True,
    ):
    '''
    Align the a pCT or MR-derived mu-map to a PET image reconstructed to chosen
    specifications (e.g., with/without attenuation and scatter corrections)

    use_sotred only works if hst or t0/t1 given but not when faff.
    '''
    log = logging.getLogger(__name__)
    if scanner_params is None:
        scanner_params = {}
    if hst is None:
        hst = []

    #> output folder
    if outpath=='':
        opth = os.path.join( datain['corepath'], 'mumap-obj' )
    else:
        opth = os.path.join( outpath, 'mumap-obj' )

    #> create the folder, if not existent
    nimpa.create_dir(opth)

    #> get the timing of PET if affine not given
    if faff!='' and 't0' in hst:
        t0 = hst['t0']
        t1 = hst['t1']

    #> file name for the output mu-map
    fnm = 'mumap-' + musrc.upper()

    #> output dictionary
    mu_dct = {}

    #---------------------------------------------------------------------------
    #> used stored if requested
    if use_stored and affine!='':
        fmu_stored = fnm + '-aligned-to_t'\
                     + str(hst['t0'])+'-'+str(hst['t1'])+'_'+petopt.upper()\
                     + fcomment
        fmupath = os.path.join(opth, fmu_stored)
        if os.path.isfile( fmupath ):
            mudct_stored = nimpa.getnii(fmupath, output='all')
            #> create output dictionary
            mu_dct['im'] = mudct_stored['im']
            mu_dct['affine'] = mudct_stored['affine']
            #mu_dct['faff'] = faff
            return mu_dct
    #---------------------------------------------------------------------------

    #> three ways of passing scanner constants <Cnt> are here decoded
    if 'Cnt' in scanner_params:
        Cnt = scanner_params['Cnt']
    elif 'SO_IMZ' in scanner_params:
        Cnt = scanner_params
    else:
        Cnt = rs.get_mmr_constants()
    Cnt   = scanner_params['Cnt']

    #> if affine not provided histogram the LM data for recon and registration
    if not os.path.isfile(faff):
        from niftypet.nipet.prj import mmrrec
        #-histogram the list data if needed
        if not hst:
            from niftypet.nipet import mmrhist
            if 'txLUT' in scanner_params:
                hst = mmrhist(datain, scanner_params, t0=t0, t1=t1)
            else:
                raise ValueError(
                    'e> full scanner parameters not provided\
                     but required for histogramming.')

    #=========================================================
    #-get hardware mu-map
    if 'hmumap' in datain and os.path.isfile(datain['hmumap']):
        muh, _, _ = np.load(datain['hmumap'], allow_pickle=True)
        log.debug('loaded hardware mu-map from file: %s' % datain['hmumap'])
    elif outpath!='':
        hmupath = os.path.join( os.path.join(outpath,'mumap-hdw'), 'hmumap.npy')
        if os.path.isfile( hmupath ):
            muh, _, _ = np.load(hmupath, allow_pickle=True)
            datain['hmumap'] = hmupath
        else:
            raise IOError('Invalid path to the hardware mu-map')
    else:
        log.error('obtain the hardware mu-map first.')
        raise IOError('Could not find the hardware mu-map.  Have you run the routine for hardware mu-map?')
    #=========================================================
    #-check if T1w image is available
    if not 'MRT1W#' in datain and not 'T1nii' in datain and not 'T1bc' in datain \
    and not 'T1N4' in datain:
        log.error('no MR T1w images required for co-registration!')
        raise IOError('T1w image could not be obtained')
    #=========================================================

    #-if the affine is not given,
    #-it will be generated by reconstructing PET image, with some or no corrections
    if not os.path.isfile(faff):
        # first recon pet to get the T1 aligned to it
        if petopt=='qnt':
            # ---------------------------------------------
            # OPTION 1 (quantitative recon with all corrections using MR-based mu-map)
            # get UTE object mu-map (may not be in register with the PET data)
            mudic = obj_mumap(
                        datain,
                        Cnt,
                        del_auxilary=del_auxilary)
            muo = mudic['im']
            # reconstruct PET image with UTE mu-map to which co-register T1w
            recout = mmrrec.osemone(
                datain, [muh, muo],
                hst, scanner_params,
                recmod=3, itr=itr, fwhm=0.,
                fcomment=fcomment+'_QNT-UTE',
                outpath=os.path.join(outpath, 'PET', 'positioning'),
                store_img=True)
        elif petopt=='nac':
            # ---------------------------------------------
            # OPTION 2 (recon without any corrections for scatter and attenuation)
            # reconstruct PET image with UTE mu-map to which co-register T1w
            muo = np.zeros(muh.shape, dtype=muh.dtype)
            recout = mmrrec.osemone(
                datain, [muh, muo],
                hst, scanner_params,
                recmod=1, itr=itr, fwhm=0.,
                fcomment=fcomment+'_NAC',
                outpath=os.path.join(outpath,'PET', 'positioning'),
                store_img=True)
        elif petopt=='ac':
            # ---------------------------------------------
            # OPTION 3 (recon with attenuation correction only but no scatter)
            # reconstruct PET image with UTE mu-map to which co-register T1w
            mudic = obj_mumap(
                    datain,
                    Cnt,
                    del_auxilary=del_auxilary)
            muo = mudic['im']
            recout = mmrrec.osemone(
                datain, [muh, muo],
                hst, scanner_params,
                recmod=1, itr=itr, fwhm=0.,
                fcomment=fcomment+'_AC-UTE',
                outpath=os.path.join(outpath,'PET', 'positioning'),
                store_img=True)

        fpet = recout.fpet
        mu_dct['fpet'] = fpet

        #------------------------------
        if musrc=='ute' and ute_name in datain and os.path.exists(datain[ute_name]):
            # change to NIfTI if the UTE sequence is in DICOM files (folder)
            if os.path.isdir(datain[ute_name]):
                fnew =  os.path.basename(datain[ute_name])
                call( [ Cnt['DCM2NIIX'], '-f', fnew, datain[ute_name] ] )
                fute = glob.glob(os.path.join(datain[ute_name], fnew+'*nii*'))[0]
            elif os.path.isfile(datain[ute_name]):
                fute = datain[ute_name]
            # get the affine transformation
            try:
                regdct = nimpa.coreg_spm(
                    fpet,
                    fute,
                    outpath=os.path.join(outpath,'PET', 'positioning')
                )
            except:
                regdct = nimpa.affine_niftyreg(
                    fpet,
                    fute,
                    outpath=os.path.join(outpath,'PET', 'positioning'),
                    #fcomment=fcomment,
                    executable = Cnt['REGPATH'],
                    omp = multiprocessing.cpu_count()/2,
                    rigOnly = True,
                    affDirect = False,
                    maxit=5,
                    speed=True,
                    pi=50, pv=50,
                    smof=0, smor=0,
                    rmsk=True,
                    fmsk=True,
                    rfwhm=15., #millilitres
                    rthrsh=0.05,
                    ffwhm = 15., #millilitres
                    fthrsh=0.05,
                    verbose=log.getEffectiveLevel() <= logging.INFO
                )

            faff_mrpet = regdct['faff']

        elif musrc=='pct':

            ft1w = nimpa.pick_t1w(datain)

            try:
                regdct = nimpa.coreg_spm(
                    fpet,
                    ft1w,
                    outpath=os.path.join(outpath,'PET', 'positioning')
                )
            except:
                regdct = nimpa.affine_niftyreg(
                    fpet,
                    ft1w,
                    outpath=os.path.join(outpath,'PET', 'positioning'),
                    executable = Cnt['REGPATH'],
                    omp = multiprocessing.cpu_count()/2,
                    rigOnly = True,
                    affDirect = False,
                    maxit=5,
                    speed=True,
                    pi=50, pv=50,
                    smof=0, smor=0,
                    rmsk=True,
                    fmsk=True,
                    rfwhm=15., #millilitres
                    rthrsh=0.05,
                    ffwhm = 15., #millilitres
                    fthrsh=0.05,
                    verbose=log.getEffectiveLevel() < logging.INFO
                )

            faff_mrpet = regdct['faff']

        else:
            raise IOError('Floating MR image not provided or is invalid.')

    else:
        faff_mrpet = faff
        regdct = {}
        if not os.path.isfile(fpet):
            raise IOError('e> the reference PET should be supplied with the affine.')

    #> output file name for the aligned mu-maps
    if musrc=='pct':
        freg = os.path.join(opth, 'pCT-res-tmp'+fcomment+'.nii.gz')
        fflo = datain['pCT']
    elif musrc=='ute':
        freg = os.path.join(opth, 'UTE-res-tmp'+fcomment+'.nii.gz')
        if 'UTE' not in datain:
            fnii = 'converted-from-DICOM_'
            tstmp = nimpa.time_stamp(simple_ascii=True)
            # convert the DICOM mu-map images to nii
            if 'mumapDCM' not in datain:
                raise IOError('DICOM with the UTE mu-map are not given.')
            call( [ Cnt['DCM2NIIX'], '-f', fnii+tstmp, '-o', opth, datain['mumapDCM'] ] )
            #files for the T1w, pick one:
            fflo = glob.glob( os.path.join(opth, '*'+fnii+tstmp+'*.nii*') )[0]
        else:
            if os.path.isfile(datain['UTE']):
                fflo = datain['UTE']
            else:
                raise IOError('The provided NIfTI UTE path is not valid.')

    #> call the resampling routine to get the pCT/UTE in place
    if 'matlab_eng' in regdct:
        nimpa.resample_spm(
            fpet,
            fflo,
            faff_mrpet,
            matlab_eng=regdct['matlab_eng'],
            fimout=freg,
            del_ref_uncmpr=True,
            del_flo_uncmpr=True,
            del_out_uncmpr=True
        )
    else:
        nimpa.resample_niftyreg(
            fpet,
            fflo,
            faff_mrpet,
            fimout=freg,
            executable=Cnt['RESPATH'],
            verbose=log.getEffectiveLevel() < logging.INFO)

    #-get the NIfTI of registered image
    nim = nib.load(freg)
    A   = nim.affine
    imreg = np.float32( nim.get_data() )
    imreg = imreg[:,::-1,::-1]
    imreg = np.transpose(imreg, (2, 1, 0))

    #-convert to mu-values; sort out the file name too.
    if musrc=='pct':
        mu = hu2mu(imreg)
    elif musrc=='ute':
        mu = np.float32(imreg)/1e4
        #-remove the converted file from DICOMs
        os.remove(fflo)
    else:
        raise NameError('Confused o_O.')

    # get rid of negatives
    mu[mu<0] = 0
    # return image dictionary with the image itself and other parameters
    mu_dct['im'] = mu
    mu_dct['affine'] = A
    mu_dct['faff'] = faff

    if store or store_npy:

        nimpa.create_dir(opth)
        if faff=='':
            fname = fnm + '-aligned-to_t'\
                    + str(hst['t0'])+'-'+str(hst['t1'])+'_'+petopt.upper()\
                    + fcomment
        else:
            fname = fnm + '-aligned-to-given-affine' + fcomment
    if store_npy:
        #> Numpy
        if store_to_npy:
            fnp = os.path.join(opth, fname + '.npy')
            np.save(fnp, (mu, A, fnp))
    if store:
        #> NIfTI
        fmu = os.path.join(opth, fname + '.nii.gz')
        nimpa.array2nii(mu[::-1,::-1,:], A, fmu)
        mu_dct['fim'] = fmu
    if del_auxilary:
        os.remove(freg)
        if not os.path.isfile(faff):
            os.remove(fute)

    return mu_dct

#=================================================================================
# PSEUDO CT MU-MAP
#---------------------------------------------------------------------------------
def pct_mumap(
        datain, scanner_params,
        hst=None, t0=0, t1=0, itr=2,
        petopt='ac', faff='', fpet='',
        fcomment='', outpath='',
        store_npy=False, store=False
    ):
    '''
    GET THE MU-MAP from pCT IMAGE (which is in T1w space)
    * the mu-map will be registered to PET which will be reconstructed for time frame t0-t1
    * it f0 and t1 are not given the whole LM dataset will be reconstructed
    * the reconstructed PET can be attenuation and scatter corrected or NOT using petopt
    '''
    log = logging.getLogger(__name__)
    if hst is None:
        hst = []

    if not os.path.isfile(faff):
        from niftypet.nipet.prj import mmrrec
        # histogram the list data if needed
        if not hst:
            from niftypet.nipet.lm import mmrhist
            hst = mmrhist.mmrhist(datain, scanner_params, t0=t0, t1=t1)

    # constants, transaxial and axial LUTs are extracted
    Cnt   = scanner_params['Cnt']
    txLUT = scanner_params['txLUT']
    axLUT = scanner_params['axLUT']

    # get hardware mu-map
    if 'hmumap' in datain and os.path.isfile(datain['hmumap']):
        muh, _, _ = np.load(datain['hmumap'], allow_pickle=True)
        log.debug('loaded hardware mu-map from file:' + datain['hmumap'])
    elif outpath!='':
        hmupath = os.path.join( os.path.join(outpath,'mumap-hdw'), 'hmumap.npy')
        if os.path.isfile( hmupath ):
            muh, _, _ = np.load(hmupath, allow_pickle=True)
            datain['hmumap'] = hmupath
        else:
            raise IOError('Invalid path to the hardware mu-map')
    else:
        log.error('obtain the hardware mu-map first.')
        raise IOError('Could not find the hardware mu-map.  Have you run the routine for hardware mu-map?')

    if not 'MRT1W#' in datain and not 'T1nii' in datain and not 'T1bc' in datain:
        log.error('no MR T1w images required for co-registration!')
        raise IOError('Missing MR data')
    # ----------------------------------

    # output dictionary
    mu_dct = {}
    if not os.path.isfile(faff):
        # first recon pet to get the T1 aligned to it
        if petopt=='qnt':
            # ---------------------------------------------
            # OPTION 1 (quantitative recon with all corrections using MR-based mu-map)
            # get UTE object mu-map (may not be in register with the PET data)
            mudic = obj_mumap(datain, Cnt)
            muo = mudic['im']
            # reconstruct PET image with UTE mu-map to which co-register T1w
            recout = mmrrec.osemone(
                datain, [muh, muo],
                hst, scanner_params,
                recmod=3, itr=itr, fwhm=0.,
                fcomment=fcomment+'_qntUTE',
                outpath=os.path.join(outpath, 'PET', 'positioning'),
                store_img=True)
        elif petopt=='nac':
            # ---------------------------------------------
            # OPTION 2 (recon without any corrections for scatter and attenuation)
            # reconstruct PET image with UTE mu-map to which co-register T1w
            muo = np.zeros(muh.shape, dtype=muh.dtype)
            recout = mmrrec.osemone(
                datain, [muh, muo],
                hst, scanner_params,
                recmod=1, itr=itr, fwhm=0.,
                fcomment=fcomment+'_NAC',
                outpath=os.path.join(outpath, 'PET', 'positioning'),
                store_img=True)
        elif petopt=='ac':
            # ---------------------------------------------
            # OPTION 3 (recon with attenuation correction only but no scatter)
            # reconstruct PET image with UTE mu-map to which co-register T1w
            mudic = obj_mumap(datain, Cnt, outpath=outpath)
            muo = mudic['im']
            recout = mmrrec.osemone(
                datain, [muh, muo],
                hst, scanner_params,
                recmod=1, itr=itr, fwhm=0.,
                fcomment=fcomment+'_AC',
                outpath=os.path.join(outpath, 'PET', 'positioning'),
                store_img=True)

        fpet = recout.fpet
        mu_dct['fpet'] = fpet
        #------------------------------
        # get the affine transformation

        ft1w = nimpa.pick_t1w(datain)
        try:
            regdct = nimpa.coreg_spm(
                fpet,
                ft1w,
                outpath=os.path.join(outpath,'PET', 'positioning')
            )
        except:
            regdct = nimpa.affine_niftyreg(
                fpet,
                ft1w,
                outpath=os.path.join(outpath,'PET', 'positioning'),
                #fcomment=fcomment,
                executable = Cnt['REGPATH'],
                omp = multiprocessing.cpu_count()/2,
                rigOnly = True,
                affDirect = False,
                maxit=5,
                speed=True,
                pi=50, pv=50,
                smof=0, smor=0,
                rmsk=True,
                fmsk=True,
                rfwhm=15., #millilitres
                rthrsh=0.05,
                ffwhm = 15., #millilitres
                fthrsh=0.05,
                verbose=log.getEffectiveLevel() < logging.INFO
            )

        faff = regdct['faff']
        #------------------------------

    # pCT file name
    if outpath=='':
        pctdir = os.path.dirname(datain['pCT'])
    else:
        pctdir = os.path.join(outpath, 'mumap-obj')
    mmraux.create_dir(pctdir)
    fpct = os.path.join(pctdir, 'pCT_r_tmp'+fcomment+'.nii.gz')

    #call the resampling routine to get the pCT in place
    if os.path.isfile( Cnt['RESPATH'] ):
        cmd = [Cnt['RESPATH'],
            '-ref', fpet,
            '-flo', datain['pCT'],
            '-trans', faff,
            '-res', fpct,
            '-pad', '0']
        if log.getEffectiveLevel() > logging.DEBUG:
            cmd.append('-voff')
        call(cmd)
    else:
        log.error('path to resampling executable is incorrect!')
        sys.exit()

    # get the NIfTI of the pCT
    nim = nib.load(fpct)
    A   = nim.get_sform()
    pct = np.float32( nim.get_data() )
    pct = pct[:,::-1,::-1]
    pct = np.transpose(pct, (2, 1, 0))
    # convert the HU units to mu-values
    mu = hu2mu(pct)
    # get rid of negatives
    mu[mu<0] = 0

    # return image dictionary with the image itself and other parameters
    mu_dct['im'] = mu
    mu_dct['affine'] = A
    mu_dct['faff'] = faff

    if store:
        # now save to numpy array and NIfTI in this folder
        if outpath=='':
            pctumapdir = os.path.join( datain['corepath'], 'mumap-obj' )
        else:
            pctumapdir = os.path.join(outpath, 'mumap-obj')
        mmraux.create_dir(pctumapdir)
        #> Numpy
        if store_npy:
            fnp = os.path.join(pctumapdir, 'mumap-pCT.npy')
            np.save(fnp, (mu, A, fnp))
        # NIfTI
        fmu = os.path.join(pctumapdir, 'mumap-pCT' +fcomment+ '.nii.gz')
        nimpa.array2nii(mu[::-1,::-1,:], A, fmu)
        mu_dct['fim'] = fmu
        datain['mumapCT'] = fmu

    return mu_dct

#*********************************************************************************
#GET HARDWARE MU-MAPS with positions and offsets
#---------------------------------------------------------------------------------
def hdr_mu(datain, Cnt):
    '''Get the headers from DICOM data file'''
    log = logging.getLogger(__name__)
    #get one of the DICOM files of the mu-map
    if 'mumapDCM' in datain:
        files = glob.glob(os.path.join(datain['mumapDCM'],'*.dcm'))
        files.extend(glob.glob(os.path.join(datain['mumapDCM'],'*.DCM')))
        files.extend(glob.glob(os.path.join(datain['mumapDCM'],'*.ima')))
        files.extend(glob.glob(os.path.join(datain['mumapDCM'],'*.IMA')))
        dcmf = files[0]
    else:
        raise NameError('no DICOM or DICOM filed <CSA Series Header Info> found!')

    if os.path.isfile( dcmf ):
        dhdr = dcm.read_file( dcmf )
    else:
        log.error('DICOM mMR mu-maps are not valid files!')
        return None
    # CSA Series Header Info
    if [0x29,0x1020] in dhdr:
        csahdr = dhdr[0x29,0x1020].value
        log.debug('got CSA mu-map info.')
    return csahdr, dhdr

def hmu_shape(hdr):
    #regular expression to find the shape
    p = re.compile(r'(?<=:=)\s*\d{1,4}')
    # x: dim [1]
    i0 = hdr.find('matrix size[1]')
    i1 = i0+hdr[i0:].find('\n')
    u = int(p.findall(hdr[i0:i1])[0])
    # x: dim [2]
    i0 = hdr.find('matrix size[2]')
    i1 = i0+hdr[i0:].find('\n')
    v = int(p.findall(hdr[i0:i1])[0])
    # x: dim [3]
    i0 = hdr.find('matrix size[3]')
    i1 = i0+hdr[i0:].find('\n')
    w = int(p.findall(hdr[i0:i1])[0])
    return w,v,u

def hmu_voxsize(hdr):
    #regular expression to find the shape
    p = re.compile(r'(?<=:=)\s*\d{1,2}[.]\d{1,10}')
    # x: dim [1]
    i0 = hdr.find('scale factor (mm/pixel) [1]')
    i1 = i0+hdr[i0:].find('\n')
    vx = float(p.findall(hdr[i0:i1])[0])
    # x: dim [2]
    i0 = hdr.find('scale factor (mm/pixel) [2]')
    i1 = i0+hdr[i0:].find('\n')
    vy = float(p.findall(hdr[i0:i1])[0])
    # x: dim [3]
    i0 = hdr.find('scale factor (mm/pixel) [3]')
    i1 = i0+hdr[i0:].find('\n')
    vz = float(p.findall(hdr[i0:i1])[0])
    return np.array([0.1*vz, 0.1*vy, 0.1*vx])

def hmu_origin(hdr):
    #regular expression to find the origin
    p = re.compile(r'(?<=:=)\s*\d{1,5}[.]\d{1,10}')
    # x: dim [1]
    i0 = hdr.find('$umap origin (pixels) [1]')
    i1 = i0+hdr[i0:].find('\n')
    x = float(p.findall(hdr[i0:i1])[0])
    # x: dim [2]
    i0 = hdr.find('$umap origin (pixels) [2]')
    i1 = i0+hdr[i0:].find('\n')
    y = float(p.findall(hdr[i0:i1])[0])
    # x: dim [3]
    i0 = hdr.find('$umap origin (pixels) [3]')
    i1 = i0+hdr[i0:].find('\n')
    z = -float(p.findall(hdr[i0:i1])[0])
    return np.array([z, y, x])

def hmu_offset(hdr):
    #regular expression to find the origin
    p = re.compile(r'(?<=:=)\s*\d{1,5}[.]\d{1,10}')
    if hdr.find('$origin offset')>0:
        # x: dim [1]
        i0 = hdr.find('$origin offset (mm) [1]')
        i1 = i0+hdr[i0:].find('\n')
        x = float(p.findall(hdr[i0:i1])[0])
        # x: dim [2]
        i0 = hdr.find('$origin offset (mm) [2]')
        i1 = i0+hdr[i0:].find('\n')
        y = float(p.findall(hdr[i0:i1])[0])
        # x: dim [3]
        i0 = hdr.find('$origin offset (mm) [3]')
        i1 = i0+hdr[i0:].find('\n')
        z = -float(p.findall(hdr[i0:i1])[0])
        return np.array([0.1*z, 0.1*y, 0.1*x])
    else:
        return np.array([0.0, 0.0, 0.0])

def rd_hmu(fh):
    #--read hdr file--
    f = open(fh, 'r')
    hdr = f.read()
    f.close()
    #-----------------
    #regular expression to find the file name
    p = re.compile(r'(?<=:=)\s*\w*[.]\w*')
    i0 = hdr.find('!name of data file')
    i1 = i0+hdr[i0:].find('\n')
    fbin = p.findall(hdr[i0:i1])[0]
    #--read img file--
    f = open(os.path.join(os.path.dirname(fh), fbin.strip()), 'rb')
    im = np.fromfile(f, np.float32)
    f.close()
    #-----------------
    return  hdr, im

def get_hmupos(datain, parts, Cnt, outpath=''):
    log = logging.getLogger(__name__)

    # check if registration executable exists
    if not os.path.isfile(Cnt['RESPATH']):
        log.error('no registration executable found!')
        sys.exit()

    #----- get positions from the DICOM list-mode file -----
    ihdr, csainfo = mmraux.hdr_lm(datain, Cnt)
    #table position origin
    fi = csainfo.find('TablePositionOrigin')
    tpostr = csainfo[fi:fi+200]
    tpo = re.sub(r'[^a-zA-Z0-9\-\.]', '', tpostr).split('M')
    tpozyx = np.array([float(tpo[-1]), float(tpo[-2]), float(tpo[-3])]) / 10
    log.debug('table position (z,y,x) (cm):%r' % tpozyx)
    #--------------------------------------------------------

    #------- get positions from the DICOM mu-map file -------
    csamu, dhdr = hdr_mu(datain, Cnt)
    #> get the indices where the table offset may reside:
    idxs = [m.start() for m in re.finditer(b'GantryTableHomeOffset(?!_)', csamu)]
    #> loop over the indices and find those which are correct
    found_off = False
    for i in idxs:
        gtostr1  = csamu[ i:i+300 ]
        gtostr2 = re.sub(b'[^a-zA-Z0-9\-\.]', b'', gtostr1)
        # gantry table offset, through conversion of string to float
        gtoxyz = re.findall(b'(?<=M)-*[\d]{1,4}\.[\d]{6,9}', gtostr2)
        gtozyx = np.float32(gtoxyz)[::-1]/10
        if len(gtoxyz)>3:
            log.warning('the gantry table offset got more than 3 entries detected--check needed.')
            gtozyx = gtozyx[-3:]
        if abs(gtozyx[0])>20 and abs(gtozyx[1])<20 and abs(gtozyx[2])<2:
            found_off = True
            break

    if found_off:
        log.debug('gantry table offset (z,y,x) (cm): {}'.format(gtozyx))
    else:
        raise ValueError('Could not find the gantry table offset or the offset is unusual.')
     #--------------------------------------------------------

    # #------- get positions from the DICOM mu-map file -------
    # csamu, dhdr = hdr_mu(datain, Cnt)
    # tmp = re.search('GantryTableHomeOffset(?!_)', csamu)
    # gtostr1  = csamu[ tmp.start():tmp.start()+300 ]
    # gtostr2 = re.sub(r'[^a-zA-Z0-9\-\.]', '', gtostr1)
    # # gantry table offset, through conversion of string to float
    # gtoxyz = re.findall(r'(?<=M)-*[\d]{1,4}\.[\d]{6,9}', gtostr2)
    # gtozyx = np.float32(gtoxyz)[::-1]/10
    # #--------------------------------------------------------

    # log.debug('gantry table offset (z,y,x) (cm):%r' % gtozyx)

    ## ----
    ## old II
    # csamu, dhdr = nipet.img.mmrimg.hdr_mu(datain, Cnt)
    # tmp = re.search('GantryTableHomeOffset(?!_)', csamu)
    # gtostr = csamu[ tmp.start():tmp.start()+300 ]
    # gto = re.sub(r'[^a-zA-Z0-9\-\.]', '', gtostr).split('M')
    # # get the first three numbers
    # zyx = np.zeros(3, dtype=np.float32)
    # c = 0
    # for i in range(len(gto)):
    #     if re.search(r'[\d]{1,3}\.[\d]{6}', gto[i])!=None and c<3:
    #         zyx[c] = np.float32(re.sub(r'[^0-9\-\.]', '', gto[i]))
    #         c+=1
    # #gantry table offset
    # gtozyx = zyx[::-1]/10
    ## ----

    ## ----
    ## old I: only worked for syngo MR B20P
    # fi = csamu.find('GantryTableHomeOffset')
    # gtostr =csamu[fi:fi+300]
    # if dhdr[0x0018, 0x1020].value == 'syngo MR B20P':
    #     gto = re.sub(r'[^a-zA-Z0-9\-\.]', '', gtostr).split('M')
    #     # get the first three numbers
    #     zyx = np.zeros(3, dtype=np.float32)
    #     c = 0
    #     for i in range(len(gto)):
    #         if re.search(r'[\d]', gto[i])!=None and c<3:
    #             zyx[c] = np.float32(re.sub(r'[^0-9\-\.]', '', gto[i]))
    #             c+=1
    #     #gantry table offset
    #     gtozyx = zyx[::-1]/10
    #     log.debug('gantry table offset (z,y,x) (cm):%r' % gtozyx)
    # # older scanner version
    # elif dhdr[0x0018, 0x1020].value == 'syngo MR B18P':
    #     zyx = np.zeros(3, dtype=np.float32)
    #     for k in range(3):
    #         tmp = re.search(r'\{\s*[\-0-9.]*\s*\}', gtostr)
    #         i0 = tmp.start()
    #         i1 = tmp.end()
    #         if gtostr[i0+1:i1-1]!=' ':  zyx[k] = np.float32(gtostr[i0+1:i1-1])
    #         gtostr = gtostr[i1:]
    #     #gantry table offset
    #     gtozyx = zyx[::-1]/10
    #     log.debug('gantry table offset (z,y,x) (cm):%r' % gtozyx)
    ## -----

    # create the folder for hardware mu-maps
    if outpath=='':
        dirhmu = os.path.join( datain['corepath'], 'mumap-hdw')
    else:
        dirhmu = os.path.join( outpath, 'mumap-hdw')
    mmraux.create_dir(dirhmu)
    # get the reference nii image
    fref = os.path.join(dirhmu, 'hmuref.nii.gz')

    #start horizontal bed position
    p = re.compile(r'start horizontal bed position.*\d{1,3}\.*\d*')
    m = p.search(ihdr)
    fi = ihdr[m.start():m.end()].find('=')
    hbedpos = 0.1*float(ihdr[m.start()+fi+1:m.end()])

    #start vertical bed position
    p = re.compile(r'start vertical bed position.*\d{1,3}\.*\d*')
    m = p.search(ihdr)
    fi = ihdr[m.start():m.end()].find('=')
    vbedpos = 0.1*float(ihdr[m.start()+fi+1:m.end()])

    log.debug('creating reference nii image for resampling')
    B = np.diag(np.array([-10*Cnt['SO_VXX'], 10*Cnt['SO_VXY'], 10*Cnt['SO_VXZ'], 1]))
    B[0,3] = 10*(.5*Cnt['SO_IMX'])*Cnt['SO_VXX']
    B[1,3] = 10*( -.5*Cnt['SO_IMY']+1)*Cnt['SO_VXY']
    B[2,3] = 10*((-.5*Cnt['SO_IMZ']+1)*Cnt['SO_VXZ'] + hbedpos )
    nimpa.array2nii(  np.zeros((Cnt['SO_IMZ'], Cnt['SO_IMY'], Cnt['SO_IMX']), dtype=np.float32), B, fref)

    #define a dictionary of all positions/offsets of hardware mu-maps
    hmupos = [None]*5
    hmupos[0] = {
        'TabPosOrg' :   tpozyx, #from DICOM of LM file
        'GanTabOff' :   gtozyx, #from DICOM of mMR mu-map file
        'HBedPos'   :   hbedpos, #from Interfile of LM file [cm]
        'VBedPos'   :   vbedpos, #from Interfile of LM file [cm]
        'niipath'   :   fref
        }

    #--------------------------------------------------------------------------
    # iteratively go through the mu-maps and add them as needed
    for i in parts:
        fh = os.path.join(Cnt['HMUDIR'], Cnt['HMULIST'][i-1])
        # get the interfile header and binary data
        hdr, im = rd_hmu(fh)
        #get shape, origin, offset and voxel size
        s = hmu_shape(hdr)
        im.shape = s
        # get the origin, offset and voxel size for the mu-map interfile data
        org = hmu_origin(hdr)
        off = hmu_offset(hdr)
        vs  = hmu_voxsize(hdr)
        # corner voxel position for the interfile image data
        vpos = (-org*vs + off + gtozyx - tpozyx)
        #add to the dictionary
        hmupos[i] = {
            'vpos'    :   vpos,
            'shape'   :   s,   #from interfile
            'iorg'    :   org, #from interfile
            'ioff'    :   off, #from interfile
            'ivs'     :   vs,  #from interfile
            'img'     :   im, #from interfile
            'niipath' :   os.path.join(dirhmu, '_'+Cnt['HMULIST'][i-1].split('.')[0]+'.nii.gz')
        }
        #save to NIfTI
        log.debug('creating mu-map for:%r' % Cnt['HMULIST'][i-1])
        A = np.diag(np.append(10*vs[::-1], 1))
        A[0,0] *= -1
        A[0,3] =  10*(-vpos[2])
        A[1,3] = -10*((s[1]-1)*vs[1] + vpos[1])
        A[2,3] = -10*((s[0]-1)*vs[0] - vpos[0])
        nimpa.array2nii(im[::-1,::-1,:], A, hmupos[i]['niipath'])

        # resample using nify.reg
        fout = os.path.join(    os.path.dirname (hmupos[0]['niipath']),
                                'r'+os.path.basename(hmupos[i]['niipath']).split('.')[0]+'.nii.gz' )
        cmd = [ Cnt['RESPATH'],
                '-ref', hmupos[0]['niipath'],
                '-flo', hmupos[i]['niipath'],
                '-res', fout,
                '-pad', '0']
        if log.getEffectiveLevel() > logging.DEBUG:
            cmd.append('-voff')
        call(cmd)

    return hmupos

#-------------------------------------------------------------------------------------
def hdw_mumap(
        datain,
        hparts,
        params,
        outpath='',
        use_stored=False,
        del_interm=True):
    '''Get hardware mu-map components, including bed, coils etc.'''
    log = logging.getLogger(__name__)

    # two ways of passing Cnt are here decoded
    if 'Cnt' in params:
        Cnt = params['Cnt']
    else:
        Cnt = params

    if outpath!='':
        fmudir = os.path.join(outpath, 'mumap-hdw')
    else:
        fmudir = os.path.join(datain['corepath'], 'mumap-hdw')

    # if requested to use the stored hardware mu_map get it from the path in datain
    if 'hmumap' in datain and os.path.isfile(datain['hmumap']) and use_stored:

        if datain['hmumap'].endswith(('.nii', '.nii.gz')):
            dct = nimpa.getnii(datain['hmumap'], output='all')
            hmu = dct['im']
            A = dct['affine']
            fmu = datain['hmumap']
        elif datain['hmumap'].endswith(('.npy')):
            hmu, A, fmu = np.load(datain['hmumap'], allow_pickle=True)
            log.debug('loaded hardware mu-map from file:' + datain['hmumap'])
            fnp = datain['hmumap']

    elif outpath!='' and os.path.isfile(os.path.join(fmudir, 'hmumap.npy')):
        fnp = os.path.join(fmudir, 'hmumap.npy')
        hmu, A, fmu = np.load(fnp, allow_pickle=True)
        datain['hmumap'] = fnp
    # otherwise generate it from the parts through resampling the high resolution CT images
    else:
        hmupos = get_hmupos(datain, hparts, Cnt, outpath=outpath)
        # just to get the dims, get the ref image
        nimo = nib.load(hmupos[0]['niipath'])
        A = nimo.get_sform()
        imo = np.float32( nimo.get_data() )
        imo[:] = 0

        for i in hparts:
            fin  = os.path.join(os.path.dirname (hmupos[0]['niipath']),
                                'r'+os.path.basename(hmupos[i]['niipath']).split('.')[0]+'.nii.gz' )
            nim = nib.load(fin)
            mu = nim.get_data()
            mu[mu<0] = 0

            imo += mu

        hdr = nimo.header
        hdr['cal_max'] = np.max(imo)
        hdr['cal_min'] = np.min(imo)
        fmu  = os.path.join(os.path.dirname (hmupos[0]['niipath']), 'hardware_umap.nii.gz' )
        nib.save(nimo, fmu)

        hmu = np.transpose(imo[:,::-1,::-1], (2, 1, 0))

        # save the objects to numpy arrays
        fnp = os.path.join(fmudir, 'hmumap.npy')
        np.save(fnp, (hmu, A, fmu))
        #update the datain dictionary (assuming it is mutable)
        datain['hmumap'] = fnp

        if del_interm:
            for fname in glob.glob(os.path.join(fmudir, '_*.nii*')):
                os.remove(fname)
            for fname in glob.glob(os.path.join(fmudir, 'r_*.nii*')):
                os.remove(fname)

    #return image dictionary with the image itself and some other stats
    hmu_dct = { 'im':hmu,
                'fim':fmu,
                'affine':A}
    if 'fnp' in locals():
        hmu_dct['fnp'] = fnp

    return hmu_dct

def rmumaps(datain, Cnt, t0=0, t1=0, use_stored=False):
    '''
    get the mu-maps for hardware and object and trim it axially for reduced rings case
    '''
    log = logging.getLogger(__name__)

    from niftypet.nipet.lm  import mmrhist
    from niftypet.nipet.prj import mmrrec

    fcomment = '(R)'

    # get hardware mu-map
    if os.path.isfile(datain['hmumap']) and use_stored:
        muh, _ = np.load(datain['hmumap'], allow_pickle=True)
        log.debug('loaded hardware mu-map from file:' + datain['hmumap'])
    else:
        hmudic = hdw_mumap(datain, [1,2,4], Cnt)
        muh = hmudic['im']

    # get pCT mu-map if stored in numpy file and then exit, otherwise do all the processing
    if os.path.isfile(datain['mumapCT']) and use_stored:
        mup, _ = np.load(datain['mumapCT'], allow_pickle=True)
        muh = muh[2*Cnt['RNG_STRT'] : 2*Cnt['RNG_END'], :, :]
        mup = mup[2*Cnt['RNG_STRT'] : 2*Cnt['RNG_END'], :, :]
        return [muh, mup]

    # get UTE object mu-map (may be not in register with the PET data)
    if os.path.isfile(datain['mumapUTE']) and use_stored:
        muo, _ = np.load(datain['mumapUTE'], allow_pickle=True)
    else:
        mudic = obj_mumap(datain, Cnt, store=True)
        muo = mudic['im']

    if os.path.isfile(datain['pCT']):
        # reconstruct PET image with default settings to be used to alight pCT mu-map
        params = mmraux.mMR_params()
        Cnt_ = params['Cnt']
        txLUT_ = params['txLUT']
        axLUT_ = params['axLUT']

        # histogram for reconstruction with UTE mu-map
        hst = mmrhist.hist(datain, txLUT_, axLUT_, Cnt_, t0=t0, t1=t1)
        # reconstruct PET image with UTE mu-map to which co-register T1w
        recute = mmrrec.osemone(
            datain, [muh, muo], hst, params,
            recmod=3, itr=4, fwhm=0., store_img=True, fcomment=fcomment+'_QNT-UTE'
        )
        # --- MR T1w
        if os.path.isfile(datain['T1nii']):
            ft1w = datain['T1nii']
        elif os.path.isfile(datain['T1bc']):
            ft1w = datain['T1bc']
        elif os.path.isdir(datain['MRT1W']):
            # create file name for the converted NIfTI image
            fnii = 'converted'
            call( [ Cnt['DCM2NIIX'], '-f', fnii, datain['T1nii'] ] )
            ft1nii = glob.glob( os.path.join(datain['T1nii'], '*converted*.nii*') )
            ft1w = ft1nii[0]
        else:
            log.error('disaster: no T1w image!')
            sys.exit()

        #output for the T1w in register with PET
        ft1out = os.path.join(os.path.dirname(ft1w), 'T1w_r'+'.nii.gz')
        #text file fo rthe affine transform T1w->PET
        faff   = os.path.join(os.path.dirname(ft1w), fcomment+'mr2pet_affine'+'.txt')  #time.strftime('%d%b%y_%H.%M',time.gmtime())
        #call the registration routine
        if os.path.isfile( Cnt['REGPATH'] ):
            cmd = [Cnt['REGPATH'],
                 '-ref', recute.fpet,
                 '-flo', ft1w,
                 '-rigOnly', '-speeeeed',
                 '-aff', faff,
                 '-res', ft1out]
            if log.getEffectiveLevel() > logging.DEBUG:
                cmd.append('-voff')
            call(cmd)
        else:
            log.error('path to registration executable is incorrect!')
            sys.exit()

        #get the pCT mu-map with the above faff
        pmudic = pct_mumap(datain, txLUT, axLUT, Cnt, faff=faff, fpet=recute.fpet, fcomment=fcomment)
        mup = pmudic['im']

        muh = muh[2*Cnt['RNG_STRT'] : 2*Cnt['RNG_END'], :, :]
        mup = mup[2*Cnt['RNG_STRT'] : 2*Cnt['RNG_END'], :, :]
        return [muh, mup]
    else:
        muh = muh[2*Cnt['RNG_STRT'] : 2*Cnt['RNG_END'], :, :]
        muo = muo[2*Cnt['RNG_STRT'] : 2*Cnt['RNG_END'], :, :]
        return [muh, muo]

#------------------------------------------------------------------------


'''This script is to be used solely for the purpose of addressing the
labeling job inconsistencies for Ground Truth semantic segmentation
tasks between June 4, 2020, and September 4, 2020.
'''

# Usage

'''
python sagemaker_ground_truth_semantic_segmentation_mask_index_correction.py --labeling-job-name example-job-name

`example-job-name` is an Amazon SageMaker Ground Truth semantic
segmentation labeling job with a single annotator with inconsistent
annotations (i.e. incorrect label indices in the pngs). This script
will locate the output for the labeling job, download the resulting
annotations into memory, fix the annotations using the label map, and
upload the corrected annotations and a new output manifest with
reference to those annotations to S3.

* The corrected annotations will be saved to an S3 location parallel
  to the original annotations:
  <labeling-job-name>/annotations/consolidated-annotation/output-corrected/
* The corrected manifest will be saved to an S3 location parallel to
  the original annotations:
  <labeling-job-name>/manifests/output/output-corrected.manifest
* The entries in `output-corrected.manifest` will contain a
  `'fixed-label-ref'` field with the S3 URI of the fixed annotations
  and an `'original-label-ref'` field with the original s3 URI of the
  label.
'''

import argparse
import base64
from collections import defaultdict
from io import BytesIO
import json
from pathlib import Path
from urllib.parse import urlparse

import boto3
import numpy as np
from PIL import Image
s3 = boto3.resource('s3')
sagemaker = boto3.client('sagemaker')
try:
    next(iter(s3.buckets.all()))
except Exception as e:
    raise Exception(
        'Could not access your s3 resources. '
        'Please verify that your AWS credentials are correctly configured and '
        'try again.'
    ) from e

class SemSegParser(object):
    def __init__(self, annotation, label_categories):
        self._annotation = annotation
        image_bytes = base64.b64decode(self._annotation)
        img = np.asarray(Image.open(BytesIO(image_bytes)).convert("RGBA"))
        self.hex_colors = defaultdict(str)
        self.hex_colors['BACKGROUND'] = '#ffffff'

        self._img_array, self._class_names, self._label_map = \
            self.get_class_masks(img, label_categories)

    def get_class_masks(self, img, label_categories):
        img_no_alpha = img[:, :, 0:3]
        masks = []
        class_names = ['BACKGROUND']
        rgb_label_maps = self.initialize_label_map()
        for idx_str, category_info in label_categories.items():
            i = int(idx_str)
            class_name = category_info['class-name']
            if class_name == 'BACKGROUND':
                continue
            class_names.append(class_name)
            class_hex_color = category_info['hex-color']
            self.hex_colors[class_name] = class_hex_color
            class_rgb_color = self.hex_to_rgb(class_hex_color)
            rgb_label_maps.append(class_rgb_color)
            class_mask = np.all(img_no_alpha == class_rgb_color, axis=-1)
            class_mask = class_mask * i
            masks.append(class_mask)
        masks = np.array(masks)
        masks = masks.sum(axis=0)

        return masks.astype(np.uint8), class_names, rgb_label_maps

    # Set background to white
    def initialize_label_map(self):
        return [(255, 255, 255)]

    @property
    def class_names(self):
        return self._class_names

    @property
    def img_array(self):
        return self._img_array

    @staticmethod
    def hex_to_rgb(hexcode):
        h = hexcode.lstrip('#')
        return tuple(int(h[i:i + 2], 16) for i in (0, 2, 4))

    @property
    def img_w_palette(self):
        im = Image.fromarray(np.uint8(self._img_array))
        num_classes = len(self._label_map)
        palette = self._label_map + [(255, 255, 255) for i
                                     in range(256 - num_classes)]

        palette = [item for rgb in palette for item in rgb] + ([255, ])
        im.putpalette(palette)
        return im

def get_bucket_and_key(s3uri):
    """Get the bucket name and key associated with an s3 object.

    Args:
      s3uri (str): The s3 uri.

    Return:
      bucket_name and key strings.
    """
    url = urlparse(s3uri)
    bucket_name = url.netloc
    key = url.path.lstrip('/')

    return bucket_name, key

def get_object_bytes(s3_obj):
    """Get bytes for an object stored in s3.

    Arg:
     s3_obj (boto3.resources.factory.s3.ObjectSummary): object for thing in s3
    """
    body = s3_obj.get()['Body']
    return body.read()

def get_metadata(entry, label_attribute_name):
    metadata_key = '{}-metadata'.format(label_attribute_name)
    try:
        metadata = entry[metadata_key]
    except KeyError as e:
        raise KeyError(
            'The metadata_key (derived from the label-attribute-name or '
            'job-name) is missing from this manifest entry. Please specify a '
            'different label-attribute-name.'
        ) from e

    return metadata

def get_output_manifest(labeling_job_name):
    description = sagemaker.describe_labeling_job(
        LabelingJobName=labeling_job_name
    )
    label_attribute_name = description['LabelAttributeName']
    manifest_path = description['LabelingJobOutput']['OutputDatasetS3Uri']

    bucket_name, key = get_bucket_and_key(manifest_path)
    manifest_bytes = get_object_bytes(s3.Bucket(bucket_name).Object(key))

    manifest = []
    for line in manifest_bytes.decode().splitlines():
        manifest.append(json.loads(line))

    return manifest, label_attribute_name, bucket_name, key

def fix_annotations(labeling_job_name):
    manifest, label_attribute_name, output_manifest_bucket_name, \
        output_manifest_key = get_output_manifest(labeling_job_name)

    fixed_manifest = []
    for entry in manifest:
        metadata = get_metadata(entry, label_attribute_name)
        try:
            job_name = metadata['job-name'].replace('labeling-job/', '')
            label_uri = entry[label_attribute_name]

            bucket_name, key = get_bucket_and_key(label_uri)
            key_for_corrected_png = key.replace(
                'consolidated-annotation/output',
                'consolidated-annotation/output-corrected'
            )
            bucket = s3.Bucket(bucket_name)
            annotation_bytes = get_object_bytes(bucket.Object(key))
            annotation_b64 = base64.b64encode(annotation_bytes)
            color_map = metadata['internal-color-map']
            parser = SemSegParser(annotation_b64, color_map)

            png_img = parser.img_w_palette
            with BytesIO() as in_mem_file:
                png_img.save(in_mem_file, format='png')
                in_mem_file.seek(0)
                obj = bucket.Object(key_for_corrected_png)
                obj.upload_fileobj(in_mem_file)

            entry['original-label-ref'] = entry[label_attribute_name]
            entry['corrected-label-ref'] = \
                's3://' + bucket_name + '/' + key_for_corrected_png

        except KeyError:
            continue

        finally:
            fixed_manifest.append(entry)

    # Upload corrected manifest to S3
    with BytesIO() as in_mem_file:
        for line in fixed_manifest:
            line_to_write = json.dumps(line) + '\n'
            in_mem_file.write(bytes(line_to_write.encode('utf-8')))
        in_mem_file.seek(0)

        corrected_manifest_bucket = s3.Bucket(output_manifest_bucket_name)
        corrected_manifest_key = output_manifest_key.replace(
            'output.manifest', 'output-corrected.manifest')
        obj = corrected_manifest_bucket.Object(corrected_manifest_key)
        obj.upload_fileobj(in_mem_file)

def parse_args():
    parser = argparse.ArgumentParser(
        description=(
            'Correct semantic segmentation masks from Amazon SageMaker Ground '
            'Truth.'
        )
    )
    parser.add_argument(
        '--labeling-job-name', type=str, required=True,
        help=(
            'Your labeling job.'
        )
    )

    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()
    print(
        'Fixing annotations from {}'.format(
            args.labeling_job_name
        )
    )
    fix_annotations(args.labeling_job_name)
    print('Done.')


# -*- coding: utf-8 -*-
# --------------------------------------------
# Author: chen hao
# Date: 2021/6/13 17:46
# Description:
# --------------------------------------------
import mkl
import faiss
import pickle
import numpy as np
from typing import List, Dict
from .constant import INDEX_PATH, FEATURES_PATH
mkl.get_max_threads()

class SearchIndex(object):
    def __init__(self, index_path:str=INDEX_PATH, feature_path:str=FEATURES_PATH):

        assert index_path, "index path is not exists"
        assert feature_path, "feature path is not exists"

        self._index = faiss.read_index(index_path)
        self._features = self.load_features(feature_path)
        self._id2label = {v.id:k for k, v in self._features.items()}

    def load_features(self, path:str):
        with open(path, 'rb') as reader:
            features = pickle.load(reader)
        return features

    def search_by_labels(self, labels:List[str], k:int=5) -> List[Dict]:
        vectors = [self._features[label].vector for label in labels if label in self._features]
        ids = [self._features[label].id for label in labels if label in self._features]
        results = self._search(ids, labels, vectors, k+1)
        return results

    def search_by_vectors(self, vectors:List[float], k:int=5) -> List[Dict]:
        ids = [-1] * len(vectors)
        labels = ['null'] * len(vectors)
        results = self._search(ids, labels, vectors, k)
        return results

    def _search(self, ids:List[int], labels:List[str], vectors:List[float], k:int) -> List[Dict]:
        def pack_neighbor(id, score):
            return {'id':int(id), 'label': str(self._id2label[id]), 'score': float(score)}
        def pack_result(id, label, vector, neighbors):
            return {'id':id, 'label': label, 'vector': vector.tolist(), 'neighbors': neighbors}

        results = []
        vectors = [np.array(vec, dtype=np.float32) for vec in vectors]
        vectors = np.atleast_2d(vectors)

        scores, neighbors = self._index.search(vectors, k) if vectors.size > 0 else ([], [])

        for id, label, vector, score, neighbor in zip(ids, labels, vectors, scores, neighbors):
            neighbor_score = zip(neighbor, score)
            neighbor_score = [(n_id, n_score) for n_id, n_score in neighbor_score if n_id != id and n_id != -1]
            neighbor_score = [pack_neighbor(n_id, n_score) for n_id, n_score in neighbor_score]
            results.append(pack_result(id, label, vector, neighbor_score))

        return results


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import numpy as np
import time
import functools
import paddle
import paddle.fluid as fluid
from paddle.fluid.framework import IrGraph
from paddle.fluid.contrib.slim.quantization import QuantizationTransformPass
from paddle.fluid.contrib.slim.quantization import QuantizationFreezePass
from paddle.fluid.contrib.slim.quantization import ConvertToInt8Pass
from paddle.fluid.contrib.slim.quantization import TransformForMobilePass
from paddle.fluid import core
import argparse
import subprocess
import sys
sys.path.append('..')
import reader
import models
from utility import add_arguments, print_arguments
from utility import save_persistable_nodes, load_persistable_nodes

parser = argparse.ArgumentParser(description=__doc__)
add_arg = functools.partial(add_arguments, argparser=parser)
# yapf: disable
add_arg('batch_size',       int,   256,                  "Minibatch size.")
add_arg('use_gpu',          bool,  True,                 "Whether to use GPU or not.")
add_arg('total_images',     int,   1281167,              "Training image number.")
add_arg('num_epochs',       int,   120,                  "number of epochs.")
add_arg('class_dim',        int,   1000,                 "Class number.")
add_arg('image_shape',      str,   "3,224,224",          "input image size")
add_arg('model_save_dir',   str,   "output",             "model save directory")
add_arg('pretrained_fp32_model', str,   None,            "Whether to use the pretrained float32 model to initialize the weights.")
add_arg('checkpoint',       str,   None,                 "Whether to resume the training process from the checkpoint.")
add_arg('lr',               float, 0.1,                  "set learning rate.")
add_arg('lr_strategy',      str,   "piecewise_decay",    "Set the learning rate decay strategy.")
add_arg('model',            str,   "SE_ResNeXt50_32x4d", "Set the network to use.")
add_arg('data_dir',         str,   "./data/ILSVRC2012",  "The ImageNet dataset root dir.")
add_arg('act_quant_type',   str,   "abs_max",            "quantization type for activation, valid value:'abs_max','range_abs_max', 'moving_average_abs_max'" )
add_arg('wt_quant_type',    str,   "abs_max",            "quantization type for weight, valid value:'abs_max','channel_wise_abs_max'" )
# yapf: enabl

def optimizer_setting(params):
    ls = params["learning_strategy"]
    if ls["name"] == "piecewise_decay":
        if "total_images" not in params:
            total_images = 1281167
        else:
            total_images = params["total_images"]
        batch_size = ls["batch_size"]
        step = int(total_images / batch_size + 1)

        bd = [step * e for e in ls["epochs"]]
        print("decay list:{}".format(bd))
        base_lr = params["lr"]
        lr = []
        lr = [base_lr * (0.1**i) for i in range(len(bd) + 1)]
        optimizer = fluid.optimizer.Momentum(
            learning_rate=fluid.layers.piecewise_decay(
                boundaries=bd, values=lr),
            momentum=0.9,
            regularization=fluid.regularizer.L2Decay(1e-4))

    elif ls["name"] == "cosine_decay":
        if "total_images" not in params:
            total_images = 1281167
        else:
            total_images = params["total_images"]

        batch_size = ls["batch_size"]
        step = int(total_images / batch_size + 1)

        lr = params["lr"]
        num_epochs = params["num_epochs"]

        optimizer = fluid.optimizer.Momentum(
            learning_rate=fluid.layers.cosine_decay(
                learning_rate=lr, step_each_epoch=step, epochs=num_epochs),
            momentum=0.9,
            regularization=fluid.regularizer.L2Decay(4e-5))
    elif ls["name"] == "exponential_decay":
        if "total_images" not in params:
            total_images = 1281167
        else:
            total_images = params["total_images"]
        batch_size = ls["batch_size"]
        step = int(total_images / batch_size +1)
        lr = params["lr"]
        num_epochs = params["num_epochs"]
        learning_decay_rate_factor=ls["learning_decay_rate_factor"]
        num_epochs_per_decay = ls["num_epochs_per_decay"]
        NUM_GPUS = 1

        optimizer = fluid.optimizer.Momentum(
            learning_rate=fluid.layers.exponential_decay(
                learning_rate = lr * NUM_GPUS,
                decay_steps = step * num_epochs_per_decay / NUM_GPUS,
                decay_rate = learning_decay_rate_factor),
            momentum=0.9,

            regularization = fluid.regularizer.L2Decay(4e-5))

    else:
        lr = params["lr"]
        optimizer = fluid.optimizer.Momentum(
            learning_rate=lr,
            momentum=0.9,
            regularization=fluid.regularizer.L2Decay(1e-4))

    return optimizer

def net_config(image, label, model, args):
    model_list = [m for m in dir(models) if "__" not in m]
    assert args.model in model_list,"{} is not lists: {}".format(
        args.model, model_list)

    class_dim = args.class_dim
    model_name = args.model

    if model_name == "GoogleNet":
        out0, out1, out2 = model.net(input=image, class_dim=class_dim)
        cost0 = fluid.layers.cross_entropy(input=out0, label=label)
        cost1 = fluid.layers.cross_entropy(input=out1, label=label)
        cost2 = fluid.layers.cross_entropy(input=out2, label=label)
        avg_cost0 = fluid.layers.mean(x=cost0)
        avg_cost1 = fluid.layers.mean(x=cost1)
        avg_cost2 = fluid.layers.mean(x=cost2)

        avg_cost = avg_cost0 + 0.3 * avg_cost1 + 0.3 * avg_cost2
        acc_top1 = fluid.layers.accuracy(input=out0, label=label, k=1)
        acc_top5 = fluid.layers.accuracy(input=out0, label=label, k=5)
        out = out2
    else:
        out = model.net(input=image, class_dim=class_dim)
        cost = fluid.layers.cross_entropy(input=out, label=label)

        avg_cost = fluid.layers.mean(x=cost)
        acc_top1 = fluid.layers.accuracy(input=out, label=label, k=1)
        acc_top5 = fluid.layers.accuracy(input=out, label=label, k=5)

    return out, avg_cost, acc_top1, acc_top5

def build_program(is_train, main_prog, startup_prog, args):
    image_shape = [int(m) for m in args.image_shape.split(",")]
    model_name = args.model
    model_list = [m for m in dir(models) if "__" not in m]
    assert model_name in model_list, "{} is not in lists: {}".format(args.model,
                                                                     model_list)
    model = models.__dict__[model_name]()
    with fluid.program_guard(main_prog, startup_prog):
        py_reader = fluid.layers.py_reader(
            capacity=16,
            shapes=[[-1] + image_shape, [-1, 1]],
            lod_levels=[0, 0],
            dtypes=["float32", "int64"],
            use_double_buffer=True)
        with fluid.unique_name.guard():
            image, label = fluid.layers.read_file(py_reader)
            out, avg_cost, acc_top1, acc_top5 = net_config(image, label, model, args)
            avg_cost.persistable = True
            acc_top1.persistable = True
            acc_top5.persistable = True
            if is_train:
                params = model.params
                params["total_images"] = args.total_images
                params["lr"] = args.lr
                params["num_epochs"] = args.num_epochs
                params["learning_strategy"]["batch_size"] = args.batch_size
                params["learning_strategy"]["name"] = args.lr_strategy

                optimizer = optimizer_setting(params)
                optimizer.minimize(avg_cost)
                global_lr = optimizer._global_learning_rate()
    if is_train:
        return image, out, py_reader, avg_cost, acc_top1, acc_top5, global_lr
    else:
        return image, out, py_reader, avg_cost, acc_top1, acc_top5

def train(args):
    # parameters from arguments
    model_name = args.model
    pretrained_fp32_model = args.pretrained_fp32_model
    checkpoint = args.checkpoint
    model_save_dir = args.model_save_dir
    data_dir = args.data_dir
    activation_quant_type = args.act_quant_type
    weight_quant_type = args.wt_quant_type
    print("Using %s as the actiavtion quantize type." % activation_quant_type)
    print("Using %s as the weight quantize type." % weight_quant_type)

    startup_prog = fluid.Program()
    train_prog = fluid.Program()
    test_prog = fluid.Program()

    _, _, train_py_reader, train_cost, train_acc1, train_acc5, global_lr = build_program(
        is_train=True,
        main_prog=train_prog,
        startup_prog=startup_prog,
        args=args)
    image, out, test_py_reader, test_cost, test_acc1, test_acc5 = build_program(
        is_train=False,
        main_prog=test_prog,
        startup_prog=startup_prog,
        args=args)
    test_prog = test_prog.clone(for_test=True)

    place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
    exe = fluid.Executor(place)
    exe.run(startup_prog)
    main_graph = IrGraph(core.Graph(train_prog.desc), for_test=False)
    test_graph = IrGraph(core.Graph(test_prog.desc), for_test=True)

    if pretrained_fp32_model:
        def if_exist(var):
            return os.path.exists(os.path.join(pretrained_fp32_model, var.name))
        fluid.io.load_vars(
            exe, pretrained_fp32_model, main_program=train_prog, predicate=if_exist)

    if args.use_gpu:
        visible_device = os.getenv('CUDA_VISIBLE_DEVICES')
        if visible_device:
            device_num = len(visible_device.split(','))
        else:
            device_num = subprocess.check_output(
                ['nvidia-smi', '-L']).decode().count('\n')
    else:
        device_num = 1

    train_batch_size = args.batch_size / device_num
    test_batch_size = 1 if activation_quant_type == 'abs_max' else 8
    train_reader = paddle.batch(
        reader.train(data_dir=data_dir), batch_size=train_batch_size, drop_last=True)
    test_reader = paddle.batch(reader.val(data_dir=data_dir), batch_size=test_batch_size)

    train_py_reader.decorate_paddle_reader(train_reader)
    test_py_reader.decorate_paddle_reader(test_reader)

    train_fetch_list = [train_cost.name, train_acc1.name, train_acc5.name, global_lr.name]
    test_fetch_list = [test_cost.name, test_acc1.name, test_acc5.name]

    # 1. Make some quantization transforms in the graph before training and testing.
    # According to the weight and activation quantization type, the graph will be added
    # some fake quantize operators and fake dequantize operators.
    transform_pass = QuantizationTransformPass(
        scope=fluid.global_scope(), place=place,
        activation_quantize_type=activation_quant_type,
        weight_quantize_type=weight_quant_type)
    transform_pass.apply(main_graph)
    transform_pass.apply(test_graph)

    if checkpoint:
        load_persistable_nodes(exe, checkpoint, main_graph)

    build_strategy = fluid.BuildStrategy()
    build_strategy.memory_optimize = False
    build_strategy.enable_inplace = False
    binary = fluid.CompiledProgram(main_graph.graph).with_data_parallel(
        loss_name=train_cost.name, build_strategy=build_strategy)
    test_prog = test_graph.to_program()
    params = models.__dict__[args.model]().params
    for pass_id in range(params["num_epochs"]):

        train_py_reader.start()

        train_info = [[], [], []]
        test_info = [[], [], []]
        train_time = []
        batch_id = 0
        try:
            while True:
                t1 = time.time()
                loss, acc1, acc5, lr = exe.run(binary, fetch_list=train_fetch_list)
                t2 = time.time()
                period = t2 - t1
                loss = np.mean(np.array(loss))
                acc1 = np.mean(np.array(acc1))
                acc5 = np.mean(np.array(acc5))
                train_info[0].append(loss)
                train_info[1].append(acc1)
                train_info[2].append(acc5)
                lr = np.mean(np.array(lr))
                train_time.append(period)
                if batch_id % 10 == 0:
                    print("Pass {0}, trainbatch {1}, loss {2}, \
                        acc1 {3}, acc5 {4}, lr {5}, time {6}"
                          .format(pass_id, batch_id, loss, acc1, acc5, "%.6f" %
                                  lr, "%2.2f sec" % period))
                    sys.stdout.flush()
                batch_id += 1
        except fluid.core.EOFException:
            train_py_reader.reset()

        train_loss = np.array(train_info[0]).mean()
        train_acc1 = np.array(train_info[1]).mean()
        train_acc5 = np.array(train_info[2]).mean()

        test_py_reader.start()

        test_batch_id = 0
        try:
            while True:
                t1 = time.time()
                loss, acc1, acc5 = exe.run(program=test_prog,
                                           fetch_list=test_fetch_list)
                t2 = time.time()
                period = t2 - t1
                loss = np.mean(loss)
                acc1 = np.mean(acc1)
                acc5 = np.mean(acc5)
                test_info[0].append(loss)
                test_info[1].append(acc1)
                test_info[2].append(acc5)
                if test_batch_id % 10 == 0:
                    print("Pass {0},testbatch {1},loss {2}, \
                        acc1 {3},acc5 {4},time {5}"
                          .format(pass_id, test_batch_id, loss, acc1, acc5,
                                  "%2.2f sec" % period))
                    sys.stdout.flush()
                test_batch_id += 1
        except fluid.core.EOFException:
            test_py_reader.reset()

        test_loss = np.array(test_info[0]).mean()
        test_acc1 = np.array(test_info[1]).mean()
        test_acc5 = np.array(test_info[2]).mean()

        print("End pass {0}, train_loss {1}, train_acc1 {2}, train_acc5 {3}, "
              "test_loss {4}, test_acc1 {5}, test_acc5 {6}".format(
                  pass_id, train_loss, train_acc1, train_acc5, test_loss,
                  test_acc1, test_acc5))
        sys.stdout.flush()

        save_checkpoint_path = os.path.join(model_save_dir,  model_name, str(pass_id))
        if not os.path.isdir(save_checkpoint_path):
            os.makedirs(save_checkpoint_path)
        save_persistable_nodes(exe, save_checkpoint_path, main_graph)

    model_path = os.path.join(model_save_dir, model_name, args.act_quant_type)
    float_path = os.path.join(model_path, 'float')
    int8_path = os.path.join(model_path, 'int8')
    mobile_path = os.path.join(model_path, 'mobile')
    if not os.path.isdir(model_path):
        os.makedirs(model_path)

    # 2. Freeze the graph after training by adjusting the quantize
    # operators' order for the inference.
    freeze_pass = QuantizationFreezePass(
        scope=fluid.global_scope(),
        place=place,
        weight_quantize_type=weight_quant_type)
    freeze_pass.apply(test_graph)
    server_program = test_graph.to_program()
    fluid.io.save_inference_model(
        dirname=float_path,
        feeded_var_names=[image.name],
        target_vars=[out], executor=exe,
        main_program=server_program)

    # 3. Convert the weights into int8_t type.
    # (This step is optional.)
    convert_int8_pass = ConvertToInt8Pass(scope=fluid.global_scope(), place=place)
    convert_int8_pass.apply(test_graph)
    server_int8_program = test_graph.to_program()
    fluid.io.save_inference_model(
        dirname=int8_path,
        feeded_var_names=[image.name],
        target_vars=[out], executor=exe,
        main_program=server_int8_program)

    # 4. Convert the freezed graph for paddle-mobile execution.
    # (This step is optional.)
    mobile_pass = TransformForMobilePass()
    mobile_pass.apply(test_graph)
    mobile_program = test_graph.to_program()
    fluid.io.save_inference_model(
        dirname=mobile_path,
        feeded_var_names=[image.name],
        target_vars=[out], executor=exe,
        main_program=mobile_program)

def main():
    args = parser.parse_args()
    print_arguments(args)
    train(args)

if __name__ == '__main__':
    main()


import os
from glob import glob
import numpy as np, pandas as pd
from numpy import array as nparr
from cdips.utils import today_YYYYMMDD

from astrobase import imageutils as iu

from astroquery.vizier import Vizier
from astroquery.eso import Eso

import astropy.units as u
from astropy.coordinates import SkyCoord
from astropy.time import Time

def get_rv_xmatch(ra, dec, G_mag=None, dr2_sourceid=None):
    """
    Given a single RA, dec, and optionally Gaia G mag, and DR2 sourceid, xmatch
    against spectroscopic catalogs to try and mine out an RV and error.

    Returns tuple of:
        (rv, rv_err, provenance)
        Presumably all RVs returned are heliocentric.

    The following spectroscopic surveys did not derive RVs:

        PASTEL
        TESS-HERMES (DR-1)

    The following surveys had non-uniform APIs to access their RVs:

        Gaia-ESO -- needs to be accessed thru ESO's query interface.

        GALAH -- needs to be accessed via
            https://github.com/svenbuder/GALAH_DR2 (or similar)

        APOGEE-2 (DR-15, now) -- also a fits file that gets queried.
            https://www.sdss.org/dr15/irspec/spectro_data/#catalogs.

    The following surveys are queried (fortunately, they're the biggest ones).
    In order of precedence:

        LAMOST DR-4: V/153. Note only the "A, F, G and K type stars catalog
        (4537436 rows)" includes RVs.

        RAVE DR-5: III/279

        Gaia DR2 RVs.
    """
    coord = SkyCoord(ra=ra, dec=dec, unit=(u.deg, u.deg), frame='icrs')
    print('begin vizier search for {}'.format(repr(coord)))

    # the following pattern enables return of all columns desired. e.g.,
    # "Source" for DR2 sourceid.
    v = Vizier(columns=['_RAJ2000', '_DEJ2000', 'HRV', 'e_HRV', 'RV', 'e_RV',
                        'Source'])
    result = v.query_region(coord, radius=10*u.arcsec,
                            catalog=["LAMOST", "RAVE", "GAIA"])

    # LAMOST, RAVE, GaiaDR2
    catkeys = ['V/153/stellar4', 'III/279/rave_dr5', 'I/345/gaia2']
    rv_given_cat = {
        'V/153/stellar4': ['HRV','e_HRV'],
        'III/279/rave_dr5': ['HRV','e_HRV'],
        'I/345/gaia2': ['RV','e_RV']
    }
    name_given_cat = {
        'V/153/stellar4': 'LAMOST-DR4',
        'III/279/rave_dr5': 'RAVE-DR5',
        'I/345/gaia2': 'GAIA-DR2'
    }

    for catkey in catkeys:
        if catkey in result.keys():
            res = result[catkey]
            if len(res) > 1:
                if catkey=='I/345/gaia2':
                    # just use dr2 source ID directly
                    res = res[res['Source'] == dr2_sourceid]
                else:
                    raise NotImplementedError(
                        'what do with multimatch? {}'.format(catkey)
                    )

            rv = float(res[rv_given_cat[catkey][0]])
            rv_err = float(res[rv_given_cat[catkey][1]])
            if not pd.isnull(rv):
                return rv, rv_err, name_given_cat[catkey]

    return np.nan, np.nan, np.nan

def wrangle_eso_for_rv_availability(ra, dec):
    """
    Checks via ESO query for available RVs on:
        ['HARPS', 'ESPRESSO', 'FORS2', 'UVES', 'XSHOOTER']

    Possible future expansion: actually get the RVs. (For now, just this is
    just used as a flag to let the user know the RVs might exist!)

    Returns tuple of:
        (nan, nan, provenance)
    """
    eso = Eso()
    eso.ROW_LIMIT = 9999

    coord = SkyCoord(ra=ra, dec=dec, unit=(u.deg, u.deg), frame='icrs')
    print('begin ESO search for {}'.format(repr(coord)))

    rastr = (
        str(coord.ra.to_string(u.hour)).
        replace('h',' ').replace('m',' ').replace('s',' ')
    )

    decstr = (
        str(coord.dec.to_string()).
        replace('d',' ').replace('m',' ').replace('s',' ')
    )

    # search within 10 arcsec of given position
    boxsize = '00 00 10'
    res = eso.query_main(
        column_filters={'ra':rastr, 'dec':decstr, 'box':boxsize}
    )

    if res is None:
        return np.nan, np.nan, np.nan

    # limit search to the following instruments, in order of preference
    instruments = ['HARPS', 'ESPRESSO', 'FORS2', 'UVES', 'XSHOOTER']
    sel = np.zeros((len(res))).astype(bool)
    for instrument in instruments:
        sel |= (nparr(res['Instrument']) == instrument)
    res = res[sel]

    # limit returned cateogires
    badcategories = ['CALIB']
    sel = np.zeros((len(res))).astype(bool)
    for badcategory in badcategories:
        sel |= (nparr(res['Category']) != badcategory)
    res = res[sel]

    if len(res) >= 1:

        # XSHOOTER doesn't seem to give archival RVs. would need to derive
        # from spectra yourself
        if np.all(nparr(res['Instrument']) == 'XSHOOTER'):
            return np.nan, np.nan, 'XSHOOTER'

        # Embargo lasts a year on all ESO observations.
        nt = Time.now()
        embargo_end = nt.mjd - 365
        if np.all(nparr(res['MJD-OBS']) > embargo_end):
            return np.nan, np.nan, np.unique(res['Instrument'])[0]

        # HARPS gives archival RVs. downloading them can be done... but for
        # s6+s7, only a few objects are viable.
        if np.all(nparr(res['Instrument']) == 'HARPS'):
            print('WARNING: SKIPPING AUTOMATION OF HARPS ARCHIVAL RV GETTING')
            return np.nan, np.nan, 'HARPS'

    else:
        return np.nan, np.nan, np.nan

import numpy,from numpy,import astropy,from astropy


#!/usr/bin/python

import cv2,cv
import numpy as np
import socket

import time

s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

host = "192.168.43.34"
port = 3335

s.bind((host, port))
# queue up to 5 requests
s.listen(5)

odroidflag=0

greencounter=0
orangecounter=0
wrongcounter=0

circlecount=0

lower_green = np.array([37,108,44]) # 1
upper_green = np.array([54,215,255]) # 11

lower_orange = np.array([4,113,51]) # 1
upper_orange = np.array([18,203,250]) # 11

prevcx=0
prevcy=0

prevcxg=0
prevcyg=0

cxg=0
cyg=0

cx=0
cy=0

sc,addr = s.accept()

while True:

    print "start"

    a=sc.recv(1024)
    print a

    if a=="ball detect":
        cv2.namedWindow("preview")
        vc = cv2.VideoCapture(0)
#	time.sleep(5)
        if vc.isOpened(): # try to get the first frame
            rval, frame = vc.read()
            frame = cv2.flip(frame, 1)
        else:
            rval = False

        while rval:
            cv2.waitKey(1)
            rval, frame = vc.read()
            frame = cv2.flip(frame, 1)
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            mask = cv2.inRange(hsv, lower_green, upper_green)
            gmask=cv2.inRange(hsv, lower_orange, upper_orange)
            mask = cv2.erode(mask, cv2.getStructuringElement(cv2.MORPH_ERODE, (3,3) ))
            mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_DILATE, (3,3) ))
            gmask = cv2.erode(gmask, cv2.getStructuringElement(cv2.MORPH_ERODE, (3,3) ))
            gmask = cv2.dilate(gmask, cv2.getStructuringElement(cv2.MORPH_DILATE, (3,3) ))
            moments=cv2.moments(mask, False)
            gmoments=cv2.moments(gmask,False)

            if(moments['m00'] >0):
                cx = moments['m10']/moments['m00']
                cy = moments['m01']/moments['m00']
            else:
                cx=0
                cy=0

            if(gmoments['m00'] >0):
                cxg = gmoments['m10']/gmoments['m00']
                cyg = gmoments['m01']/gmoments['m00']
            else:
                cxg=0
                cyg=0

            if (abs(cxg-prevcxg)<25 and abs(cyg-prevcyg)<25 and cxg!=0 and cyg!=0):
                cv2.circle(frame,(int(cxg),int(cyg)),10,(0,0,255),2)
                orangecounter+=1
                greencounter=0
                wrongcounter=0
                if(orangecounter >= 2):
#		   print str(cxg)+" "+str(cyg)
                    if(cxg >= 480):
                        sc.send("forward".encode('ascii'))
                    if(cxg < 480):
                        sc.send("right".encode('ascii'))
                    orangecounter=0

            elif (abs(cx-prevcx)<25 and abs(cy-prevcy)<25 and cx!=0 and cy!=0):
#            print ("arm green"+str(int(cx))+" "+str(int(cy)))
                cv2.circle(frame,(int(cx),int(cy)),10,(0,0,255),2)
                greencounter+=1
                orangecounter=0
                wrongcounter=0
                if(greencounter >= 2):
#		   print(str(cx)+" "+str(cy))
                    if(cx >= 480):
                        sc.send("forward".encode('ascii'))
                    if(cx < 480):
                        sc.send("right".encode('ascii'))
                    greencounter=0
            else:
#       cv2.imshow("preview",frame)
                wrongcounter+=1
                greencounter=0
                orangecounter=0
                if(wrongcounter >=10):
                    circlecount+=10
                    sc.send("right 10".encode('ascii'))
                    if(circlecount==360):
                        sc.send("forward".encode('ascii'))
                        circlecount=0

            prevcx=cx
            prevcy=cy
            prevcxg=cxg
            prevcyg=cyg

#            vc.release()
#            break
            cv2.imshow("preview", frame)

    else:
        print "third cam mode"
#            break
#            sc.close()
#    cv2.imshow("preview", frame)
#    cv2.destroyAllWindows()
#    else:

sc.close()