#!/usr/bin/env python3

# -*- coding: utf-8 -*-

"""

Created on Tue Oct 3 10:12:26 2017

@author: throop

"""

import pdb

import glob

import math # We use this to get pi. Documentation says math is 'always available'

# but apparently it still must be imported.

from subprocess import call

import warnings

import pdb

import os.path

import os

import subprocess

import astropy

from astropy.io import fits

from astropy.table import Table

import astropy.table # I need the unique() function here. Why is in in table and not Table??

import matplotlib

import matplotlib.pyplot as plt # pyplot

from matplotlib.figure import Figure

import numpy as np

import astropy.modeling

from scipy.optimize import curve_fit

#from pylab import * # So I can change plot size.

# Pylab defines the 'plot' command

import spiceypy as sp

#from itertools import izip # To loop over groups in a table -- see astropy tables docs

from astropy.wcs import WCS

from astropy import units as u # Units library

from astropy.coordinates import SkyCoord # To define coordinates to use in star search

#from photutils import datasets

from scipy.stats import linregress

from astropy.visualization import wcsaxes

import time

from scipy.interpolate import griddata

from importlib import reload # So I can do reload(module)

import imreg_dft as ird # Image translation

import scipy.ndimage # For

import skimage.transform as skt # This 'resize' function is more useful than np's

from astropy.convolution import convolve, convolve_fft, Gaussian2DKernel, Box2DKernel

import re # Regexp

import pickle # For load/save

import scipy

from matplotlib.figure import Figure

# HBT imports

import hbt

# =============================================================================

# Compute brightness of MU69, in DN, based on distance to body and exptime.

# =============================================================================

def nh_dn_lorri_mu69(mode = '1X1', dist = 0.1*u.au, exptime = 10):

return(S) # Return the total counts

# =============================================================================

# Compute conversion constant from I/F to DN.

# mode = {'1X1', '4X4'}

# dist = distance to sun, in astropy units.

# exptime = exposure time in seconds

# This ignores any ring-specific geometry

# =============================================================================

def nh_iof_2_dn_lorri_extended(mode = '1X1', dist = 43*u.au, exptime = 10):

return factor

# =============================================================================

# Make a simulated image.

# For now, this image has zero DN of background readnoise, stray light, etc.

# It is just the ring and/or MU69.

#

# The returned image is in DN (not I/F).

# =============================================================================

def nh_make_simulated_image_lorri(do_ring = False, # Flag: do we draw a ring?

return arr

# =============================================================================

# Now do some tests on this

# =============================================================================

plt.set_cmap('Greys_r')

do_destripe = True # Flag: Do we run our destripe function on the raw LORRI data?

sp.furnsh('kernels_kem.tm')

hbt.figsize((12,12))

hbt.set_fontsize(15)

iof_ring = 1e-7

exptime = 30

mode = '4X4'

#mode = '1X1'

pos = (None, None)

#pos = (300, 700)

pos = (100, 200) # y, x in normal imshow() coordinates.

#dist_target = 0.01*u.au

dist_solar = 43.2*u.au # MU69 dist at encounter: 43.2 AU, from KEM Wiki page

do_psf = True # Flag: Do we convolve result with NH LORRI PSF?26.7 + 0.65

dt_obs = -22*u.day # Time relative to MU69 C/A

utc_ca = '2019 1 Jan 05:33:00'

et_ca = sp.utc2et(utc_ca)

et_obs = et_ca + dt_obs.to('s').value

utc_obs = sp.et2utc(et_obs, 'C', 0)

utc_obs_human = 'K{:+}d'.format(dt_obs.to('day').value)

vec,lt = sp.spkezr('2014 MU69', et_obs, 'J2000', 'LT', 'New Horizons')

vec_sc_targ = vec[0:3]

dist_target = np.sqrt(np.sum(vec_sc_targ**2))*u.km.to('AU')*u.au

arr = nh_make_simulated_image_lorri(do_ring=True,

do_psf = True)

# Calculate the max DN value in the array. This is (more-or-less) the ring target, converted to DN.

# Though it might be diminished a bit due to the PSF.

dn_max = np.amax(arr[arr > 0])

plt.imshow(arr)

plt.title('I/F = {}, max = {:0.2f} DN, t = {} s, mode = {}, {}, {:.2f}, {}'.format(

iof_ring, dn_max, exptime, mode,

dist_solar,

dist_target,

utc_obs_human))

plt.show()

# Do a test of the I/F -> DN conversion. This value can be compared to that from Exposure Time Calculator

iof2dn = nh_iof_2_dn_lorri_extended(mode = mode, dist = dist_solar, exptime = exptime)

print('For ring: I/F = {}, exptime = {} s -> {:0.2f} DN/pixel'.format(iof_ring, exptime, iof_ring * iof2dn))

print('For MU69: total signal = {:0.2f} DN, pre-PSF'.format(nh_dn_lorri_mu69(dist = dist_target, exptime = exptime)))

# =============================================================================

# Now that we have the routine set up:

# Load all of the KEM MU69 frames.

# Add a ring to them, at several sizes and I/F's

# Write these to disk.

# =============================================================================

# First load the pickle file that describes all of the navigation info, header info, etc.

# for Simon's version of the KEM frames.

dir_buie = '/Users/throop/Data/NH_KEM_Hazard/Buie_Sep17' # Very optimized WCS coords

dir_porter = '/Users/throop/Data/NH_KEM_Hazard/Porter_Sep17' # Basic WCS coords

dir_out = '/Users/throop/Data/NH_KEM_Hazard/synthetic/'

file_save = dir_porter + '/kem_hazard_picksar_n344.pkl'

lun = open(file_save, 'rb')

(t, data, indices) = pickle.load(lun)

print("Read: " + file_save)

lun.close()

stretch_percent = 95

stretch = astropy.visualization.PercentileInterval(stretch_percent) # PI(90) scales to 5th..95th %ile.

# Set up various indices so we can extract different image sets

# We don't use all of these -- but just here for reference if we need them.

indices_sep17 = t['et'] > sp.utc2et('15 sep 2017') # The positon of MU69 has changed a few pixels.

# We can't blindly co-add between sep and pre-sep

indices_jan17 = t['et'] < sp.utc2et('1 sep 2017')

indices_rot0 = t['angle'] < 180 # One rotation angle

indices_rot90 = t['angle'] > 180 # The other rotation angle

indices_10sec = np.logical_and( t['exptime'] < 10, t['exptime'] > 5 )

indices_20sec = np.logical_and( t['exptime'] < 20, t['exptime'] > 10 )

indices_30sec = np.logical_and( t['exptime'] < 30, t['exptime'] > 20 )

indices_1x1 = t['naxis1'] == 1024

indices_4x4 = t['naxis1'] == 256

indices_30sec_rot0 = np.logical_and(indices_30sec, indices_rot0) # 94

indices_30sec_rot90 = np.logical_and(indices_30sec, indices_rot90) # 0

# Now go through and extract all of the images we want

indices = indices_30sec_rot0

files_short = t['filename_short'][indices]

# Set up the iteration parameters. These are the values we will loop over

#iof_ring = [1e-7, 3e-7, 1e-6, 3e-6, 1e-5, 1e-4]

iof_ring = [3e-6]

a_ring = [(3000*u.km, 4000*u.km), (8000*u.km, 10000*u.km)]

a_ring_name = ['small', 'large']

#a_ring = [(3000*u.km, 4000*u.km)]

#a_ring_name = ['small']

# Loop over the files that we have downselected to

for file_short in files_short:

# Reconstruct the full filename from the short one

file = glob.glob(dir_porter + '/*/*{}*'.format(file_short))[0]

# Load the image

hdu = fits.open(file)

im_data = hdu['PRIMARY'].data.copy()

header = hdu['PRIMARY'].header

# Grab values from the FITS header

et = header['SPCSCET']

exptime = header['EXPTIME']

# Get the image size

if (header['NAXIS1'] == 256):

mode = '4X4'

else:

mode = '1X1'

# Parse the WCS

w = WCS(file)

# Get the XY position of MU69

(vec, lt) = sp.spkezr('MU69', et, 'J2000', 'LT', 'New Horizons')

vec = vec[0:3]

vec_sc_targ = vec[0:3]

(junk,ra,dec) = sp.recrad(vec_sc_targ) # Get the RA / Dec of the object

x_pix, y_pix = w.wcs_world2pix(ra*hbt.r2d, dec*hbt.r2d, 0) # Convert to pixels.

# XXX weird: world2pix returns a zero-dim array. Not an ndarray.

# Have to do this float() business to extract the values. Or float().

x_pix = float(x_pix)

y_pix = float(y_pix)

# Make the simulated image

for iof_ring_i in iof_ring:

for i,a_ring_i in enumerate(a_ring):

a_ring_name_i = a_ring_name[i]

im_simulated = nh_make_simulated_image_lorri(do_ring=True,

dist_ring_smoothing = 2000*u.km,

iof_ring = iof_ring_i,

a_ring = a_ring_i,

exptime = exptime,

mode = mode,

pos = (y_pix, x_pix), # Swap x and y here

dist_solar = dist_solar,

dist_target = dist_target, # *Not* the distance from FITS header, of course

met = file_short,

do_mu69 = False,

do_psf = True) # NB: The slight off-center ring is caused by the PSF.

# Merge the two images

im_out = im_data + im_simulated

plt.imshow(stretch(im_out))

plt.plot(x_pix, y_pix, marker = 'o', color='red', ms = 5, alpha=0.5)

plt.title('{}, t={} s, {}, {}, I/F = {:.0e}, K{}'.format(

file_short, exptime, mode, a_ring_name_i, iof_ring_i, dt_obs))

plt.show()

# Create the output filename

file_out = dir_out + (file.split('/')[-1]).replace('.fits', '') + \

'_ring_{}_iof{:.0e}_K{:+.0f}d'.format(a_ring_name_i, iof_ring_i, int(dt_obs.to('d').value)) \

+ '.fits'

# Copy the new image into the existing FITS HDU. This does not modify the original -- only the one in memory.

hdu['PRIMARY'].data = im_out

# And write a new FITS file -- same as original, but with the new simulated image array

hdu.writeto(file_out, overwrite=True)

print("Created: {}".format(file_out))

# Finally, close this input FITS file

hdu.close()