#!/usr/bin/python

# Sine wave fitting.  This runs about 5x faster than realtime on raw data.

import nilmtools.filter
import nilmdb.client
from numpy import *
from scipy import *
#import pylab as p
import operator

def main(argv = None):
    f = nilmtools.filter.Filter()
    parser = f.setup_parser("Sine wave fitting")
    group = parser.add_argument_group("Sine fit options")
    group.add_argument('-c', '--column', action='store', type=int,
                       help='Column number (first data column is 1)')
    group.add_argument('-f', '--frequency', action='store', type=float,
                       default=60.0,
                       help='Approximate frequency (default: %(default)s)')

    # Parse arguments
    try:
        args = f.parse_args(argv)
    except nilmtools.filter.MissingDestination as e:
        rec = "float32_4"
        print "Source is %s (%s)" % (e.src.path, e.src.layout)
        print "Destination %s doesn't exist" % (e.dest.path)
        print "You could make it with a command like:"
        print "  nilmtool -u %s create %s %s" % (e.dest.url, e.dest.path, rec)
        raise SystemExit(1)

    if args.column is None or args.column < 1:
        parser.error("need a column number >= 1")
    if args.frequency < 0.1:
        parser.error("frequency must be >= 0.1")

    f.check_dest_metadata({ "sinefit_source": f.src.path,
                            "sinefit_column": args.column })
    f.process_numpy(process, args = (args.column, args.frequency))

def process(data, interval, args, insert_function, final):
    (column, f_expected) = args
    rows = data.shape[0]

    # Estimate sampling frequency from timestamps
    fs = 1e6 * (rows-1) / (data[-1][0] - data[0][0])

    # Pull out about 3.5 periods of data at once;
    # we'll expect to match 3 zero crossings in each window
    N = max(int(3.5 * fs / f_expected), 10)

    # If we don't have enough data, don't bother processing it
    if rows < N:
        return 0

    # Process overlapping windows
    start = 0
    num_zc = 0
    while start < (rows - N):
        this = data[start:start+N, column]
        t_min = data[start, 0]/1e6
        t_max = data[start+N-1, 0]/1e6

        # Do 4-parameter sine wave fit
        (A, f0, phi, C) = sfit4(this, fs)

        # Check bounds.  If frequency is too crazy, ignore this window
        if f0 < (f_expected/2) or f0 > (f_expected*2):
            print "frequency", f0, "too far from expected value", f_expected
            start += N
            continue

        #p.plot(arange(N), this)
        #p.plot(arange(N), A * cos(f0/fs * 2 * pi * arange(N) + phi) + C, 'g')

        # Period starts when the argument of cosine is 3*pi/2 degrees,
        # so we're looking for sample number:
        #     n = (3 * pi / 2 - phi) / (f0/fs * 2 * pi)
        zc_n = (3 * pi / 2 - phi) / (f0 / fs * 2 * pi)
        period_n = fs/f0

        # Add periods to make N positive
        while zc_n < 0:
            zc_n += period_n

        last_zc = None
        # Mark the zero crossings until we're a half period away
        # from the end of the window
        while zc_n < (N - period_n/2):
            #p.plot(zc_n, C, 'ro')
            t = t_min + zc_n / fs
            insert_function([[t * 1e6, f0, A, C]])
            num_zc += 1
            last_zc = zc_n
            zc_n += period_n

        # Advance the window one quarter period past the last marked
        # zero crossing, or advance the window by half its size if we
        # didn't mark any.
        if last_zc is not None:
            advance = min(last_zc + period_n/4, N)
        else:
            advance = N/2
        #p.plot(advance, C, 'go')
        #p.show()

        start = int(round(start + advance))

    # Return the number of rows we've processed
    print "Marked", num_zc, "zero-crossings in", start, "rows"
    return start

def sfit4(data, fs):
    """(A, f0, phi, C) = sfit4(data, fs)

    Compute 4-parameter (unknown-frequency) least-squares fit to
    sine-wave data, according to IEEE Std 1241-2010 Annex B

    Input:
      data  vector of input samples
      fs    sampling rate (Hz)

    Output:
      Parameters [A, f0,  phi, C] to fit the equation
        x[n] = A * cos(f0/fs * 2 * pi * n + phi) + C
      where n is sample number.  Or, as a function of time:
        x(t) = A * cos(f0 * 2 * pi * t + phi) + C

    by Jim Paris
    (Verified to match sfit4.m)
    """
    N = len(data)
    t = linspace(0, (N-1) / fs, N)

    ## Estimate frequency using FFT (step b)
    Fc = fft(data)
    F = abs(Fc)
    F[0] = 0   # eliminate DC

    # Find pair of spectral lines with largest amplitude:
    # resulting values are in F(i) and F(i+1)
    i = argmax(F[0:int(N/2)] + F[1:int(N/2+1)])

    # Interpolate FFT to get a better result (from Markus [B37])
    U1 = real(Fc[i])
    U2 = real(Fc[i+1])
    V1 = imag(Fc[i])
    V2 = imag(Fc[i+1])
    n = 2 * pi / N
    ni1 = n * i
    ni2 = n * (i+1)
    K = ((V2-V1)*sin(ni1) + (U2-U1)*cos(ni1)) / (U2-U1)
    Z1 = V1 * (K - cos(ni1)) / sin(ni1) + U1
    Z2 = V2 * (K - cos(ni2)) / sin(ni2) + U2
    i = arccos((Z2*cos(ni2) - Z1*cos(ni1)) / (Z2-Z1)) / n

    # Convert to Hz
    f0 = i * fs / N

    ## Fit it
    # first guess for A0, B0 using 3-parameter fit (step c)
    w = 2*pi*f0
    D = c_[cos(w*t), sin(w*t), ones(N)]
    s = linalg.lstsq(D, data)[0]

    # Now iterate 6 times (step i)
    for idx in range(6):
        D = c_[cos(w*t), sin(w*t), ones(N),
              -s[0] * t * sin(w*t) + s[1] * t * cos(w*t) ] # eqn B.16
        s = linalg.lstsq(D, data)[0] # eqn B.18
        w = w + s[3]	# update frequency estimate

    ## Extract results
    A = sqrt(s[0]*s[0] + s[1]*s[1]) # eqn B.21
    f0 = w / (2*pi)
    try:
        phi = -arctan2(s[1], s[0]) # eqn B.22
    except TypeError:
        # something broke down, just return zeros
        return (0, 0, 0, 0)
    C = s[2]

    return (A, f0, phi, C)

if __name__ == "__main__":
    main()