#!/usr/bin/env python3
import nilmtools.filter
import 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)

    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()