function [trefout,zmaxsecout] = skewview(varargin)
% skewview(ref, targ, varargin)   Visualize alignment between audio files
%   Audio files <ref> and <targ> are loaded into memory, then the
%   maxima of the short-time cross correlation is calculated and
%   plotted, to facilitate analysis of the time alignment between
%   files.  vararg options include:
%     -start <time>  start analysis at this point in the files (0)
%     -end   <time>  end analysis at this point in the files (end)
%     -disp  <0/1>  generate graphical representation (1)
%     -win   <time>  window length for xcorr analysis in sec (4.0)
%     -hop   <time>  hop between successive windows in sec (0.5)
%     -maxlag <time> largest lag time to consider (2.0)
%     -peakth <val>  xcorr threshold to be considered a match (0.1)
% 2011-03-15 Dan Ellis dpwe@ee.columbia.edu

% $Header: /u/drspeech/data/RATS/code/skewview/RCS/skewview.m,v 1.3 2011/07/19 19:03:40 dpwe Exp dpwe $

VERSION = 0.90;
DATE = 20140219;

PROGNAME = ['skewview v',num2str(VERSION),' of ',num2str(DATE)];

newargs = 1;
if length(varargin) > 0
  ref = varargin{1};
  if ref(1) ~= '-'
    % looks like old-style unlabeled first args
    ref = varargin{1}; targ = varargin{2}; 
    varargin = varargin(3:end);
    newargs = 0;
  end
end

[P,XTRA] = get_params(varargin, ...
                      { 'ref'  , '', 'reference audio wavfile', ...
                        'targ' , '', 'target audio wavfile', ...
                        'start',  0, 'start at this point in files', ...
                        'end',    0, 'end analysis at this point', ...
                        'dur',   -1, 'limit analysis to this much audio', ...
                        'win', 10.0, 'xcorr analysis window in sec', ...
                        'hop',  2.0, 'hop between success windows in sec', ...
                        'maxlag', 0, 'largest lag to consider (dlft win)', ...
                        'peakth', 0.2, 'threshold of max to count as peak', ...
                        'initialdelay', NaN, 'center around this t_targ', ...
                        'estinitialdelay', 1, 'estimate targ-ref by global xcorr', ...
                        'initialdownsamp', 4, 'downsample by before initial xcorr', ...
                        'samplerate', 0, 'resample to this before comparison', ...
                        'fitthresh', 2.0, 'controls inclusion of outliers in lin fit', ...
                        'alignout', '', 'name for time-warped target audio output', ...
                        'pngout', '', 'name for PNG-format screen dump', ...
                        'textout', '', 'name for text-format <time skew> pairs', ...
                        'corrout', 0, 'include normalized corr vals in textout', ...
                        'minspread', 0.1, 'minimum spread of Y axis (sec)', ...
                        'plotsgrams', 0, 'add specgram plots above xcorr', ...
                        'disp',   1, 'enable (disable) graphic display' ...
                      }, ...
                      PROGNAME);

HELP = 0;
if length(XTRA) > 0
  HELP = length(strmatch('-help',XTRA,'exact')) > 0;
  if ~HELP
    disp(['Unrecognized options:',sprintf(' %s',XTRA{1:end})]);
    HELP = 1;
  end
end

if newargs
  if length(P)
    ref = P.ref;
    targ = P.targ;
  else
    ref = '';
    targ = '';
  end
end

if length(ref) == 0 || strcmp(ref,'-help') == 1; HELP = 1; end
  
if HELP
  %disp('usage...');
  return
end

%
start = P.start; endt = P.end; durt = P.dur; doplot = P.disp; 
xcorrwinsec = P.win; xcorrhopsec = P.hop; xcorrmaxlagsec = P.maxlag; 
xcorrpeakthresh = P.peakth;
initialdelay = P.initialdelay; estinitialdelay = P.estinitialdelay;
initialdownsamp = P.initialdownsamp; samplerate = P.samplerate; 
fitthresh = P.fitthresh; alignout = P.alignout; 
pngout = P.pngout; textout = P.textout; corrout = P.corrout; 
minspread = P.minspread; plotsgrams = P.plotsgrams;
%
%if length(alignout) == 0; alignout = '<alignout>.wav'; end

%if nargin < 3; start = 0; end
%if nargin < 4; endt = 0; end

%xcorrwinsec = 4.0;
%xcorrhopsec = 0.5;
%xcorrmaxlagsec  = 2.0;
%xcorrpeakthresh = 0.1;

if xcorrmaxlagsec == 0
  xcorrmaxlagsec = xcorrwinsec;
end

if isnan(initialdelay)
  initialdelay = 0;
else
  % specifying an initial delay defeats estinitialdelay
  estinitialdelay = 0;
end

if start+initialdelay < 0
  start = -initialdelay;
end

if samplerate == 0
  sr = [];
else
  sr = samplerate;
end

ddisp(['Reading ref  ',ref,' ...'], doplot);
if durt < 0
  durt = max(0,endt-start);
else
  endt = start+durt;
end
[dr,sr] = audioread(ref,sr,1,start,durt);
refstart = start;

ddisp(['Reading targ ',targ,' ...'], doplot);
[dt,sr] = audioread(targ,sr,1,start+initialdelay, durt);
targstart = start+initialdelay;

%if samplerate > 0 && samplerate ~= sr
%  ddisp(['Resampling both to ',num2str(samplerate), ' Hz...'], doplot);
%  dt = resample(dt, samplerate, sr);
%  dr = resample(dr, samplerate, sr);
%  sr = samplerate;
%end

if estinitialdelay
  % first-pass xcorr to find ref in targ
  ddisp(['Making initial coarse alignment with initialdownsamp=', ...
         num2str(initialdownsamp),'...'], doplot);
  [n,xc] = find_skew(dt, dr, [], initialdownsamp);
  newinitialdelay = n/sr;
  if n > 0
    % dr starts midway through dt, chop off first part of dt
    dt = dt((n+1):end);
    targstart = targstart+newinitialdelay;
  else
    % chop off first part of ref
    dr = dr(((-n)+1):end);
    refstart = refstart-newinitialdelay;
  end
  initialdelay = initialdelay + newinitialdelay;
  ddisp(sprintf('New initial delay = %.3f sec', initialdelay), doplot);
else
  newinitialdelay = 0;
end

% truncate to shorter of the two
% AFTER initialdelay estimation
if length(dr) > length(dt)
  dr = dr(1:length(dt));
elseif length(dt) > length(dr)
  dt = dt(1:length(dr));
end

%disp(['refstart=',num2str(refstart),' targstart=',num2str(targstart)]);

xcorrwin = round(sr * xcorrwinsec);
xcorrmaxlag = round(sr * xcorrmaxlagsec);
xcorrhop = round(sr * xcorrhopsec);

ddisp(['Calculating short-time cross-correlation...'], doplot);
%[Z,E] = stxcorr(dt,dr,xcorrwin,xcorrhop,xcorrmaxlag);
[Z,E] = new_stxcorr(dr,dt,xcorrwin,xcorrhop,xcorrmaxlag);

% normalized xcorr
ZN = Z.*repmat(1./E,size(Z,1),1);

if -min(ZN(:))>max(ZN(:))  % -ve correlation dominates
  zsgn = -1;
else
  zsgn = 1;
end

[zmax,zmaxpos] = max(zsgn*ZN);

% remove points where correlation is much lower than peak
zmaxpos(find(zmax<(xcorrpeakthresh*max(zmax)))) = NaN;

% actual times that corresponds to
%zmaxsec = initialdelay + (zmaxpos-xcorrmaxlag-1)/sr;
zmaxsec = (targstart - refstart) + (zmaxpos - xcorrmaxlag - 1)/sr;
% hence best linear fit?
%tref = start + newinitialdelay + [1:length(zmaxpos)]*xcorrhop/sr;
tref = refstart + (xcorrwinsec/2) + [0:length(zmaxpos)-1]*xcorrhopsec;

% make the specgrams line up - don't use the initial offset
if plotsgrams 
%  tref = tref - (start - initialdelay);
end

[a,b,lfs,lfp] = linfit(tref, zmaxsec, fitthresh); %,1 for debug in linfit
% [a,b,s] = linfit(x,y,t,viz)
%   Try to find a linear regression of y onto x
%   s.t. y[n] = a x[n] + b
%  disp(sprintf(' Lin fit: t_targ = %.6f t_ref %s %.3f', 1+a, ...
%               '+'+('-'-'+')*(b<0), abs(b)));
%disp(['lin fit sd = ',num2str(s),' prop pts = ',num2str(p)]);

if doplot

  ddisp('Plotting...', doplot);

  if plotsgrams
    fftwin = 0.030;
    fftlen = 2^round(log(fftwin*sr)/log(2));
    pl1 = subplot(411);
    specgram(dr, fftlen, sr);
    caxis(max(caxis())+[-80 0]);
    colorbar;
    title(['ref - ',ref], 'interpreter', 'none');
    pl2 = subplot(412);
    specgram([zeros(round(initialdelay*sr),1);dt], fftlen, sr);
    caxis(max(caxis())+[-80 0]);
    colorbar;
    title(['targ - ',targ], 'interpreter', 'none');
    pl3 = subplot(212);
  end
  
  global greenline
  greenline = [];
  
  ll = initialdelay + [-xcorrmaxlag:xcorrmaxlag]/sr;

  % plot
%  subplot(212)
%  hist(zmaxsec, [min(zmaxsec):.02:max(zmaxsec)]);
%  subplot(211)

%  editlist = editlist_find(tref, zmaxsec);
  
  imagesc(tref,ll,ZN); axis('xy');
  colormap(1-gray);
  colorbar

  hold on; 
  plot(tref, zmaxsec,'.r');

  %attb = a.*tref+b;
  %greenline = plot(tref, attb, '-g');
  % interactive
  pt1 = [tref(1),a*tref(1)+b];
  pt2 = [tref(end), a*tref(end)+b];
  hh(1) = plot(pt1(1), pt1(2), '.g');
  hh(2) = plot(pt2(1), pt2(2), '.g');
  replot(pt1, pt2);

  if plotsgrams
    linkaxes([pl1, pl2, pl3], 'x');
  end
  
%  moveit2skewview(hh(1), @(pt) fprintf(1,'A %d %d\n',pt(1),pt(2)) , 0, 1);
%  moveit2skewview(hh(2), @(pt) fprintf(1,'B %d %d\n',pt(1),pt(2)) , 0, 1);
% pt(2) end used to be a*tref(end) + b
% to make it larger, we'll need to increase a by Dpt/tref(end)
  moveit2skewview(hh(1), 0, 1, @(pt) replot(pt, []), @finalcallback);
  moveit2skewview(hh(2), 0, 1, @(pt) replot([], pt), @finalcallback);
  mnattb = min(pt1(2),pt2(2));
  mxattb = max(pt1(2),pt2(2));
  axofact = 8;
  dttb = max(axofact/2*minspread, mxattb - mnattb);
  axis([min(tref), max(tref), mnattb-dttb/axofact, mxattb+dttb/axofact]);
  % -----------
  hold off;
  grid
  [pr,nr,er] = fileparts(ref);
  [pt,nt,et] = fileparts(targ);
  title(['Ref: ', nr, ' Targ: ', nt],'interpreter','none'); % don't map "_"
  ylabel('t_targ - t_ref / sec', 'interpreter', 'none')
  xlabel('t_ref / sec', 'interpreter', 'none');
  %xlabel('t_targ / sec', 'interpreter', 'none');
  %disp('Press any key to finish...');
  %pause;

  if length(pngout)
    print('-dpng',pngout);
    close();
  end
  
  % Make it so the output audio is written when the figure is
  % closed
  if length(alignout)
    %rewrite_aligned_audio(targ, alignout, b, 1+a);
    % Setup params
    global closeData
    closeData.targ = targ; closeData.alignout = alignout;
    closeData.b = b; closeData.a = a;
    set(gcf, 'CloseRequestFcn',@rewrite_on_close)
  end
  
end

if length(textout) > 0
  % dump actual skew times as output
  STDOUT = 1;  % constant value for fprintf
  if strcmp(textout, '-')
    fp = STDOUT;
  else
    fp = fopen(textout, 'w');
    if fp == 0
      error(['unable to write to "',textout,'"']);
    end
  end
%  tref = start + max(0, newinitialdelay) + [1:length(zmaxpos)]*xcorrhop/sr;
  gg = initialdelay + (zmaxpos-xcorrmaxlag-1)/sr;
  for i = 1:length(tref)
    if ~isnan(gg(i))
      if corrout
        fprintf(fp,'%.3f\t%.3f\t%.3f\n',tref(i),gg(i),zmax(i));
      else
        fprintf(fp,'%.3f\t%.3f\n',tref(i),gg(i));
      end
    end
  end
  if fp ~= STDOUT
    fclose(fp);
    ddisp(['Timing skews written to ',textout], doplot);
  end

end

if strcmp(textout, '-') == 0
  % text report (unless we wrote the textout to stdout
  disp(['+++ SkewView v',num2str(VERSION),' for ',targ]);
  disp([' ref=',ref]);
  disp([' times start=',num2str(start),' end=',num2str(endt)]);
  disp([' win=',num2str(xcorrwinsec),' hop=',num2str(xcorrhopsec), ...
        ' maxlag=',num2str(xcorrmaxlagsec), ' peakth=',num2str(xcorrpeakthresh)]);
  goodlags = find(~isnan(zmaxpos));
  lagts = (zmaxpos(goodlags)-xcorrmaxlag-1)/sr;
  dlagts = lagts - a * tref(goodlags);
  %disp([' proportion of frames above threshold=',num2str(mean(~isnan(zmaxpos)))]);
  disp(sprintf(' Lin fit stats: sd = %.6f prop pts = %.3f',lfs,lfp));
  disp(sprintf(' Lin fit: t_targ = (1 %c %.6f) t_ref %s %.3f', ...
               '+'+('-'-'+')*(a<0), abs(a), ...
               '+'+('-'-'+')*(b<0), abs(b)));
  disp([' MEDIAN LAG = ',sprintf('%.3f',median(lagts+initialdelay)), ...
        ' s, STDDEV = ', ...
        sprintf('%.3f', std(dlagts)),' s, ',...
        ' ABVTHRESH = ', ...
        sprintf('%.3f', mean(~isnan(zmaxpos)))]);
end

if doplot == 0 && length(alignout)
  % No graphics window to wait to close, so write output right now
  rewrite_aligned_audio(targ, alignout, b, 1+a);
end

if nargout > 0
  trefout = tref;
  zmaxsecout = zmaxsec;
end


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function y = applystartend(x,sr,START,END)
% Trim a waveform to match specified start/end times
y = x;
if START > 0 || END > START
  if END < START; END = length(y)/sr; end
  ENDS = round(END*sr);
  if ENDS > length(y); ENDS = length(y); end
  y = y((round(START*sr)+1):ENDS);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function ddisp(str,doit)
% Optional display
if doit
  disp(str)
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function replot(pt1, pt2)
global greenline 
persistent p1 p2
if ~isempty(greenline); delete(greenline); end
if ~isempty(pt1); p1 = pt1; end
if ~isempty(pt2); p2 = pt2; end
dp = p2 - p1;
ts = [0.1 0.9];
hold on;
greenline = plot(p1(1) + dp(1) * ts, p1(2) + dp(2) * ts, '-g');
hold off;

% record the new points for the end, later
finalcallback(p1, p2)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function finalcallback(pt1, pt2)
persistent p1 p2
if nargin > 0
  % finalcallback with args is just replot sharing its points
  if ~isempty(pt1); p1 = pt1; end
  if ~isempty(pt2); p2 = pt2; end
else
  % finalcallback with no args means to perform final action

  % slope is dy/dx
  a = ( p2(2) - p1(2) ) / ( p2(1) - p1(1) );
  % offset is y|x = 0
  b = p1(2) - a*p1(1);
  
  disp(sprintf(' Lin fit: t_targ = %.6f t_ref %s %.3f', 1+a, ...
               '+'+('-'-'+')*(b<0), abs(b)));

  % Record new values where rewrite_on_close can find them
  global closeData
  closeData.b = b; closeData.a = a;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function rewrite_on_close(src,event)
% function to set as CloseRequestFcn on the figure 
% which will rewrite resampled audio only when the main figure is
% closed.  
% see
% http://www.mathworks.com/help/matlab/creating_plots/the-figure-close-request-function.html

% Do actually close the window
delete(gcf);
drawnow;

% Call the rewrite function
global closeData
if length(closeData.alignout)
  rewrite_aligned_audio(closeData.targ, closeData.alignout, ...
                        closeData.b, 1+closeData.a);
end
