import numpy as np
import imageio
from butterworth_filter import butterworth_filter
from clean_negbgr import clean_negbgr

def lup_func(x, scale=1.5e-2, nonlinearity=10, base=0) :
   flx = scale * nonlinearity * (x + base)
   np.arcsinh(flx, out=flx) 
   flx /= nonlinearity
   return flx

def make_rgb_lup(rgb_images, scales=[1,1,1], base0=[0,0,0], out_name='rgb_ukidss'):
   import os
   c = rgb_images * scales
   I = np.zeros((c.shape[0], c.shape[1]), dtype=np.float)
   np.sum(c, out=I, axis=-1)
   I[I<1e-10] = 1e-10
   L = lup_func(I, scale=0.8e-2, nonlinearity=1.30, base=0)
   r = 255.0 * scales[0] * (rgb_images[...,0] - base0[0]) * L / I 
   r = butterworth_filter(r, omega_c=0.18*c.shape[0])
   r -= 15.0
   r[r < 0] = 0
   r += 15.0
   r[r > 255] = 255
   g = 255.0 * scales[1] * (rgb_images[...,1] - base0[1])* L / I 
   g = butterworth_filter(g, omega_c=0.18*c.shape[0])
   g -= 15.0
   g[g < 0] = 0
   g += 15.0
   g[g > 255] = 255
   b = 255.0 * scales[2] * (rgb_images[...,2] - base0[2])* L / I 
   b = butterworth_filter(b, omega_c=0.18*c.shape[0])
   b -= 15.0
   b[b < 0] = 0
   b += 15.0
   b[b > 255] = 255
   rgb_out = np.zeros(rgb_images.shape).astype(np.uint8)
   rgb_out[...,0] = r.astype(np.uint8)
   rgb_out[...,1] = g.astype(np.uint8)
   rgb_out[...,2] = b.astype(np.uint8)
   imageio.imwrite(out_name+'.jpg', np.flip(rgb_out, axis=0), quality=98)


def YHKfits2RGBjpeg(swarp_dir, out_name):
   from astropy.io import fits
   import os
   hr = fits.open(os.path.join(swarp_dir, 'K_COADD.fits'))
   hg = fits.open(os.path.join(swarp_dir, 'H_COADD.fits'))
   hb = fits.open(os.path.join(swarp_dir, 'Y_COADD.fits'))
   s = hr[0].data.shape
   #cube_rgb = np.zeros((1000,1000,3))
   cube_rgb = np.zeros((s[0],s[1],3))
   bgr_r = clean_negbgr(hr[0].data, hr[0].data*0, hr[0].data*0)
   bgr_g = clean_negbgr(hg[0].data, hg[0].data*0, hg[0].data*0)
   bgr_b = clean_negbgr(hb[0].data, hb[0].data*0, hb[0].data*0)
   cube_rgb[...,0] = hr[0].data - bgr_r #[5000:6000,3400:4400]
   cube_rgb[...,1] = hg[0].data - bgr_g #[5000:6000,3400:4400]
   cube_rgb[...,2] = hb[0].data - bgr_b #[5000:6000,3400:4400]
   hr[0].header.tofile(out_name+'.hhh', sep='\n', padding=False, overwrite=True)
   make_rgb_lup(cube_rgb, scales=[0.53,0.63,0.95], base0=[0.0, 0.0, 0.0], out_name=out_name)
   hr.close()
   hg.close()
   hb.close()


if __name__=='__main__':
   from astropy.io import fits
   import os
   swarp_dir = 'images/UKIDSS_J110731.23+134712.88/swarp/'
   YHKfits2RGBjpeg(swarp_dir, 'UKIDSS_J110731.23+134712.88')
'''
   hr = fits.open(os.path.join(swarp_dir, 'K_COADD.fits'))
   hg = fits.open(os.path.join(swarp_dir, 'H_COADD.fits'))
   hb = fits.open(os.path.join(swarp_dir, 'Y_COADD.fits'))
   s = hr[0].data.shape
   #cube_rgb = np.zeros((1000,1000,3))
   cube_rgb = np.zeros((s[0],s[1],3))
   bgr_r = clean_negbgr(hr[0].data, hr[0].data*0, hr[0].data*0)
   bgr_g = clean_negbgr(hg[0].data, hg[0].data*0, hg[0].data*0)
   bgr_b = clean_negbgr(hb[0].data, hb[0].data*0, hb[0].data*0)
   cube_rgb[...,0] = hr[0].data - bgr_r #[5000:6000,3400:4400]
   cube_rgb[...,1] = hg[0].data - bgr_g #[5000:6000,3400:4400]
   cube_rgb[...,2] = hb[0].data - bgr_b #[5000:6000,3400:4400]
   make_rgb_lup(cube_rgb, scales=[0.53,0.63,0.95], base0=[0.0, 0.0, 0.0])
   hr.close()
   hg.close()
   hb.close()
'''