#!python
# coding: utf-8
"""Clean dirty rows and columns from fits image."""


import os
import argparse
import numpy as np
from astropy.io import fits


parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("filename", type=str)
parser.add_argument("-r", "--row", nargs="?", type=int)
parser.add_argument("-c", "--col", nargs="?", type=int)
parser.add_argument("-o", "--out", nargs="?", type=str)
parser.add_argument("-y", "--yes", action="store_true")
parser.add_argument("-s", "--step", nargs="?", type=int, default=0)
parser.add_argument("-n", "--nighbor", nargs="?", type=int, default=2)
parser.add_argument("-f", "--function", nargs="?", type=str,
                    default="median", choices=["mean", "median"])
args = parser.parse_args()

if args.row is None and args.col is None:
    parser.error("Please specify dirty row or column to clean.")

if args.row is not None and args.col is not None:
    parser.error("Please specify only one from row or column.")

if args.out is None and args.yes is False:
    parser.error("Please specify output filename or enable overwrite.")

if args.out is not None and os.path.exists(args.out):
    parser.error("I am so sorry, but output file can not exists yet.")

if args.step < 0:
    parser.error("Please notify that number of steps must be >= 0.")

if args.nighbor < 1:
    parser.error("Please notify that number of nighbors must be >= 1.")

if args.function == "mean":
    fce = np.mean
elif args.function == "median":
    fce = np.median
else:
    pass

if args.yes:
    image = fits.open(args.filename, mode="update")
else:
    image = fits.open(args.filename, mode="readonly")

head = image[0].header
data = image[0].data

s = args.step
n = args.nighbor

if args.row is not None:
    r = args.row
    data[r-1, :] = fce(np.vstack((data[r-1-n:r-1, :], data[r+s:r+s+n+1, :])),
                       axis=0)
    head["HISTORY"] = "Clean row {} using {} of +-{} pix with step {}"\
                      .format(r, args.function, n, s)

if args.col is not None:
    c = args.col
    data[:, c-1] = fce(np.hstack((data[:, c-1-n:c-1], data[:, c+s:c+s+n+1])),
                       axis=1)
    head["HISTORY"] = "Clean column {} using {} of +-{} pix with step {}"\
                      .format(c, args.function, n, s)

image[0].header = head
image[0].data = data

if args.yes:
    image.writeto(args.filename, overwrite=True)
else:
    image.writeto(args.out)

image.close()