function model = img2model( dimensionality, varargin )
%IMG2MODEL Trains a generative model of protein subcellular location from a
%collection of microscope images.
%
%Inputs
%dimensionality    either '2D' or '3D'
%param             a structure holding possible parameter options
%
%See also IMG2SLML

% Ivan E. Cao-Berg
%
% Copyright (C) 2007-2013 Murphy Lab
% Carnegie Mellon University
%
% ?? ??, 2011 I. Cao-Berg Added 3D model training functionality
% March 23, 2012 I. Cao-Berg Added the creation/deletion of temporary folder
% March 28, 2012 I. Cao-Berg Added a control structure under which if one or more of the 
%                image folders are nonexistent or they do not contain images,
%                the method exits
% March 28, 2012 I. Cao-Berg Added verification of input arguments when training a 2D
%                            generative model
% March 28, 2012 I. Cao-Berg Added verification of input arguments when training a 3D
%                            generative model
% April 10, 2012 I. Cao-Berg Added debug flag to the method. If flag is true, temporary
%                            files will not be deleted
% April 11, 2012 I. Cao-Berg Added verbose flag to the method
% April 17, 2012 I. Cao-Berg Returns an empty model when model cannot be trained
% July 5, 2012 I. Cao-Berg Added training flags to method so that users can train whatever component they wish
% July 26, 2012 Y.Yu Fixed a bug of the order of input argument for ml_traingenmodel2D method
% August 29, 2012 G. Johnson Modified method call to include parameter structure
% May 7, 2013 I. Cao-Berg Included support of masks when training a model
% only for the 2D case
% May 8, 2013 I. Cao-Berg Removed check for the existence of image folder since the check will happen later in the code
% May 15, 2013 I. Cao-Berg Updated method to support wildcards
% June 7-13 2013 D. Sullivan Major refactoring to support parallel/per-cell
%                            parameter calcultaions for 3D
%
% Jul 22, 2013 G. Johnson    Added parameter to skip preprocessing entirely
%                            and use only currently existing preprocessing
%                            results
% Aug 2, 2013 G. Johnson     Fixed logic so that prot_image_files are not
%                            overwritten by an empty cell array if they exist
% Aug 2, 2013 G. Johnson     Implemented chunk_start parallelization on
%                            per-cell parameterization
% Aug 30, 2013 G. Johnson    Changed they way files are input into the
%                            diffeomorphic model function
% March 14, 2014 I. Cao-Berg Changed method so that if param.masks is empty
%                            or nonvalid it displays a warning
%
% This program is free software; you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published
% by the Free Software Foundation; either version 2 of the License,
% or (at your option) any later version.
%
% This program is distributed in the hope that it will be useful, but
% WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
% General Public License for more details.
%
% You should have received a copy of the GNU General Public License
% along with this program; if not, write to the Free Software
% Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
% 02110-1301, USA.
%
% For additional information visit http://murphylab.web.cmu.edu/ or
% send email to murphy@cmu.edu

% Get collection of images
nuclearModelImages = [];
cellModelImages = [];
proteinModelImages = [];

model = [];
switch lower(dimensionality)
    case '2d'
        %icaoberg march 28, 2012
        %check number of input arguments. if there are not 5 input
        %arguments, then the method returns an empty model
        if nargin ~= 5
           warning('CellOrganizer: Wrong number of input arguments.');
           return;
        end

        %check the image directories. if they don't exist or they are
        %empty, the method returns an empty model
        nuclearModelImagesDirectory = varargin{1};

        cellModelImagesDirectory = varargin{2};
 
        proteinModelImagesDirectory = varargin{3};
    
        nuclearModelImages = ml_ls( ...
            nuclearModelImagesDirectory );
        cellModelImages = ml_ls( ...
            cellModelImagesDirectory );
        proteinModelImages = ml_ls( ...
            proteinModelImagesDirectory );

        
        %check that the input parameter is a structure
        param = varargin{4};
        if ~isa( param, 'struct' )
           warning('CellOrganizer: Input argument parameter must be a struct.');
           return;
        end

        %grj 14/5/2013
        debug = setDebugParam(param);
        verbose = setVerboseParam(param);
        
        %icaoberg 7/5/2013
        try
            masks = ml_ls( ...
            param.masks );
        catch
            masks = [];
        end
        
        %icaoberg 24/3/2014
        if isempty( masks ) && ~isempty( param.masks )
            warning( [ 'Directory of masks was defined but method returned ' ...
                'an empty list. Check the folder ' param.masks ' and try again.'] );
        end
        
        %check existence of temporary folder and make on if it doesn't exist
        if ~exist([pwd filesep 'temp' filesep 'preprocessed'])
          mkdir([pwd filesep 'temp' filesep 'preprocessed']);
        end
        
        % train generative model of protein subcellular location
        param.disp = 'false';
        model = ml_traingenmodel2D( ...
            proteinModelImages,...
            nuclearModelImages, ...
            cellModelImages, ...
            masks, param );
    case '3d'
        %icaoberg march 28, 2012
        %check the existence of the image directory
        dnaImagesDirectoryPath = varargin{1};
        cellImagesDirectoryPath = varargin{2};
        proteinImagesDirectoryPath = varargin{3};
        param = varargin{4};
        
        param = ml_initparam(param, ...
            struct('skip_ls', false));
        
        
        if ~isa( param , 'struct' )
            warning('CellOrganizer: Input parameter list must be a structure.');
            return
        end

        %grj 14/5/2013
        debug = setDebugParam(param);
        verbose = setVerboseParam(param);
        
        %mmackie july 3 2012
        try
            trainFlag = param.train.flag;
            if ~isa( trainFlag, 'char' );
                error('CellOrganizer: training flag must be a string');
            end
        catch
            param.train.flag = 'all';
            trainFlag = param.train.flag;
        end
        
        %gj 11/3/2014
        if ~ismember(trainFlag, {'nuclear', 'framework', 'all', 'microtubule'})
            error('CellOrganizer: Unrecognized training flag');
        end
        
        %grj 7/26/13 - Check to see eith cell or nucleus are diffeomorphic
        % and if so, use the diffeomorphic model
        if (isfield( param, 'nucleus') && isfield( param.nucleus, 'type') && strcmpi(param.nucleus.type, 'diffeomorphic')) ...
           || (isfield( param, 'cell')   && isfield( param.cell, 'type')    && strcmpi(param.cell.type, 'diffeomorphic'))
       
            isdiffeomorphic = true; 
        else
            isdiffeomorphic = false;
        end
        
        %grj 10/24/13 - If these contain characters, retreive filepaths
        %associated with them, otherwise assume they're function handles
        if isempty(dnaImagesDirectoryPath)
            dna_image_files = [];
        elseif ischar(dnaImagesDirectoryPath(1)) | (iscell(dnaImagesDirectoryPath) & ischar(dnaImagesDirectoryPath{1})) & ~param.skip_ls
            dna_image_files = ml_ls( dnaImagesDirectoryPath );
        else
            dna_image_files = dnaImagesDirectoryPath;
        end
        
        if isempty(cellImagesDirectoryPath)
            cell_image_files = [];
        elseif ischar(cellImagesDirectoryPath(1)) | (iscell(cellImagesDirectoryPath) & ischar(cellImagesDirectoryPath{1})) & ~param.skip_ls
            cell_image_files = ml_ls( cellImagesDirectoryPath );
        else
            cell_image_files = cellImagesDirectoryPath;
        end
        
        if isempty(proteinImagesDirectoryPath)
            prot_image_files = [];
        elseif ischar(proteinImagesDirectoryPath(1)) | (iscell(proteinImagesDirectoryPath) & ischar(proteinImagesDirectoryPath{1})) & ~param.skip_ls
            prot_image_files = ml_ls( proteinImagesDirectoryPath );
        else
            prot_image_files = proteinImagesDirectoryPath;
        end

        %%%%%%%%
        %D. Sullivan 6/5/13 refactor the code to produce cell
        %parameterizations first and then create models. Can put this into
        %separate functions once complete. 
        
        %G. Johnson 9/17/14 
        if isempty(dna_image_files) & isempty(cell_image_files)
            warning('No images. Exiting model building')
        elseif ~isempty(dna_image_files) & ~isempty(cell_image_files) & (length(dna_image_files) ~= length(cell_image_files))
            warning('Number of cell images is not the same as the number of dna images. Exiting model building')
        elseif isempty(dna_image_files)
            dna_image_files = cell(size(cell_image_files));
        elseif isempty(cell_image_files)
            cell_image_files = cell(size(dna_image_files));
        end
        
        numimgs = length(dna_image_files);
        
%         %icaoberg 7/3/2013
%         if isempty( dna_image_files )
%             warning('Could not find any images in the DNA images directory. Using DNA hole finding.' );
%             dna_image_files = cell(size(cell_image_files));
% %             model = [];
% %             return
%         end
        
%         if isempty( cell_image_files )
%             warning('Could not find any images in the cell images directory. Exiting method.' );
%             model = [];
%             return
%         end

        %grj 8/2/13 fixed logic so that prot_image_files are not
        %overwritten by an empty cell array if they exist
        if isempty(prot_image_files)
            if strcmpi(param.train.flag,'framework')
                prot_image_files = cell(size(cell_image_files));
            else
                warning('Could not find any images in the protein images directory. Exiting method.' );
                model = [];
                return
            end
        end
        
        %D. Sullivan 6/5/13 get all the masks if they exist
        
        if ~isfield(param, 'masks')
            mask_image_files = cell(1,length(dna_image_files));
        elseif ischar(param.masks(1)) | (iscell(param.masks) & ischar(param.masks{1}))
            mask_image_files = ml_ls(param.masks);
        else
            mask_image_files = param.masks;
        end

        param.documentation.numimgs = numimgs;
        param.documentation.im_cell = cell_image_files;
        param.documentation.im_dna = dna_image_files;
        param.documentation.im_prot = prot_image_files;
        param.documentation.im_mask = mask_image_files;
        
        %%%%%%
        %D. Sullivan June/2013 - Refactoring to per-cell oriented feature
        %extraction
        %setup param folders
        param = set_temp_result_folders(param);
        param = ml_initparam(param,struct('downsample',[1,1,1], 'preprocess', true, 'display', false));
        
        %D. Sullivan 6/7/13-6/13/13
        %Do all the cells in parallel
        %GRJ 6/17/13
            % changed the percellparam_(for/parfor) to work on single
            % images to improve maintainability
        startmodel = true;
        if param.preprocess
            if ~isfield(param,'parallel') || param.parallel>=1
                try 
                    matlabpool('open', param.parallel)

                    parfor i = 1:param.documentation.numimgs
                        percellparam(dna_image_files{i},cell_image_files{i},...
                        prot_image_files{i},mask_image_files{i}, i, param)
                    end
                    parallelflag = true;

                catch
                    disp('Parallel code failed, trying linear method');
                    parallelflag = false;
                end
            else
                parallelflag = false;
            end

            %grj Implementing chunk_start to perform parallel computing
            if ~parallelflag
                cellCounter = [];
                c = 1;
                for i = 1:param.documentation.numimgs
                    tmpfile = [param.tempparent filesep 'image_lock_' num2str(i)];
                    
                    isdonefile = [tmpfile '_done.tmp'];
                    
                    if i == 19
                        1;
                    end
                    
                    if ~exist(isdonefile, 'file')
                    
                        [startimage, ~, ~, tmpfile] = chunk_start(tmpfile);
                        if startimage
                            try
                                percellparam(dna_image_files{i},cell_image_files{i},...
                                    prot_image_files{i},mask_image_files{i}, i, param)
                            catch err
                                disp(['Skipping image ' num2str(i) ' due to error'])
                                getReport( err, 'extended' )
                            end
                                chunk_finish(tmpfile)
                                system(['touch ' isdonefile]);

                        else
                            disp(['Image ' num2str(i) ' currently being operated on. Skipping.']);
                            startmodel = false;

                            cellCounter(c) = i;
                            tmpCounter{c} = tmpfile;
                            c = c+1;
                        end

                    end
                    
                end
            end
        else
            disp('Preprocessing flag set to false. Using only currently existing preprocessing results in the temp directory')
        end
        
        if ~startmodel
            %check to see if the tmp files finished just incase
            if any(cellfun(@(x) exist(x, 'file'), tmpCounter))
                disp('The following images are still processing:')
                for i = 1:length(cellCounter)
                    disp([num2str(cellCounter(i)) ': ' tmpCounter{i}])
                end
                
                %if the model is diffeomorphic, we can still build
                %intermediate parts with the data that currently exists
                if ~isdiffeomorphic
                    model = [];
                    return;
                end
            end
        end
        
        %D. Sullivan 6/17/13 
        %With the percell features computed, visualize the distributions of
        %some interesting parameters
        if isfield(param,'percellreport') && ...
            (param.percellreport == true || param.percellreport == 1)
            
            model2report_percell(param);
        end

        %Now load the relevant files and create single cell arrays for each
        %compartment
 
        %D. Sullivan 6/12/13 all this is now taken care of by set_temp_result_folders.m
        %check existence of temporary folder and make on if it doesn't exist
%         if ~exist( [ pwd filesep 'temp'], 'dir' )
%             mkdir( [ pwd filesep 'temp'] );
%         end
        
%         try
            %icaoberg april 17, 2012
        model.dimensionality = '3D';

        if strcmpi(trainFlag, 'microtubule')
            model.proteinShape = img2microtubule_model(prot_image_files, param);
        end
        
        %gj jul 23, 2013 add diffeomorphic model
        if isdiffeomorphic

            diff_model = train_diffeomorphic_model(param);
            if strcmpi( trainFlag, 'all' ) | strcmpi(trainFlag, 'framework');
                model.nuclearShapeModel = diff_model;
                model.cellShapeModel = diff_model;
            elseif strcmpi(trainFlag, 'nuc')
                model.nuclearShapeModel = diff_model;
            elseif strcmpi(trainFlag, 'cell')
                model.cellShapeModel = diff_model;
            end

        else
            if verbose
               %clc; 
               fprintf( 1, '%s\n', 'Training nuclear shape model' );
            end

            %D. Sullivan 6/12/13 refactored to use per-cell parameters
            %Nuclear model
            model.nuclearShapeModel = train_nuc_shape_model( param.nuctemppath,...
                    param.tempparent,param );
            
            %gj aug 29, 2012
            %passes in 'param' now
%             model.nuclearShapeModel = train_nuc_shape_model( ...
%                 dnaImagesDirectoryPath, ...
%                 cellImagesDirectoryPath, ...
%                 proteinImagesDirectoryPath, ...
%                 param );
            %D. Sullivan 6/12/13 refactored to use per-cell features
            %Cell model
            if strcmpi(param.train.flag,'all')||strcmpi(param.train.flag,'framework') || strcmpi(param.train.flag,'microtubule')
                fprintf( 1, '%s\n', 'Training cell shape model' );

                %G. Johnson 11/1/14
                %Moved the protein model file checking and loading into
                model.cellShapeModel = train_cell_shape_model3(param.celltemppath,...
                   param.tempparent);

            end

            if strcmpi( trainFlag, 'all' )
                if verbose
                  %clc; 
                  fprintf( 1, '%s\n', 'Training protein model' );
                end
                %D. Sullivan 6/12/13 refactored to use percell features 
                %Note: param contains all temp path info already from
                %set_temp_result_folders.m
                %Prot model

                %G. Johnson 11/1/14
                %Moved the protein model file checking and loading into
                %the train_protein_model2 function
                model.proteinShape = train_protein_model2( param );
%                   
%                 %D. Sullivan 2/22/13 added param structure to pass resolution
%                 model.proteinShape = train_protein_model( ...
%                     dnaImagesDirectoryPath, ...
%                     cellImagesDirectoryPath, ...
%                     proteinImagesDirectoryPath, ...
%                     param );
            end
        end
        %grj 7/9/13 check for fields so the model wont crash

       %icaoberg 22/02/2013
        if isfield(param.model, 'original_resolution')
            model.info.original_resolution = param.model.original_resolution;
        else
            model.info.original_resolution = 'n/a';
        end

        if isfield(param.model, 'downsampling')
            model.info.downsampling_vector = param.model.downsampling;
        else
            model.info.downsampling_vector = 1;
        end
       %D. Sullivan 6/12/13 removed. already set and misspelled 
%            model.nuclearShapeModel = struct('resoluton', param.model.resolution);
        if isfield(param.model, 'resolution')
            model.cellShapeModel.resolution = param.model.resolution;
        else
            model.cellShapeModel.resolution = 'n/a';
        end
       %D. Sullivan 2/24/2013 This should be set already in
       %train_protein_model
       %model.proteinShape.resolution = param.model.protein_resolution;
%         catch err
%            %icaoberg april 17, 2012
%            %returns empty model if model cannot be trained
%                      
%            model = [];
%            warning('CellOrganizer: Unable to train 3D generative model.');
%           
%            %icaoberg 06/02/2013
%            if debug
%                getReport( err, 'extended' )
%            end
%            
%            if ~debug
%               rmdir( 'temp', 's' );
%            end
%            
%            return
%         end
    otherwise
        warning(['Unknown dimensionality ' ...
            dimensionality '. Exiting method.' ]);
        model = [];
        return
end

end%img2model

function [debug] = setDebugParam(param)
%set debug default to true
    try
        debug = param.debug;
        if ischar(debug)
            if strcmpi(debug, 'false')
               debug = false;
            else 
               debug = true;
            end
        else
            debug = logical(debug);
        end
    catch
        debug = true;
    end 
end

function [verbose] = setVerboseParam(param)
%set verbose default to true
    try
        verbose = param.verbose;
        if ischar(verbose)
            if strcmpi(verbose, 'false')
               verbose = false;
            else 
               verbose = true;
            end
        else
            verbose = logical(verbose);
        end
    catch
        verbose = true;
    end 
end
