function [te, dfi, FIT, TEQe, TELe, FITQe, FITLe] = compute_FIT_TE(feature, X, Y, hY,doQe,xtrap)
% This function computes the Feature-specific Information Transfer (FIT)
% Transfer Entropy (TE) and directed feature information (DFI) 
% Quadratic extrapolation bias correction is available for TE and FIT

% input: (N = number of experimental trials)
% feature = discrete feature value (1 x N)
% X = discrete past activity of the sender X_past (1 x N)
% Y = discrete present activity of the receiver Y_pres (1 x N)
% hY = discrete past activity of the receiver Y_past (1 x N)
% doQe = set to 1 if you want to use the Quadratic Extrapolation technique to correct for the limited sampling bias (by default is 0)
% xtrap = number of data resamplings when estimating FIT and TE on datasets with N/2 and N/4 trials (if no value is specified, the default is xtrap = 20)

% output:
% te = transfer entropy value (from X to Y; see Schreiber T. (2000) Phys Rev Letters)
% dfi = DFI value (see Ince R., et al. (2015) Sci Reports)
% FIT = Feature-specific Information Transfer value (from X to Y about S)
% TEQe = bias-corrected TE value using quadratic extrapolation
% TELe = bias-corrected TE value using linear extrapolation 
% FITQe = bias-corrected FIT value using quadratic extrapolation
% FITLe = bias-corrected FIT value using linear extrapolation

% Note: linear extrapolation values are computed when doQe = 1 since they
% don't require additional calculations. However, Qe estimates are more reliable

% License:
% MIT License
% 
% Copyright (c) 2023 M. Celotto, J. Bím, A. Tlaie, V. De Feo, A.Toso, S. Lemke, D. Chicharro, 
% H. Nili, M. Bieler, I.L. Hanganu-Opatz, T.H. Donner, A. Brovelli, and S. Panzeri
% 
% Permission is hereby granted, free of charge, to any person obtaining a copy
% of this software and associated documentation files (the "Software"), to deal
% in the Software without restriction, including without limitation the rights
% to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
% copies of the Software, and to permit persons to whom the Software is
% furnished to do so, subject to the following conditions:
% 
% The above copyright notice and this permission notice shall be included in all
% copies or substantial portions of the Software.
% 
% THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
% IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
% FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
% AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
% LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
% OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
% SOFTWARE.

if nargin == 4 % by default, don't use the quadratic extrapolation (computationally expensive)
    doQe = 0;
    xtrap = 0;
elseif nargin == 5
    xtrap = 20;
end

% Build the two four-variables probability distributions needed to compute FIT
pXYhYS = probabilityDist(X, Y, hY, feature); % probability distribution for the PID with (Xp, Yp, Yt) as sources and S as target
pXShYY = probabilityDist(X, feature, hY, Y); % probability distribution for the PID with (Xp, Yp, S) as sources and Yt as target

% Compute the two FIT atoms
SUI_S = compute_SUI(pXYhYS);
SUI_Y = compute_SUI(pXShYY);

FIT = min(SUI_S,SUI_Y);

% Compute TE
te = TE(pXYhYS);

% Compute DFI
dfi = DFI(pXYhYS);

% Compute quadratic extrapolation bias correction for FIT and TE
if doQe
    FITAll = FIT;
    TEAll = te;
    
    for xIdx = 1:xtrap
        % Partition trials into subsamples of N/2 and N/4
        numberOfTrials = length(X);
        
        rIdx = randperm(numberOfTrials);
        
        idx21 = rIdx(1:round(numberOfTrials/2));
        idx22 = rIdx(round(numberOfTrials/2)+1:numberOfTrials);
        idx41 = rIdx(1:round(numberOfTrials/4));
        idx42 = rIdx(round(numberOfTrials/4)+1:round(numberOfTrials/2));
        idx43 = rIdx(round(numberOfTrials/2)+1:round(3*numberOfTrials/4));
        idx44 = rIdx(round(3*numberOfTrials/4)+1:numberOfTrials);
        
        X21 = X(idx21);
        X22 = X(idx22);
        Y21 = Y(idx21);
        Y22 = Y(idx22);
        hY21 = hY(idx21);
        hY22 = hY(idx22);
        feat21 = feature(idx21);
        feat22 = feature(idx22);

        X41 = X(idx41);
        X42 = X(idx42);
        X43 = X(idx43);
        X44 = X(idx44);
        Y41 = Y(idx41);
        Y42 = Y(idx42);
        Y43 = Y(idx43);
        Y44 = Y(idx44);
        hY41 = hY(idx41);
        hY42 = hY(idx42);
        hY43 = hY(idx43);
        hY44 = hY(idx44);
        feat41 = feature(idx41);
        feat42 = feature(idx42);
        feat43 = feature(idx43);
        feat44 = feature(idx44);

        pXYhYS21 = probabilityDist(X21, Y21, hY21, feat21);
        pXShYY21 = probabilityDist(X21, feat21, hY21, Y21);
        FIT21S = compute_SUI(pXYhYS21);
        FIT21Y = compute_SUI(pXShYY21);
        FIT21 = min(FIT21S, FIT21Y);
        TE21 = TE(pXYhYS21);

        pXYhYS22 = probabilityDist(X22, Y22, hY22, feat22);
        pXShYY22 = probabilityDist(X22, feat22, hY22, Y22);
        FIT22S = compute_SUI(pXYhYS22);
        FIT22Y = compute_SUI(pXShYY22);
        FIT22 = min(FIT22S, FIT22Y);
        TE22 = TE(pXYhYS22);
        
        pXYhYS41 = probabilityDist(X41, Y41, hY41, feat41);
        pXShYY41 = probabilityDist(X41, feat41, hY41, Y41);
        FIT41S = compute_SUI(pXYhYS41);
        FIT41Y = compute_SUI(pXShYY41);
        FIT41 = min(FIT41S, FIT41Y);
        TE41 = TE(pXYhYS41);

        pXYhYS42 = probabilityDist(X42, Y42, hY42, feat42);
        pXShYY42 = probabilityDist(X42, feat42, hY42, Y42);
        FIT42S = compute_SUI(pXYhYS42);
        FIT42Y = compute_SUI(pXShYY42);
        FIT42 = min(FIT42S, FIT42Y);
        TE42 = TE(pXYhYS42);

        pXYhYS43 = probabilityDist(X43, Y43, hY43, feat43);
        pXShYY43 = probabilityDist(X43, feat43, hY43, Y43);
        FIT43S = compute_SUI(pXYhYS43);
        FIT43Y = compute_SUI(pXShYY43);
        FIT43 = min(FIT43S, FIT43Y);
        TE43 = TE(pXYhYS43);
        
        pXYhYS44 = probabilityDist(X44, Y44, hY44, feat44);
        pXShYY44 = probabilityDist(X44, feat44, hY44, Y44);
        FIT44S = compute_SUI(pXYhYS44);
        FIT44Y = compute_SUI(pXShYY44);
        FIT44 = min(FIT44S, FIT44Y);
        TE44 = TE(pXYhYS44);
        
        FIT2(xIdx) = (FIT21 + FIT22) / 2;
        FIT4(xIdx) = (FIT41 + FIT42 + FIT43 + FIT44) / 4;

        TE2(xIdx) = (TE21 + TE22) / 2;
        TE4(xIdx) = (TE41 + TE42 + TE43 + TE44) / 4;
    end
    x = [1/length(X41) 1/length(X21) 1/length(X)];
    y = [mean(FIT4) mean(FIT2) FITAll];

    p2 = polyfit(x, y, 2); 
    p1 = polyfit(x, y, 1);  
    FITQe = p2(3);
    FITLe = p1(2);
    
    y = [mean(TE4) mean(TE2) TEAll];

    p2 = polyfit(x, y, 2); 
    p1 = polyfit(x, y, 1);  
    TEQe = p2(3);
    TELe = p1(2);
end

end



