def fits_read(file):
    hdu = fits.open(file)
    data = hdu[0].data
    hdr = hdu[0].header
    return data,hdr

def hastrom(image,hdr,refhdr):
    from reproject import reproject_interp
    hdu = fits.PrimaryHDU()
    hdu.header = hdr
    hdu.data = image
    array, footprint = reproject_interp(hdu, refhdr)
    return array

def getstats(img,ff):
    gd = np.where(img != 0)
    arr = img[gd]
    arr = sorted(arr)
    n = len(arr)
    print('there are ',n,' elements')
    i = round(ff*n)
    vmax = arr[i]
    print(ff,' signal range value is ',vmax)
    
    print('making mask')
    mask = make_source_mask(img, nsigma=2, npixels=5, dilate_size=11)
    print('calculating stats')
    vmean, vmedian, vstd = sigma_clipped_stats(img, sigma=3.0, mask=mask,mask_value=0.)
    print('mean: ',vmean)
    print('median: ',vmedian)
    print('sigma: ',vstd)
    return vmean,vmedian,vstd,vmax

def mkcol(b,v,r,ff,gamma):

    # get statitics from the following section of the image(s)
    xlo = 100
    xhi = 900
    ylo = 100
    yhi = 800

    #print(b.shape)
    #print(v.shape)
    #print(r.shape)

    # get statistics
    bmean,bmedian,bstd,bmax = getstats(b[ylo:yhi,xlo:xhi],ff)
    vmean,vmedian,vstd,vmax = getstats(v[ylo:yhi,xlo:xhi],ff)
    rmean,rmedian,rstd,rmax = getstats(r[ylo:yhi,xlo:xhi],ff)
    print('stats done')

    # set black=mean
    bmin = bmean
    vmin = vmean
    rmin = rmean

    # renromalise
    b = (b-bmin)/(bmax-bmin)
    v = (v-vmin)/(vmax-vmin)
    r = (r-rmin)/(rmax-rmin)

    #print('setting limits')
    lo = 0.
    hi = 1.

    bad = np.where(b <= lo)
    b[bad] = 0.
    bad = np.where(b >= hi)
    b[bad] = 1.

    bad = np.where(v <= lo)
    v[bad] = 0.
    bad = np.where(v >= hi)
    v[bad] = 1.

    bad = np.where(r <= lo)
    r[bad] = 0.
    bad = np.where(r >= hi)
    r[bad] = 1.

    print('applying gamma')
    b = b**gamma
    v = v**gamma
    r = r**gamma

    sz = b.shape
    print('x/y output size: ',sz[1],sz[0])

    col = np.zeros((sz[0],sz[1],3))
    col[:,:,0] = b
    col[:,:,1] = v
    col[:,:,2] = r
    print('done')
    return col

import numpy as np
from matplotlib import pyplot as plt
from astropy.io import fits
from photutils import make_source_mask
from astropy.stats import sigma_clipped_stats
from scipy import ndimage

b,bhdr = fits_read('B.fits')
v,vhdr = fits_read('V.fits')
r,rhdr = fits_read('R.fits')

v = hastrom(v,vhdr,bhdr)
r = hastrom(r,rhdr,bhdr)

# frac is the fraction of pixels to define black-to-white
# gamma is the gamma factor applied to the pixel distribution to adjust the contrast
frac = 0.995
gamma = 0.8
col = mkcol(r,v,b,frac,gamma)

plt.imshow(col,origin='lower')
plt.show()