#!/usr/bin/python
import nilmtools.filter, scipy.signal

def main(argv = None):
    f = nilmtools.filter.Filter()
    parser = f.setup_parser("Median Filter")
    group = parser.add_argument_group("Median filter options")
    group.add_argument("-z", "--size", action="store", type=int, default=25,
                       help = "median filter size (default %(default)s)")
    group.add_argument("-d", "--difference", action="store_true",
                       help = "store difference rather than filtered values")

    try:
        args = f.parse_args(argv)
    except nilmtools.filter.MissingDestination as e:
        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, e.src.layout))
        raise SystemExit(1)

    meta = f.client_src.stream_get_metadata(f.src.path)
    f.check_dest_metadata({ "median_filter_source": f.src.path,
                            "median_filter_size": args.size,
                            "median_filter_difference": repr(args.difference) })

    f.process_numpy(median_filter, args = (args.size, args.difference))

def median_filter(data, interval, args, insert, final):
    (size, diff) = args
    (rows, cols) = data.shape
    for i in range(cols - 1):
        filtered = scipy.signal.medfilt(data[:, i+1], size)
        if diff:
            data[:, i+1] -= filtered
        else:
            data[:, i+1] = filtered
    insert(data)
    return rows

if __name__ == "__main__":
    main()
