#!/usr/local/bin/python3

"""

   Write a rgb fits images from a fits Bayer masked fits image file
   
   Input: 
     infile.fits
     outfile_base
     
   Output:
     outfile_r.fits outfile_g.fits outfile_b.fits 
     
   This version combines the Bayer r, g, and b pixels into a color superpixel
   with r, g, and b colors assigned from the original filtered pixels 
   
   RGB pixel mask is set for the protocol of an Allied Vision G319C camera 
     

"""

import os
import sys
import numpy as np
import astropy.io.fits as pyfits
from time import gmtime, strftime  # for utc

if len(sys.argv) != 3:
  print(" ")
  print("Usage: fits_to_rgb_fits.py infile.fits outfile_base")
  print(" ")
  sys.exit("Create FITS images in each color\n")
else: 
  infile = sys.argv[1]
  outfile_base = sys.argv[2]
  
# Flags

overwriteflag = True

# Open the fits file and create an hdulist

inlist = pyfits.open(infile) 

# Assign image data to a numpy array

inimage =  inlist[0].data.astype('float32')

print('Image size: (pixels) ', inimage.size)
print('Image shape: (h, w)  ', inimage.shape)

outimage = inimage

h, w = outimage.shape
oh = h//2
ow = w//2

# Pick up RGB samples to match Allied Vision Bayer camera mask
# FITS images enumerate pixels starting at [1,1] as does Julia
# Numpy images enumerate arrays starting at [0,0]

# The pattern in all Allied Vision Manta cameras is 
#
#
#           Column 1    Column 2    Column 3
#
#   Row 1      R           G           R
#   Row 2      G           B           G
#   Row 3      R           G           R
#
# When x is the column number (first index) and y is the row number (second index)
#  in a 1-based enumeration (e.g. Julia)
#
#    FITS B  [x, y] is [even, even]
#    FITS R  [x, y] is [odd,   odd]
#    FITS G  [x, y] is both [even, odd] and [odd, even]
#
#  and in a 0-based enumeration even/odd are swapped.
#
# Most Sony CMOS sensors are rectangular and wider than high (w>h)
# In Numpy arrays the row number is the first index
#   and the column number is the second, i.e. [y,x]

# This algorithm assigns the color to its pixel and does not interpolate
# The resulting images are true to color with 1-pixel spatial dithering

r_image  = outimage[0::2, 0::2]     # [odd, odd]   -> rows 0,2,4 columns 0,2,4
b_image  = outimage[1::2, 1::2]     # [even, even] -> rows 1,3,5 columns 1,3,5
g0_image = outimage[0::2, 1::2]     # [odd, even]  -> rows 0,2,4 columns 1,3,5
g1_image = outimage[1::2, 0::2]     # [even, odd]  -> rows 1,3,5 columns 0,2,4

# Trim to size and average the g images

r_image = r_image[:oh,:ow]
b_image = b_image[:oh,:ow]
g_image = g0_image[:oh,:ow]//2 + g1_image[:oh,:ow]//2

r_outfile = outfile_base+"_r.fits"
g_outfile = outfile_base+"_g.fits"
b_outfile = outfile_base+"_b.fits"


# Save the images as FITS

file_time = strftime("%Y-%m-%d %H:%M:%S", gmtime())

outlist = pyfits.PrimaryHDU(r_image)
outhdr = outlist.header
outhdr['DATE'] = file_time
outhdr['history'] = 'Red image generated by fits_to_rgb_fits' 
outhdr['history'] = 'Image file '+  infile
outlist.writeto(r_outfile, overwrite = overwriteflag)


outlist = pyfits.PrimaryHDU(g_image)
outhdr = outlist.header
outhdr['DATE'] = file_time
outhdr['history'] = 'Green image generated by fits_to_rgb_fits' 
outhdr['history'] = 'Image file '+  infile
outlist.writeto(g_outfile, overwrite = overwriteflag)


outlist = pyfits.PrimaryHDU(b_image)
outhdr = outlist.header
outhdr['DATE'] = file_time
outhdr['history'] = 'Blue image generated by fits_to_rgb_fits' 
outhdr['history'] = 'Image file '+  infile
outlist.writeto(b_outfile, overwrite = overwriteflag)



# Close the input image file

inlist.close()

exit()