from __future__ import division, print_function

%matplotlib inline
%config InlineBackend.figure_format = "retina"

from matplotlib import rcParams
rcParams["savefig.dpi"] = 100
rcParams["font.size"] = 20

import h5py
import fitsio
import numpy as np
import matplotlib.pyplot as pl
from scipy.linalg import cho_solve, cho_factor

from astropy.io import fits
from astropy.wcs import WCS

tpf = fitsio.read("ktwo201466124-c01_lpd-targ.fits.gz")

img = tpf["FLUX"][-1]
pl.imshow(img.T, cmap="gray", interpolation="nearest");

# These are some useful things to pre-compute and use later.
_x, _y = np.meshgrid(range(-1, 2), range(-1, 2), indexing="ij")
_x, _y = _x.flatten(), _y.flatten()
_AT = np.vstack((_x*_x, _y*_y, _x*_y, _x, _y, np.ones_like(_x)))
_ATA = np.dot(_AT, _AT.T)
factor = cho_factor(_ATA, overwrite_a=True)

# This function finds the centroid and second derivatives in a 3x3 patch.
def fit_3x3(img):
    a, b, c, d, e, f = cho_solve(factor, np.dot(_AT, img.flatten()))
    m = 1. / (4 * a * b - c*c)
    x = (c * e - 2 * b * d) * m
    y = (c * d - 2 * a * e) * m
    dx2, dy2, dxdy = 2 * a, 2 * b, c
    return [x, y, dx2, dy2, dxdy]

# This function finds the centroid in an image.
# You can provide an estimate of the centroid using WCS.
def find_centroid(img, init=None):
    if init is None:
        xi, yi = np.unravel_index(np.argmax(img), img.shape)
    else:
        xi, yi = map(int, map(np.round, init))
        ox, oy = np.unravel_index(np.argmax(img[xi-1:xi+2, yi-1:yi+2]), (3, 3))
        xi += ox - 1
        yi += oy - 1
    assert (xi >= 1 and xi < img.shape[0] - 1), "effed, x"
    assert (yi >= 1 and yi < img.shape[1] - 1), "effed, y"
    pos = fit_3x3(img[xi-1:xi+2, yi-1:yi+2])
    pos[0] += xi
    pos[1] += yi
    return pos

with fits.open("ktwo201466124-c01_lpd-targ.fits.gz") as hdus:
    hdr = hdus[2].header
    wcs = WCS(hdr)
    # The order is the opposite of what I normally use...
    init = wcs.wcs_world2pix(hdr["RA_OBJ"], hdr["DEC_OBJ"], 0.0)[::-1]

img = tpf["FLUX"][-1]
pl.imshow(img.T, cmap="gray", interpolation="nearest")
pl.plot(init[0], init[1], "or");

cx, cy = find_centroid(img, init)[:2]

img = tpf["FLUX"][-1]
pl.imshow(img.T, cmap="gray", interpolation="nearest")
pl.plot(init[0], init[1], "or")
pl.plot(cx, cy, ".b");

coords = np.nan + np.zeros((len(tpf), 5))
for i, img in enumerate(tpf["FLUX"]):
    msk = np.isfinite(img)
    if np.any(msk) and np.any(img[msk] > 0.0):
        coords[i] = find_centroid(img, init=init)

print(coords.shape)

data = fitsio.read("ktwo201466124-c01_lpd-lc.fits")
aps = fitsio.read("ktwo201466124-c01_lpd-lc.fits", 2)

# Load the data.
time = data["time"]
flux = data["flux"][:, np.argmin(aps["cdpp6"])]  # Lowest CDPP6.
qual = data["quality"]

# Mask missing/bad data.
good_mask = np.isfinite(time) & np.isfinite(flux) & (qual == 0)
good_time = time[good_mask]
good_flux = flux[good_mask] / np.median(flux[good_mask])

pl.plot(good_time, good_flux, "k");

pl.scatter(coords[good_mask, 0], coords[good_mask, 1], c=good_time, edgecolor="none", alpha=0.5)
pl.xlabel("x [pix]")
pl.ylabel("y [pix]");