"""This python script fixes the FITS filter information based on the fits file
name. One need to install astropy python module before running
this script.
"""
from astropy.io import fits
from astropy.utils.data import get_pkg_data_filename
import os

#All the fits file inside this directory will be checked and fixed if needed
base_dir = "/home/astroberry/Pictures/Astronomy"

for root, dirs, files in os.walk(base_dir, topdown=False):
    for name in files:
        if os.path.join(root, name).endswith("fits"):
            #print(os.path.join(root, name))
            fits_file = os.path.join(root, name)
            d = dict(fits.getheader(fits_file))
            #print(d)
            # if the FILTER attribute is not present iht the header...
            if not "FILTER" in d:
                if 'Red' in name:
                    print("set Red for {}".format(name))
                    fits.setval(fits_file, 'FILTER', value='Red')
                elif 'Green' in name:
                    print("set Green for {}".format(name))
                    fits.setval(fits_file, 'FILTER', value='Green')
                elif 'Blue' in name:
                    print("set Blue for {}".format(name))
                    fits.setval(fits_file, 'FILTER', value='Blue')
                elif 'Luminance' in name:
                     print("set Luminance for {}".format(name))
                     fits.setval(fits_file, 'FILTER', value='Luminance')
                elif '_Ha_' in name:
                     print("set Ha for {}".format(name))
                     fits.setval(fits_file, 'FILTER', value='Ha')
                elif '_OIII_' in name:
                     print("set OIII for {}".format(name))
                     fits.setval(fits_file, 'FILTER', value='OIII')
                elif '_SII_' in name:
                     print("set SII for {}".format(name))
                     fits.setval(fits_file, 'FILTER', value='SII')
                else:
                    print("No changes in {}".format(name))