import numpy as np
import argparse
import os
from skimage.draw import ellipse
from astropy.io import fits

def define_options():
    parser = argparse.ArgumentParser()
    parser.add_argument('--catalog', dest='catalog_path', type=str, required=False,
                        default=None, help='path to the object catalog (default: "%(default)s"')
    parser.add_argument('--flag_image', dest='flag_path', type=str, required=False,
                        default=None, help='path to the flag image (default: "%(default)s"')
    parser.add_argument('--output_dir', dest='output_dir', type=str,
                        default=None, help='path to the output mask image (default: "%(default)s"')
    parser.add_argument('--margin', dest='margin', type=float,
                        default=0, help='margin for major and minor axes')
    return parser

def mask_object(img, x_image, y_image, A_image, B_image, theta_image, margin=0):
    rr, cc = ellipse(
        r=y_image,
        c=x_image,
        r_radius=B_image,
        c_radius=A_image,
        shape=img.shape,
        rotation=np.deg2rad(-theta_image)
    )
    img[rr, cc] = 1

def create_mask_catalog(catalog_path, flag_path, output_dir, margin=0):
    catalog_dir = os.path.dirname(catalog_path)
    catalog_name = os.path.basename(catalog_path)
    flag_dir = os.path.dirname(flag_path)
    flag_name = os.path.basename(flag_path)

    # Read header and data from flag image
    hdu = fits.open(flag_path)
    header0 = hdu[0].header
    header1 = hdu[1].header
    flag_img = fits.getdata(flag_path)
    hdu.close()

    # Read data from detection catalog
    hdu = fits.open(catalog_path)
    catalog = hdu[1].data
    hdu.close()

    # Setup mask image
    mask = flag_img.copy()
    mask[mask != 0] = 1

    # Columns used for mask creation
    x_image = catalog["X_IMAGE"]
    y_image = catalog["Y_IMAGE"]
    A_image = catalog["A_IMAGE"]
    B_image = catalog["B_IMAGE"]
    theta_image = catalog["theta_image"]

    for i in range(len(catalog)):
        mask_object(
            img=mask,
            x_image=x_image[i],
            y_image=y_image[i],
            A_image=A_image[i],
            B_image=B_image[i],
            theta_image=theta_image[i],
            margin=margin
        )
    
    # Save flag image
    fname = os.path.join(output_dir, flag_name.replace("flg_L1", "mask_L1"))
    hdu1 = fits.PrimaryHDU(header=header0)
    hdu2 = fits.ImageHDU(mask, header=header1)
    hdu1 = fits.HDUList([hdu1, hdu2])
    hdu1.writeto(fname, output_verify='ignore', overwrite=True)

if __name__ == "__main__":
    args = define_options().parse_args()
    create_mask_catalog(
        catalog_path=args.catalog_path,
        flag_path=args.flag_path,
        output_dir=args.output_dir,
        margin=args.margin
    )