# Rolling shutter measurement from flickering lights
# Idea: https://blog.kasson.com/nikon-z6-7/how-fast-is-the-z7-silent-shutter/
#
# Tested on octave 4.x with the signal and image packages installed
pkg load signal
pkg load image

# some options
vignette_fix = false;
method = "pwelch";

args = {};
for i = 1:length(argv())
    if argv(){i}(1) == "-"
        switch (argv(){i})
            case "-o"
                method = "overlap";
            case "-z"
                method = "zerocross";
            case "-v"
                vignette_fix = true;
            otherwise
                assert(["unknown option %s\n", argv(){i}]);
        end
    else
        args{end+1} = argv(){i};
    end
end

# positional arguments
if length(args) == 2
    # frequency of the light source used in the test
    # for a plain old incandescent bulb in PAL land, use 100
    measure_freq = false;
    light_freq = str2num(args{2});

elseif length(args) == 3 || length(args) == 4
    # if unknown, light frequency can be measured from a camera with known timings
    # (e.g. any camera that runs Magic Lantern)
    measure_freq = true;

    # Main clock: camera-specific constant
    # - 5D2/5D3: 24 MHz
    # - 6D: 25.6 MHz
    # - 550D/60D/600D: 28.8 MHz
    # - 700D/650D/100D/M: 32 MHz
    # - 5D4: 50 MHz
    main_clock = str2num(args{2});

    # FPS timers (from logs, adtg_gui, FPS override submenu etc):
    # ML menu: to see default Canon values, open the submenu without enabling FPS override.
    # Logs or adtg_gui: look for the following registers:
    # - Timer A (register C0F06008 / D0006008)
    # - Timer B (register C0F06014 / D0006014)
    # Add 1 to raw register values.
    timer_a = str2num(args{3});

    # Line clock: Main clock / Timer A
    # Frame rate: Main clock / Timer A / Timer B
    # One tick of timer A reads out an integer number of pixels (usually 2, 4 or 8)
    # One tick of timer B reads out one line (at least until DIGIC 5)
    # DIGIC >= 6: one tick of timer B apparently reads two lines?
    # Or maybe there are two rolling shutter readouts going in parallel?

    if length(args) == 3,
        lines_per_tick = 1;
    else
        lines_per_tick = str2num(args{4});
    end
else
    # help
    printf("\n")
    printf("Rolling shutter measurement from flickering lights\n")
    printf("==================================================\n")
    printf("\n")
    printf("Test picture: \n")
    printf("\n")
    printf(" - DNG frame from LiveView (e.g. silent picture or raw video frame)\n")
    printf(" - scene: blank wall illuminated by some flickering light\n")
    printf("\n")
    printf("Usage: \n")
    printf("\n")
    printf("1) to measure rolling shutter, if you know the frequency of the test light:\n")
    printf("\n")
    printf("       octave rolling.m 1234.DNG 100   #  incandescent bulb, 50 Hz; light pulses at 100 Hz\n")
    printf("\n")
    printf("2) to measure frequency of the test light, you need to specify:\n")
    printf(" - main FPS clock (5D2/5D3: 24 MHz; 550D/60D/600D: 28.8 MHz; 700D/650D/100D/M: 32 MHz; 5D4: 50 MHz)\n")
    printf(" - FPS timer A (from FPS override submenu with main menu disabled, or register C0F06008/D0006008)\n")
    printf("\n")
    printf("       octave rolling.m 1234.DNG 24e6 600      # 5D2 1080p25\n")
    printf("       octave rolling.m 1234.DNG 24e6 572      # 5D2 1080p24/30\n")
    printf("       octave rolling.m 1234.DNG 24e6 480      # 5D3 1080p25\n")
    printf("       octave rolling.m 1234.DNG 24e6 440      # 5D3 1080p24/30\n")
    printf("\n")
    printf("Options:\n")
    printf(" -o: measure repeated pattern by overlapping it onto itself\n")
    printf(" -z: measure repeated pattern by looking at zero crossings\n")
    printf(" -v: attempt to fix vignette in the image\n")
    return
end

# test image
switch (args{1}(end-3:end))
    case { ".dng", ".DNG", ".cr2", ".CR2" }
        # assume Bayer RGGB
        im = read_raw(args{1});
        r  = im(1:2:end,1:2:end);
        g1 = im(1:2:end,2:2:end);
        g2 = im(2:2:end,1:2:end);
        b  = im(2:2:end,2:2:end);
        r = reshape([r  r ]', [size(im,2)/2, size(im,1)])';
        g = reshape([g1 g2]', [size(im,2)/2, size(im,1)])';
        b = reshape([b  b ]', [size(im,2)/2, size(im,1)])';
    otherwise
        # assume 8-bit RGB
        im = im2double(imread(args{1}));
        r  = im(:,:,1);
        g  = im(:,:,2);
        b  = im(:,:,3);
end

# Gaussian blur approximation (box blur on each axis, applied 3 times)
function x = fix_vignette(x)
    strength = round(min(size(x)) / 10);
    bb = ones(strength, 1) / strength;
    f = x;
    f = imfilter(f, bb, 'replicate');
    f = imfilter(f, bb, 'replicate');
    f = imfilter(f, bb, 'replicate');
    f = imfilter(f, bb', 'replicate');
    f = imfilter(f, bb', 'replicate');
    f = imfilter(f, bb', 'replicate');
    x = x - f;
end

# pick the strongest channel
if mean(r(:)) > mean(g(:)) && mean(r(:)) > mean(b(:))
    printf("Using red channel.\n");
    d = r;
elseif mean(b(:)) > mean(g(:)) && mean(b(:)) > mean(r(:))
    printf("Using blue channel.\n");
    d = b;
else
    printf("Using green channel.\n");
    d = g;
end



if vignette_fix,
    printf("Vignette fix...\n");
    d = fix_vignette(d);
end

# extract vertical stripe from the image
p = mean(d');                             # row average
p = p - sum(prctile(p, [10 90])) / 2;     # center the signal
x = 1:length(p);                          # pixel position

# plot the pattern, for debugging
plot(x, p, 'r', 'linewidth', 2); hold on

if 1
    # from https://www.mathworks.com/matlabcentral/answers/160059-finding-the-frequency-value-of-im-signal
    Nfft = 1024 * 1024;
    [Pxx,f] = pwelch(p, gausswin(Nfft), 0.5, Nfft, 1);

    # filter out frequencies outside the expected range
    Pxx(f > 1/20 | f < 1/2000) = 0;

    # find peak frequency
    [~,loc] = max(Pxx);
    num_lines_pwelch = 1 / f(loc);
    printf("Pattern repeats every %d lines (method: pwelch).\n", round(num_lines_pwelch));

    # plot the identified sine, for debugging
    a = 2 * mean([sin(2 * pi * f(loc) * x) .* p]);
    b = 2 * mean([cos(2 * pi * f(loc) * x) .* p]);
    plot(x, a * sin(2 * pi * f(loc) * x) + b * cos(2 * pi * f(loc) * x), 'g', 'linewidth', 1);
    drawnow;
end

if 1
    # find the best overlap
    x = 1:length(p);
    E = [];
    f = 1 ./ [20:length(p)/2];
    for ov = round(1 ./ f)
        E(end+1) = norm(p(1+ov:end) - p(1:end-ov));
    end
    [~,loc] = min(E);
    num_lines_overlap = 1 / f(loc);
    printf("Pattern repeats every %d lines (method: overlap).\n", round(num_lines_overlap));

    ov = round(num_lines_overlap);
    plot(x(1:end-ov), p(1+ov:end), 'b', 'linewidth', 1)
end

if 1
    # look at zero crossings
    z = zerocrossing(x, p);
    num_lines_zerocross = median(z(3:end) - z(1:end-2));
    printf("Pattern repeats every %d lines (method: zerocross).\n", round(num_lines_zerocross));
    plot(z, z * 0, 'or');
    grid on
end

switch (method)
    case "pwelch"
        num_lines = num_lines_pwelch;
        printf("Method: pwelch => %.02f lines\n", num_lines);
    case "overlap"
        num_lines = num_lines_overlap;
        printf("Method: overlap => %.02f lines\n", num_lines);
    case "zerocross"
        num_lines = num_lines_zerocross;
        printf("Method: zerocross => %.02f lines\n", num_lines);
    otherwise
        assert("invalid method")
end

# 5D2 H.264 1080p is recorded from 1872x1053 (best guess)
%~ num_lines = num_lines * 1053 / 1080

if measure_freq,
    printf("Line readout clock: %.2f kHz, i.e. %.2f μs/line (known).\n", main_clock / timer_a / 1000, 1e6 / (main_clock / timer_a));
    printf("Light source frequency: %.2f Hz (measured).\n", 1 / (1 / (main_clock / timer_a) * num_lines));
else
    line_clock = 1 / ((1 / light_freq) / num_lines);
    printf("Light source frequency: %.2f Hz (known).\n", light_freq);
    printf("Line readout clock: %.2f kHz, i.e. %.2f μs/line (measured).\n", line_clock, 1e6 / line_clock);
    printf("Rolling shutter: %.2f ms for %d lines.\n", (1e3 / line_clock) * size(im, 1), size(im, 1));
end

# debug info
print -dpng rolling.png