seq.jl

include("util.jl")

using SampledSignals

"""
    known_fixed_grid_seq_i(i, qlarge, qsmall, period, phase=1)

If the given index `i` coincides with the given `phase` and `period`, returns 1
with probability `qlarge`. Otherwise, returns 1 with probability `qsmall`.
"""
function known_fixed_grid_seq_i(i, qlarge, qsmall, period, phase=1)
    @assert(phase >= 0 && phase <= period)
    phase = phase % period
    return i % period == phase ? bernoulli(qlarge) : bernoulli(qsmall)
end

"""
    known_fixed_grid_seq(length, qlarge, qsmall, period, phase)

Returns a random repetitive sequence generated by sampling from 2 Bernoulli
random variables with probabilities `qlarge` and `qsmall`. Every `period` entries,
a 1 appears with probability `qlarge` (otherwise 0). In between, the value is 1
with probability `qsmall`.

See Rhythm and Transforms p. 179-180
"""
function known_fixed_grid_seq(length, qlarge, qsmall, period, phase=1)
    return [known_fixed_grid_seq_i(i, qlarge, qsmall, period, phase) for i=1:length]
end

"""
    known_fixed_grid_seq_ll(seq, qlarge, qsmall, period, phase)

Returns the log likelihood that the data in seq was generated by calling known_fixed_grid_seq with the
given parameters.
"""
function known_fixed_grid_seq_ll(seq, qlarge, qsmall, period, phase)
    phase = phase % period

    num_ones_large = 0
    num_zeros_large = 0
    num_ones_small = 0
    num_zeros_small = 0
    for i = 1:lastindex(seq)
        if i % period == phase
            if seq[i] == 1
                num_ones_large += 1
            else
                num_zeros_large += 1
            end
        else
            if seq[i] == 1
                num_ones_small += 1
            else
                num_zeros_small += 1
            end
        end
    end

    # Since the larges and smalls are repeated Bernoulli trials, the
    # probability of observing n 1s and m 0s is p^n(1-p)^m, where p
    # is either `qlarge` or `qsmall`, respectively. Since the larges
    # and smalls are assumed independent, the overall probability of
    # observing `seq` is obtained by multiplying the terms together.
    # The log of each term is of the form nlog(p) + mlog(1-p).
    function ll(p, num_ones, num_zeros)
        return num_ones*log(p) + num_zeros*log(1-p)
    end

    ll_large = ll(qlarge, num_ones_large, num_zeros_large)
    ll_small = ll(qsmall, num_ones_small, num_zeros_small)

    return ll_large + ll_small
end

function unknown_fixed_grid_seq_q(t, qlarge, qsmall, period, phase; width=period/20)
    phase = phase % period
    # TODO: this is asymmetrical if right on a spike, but the
    #       contribution of the next spike should be ~0 anyway...
    prev_spike_idx = div((t-phase), period)
    next_spike_idx = prev_spike_idx + 1
    tprev = phase + prev_spike_idx * period
    tnext = phase + next_spike_idx * period
    prevpulse = gaussianpulse(t; u=tprev, s=width, shift=0, height=qlarge-qsmall)
    nextpulse = gaussianpulse(t; u=tnext, s=width, shift=0, height=qlarge-qsmall)
    return qsmall + prevpulse + nextpulse
end

function unknown_fixed_grid_seq_t(t, qlarge, qsmall, period, phase; width=period/20)
    prob = unknown_fixed_grid_seq_q(t, qlarge, qsmall, period, phase; width=width)
    return bernoulli(prob)
end

"""
    unknown_fixed_grid_seq(duration, qlarge, qsmall, period, phase, width=period/20; samplerate=86)

This can be thought of as a version of known_fixed_grid_seq() where the onsets don't always fall exactly
on grid points. To make that work, the sequence is defined in continuous (sampled) time. Onsets will
tend to cluster around grid points, but the clustering will be inexact unless width is 0.

Unlike known_fixed_grid_seq(), this returns a SampleBuf.

See Rhythm and Transforms p. 182-183
"""
function unknown_fixed_grid_seq(duration, qlarge, qsmall, period, phase, width=period/20; samplerate=86)
    length = Int(ceil(duration * samplerate))
    t = (0:length-1) / samplerate
    samples = unknown_fixed_grid_seq_t.(t, qlarge, qsmall, period, phase; width=width)
    return SampleBuf(samples, samplerate)
end

"""
    unknown_fixed_grid_seq_ll(buf, qlarge, qsmall, period, phase, width=period/20)

Returns the log likelihood that the data in buf was obtained by calling unknown_fixed_grid_seq with the
given parameters.
"""
function unknown_fixed_grid_seq_ll(data, samplerate, qlarge, qsmall, period, phase, width=period/20)
    t = (0:length(data)-1)/samplerate
    probs = unknown_fixed_grid_seq_q.(t, qlarge, qsmall, period, phase; width=width)
    return sum(data .* log.(probs) .+ (1 .- data) .* log.(1 .- probs))
end

function infer_unknown_fixed_grid_seq(buf::SampleBuf; minbpm=40, maxbpm=200, bpmres=5, phaseres=60/(maxbpm*4))
    periods = 60 ./ (minbpm:bpmres:maxbpm)
    bestperiod = nothing
    bestphase = nothing
    bestll = -Inf
    for period in periods
        for phase in 0:phaseres:period
            ll = unknown_fixed_grid_seq_ll(buf.data, buf.samplerate, 0.75, 0.1, period, phase)
            if ll > bestll
                bestll = ll
                bestperiod = period
                bestphase = phase
                println(60/bestperiod, '\t', bestphase, '\t', bestll)
            end
        end
    end

    return (60/bestperiod, bestphase, bestll)
end

function tosymbolic!(featurevec::SampleBuf; percentile=0.5)
    map!(abs, featurevec, featurevec)
    threshold = quantile(featurevec, percentile)
    threshold!(featurevec, threshold)
end

function tosymbolic(featurevec::SampleBuf; percentile=0.5)
    ret = copy(featurevec)
    tosymbolic!(ret; percentile=percentile)
    return ret
end