function varargout = domaincolor(f, varargin)
%DOMAINCOLOR   Domain color plot of a complex function.
%
%   DOMAINCOLOR(f, __), with f a function handle, draws a domain color plot
%      of f(z) in the complex plane.
%
%   DOMAINCOLOR(f, [rmin rmax imin imax], __) uses the provided range. If
%      none is provided, [-1 1 -1 1] is used.
%
%   Further options:
%    - 'grid'  or  'g'
%         Additionally show the integer grid.
%    - 'magnitude'  or  'm'
%         Additionally show a contour of the function's magnitude.
%    - 'logmagnitude'  or  'l'
%         Use the log of the magnitude.
%    - 'domaincolor'  or  'd'
%         Show integer grid and magnitude.
%    - 'pixels', [N M]  or  'p', [N M]
%         Sample N x M pixels, default [720 720].

%  Copyright (c) 2023, Evert Provoost,
%     Department of Computer Science, KU Leuven.

%% Parse input
% set defaults
ax = [-1 1 -1 1]; % default axis
pxs = [720 720];  % default number of pixels
intg = 0;         % don't show integer grid by default
magn = 0;         % don't show magnitude by default
lmagn = 0;        % don't take the log by default

setax = 0;
if nargin < 1
    error('DOMAINCOLOR:inputs', 'No function f(z) provided.')
end
j = 1;
while j < nargin
    j = j+1;
    v = varargin{j-1};
    % user specified axis
    if ~ischar(v) && ~isstring(v)
        if setax
            warning('DOMAINCOLOR:inputs', 'Axes already provided.')
        end
        ax = v;
        if length(ax) == 1
            ax = ax * [-1 1 -1 1];
        end
        setax = 1;
    % user wants integer grid
    elseif isequal(v, 'grid') || isequal(v, 'g')
        intg = 1;
    % user wants magnitude on plot
    elseif isequal(v, 'magnitude') || isequal(v, 'm')
        magn = 1;
    % user wants both
    elseif isequal(v, 'domaincolor') || isequal(v, 'd')
        intg = 1;
        magn = 1;
    % user wants log of magnitude
    elseif isequal(v, 'logmagnitude') || isequal(v, 'l')
        magn = 1;
        lmagn = 1;
    % user specified number of pixels
    elseif isequal(v, 'pixels') || isequal(v, 'p')
        if j >= nargin || ~isnumeric(varargin{j})
            error('DOMAINCOLOR:inputs', 'Number of pixels not provided.')
        end
        pxs = varargin{j};
        if length(pxs) == 1
            pxs = pxs * [1 1];
        end
        j = j+1;
    else
        error('DOMAINCOLOR:inputs', ['Unrecognized input "',v,'".'])
    end
end

%% Generate domain coloring
% check if needed function exists
if ~exist('lab2rgb', 'file')
    error(['lab2rgb is not available. ' ...
           'In Octave: install and/or load the image package.'])
end

% compute f(z)
x = linspace(ax(1), ax(2), pxs(1));
y = linspace(ax(3), ax(4), pxs(2));
[xx,yy] = meshgrid(x, y);
zz = xx + 1i*yy;
try
    fzz = f(zz);
catch
    warning('DOMAINCOLOR:inputs', ['f fails on array inputs, ' ...
            'plotting could be slow.'])
    fzz = arrayfun(f, zz);
end

% custom color map in L*a*b* space
mp = @(A) cat(3,12*cos(3*A-pi)+67,46*cos(A+.4)-3,46*sin(A+.4)+16);
% generate phase coloring
C = mp(mod(angle(fzz), 2*pi));
% add magnitude if requested
if magn
    mzz = abs(fzz);
    if lmagn
        mzz = log(mzz);
    end
    C(:,:,1) = C(:,:,1) + 20*mod(mzz,1) - 10;
end
% add integer grid if requested
if intg
    C = abs(sin(pi*real(fzz)).*sin(pi*imag(fzz))).^.06 .* C;
end

% now show plot
h = image([ax(1), ax(2)], [ax(3), ax(4)], lab2rgb(C));
% image defaults to y inverted, fix that
set(gca, 'YDir', 'normal') 
% smooth image on newer versions
if isprop(h, 'Interpolation')
    set(h, 'Interpolation', 'bilinear')
end
% keep axes proportional
axis equal

% return image if asked
if nargout > 0
    varargout = {h};
end
