
Licence = """
<OWNER> = Barak Zackay (Weizmann Institute of Science)
<YEAR> = 2014

In the original BSD license, both occurrences of the phrase "COPYRIGHT HOLDERS AND CONTRIBUTORS" in the disclaimer read "REGENTS AND CONTRIBUTORS".

Here is the license template:

Copyright (c) 2014, Barak Zackay (Weizmann Institute of Science)
All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""

import numpy as np
import time
import sys

# here you should change the path to the path where cart.py (also downloadable from my website) exists
sys.path.append(r'C:/Users/Ora/Dropbox/python/utilities/')
sys.path.append('/Users/barakzackay/Dropbox/python/utilities/')
sys.path.append('C:/Users/physics/Dropbox/python/utilities/')

# This package is a propietary library for convenient plotting and data manipulation
# it is imported and not execfiled in order not to litter the user's workspace
# if you use it for other things (not recommended by me, and I wrote it!) use by execfile, it will set up MY preferred workspace
import cart

#This is the standard dispersion constant, with units that fit coherent dedispersion
DispersionConstant = 4.148808*10**9; ## Mhz * pc^-1 * cm^3
Verbose = False

# This file contains code for main functions:
# 1) FDMT
# 2) FDMTFFT
# 3) HybridDedispersion
# To test their working (and to make what they do absolutely clear) there are 3 test functions
# FDMT_test_ones - used for timing of the FDMT/FDMTFFT functions, and sanity check
# FDMT_test_curve - used for testing the compatability of FDMT with coherently dispersed pulses
# FDMT_test_hitting_efficiency - a 3 pulse caricature is generated by the coherent dispersion recipe
#                                then you test how many of the added bins belongs to the pulse caricature.
# HybridDedispersion_test - used for verifying the HybridDedispersion function works

# possible improvements to runtime:
# 1) Use Cython
# 2) Write in C
# 3) make FDMTFFT be efficient in C/Cython

# Full description is available in Zackay & Ofek (2014, submitted to ApJ)
# ArXiv URL:
# ADS URL:


def FDMT_test_ones(i, N_d = 1024, N_f = 1024,f_min = 1200,f_max = 1600,TestFDMTFFT=False):
     
    if TestFDMTFFT:
        FUNC = FDMTFFT
        dataType = 'complex128'
    else :
        dataType = 'int64'
        FUNC = FDMT
    I =  np.ones([N_f,N_f*i],dataType)
    xx = time.time()
    F = np.real(FUNC(I, f_min, f_max, N_d, dataType))
    print time.time()-xx
    G = cart.cview(F)
    return G

def PDB(message, var = ''):
    if Verbose :
        print message + str(var) 

def FDMT_test_curve(TestFDMTFFT = False):
    f_min = 1200 #Mhz
    f_max = 1600 #MHz
    
    N_bins = 40
    N_t = 512
    N_f = 512
    N_total = N_f*N_t*N_bins
    PulseLength = N_f*N_bins
    PulseSig = 0.4
    PulsePosition = 4134567
    D = 5
    maxDT = N_t
    
    if TestFDMTFFT:
        FUNC = FDMTFFT
        dataType = 'complex128'
    else :
        dataType = 'int64'
        FUNC = FDMT
    Verbose = False
    
    PDB("Signal preparation")
    practicalD = DispersionConstant * D
    I = np.random.normal(0,1,N_total)
    I[PulsePosition:PulsePosition+PulseLength] += np.random.normal(0,PulseSig,PulseLength)
    print "MAX Thoretical SNR:", np.sum(np.abs(I[PulsePosition:PulsePosition+PulseLength])**2 - np.mean(abs(I)**2)) / (np.sqrt(PulseLength*np.var(abs(I)**2)))
    
    X = CoherentDedispersion(I, -D, f_min,f_max,False)    
    #F = np.fft.fft(I)/np.sqrt(len(I))
    #f = np.arange(0,f_max-f_min, float(f_max-f_min)/N_total)
    #PDB("f_shape",f.shape)
    #F*= np.e**(2*np.pi*complex(0,1) * practicalD /(f_min + f) )
    #X = np.fft.ifft(F) # The input raw voltage signal...
    
    
    PDB("STFT, shape = ",X.shape)
    XX = np.abs(np.fft.fft(X.reshape([N_f,N_t*N_bins]),axis = 1))**2
    PDB("Transposing, shape = ", XX.shape)
    XX = np.transpose(XX)
    XX = np.sum(XX.reshape(N_f,N_bins,N_t),axis=1)
    
    E = np.mean(XX[:,:10])
    XX -= E
    V = np.var(XX[:,:10])
    
    XX /= (0.25*np.sqrt(V))
    V = np.var(XX[:,:10])
    PDB("E = ",E)
    PDB("V = ",V)
    
    G0 = cart.cview(XX)
    DM0 = np.real(FUNC(np.ones(XX.shape,dataType),f_min,f_max,maxDT,dataType,Verbose))
    
    
    PDB("Applying FDMT")
    DM = np.real(FUNC(XX, f_min, f_max, maxDT, dataType, Verbose))
    Res = DM/np.sqrt((DM0+0.000001) * V )
    G = cart.cview(Res)
    print "Maximum acieved SNR:", np.max(Res)
    print "Maximum Position:", cart.argmaxnd(Res)
    return G0,G

def FDMT_test_hitting_efficiency(test_FDMT_FFT = False):
    # This functions doesnt really have units.
    # But it is used only for seeing that the dispersion is really tracing the correct curve.
    
    N_freqs = 1024;
    N_times = 16*1024;
    f_min = 400;
    f_max = 800;
    N_DM = 1024;
    N_show = 1024;
    
    if test_FDMT_FFT :
        FUNC = FDMTFFT
        dataType = 'complex64';
    else :
        FUNC = FDMT
        dataType = 'int64'
    
    x_start = 0
    x_end = 10**3*N_show * N_freqs / ((f_max - f_min) *10**6.)
    y_start = f_min
    y_end = f_max
    xticklabels = np.round([x_start] + list(np.linspace(x_start,x_end,6)),2)
    yticklabels = np.round([y_start] + list(np.linspace(y_start,y_end,6)))

    Const1 = 1
    Const2 = 2
    Const3 = 3
    Thickness = 2
    
    T_position = 900
    
    [F,T] = np.mgrid[f_min:f_max:(f_max - f_min)/float(N_freqs),1:1024*16]
    
    R0 = np.ones([N_freqs,N_times],dataType);
    RT0 = np.real(FUNC(R0,f_min,f_max,N_DM,dataType));
    G0 = cart.cview(RT0[:N_show,:N_show], 'color')



    R1 = np.zeros([N_freqs,N_times],dataType);
    S = T-T_position  + 4*40*1000000*((1./f_min)**2 - (1./(F)**2));
    R1[S**2 < (Thickness/2.)**2] = Const1;
    S = T-T_position  + 4*39*1000000*((1./f_min)**2 - (1./(F)**2));
    R1[S**2 < (Thickness/2.)**2] = Const2;
    S = T-T_position  + 4*41*1000000*((1./f_min)**2 - (1./(F)**2));
    R1[S**2 < (Thickness/2.)**2] = Const3;


    G1 = cart.cview(np.real(R1[:N_show,:N_show]),'color',xlabel = "Time[ms]", ylabel = "Frequency[MHz]", fontsize = 20, colorbar = False, xticklabels = xticklabels, yticklabels = yticklabels)
    
    RT1 = np.real(FUNC(R1,f_min,f_max,N_DM,dataType))
    G2 = cart.cview(RT1[:N_show,:N_show], 'color')

    hitting_efficiency = cart.maxnd(RT1/(RT0+0.0000000001))/Const3
    
    print 'hitting_efficiency', hitting_efficiency
    return G0,G1,G2
    
def HybridDedispersion_test():
    sig = np.random.normal(0,1,2**20)
    sig[1000:1128] = np.random.normal(0,1.8,128)
    # Dispersion is just like dedispersion with a minus sign...
    sig2 = CoherentDedispersion(sig,-1,1200,1600,False)
    H = HybridDedispersion(sig2,128,1.5,1200,1600,5)
    
    return cart.cview(H[1][:128,:128])



def FDMT(Image, f_min, f_max,maxDT ,dataType, Verbose = True):
    """
    This function implements the  FDMT algorithm.
    Input: Input power matrix I(f,t)
           f_min,f_max are the base-band begin and end frequencies.
                   The frequencies should be entered in MHz 
           maxDT - the maximal delay (in time bins) of the maximal dispersion.
                   Appears in the paper as N_{\Delta}
                   A typical input is maxDT = N_f
           dataType - a valid numpy dtype.
                      reccomended: either int32, or int64.
    Output: The dispersion measure transform of the Input matrix.
            The output dimensions are [Input.shape[1],maxDT]
    
    For details, see algorithm 1 in Zackay & Ofek (2014)
    """
    F,T = Image.shape
    f = int(cart.log2(F))
    if (F not in [2**i for i in range(1,30)]) or (T not in [2**i for i in range(1,30)]) :
        raise NotImplementedError("Input dimensions must be a power of 2")

    x = time.time()
    State = FDMT_initialization(Image,f_min,f_max,maxDT,dataType)
    PDB('initialization ended')
    
    for i_t in range(1,f+1):
        State = FDMT_iteration(State,maxDT,F,f_min,f_max,i_t,dataType, Verbose)
    PDB('total_time:', time.time() - x)
    [F,dT,T] = State.shape;
    DMT= np.reshape(State,[dT,T]);
    return DMT

def FDMTFFT(Input, f_min, f_max,maxDT ,dataType, Verbose = True):
    """
    This function implements the  FDMT-FFT algorithm.
    Input: Input power matrix I(f,t)
           f_min,f_max are the base-band begin and end frequencies.
                   The frequencies can be entered in both MHz and GHz, units are factored out in all uses.
           maxDT - the maximal delay (in time bins) of the maximal dispersion.
                   Appears in the paper as N_{\Delta}
                   A typical input is maxDT = N_f
           dataType - To naively use FFT, one must use floating point types.
                      Due to casting, use either complex64 or complex128.
    Output: The dispersion measure transform of the Input matrix.
            The output dimensions are [Input.shape[1],maxDT]
    For details, see algorithm 2 in Zackay & Ofek (2014)

    """
    F,T = Input.shape
    f = int(cart.log2(F))
    if (F not in [2**i for i in range(1,30)]) or (T not in [2**i for i in range(1,30)]) :
        raise NotImplementedError("Input dimensions must be a power of 2")

    x = time.time()
    State = FDMTFFT_initialization(Input,f_min,f_max,maxDT,dataType)
    PDB('initialization ended')
    
    for i_t in range(1,f+1):
        #print i_t
        #xx = time.time()
        State = FDMTFFT_iteration(State,maxDT,F,f_min,f_max,i_t,dataType,Verbose)
        #print time.time() - xx
    PDB('total_time:', time.time() - x)
    [T,F,dT] = State.shape;
    State = np.transpose(State,axes = [1,2,0])
    DMT= np.reshape(np.fft.ifft(State,axis = 2),[dT,T]);
    return DMT

    

def FDMT_initialization(Image,f_min,f_max,maxDT,dataType):
    """
    Input: Image - power matrix I(f,t)
        f_min,f_max - are the base-band begin and end frequencies.
            The frequencies can be entered in both MHz and GHz, units are factored out in all uses.
        maxDT - the maximal delay (in time bins) of the maximal dispersion.
            Appears in the paper as N_{\Delta}
            A typical input is maxDT = N_f
        dataType - To naively use FFT, one must use floating point types.
            Due to casting, use either complex64 or complex128.
    Output: 3d array, with dimensions [N_f,N_d0,Nt]
            where N_d0 is the maximal number of bins the dispersion curve travels at one frequency bin
    
    For details, see algorithm 1 in Zackay & Ofek (2014)
    """
    # Data initialization is done prior to the first FDMT iteration
    # See Equations 17 and 19 in Zackay & Ofek (2014)

    [F,T] = Image.shape

    deltaF = (f_max - f_min)/float(F)
    deltaT = int(np.ceil((maxDT-1) *(1./f_min**2 - 1./(f_min + deltaF)**2) / (1./f_min**2 - 1./f_max**2)))

    Output = np.zeros([F,deltaT+1,T],dataType)
    Output[:,0,:] = Image
    
    for i_dT in xrange(1,deltaT+1):
        Output[:,i_dT,i_dT:] = Output[:,i_dT-1,i_dT:] + Image[:,:-i_dT]
    return Output


def FDMTFFT_initialization(Image,f_min,f_max,maxDT,dataType):
    """
        Input: Image - power matrix I(f,t)
            f_min,f_max - are the base-band begin and end frequencies.
                    The frequencies can be entered in both MHz and GHz, units are factored out in all uses.
            maxDT - the maximal delay (in time bins) of the maximal dispersion.
                    Appears in the paper as N_{\Delta}
                    A typical input is maxDT = N_f
            dataType - To naively use FFT, one must use floating point types.
                    Due to casting, use either complex64 or complex128.
        Output: 3d array, with dimensions [N_f,N_d0,Nt]
                where N_d0 is the maximal number of bins the dispersion curve travels at one frequency bin
                the difference from FDMT_FFT_initialization is that the time axis is FFT'ed
        
        For details, see algorithm 2 in Zackay & Ofek (2014)
    """

    [F,T] = Image.shape

    deltaF = (f_max - f_min)/float(F)
    # determining the maximal deltaT that we will encounter in the first iteration.
    # if deltaT is too large, consider binning
    deltaT = int(np.ceil((maxDT-1) *(1./f_min**2 - 1./(f_min + deltaF)**2) / (1./f_min**2 - 1./f_max**2)))

    Output = np.zeros([F,deltaT+1,T],dataType)
    
    # Initializing the "A_f^{f + \delta f} (t_0,\Delta t)" array
    Output[:,0,:] = Image    
    for i_dT in xrange(1,deltaT+1):
        Output[:,i_dT,i_dT:] = Output[:,i_dT-1,i_dT:] + Image[:,:-i_dT]
    
    # FFT-ing the time axis and transposing the data
    return np.transpose(np.fft.fft(Output,axis = 2),axes = [2,0,1])

def FDMT_iteration(Input,maxDT,F,f_min,f_max,iteration_num,dataType, Verbose = False):
    """
        Input: 
            Input - 3d array, with dimensions [N_f,N_d0,Nt]
            f_min,f_max - are the base-band begin and end frequencies.
                The frequencies can be entered in both MHz and GHz, units are factored out in all uses.
            maxDT - the maximal delay (in time bins) of the maximal dispersion.
                Appears in the paper as N_{\Delta}
                A typical input is maxDT = N_f
            dataType - To naively use FFT, one must use floating point types.
                Due to casting, use either complex64 or complex128.
            iteration num - Algorithm works in log2(Nf) iterations, each iteration changes all the sizes (like in FFT)
        Output: 
            3d array, with dimensions [N_f/2,N_d1,Nt]
        where N_d1 is the maximal number of bins the dispersion curve travels at one output frequency band
        
        For details, see algorithm 1 in Zackay & Ofek (2014)
    """

    input_dims = Input.shape;
    output_dims = list(input_dims);
    
    deltaF = 2**(iteration_num) * (f_max - f_min)/float(F);
    dF = (f_max - f_min)/float(F)
    # the maximum deltaT needed to calculate at the i'th iteration
    deltaT = int(np.ceil((maxDT-1) *(1./f_min**2 - 1./(f_min + deltaF)**2) / (1./f_min**2 - 1./f_max**2)));
    PDB("deltaT = ",deltaT)
    PDB("N_f = ",F/2.**(iteration_num))
    PDB('input_dims', input_dims)
    
    output_dims[0] = output_dims[0]/2;
    
    
    output_dims[1] = deltaT + 1;
    PDB('output_dims', output_dims)
    Output = np.zeros(output_dims,dataType);
    
    # No negative D's are calculated => no shift is needed
    # If you want negative dispersions, this will have to change to 1+deltaT,1+deltaTOld
    # Might want to calculate negative dispersions when using coherent dedispersion, to reduce the number of trial dispersions by a factor of 2 (reducing the complexity of the coherent part of the hybrid)
    ShiftOutput = 0
    ShiftInput = 0
    T = output_dims[2]
    F_jumps = output_dims[0]
    
    # For some situations, it is beneficial to play with this correction.
    # When applied to real data, one should carefully analyze and understand the effect of 
    # this correction on the pulse he is looking for (especially if convolving with a specific pulse profile)
    if iteration_num>0:
        correction = dF/2.
    else:
        correction = 0
    for i_F in range(F_jumps):
        
        f_start = (f_max - f_min)/float(F_jumps) * (i_F) + f_min;
        f_end = (f_max - f_min)/float(F_jumps) *(i_F+1) + f_min;
        f_middle = (f_end - f_start)/2. + f_start - correction;
        # it turned out in the end, that putting the correction +dF to f_middle_larger (or -dF/2 to f_middle, and +dF/2 to f_middle larger)
        # is less sensitive than doing nothing when dedispersing a coherently dispersed pulse.
        # The confusing part is that the hitting efficiency is better with the corrections (!?!).
        f_middle_larger = (f_end - f_start)/2 + f_start + correction;
        deltaTLocal = int(np.ceil((maxDT-1) *(1./f_start**2 - 1./(f_end)**2) / (1./f_min**2 - 1./f_max**2)))
        
        for i_dT in range(deltaTLocal+1):
            dT_middle = round(i_dT * (1./f_middle**2 - 1./f_start**2)/(1./f_end**2 - 1./f_start**2))
            dT_middle_index = dT_middle + ShiftInput;
            
            dT_middle_larger = round(i_dT * (1./f_middle_larger**2 - 1./f_start**2)/(1./f_end**2 - 1./f_start**2))            
                     
            
            dT_rest = i_dT - dT_middle_larger;
            dT_rest_index = dT_rest + ShiftInput;
            
            i_T_min = 0;
            
            i_T_max = dT_middle_larger
            Output[i_F,i_dT + ShiftOutput,i_T_min:i_T_max] = Input[2*i_F, dT_middle_index,i_T_min:i_T_max];
            
            
            i_T_min = dT_middle_larger;
            i_T_max = T;
            
            
            
            Output[i_F,i_dT + ShiftOutput,i_T_min:i_T_max] = Input[2*i_F, dT_middle_index,i_T_min:i_T_max] + Input[2*i_F+1, dT_rest_index,i_T_min - dT_middle_larger:i_T_max-dT_middle_larger]
    
    return Output

def FDMTFFT_iteration(Input,maxDT,F,f_min,f_max,iteration_num,dataType,Verbose = False):
    """
        Input:
            Input - 3d array, with dimensions [N_f,N_d0,Nt]
            f_min,f_max - are the base-band begin and end frequencies.
                The frequencies can be entered in both MHz and GHz, units are factored out in all uses.
            maxDT - the maximal delay (in time bins) of the maximal dispersion.
                Appears in the paper as N_{\Delta}
                A typical input is maxDT = N_f
            dataType - To naively use FFT, one must use floating point types.
                Due to casting, use either complex64 or complex128.
            iteration num - Algorithm works in log2(Nf) iterations, each iteration changes all the sizes (like in FFT)
        Output:
            3d array, with dimensions [N_f/2,N_d1,Nt]
            where N_d1 is the maximal number of bins the dispersion curve travels at one output frequency band
        
        For details, see algorithm 2 in Zackay & Ofek (2014)
        """
    
    input_dims = Input.shape;
    output_dims = list(input_dims);
    
    deltaF = 2**(iteration_num) * (f_max - f_min)/float(F);
    dF = (f_max - f_min)/float(F)
    # the maximum deltaT needed to calculate at the i'th iteration
    deltaT = int(np.ceil((maxDT-1) *(1./f_min**2 - 1./(f_min + deltaF)**2) / (1./f_min**2 - 1./f_max**2)));
    PDB("deltaT = ",deltaT)
    PDB("N_f = ",F/2.**(iteration_num))
    
    output_dims[1] = output_dims[1]/2;
    
    
    output_dims[2] = deltaT + 1;
    Output = np.zeros(output_dims,dataType);
    
    # No negative K's are calculated => no shift is needed
    # If you want negative shifts, this will have to change to 1+deltaT,
    # 1+deltaTOld
    ShiftOutput = 0
    ShiftInput = 0
    T = output_dims[0]
    F_jumps = output_dims[1]
    
    # see remark about this correction in the FDMT implementation.
    correction = dF/2.    
    
    deltaTShift = np.ceil((maxDT-1) *(1./f_min**2. - 1./(f_min + deltaF/2. + dF/2.)**2.) / (1./f_min**2. - 1./f_max**2.))+3;
    ShiftRow = (np.fft.fft((np.eye(deltaTShift,T)),axis = 1))
    for i_F in range(F_jumps):
        f_start = (f_max - f_min)/float(F_jumps) * (i_F) + f_min;
        f_end = (f_max - f_min)/float(F_jumps) *(i_F+1) + f_min;
        f_middle = (f_end - f_start)/2. + f_start - correction;
        # correction was removed. see the explanation in FDMT code.
        f_middle_larger = (f_end - f_start)/2. + f_start + correction
        deltaTLocal = int(np.ceil((maxDT-1) *(1./f_start**2 - 1./(f_end)**2) / (1./f_min**2 - 1./f_max**2)))
        for i_dT in range(deltaTLocal+1):
            dT_middle = round(i_dT * (1./f_middle**2 - 1./f_start**2)/(1./f_end**2 - 1./f_start**2))
            dT_middle_index = dT_middle + ShiftInput;
            dT_middle_larger = round(i_dT * (1./f_middle_larger**2 - 1./f_start**2)/(1./f_end**2 - 1./f_start**2))            
            
            dT_rest = i_dT - dT_middle_larger;
            dT_rest_index = dT_rest + ShiftInput;
            
            Output[:,i_F,i_dT + ShiftOutput] = Input[:,2*i_F, dT_middle_index] + Input[:,2*i_F+1, dT_rest_index] * ShiftRow[dT_middle_larger,:];
    
    return Output

def CoherentDedispersion(raw_signal,d, f_min, f_max, alreadyFFTed = False):
    """
    Will perform coherent dedispersion.
    raw signal   - is assumed to be a one domensional signal
    d            - is the dispersion measure. units: pc*cm^-3
    f_min        - the minimum freq, given in Mhz
    f_max        - the maximum freq, given in Mhz
    alreadyFFTed - to reduce complexity, insert fft(raw_signal) instead of raw_signal, and indicate by this flag
    
    For future improvements:
    1) Signal partition to chunks of length N_d is not applied, and maybe it should be.
    2) No use of packing is done, though it is obvious it should be done (either in the coherent stage (and take special care of the abs()**2 operation done by other functions) or in the incoherent stage)
    
    """
    N_total = len(raw_signal)
    practicalD = DispersionConstant * d
    f = np.arange(0,f_max-f_min, float(f_max-f_min)/N_total)
    
    # The added linear term makes the arrival times of the highest frequencies be 0
    H = np.e**(-(2*np.pi*complex(0,1) * practicalD /(f_min + f) + 2*np.pi*complex(0,1) * practicalD*f /(f_max**2)))
    if not alreadyFFTed:
        CD = np.fft.ifft(np.fft.fft(raw_signal) * H)
    else :
        CD = np.fft.ifft(raw_signal * H)    
    return CD


def STFT(raw_signal,block_size):
    """
        Raw signal will be devided to blocks, each block will be fourier transformed
        Input: 
            raw_signal - raw antenna voltage time series
            block_size - number of bins in each block
        Output:
            frequency vs. time matrix
        Note: absolute value squared is not performed!
    """
    S = np.transpose(raw_signal[:len(raw_signal)/block_size * block_size].reshape([len(raw_signal)/block_size,block_size]))
    return np.fft.fft(S,axis = 0)
        
    

def HybridDedispersion(raw_signal, N_p, D_max, f_min, f_max, SigmaBound = 7):
    """
    Will perform the coherent FDMT hybrid algorithm.
    see algorithm 3 in Zackay & Ofek (2014)
    Input:
            raw_signal - raw antenna voltage time series
            N_p - length of the pulse in time bins, i.e (t_p/\tau), or N_p in the paper.
            D_max - maximal dispersion to scan, in units of pc cm^-3
            f_min,f_max - minimal and maximal frequency band of the signal.
                        - signal is assumed to be base-band sampled
            SigmaBound - The minimal statistical significane to trigger a saving of a result.
                       - You may want to change the storage decision process for your own purposes.
    Output:
            frequency vs. time matrix
    
    See example usage in the test function HybridDedispersion_test
    
    Subjects for future improvement:
    Might want to reduce trial dispersions by a factor of 2 by letting FDMT scan negative dispersions.
    Might want to use packing (either in the coherent stage, or in the incoherent stage)
    """
    # ConversionConst is used to convert from Dispersion measure in units of pc*cm^-3 to time bins
    ConversionConst = DispersionConstant *(1./f_min**2 - 1./f_max**2) * (f_max - f_min)
    N_d = D_max * ConversionConst

    n_coherent = int(np.ceil(N_d/(N_p**2)))
    print "number of coherent iterations:", n_coherent
    ffted_signal = np.fft.fft(raw_signal)

    print "FDMT parameters:", [N_p,len(raw_signal)/N_p],f_min,f_max,N_p, 'int64',False
    FDMT_normalization = FDMT(np.ones([N_p,len(raw_signal)/N_p]),f_min,f_max,N_p, 'int64',False)
    for outer_d in range(n_coherent):
        print "coherent iteration", outer_d
        cur_coherent_d = outer_d * (D_max/float(n_coherent))
        print cur_coherent_d
        FDMT_input = abs(STFT(CoherentDedispersion(ffted_signal, cur_coherent_d, f_min, f_max, alreadyFFTed = True), N_p))**2
        FDMT_input -= np.mean(FDMT_input)
        FDMT_input /= 0.25*np.std(FDMT_input)
        V = np.var(FDMT_input)
        FDMT_output = FDMT(FDMT_input, f_min, f_max, N_p, 'int64', False)
        FDMT_output /= np.sqrt(FDMT_normalization * V + 0.000001)
        if np.max(FDMT_output) > SigmaBound :
            SigmaBound = np.max(FDMT_output)
            print "achieved score with", SigmaBound, "sigmas"
            res = [cur_coherent_d,cart.copy.deepcopy(FDMT_output)]
    return res
        
def BitPack(Inputs, dataType, safe = True):
    """
    This function packs different insances of FDMT inputs to one output
    It does that by shifting each instance with a different amount of bits.
    
    If dataType is complex then there is also packing of half the instances in
    the imaginary field (this is useful for FDMTFFT)
    
    Not completely debugged yet, beware with the float datatypes.
    """   
    if safe :
        raise NotImplementedError("Not debugged yet - Use with caution")
    Output = Inputs[0].astype(dataType)
    
    # check the size of the mantissa and the required number of bits for precision in FFT
    N_float_free64 = 18 #(11 for exp field, 7 for precision, verify with FFT how much is enough)
    N_float_free32 = 14 #(8 for exp field, 6 for precision, verify with FFT how much is enough)
    if np.dtype(dataType) == np.dtype('int64') :
        N_bits = 64
        pack_complex = False
    if np.dtype(dataType) == np.dtype('int32') :
        N_bits = 32
        pack_complex = False
    if np.dtype(dataType) == np.dtype('complex128'):
        if len(Inputs)%2 != 0:
            raise "On complex data types the number of inputs should be even!"
        N_bits = 64-N_float_free64
        pack_complex = True
    if np.dtype(dataType) == np.dtype('complex64'):
        if len(Inputs)%2 != 0:
            raise "On complex data types the number of inputs should be even!"
        N_bits = 32-N_float_free32
        pack_complex = True
    
    if not pack_complex :
        N_shift = N_bits / len(Inputs)
        N_pack = len(Inputs)
    else :
        N_shift = N_bits / (len(Inputs)/2)
        N_pack = len(Inputs)/2
    for i in range(1,N_pack):
        Output += Inputs[i] * (2**(N_shift*i))
    
    if pack_complex :
        for i in range(N_pack,len(Inputs)):
            Output += complex(0,1)*Inputs[i] * (2**(N_shift*i))
    
    return Output
                
    
        



    
    

