function [results] = psnr_search(clip_dir, test, varargin)
% PSNR_SEARCH
%   Estimate the Y-channel PSNR (PSNR) of all clips and HRCs (Hypothetical 
%   Reference Circuits) in a video test (input argument test) where the 
%   video clips are stored in the specified directory (clip_dir).  The 
%   video clips must have names that conform to the naming convention 
%   test_scene_hrc.yuv, with no extra '_' or '.' in the file names.  "test"
%   is the name of the test, "scene" is the name of the scene, and "hrc" is
%   the name of the HRC.  The name of the original reference clip for the 
%   PSNR calculation must be "test_scene_original.yuv".
%
%   Files must be stored in "Big YUV format", which is a binary format for
%   storing ITU-R Recommendation BT.601 video sequences.  The format can
%   be used for any image size.  In the Big YUV format, all the frames are
%   stored sequentially in one big binary file. The sampling is 4:2:2 and 
%   image pixels are stored sequentially by video scan line as bytes in the
%   following order: Cb1 Y1 Cr1 Y2 Cb3 Y3 Cr3 Y4…, where Y is the 
%   luminance component, Cb is the blue chrominance component, Cr is the 
%   red chrominance component, and the subscript is the pixel number. The 
%   Y signal is quantized into 220 levels where black = 16 and white = 235,
%   while the Cb and Cr signals are quantized into 225 levels with zero 
%   signal corresponding to 128.  Occasional excursions beyond these levels
%   may occur. For example, Y=15 may be produced by a real system, but in 
%   any case, Y, Cb, and Cr values are always between 0 and 255.  The
%   original and processed Big YUV files must have the same number of frames.
%
%   A peak signal of 255 is used for calculation of PSNR.  Double precision
%   calculations are used everywhere.  A 64-bit operating system with at
%   least 4 GB of free memory is recommended since the entire double
%   precision versions of the original and processed sequences must be held
%   in memory.
%
% SYNTAX
%   [results] = psnr_search('clip_dir', 'test', option);
%
% DESCRIPTION
%   This function will process all video clips in the user specified
%   clip_dir and test, estimate the Y-channel PSNR of each clip, and then
%   average these clip results to produce an estimate for each HRC.  The
%   algorithm performs an exhaustive search for max PSNR over plus or minus
%   the spatial_uncertainty (in pixels) and plus or minus the 
%   temporal_uncertainty (in frames).  The processed video segment is fixed
%   and the original video segment is shifted over the search range.  For
%   each spatial-temporal shift, a linear fit between the processed pixels 
%   and the original pixels is performed such that the mean square error of
%   [original-gain*processed+offset] is minimized (hence maximizing PSNR).
%
%   Any or all of the following optional properties may be requested (the
%   first option is required for yuv files).
%
%   'yuv',rows,cols    Specifies the number of rows and cols for the Big
%                      YUV files.
%
%   'sroi',top,left,bottom,right,   Only use the specified spatial region 
%                                   of interest (sroi) for the PSNR
%                                   calculation.  This is the sroi of the
%                                   processed sequence, which remains fixed
%                                   over all spatial shifts.  By default,
%                                   sroi is the entire image reduced by the
%                                   spatial uncertainty.  If the user
%                                   inputs a sroi, allowance must be made
%                                   for the spatial search specified by
%                                   'spatial_uncertainty'.
%
%   'frames',fstart,fstop  Only use the frames from fstart to fstop
%                          (inclusive) to perform the PSNR estimate.  This
%                          specifies the temporal segment of the processed
%                          sequence, which remains fixed over all temporal
%                          shifts.  By default, the temporal segment is the
%                          entire file reduced by the temporal uncertainty.
%                          If the user inputs an fstart and fstop,
%                          allowance must be made for the temporal search
%                          specified by 'temporal_uncertainty'. 
%
%                          
%   'spatial_uncertainty',x,y   Specifies the spatial uncertainty (plus 
%                               or minus, in pixels) over which to 
%                               search.  The processed remains fixed and
%                               the original is shifted.  By default,
%                               this is set to zero.
%
%   'temporal_uncertainty',t  Specifies the temporal uncertainty
%                             (plus or minus, in frames) over which
%                             to search.  The processed remains fixed
%                             and the original is shifted.  By
%                             default, this is set to zero.
%
%   'verbose'   Display output during processing.
%
%
%   The returned variable [results] is a struct that contains the following
%   information for each processed clip i:
%
%   results(i).test    The test name for the video clip.
%   results(i).scene   The scene name for the video clip.
%   results(i).hrc     The HRC name for the video clip.
%   results(i).yshift  The y shift for max PSNR.
%   results(i).xshift  The x shift for max PSNR.
%   results(i).tshift  The time shift for max PSNR.
%   results(i).gain    The gain*processed+offset for max PSNR.
%   results(i).offset
%   results(i).psnr    The maximum PSNR observed over the search.
%
% EXAMPLES
%   These examples illustrate how to call the routine to process the VQEG
%   MM Phase I test scenes, where test scenes from each subjective
%   experiment are stored in a unique directory.  These sequences must
%   first be converted from AVI format to Big YUV format.
%
%   q01 = psnr_search('d:\q01\','q01','yuv',144,176,'sroi',5,5,140,172,...
%       'spatial_uncertainty',1,1,'temporal_uncertainty',8,'verbose');
%   c01 = psnr_search('d:\c01\','c01','yuv',288,352,'sroi',8,8,281,345,...
%       'spatial_uncertainty',1,1,'temporal_uncertainty',8,'verbose');
%   v01 = psnr_search('d:\v01\','v01','yuv',480,640,'sroi',14,14,467,627,...
%       'spatial_uncertainty',1,1,'temporal_uncertainty',8,'verbose');
%

% Define the peak signal level
peak = 255.0;
 
% Add extra \ in clip_dir in case user did not
clip_dir = strcat(clip_dir,'\');
 
% Validate input arguments and set their defaults
is_yuv = 0;
is_whole_image = 1;
is_whole_time = 1;
x_uncert = 0;
y_uncert = 0;
t_uncert = 0;
verbose = 0;
dx=1;  % dx, dy, and dt sizes to use for gain and level offset calculations
dy=1;
dt=1;
cnt=1;
while cnt <= length(varargin),
    if ~isstr(varargin{cnt}),
        error('Property value passed into psnr_search is not recognized');
    end
    if strcmpi(varargin(cnt),'yuv') == 1
        rows = varargin{cnt+1};
        cols = varargin{cnt+2};
        is_yuv = 1;
        cnt = cnt + 3;
    elseif strcmpi(varargin(cnt),'sroi') == 1
        top = varargin{cnt+1};
        left = varargin{cnt+2};
        bottom = varargin{cnt+3};
        right = varargin{cnt+4};
        is_whole_image = 0;
        cnt = cnt + 5;
    elseif strcmpi(varargin(cnt),'frames') == 1
        fstart = varargin{cnt+1};
        fstop = varargin{cnt+2};
        is_whole_time = 0;
        cnt = cnt + 3;
    elseif strcmpi(varargin(cnt), 'spatial_uncertainty') ==1
        x_uncert = varargin{cnt+1};
        y_uncert = varargin{cnt+2};
        cnt = cnt + 3;
    elseif strcmpi(varargin(cnt), 'temporal_uncertainty') ==1
        t_uncert = varargin{cnt+1};
        cnt = cnt + 2;
    elseif strcmpi(varargin(cnt),'verbose') == 1
        verbose = 1;
        cnt = cnt +1;
    else
        error('Property value passed into psnr_search not recognized');
    end
end
 
%  Get a directory listing
files = dir(clip_dir);  % first two files are '.' and '..'
num_files = size(files,1);
 
% Find the HRCs and their scenes for the specified video test
hrc_list = {};
scene_list = {};
for i=3:num_files
    this_file = files(i).name;
    und = strfind(this_file,'_'); % find underscores and period
    dot = strfind(this_file,'.');
    if(size(und,2)==2) % possible standard naming convention file found
        this_test = this_file(1:und(1)-1);  % pick off the test name
        if(strmatch(test,this_test,'exact'))  % test clip found
            this_scene = this_file(und(1)+1:und(2)-1);
            this_hrc = this_file(und(2)+1:dot(1)-1);
            % See if this HRC already exists and find its list location
            loc = strmatch(this_hrc,hrc_list,'exact');
            if(loc)  % HRC already present, add to scene list for that HRC
                if(size(strmatch(this_scene,scene_list{loc},'exact'),1)==0)
                    scene_list{loc} = [scene_list{loc} this_scene];
                end
            else  % new HRC found
                hrc_list = [hrc_list;{this_hrc}];
                this_loc = size(hrc_list,1);
                scene_list(this_loc) = {{this_scene}};
            end
        end
    end
end
 
scene_list = scene_list';
num_hrcs = size(hrc_list,1);
 
% Results struct to store results, shifts are how much the original must be
% shifted with respect to the processed
results = struct('test', {}, 'scene', {}, 'hrc', {}, 'yshift', {}, ...
    'xshift', {}, 'tshift', {}, 'gain', {}, 'offset', {}, 'psnr', {});
 
% Process one HRC at a time to compute average PSNR for that HRC
index = 1;  % index to store results
for i = 1:num_hrcs
    
    psnr_ave = 0;  % initialize the psnr average summer for this HRC
    this_hrc = hrc_list{i};
    if(strmatch('original',this_hrc,'exact')) % Don't process original
        continue;
    end
    num_scenes = size(scene_list{i},2);  % Number of scenes in this HRC
    
    for j = 1:num_scenes
        
        this_scene = scene_list{i}{j};
        results(index).test = test;
        results(index).scene = this_scene;
        results(index).hrc = this_hrc;
        
        % Read original and processed video files
        if (~is_yuv)  % YUV file parameters not specified
            display('Must specify Big YUV rows and cols.  Use yuv input option.');
            return
        else  % YUV file
            % Re-generate the original and processed YUV file name
            orig = strcat(clip_dir, test,'_', this_scene, '_', 'original', '.yuv');
            proc = strcat(clip_dir, test,'_', this_scene, '_', this_hrc, '.yuv');
            % Set/Validate the ROI
            if (is_whole_image) % make ROI whole image less uncertainty
                top = 1+y_uncert;
                left = 1+x_uncert;
                bottom = rows-y_uncert;
                right = cols-x_uncert;
            elseif (top<1 || left<1 || bottom>rows || right>cols)
                display('Requested SROI too large for image size.');
                return;
            end
            % Find the total frames of the input original file
            [fid, message] = fopen(orig, 'r');
            if fid == -1
                fprintf(message);
                error('Cannot open this clip''s bigyuv file, %s', orig);
                return
            end
            % Find last frame.
            fseek(fid,0, 'eof');
            tframes = ftell(fid) / (2 * rows * cols);
            fclose(fid);
            % Find the total frames of the processed file
            [fid, message] = fopen(proc, 'r');
            if fid == -1
                fprintf(message);
                error('Cannot open this clip''s bigyuv file, %s', proc);
                return
            end
            % Find last frame.
            fseek(fid,0, 'eof');
            tframes_proc = ftell(fid) / (2 * rows * cols);
            fclose(fid);
            % Validate that orig and proc have the same number of frames
            if (tframes ~= tframes_proc)
                display('The orig & proc files must have the same length.');
                return
            end
            % Set/Validate the time segment to use
            if (is_whole_time) % use whole time segment less uncertainty
                fstart= 1+t_uncert;
                fstop = tframes-t_uncert;
            elseif (fstart<1 || fstop>tframes)
                display('Requested Temporal segment too large for file size.');
                return
            end
             % Validate the spatial uncertainty search bounds
            if (left-x_uncert < 1 || right+x_uncert > cols)
                display('Spatial x-uncertainty too large for SROI.');
                return;
            end
            if (top-y_uncert < 1 || bottom+y_uncert > rows)
                display('Spatial y-uncertainty too large for SROI.');
                return;
            end
            % Validate the temporal uncertainty search bounds
            if(fstart-t_uncert < 1 || fstop+t_uncert > tframes)
                display('Temporal uncertainty too large for fstart or fstop.');
                return;
            end
            % Read in video and clear color planes to free up memory
            [y_orig,cb,cr] = read_bigyuv(orig,'frames',fstart-t_uncert,...
                fstop+t_uncert,'size',rows,cols,'sroi',top-y_uncert,...
                left-x_uncert,bottom+y_uncert,right+x_uncert);
            clear cb cr;
            [y_proc,cb,cr] = read_bigyuv(proc,'frames',fstart,fstop,...
                'size',rows,cols,'sroi',top,left,bottom,right);
            clear cb cr;
        end
        
        % Convert images to double precision
        y_orig = double(y_orig);
        y_proc = double(y_proc);
        
        [nrows, ncols, nsamps] = size(y_proc);
        % Reshape y_proc for the PSNR calculation:  this stays fixed
        y_proc = reshape(y_proc,nrows*ncols*nsamps,1); % make column vector
            
        % Compute PSNR for each spatial-temporal shift
        best_psnr = -inf;
        best_xshift = 0;
        best_yshift = 0;
        best_tshift = 0;
        best_gain = 1;
        best_offset = 0;
        if(verbose)
            fprintf('\nTest = %s,   Scene = %s,   HRC = %s\n',test, this_scene, this_hrc);
        end
        
        for k = -t_uncert:t_uncert
            for m = -x_uncert:x_uncert
                for n = -y_uncert:y_uncert
                    
                    % Perform gain and level offset calculation
                    this_fit = polyfit(y_proc,reshape(y_orig(1+n+y_uncert:n+y_uncert+nrows,...
                        1+m+x_uncert:m+x_uncert+ncols,1+k+t_uncert:k+t_uncert+nsamps),...
                        nrows*ncols*nsamps,1),1);
                    
                    % Calculate the PSNR
                    this_psnr = 10*(log10(peak*peak)-log10(sum(((this_fit(1)*y_proc+this_fit(2))-...
                        reshape(y_orig(1+n+y_uncert:n+y_uncert+nrows,1+m+x_uncert:m+x_uncert+ncols,...
                        1+k+t_uncert:k+t_uncert+nsamps),nrows*ncols*nsamps,1)).^2)/(nrows*ncols*nsamps)));
                    
                    if(this_psnr > best_psnr)
                        best_psnr = this_psnr;
                        best_yshift = n;
                        best_xshift = m;
                        best_tshift = k;
                        best_gain = this_fit(1);
                        best_offset = this_fit(2);
                        if(verbose)
                            fprintf('dy =%3i, dx =%3i, dt =%3i, gain = %5.4f, off = %5.4f, PSNR = %5.4f\n',...
                                best_yshift,best_xshift,best_tshift,best_gain,best_offset,best_psnr);
                        end
                    end
                    
                end
            end
        end
        
        results(index).yshift = best_yshift;
        results(index).xshift = best_xshift;
        results(index).tshift = best_tshift;
        results(index).gain = best_gain;
        results(index).offset = best_offset;
        results(index).psnr = best_psnr;
        psnr_ave = psnr_ave+best_psnr;
        index = index+1;
        
    end
    
    % Compute average PSNR for this HRC
    psnr_ave = psnr_ave/(num_scenes);
    if(verbose)
        fprintf('HRC = %s, psnr_ave = %5.4f\n',this_hrc, psnr_ave);
    end
    
end