% deconvolveSamples
% Auth: Rogan Carr <rogan@uw.edu>
% Date: July 2013
%
% deconvolves Sample Taxa Abundances and Sample Element Abundances
%  to predict the most likely distribution of the elements among the taxa
%
% Uses the least squares regression and non-negative least squares regression implementations that
%  are built in the base install of Matlab.
%
%
%  Inputs
%  ------
%
% TaxaAbundaceFile - nSamples x nTaxa - tab delimited file containing taxa abundances for each sample
%
% ElementAbundaceFile - nSamples x nElements - tab delimited file containing element abundances for each sample
%
% ElementFile - 1 x nElements - tab delimeted file containing all element names
%
% constantElementName - the name of the constant element
%
% constantElementLength - the (approximate) length (or copy number) of the constant element in each taxa
%
% outPrefix - the file prefix for all output
%
% (optional) 'MaskAbundance', 'number between 0 and 1'
%             only deconvolve elements with a mean abundance at least the product of this number and the abundance of the most abundant element in any samples.
%             example: 'MaskAbundance',0.001 -> only deconvolve elements with a mean abundance at least 1/1000th that of the most abundant element in any sample.
%
%  Outputs
%  -------
%
% outPrefix.LeastSquares.dat - the genome reconstruction for each taxa using least squares
%
% outPrefix.NNLeastSquares.dat - the genome reconstruction for each taxa using non-negative least squares
%
% outPrefix.Mask.dat - 1 x nElements - a list of 1s or 0s describing whether the element was solved for or not, respectively
%

function deconvolveSamples(TaxaAbundanceFile,ElementAbundanceFile,ElementFile,constantElementName,constantElementLength,outPrefix,varargin)

    % Check for options
    v=1;
    MaskAbundance=0;
    while v <= length(varargin)
        if isstr(varargin{v}) && strcmpi('maskabundance',varargin{v})
            MaskAbundance=varargin{v+1};
            varargin(v+1)=[]; % Order counts!
            varargin(v)=[];
        end
    end

    % --- Read in the parameters --- %

    A=dlmread(TaxaAbundanceFile,'\t');
    P=dlmread(ElementAbundanceFile,'\t');
    E=dlmread(ElementFile,'\t');
    
    % --- --- --- ------ --- --- --- %



    % --- Determine the solution set --- %

    nSamples=size(A,1);
    if nSamples ~= size(P,1)
        fprintf('Number of samples in disagreement between TaxaAbundanceFile and ElementAbundanceFile.  Exiting.\n');
        exit;
    end
    
    nElements=size(P,2);
    nTaxa=size(A,2);
    
    % Get the constant element abundances
    C = P(:,E==constantElementName)';

    % We may not want to solve for every element
    mask=ones(1,nElements);
    
    % Don't solve for the constant element
    mask(E==constantElementName)=0;

    % Use an abundance metric to ignore some elements
    if MaskAbundance

        noiseFloor=max(reshape(P(:,E~=constantElementName),1,[])) * MaskAbundance;
        
        for k=1:nElements
            if mean(P(:,k)) < noiseFloor
                mask(k) = 0;
            end
        end
        
        fprintf('Ignoring KOs with mean reads fewer than %.2g (Keeping %d, or %.2f%% of all KOs)\n',noiseFloor,sum(mask),100*mean(mask));
        
        % Write down the suggested mask
        fileName = sprintf('%s.Mask.dat',outPrefix);
        dlmwrite(fileName,mask,'\t');
    end 
    % --- --- --- --- -- --- --- --- --- %
    
    
    
    % --- Normalize the samples --- %
    
    for i=1:nSamples
        
        G = ( 1 / C(i) ) * constantElementLength * sum(A(i,:));
        
        P(i,:) = P(i,:) .* G;
        
    end
    
    % --- --- --- -- -- --- --- --- %
    
    
    
    % --- Deconvolve --- %
    
    fprintf('Deconvolving %d elements into %d groups using a constant element of length %g.\n',nElements,nTaxa,constantElementLength);

    % Rows are taxa, columns are elements
    L = zeros(nTaxa,nElements);
    N = zeros(nTaxa,nElements);
   
    % Do each element independently, optionally ignoring elements
    for k=1:nElements
        if mask(k)
            % least squares
            L(:,k) = regress(P(:,k),A);
            
            % nonegative least squares
            N(:,k) = lsqnonneg(A,P(:,k));
            
        else
            L(:,k) = zeros(nTaxa,1);
            N(:,k) = zeros(nTaxa,1);
        end
    end

    % Write out the files
    fileName = sprintf('%s.LeastSquares.dat',outPrefix);
    dlmwrite(fileName,L,'\t');
    
    fileName = sprintf('%s.NNLeastSquares.dat',outPrefix);
    dlmwrite(fileName,N,'\t');
    
    % --- --- -- --- --- %
  
end
