#!/usr/bin/env python
# coding: utf-8

import os
import sys
import subprocess

from pathlib import Path

import csv

import numpy as np

import itk
from itk import TubeTK as tube

def scv_is_bundled():
    return getattr(sys, 'frozen', False) and hasattr(sys, '_MEIPASS')

def scv_get_perfusion_toolbox_path():
    if scv_is_bundled():
        return os.path.join(sys._MEIPASS, 'StroCoVess', 'perfusion_toolbox')
    return os.path.join(os.path.dirname(os.path.realpath(__file__)),
        'perfusion_toolbox')

#################
#################
#################
#################
#################
def scv_convert_ctp_to_cta(filenames,
                           report_progress=print,
                           debug=False,
                           output_dirname="."):

    filenames.sort()
    num_images = len(filenames)

    base_im = itk.imread(filenames[num_images//2],itk.F)
    base_spacing = base_im.GetSpacing()

    progress_percent = 10
    report_progress("Reading images",progress_percent)

    Dimension = 3
    PixelType = itk.ctype('float')
    ImageType = itk.Image[PixelType,Dimension]

    imdatamax = itk.GetArrayFromImage(base_im)
    imdatamin = imdatamax

    if output_dirname!=None and not os.path.exists(output_dirname):
        os.mkdir(output_dirname)

    progress_percent = 20
    progress_per_file = 70/num_images
    for imNum in range(num_images):
        imMoving = itk.imread(filenames[imNum],itk.F)
        if imMoving.shape != base_im.shape:
            resample = tube.ResampleImage.New(Input=imMoving)
            resample.SetMatchImage(base_im)
            resample.Update()
            imMovingIso = resample.GetOutput()
            progress_label = "Resampling "+str(imNum)+" of "+str(num_images)
            report_progress(progress_label,progress_percent)
        else:
            imMovingIso = imMoving
        imdataTmp = itk.GetArrayFromImage(imMovingIso)
        imdatamax = np.maximum(imdatamax,imdataTmp)
        imdataTmp = np.where(imdataTmp==-1024,imdatamin,imdataTmp)
        imdatamin = np.minimum(imdatamin,imdataTmp)
        progress_percent += progress_per_file
        progress_label = "Integrating "+str(imNum)+" of "+str(num_images)
        report_progress(progress_label,progress_percent)

    report_progress("Generating CT, CTA, and CTP",90)

    ct = itk.GetImageFromArray(imdatamin)
    ct.CopyInformation(base_im)

    cta = itk.GetImageFromArray(imdatamax)
    cta.CopyInformation(base_im)

    diff = imdatamax-imdatamin
    diff[:4,:,:] = 0
    diff[-4:,:,:] = 0
    diff[:,:4,:] = 0
    diff[:,-4:,:] = 0
    diff[:,:,:4] = 0
    diff[:,:,-4:] = 0
    dsa = itk.GetImageFromArray(diff)
    dsa.CopyInformation(base_im)

    report_progress("Done",100)
    return ct,cta,dsa

#################
#################
#################
#################
#################
def scv_segment_brain_from_ct(ct_image,
                               report_progress=print,
                               debug=False):

    ImageType = itk.Image[itk.F,3]
    LabelMapType = itk.Image[itk.UC,3]

    report_progress("Threshold",5)
    thresh = tube.ImageMath.New(Input=ct_image)
    thresh.IntensityWindow(-50,6000,0,6000)
    imgt = thresh.GetOutput()
    thresh.ReplaceValuesOutsideMaskRange(imgt,1,6000,0)
    thresh.ReplaceValuesOutsideMaskRange(imgt,0,700,1)
    tmpimg = thresh.GetOutput()
    thresh.ReplaceValuesOutsideMaskRange(tmpimg,0,1,2)
    ct_tmp = thresh.GetOutput()

    report_progress("Initial Mask",10)
    maskMath = tube.ImageMath.New(Input=ct_tmp)
    # remove skin
    maskMath.Dilate(15,0,1)
    # shrink brain
    maskMath.Erode(6,2,0)
    # restore skull
    maskMath.ReplaceValueWithinMaskRange(ct_tmp,1,1,0,1)
    # shrink skull
    maskMath.Dilate(3,1,0)
    maskMath.Erode(1,1,0)
    comboSeed = maskMath.GetOutputUChar()

    report_progress("Connected Component",20)
    segmenter = tube.SegmentConnectedComponentsUsingParzenPDFs[ImageType,
                    LabelMapType].New()
    segmenter.SetFeatureImage( ct_image )
    segmenter.SetInputLabelMap( comboSeed )
    segmenter.SetObjectId( 2 )
    segmenter.AddObjectId( 1 )
    segmenter.SetVoidId( 0 )
    segmenter.SetErodeDilateRadius( 20 )
    segmenter.SetHoleFillIterations( 40 )
    segmenter.Update()
    segmenter.ClassifyImages()
    brainMaskRaw = segmenter.GetOutputLabelMap()

    report_progress("Masking",60)

    maskMath = itk.CastImageFilter[LabelMapType, ImageType].New()
    maskMath.SetInput(brainMaskRaw)
    maskMath.Update()
    brainMaskF = maskMath.GetOutput()
    maskMath = tube.ImageMath.New(Input = brainMaskF)
    maskMath.Threshold(2,2,1,0)
    maskMath.Dilate(2,1,0)
    maskMath.Erode(3,1,0)
    brainMaskRaw2 = maskMath.GetOutputUChar()

    connComp = tube.SegmentConnectedComponents.New(Input=brainMaskRaw2)
    connComp.SetKeepOnlyLargestComponent(True)
    connComp.Update()
    brainMask = connComp.GetOutput()

    report_progress("Finishing",90)
    cast = itk.CastImageFilter[LabelMapType,ImageType].New()
    cast.SetInput(brainMask)
    cast.Update()
    brainMaskF = cast.GetOutput()

    brainMath = tube.ImageMath[ImageType].New(Input=ct_image)
    brainMath.ReplaceValuesOutsideMaskRange( brainMaskF,1,1,-1024)
    ct_brain_image = brainMath.GetOutput()

    report_progress("Done",100)
    return ct_brain_image, brainMask

#################
#################
#################
#################
#################
def scv_enhance_vessels_in_cta(cta_image,
                               cta_roi_image,
                               report_progress=print,
                               debug=False ):

    ImageType = itk.Image[itk.F,3]
    LabelMapType = itk.Image[itk.UC,3]

    report_progress("Masking",5)
    imMath = tube.ImageMath.New(Input=cta_roi_image)
    imMath.Threshold( 0.00001,4000,1,0)
    imMath.Erode(10,1,0)
    imBrainMaskErode = imMath.GetOutput()
    imMath.SetInput(cta_roi_image)
    imMath.IntensityWindow(0,300,0,300)
    imMath.ReplaceValuesOutsideMaskRange(imBrainMaskErode,0.5,1.5,0)
    imBrainErode = imMath.GetOutput()

    spacing = cta_image.GetSpacing()[0]

    report_progress("Blurring",10)
    imMath = tube.ImageMath[ImageType].New()
    imMath.SetInput(imBrainErode)
    imMath.Blur(1.5*spacing)
    imBlur = imMath.GetOutput()
    imBlurArray = itk.GetArrayViewFromImage(imBlur)

    report_progress("Generating Seeds",20)
    numSeeds = 15
    seedCoverage = 20
    seedCoord = np.zeros([numSeeds,3])
    for i in range(numSeeds):
        seedCoord[i] = np.unravel_index(np.argmax(imBlurArray,
                           axis=None),imBlurArray.shape)
        indx = [int(seedCoord[i][0]),int(seedCoord[i][1]),
                int(seedCoord[i][2])]
        minX = max(indx[0]-seedCoverage,0)
        maxX = max(indx[0]+seedCoverage,imBlurArray.shape[0])
        minY = max(indx[1]-seedCoverage,0)
        maxY = max(indx[1]+seedCoverage,imBlurArray.shape[1])
        minZ = max(indx[2]-seedCoverage,0)
        maxZ = max(indx[2]+seedCoverage,imBlurArray.shape[2])
        imBlurArray[minX:maxX,minY:maxY,minZ:maxZ]=0
        indx.reverse()
        seedCoord[:][i] = cta_roi_image.TransformIndexToPhysicalPoint(indx)

    report_progress("Segmenting Initial Vessels",30)
    vSeg = tube.SegmentTubes.New(Input=cta_roi_image)
    vSeg.SetVerbose(debug)
    vSeg.SetMinRoundness(0.4)
    vSeg.SetMinCurvature(0.002)
    vSeg.SetRadiusInObjectSpace( 1 )
    for i in range(numSeeds):
        progress_label = "Vessel "+str(i)+" of "+str(numSeeds)
        progress_percent = i/numSeeds*20+30
        report_progress(progress_label,progress_percent)
        vSeg.ExtractTubeInObjectSpace( seedCoord[i],i )
    tubeMaskImage = vSeg.GetTubeMaskImage()

    imMath.SetInput(tubeMaskImage)
    imMath.AddImages(cta_roi_image,200,1)
    blendIm = imMath.GetOutput()

    report_progress("Computing Training Mask",50)
    trMask = tube.ComputeTrainingMask[ImageType,LabelMapType].New()
    trMask.SetInput( tubeMaskImage )
    trMask.SetGap( 4 )
    trMask.SetObjectWidth( 1 )
    trMask.SetNotObjectWidth( 1 )
    trMask.Update()
    fgMask = trMask.GetOutput()

    report_progress("Enhancing Image",70)
    enhancer = tube.EnhanceTubesUsingDiscriminantAnalysis[ImageType,
                   LabelMapType].New()
    enhancer.AddInput( cta_image )
    enhancer.SetLabelMap( fgMask )
    enhancer.SetRidgeId( 255 )
    enhancer.SetBackgroundId( 128 )
    enhancer.SetUnknownId( 0 )
    enhancer.SetTrainClassifier(True)
    enhancer.SetUseIntensityOnly(True)
    enhancer.SetScales([0.75*spacing,2*spacing,6*spacing])
    enhancer.Update()
    enhancer.ClassifyImages()

    report_progress("Finalizing",90)
    imMath = tube.ImageMath[ImageType].New()
    imMath.SetInput(enhancer.GetClassProbabilityImage(0))
    imMath.Blur(0.5*spacing)
    prob0 = imMath.GetOutput()
    imMath.SetInput(enhancer.GetClassProbabilityImage(1))
    imMath.Blur(0.5*spacing)
    prob1 = imMath.GetOutput()
    cta_vess = itk.SubtractImageFilter(Input1=prob0, Input2=prob1)

    imMath.SetInput(cta_roi_image)
    imMath.Threshold(0.0000001,2000,1,0)
    imMath.Erode(2,1,0)
    imBrainE = imMath.GetOutput()

    imMath.SetInput(cta_vess)
    imMath.ReplaceValuesOutsideMaskRange(imBrainE,1,1,-0.001)
    cta_roi_vess = imMath.GetOutput()

    report_progress("Done",100)
    return cta_vess,cta_roi_vess

#################
#################
#################
#################
#################
def scv_extract_vessels_from_cta(cta_image,
                                 cta_roi_vessels_image,
                                 report_progress=print,
                                 debug=False,
                                 output_dirname="."):

    if output_dirname!=None and not os.path.exists(output_dirname):
        os.mkdir(output_dirname)

    spacing = cta_image.GetSpacing()[0]

    report_progress("Thresholding",5)
    imMath = tube.ImageMath.New(cta_roi_vessels_image)
    imMath.MedianFilter(1)
    imMath.Threshold(0.00000001,9999,1,0)
    vess_mask_im = imMath.GetOutputShort()

    if debug and output_dirname!=None:
        itk.imwrite(vess_mask_im,
            output_dirname+"/extract_vessels_mask.mha",
            compression=True)

    report_progress("Connecting",10)
    ccSeg = tube.SegmentConnectedComponents.New(vess_mask_im)
    ccSeg.SetMinimumVolume(50)
    ccSeg.Update()
    vess_mask_cc_im = ccSeg.GetOutput()

    if debug and output_dirname!=None:
        itk.imwrite(vess_mask_cc_im,
            output_dirname+"/extract_vessels_mask_cc.mha",
            compression=True)

    imMathSS = tube.ImageMath.New(vess_mask_cc_im)
    imMathSS.Threshold(0,0,1,0)
    vess_mask_inv_im = imMathSS.GetOutputFloat()

    report_progress("Filling in",20)
    distFilter = itk.DanielssonDistanceMapImageFilter.New(vess_mask_inv_im)
    distFilter.Update()
    dist_map_im = distFilter.GetOutput()

    report_progress("Generating seeds",30)
    imMath.SetInput(dist_map_im)
    imMath.Blur(0.5*spacing)
    tmp = imMath.GetOutput()
    # Distance map's distances are in index units, not spacing
    imMath.ReplaceValuesOutsideMaskRange(tmp,0.333,10,0)
    initial_radius_im = imMath.GetOutput()

    if debug and output_dirname!=None:
        itk.imwrite(initial_radius_im,
            output_dirname+"/vessel_extraction_initial_radius.mha",
            compression=True)

    report_progress("Generating input",30)
    imMath.SetInput(cta_image)
    imMath.ReplaceValuesOutsideMaskRange(cta_roi_vessels_image,0,1000,0)
    imMath.Blur(0.4*spacing)
    imMath.NormalizeMeanStdDev()
    imMath.IntensityWindow(-4,4,0,1000)
    input_im = imMath.GetOutput()

    if debug and output_dirname!=None:
        itk.imwrite(input_im,
            output_dirname+"/vessel_extraction_input.mha",
            compression=True)

    report_progress("Extracting vessels",40)
    vSeg = tube.SegmentTubes.New(Input=input_im)
    vSeg.SetVerbose(debug)
    vSeg.SetMinCurvature(0.0001) #0
    vSeg.SetMinRoundness(0.02)
    vSeg.SetMinRidgeness(0.5)
    vSeg.SetMinLevelness(0.0)
    vSeg.SetRadiusInObjectSpace( 0.8*spacing )
    vSeg.SetBorderInIndexSpace(3)
    vSeg.SetSeedMask( initial_radius_im )
    #vSeg.SetSeedRadiusMask( initial_radius_im )
    vSeg.SetOptimizeRadius(True)
    vSeg.SetUseSeedMaskAsProbabilities(True)
    vSeg.SetSeedExtractionMinimumProbability(0.95) #0.99
    vSeg.ProcessSeeds()

    report_progress("Finalizing",90)
    tubeMaskImage = vSeg.GetTubeMaskImage()

    if debug and output_dirname!=None:
        itk.imwrite(tubeMaskImage,
            output_dirname+"/vessel_extraction_output.mha",
            compression=True)

    report_progress("Done",100)
    return tubeMaskImage,vSeg.GetTubeGroup()

#################
#################
#################
#################
#################
def scv_register_ctp_images(fixed_image_filename,
                        moving_image_filenames,
                        output_dirname,
                        report_progress=print,
                        debug=False):
    ImageType = itk.Image[itk.F,3]

    num_images = len(moving_image_filenames)
    progress_percent = 10
    progress_per_file = 70/num_images

    fixed_im = itk.imread(fixed_image_filename,itk.F)
    fixed_im_spacing = fixed_im.GetSpacing()
    if fixed_im_spacing[0] != fixed_im_spacing[1] or \
       fixed_im_spacing[1] != fixed_im_spacing[2]:
        report_progress("Resampling",progress_percent)
        resample = tube.ResampleImage.New(Input=fixed_im)
        resample.SetMakeIsotropic(True)
        resample.Update()
        fixed_im = resample.GetOutput()
        if debug:
            progress_label = "DEBUG: Resampling to "+str(
                fixed_im.GetSpacing())
            report_progress(progress_label,progress_percent)

    imMath = tube.ImageMath.New(fixed_im)
    imMath.Threshold(150,800,1,0)
    imMath.Dilate(10,1,0)
    mask_im = imMath.GetOutputUChar()
    mask_array = itk.GetArrayViewFromImage(mask_im)
    mask_array[:4,:,:] = 0
    mask_array[-4:,:,:] = 0
    mask_obj = itk.ImageMaskSpatialObject[3].New()
    mask_obj.SetImage(mask_im)
    mask_obj.Update()

    for imNum in range(num_images):
        progress_percent += progress_per_file
        progress_label = "Registering "+str(imNum)+" of "+str(num_images)
        report_progress(progress_label,progress_percent)

        if moving_image_filenames[imNum] != fixed_image_filename:
            moving_im = itk.imread(moving_image_filenames[imNum],itk.F)

            imreg = tube.RegisterImages[ImageType].New()
            imreg.SetFixedImage(fixed_im)
            imreg.SetMovingImage(moving_im)
            imreg.SetRigidMaxIterations(100)
            imreg.SetRegistration("RIGID")
            imreg.SetExpectedOffsetMagnitude(5)
            imreg.SetExpectedRotationMagnitude(0.05)
            imreg.SetFixedImageMaskObject(mask_obj)
            imreg.SetUseEvolutionaryOptimization(False)
            if debug:
                imreg.SetReportProgress(True)
            imreg.Update()

            tfm = imreg.GetCurrentMatrixTransform()
            moving_reg_im = imreg.ResampleImage("SINC_INTERPOLATION",
                                                moving_im,tfm,-1024)
            if output_dirname!=None:
                pname,fname = os.path.split(moving_image_filenames[imNum])
                rename_file_fname = os.path.splitext(fname)
                new_fname = str(rename_file_fname[0])+"_reg.nii"
                new_filename = os.path.join(output_dirname,new_fname)
                itk.imwrite(moving_reg_im,new_filename,compression=True)
        elif output_dirname!=None:
            pname,fname = os.path.split(moving_image_filenames[imNum])
            rename_file_fname = os.path.splitext(fname)
            new_fname = str(rename_file_fname[0])+"_reg.nii"
            new_filename = os.path.join(output_dirname,new_fname)
            itk.imwrite(moving_reg_im,new_filename,compression=True)
    report_progress("Done",100)

#################
#################
#################
#################
#################
def scv_register_atlas_to_image(atlas_im, atlas_mask_im, in_im):
    ImageType = itk.Image[itk.F,3]

    regAtlasToIn = tube.RegisterImages[ImageType].New(FixedImage=in_im,
        MovingImage=atlas_im)
    regAtlasToIn.SetReportProgress(True)
    regAtlasToIn.SetRegistration("PIPELINE_AFFINE")
    regAtlasToIn.SetMetric("MATTES_MI_METRIC")
    regAtlasToIn.SetInitialMethodEnum("INIT_WITH_IMAGE_CENTERS")
    regAtlasToIn.Update()
    atlas_reg_im = regAtlasToIn.ResampleImage()
    atlas_mask_reg_im = regAtlasToIn.ResampleImage("NEAREST_NEIGHBOR",
        atlas_mask_im)

    return atlas_reg_im,atlas_mask_reg_im

#################
#################
#################
#################
#################
def scv_compute_atlas_region_stats(atlas_im,
                                   time_im,
                                   vess_im,
                                   number_of_time_bins=100,
                                   report_progress=print,
                                   debug=False):

    atlas_arr = itk.GetArrayFromImage(atlas_im)
    time_arr = itk.GetArrayFromImage(time_im)
    vess_arr = itk.GetArrayFromImage(vess_im)

    num_regions = int(atlas_arr.max())
    time_max = float(time_arr.max())
    time_min = float(time_arr.min())
    nbins = int(number_of_time_bins)
    time_factor = (time_max-time_min)/(nbins+1)
    print("Time range =",time_min,"-",time_max)

    bin_value = np.zeros([num_regions,nbins])
    bin_count = np.zeros([num_regions,nbins])

    for atlas_region in range(num_regions):
        report_progress("Masking",(atlas_region+1)*(100/num_regions))
        indx_arr = np.where(atlas_arr==atlas_region)
        indx_list = list(zip(indx_arr[0],indx_arr[1],indx_arr[2]))
        for indx in indx_list:
            time_bin = int((time_arr[indx]-time_min)/time_factor)
            time_bin = min(max(0,time_bin),nbins-1)
            if np.isnan(vess_arr[indx]) == False:
                bin_count[atlas_region,time_bin] += 1
                bin_value[atlas_region,time_bin] += vess_arr[indx]

    bin_label = np.arange(nbins) * time_factor - time_min
    bin_value = np.divide(bin_value,bin_count,where=bin_count!=0)

    report_progress("Done",100)
    return bin_label,bin_value,bin_count

#################
#################
#################
#################
#################
def scv_convert_3d_files_to_4d_file(in_filenames,out_filename):
    num_3d_files = len(in_filenames)

    ImageType = itk.Image[itk.F, 3]
    Write4D = tube.Write4DImageFrom3DImages[ImageType].New()
    Write4D.SetNumberOfInputImages(num_3d_files)
    Write4D.SetFileName(out_filename)
    for i,file in enumerate(in_filenames):
        img = itk.imread(file,itk.F)
        Write4D.SetNthInputImage(i, img)
    Write4D.Update()

#################
#################
#################
#################
#################
def scv_prepare_3d_for_perfusion_toolbox(prep_3d_in_filenames,
                                         prep_3d_out_dirname,
                                         report_progress=print,
                                         report_subprogress=print,
                                         debug=False):
    num_3d_files = len(prep_3d_in_filenames)

    reg_fixed_image = prep_3d_in_filenames[num_3d_files//2]
    reg_in_filenames = prep_3d_in_filenames
    report_progress("Registering CTP",40)
    scv_register_ctp_images(reg_fixed_image,
        reg_in_filenames,
        output_dirname=prep_3d_out_dirname,
        report_progress=report_subprogress,
        debug=debug)

    # update filenames to registered ctp
    report_progress("Saving 4D CTP",60)
    new_ctp_3d_filenames = []
    for i in range(num_3d_files):
        file_path,file_name = os.path.split(prep_3d_in_filenames[i])
        rename_file_name = os.path.splitext(str(file_name))
        new_ctp_3d_filenames.append( os.path.realpath(os.path.join(
            prep_3d_out_dirname, str(rename_file_name[0]) + "_reg.nii")))

    # Write 4D ctp registered
    file_path,file_name = os.path.split(prep_3d_in_filenames[0])
    rename_file_name = os.path.splitext(str(file_name))
    ctp_base_filename = str(rename_file_name[0])

    ctp_filename = ctp_base_filename + "-4D_reg.nii"
    ctp_4d_out_filename = os.path.realpath(os.path.join(
        prep_3d_out_dirname, ctp_filename))
    scv_convert_3d_files_to_4d_file(new_ctp_3d_filenames,
        ctp_4d_out_filename)

    # Compute CT, CTA, DSA
    report_progress("Computing CT, CTA, DSA",70)
    ct_im,cta_im,dsa_im = scv_convert_ctp_to_cta(new_ctp_3d_filenames,
        report_progress = report_subprogress,
        debug=debug,
        output_dirname=prep_3d_out_dirname)
    ct_filename = ctp_base_filename + "_ct.nii"
    ct_out_filename = os.path.realpath(os.path.join(
        prep_3d_out_dirname, ct_filename))
    itk.imwrite(ct_im, ct_out_filename,compression=True)
    cta_filename = ctp_base_filename + "_cta.nii"
    cta_out_filename = os.path.realpath(os.path.join(
        prep_3d_out_dirname, cta_filename))
    itk.imwrite(cta_im, cta_out_filename,compression=True)
    dsa_filename = ctp_base_filename + "_dsa.nii"
    dsa_out_filename = os.path.realpath(os.path.join(
        prep_3d_out_dirname, dsa_filename))
    itk.imwrite(dsa_im, dsa_out_filename,compression=True)

    # Segment brain from CT
    report_progress("Segmenting Brain",80)
    ct_brain,mask_brain = scv_segment_brain_from_ct(ct_im,
        report_progress = report_subprogress,
        debug=debug)
    rename_file_ct = str(ct_filename)
    rename_file_ct = os.path.splitext(rename_file_ct)
    mask_brain_filename = str(rename_file_ct[0]) + "_brain_mask.nii"
    mask_brain_out_filename = os.path.realpath(os.path.join(
        prep_3d_out_dirname, mask_brain_filename))
    itk.imwrite(mask_brain, mask_brain_out_filename)

    report_progress("Done",100)

    return ctp_4d_out_filename, ct_out_filename, \
        cta_out_filename, dsa_out_filename, \
        mask_brain_out_filename

#################
#################
#################
#################
#################
def scv_prepare_4d_for_perfusion_toolbox(prep_4d_in_filename,
                                         prep_4d_out_dirname,
                                         report_progress=print,
                                         report_subprogress=print,
                                         debug=False):
    # convert 4D ctp to 3D images
    ctp_dirname,ctp_filename = os.path.split(prep_4d_in_filename)

    rename_file_ctp = os.path.splitext(str(ctp_filename))

    report_progress("Reading image",10)

    img4d_im = itk.imread(prep_4d_in_filename, itk.F)
    img4d_array = itk.GetArrayFromImage(img4d_im)
    img4d_shape = img4d_array.shape
    img4d_spacing = np.array(img4d_im.GetSpacing())
    img4d_direction = np.array(img4d_im.GetDirection())
    img4d_origin = np.array(img4d_im.GetOrigin())
    img3d_spacing = img4d_spacing[0:3]
    img3d_origin = img4d_origin[0:3]
    img3d_direction = img4d_direction[0:3, 0:3]
    num_3d_files = img4d_shape[0]
    new_filenames = []
    report_progress("Creating CTP",15)
    subprogress = 15
    subprogress_per_file = 100 / num_3d_files
    for i in range(num_3d_files):
        report_subprogress("Writing",
            subprogress+subprogress_per_file*i)
        tmp_filename = f'CTP_{i:03}.mha'
        new_filename = os.path.join(prep_4d_out_dirname,tmp_filename)
        img3d_array = img4d_im[i,:,:,:]
        img3d_im = itk.GetImageFromArray(img3d_array)
        img3d_im.SetSpacing(img3d_spacing)
        img3d_im.SetOrigin(img3d_origin)
        img3d_im.SetDirection(img3d_direction)
        itk.imwrite(img3d_im,new_filename,compression=True)
        new_filenames.append(new_filename)

    results = [new_filenames]
    results += scv_prepare_3d_for_perfusion_toolbox( new_filenames, \
        prep_4d_out_dirname, report_progress, report_subprogress, debug)

    return results

#################
#################
#################
#################
#################
def scv_fix_image_info(filename, match_image):
    im = itk.imread(filename,itk.F)
    im.SetSpacing(match_image.GetSpacing())
    im.SetOrigin(match_image.GetOrigin())
    im.SetDirection(match_image.GetDirection())
    itk.imwrite(im,filename,compression=True)

#################
#################
#################
#################
#################
#matlab -r "addpath('./perfusion_toolbox');DSC_report('/data/UNC-Stroke/UNC/CTP/CTAT-001-PTO/CTAT-001-Perf-4D_reg.nii','/data/UNC-Stroke/UNC/CTP/CTAT-001-PTO/CTAT-001-Perf-4D_reg_ct_brain_mask.nii','./test');quit"
def scv_run_perfusion_toolbox(ctp_4d_in_filename,
                              ctp_mask_filename,
                              perfusion_out_dirname):
    cmd = "addpath('"+scv_get_perfusion_toolbox_path()+"');DSC_report('"+ \
          ctp_4d_in_filename+"','"+ctp_mask_filename+"','"+ \
          perfusion_out_dirname+"');quit"
    subprocess.run(["matlab", "-wait", "-r", cmd])

    mask_image = itk.imread(ctp_mask_filename,itk.F)
    cbf_filename = os.path.join(perfusion_out_dirname,"CBF_SVD.nii")
    scv_fix_image_info(cbf_filename,mask_image)

    cbv_filename = os.path.join(perfusion_out_dirname, "CBV.nii")
    scv_fix_image_info(cbv_filename,mask_image)

    cbv_lc_filename = os.path.join(perfusion_out_dirname, "CBV_LC.nii")
    scv_fix_image_info(cbv_lc_filename,mask_image)

    mtt_filename = os.path.join(perfusion_out_dirname, "MTT_SVD.nii")
    scv_fix_image_info(mtt_filename,mask_image)

    tmax_filename = os.path.join(perfusion_out_dirname, "Tmax_SVD.nii")
    scv_fix_image_info(tmax_filename,mask_image)

    ttp_filename = os.path.join(perfusion_out_dirname, "TTP.nii")
    scv_fix_image_info(ttp_filename,mask_image)

    return cbf_filename,cbv_filename,mtt_filename,tmax_filename,ttp_filename

#################
#################
#################
#################
#################
def scv_generate_vessel_report(ctp_3d_filenames,
                               ctp_4d_filename,
                               ct_filename,
                               cta_fielname,
                               dsa_filename,
                               mask_brain_filename,
                               atlas_path,
                               report_out_dirname,
                               report_progress=print,
                               report_subprogress=print,
                               debug=False):
    new_ctp_filenames = []
    base_3d_image = itk.imread(ctp_3d_filenames[0], itk.F)
    ImageMath = tube.ImageMath.New(base_3d_image)
    for filename in ctp_3d_filenames:
        img = itk.imread(filename, itk.F)
        ImageMath.SetInput(img)
        ImageMath.Blur(0.5)
        ImageMath.BlurOrder(2.0,0,2)
        new_img = ImageMath.GetOutput()
        Resample = tube.ResampleImage.New(Input=new_img)
        Resample.SetSpacing([1.5,1.5,5])
        Resample.SetInterpolator("Sinc")
        Resample.Update()
        img = Resample.GetOutput()
        filename_path,filename_name = os.path.split(filename)
        org_filename = os.path.splitext(filename_name)
        base_filename = os.path.join(report_out_dirname,org_filename[0])
        new_filename = str(base_filename)+'_15x15x5.nii'
        itk.imwrite(img,new_filename,compression=True)
        new_ctp_filenames.append(new_filename)

    filename_path,filename_name = os.path.split(ctp_4d_filename)
    org_filename = os.path.splitext(filename_name)
    base_filename = os.path.join(report_out_dirname,org_filename[0])
    new_ctp_4d_filename = str(base_filename)+'_15x15x5.nii'
    scv_convert_3d_files_to_4d_file(new_ctp_filenames,
        new_ctp_4d_filename)

    match_image = itk.imread(new_ctp_filenames[0], itk.UC)
    mask = itk.imread(mask_brain_filename,itk.UC)
    ResampleMask = tube.ResampleImage.New(Input=mask)
    ResampleMask.SetMatchImage(match_image)
    ResampleMask.SetInterpolator("NearestNeighbor")
    ResampleMask.Update()
    new_mask = ResampleMask.GetOutput()
    filename_path,filename_name = os.path.split(mask_brain_filename)
    org_filename = os.path.splitext(filename_name)
    base_filename = os.path.join(report_out_dirname,org_filename[0])
    new_mask_filename = str(base_filename)+'_15x15x5.nii'
    itk.imwrite(new_mask,new_mask_filename,compression=True)

    # Call matlab and pass it the 4D ctp filename and the output directory
    cbf_filename,cbv_filename,mtt_filename,tmax_filename,ttp_filename = \
        scv_run_perfusion_toolbox(new_ctp_4d_filename, \
            new_mask_filename, report_out_dirname)

    # Vessel enhancement and extraction
    in_im = itk.imread(dsa_filename, itk.F)

    in_filename_base = str(os.path.splitext(dsa_filename)[0])

    brain_mask_im = itk.imread(mask_brain_filename, itk.F)
    ImageMath.SetInput(in_im)
    ImageMath.ReplaceValuesOutsideMaskRange(brain_mask_im,0.9,1.1,0)
    in_brain_im = ImageMath.GetOutput()

    report_progress("Enhancing vessels",40)
    # Enhancing vessels creates an image in which intensity is
    #    related to "vesselness" instead of being related to the
    #    amount of contrast agent in the vessel.  This simplifies
    #    subsequent vessel seeding and traversal stopping criteria.
    in_vess_im,in_brain_vess_im = scv_enhance_vessels_in_cta(
        in_im,
        in_brain_im,
        report_progress=report_subprogress,
        debug=debug)
    itk.imwrite(in_vess_im,
        os.path.join(report_out_dirname,
            in_filename_base+"_vessels_enhanced.mha"),
        compression=True)
    itk.imwrite(in_brain_vess_im,
        os.path.join(report_out_dirname,
            in_filename_base+"_brain_vessels_enhanced.mha"),
        compression=True)

    report_progress("Extracting vessels",60)
    vess_mask_im,vess_so = scv_extract_vessels_from_cta(
        in_vess_im,
        in_brain_vess_im,
        report_progress=report_subprogress,
        debug=debug,
        output_dirname=report_out_dirname)

    itk.imwrite(in_brain_vess_im,
        os.path.join(report_out_dirname,
            in_filename_base+"_vessels_extracted.mha"),
        compression=True)

    SOWriter = itk.SpatialObjectWriter[3].New()
    SOWriter.SetInput(vess_so)
    SOWriter.SetBinaryPoints(True)
    SOWriter.SetFileName(os.path.join(report_out_dirname,
        in_filename_base+"_vessels_extracted.tre"))
    SOWriter.Update()

    VTPWriter = itk.WriteTubesAsPolyData.New()
    VTPWriter.SetInput(vess_so)
    VTPWriter.SetFileName(os.path.join(report_out_dirname,
        in_filename_base+"_vessels_extracted.vtp"))
    VTPWriter.Update()

    report_progress("Generating Perfusion Stats",80)

    script_dirname = os.path.dirname(os.path.realpath(__file__))
    atlas_im = itk.imread(
        os.path.join(atlas_path,'atlas_brainweb.mha'), itk.F)
    atlas_mask_im = itk.imread(
        os.path.join(atlas_path,'atlas_brainweb_mask.mha'), itk.F)
    atlas_reg_im,atlas_mask_reg_im = scv_register_atlas_to_image(
        atlas_im,
        atlas_mask_im,
        in_brain_im)
    ImageMath = tube.ImageMath.New(Input=atlas_mask_reg_im)
    ImageMath.ReplaceValuesOutsideMaskRange(vess_mask_im,
        0.000001,9999,4)
    ImageMath.ReplaceValuesOutsideMaskRange(in_brain_im,
        0.000001,9999,0)
    vess_atlas_mask_im = ImageMath.GetOutput()
    itk.imwrite(vess_atlas_mask_im,
        os.path.join(report_out_dirname,
            in_filename_base+"_vessels_atlas_mask.mha"),
        compression=True)

    TubeMath = tube.TubeMath[3,itk.F].New()
    TubeMath.SetInputTubeGroup(vess_so)
    TubeMath.SetUseAllTubes()
    TubeMath.ComputeTubeRegions(vess_atlas_mask_im)

    graph_label = None
    graph_data = None
    ttp_im = None
    if len(ttp_filename) > 0:
        report_progress("Generating TTP Graphs",92)
        ttp_im = itk.imread(ttp_filename, itk.F)
        Resample = tube.ResampleImage.New(Input=ttp_im)
        Resample.SetMatchImage(vess_atlas_mask_im)
        Resample.Update()
        fit_ttp_im = Resample.GetOutput()
        TubeMath.SetPointValuesFromImage(fit_ttp_im, "TTP")
        TubeMath.SetPointValuesFromTubeRegions(fit_ttp_im,
            "TTP_Tissue",
            1.5,
            4)
        TubeMath.SmoothTubeProperty("TTP",4)
        TubeMath.SmoothTubeProperty("TTP_Tissue",16)
        time_bin,ttp_bin,ttp_count = scv_compute_atlas_region_stats(
            vess_atlas_mask_im,
            fit_ttp_im,
            fit_ttp_im,
            100,
            report_subprogress)
        graph_label = ["Bin_Num"]
        graph_data = np.arange(len(time_bin))
        graph_label = np.append(graph_label,"TPP")
        graph_data = np.stack((graph_data,time_bin))
        for r in range(1,ttp_bin.shape[0]):
            graph_label = np.append(graph_label,
                "Count_Region_"+str(r))
            graph_data = np.concatenate((graph_data,
                [ttp_count[r,:]]))

    if len(cbf_filename) > 0:
        report_progress("Generating CBF Graphs",94)
        cbf_im = itk.imread(cbf_filename, itk.F)
        Resample = tube.ResampleImage.New(Input=cbf_im)
        Resample.SetMatchImage(vess_atlas_mask_im)
        Resample.Update()
        fit_cbf_im = Resample.GetOutput()
        TubeMath.SetPointValuesFromImage(fit_cbf_im, "CBF")
        TubeMath.SetPointValuesFromTubeRegions(fit_cbf_im,
            "CBF_Tissue",
            1.5,
            4)
        TubeMath.SmoothTubeProperty("CBF",4)
        TubeMath.SmoothTubeProperty("CBF_Tissue",16)
        if ttp_im!=None:
            time_bin,cbf_bin,cbf_count = scv_compute_atlas_region_stats(
                vess_atlas_mask_im,
                fit_ttp_im,
                fit_cbf_im,
                100,
                report_subprogress)
            for r in range(1,cbf_bin.shape[0]):
                graph_label = np.append(graph_label,
                    "CBF_Region"+str(r))
                graph_data = np.concatenate((graph_data,
                    [cbf_bin[r,:]]))

    if len(cbv_filename) > 0:
        report_progress("Generating CBV Graphs",96)
        cbv_im = itk.imread(cbv_filename, itk.F)
        Resample = tube.ResampleImage.New(Input=cbv_im)
        Resample.SetMatchImage(vess_atlas_mask_im)
        Resample.Update()
        fit_cbv_im = Resample.GetOutput()
        TubeMath.SetPointValuesFromImage(fit_cbv_im, "CBV")
        TubeMath.SetPointValuesFromTubeRegions(fit_cbv_im,
            "CBV_Tissue",
            1.5,
            4)
        TubeMath.SmoothTubeProperty("CBV",4)
        TubeMath.SmoothTubeProperty("CBV_Tissue",16)
        if ttp_im!=None:
            time_bin,cbv_bin,cbv_count = scv_compute_atlas_region_stats(
                vess_atlas_mask_im,
                fit_ttp_im,
                fit_cbv_im,
                100,
                report_subprogress)
            for r in range(1,cbv_bin.shape[0]):
                graph_label = np.append(graph_label,
                    "CBV_Region"+str(r))
                graph_data = np.concatenate((graph_data,
                    [cbv_bin[r,:]]))
    if len(tmax_filename) > 0:
        report_progress("Generating TMax Graphs",98)
        tmax_im = itk.imread(tmax_filename, itk.F)
        Resample = tube.ResampleImage.New(Input=tmax_im)
        Resample.SetMatchImage(vess_atlas_mask_im)
        Resample.Update()
        fit_tmax_im = Resample.GetOutput()
        TubeMath.SetPointValuesFromImage(fit_tmax_im, "TMax")
        TubeMath.SetPointValuesFromTubeRegions(fit_tmax_im,
            "TMax_Tissue",
            1.5,
            4)
        TubeMath.SmoothTubeProperty("TMax",4)
        TubeMath.SmoothTubeProperty("TMax_Tissue",16)
        if ttp_im!=None:
            time_bin,tmax_bin,tmax_count = scv_compute_atlas_region_stats(
                vess_atlas_mask_im,
                fit_ttp_im,
                fit_tmax_im,
                100,
                report_subprogress)
            for r in range(1,tmax_bin.shape[0]):
                graph_label = np.append(graph_label,
                    "TMax_Region"+str(r))
                graph_data = np.concatenate((graph_data,
                    [tmax_bin[r,:]]))

    report_progress("Saving results",99)
    SOWriter = itk.SpatialObjectWriter[3].New()
    SOWriter.SetInput(vess_so)
    SOWriter.SetBinaryPoints(True)
    SOWriter.SetFileName(os.path.join(report_out_dirname,
        in_filename_base+"_vessels_extracted_perf.tre"))
    SOWriter.Update()

    VTPWriter = itk.WriteTubesAsPolyData.New()
    VTPWriter.SetInput(vess_so)
    VTPWriter.SetFileName(os.path.join(report_out_dirname,
        in_filename_base+"_vessels_extracted_perf.vtp"))
    VTPWriter.Update()

    csvfilename = os.path.join(report_out_dirname,
        in_filename_base+"_vessels_extracted_perf.csv")
    csvfile = open(csvfilename,'w',newline='')
    csvwriter = csv.writer(csvfile, dialect='excel',
        quoting=csv.QUOTE_NONE)
    csvwriter.writerow(graph_label)
    for r in range(graph_data.shape[1]):
        csvwriter.writerow(['{:f}'.format(x) for x in graph_data[:,r]])
    csvfile.close()
    report_progress("Done",100)

#################
#################
#################
#################
#################
def scv_generate_4d_ctp_vessel_report(ctp_4d_filename,
                                      atlas_path,
                                      report_out_dirname,
                                      report_progress=print,
                                      report_subprogress=print,
                                      debug=False):

    ctp_3d_filenames,ctp_4d_filename,ct_filename, \
        cta_filename,dsa_filename,mask_brain_filename \
            = scv_prepare_4d_for_perfusion_toolbox( \
                ctp_4d_filename, report_out_dirname, \
                report_progress, report_subprogress, debug)

    scv_generate_vessel_report(ctp_3d_filenames, \
        ctp_4d_filename,ct_filename, \
        cta_filename,dsa_filename, \
        mask_brain_filename,atlas_path,report_out_dirname, \
        report_progress,report_subprogress,debug)

#################
#################
#################
#################
#################
def scv_generate_3d_ctp_vessel_report(ctp_3d_filenames,
                                      atlas_path,
                                      report_out_dirname,
                                      report_progress=print,
                                      report_subprogress=print,
                                      debug=False):
    ctp_4d_filename,ct_filename, \
        cta_filename,dsa_filename,mask_brain_filename \
            = scv_prepare_3d_for_perfusion_toolbox( \
                ctp_3d_filenames, report_out_dirname, \
                report_progress, report_subprogress, debug)

    scv_generate_vessel_report(ctp_3d_filenames, \
        ctp_4d_filename,ct_filename, \
        cta_filename,dsa_filename, \
        mask_brain_filename,atlas_path,report_out_dirname, \
        report_progress,report_subprogress,debug)


"""
Module containing customized antenna classes for ARA.

Many of the methods here mirror methods used in the antennas in AraSim, to
ensure that AraSim results can be matched.

"""

import logging
import os.path
import pickle
import numpy as np
import scipy.constants
import scipy.signal
from pyrex.internal_functions import (normalize, complex_bilinear_interp,
                                      complex_interp)
from pyrex.signals import Signal, FunctionSignal
from pyrex.antenna import Antenna
from pyrex.detector import AntennaSystem

logger = logging.getLogger(__name__)

def _read_arasim_antenna_data(filename):
    """
    Gather antenna directionality data from an AraSim-formatted data file.

    The data file should have columns for theta, phi, dB gain, non-dB gain, and
    phase (in degrees). This should be divided into sections for each frequency
    with a header line "freq : X MHz", optionally followed by a second line
    "SWR : Y".

    Parameters
    ----------
    filename : str
        Name of the data file.

    Returns
    -------
    response : ndarray
        3-D array of complex-valued voltage gains as a function of frequency
        along axis 0, zenith along axis 1, and azimuth along axis 2.
    frequencies : ndarray
        Frequencies (Hz) corresponding to axis 0 of `response`.
    thetas : ndarray
        Zenith angles (degrees) corresponding to axis 1 of `response`.
    phis : ndarray
        Azimuth angles (degrees) corresponding to axis 2 of `response`.

    """
    data = {}
    freqs = set()
    thetas = set()
    phis = set()
    freq = 0
    with open(filename) as f:
        for line in f:
            words = line.split()
            if line.startswith('freq'):
                freq = 1
                if words[-1]=="Hz":
                    pass
                elif words[-1]=="kHz":
                    freq *= 1e3
                elif words[-1]=="MHz":
                    freq *= 1e6
                elif words[-1]=="GHz":
                    freq *= 1e9
                else:
                    raise ValueError("Cannot parse line: '"+line+"'")
                freq *= float(words[-2])
                freqs.add(freq)
            elif line.startswith('SWR'):
                swr = float(words[-1])
            elif len(words)==5 and words[0]!="Theta":
                theta = int(words[0])
                thetas.add(theta)
                phi = int(words[1])
                phis.add(phi)
                db_gain = float(words[2])
                # AraSim actually only seems to use the sqrt of the gain
                # (must be gain in power, not voltage)
                # gain = np.sqrt(float(words[3]))
                gain = np.sqrt(10**(db_gain/10))
                phase = np.radians(float(words[4]))
                data[(freq, theta, phi)] = (gain, phase)

    # Convert data dictionary into 3-D array of responses
    response = np.empty((len(freqs), len(thetas), len(phis)),
                        dtype=np.complex_)
    for i, freq in enumerate(sorted(freqs)):
        for j, theta in enumerate(sorted(thetas)):
            for k, phi in enumerate(sorted(phis)):
                gain, phase = data[(freq, theta, phi)]
                response[i, j, k] = gain * np.exp(1j*phase)

    response_data = (response, np.array(sorted(freqs)),
                     np.array(sorted(thetas)), np.array(sorted(phis)))
    return _fix_response_wrapping(response_data)

# If the antenna responses don't go all the way to phi=360, add the extra
# column for the sake of the interpolators
def _fix_response_wrapping(response_data):
    """
    Add phi=360 degrees column to antenna response data.

    The interpolators require that the full azimuth range of the antennas is
    described, so this function duplicates the phi=0 column of the antenna
    response into a phi=360 column, as long as it matches with the rest of
    the phi spacings.

    Parameters
    ----------
    response_data : tuple of ndarray
        Tuple containing the response data for the antenna. The first element
        should contain a 3-D array of the antenna response model as a function
        of frequency (axis 0), zenith (axis 1), and azimuth (axis 2). The
        remaining elements should be the values of the frequency, zenith, and
        azimuth axes, respectively.

    Returns
    -------
    response : ndarray
        Corrected 3-D array of antenna response values, including the phi=360
        degree column if possible.
    frequencies : ndarray
        Frequencies (Hz) corresponding to axis 0 of `response`.
    thetas : ndarray
        Zenith angles (degrees) corresponding to axis 1 of `response`.
    phis : ndarray
        Azimuth angles (degrees) corresponding to axis 2 of `response`.

    """
    response, freqs, thetas, phis = response_data
    if phis[-1]==360:
        return response_data
    if phis[0]==0 and phis[-1]==360-phis[1]:
        phis = np.concatenate((phis, [360]))
        response = np.concatenate((response, response[:, :, 0:1]), axis=2)
    return response, freqs, thetas, phis

def _read_arasim_antenna_pickle(filename):
    """
    Gather antenna directional response data from a pickled data file.

    The data file should be a pickled file containing the antenna directional
    response data from an AraSim-formatted data file as returned by the
    `_read_arasim_antenna_data` function.

    Parameters
    ----------
    filename : str
        Name of the data file without the ``.pkl`` extension.

    Returns
    -------
    response : ndarray
        3-D array of complex-valued voltage gains as a function of frequency
        along axis 0, zenith along axis 1, and azimuth along axis 2.
    frequencies : ndarray
        Frequencies (Hz) corresponding to axis 0 of `response`.
    thetas : ndarray
        Zenith angles (degrees) corresponding to axis 1 of `response`.
    phis : ndarray
        Azimuth angles (degrees) corresponding to axis 2 of `response`.

    See Also
    --------
    _read_arasim_antenna_data : Gather antenna directionality data from an
                                AraSim-formatted data file.

    """
    # Quick fix for filenames with one of the approved extensions already
    # (just strip it)
    if filename.endswith(".txt") or filename.endswith(".pkl"):
        filename = filename[:-4]

    # If there is no pickle file, read the response data using the
    # _read_arasim_antenna_data function, and then make a pickle file
    if not os.path.isfile(filename+".pkl"):
        logger.warning("Antenna model file %s.pkl not found. "+
                       "Generating a new file now", filename)
        response_data = _read_arasim_antenna_data(filename+".txt")
        with open(filename+".pkl", 'wb') as f:
            pickle.dump(response_data, f)
        return response_data

    # Otherwise, read from the pickle file
    else:
        with open(filename+".pkl", 'rb') as f:
            return pickle.load(f)

def _read_filter_data(filename):
    """
    Gather frequency-dependent filtering data from a data file.

    The data file should have columns for frequency, non-dB gain, and phase
    (in radians).

    Parameters
    ----------
    filename : str
        Name of the data file.

    Returns
    -------
    gains : ndarray
        Complex-valued voltage gains as a function of frequency.
    frequencies : ndarray
        Frequencies (Hz) corresponding to the values of `gains`.

    """
    gains = []
    freqs = []
    freq_scale = 0
    with open(filename) as f:
        for line in f:
            words = line.split()
            if line.startswith('Freq'):
                _, scale = words[0].split("(")
                scale = scale.rstrip(")")
                if scale=="Hz":
                    freq_scale = 1
                elif scale=="kHz":
                    freq_scale = 1e3
                elif scale=="MHz":
                    freq_scale = 1e6
                elif scale=="GHz":
                    freq_scale = 1e9
                else:
                    raise ValueError("Cannot parse line: '"+line+"'")
            elif len(words)==3 and words[0]!="Total":
                f, g, p = line.split(",")
                freq = float(f) * freq_scale
                gain = float(g)
                phase = float(p)
                freqs.append(freq)
                gains.append(gain * np.exp(1j*phase))

    return np.array(gains), np.array(freqs)

ARA_DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
VPOL_DATA_FILE = os.path.join(ARA_DATA_DIR,
                              "Vpol_original_CrossFeed_150mmHole_Ice_ARASim.txt")
HPOL_DATA_FILE = os.path.join(ARA_DATA_DIR,
                              "Hpol_original_150mmHole_Ice_ARASim.txt")
FILT_DATA_FILE = os.path.join(ARA_DATA_DIR,
                              "ARA_Electronics_TotalGain_TwoFilters.txt")
# Vpol data file contains only the theta responses
VPOL_THETA_RESPONSE_DATA = _read_arasim_antenna_pickle(VPOL_DATA_FILE)
VPOL_RESPONSE_DATA = (
    VPOL_THETA_RESPONSE_DATA[0],
    np.zeros(VPOL_THETA_RESPONSE_DATA[0].shape),
    *VPOL_THETA_RESPONSE_DATA[1:]
)
# Hpol data file contains only the phi responses
HPOL_PHI_RESPONSE_DATA = _read_arasim_antenna_pickle(HPOL_DATA_FILE)
HPOL_RESPONSE_DATA = (
    np.zeros(HPOL_PHI_RESPONSE_DATA[0].shape),
    *HPOL_PHI_RESPONSE_DATA
)
ALL_FILTERS_DATA = _read_filter_data(FILT_DATA_FILE)

class ARAAntenna(Antenna):
    """
    Antenna class to be used for ARA antennas.

    Stores the attributes of an antenna as well as handling receiving,
    processing, and storing signals and adding noise. Antenna response based on
    provided models.

    Parameters
    ----------
    response_data : tuple of array_like
        Tuple containing the response data for the antenna along the theta
        and phi polarization directions. The first and second elements should
        contain 3-D arrays of the antenna response model in the theta and phi
        polarizations, respectively, as a function of frequency (axis 0),
        zenith (axis 1), and azimuth (axis 2). The remaining elements should be
        the values of the frequency, zenith, and azimuth axes, respectively.
    position : array_like
        Vector position of the antenna.
    center_frequency : float
        Frequency (Hz) at the center of the antenna's frequency range.
    bandwidth : float
        Bandwidth (Hz) of the antenna.
    temperature : float
        The noise temperature (K) of the antenna. Used in combination with
        `resistance` to calculate the RMS voltage of the antenna noise.
    resistance : float
        The noise resistance (ohm) of the antenna. Used in combination with
        `temperature` to calculate the RMS voltage of the antenna noise.
    orientation : array_like, optional
        Vector direction of the z-axis of the antenna.
    efficiency : float, optional
        Antenna efficiency applied to incoming signal values.
    noisy : boolean, optional
        Whether or not the antenna should add noise to incoming signals.
    unique_noise_waveforms : int, optional
        The number of expected noise waveforms needed for each received signal
        to have its own noise.

    Attributes
    ----------
    position : array_like
        Vector position of the antenna.
    z_axis : ndarray
        Vector direction of the z-axis of the antenna.
    x_axis : ndarray
        Vector direction of the x-axis of the antenna.
    antenna_factor : float
        Antenna factor used for converting fields to voltages.
    efficiency : float
        Antenna efficiency applied to incoming signal values.
    noisy : boolean
        Whether or not the antenna should add noise to incoming signals.
    unique_noises : int
        The number of expected noise waveforms needed for each received signal
        to have its own noise.
    freq_range : array_like
        The frequency band in which the antenna operates (used for noise
        production).
    temperature : float or None
        The noise temperature (K) of the antenna. Used in combination with
        `resistance` to calculate the RMS voltage of the antenna noise.
    resistance : float or None
        The noise resistance (ohm) of the antenna. Used in combination with
        `temperature` to calculate the RMS voltage of the antenna noise.
    noise_rms : float or None
        The RMS voltage (V) of the antenna noise. If not ``None``, this value
        will be used instead of the RMS voltage calculated from the values of
        `temperature` and `resistance`.
    signals : list of Signal
        The signals which have been received by the antenna.
    is_hit
    is_hit_mc_truth
    waveforms
    all_waveforms

    See Also
    --------
    pyrex.Antenna : Base class for antennas.

    """
    def __init__(self, response_data, position, center_frequency, bandwidth,
                 temperature, resistance, orientation=(0,0,1), efficiency=1,
                 noisy=True, unique_noise_waveforms=10):
        # Parse the response data
        self._theta_response = response_data[0]
        self._phi_response = response_data[1]
        self._response_freqs = response_data[2]
        self._response_zens = response_data[3]
        self._response_azis = response_data[4]

        # Get the critical frequencies in Hz
        f_low = center_frequency - bandwidth/2
        f_high = center_frequency + bandwidth/2

        # Get arbitrary x-axis orthogonal to orientation
        tmp_vector = np.zeros(3)
        while np.array_equal(np.cross(orientation, tmp_vector), (0,0,0)):
            tmp_vector = np.random.rand(3)
        ortho = np.cross(orientation, tmp_vector)
        # Note: ortho is not normalized, but will be normalized by Antenna's init

        super().__init__(position=position, z_axis=orientation, x_axis=ortho,
                         efficiency=efficiency, freq_range=(f_low, f_high),
                         temperature=temperature, resistance=resistance,
                         noisy=noisy,
                         unique_noise_waveforms=unique_noise_waveforms)

    def directional_gain(self, theta, phi):
        raise NotImplementedError("Directional gain is not defined for "+
                                  self.__class__.__name__+". Use the "+
                                  "directional_response method instead.")

    def polarization_gain(self, polarization):
        raise NotImplementedError("Polarization gain is not defined for "+
                                  self.__class__.__name__+". Use the "+
                                  "directional_response method instead.")

    def directional_response(self, theta, phi, polarization):
        """
        Generate the (complex) frequency-dependent directional response.

        For given angles and polarization direction, use the model of the
        directional and polarization gains of the antenna to generate a
        function for the interpolated response of the antenna with respect to
        frequency. Used with the `frequency_response` method to calculate
        effective heights.

        Parameters
        ----------
        theta : float
            Polar angle (radians) from which a signal is arriving.
        phi : float
            Azimuthal angle (radians) from which a signal is arriving.
        polarization : array_like
            Normalized polarization vector in the antenna coordinate system.

        Returns
        -------
        function
            A function which returns complex-valued voltage gains for given
            frequencies, using the values of incoming angle and polarization.

        See Also
        --------
        ARAAntenna.frequency_response : Calculate the (complex) frequency
                                        response of the antenna.

        """
        e_theta = [np.cos(theta) * np.cos(phi),
                   np.cos(theta) * np.sin(phi),
                   -np.sin(theta)]
        e_phi = [-np.sin(phi), np.cos(phi), 0]
        theta_factor = np.dot(polarization, e_theta)
        phi_factor = np.dot(polarization, e_phi)
        theta_gains = complex_bilinear_interp(
            x=np.degrees(theta), y=np.degrees(phi),
            xp=self._response_zens,
            yp=self._response_azis,
            fp=self._theta_response,
            method='cartesian'
        )
        phi_gains = complex_bilinear_interp(
            x=np.degrees(theta), y=np.degrees(phi),
            xp=self._response_zens,
            yp=self._response_azis,
            fp=self._phi_response,
            method='cartesian'
        )
        freq_interpolator = lambda frequencies: complex_interp(
            x=frequencies, xp=self._response_freqs,
            fp=theta_factor*theta_gains + phi_factor*phi_gains,
            method='euler', outer=0
        )
        return freq_interpolator

    def frequency_response(self, frequencies):
        """
        Calculate the (complex) frequency response of the antenna.

        Rather than handling the entire frequency response of the antenna, this
        method is being used to convert the frequency-dependent gains from the
        `directional_response` method into effective heights.

        Parameters
        ----------
        frequencies : array_like
            1D array of frequencies (Hz) at which to calculate gains.

        Returns
        -------
        array_like
            Complex gains in voltage for the given `frequencies`.

        See Also
        --------
        ARAAntenna.directional_response : Generate the (complex) frequency
                                          dependent directional response.

        """
        # From AraSim GaintoHeight function, with gain calculation moved to
        # the directional_response method.
        # gain=4*pi*A_eff/lambda^2 and h_eff=2*sqrt(A_eff*Z_rx/Z_air)
        # Then 0.5 to calculate power with heff (cancels 2 above)
        heff = np.zeros(len(frequencies))
        # The index of refraction in this calculation should be the index of
        # the ice used in the production of the antenna model.
        n = 1.78
        heff[frequencies!=0] = np.sqrt((scipy.constants.c
                                        /frequencies[frequencies!=0]/n)**2
                                       * n*50/377 /(4*np.pi))
        return heff

    def apply_response(self, signal, direction=None, polarization=None,
                       force_real=True):
        """
        Process the complete antenna response for an incoming signal.

        Processes the incoming signal according to the frequency response of
        the antenna, the efficiency, and the antenna factor. May also apply the
        directionality and the polarization gain depending on the provided
        parameters. Subclasses may wish to overwrite this function if the
        full antenna response cannot be divided nicely into the described
        pieces.

        Parameters
        ----------
        signal : Signal
            Incoming ``Signal`` object to process.
        direction : array_like, optional
            Vector denoting the direction of travel of the signal as it reaches
            the antenna (in the global coordinate frame). If ``None`` no
            directional response will be applied.
        polarization : array_like, optional
            Vector denoting the signal's polarization direction (in the global
            coordinate frame). If ``None`` no polarization gain will be applied.
        force_real : boolean, optional
            Whether or not the frequency response should be redefined in the
            negative-frequency domain to keep the values of the filtered signal
            real.

        Returns
        -------
        Signal
            Processed ``Signal`` object after the complete antenna response has
            been applied. Should have a ``value_type`` of ``voltage``.

        Raises
        ------
        ValueError
            If the given `signal` does not have a ``value_type`` of ``voltage``
            or ``field``.

        See Also
        --------
        pyrex.Signal : Base class for time-domain signals.

        """
        new_signal = signal.copy()
        new_signal.value_type = Signal.Type.voltage
        freq_response = self.frequency_response

        if direction is not None and polarization is not None:
            # Calculate theta and phi relative to the orientation
            origin = self.position - normalize(direction)
            r, theta, phi = self._convert_to_antenna_coordinates(origin)
            # Calculate polarization vector in the antenna coordinates
            y_axis = np.cross(self.z_axis, self.x_axis)
            transformation = np.array([self.x_axis, y_axis, self.z_axis])
            ant_pol = np.dot(transformation, normalize(polarization))
            # Calculate directional response as a function of frequency
            directive_response = self.directional_response(theta, phi, ant_pol)
            freq_response = lambda f: (self.frequency_response(f)
                                       * directive_response(f))

        elif (direction is not None and polarization is None
              or direction is None and polarization is not None):
            raise ValueError("Direction and polarization must be specified together")

        # Apply (combined) frequency response
        new_signal.filter_frequencies(freq_response, force_real=force_real)

        signal_factor = self.efficiency

        if signal.value_type==Signal.Type.voltage:
            pass
        elif signal.value_type==Signal.Type.field:
            signal_factor /= self.antenna_factor
        else:
            raise ValueError("Signal's value type must be either "
                             +"voltage or field. Given "+str(signal.value_type))

        new_signal *= signal_factor

        return new_signal

    # Redefine receive method to use force_real as True by default
    def receive(self, signal, direction=None, polarization=None,
                force_real=True):
        """
        Process and store one or more incoming (polarized) signals.

        Processes the incoming signal(s) according to the ``apply_response``
        method, then stores the total processed signal to the signals list. If
        more than one signal is given, they should be logically connected as
        separately polarized portions of the same signal.

        Parameters
        ----------
        signal : Signal or array_like
            Incoming ``Signal`` object(s) to process and store. May be separate
            polarization representations, but therefore should have the same
            times.
        direction : array_like, optional
            Vector denoting the direction of travel of the signal(s) as they
            reach the antenna (in the global coordinate frame). If ``None`` no
            directional gain will be applied.
        polarization : array_like, optional
            Vector(s) denoting the signal's polarization direction (in the
            global coordinate frame). Number of vectors should match the number
            of elements in `signal` argument. If ``None`` no polarization gain
            will be applied.
        force_real : boolean, optional
            Whether or not the frequency response should be redefined in the
            negative-frequency domain to keep the values of the filtered signal
            real.

        Raises
        ------
        ValueError
            If the number of polarizations does not match the number of signals.
            Or if the signals do not have the same `times` array.

        See Also
        --------
        pyrex.Signal : Base class for time-domain signals.

        """
        super().receive(signal=signal, direction=direction,
                        polarization=polarization,
                        force_real=force_real)

class ARAAntennaSystem(AntennaSystem):
    """
    Antenna system extending base ARA antenna with front-end processing.

    Applies as the front end a filter representing the full ARA electronics
    chain (including amplification) and signal clipping. Additionally provides
    a method for passing a signal through the tunnel diode.

    Parameters
    ----------
    response_data : tuple of array_like
        Tuple containing the response data for the antenna along the theta
        and phi polarization directions. The first and second elements should
        contain 3-D arrays of the antenna response model in the theta and phi
        polarizations, respectively, as a function of frequency (axis 0),
        zenith (axis 1), and azimuth (axis 2). The remaining elements should be
        the values of the frequency, zenith, and azimuth axes, respectively.
    name : str
        Name of the antenna.
    position : array_like
        Vector position of the antenna.
    power_threshold : float
        Power threshold for trigger condition. Antenna triggers if a signal
        passed through the tunnel diode exceeds this threshold times the noise
        RMS of the tunnel diode.
    orientation : array_like, optional
        Vector direction of the z-axis of the antenna.
    amplification : float, optional
        Amplification to be applied to the signal pre-clipping. Note that the
        usual ARA electronics amplification is already applied without this.
    amplifier_clipping : float, optional
        Voltage (V) above which the amplified signal is clipped (in positive
        and negative values).
    noisy : boolean, optional
        Whether or not the antenna should add noise to incoming signals.
    unique_noise_waveforms : int, optional
        The number of expected noise waveforms needed for each received signal
        to have its own noise.

    Attributes
    ----------
    antenna : Antenna
        ``Antenna`` object extended by the front end.
    name : str
        Name of the antenna.
    position : array_like
        Vector position of the antenna.
    power_threshold : float
        Power threshold for trigger condition. Antenna triggers if a signal
        passed through the tunnel diode exceeds this threshold times the noise
        RMS of the tunnel diode.
    amplification : float
        Amplification to be applied to the signal pre-clipping. Note that the
        usual ARA electronics amplification is already applied without this.
    amplifier_clipping : float
        Voltage (V) above which the amplified signal is clipped (in positive
        and negative values).
    lead_in_time : float
        Lead-in time (s) required for the front end to equilibrate.
        Automatically added in before calculation of signals and waveforms.
    is_hit
    is_hit_mc_truth
    signals
    waveforms
    all_waveforms

    See Also
    --------
    pyrex.AntennaSystem : Base class for antenna system with front-end
                          processing.
    ARAAntenna : Antenna class to be used for ARA antennas.

    """
    lead_in_time = 5e-9

    def __init__(self, response_data, name, position, power_threshold,
                 orientation=(0,0,1), amplification=1, amplifier_clipping=1,
                 noisy=True, unique_noise_waveforms=10,
                 **kwargs):
        super().__init__(ARAAntenna)

        self.name = str(name)
        self.position = position

        self.amplification = amplification
        self.amplifier_clipping = amplifier_clipping

        self.setup_antenna(response_data=response_data,
                           orientation=orientation, noisy=noisy,
                           unique_noise_waveforms=unique_noise_waveforms,
                           **kwargs)

        self.power_threshold = power_threshold
        self._power_mean = None
        self._power_std = None

        self._filter_response = ALL_FILTERS_DATA[0]
        self._filter_freqs = ALL_FILTERS_DATA[1]

    @property
    def _metadata(self):
        """Metadata dictionary for writing `ARAAntennaSystem` information."""
        meta = super()._metadata
        meta.update({
            "name": self.name,
            "lead_in_time": self.lead_in_time,
            "amplification": self.amplification,
            "amplifier_clipping": self.amplifier_clipping,
            "power_threshold": self.power_threshold,
        })
        return meta

    def setup_antenna(self, response_data, center_frequency=500e6,
                      bandwidth=800e6, temperature=325, resistance=50,
                      orientation=(0,0,1), efficiency=1, noisy=True,
                      unique_noise_waveforms=10, **kwargs):
        """
        Setup the antenna by passing along its init arguments.

        Any arguments passed to this method are directly passed to the
        ``__init__`` methods of the ``antenna``'s class.

        Parameters
        ----------
        response_data : tuple of array_like
            Tuple containing the response data for the antenna along the theta
            and phi polarization directions. The first and second elements
            should contain 3-D arrays of the antenna response model in the
            theta and phi polarizations, respectively, as a function of
            frequency (axis 0), zenith (axis 1), and azimuth (axis 2). The
            remaining elements should be the values of the frequency, zenith,
            and azimuth axes, respectively.
        center_frequency : float, optional
            Frequency (Hz) at the center of the antenna's frequency range.
        bandwidth : float, optional
            Bandwidth (Hz) of the antenna.
        temperature : float, optional
            The noise temperature (K) of the antenna. Used in combination with
            `resistance` to calculate the RMS voltage of the antenna noise.
        resistance : float, optional
            The noise resistance (ohm) of the antenna. Used in combination with
            `temperature` to calculate the RMS voltage of the antenna noise.
        orientation : array_like, optional
            Vector direction of the z-axis of the antenna.
        efficiency : float, optional
            Antenna efficiency applied to incoming signal values.
        noisy : boolean, optional
            Whether or not the antenna should add noise to incoming signals.
        unique_noise_waveforms : int, optional
            The number of expected noise waveforms needed for each received
            signal to have its own noise.

        """
        # Noise rms should be about 40 mV (after filtering with gain of ~5000).
        # This is mostly satisfied by using the default noise temperature from
        # AraSim, 325 K, along with a 50 ohm resistance
        # Additionally, the bandwidth of the antenna is set slightly larger
        # than the nominal bandwidth of the true ARA antenna system (700 MHz),
        # but the extra frequencies should be killed by the front-end filter
        super().setup_antenna(response_data=response_data,
                              position=self.position,
                              center_frequency=center_frequency,
                              bandwidth=bandwidth,
                              temperature=temperature,
                              resistance=resistance,
                              orientation=orientation,
                              efficiency=efficiency,
                              noisy=noisy,
                              unique_noise_waveforms=unique_noise_waveforms,
                              **kwargs)

    # Tunnel diode response functions pulled from arasim
    _td_args = {
        'down1': (-0.8, 15e-9, 2.3e-9, 0),
        'down2': (-0.2, 15e-9, 4e-9, 0),
        'up': (1, 18e-9, 7e-9, 1e9)
    }
    # Set td_args['up'][0] based on the other args, like in arasim
    _td_args['up'] = (-np.sqrt(2*np.pi) *
                      (_td_args['down1'][0]*_td_args['down1'][2] +
                       _td_args['down2'][0]*_td_args['down2'][2]) /
                      (2e18*_td_args['up'][2]**3),) + _td_args['up'][1:]

    # Set "down" and "up" functions as in arasim
    @classmethod
    def _td_fdown1(cls, x):
        return (cls._td_args['down1'][3] + cls._td_args['down1'][0] *
                np.exp(-(x-cls._td_args['down1'][1])**2 /
                       (2*cls._td_args['down1'][2]**2)))

    @classmethod
    def _td_fdown2(cls, x):
        return (cls._td_args['down2'][3] + cls._td_args['down2'][0] *
                np.exp(-(x-cls._td_args['down2'][1])**2 /
                       (2*cls._td_args['down2'][2]**2)))

    @classmethod
    def _td_fup(cls, x):
        return (cls._td_args['up'][0] *
                (cls._td_args['up'][3] * (x-cls._td_args['up'][1]))**2 *
                np.exp(-(x-cls._td_args['up'][1])/cls._td_args['up'][2]))

    def tunnel_diode(self, signal):
        """
        Calculate a signal as processed by the tunnel diode.

        The given signal is convolved with the tunnel diode response as in
        AraSim.

        Parameters
        ----------
        signal : Signal
            Signal to be processed by the tunnel diode.

        Returns
        -------
        Signal
            Signal output of the tunnel diode for the input `signal`.

        Raises
        ------
        ValueError
            If the input `signal` doesn't have a ``value_type`` of ``voltage``.

        Notes
        -----
        The tunnel diode response is based on the response parameterized in
        AraSim, as developed by ANITA [1]_.

        References
        ----------
        .. [1] A. Connolly & R. Nichol, ANITA Note #411, "A Power-Based Time
            Domain Trigger Simulation."
            https://elog.phys.hawaii.edu/elog/anita_notes/080827_041639/powertrigger.pdf

        """
        if signal.value_type!=Signal.Type.voltage:
            raise ValueError("Tunnel diode only accepts voltage signals")
        t_max = 1e-7
        n_pts = int(t_max/signal.dt)
        times = np.linspace(0, t_max, n_pts+1)
        diode_resp = self._td_fdown1(times) + self._td_fdown2(times)
        t_slice = times>self._td_args['up'][1]
        diode_resp[t_slice] += self._td_fup(times[t_slice])
        conv = scipy.signal.convolve(signal.values**2 / self.antenna.resistance,
                                     diode_resp, mode='full')
        # Signal class will automatically only take the first part of conv,
        # which is what we want.
        # conv multiplied by dt so that the amplitude stays constant for
        # varying dts (determined empirically, see ARZAskaryanSignal comments)
        output = Signal(signal.times, conv*signal.dt,
                        value_type=Signal.Type.power)
        return output

    def interpolate_filter(self, frequencies):
        """
        Generate interpolated filter values for given frequencies.

        Calculate the interpolated values of the antenna system's filter gain
        data for some frequencies.

        Parameters
        ----------
        frequencies : array_like
            1D array of frequencies (Hz) at which to calculate gains.

        Returns
        -------
        array_like
            Complex filter gain in voltage for the given `frequencies`.

        """
        return complex_interp(
            x=frequencies, xp=self._filter_freqs, fp=self._filter_response,
            method='euler', outer=0
        )

    def front_end(self, signal):
        """
        Apply front-end processes to a signal and return the output.

        The front-end consists of the full ARA electronics chain (including
        amplification) and signal clipping.

        Parameters
        ----------
        signal : Signal
            ``Signal`` object on which to apply the front-end processes.

        Returns
        -------
        Signal
            Signal processed by the antenna front end.

        """
        base_signal = signal.copy()
        base_signal.filter_frequencies(self.interpolate_filter,
                                       force_real=True)
        # Apply sqrt(2) for 3dB splitter for TURF, SURF
        base_signal *= self.amplification / np.sqrt(2)
        clip_values = lambda times: np.clip(
            base_signal.with_times(times).values,
            a_min=-self.amplifier_clipping,
            a_max=self.amplifier_clipping
        )
        return FunctionSignal(signal.times, clip_values,
                              value_type=signal.value_type)

    def trigger(self, signal):
        """
        Check if the antenna system triggers on a given signal.

        Passes the signal through the tunnel diode. Then compares the maximum
        and minimum values to a tunnel diode noise signal. Triggers if one of
        the maximum or minimum values exceed the noise mean +/- the noise rms
        times the power threshold.

        Parameters
        ----------
        signal : Signal
            ``Signal`` object on which to test the trigger condition.

        Returns
        -------
        boolean
            Whether or not the antenna triggers on `signal`.

        """
        if self._power_mean is None or self._power_std is None:
            # Prepare for antenna trigger by finding mean and standard
            # deviation of the full noise waveform convolved with the tunnel
            # diode response
            if len(self.antenna.signals)>0:
                times = self.antenna.signals[0].times
            else:
                times = signal.times
            n = len(times)
            dt = times[1]-times[0]
            duration = times[-1]-times[0] + dt
            full_times = np.linspace(0, duration*self.antenna.unique_noises,
                                     n*self.antenna.unique_noises)
            if self.antenna._noise_master is None:
                # Make sure the noise_master has the appropriate length
                # (automatically gets set to N*len(times) the first time it is
                # called, so make sure the first `times` is not the expanded
                # array but the single-signal array)
                self.antenna.make_noise(times)
            long_noise = self.antenna.make_noise(full_times)
            power_noise = self.tunnel_diode(self.front_end(long_noise))
            self._power_mean = np.mean(power_noise.values)
            self._power_std = np.std(power_noise.values)

        power_signal = self.tunnel_diode(signal)
        # Use the absolute value of the power_threshold value so that the value
        # can be specified as positive or negative (compatible with AraSim
        # which only works with negative values, resulting in some confusion)
        low_trigger = (self._power_mean -
                       self._power_std*np.abs(self.power_threshold))
        return np.min(power_signal.values)<low_trigger

    # Redefine receive method to use force_real as True by default
    def receive(self, signal, direction=None, polarization=None,
                force_real=True):
        """
        Process and store one or more incoming (polarized) signals.

        Processes the incoming signal(s) according to the ``apply_response``
        method, then stores the total processed signal to the signals list. If
        more than one signal is given, they should be logically connected as
        separately polarized portions of the same signal.

        Parameters
        ----------
        signal : Signal or array_like
            Incoming ``Signal`` object(s) to process and store. May be separate
            polarization representations, but therefore should have the same
            times.
        direction : array_like, optional
            Vector denoting the direction of travel of the signal(s) as they
            reach the antenna (in the global coordinate frame). If ``None`` no
            directional gain will be applied.
        polarization : array_like, optional
            Vector(s) denoting the signal's polarization direction (in the
            global coordinate frame). Number of vectors should match the number
            of elements in `signal` argument. If ``None`` no polarization gain
            will be applied.
        force_real : boolean, optional
            Whether or not the frequency response should be redefined in the
            negative-frequency domain to keep the values of the filtered signal
            real.

        Raises
        ------
        ValueError
            If the number of polarizations does not match the number of signals.
            Or if the signals do not have the same `times` array.

        See Also
        --------
        pyrex.Signal : Base class for time-domain signals.

        """
        super().receive(signal=signal, direction=direction,
                        polarization=polarization,
                        force_real=force_real)

class HpolAntenna(ARAAntennaSystem):
    """
    ARA Hpol ("quad-slot") antenna system with front-end processing.

    Applies as the front end a filter representing the full ARA electronics
    chain (including amplification) and signal clipping. Additionally provides
    a method for passing a signal through the tunnel diode.

    Parameters
    ----------
    name : str
        Name of the antenna.
    position : array_like
        Vector position of the antenna.
    power_threshold : float
        Power threshold for trigger condition. Antenna triggers if a signal
        passed through the tunnel diode exceeds this threshold times the noise
        RMS of the tunnel diode.
    amplification : float, optional
        Amplification to be applied to the signal pre-clipping. Note that the
        usual ARA electronics amplification is already applied without this.
    amplifier_clipping : float, optional
        Voltage (V) above which the amplified signal is clipped (in positive
        and negative values).
    noisy : boolean, optional
        Whether or not the antenna should add noise to incoming signals.
    unique_noise_waveforms : int, optional
        The number of expected noise waveforms needed for each received signal
        to have its own noise.

    Attributes
    ----------
    antenna : Antenna
        ``Antenna`` object extended by the front end.
    name : str
        Name of the antenna.
    position : array_like
        Vector position of the antenna.
    power_threshold : float
        Power threshold for trigger condition. Antenna triggers if a signal
        passed through the tunnel diode exceeds this threshold times the noise
        RMS of the tunnel diode.
    amplification : float
        Amplification to be applied to the signal pre-clipping. Note that the
        usual ARA electronics amplification is already applied without this.
    amplifier_clipping : float
        Voltage (V) above which the amplified signal is clipped (in positive
        and negative values).
    is_hit
    is_hit_mc_truth
    signals
    waveforms
    all_waveforms

    See Also
    --------
    ARAAntennaSystem : Antenna system extending base ARA antenna with front-end
                       processing.

    """
    def __init__(self, name, position, power_threshold,
                 amplification=1, amplifier_clipping=1, noisy=True,
                 unique_noise_waveforms=10):
        super().__init__(response_data=HPOL_RESPONSE_DATA,
                         name=name, position=position,
                         power_threshold=power_threshold,
                         orientation=(0,0,1),
                         amplification=amplification,
                         amplifier_clipping=amplifier_clipping,
                         noisy=noisy,
                         unique_noise_waveforms=unique_noise_waveforms)

class VpolAntenna(ARAAntennaSystem):
    """
    ARA Vpol ("bicone" or "birdcage") antenna system with front-end processing.

    Applies as the front end a filter representing the full ARA electronics
    chain (including amplification) and signal clipping. Additionally provides
    a method for passing a signal through the tunnel diode.

    Parameters
    ----------
    name : str
        Name of the antenna.
    position : array_like
        Vector position of the antenna.
    power_threshold : float
        Power threshold for trigger condition. Antenna triggers if a signal
        passed through the tunnel diode exceeds this threshold times the noise
        RMS of the tunnel diode.
    amplification : float, optional
        Amplification to be applied to the signal pre-clipping. Note that the
        usual ARA electronics amplification is already applied without this.
    amplifier_clipping : float, optional
        Voltage (V) above which the amplified signal is clipped (in positive
        and negative values).
    noisy : boolean, optional
        Whether or not the antenna should add noise to incoming signals.
    unique_noise_waveforms : int, optional
        The number of expected noise waveforms needed for each received signal
        to have its own noise.

    Attributes
    ----------
    antenna : Antenna
        ``Antenna`` object extended by the front end.
    name : str
        Name of the antenna.
    position : array_like
        Vector position of the antenna.
    power_threshold : float
        Power threshold for trigger condition. Antenna triggers if a signal
        passed through the tunnel diode exceeds this threshold times the noise
        RMS of the tunnel diode.
    amplification : float
        Amplification to be applied to the signal pre-clipping. Note that the
        usual ARA electronics amplification is already applied without this.
    amplifier_clipping : float
        Voltage (V) above which the amplified signal is clipped (in positive
        and negative values).
    is_hit
    is_hit_mc_truth
    signals
    waveforms
    all_waveforms

    See Also
    --------
    ARAAntennaSystem : Antenna system extending base ARA antenna with front-end
                       processing.

    """
    def __init__(self, name, position, power_threshold,
                 amplification=1, amplifier_clipping=1, noisy=True,
                 unique_noise_waveforms=10):
        super().__init__(response_data=VPOL_RESPONSE_DATA,
                         name=name, position=position,
                         power_threshold=power_threshold,
                         orientation=(0,0,1),
                         amplification=amplification,
                         amplifier_clipping=amplifier_clipping,
                         noisy=noisy,
                         unique_noise_waveforms=unique_noise_waveforms)


# connected component analysis of segmentation compared to ground truth
import numpy as np
from scipy import ndimage

_thresh = 0.50
_smooth = 1

_LESION_UNC_VA = {
    'bald': {'ymin': -304332.435036, 'ymax': -165.66824439, 'xmin': -287025.719183, 'xmax': -60.0155142638},
    'ent': {'ymin': -264211.01785, 'ymax': -120.603669105, 'xmin': -246545.542156, 'xmax': -42.3070650557},
    'prdvar': {'ymin': -77261.4, 'ymax': -63.0317, 'xmin': -73166.6, 'xmax': -32.3993},
    'varmcs': {'ymin': -357833.52626, 'ymax': -254.615228264, 'xmin': -340539.075182, 'xmax': -92.6999732885}}

def remove_tiny_les(lesion_image, nvox=2):
    labels, nles = ndimage.label(lesion_image)
    for i in range(1, nles + 1):
        nb_vox = np.sum(lesion_image[labels == i])
        if nb_vox <= nvox:
            lesion_image[labels == i] = 0
    return lesion_image

def global_dice(h, t):
    h = h.flatten()
    t = t.flatten()
    intersection = np.sum(h * t)
    union = np.sum(h) + np.sum(t)
    dice = (2. * intersection + _smooth) / (union + _smooth)
    return dice

def get_lesion_bin(nvox):
    if 3 <= nvox <= 10:
        return 'small'
    elif 11 <= nvox <= 50:
        return 'med'
    elif nvox >= 51:
        return 'large'
    else:
        return "small"

def cca_img_bin(h, t, u, uth, metric):
    t = remove_tiny_les(t, nvox=2)
    h0 = h.copy()
    h = ndimage.binary_dilation(h, structure=ndimage.generate_binary_structure(3, 2))
    t_unc_labels = get_unc_labels(t, u, metric, 'y')
    h_unc_labels = get_unc_labels(h0, u, metric, 'x')

    labels = {}
    nles = {}
    labels['h'], nles['h'] = ndimage.label(h)
    labels['t'], nles['t'] = ndimage.label(t)
    found_h = np.ones(nles['h'], np.int16)
    ntp = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nfp = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nfn = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nb_les = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nles_gt = {'all': nles['t'], 'small': 0, 'med': 0, 'large': 0}
    for i in range(1, nles['t'] + 1):
        t_unc = np.max(t_unc_labels[labels['t'] == i])
        lesion_size = np.sum(t[labels['t'] == i])
        nles_gt[get_lesion_bin(lesion_size)] += 1
        if t_unc < uth:
            # list of detected lesions in this area
            h_lesions = np.unique(labels['h'][labels['t'] == i])
            # all the voxels in this area contribute to detecting the lesion
            nb_overlap = h[labels['t'] == i].sum()
            nb_les[get_lesion_bin(lesion_size)] += 1
            if nb_overlap >= 3 or nb_overlap >= 0.5 * lesion_size:
                ntp[get_lesion_bin(lesion_size)] += 1
                for h_lesion in h_lesions:
                    if h_lesion != 0:
                        found_h[h_lesion - 1] = 0
            else:
                nfn[get_lesion_bin(lesion_size)] += 1

    for i in range(1, nles['h'] + 1):
        if found_h[i - 1] == 1:
            h_unc = np.max(h_unc_labels[labels['h'] == i])
            if h_unc < uth:
                nb_vox = np.sum(h0[labels['h'] == i])
                nfp[get_lesion_bin(nb_vox)] += 1

    nb_les['all'] = nb_les['small'] + nb_les['med'] + nb_les['large']
    ntp['all'] = ntp['small'] + ntp['med'] + ntp['large']
    nfp['all'] = nfp['small'] + nfp['med'] + nfp['large']
    nfn['all'] = nfn['small'] + nfn['med'] + nfn['large']

    tpr = {}
    fdr = {}
    for s in ntp.keys():
        # tpr (sensitivity)
        if nb_les[s] != 0:
            tpr[s] = ntp[s] / nb_les[s]
        elif nb_les[s] == 0 and ntp[s] == 0:
            tpr[s] = 1
        else:
            tpr[s] = 0
        # ppv (1-fdr)
        if ntp[s] + nfp[s] != 0:
            ppv = ntp[s] / (ntp[s] + nfp[s])
        else:
            ppv = 1
        fdr[s] = 1 - ppv

    return {'ntp': ntp, 'nfp': nfp, 'nfn': nfn, 'fdr': fdr, 'tpr': tpr, 'nles': nb_les, 'nles_gt': nles_gt}

def cca_img(h, t, u, uth, metric):
    t = remove_tiny_les(t, nvox=2)
    h0 = h.copy()
    h = ndimage.binary_dilation(h, structure=ndimage.generate_binary_structure(3, 2))
    t_unc_labels = get_unc_labels(t, u, metric, 'y')
    h_unc_labels = get_unc_labels(h0, u, metric, 'x')

    labels = {}
    nles = {}
    labels['h'], nles['h'] = ndimage.label(h)
    labels['t'], nles['t'] = ndimage.label(t)
    found_h = np.ones(nles['h'], np.int16)
    ntp = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nfp = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nfn = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nb_les = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nles_gt = {'all': nles['t'], 'small': 0, 'med': 0, 'large': 0}
    for i in range(1, nles['t'] + 1):
        t_unc = np.max(t_unc_labels[labels['t'] == i])
        lesion_size = np.sum(t[labels['t'] == i])
        nles_gt[get_lesion_bin(lesion_size)] += 1
        if t_unc < uth:
            # list of detected lesions in this area
            h_lesions = np.unique(labels['h'][labels['t'] == i])
            # all the voxels in this area contribute to detecting the lesion
            nb_overlap = h[labels['t'] == i].sum()
            nb_les[get_lesion_bin(lesion_size)] += 1
            if nb_overlap >= 3 or nb_overlap >= 0.5 * lesion_size:
                ntp[get_lesion_bin(lesion_size)] += 1
                for h_lesion in h_lesions:
                    if h_lesion != 0:
                        found_h[h_lesion - 1] = 0
            else:
                nfn[get_lesion_bin(lesion_size)] += 1

    for i in range(1, nles['h'] + 1):
        if found_h[i - 1] == 1:
            h_unc = np.max(h_unc_labels[labels['h'] == i])
            if h_unc < uth:
                nb_vox = np.sum(h0[labels['h'] == i])
                nfp[get_lesion_bin(nb_vox)] += 1

    nb_les['all'] = nb_les['small'] + nb_les['med'] + nb_les['large']
    ntp['all'] = ntp['small'] + ntp['med'] + ntp['large']
    nfp['all'] = nfp['small'] + nfp['med'] + nfp['large']
    nfn['all'] = nfn['small'] + nfn['med'] + nfn['large']

    tpr = {}
    fdr = {}
    for s in ntp.keys():
        # tpr (sensitivity)
        if nb_les[s] != 0:
            tpr[s] = ntp[s] / nb_les[s]
        elif nb_les[s] == 0 and ntp[s] == 0:
            tpr[s] = 1
        else:
            tpr[s] = 0
        # ppv (1-fdr)
        if ntp[s] + nfp[s] != 0:
            ppv = ntp[s] / (ntp[s] + nfp[s])
        else:
            ppv = 1
        fdr[s] = 1 - ppv

    return {'ntp': ntp, 'nfp': nfp, 'nfn': nfn, 'fdr': fdr, 'tpr': tpr, 'nles': nb_les, 'nles_gt': nles_gt}

def cca_img_no_unc(bbox, t, th):
    """
    Connected component analysis of between prediction `h` and ground truth `t` across lesion bin sizes.
    :param h: network output on range [0,1], shape=(NxMxO)
    :type h: float16, float32, float64
    :param t: ground truth labels, shape=(NxMxO)
    :type t: int16
    :param th: threshold to binarize prediction `h`
    :type th: float16, float32, float64
    :return: dict
    """
    h[h >= th] = 1
    h[h < th] = 0
    h = h.astype(np.int16)

    t = remove_tiny_les(t, nvox=2)
    h = ndimage.binary_dilation(h, structure=ndimage.generate_binary_structure(3, 2))

    labels = {}
    nles = {}
    labels['h'], nles['h'] = ndimage.label(h)
    labels['t'], nles['t'] = ndimage.label(t)
    found_h = np.ones(nles['h'], np.int16)
    ntp = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nfp = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nfn = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nb_les = {'all': 0, 'small': 0, 'med': 0, 'large': 0}
    nles_gt = {'all': nles['t'], 'small': 0, 'med': 0, 'large': 0}
    for i in range(1, nles['t'] + 1):
        lesion_size = np.sum(t[labels['t'] == i])
        nles_gt[get_lesion_bin(lesion_size)] += 1
        # list of detected lesions in this area
        h_lesions = np.unique(labels['h'][labels['t'] == i])
        # all the voxels in this area contribute to detecting the lesion
        nb_overlap = h[labels['t'] == i].sum()
        nb_les[get_lesion_bin(lesion_size)] += 1
        if nb_overlap >= 3 or nb_overlap >= 0.5 * lesion_size:
            ntp[get_lesion_bin(lesion_size)] += 1
            for h_lesion in h_lesions:
                if h_lesion != 0:
                    found_h[h_lesion - 1] = 0
        else:
            nfn[get_lesion_bin(lesion_size)] += 1

    for i in range(1, nles['h'] + 1):
        if found_h[i - 1] == 1:
            nb_vox = np.sum(h[labels['h'] == i])
            nfp[get_lesion_bin(nb_vox)] += 1

    nb_les['all'] = nb_les['small'] + nb_les['med'] + nb_les['large']
    ntp['all'] = ntp['small'] + ntp['med'] + ntp['large']
    nfp['all'] = nfp['small'] + nfp['med'] + nfp['large']
    nfn['all'] = nfn['small'] + nfn['med'] + nfn['large']

    tpr = {}
    fdr = {}
    for s in ntp.keys():
        # tpr (sensitivity)
        if nb_les[s] != 0:
            tpr[s] = ntp[s] / nb_les[s]
        elif nb_les[s] == 0 and ntp[s] == 0:
            tpr[s] = 1
        else:
            tpr[s] = 0
        # ppv (1-fdr)
        if ntp[s] + nfp[s] != 0:
            ppv = ntp[s] / (ntp[s] + nfp[s])
        else:
            ppv = 1
        fdr[s] = 1 - ppv

    return {'ntp': ntp, 'nfp': nfp, 'nfn': nfn, 'fdr': fdr, 'tpr': tpr, 'nles': nb_les, 'nles_gt': nles_gt}

def get_unc_labels(x, unc, metric, xy):
    x_big = ndimage.binary_dilation(x, structure=ndimage.generate_binary_structure(3, 2))
    labels, nles = ndimage.label(x_big)
    unc_labels = np.zeros_like(unc)
    for i in range(1, nles + 1):
        unc_labels[labels == i] = np.sum(np.log(unc[labels == i] + 1e-5))

    unc_labels[unc_labels != 0] = (unc_labels[unc_labels != 0] - _LESION_UNC_VA[metric][xy + 'min']) / (
        _LESION_UNC_VA[metric][xy + 'max'] - _LESION_UNC_VA[metric][xy + 'min'])
    return unc_labels

def paint_cca_img(h, t):
    TP_COLOUR=1
    FP_COLOUR=2
    FN_COLOUR=3

    h_paint = h.copy() # this should be the un-connected labels

    t = remove_tiny_les(t)

    # get the 18 neighbourhood of the hypothesis
    neighbourhood18 = ndimage.generate_binary_structure(3, 2)
    h = ndimage.binary_dilation(h, structure=neighbourhood18)

    labels = {}
    nles = {}
    labels['h'], nles['h'] = ndimage.label(h)
    labels['t'], nles['t'] = ndimage.label(t)
    ntp = 0; nfp = 0; nfn = 0
    found_h = np.ones(nles['h'], np.int16)

    for i in range(1,nles['t']+1):
        got_les = False
        h_lesions = np.unique(labels['h'][labels['t']==i])
        for h_lesion in h_lesions:
            if h_lesion != 0:
                nb_overlap = h[labels['h']==h_lesion].sum()
                if nb_overlap >= 3 or nb_overlap >= 0.5*np.sum(t[labels['t']==i]):
                    # this is a a true positive
                    h_paint[labels['h']==h_lesion] *= TP_COLOUR
                    found_h[h_lesion-1] = 0
                    got_les = True
        # if we didn't get it --> false negative
        if not got_les:
            h_paint[labels['t']==i] = FN_COLOUR

    # any remaining that we found are false positivse
    for i, fp in enumerate(found_h):
        if fp:
            h_paint[labels['h']==(i+1)] *= FP_COLOUR

    labels_hpaint, nles = ndimage.label(h_paint)
    for i in range(1,nles+1):
        nb_vox = np.size(labels_hpaint[labels_hpaint==i])
        if nb_vox < 3:
            h_paint[labels_hpaint==i] = 0
    return h_paint

def ohe(x):
    ohe = np.zeros_like(x)
    ohe = np.repeat(np.expand_dims(x,-1),3,-1)
    ohe[:,:,:,0][x!=2] = 0 # red --> FP(2)
    ohe[:,:,:,1][x!=1] = 0 # gre --> TP(1)
    ohe[:,:,:,2][x!=3] = 0 # blu --> FN(3)
    ohe[ohe>0]=1
    return ohe

# def cca_img(h, t, thresh, read_img=True):
#     """
#     1. binarize segmentation 'h' based on sigmoid 0.5 thresh
#     2. generate blob (18 neighbourhood) of the 'h'
#     3. remove tiny lesions in truth 't' (<=2 voxels in size)
#     4. count true positives:
#         for lesion in t:
#             overlap  = _get_overlap(h, t)
#             if overap >= voxels:
#                 tp ++
#             else if ( overlap > 0.5 * size(lesion))
#                 tp ++
#             else:
#                 fn ++

#     5. count false positives:
#         for lesion in h:
#             overlap  = _get_overlap(h, t)
#             if overlap == 0:
#                 fp ++
#     """
#     if read_img:
#         h, header = nrrd.read(h)
#         t, header = nrrd.read(t)

#     h[h < thresh] = 0
#     h[h >= thresh] = 1
#     t = t.astype(np.int16)
#     h = h.astype(np.int16)

#     t = remove_tiny_les(t)

#     # get the 18 neighbourhood of the hypothesis
#     neighbourhood18 = ndimage.generate_binary_structure(3, 2)
#     h = ndimage.binary_dilation(h, structure=neighbourhood18)

#     labels = {}
#     nles = {}
#     labels['h'], nles['h'] = ndimage.label(h)
#     labels['t'], nles['t'] = ndimage.label(t)

#     ntp = 0; nfp = 0; nfn = 0
#     found_h = np.ones(nles['h'], np.int16)

#     for i in range(1,nles['t']):
#         got_les = False
#         h_lesions = np.unique(labels['h'][labels['t']==i])
#         for h_lesion in h_lesions:
#             if h_lesion != 0:
#                 nb_overlap = np.size(labels['h']==h_lesion)
#                 if nb_overlap >= 3 or nb_overlap >= 0.5*np.sum(t[labels['t']==i]):
#                     got_les = True
#                     found_h[h_lesion-1] = 0
#         if got_les:
#             ntp += 1
#         else:
#             nfn += 1

#     nfp = np.sum(found_h)

#     # tpr
#     if nles['t']!= 0:
#         tpr = ntp / nles['t']
#     elif nles['t']==0 and ntp ==0:
#         tpr = 1
#     else:
#         tpr = 0

#     # ppv, fdr
#     if ntp+nfp != 0:
#         ppv = ntp / (ntp+nfp)
#     else:
#         ppv = 1

#     # fdr
#     # if ntp+nfp != 0:
#     #     fdr = nfp / (ntp+nfp)
#     # else:
#     #     fdr = 0

#     # f1 score
#     # denom = (2*ntp + nfp + nfn)
#     # if denom != 0:
#     #     f1_score = 2*ntp / denom
#     # else:
#     #     f1_score = 1

#     return {'ntp':ntp, 'nfp':nfp, 'nfn':nfn, 'ppv':ppv, 'tpr':tpr, 'fdr':1-ppv}