%matplotlib inline
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import CombineCCFs
import numpy as np
from astropy import units as u, constants
from HelperFunctions import Gauss, integral
import os
import lmfit
import emcee
import triangle

from scipy.interpolate import InterpolatedUnivariateSpline as spline

sns.set_context('paper', font_scale=2.0)

home = os.environ['HOME']


"""
Explanation: Run the CCF analysis
You should have already cross-correlated all your spectra against a suitable model using Search.py. This notebook goes through the shifting and subtracting of the average CCF, and measuring the companion velocities.
End of explanation
"""


hdf_file = '{}/School/Research/McDonaldData/PlanetData/PsiDraA/Cross_correlations/CCF.hdf5'.format(home)
output_dir = '{}/School/Research/McDonaldData/PlanetData/Paper/Figures/'.format(home)

T = 4400
vsini = 5
logg = 4.5
metal = 0.0
dV = 0.1
c = constants.c.cgs.to(u.m/u.s).value
xgrid = np.arange(-400, 400+dV/2., dV)
ccfs, original_files = CombineCCFs.get_ccfs(T=T, vsini=vsini, logg=logg, metal=metal,
                                        hdf_file=hdf_file, xgrid=xgrid)

# Plot all the CCFs
cmap = sns.cubehelix_palette(reverse=False, as_cmap=True, gamma=1, rot=0.7, start=2)
fig, ax = plt.subplots(1, 1)
out = ax.imshow(ccfs, cmap=cmap, aspect='auto', origin='lower')#, vmin=vmin, vmax=vmax)
min_v = -75.
max_v = 75.
dv_ticks = 50.0/dV
ax.set_xlim(((min_v+400)/dV, (max_v+400)/dV))
ticks = np.arange((min_v+400)/dV, (max_v+400)/dV+1, dv_ticks)
ax.set_xticks((ticks))
#ax.set_xticklabels((-150, -100, -50, 0, 50, 100, 150))
ax.set_xticklabels((-75, -25, 25, 75))
ax.set_xlabel('Velocity in primary star rest frame (km/s)')
ax.set_yticklabels(())
ax.set_ylabel('Observation Date')

# Colorbar
cb = plt.colorbar(out)
cb.set_label('CCF Power')

# Save
plt.savefig('{}Original_CCFs.pdf'.format(output_dir))

avg_ccf = np.mean(ccfs, axis=0)
normed_ccfs = ccfs - avg_ccf

# Get the time each observation was made (in julian date)
dates = np.array([CombineCCFs.fits.getheader(fname)['HJD'] for fname in original_files])

# Set up the scaling manually
low, high = np.min(normed_ccfs), np.max(normed_ccfs)
rng = max(abs(low), abs(high))
vmin = np.sign(low) * rng
vmax = np.sign(high) * rng

# Make the actual plot
#cmap = sns.cubehelix_palette(reverse=False, as_cmap=True, gamma=1, rot=0.7, start=2) #defined above now...
fig, ax = plt.subplots(1, 1)
out = ax.imshow(normed_ccfs, cmap=cmap, aspect='auto', origin='lower')#, vmin=vmin, vmax=vmax)
ax.set_xlim(((min_v+400)/dV, (max_v+400)/dV))
ticks = np.arange((min_v+400)/dV, (max_v+400)/dV+1, dv_ticks)
ax.set_xticks((ticks))
#ax.set_xticklabels((-150, -100, -50, 0, 50, 100, 150))
ax.set_xticklabels((-75, -25, 25, 75))
ax.set_xlabel('Velocity (km/s)')
ax.set_yticklabels(())
ax.set_ylabel('Observation Date')

# Colorbar
cb = plt.colorbar(out)
cb.set_label('CCF Power')

fig.subplots_adjust(bottom=0.18, left=0.10, top=0.95, right=0.90)

plt.savefig('{}Resid_CCFs.pdf'.format(output_dir))

# Make the same plot, but with the y-axis linearized so that I can give dates.

X,Y = np.meshgrid(xgrid, dates-2450000)

fig, ax = plt.subplots(1, 1, figsize=(6,4))
out = ax.pcolormesh(X,Y,normed_ccfs, cmap=cmap, rasterized=True)
ax.set_xlabel('Velocity (km/s)')
ax.set_ylabel('JD - 2450000')
ax.set_xlim((-75, 75))
ax.set_ylim((Y.min(), Y.max()))


# Colorbar
cb = plt.colorbar(out)
cb.set_label('CCF Power')

fig.subplots_adjust(bottom=0.18, left=0.15, top=0.95, right=0.90)
plt.savefig('{}Resid_CCFs.pdf'.format(output_dir))
print(output_dir)

ax.pcolormesh?

def fwhm(x, y, search_range=(-500, 500)):
    good = (x > search_range[0]) & (x < search_range[1])
    x = x[good].copy()
    y = y[good].copy()
    idx = np.argmax(y)
    ymax = y[idx]
    half_max = ymax / 2.0
    
    # Find the first pixels less than half_max to the left of idx
    for di in range(1, idx):
        if y[idx-di] < half_max:
            break
    slope = (y[idx-(di+1)] - y[idx-di])/(x[idx-(di+1)] - x[idx-di])
    left = x[idx-di] + (half_max - y[idx-di])/slope
    
    # Find the first pixels less than half_max to the right of idx
    for di in range(1, len(x)-idx-1):
        if y[idx+di] < half_max:
            break
    slope = (y[idx+(di+1)] - y[idx+di])/(x[idx+(di+1)] - x[idx+di])
    right = x[idx+di] + (half_max - y[idx+di])/slope
    
    return left, x[idx], right

# Make plot of a normal residual CCF
sns.set_style('white')
sns.set_style('ticks')

i = 50
corr = normed_ccfs[i]
fig, ax = plt.subplots()
ax.plot(xgrid, corr, 'k-', lw=3)

l, m, h = fwhm(xgrid, corr, search_range=(-50, 10))
ylim = ax.get_ylim()
ax.plot([m, m], ylim, 'g--', alpha=0.7)
ax.plot([l, l], ylim, 'r--', alpha=0.7)
ax.plot([h, h], ylim, 'r--', alpha=0.7)

ax.set_xlim((-60, 20))
ax.set_ylim(ylim)

ax.set_xlabel('Velocity (km/s)')
ax.set_ylabel('CCF Power')
fig.subplots_adjust(bottom=0.18, left=0.20, top=0.95, right=0.95)

plt.savefig('{}Typical_CCF.pdf'.format(output_dir))

"""
Explanation: Get and shift the Cross-correlation functions to the primary star rest frame
Run with T > 6000 first, to measure the primary star radial velocity in my CCFs. Then, run with T = 4000 or so to detect the companion.
End of explanation
"""


# Measure the radial velocities of the companion as the peak and FWHM of the residual CCF

date = []
rv1 = []
rv1_err = []
rv2 = []
rv2_err = []
import time
for i in range(10, len(original_files)):
#for i in range(20, 21):
    header = CombineCCFs.fits.getheader(original_files[i])
    jd = header['HJD']
    prim_rv = CombineCCFs.get_prim_rv(original_files[i], data_shift=0.0)
    measurements = CombineCCFs.get_measured_rv(original_files[i])
    rv1.append(measurements[0])
    rv1_err.append(measurements[1])
    
    try:
        l, m, h = fwhm(xgrid.copy(), normed_ccfs[i], search_range=(-50, 10))
        print('i = {}, HJD = {}\n\t{:.1f} +{:.1f}/-{:.1f}\n\t{}'.format(i, jd, m, h-m, m-l, original_files[i]))
        date.append(jd)
        rv2.append((h+l)/2.)
        rv2_err.append((h-l)/2.355)
    except:
        date.append(jd)
        rv2.append(np.nan)
        rv2_err.append(np.nan)
        continue
        
    plt.plot(xgrid, normed_ccfs[i])
    plt.xlim((-100, 50))
    ylim = plt.ylim()
    plt.plot([(h+l)/2., (h+l)/2.], ylim, 'g--')
    plt.savefig('Figures/CCF_{}.pdf'.format(original_files[i][:-5]))
    plt.cla()
    
    
rv1 = np.array(rv1)
rv1_err = np.array(rv1_err)
rv2 = np.array(rv2)
rv2_err = np.array(rv2_err)

# I don't trust the very early measurements
rv2[:6] = np.nan*np.ones(6)
rv2_err[:6] = np.nan*np.ones(6)

# Save the RVs
np.savetxt('rv_data.txt', (date, rv1, rv1_err, rv2, rv2_err))

"""
Explanation: Measure the companion RVs.
The measured CCF velocities are in the primary star rest frame + some $\Delta v$ caused by the mismatch between the modeled index of refraction and the real index of refraction at McDonald observatory. That is a constant though. So, we get
$v_m(t) = v_2(t) - v_1(t) + \Delta v$
and
$v_2(t) = v_m(t) + v_1(t) - \Delta v$
where $v_m(t)$ are the measured radial velocities in the residual CCFS. But, all that get's done in the MCMC fitting code. Let's just measure $v_m$ for as many times as possible.
End of explanation
"""


import pandas as pd
rv1_data = pd.read_fwf('../Planet-Finder/data/psi1draa_140p_28_37_ASW.dat', header=None)
t1 = rv1_data[0].values
rv1 = rv1_data[2].values / 1000. # Convert from m/s to km/s
rv1_err = rv1_data[3].values / 1000.

new_rv2 = np.empty_like(t1)
new_rv2_err = np.empty_like(t1)
for i, t1_i in enumerate(t1):
    #idx = np.searchsorted(t1, t2)
    idx = np.argmin(np.abs(date-t1_i))
    if abs(t1_i - date[idx]) < 0.001:
        print i, t1_i, date[idx], t1_i-date[idx]
        new_rv2[i] = rv2[idx]
        new_rv2_err[i] = rv2_err[idx]
    else:
        print i, t1_i, date[idx], np.nan
        new_rv2[i] = np.nan
        new_rv2_err[i] = np.nan

np.savetxt('../Planet-Finder/data/rv_data.txt', (t1, rv1, rv1_err, new_rv2, new_rv2_err))

"""
Explanation: Fix the dates to line up with the RV output file
End of explanation
"""


