## Observation correction model

import numpy as np

import astropy.time
import astropy.coordinates
import astropy.units as u
import astropy.units.imperial
import astropy.units.cds

import bossdata
import specsim

import tpcorr.pointing
import tpcorr.guider
import tpcorr.acceptance_model

class Observation(object):
    def __init__(self, plate, mjd, guide_wlen=5400 * u.Angstrom, offset_wlen=4000 * u.Angstrom,
        std_wlen=5400 * u.Angstrom, wlen_grid_steps=15, steps_per_exposure=5,
        pressure0=None, temperature0=None):
        print 'Calculating corrections for {} observed on MJD {}'.format(plate, mjd)

        self.plate = plate
        self.mjd = mjd
        self.guide_wlen = guide_wlen
        self.std_wlen = std_wlen
        self.offset_wlen = offset_wlen
        self.steps_per_exposure = steps_per_exposure
        self.wlen_grid_steps = wlen_grid_steps

        self.finder = bossdata.path.Finder()
        self.mirror = bossdata.remote.Manager()

        # Get the list of exposures used in this observation's coadd from a spec lite file.
        spec_name = self.finder.get_spec_path(plate, mjd, fiber=1, lite=True)
        self.spec_file = bossdata.spec.SpecFile(self.mirror.get(spec_name))

        # Read the first b1 raw science exposure to find this plate's plug map.
        raw = self.spec_file.get_raw_image(0, 'blue', finder=self.finder, mirror=self.mirror)
        plug_map = raw.read_plug_map()

        # Look up the plate design pointing from the raw header and convert to
        # an index A,B,C,... -> 0,1,2,...
        pointing_label = raw.header['POINTING'].strip()
        pointing_index = ord(pointing_label) - ord('A')

        # Initialize a pointing object for this plate's sky location.
        ra_center = float(plug_map['raCen']) * u.deg
        self.ra_center = ra_center
        dec_center = float(plug_map['decCen']) * u.deg
        print 'Plate center is RA={:.3f}, DEC={:.3f} for {}-{}'.format(ra_center, dec_center, plate, pointing_label)
        self.pointing = tpcorr.pointing.Pointing(ra_center, dec_center)

        # Find the nominal observing temperature and time that this plate's holes are drilled for.
        self.design_temp = float(plug_map['temp'])*u.deg_C
        self.design_pressure = None # Calculate based on elevation and temperature
        self.design_ha = float(plug_map['ha'].split()[pointing_index]) * u.deg
        midnight = astropy.time.Time(mjd, format='mjd', scale='tai', location=self.pointing.where)
        design_time = specsim.transform.adjust_time_to_hour_angle(midnight, ra_center, self.design_ha, max_iterations=100)
        self.design_tai = design_time.mjd * 86400.
        print 'Holes drilled for T={:.1f} and HA={:.1f} (TAI={:.1f})'.format(self.design_temp, self.design_ha, self.design_tai)

        # design_time.mjd
        # when = astropy.time.Time(tai/86400., format='mjd', scale='tai', location=self.where)
        self.design_alt = self.pointing.plate_center.transform_to(astropy.coordinates.AltAz(
            obstime=design_time, location=self.pointing.where)).alt.to(u.deg)

        # Find this plate's guide stars.
        plugging = plug_map['PLUGMAPOBJ']
        guide_fibers = plugging['holeType'] == 'GUIDE'

        guide_ra, guide_dec = plugging['ra'][guide_fibers], plugging['dec'][guide_fibers]
        self.guide_targets = astropy.coordinates.ICRS(guide_ra * u.deg, guide_dec * u.deg)

        # Calculate the nominal guide fiber positions.
        self.guide_x0, self.guide_y0, _, _ = self.pointing.transform(
            self.guide_targets, self.design_tai, guide_wlen, self.design_temp, self.design_pressure)

        # Find this plate's offset fibers. We have to use spAll for this since the plug map does
        # not record the design wavelengths.
        self.plugmap = self.get_plugmap_from_spframes()

        offset_fibers_mask = self.plugmap['LAMBDA_EFF'] == offset_wlen.to(u.Angstrom).value
        offset_fibers = self.plugmap[offset_fibers_mask]

        offset_xfocal = offset_fibers['XFOCAL'] * u.mm
        offset_yfocal = offset_fibers['YFOCAL'] * u.mm
        self.fiber_ids = offset_fibers['FIBERID']
        self.offset_targets = astropy.coordinates.ICRS(ra=offset_fibers['RA'] * u.deg, dec=offset_fibers['DEC'] * u.deg)
        self.num_offset_targets = np.count_nonzero(self.offset_targets)
        print 'Plate has {:d} guide fibers and {:d} offset targets.'.format(len(self.guide_targets), self.num_offset_targets)

        if self.num_offset_targets > 0:
            # Calculate the nominal science fiber positions. These will not match XFOCAL, YFOCAL
            # exactly since we do not exactly replicate the IDL transforms, but they should be
            # close (within ~0.2 arcsec) and we only use offsets calculated consistently with
            # transform() in the following.
            self.offset_x0, self.offset_y0, offset_alt, offset_az = self.pointing.transform(
                self.offset_targets, self.design_tai, offset_wlen, self.design_temp, self.design_pressure)

            # Calculate where the offset target fibers would have been positioned if they were
            # designed for the same wavelength as the standard stars.
            self.offset_x0_std, self.offset_y0_std, _, _ = self.pointing.transform(
                self.offset_targets, self.design_tai, std_wlen, self.design_temp, self.design_pressure)

        # Initialize the wavelength grid to use for calculating corrections.
        self.wlen_grid = np.linspace(3500., 10500., wlen_grid_steps)[:, np.newaxis] * u.Angstrom

        # Initialize guided target centroid list
        self.guided_centroids = []

        # Initialize exposure meta data lists
        self.seeing = np.empty((self.spec_file.num_exposures)) * u.arcsec
        self.ha = np.empty((self.spec_file.num_exposures)) * u.degree
        self.pressure = np.empty((self.spec_file.num_exposures)) * u.kPa
        self.temperature = np.empty((self.spec_file.num_exposures)) * u.deg_C
        self.tai_beg = np.empty((self.spec_file.num_exposures)) # seconds
        self.tai_end = np.empty((self.spec_file.num_exposures)) # seconds
        self.alt = np.empty((self.spec_file.num_exposures)) * u.degree

        self.init_exposure_meta(pressure=pressure0, temperature=temperature0)

    def init_exposure_meta(self, seeing=None, pressure=None, temperature=None):
        # Precompute the conversion from inches of Hg to kPa.
        pconv = (1 * u.cds.mmHg * u.imperial.inch / u.mm).to(u.kPa).value

        # Loop over exposures
        for exp_index in range(self.spec_file.num_exposures):

            # Open the b1 frame for this exposure, to access its metadata.
            b1_frame_name = self.finder.get_plate_path(
                self.plate, self.spec_file.get_exposure_name(exp_index, 'blue', 'spFrame'))
            b1_frame = bossdata.plate.FrameFile(self.mirror.get(b1_frame_name))
            exp_id = b1_frame.exposure_id

            # Lookup this exposure's observing time, seeing, and temperature.
            self.tai_beg[exp_index] = b1_frame.header['TAI-BEG']
            self.tai_end[exp_index] = b1_frame.header['TAI-END']
            tai_mid = 0.5 * (self.tai_beg[exp_index] + self.tai_end[exp_index])

            # Convert tai to hour angle
            self.ha[exp_index] = tpcorr.pointing.normalize_angle(
                self.pointing.hour_angle(tai_mid).to(u.deg).value)*u.deg

            if seeing is not None:
                self.seeing[exp_index] = seeing
            else:
                if b1_frame.header['SEEING50'] == 0:
                    self.seeing[exp_index] = 1.49 * u.arcsec
                    print 'Warning: SEEING50=0. Using nominal value: ', self.seeing[exp_index]
                else:
                    self.seeing[exp_index] = b1_frame.header['SEEING50'] * u.arcsec

            if temperature is not None:
                self.temperature[exp_index] = temperature
            else:
                try:
                    self.temperature[exp_index] = b1_frame.header['AIRTEMP'] * u.deg_C
                except ValueError, e:
                    self.temperature[exp_index] = 5 * u.deg_C
                    print 'Warning: AIRTEMP not available in exp header. Using nominal value: ', self.temperature[exp_index]

            if pressure is not None:
                self.pressure[exp_index] = pressure
            else:
                try:
                    self.pressure[exp_index] = b1_frame.header['PRESSURE'] * pconv * u.kPa
                except ValueError, e:
                    self.pressure[exp_index] = 71.890 * u.kPa
                    print 'Warning: PRESSURE not available in exp header. Using nominal value: ', self.pressure[exp_index]

            obstime = astropy.time.Time(tai_mid/86400., format='mjd', scale='tai', location=self.pointing.where)
            self.alt[exp_index] = self.pointing.plate_center.transform_to(astropy.coordinates.AltAz(
                obstime=obstime, location=self.pointing.where)).alt.to(u.deg)

            print 'Exp[{:02d}] #{:08d} seeing {:.3f}, T={:+5.1f}, P={:.1f}, TAI {:.1f} ({:+7.3f} days, HA {:+.1f})'.format(
                exp_index, exp_id, self.seeing[exp_index], self.temperature[exp_index], self.pressure[exp_index],
                tai_mid, (tai_mid - self.design_tai)/86400., self.ha[exp_index])

    def get_plugmap_from_spframes(self):
        # Read frame files for both spectrographs
        frames = []
        for fiber in (1,501):
            spec_name = self.finder.get_spec_path(self.plate, self.mjd, fiber=fiber, lite=True)
            spec_file = bossdata.spec.SpecFile(self.mirror.get(spec_name))
            frame_name = self.finder.get_plate_path(self.plate, spec_file.get_exposure_name(0, 'blue', 'spFrame'))
            frames.append(bossdata.plate.FrameFile(self.mirror.get(frame_name)))
        # Stack frame plugmaps
        return astropy.table.vstack([frame.plug_map for frame in frames])

    def get_exp_centroids(self, exp_index, guide_plot_name=None):
        # Create time steps covering this exposure.
        tai_steps = np.linspace(self.tai_beg[exp_index], self.tai_end[exp_index], self.steps_per_exposure)

        # Calculate the actual guide target positions on the focal plane without any guiding.
        guide_x, guide_y, _, _ = self.pointing.transform(
            self.guide_targets[:, np.newaxis], tai_steps, self.guide_wlen, self.temperature[exp_index], self.pressure[exp_index])

        # Solve for the optimal guider corrections.
        guider = tpcorr.guider.Guider(self.guide_x0, self.guide_y0, guide_x, guide_y)
        if guide_plot_name:
            guider.plot(tai_steps, field_radius=340 * u.mm, zoom=5000.,
                fiber_radius=0.1 * u.arcsec * self.pointing.platescale, save=guide_plot_name)

        # Calculate the offset target paths on the focal plane without any guiding, for the actual observing conditions.
        offset_x, offset_y, _, _ = self.pointing.transform(
            self.offset_targets[:, np.newaxis, np.newaxis], tai_steps, self.wlen_grid, self.temperature[exp_index], self.pressure[exp_index])

        # Apply guiding corrections to estimate the actual offset target paths during the exposure.
        return guider.correct(offset_x, offset_y)

    def get_mean_exp_centroids(self, extrap_wlen=False):
        # Create time steps covering this exposure.
        midnight = astropy.time.Time(self.mjd, format='mjd', scale='tai', location=self.pointing.where)
        ha = np.mean(self.ha)
        time = specsim.transform.adjust_time_to_hour_angle(midnight, self.ra_center, ha, max_iterations=100)
        tai = time.mjd * 86400.

        temperature = np.mean(self.temperature)
        pressure = np.mean(self.pressure)

        print 'Mean Exposure: seeing {:.3f}, T={:+5.1f}, P={:.1f}, TAI {:.1f} ({:+7.3f} days, HA {:+.1f})'.format(
                np.mean(self.seeing), temperature, pressure, tai, (tai - self.design_tai)/86400., ha)

        # Calculate the actual guide target positions on the focal plane without any guiding.
        guide_x, guide_y, _, _ = self.pointing.transform(
            self.guide_targets[:, np.newaxis], tai, self.guide_wlen, self.design_temp, self.design_pressure)

        # Solve for the optimal guider corrections.
        guider = tpcorr.guider.Guider(self.guide_x0, self.guide_y0, guide_x, guide_y)

        # Calculate the offset target paths on the focal plane without any guiding, for the actual observing conditions.
        offset_x, offset_y, _, _ = self.pointing.transform(
            self.offset_targets[:, np.newaxis, np.newaxis], tai, self.wlen_grid, self.design_temp, self.design_pressure, extrap_wlen=extrap_wlen)

        return guider.correct(offset_x, offset_y)

    def get_mean_correction(self, extrap_wlen=False):

        corrections = np.empty((self.num_offset_targets, self.wlen_grid_steps, 1))

        # Estimate the actual offset target paths during the exposure
        # (offset_x0, offset_y0), (offset_x0_std, offset_y0_std), (guided_x, guided_y) = self.get_mean_exp_centroids()
        guided_x, guided_y = self.get_mean_exp_centroids(extrap_wlen=extrap_wlen)

        # Calculate centroid offsets for each offset target, relative to its nominal fiber center.
        offset = np.sqrt(
            (guided_x - self.offset_x0[:, np.newaxis, np.newaxis])**2 +
            (guided_y - self.offset_y0[:, np.newaxis, np.newaxis])**2)

        # Calculate centroid offsets for each offset target, relative to where its fiber center would
        # be if it were designed for the same wavelength as the standard stars.
        offset_std = np.sqrt(
            (guided_x - self.offset_x0_std[:, np.newaxis, np.newaxis])**2 +
            (guided_y - self.offset_y0_std[:, np.newaxis, np.newaxis])**2)

        seeing = np.mean(self.seeing)

        max_offset = 1.1/2.0*max(np.max((offset / self.pointing.platescale).to(u.arcsec)).value,
                                 np.max((offset_std / self.pointing.platescale).to(u.arcsec)).value)

        acceptance_model = tpcorr.acceptance_model.AcceptanceModel(seeing, max_offset=max_offset)

        # Calculate the acceptance fractions for both sets of centroid offsets.
        acceptance = acceptance_model((offset / self.pointing.platescale).to(u.arcsec))
        acceptance_std = acceptance_model((offset_std / self.pointing.platescale).to(u.arcsec))

        # Calculate the acceptance fraction ratios, tabulated for each offset target, wavelength and time.
        # The ratio calculated this way gives the correction of eqn (13).
        corrections = acceptance_std / acceptance

        mean_correction = np.mean(corrections, axis=-1)

        return mean_correction, guided_x, guided_y

    def get_corrections(self, seeing_wlen=5400.*u.Angstrom):

        # Precompute wlen ratio for wavelength dependent seeing adjustment
        wlen_ratio = (self.wlen_grid / seeing_wlen).si

        # Initialize acceptance ratio grid
        corrections = np.empty(
            (self.spec_file.num_exposures, self.num_offset_targets, self.wlen_grid_steps, self.steps_per_exposure),
            dtype=float)

        guided_centroids = []

        # Loop over exposures
        for exp_index in range(self.spec_file.num_exposures):

            # Estimate the actual offset target paths during the exposure
            guided_x, guided_y = self.get_exp_centroids(exp_index)
            guided_centroids.append((guided_x, guided_y))

            # Calculate centroid offsets for each offset target, relative to its nominal fiber center.
            offset = np.sqrt(
                (guided_x - self.offset_x0[:, np.newaxis, np.newaxis])**2 +
                (guided_y - self.offset_y0[:, np.newaxis, np.newaxis])**2)

            # Calculate centroid offsets for each offset target, relative to where its fiber center would
            # be if it were designed for the same wavelength as the standard stars.
            offset_std = np.sqrt(
                (guided_x - self.offset_x0_std[:, np.newaxis, np.newaxis])**2 +
                (guided_y - self.offset_y0_std[:, np.newaxis, np.newaxis])**2)

            seeing = self.seeing[exp_index]
            # psf = sdss_25m.get_atmospheric_psf(seeing_wlen, seeing, gauss=False)
            # acceptance_model = sdss_25m.calculate_fiber_acceptance(psf)

            seeing_wlen_adjusted = seeing*wlen_ratio**(-0.2)
            # acceptance_model_grid = map(tpcorr.acceptance_model.AcceptanceModel, seeing_wlen_adjusted)
            max_offset = 1.1/2.0*max(np.max((offset / self.pointing.platescale).to(u.arcsec)).value,
                    np.max((offset_std / self.pointing.platescale).to(u.arcsec)).value)

            for wlen_index in range(self.wlen_grid_steps):
                # Build acceptance model for this wavelength
                acceptance_model = tpcorr.acceptance_model.AcceptanceModel(seeing_wlen_adjusted[wlen_index], max_offset=max_offset)

                # Calculate the acceptance fractions for both sets of centroid offsets.
                acceptance = acceptance_model((offset[:,wlen_index,:] / self.pointing.platescale).to(u.arcsec))
                acceptance_std = acceptance_model((offset_std[:,wlen_index,:] / self.pointing.platescale).to(u.arcsec))

                # Calculate the acceptance fraction ratios, tabulated for each offset target, wavelength and time.
                # The ratio calculated this way gives the correction of eqn (13).
                corrections[exp_index,:,wlen_index,:] = acceptance_std / acceptance

        # Average the correction over each exposure time slice.
        avg_corrections = np.mean(np.mean(corrections, axis=-1), axis=0)

        return corrections, avg_corrections, guided_centroids

if __name__ == '__main__':
    pass


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Feb  2 16:06:09 2019

@author:
Maximilian N. Günther
MIT Kavli Institute for Astrophysics and Space Research,
Massachusetts Institute of Technology,
77 Massachusetts Avenue,
Cambridge, MA 02109,
USA
Email: maxgue@mit.edu
Web: www.mnguenther.com
"""

from __future__ import print_function, division, absolute_import

#::: plotting settings
import seaborn as sns
sns.set(context='paper', style='ticks', palette='deep', font='sans-serif', font_scale=1.5, color_codes=True)
sns.set_style({"xtick.direction": "in","ytick.direction": "in"})
sns.set_context(rc={'lines.markeredgewidth': 1})

#::: modules
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
from matplotlib.patches import Circle
import matplotlib.ticker as plticker
from matplotlib.collections import LineCollection
from matplotlib.colors import LinearSegmentedColormap
import random
import warnings
try:
    import rebound
    from rebound.particle import Particle
except:
    warnings.warn('Module "rebound" could not be imported. Orbital plots are not available.')
from itertools import cycle

#::: allesfitter codes
from allesfitter import config
from allesfitter.exoworlds_rdx.lightcurves.index_transits import get_first_epoch
from allesfitter.exoworlds_rdx.lightcurves.lightcurve_tools import calc_phase

def OrbitPlot(sim, figsize=None, lim=None, limz=None, Narc=100, xlabel='x', ylabel='y', zlabel='z', color=False, periastron=False, trails=True, show_orbit=True, lw=1., glow=False, slices=False, plotparticles=[], primary=None, fancy=False, ax=None):
    """
    Convenience function for plotting instantaneous orbits.

    Parameters
    ----------
    slices          : bool, optional
        Plot all three slices if set to True. Default is False and plots orbits only in the xy plane.
    figsize         : tuple of float, optional
        Tuple defining the figure size (default: (5,5))
    lim             : float, optional
        Limit for axes (default: None = automatically determined)
    limz            : float, optional
        Limit for z axis, only used if slices=True (default: None = automatically determined)
    unitlabel       : str, optional
        String describing the units, shown on axis labels (default: None)
    color           : bool, str or list, optional
        By default plots in black. If set to True, plots using REBOUND color cycle. If a string or list of strings, e.g. ['red', 'cyan'], will cycle between passed colors.
    periastron  : bool, optional
        Draw a marker at periastron (default: False)
    trails          : bool, optional
        Draw trails instead of solid lines (default: False)
    show_orbit      : bool, optional
        Draw orbit trails/lines (default: True)
    lw              : float, optional
        Linewidth (default: 1.)
    glow            : bool (default: False)
        Make lines glow
    fancy           : bool (default: False)
        Changes various settings to create a fancy looking plot
    plotparticles   : list, optional
        List of particles to plot. Can be a list of any valid keys for accessing sim.particles, i.e., integer indices or hashes (default: plot all particles)
    primary         : rebound.Particle, optional
        Pimrary to use for the osculating orbit (default: Jacobi center of mass)

    Returns
    -------
    fig
        A matplotlib figure

    Examples
    --------
    The following example illustrates a typical use case.

    >>> sim = rebound.Simulation()
    >>> sim.add(m=1)
    >>> sim.add(a=1)
    >>> fig = rebound.OrbitPlot(sim)
    >>> fig.savefig("image.png") # save figure to file
    >>> fig.show() # show figure on screen

    """
    if slices:
        if figsize is None:
            figsize = (8,8)
        if ax is None:
            fig, ax = plt.subplots(2, 2, figsize=figsize)
        gs = gridspec.GridSpec(2, 2, width_ratios=[3., 2.], height_ratios=[2.,3.],wspace=0., hspace=0.)
        OrbitPlotOneSlice(sim, plt.subplot(gs[2]), lim=lim, Narc=Narc, color=color, periastron=periastron, trails=trails, show_orbit=show_orbit, lw=lw, axes="xy",fancy=fancy, plotparticles=plotparticles, primary=primary, glow=glow)
        OrbitPlotOneSlice(sim, plt.subplot(gs[3]), lim=lim, limz=limz, Narc=Narc, color=color, periastron=periastron, trails=trails, show_orbit=show_orbit, lw=lw,fancy=fancy, axes="zy", plotparticles=plotparticles, primary=primary, glow=glow)
        OrbitPlotOneSlice(sim, plt.subplot(gs[0]), lim=lim, limz=limz, Narc=Narc, color=color, periastron=periastron, trails=trails, show_orbit=show_orbit, lw=lw,fancy=fancy, axes="xz", plotparticles=plotparticles, primary=primary, glow=glow)
        plt.subplot(gs[2]).set_xlabel(xlabel)
        plt.subplot(gs[2]).set_ylabel(ylabel)

        plt.setp(plt.subplot(gs[0]).get_xticklabels(), visible=False)
        plt.subplot(gs[0]).set_ylabel(zlabel)

        plt.subplot(gs[3]).set_xlabel(zlabel)
        plt.setp(plt.subplot(gs[3]).get_yticklabels(), visible=False)
    else:
        if figsize is None:
            figsize = (5,5)
        if ax is None:
            fig, ax = plt.subplots(1, 1, figsize=figsize)
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)
        OrbitPlotOneSlice(sim, ax, lim=lim, Narc=Narc, color=color, periastron=periastron, trails=trails, show_orbit=show_orbit, lw=lw,fancy=fancy, plotparticles=plotparticles, primary=primary, glow=glow)
    return plt.gcf(), ax

def get_color(color):
    """
    Takes a string for a color name defined in matplotlib and returns of a 3-tuple of RGB values.
    Will simply return passed value if it's a tuple of length three.

    Parameters
    ----------
    color   : str
        Name of matplotlib color to calculate RGB values for.
    """

    if isinstance(color, tuple) and len(color) == 3: # already a tuple of RGB values
        return color

    hexcolor = sns.colors.xkcd_rgb[color]

    hexcolor = hexcolor.lstrip('#')
    lv = len(hexcolor)
    return tuple(int(hexcolor[i:i + lv // 3], 16)/255. for i in range(0, lv, lv // 3)) # tuple of rgb values

def fading_line(x, y, color='black', alpha_initial=1., alpha_final=0., glow=False, **kwargs):
    """
    Returns a matplotlib LineCollection connecting the points in the x and y lists, with a single color and alpha varying from alpha_initial to alpha_final along the line.
    Can pass any kwargs you can pass to LineCollection, like linewidgth.

    Parameters
    ----------
    x       : list or array of floats for the positions on the (plot's) x axis
    y       : list or array of floats for the positions on the (plot's) y axis
    color   : matplotlib color for the line. Can also pass a 3-tuple of RGB values (default: 'black')
    alpha_initial:  Limiting value of alpha to use at the beginning of the arrays.
    alpha_final:    Limiting value of alpha to use at the end of the arrays.
    """
    if glow:
        glow = False
        kwargs["lw"] = 1
        fl1 = fading_line(x, y, color, alpha_initial, alpha_final, glow=False, **kwargs)
        kwargs["lw"] = 2
        alpha_initial *= 0.5
        alpha_final *= 0.5
        fl2 = fading_line(x, y, color, alpha_initial, alpha_final, glow=False, **kwargs)
        kwargs["lw"] = 6
        alpha_initial *= 0.5
        alpha_final *= 0.5
        fl3 = fading_line(x, y, color, alpha_initial, alpha_final, glow=False, **kwargs)
        return [fl3,fl2,fl1]

    color = get_color(color)
    cdict = {'red': ((0.,color[0],color[0]),(1.,color[0],color[0])),
             'green': ((0.,color[1],color[1]),(1.,color[1],color[1])),
             'blue': ((0.,color[2],color[2]),(1.,color[2],color[2])),
             'alpha': ((0.,alpha_initial, alpha_initial), (1., alpha_final, alpha_final))}

    Npts = len(x)
    if len(y) != Npts:
        raise AttributeError("x and y must have same dimension.")

    segments = np.zeros((Npts-1,2,2))
    segments[0][0] = [x[0], y[0]]
    for i in range(1,Npts-1):
        pt = [x[i], y[i]]
        segments[i-1][1] = pt
        segments[i][0] = pt
    segments[-1][1] = [x[-1], y[-1]]

    individual_cm = LinearSegmentedColormap('indv1', cdict)
    lc = LineCollection(segments, cmap=individual_cm, **kwargs)
    lc.set_array(np.linspace(0.,1.,len(segments)))
    return lc

def OrbitPlotOneSlice(sim, ax, lim=None, limz=None, Narc=100, color=False, periastron=False, trails=False, show_orbit=True, lw=1., axes="xy", plotparticles=[], primary=None, glow=False, fancy=False):
    p_orb_pairs = []
    if not plotparticles:
        plotparticles = range(1, sim.N_real)
    for i in plotparticles:
        p = sim.particles[i]
        p_orb_pairs.append((p, p.calculate_orbit(primary=primary)))

    if lim is None:
        lim = 0.
        for p, o in p_orb_pairs:
            if o.a>0.:
                r = (1.+o.e)*o.a
            else:
                r = o.d
            if r>lim:
                lim = r
        lim *= 1.15
    if limz is None:
        z = [p.z for p,o in p_orb_pairs]
        limz = 2.0*max(z)
        if limz > lim:
            limz = lim
        if limz <= 0.:
            limz = lim

    if axes[0]=="z":
        ax.set_xlim([-limz,limz])
    else:
        ax.set_xlim([-lim,lim])
    if axes[1]=="z":
        ax.set_ylim([-limz,limz])
    else:
        ax.set_ylim([-lim,lim])

    if fancy:
        ax.set_facecolor((0.,0.,0.))
        for pos in ['top', 'bottom', 'right', 'left']:
            ax.spines[pos].set_edgecolor((0.3,0.3,0.3))

    if color is not False:
        if isinstance(color, list):
            colors = []
            for c in color:
                colors.append(get_color(c))
        elif isinstance(color, str):
            colors = [get_color(color)]
        elif color == True:
            colors = [(1.,0.,0.),(0.,0.75,0.75),(0.75,0.,0.75),(0.75, 0.75, 0,),(0., 0., 0.),(0., 0., 1.),(0., 0.5, 0.)]
    else:
        if fancy:
            colors = [(181./206.,66./206.,191./206.)]
            glow = True
        else:
            colors = ["black"]
    coloriterator = cycle(colors)

#    coords = {'x':0, 'y':1, 'z':2}
#    axis0 = coords[axes[0]]
#    axis1 = coords[axes[1]]

    prim = sim.particles[0] if primary is None else primary
    if fancy:
        sun = (256./256.,256./256.,190./256.)
        opa = 0.020
        size = 6000.
        for i in range(256):
            ax.scatter(getattr(prim,axes[0]),getattr(prim,axes[1]), alpha=opa, s=size*lw, facecolor=sun, edgecolor=None, zorder=3)
            size *= 0.95

        starcolor = (1.,1.,1.)
        mi, ma = ax.get_xlim()
        prestate = random.getstate()
        random.seed(1) #always same stars
        x, y = [], []
        #small stars
        for i in range(64):
            x.append(random.uniform(mi,ma))
            y.append(random.uniform(mi,ma))
        ax.scatter(x,y, alpha=0.05, s=8*lw, facecolor=starcolor, edgecolor=None, zorder=3)
        ax.scatter(x,y, alpha=0.1, s=4*lw, facecolor=starcolor, edgecolor=None, zorder=3)
        ax.scatter(x,y, alpha=0.2, s=0.5*lw, facecolor=starcolor, edgecolor=None, zorder=3)
        #medium stars
        x, y = [], []
        for i in range(16):
            x.append(random.uniform(mi,ma))
            y.append(random.uniform(mi,ma))
        ax.scatter(x,y, alpha=0.1, s=15*lw, facecolor=starcolor, edgecolor=None, zorder=3)
        ax.scatter(x,y, alpha=0.1, s=5*lw, facecolor=starcolor, edgecolor=None, zorder=3)
        ax.scatter(x,y, alpha=0.5, s=2*lw, facecolor=starcolor, edgecolor=None, zorder=3)
        random.setstate(prestate)

    else:
        ax.scatter(getattr(prim,axes[0]),getattr(prim,axes[1]), marker="*", s=35*lw, facecolor="black", edgecolor=None, zorder=3)

    proj = {}
    for p, o in p_orb_pairs:
        colori = next(coloriterator)

        prim = p.jacobi_com if primary is None else primary
        if fancy:
            ax.scatter(getattr(p,axes[0]), getattr(p,axes[1]), s=25*lw, facecolor=colors, edgecolor=None, zorder=3)
        else:
            pass
            #::: !!!
            #::: !!!
            #::: plot the planet symbols
            #::: !!!
            #::: !!!
#            ax.scatter(getattr(p,axes[0]), getattr(p,axes[1]), s=p.r*lw, facecolor=colori, edgecolor=None, zorder=3)

        if show_orbit is True:
            alpha_final = 0. if trails is True else 1. # fade to 0 with trails

            hyperbolic = o.a < 0. # Boolean for whether orbit is hyperbolic
            if hyperbolic is False:
                pts = np.array(p.sample_orbit(Npts=Narc+1, primary=prim))
                proj['x'],proj['y'],proj['z'] = [pts[:,i] for i in range(3)]
                lc = fading_line(proj[axes[0]], proj[axes[1]], colori, alpha_final=alpha_final, lw=lw, glow=glow)
                if type(lc) is list:
                    for l in lc:
                        ax.add_collection(l)
                else:
                    ax.add_collection(lc)

            else:
                pts = np.array(p.sample_orbit(Npts=Narc+1, primary=prim, useTrueAnomaly=False))
                # true anomaly stays close to limiting value and switches quickly at pericenter for hyperbolic orbit, so use mean anomaly
                proj['x'],proj['y'],proj['z'] = [pts[:,i] for i in range(3)]
                lc = fading_line(proj[axes[0]], proj[axes[1]], colori, alpha_final=alpha_final, lw=lw, glow=glow)
                if type(lc) is list:
                    for l in lc:
                        ax.add_collection(l)
                else:
                    ax.add_collection(lc)

                alpha = 0.2 if trails is True else 1.
                pts = np.array(p.sample_orbit(Npts=Narc+1, primary=prim, trailing=False, useTrueAnomaly=False))
                proj['x'],proj['y'],proj['z'] = [pts[:,i] for i in range(3)]
                lc = fading_line(proj[axes[0]], proj[axes[1]], colori, alpha_initial=alpha, alpha_final=alpha, lw=lw, glow=glow)
                if type(lc) is list:
                    for l in lc:
                        ax.add_collection(l)
                else:
                    ax.add_collection(lc)

        if periastron:
            newp = Particle(a=o.a, f=0., inc=o.inc, omega=o.omega, Omega=o.Omega, e=o.e, m=p.m, primary=prim, simulation=sim)
            ax.plot([getattr(prim,axes[0]), getattr(newp,axes[0])], [getattr(prim,axes[1]), getattr(newp,axes[1])], linestyle="dotted", c=colori, zorder=1, lw=lw)
            ax.scatter([getattr(newp,axes[0])],[getattr(newp,axes[1])], marker="o", s=5.*lw, facecolor="none", edgecolor=colori, zorder=1)

def plot_top_down_view(params_median, params_star, a=None, timestep=None, scaling=5., colors=sns.color_palette('deep'), linewidth=2, plot_arrow=False, ax=None):

    sim = rebound.Simulation()
    sim.add(m=1)

    for i, companion in enumerate(config.BASEMENT.settings['companions_all']):
        if (i==0) and (timestep is None):
            timestep = params_median[companion+'_epoch'] #calculate it for the timestep where the first companion is in transit
        first_epoch = get_first_epoch(timestep, params_median[companion+'_epoch'], params_median[companion+'_period'])
        phase = calc_phase(timestep, params_median[companion+'_period'], first_epoch)
        ecc = params_median[companion+'_f_s']**2 + params_median[companion+'_f_c']**2
        w = np.arccos( params_median[companion+'_f_c'] / np.sqrt(ecc) ) #in rad
        inc = params_median[companion+'_incl']/180.*np.pi
        if a is None:
            a1 = params_star['R_star'] / params_median[companion+'_radius_1'] #in Rsun
            a1 *= 0.004650467260962157 #in AU
        else:
            a1 = a[i]
#        print(a, inc, ecc, w, phase*2*np.pi)
        if ecc>0:
            sim.add(a=a1, inc=inc-np.pi/2., e=ecc, omega=w, f=phase*2*np.pi)
        else:
            sim.add(a=a1, inc=inc--np.pi/2., f=phase*2*np.pi)
#    print(len(sim.particles))

#    print('Epoch, Period and mean anomaly, b:', sim.particles[0].M )
#    print('Mean anomaly, c:', sim.particles[0].M )
#    print('Mean anomaly, c:', sim.particles[0].M )
#    err

    fig, ax = OrbitPlot(sim, xlabel='AU', ylabel='AU', color=colors, lw=linewidth, ax=ax) #color=[sns.color_palette('deep')[i] for i in [0,1,3]],

    for i, companion in enumerate(config.BASEMENT.settings['companions_all']):

        R_companion = params_star['R_star'] * params_median[companion+'_rr'] # in Rsun
        R_companion *= 0.004650467260962157 #in AU
        R_companion *= scaling

        x = sim.particles.get(i+1).x
        y = sim.particles.get(i+1).y
        p = Circle((x,y), R_companion, color=colors[i])
        ax.add_artist(p)

    if plot_arrow:
        x0, x1 = ax.get_xlim()
        plt.arrow( 0.1*x1, 0, 0.7*x1, 0, color='silver', zorder=1 )

    plt.axis('equal')
#    ax.set(xlim=[-0.12,0.12], ylim=[-0.12,0.12])
#    loc = plticker.MultipleLocator(base=0.05) # this locator puts ticks at regular intervals
#    ax.xaxis.set_major_locator(loc)
#    ax.yaxis.set_major_locator(loc)

    return fig, ax


#!/usr/bin/env python
# -*- coding: utf8 -*-
"""
    SAMI XJoin

    This script simply joins the four existing extensions inside a FITS file
    created during observations with SAMI (SAM Imager). During the reduce,
    it also fits a 2nd degree polynomium to the OVERSCAN region that is
    subtracted from the corresponding image.

    The user also may want to add flags in order to reduce the images
    according to the following options (in order):

    - BIAS subtraction;
    - DARK subtraction;
    - Remove hot pixels and cosmic rays;
    - Remove overglow using a long exposure DARK image;
    - Divide by the FLAT;
    - Divide by the exposure time;

    The documentation for each reduce is shown in the corresponding function.

    Todo
    ----
    - Use multithread or multiprocessing to run this script faster.
    - Use astropy.ccdproc to reduce the data.

    Bruno Quint (bquint at ctio.noao.edu)
    May 2016

    Thanks to Andrei Tokovinin and Claudia M. de Oliveira for the ideas that
    were implemented here.
"""

import numpy as _np

from ccdproc import cosmicray_lacosmic as _cosmicray_lacosmic
from scipy import stats

from astropy import wcs
from astropy.coordinates import SkyCoord
from astropy import units as u

from soar_simager.io import pyfits as _pyfits
from soar_simager.io.logging import get_logger
from soar_simager.tools import slices

logger = get_logger(__name__)

# Piece of code from cosmics.py
# We define the laplacian kernel to be used
_laplkernel = _np.array([[0.0, -1.0, 0.0], [-1.0, 4.0, -1.0], [0.0, -1.0, 0.0]])

# Other kernels :
_growkernel = _np.ones((3, 3))

# dilation structure for some morphological operations
_dilstruct = _np.ones((5, 5))
_dilstruct[0, 0] = 0
_dilstruct[0, 4] = 0
_dilstruct[4, 0] = 0
_dilstruct[4, 4] = 0

class Reducer:
    """
    This class holds all the methods used to join the extensions within a
    FITS file obtained with SAMI.

    Parameters
    ----------
        zero_file : str
            The filename of the master zero that will be used in subtraction.

        clean : bool
            Clean bad collumns by taking the _median value of the pixels around
            them.

        cosmic_rays : bool
            Clean cosmic rays using LACosmic package. See noted bellow for
            reference.

        dark_file : str
            Master Dark's filename to be used for dark subtraction.

        debug : bool
            Turn on debug mode with lots of printing.

        flat_file : str
            Master Flat filename to be used for normalization.

        glow_file : str
            Master file that contains the lateral glowings sometimes present in
            SAMI's data.

        time : bool
            Divide each pixel's values by the exposure time and update header.

        verbose : bool
            Turn on verbose mode (not so talktive as debug mode).

    Attributes
    ----------
        gain : list
            A list containing the gain that converts ADU values to eletrons for
            each simager amplifier.

        read_noise : list
            A list containing the read noise on each simager amplifier.

    See also
    --------
        LACosmic - http://www.astro.yale.edu/dokkum/lacosmic/
    """

    gain = [2.6, 2.6, 2.6, 2.6]
    read_noise = [10., 10., 10., 10.]

    def __init__(self, clean=False, cosmic_rays=False, dark_file=None,
                 debug=False, flat_file=None, glow_file=None, merge=False,
                 overscan=False, norm_flat=False, time=False, verbose=False,
                 zero_file=None):

        logger.setLevel("ERROR")

        if verbose:
            logger.setLevel("INFO")

        if debug:
            logger.setLevel("DEBUG")

        self.clean = clean
        self.cosmic_rays = cosmic_rays
        self.dark_file = dark_file
        self.flat_file = flat_file
        self.glow_file = glow_file
        self._merge = merge
        self.norm_flat = norm_flat
        self.overscan = overscan
        self.time = time
        self.zero_file = zero_file

        return

    def reduce(self, hdu_list, prefix=""):

        # If the number of extensions is just 1, then the file is already
        # processed.
        if len(hdu_list) == 1:
            return hdu_list, ''

        # Merge file
        data, header, prefix = self.merge(hdu_list)

        # Correct ZERO
        data, header, prefix = self.correct_zero(
            data, header, prefix, self.zero_file
        )

        # Correct DARK
        data, header, prefix = self.correct_dark(
            data, header, prefix, self.dark_file
        )

        # Remove cosmic rays and hot pixels
        data, header, prefix = self.remove_cosmic_rays(
            data, header, prefix, self.cosmic_rays
        )

        # Remove lateral glows
        data, header, prefix = self.correct_lateral_glow(
            data, header, prefix, self.glow_file
        )

        # Correct FLAT
        data, header, prefix = self.correct_flat(
            data, header, prefix, self.flat_file
        )

        # Normalize by the EXPOSURE TIME
        data, header, prefix = self.divide_by_exposuretime(
            data, header, prefix, self.time
        )

        # Clean known bad columns and lines
        data, header, prefix = self.clean_hot_columns_and_lines(
            data, header, prefix, self.clean
        )

        # Add WCS
        data, header = self.create_wcs(
            data, header
        )

        return data, header, prefix

    @staticmethod
    def create_wcs(data, header):
        """
        Creates a first guess of the WCS using the telescope coordinates, the
        CCDSUM (binning), position angle and plate scale.

        Parameters
        ----------
            data : numpy.ndarray
                2D array with the data.

            header : astropy.io.fits.Header
                Primary Header to be updated.

        Returns
        -------
            header : astropy.io.fits.Header
                Primary Header with updated WCS information.
        """
        h = header

        if 'EQUINOX' not in h:
            h['EQUINOX'] = 2000.

        if 'EPOCH' not in h:
            h['EPOCH'] = 2000.

        if h['PIXSCAL1'] != h['PIXSCAL2']:
            logger.warning('Pixel scales for X and Y do not mach.')

        if h['OBSTYPE'] != 'OBJECT':
            return data, header

        binning = _np.array([int(b) for b in h['CCDSUM'].split(' ')])
        plate_scale = h['PIXSCAL1'] * u.arcsec
        p = plate_scale.to('degree').value
        w = wcs.WCS(naxis=2)

        try:
            coordinates = SkyCoord(ra=h['RA'], dec=h['DEC'],
                                   unit=(u.hourangle, u.deg))

        except ValueError:

            logger.error(
                '"RA" and "DEC" missing. Using "TELRA" and "TELDEC" instead.')

            coordinates = SkyCoord(ra=h['TELRA'], dec=h['TELDEC'],
                                   unit=(u.hourangle, u.deg))

        ra = coordinates.ra.to('degree').value
        dec = coordinates.dec.to('degree').value

        w.wcs.crpix = [data.shape[1] / 2, data.shape[0] / 2]
        w.wcs.cdelt = p * binning
        w.wcs.crval = [ra, dec]
        w.wcs.ctype = ["RA---TAN", "DEC--TAN"]

        wcs_header = w.to_header()

        theta = _np.deg2rad(h['DECPANGL'])
        wcs_header['cd1_1'] = p * binning[0] * _np.cos(theta)
        wcs_header['cd2_2'] = p * binning[0] * _np.cos(theta)
        wcs_header['cd1_2'] = p * binning[0] * _np.sin(theta)
        wcs_header['cd2_1'] = - p * binning[0] * _np.sin(theta)

        for key in wcs_header.keys():
            header[key] = wcs_header[key]

        return data, header

    @staticmethod
    def check_header(hdu_list, prefix):

        for i in range(5):

            h = hdu_list[i].header

            try:
                h['RADESYSa'] = h['RADECSYS']
                del h['RADECSYS']
            except KeyError:
                pass

            if 'EQUINOX' in h and 'unavail' in h['EQUINOX']:
                h['EQUINOX'] = 2000.

            if 'EPOCH' not in h:
                h['EPOCH'] = 2000.

        return hdu_list, prefix

    @staticmethod
    def clean_column(_data, x0, y0, yf, n=5):
        """
        Substitutes a single column by the _median of the neighbours columns.

        Args:

            _data (numpy.ndarray) : A 2D numpy array that contains the data.

            x0 (int) : X position of the pixel to be cleaned.

            y0 (int) : Start position of the column.

            yf (int) : Final position of the column.

            n (int, optional) : Number of neighbour columns (Default=5).

        Returns:

            _data (numpy.ndarray) : Processed 2D numpy array.

        See also:

            Reducer.clean_columns
            Reducer.clean_line
            Reducer.clean_lines
        """

        if not isinstance(_data, _np.ndarray):
            raise (TypeError, 'Please, use a np.array as input')

        if _data.ndim is not 2:
            raise (TypeError, 'Data contains %d dimensions while it was '
                              'expected 2 dimensions.')

        t1 = _data[y0:yf, x0 - n:x0]
        t2 = _data[y0:yf, x0 + 1:x0 + n]
        t = _np.hstack((t1, t2))
        _data[y0:yf, x0] = _np.median(t, axis=1)

        return _data

    def clean_columns(self, data, header):
        """
        Clean the known bad columns that exists in most of SAMI's, SOI's or
        SIFS's data. This method is meant to be overwritten via inheritance.

        Args:

            data (numpy.ndarray) : A 2D numpy array that contains the data.

            header (astropy.io.fits.Header) : A header that will be updated.

        See also:

            Reducer.clean_column
            Reducer.clean_line
            Reducer.clean_lines
        """
        return data, header

    @staticmethod
    def clean_line(_data, x0, xf, y, n=5):
        """
        Substitutes a single column by the _median of the neighbours columns.

        Args:

            _data (numpy.ndarray) : A 2D numpy array that contains the data.

            x0 (int) : Start position of the line.

            xf (int) : Final position of the line.

            y (int) : Y position of the pixel to be cleaned.

            n (int) : Number of neighbour columns. (Default=5)

        See also:

            Reducer.clean_column
            Reducer.clean_columns
            Reducer.clean_lines
        """
        if not isinstance(_data, _np.ndarray):
            raise (TypeError, 'Please, use a np.array as input')

        if _data.ndim is not 2:
            raise (TypeError, 'Data contains %d dimensions while it was '
                              'expected 2 dimensions.')

        t1 = _data[y - n:y, x0:xf]
        t2 = _data[y + 1:y + n, x0:xf]
        t = _np.vstack((t1, t2))
        _data[y, x0:xf] = _np.median(t, axis=0)

        return _data

    def clean_lines(self, data, header):
        """
        Clean the known bad lines that exists in most of SAMI's, SOI's or
        SIFS's data. This method is meant to be overwritten via inheritance.

        Args:

            data (numpy.ndarray) : A 2D numpy array that contains the data.

            header (astropy.io.fits.Header) : A header that will be updated.

        See also:

            Reducer.clean_column
            Reducer.clean_line
            Reducer.clean_lines
        """
        return data, header

    def clean_hot_columns_and_lines(self, data, header, prefix, clean):
        """
        Clean known hot columns and lines from SAMI's images.

        Args:

            data (numpy.ndarray) : A 2D numpy array that contains the data.

            header (astropy.io.fits.Header) : A header that will be updated.

            prefix (str) : File prefix that is added after each reduce.

            clean (bool) : Should I perform the clean?

        See also:

            Reducer.clean_column
            Reducer.clean_columns
            Reducer.clean_line
            Reducer.clean_lines
        """
        if clean is True:

            data = self.clean_columns(data, header)
            data = self.clean_lines(data, header)
            header.add_history('Cleaned bad columns and lines.')
            prefix = 'c' + prefix

        return data, header, prefix

    @staticmethod
    def correct_dark(data, header, prefix, dark_file=None):
        """
        Subtract the dark file from data and add HISTORY to header.

        Args:

            data (numpy.ndarray) : A 2D numpy array that contains the data.

            header (astropy.io.fits.Header) : A header that will be updated.

            prefix : str
                File prefix that is added after each reduce.

            dark_file: str | None
                Master Dark filename. If None is given, nothing is done.
        """

        if not isinstance(prefix, str):
            raise (TypeError, 'Expected string but found %s instead.' %
                   prefix.__class__)

        if dark_file is not None:

            dark = _pyfits.open(dark_file)[0]
            dark.data = dark.data / float(dark.header['EXPTIME'])

            data = data - dark.data * header['EXPTIME']
            header['DARKFILE'] = dark_file
            prefix = 'd' + prefix

        return data, header, prefix

    @staticmethod
    def correct_flat(data, header, prefix, flat_file):
        """
        Divide the image by the master flat file and add HISTORY to header.

        Args:

            data (numpy.ndarray) : A 2D numpy array that contains the data.

            header (astropy.io.fits.Header) : A header that will be updated.

            prefix (str) : File prefix that is added after each reduce.

            flat_file (str or None) : Master flat filename. If None is given,
            nothing is done.
        """
        if not isinstance(prefix, str):
            raise (TypeError, 'Expected string but found %s instead.' %
                   prefix.__class__)

        if flat_file is not None:
            flat = _pyfits.open(flat_file)[0]

            data /= flat.data
            header['FLATFILE'] = flat_file
            prefix = 'f' + prefix

        return data, header, prefix

    def correct_lateral_glow(self, data, header, prefix, glow_file):
        """
        Remove lateral glows by scaling the glows in the `glow_file` based
         on `data` and subtracting it.

        Args:

            data (numpy.ndarray) : A 2D numpy array that contains the data.

            header (astropy.io.fits.Header) : A header that will be updated.

            prefix : str
                Filename prefix to flag images that were clean.

            glow_file : str
                Path to a long dark file that contains the lateral glow.
        """

        if glow_file is not None:

            # Create four different regions.
            regions = [
                [_np.median(data[539:589, 6:56]),  # Top Left
                 _np.median(data[539:589, 975:1019])],  # Top Right
                [_np.median(data[449:506, 6:56]),  # Bottom Left
                 _np.median(data[449:506, 975:1019])]  # Bottom Right
            ]

            min_std_region = _np.argmin(regions) % 2

            # The upper reg has background lower or equal to the lower reg
            midpt1 = regions[0][min_std_region]
            midpt2 = regions[1][min_std_region]
            diff = midpt2 - midpt1

            dark = _pyfits.getdata(glow_file)
            dark = self.clean_columns(dark)
            dark = self.clean_lines(dark)

            dark_regions = [
                [_np.median(dark[539:589, 6:56]),  # Top Left
                 _np.median(dark[539:589, 975:1019])],  # Top Right
                [_np.median(dark[449:506, 6:56]),  # Bottom Left
                 _np.median(dark[449:506, 975:1019])]  # Bottom Right
            ]

            dark_midpt1 = dark_regions[0][min_std_region]
            dark_midpt2 = dark_regions[1][min_std_region]

            dark_diff = dark_midpt2 - dark_midpt1
            dark -= dark_midpt1

            k = diff / dark_diff
            temp_dark = dark * k
            data -= midpt1
            data -= temp_dark

            header.add_history('Lateral glow removed using %s file' % glow_file)
            prefix = 'g' + prefix

        return data, header, prefix

    @staticmethod
    def correct_zero(data, header, prefix, zero_file):
        """
        Subtract zero from data.

        Args:

            data (numpy.ndarray) : A 2D numpy array that contains the data.

            header (astropy.io.fits.Header) : A header that will be updated.

            prefix (str) : File prefix that is added after each reduce.

            zero_file (str | None) : Master Bias filename. If None is given,
            nothing is done.

        """
        from os.path import abspath

        if zero_file is not None:

            zero = _pyfits.open(abspath(zero_file))[0]
            data = data - zero.data
            header['BIASFILE'] = zero_file
            prefix = 'z' + prefix

        return data, header, prefix

    @staticmethod
    def divide_by_exposuretime(data, header, prefix, time):
        """
            Divide the image by the exposure time and add HISTORY to header.

            Args:

                data (numpy.ndarray) : A 2D numpy array that contains the data.

                header (astropy.io.fits.Header) : A header that will be updated.

                prefix : str
                    File prefix that is added after each reduce.

                time: bool
                    Divide image by exposure time?
        """
        if time is True:

            h = header

            try:

                h['UNITS'] = 'adu / s'
                t = float(h['EXPTIME'])
                d = data / t

                header = h
                data = d

            except AttributeError:
                header = h

            except KeyError:
                pass

            prefix = 't' + prefix

        return data, header, prefix

    def get_header(self, hdu_source):
        """
        Return the header of the primary HDU extension of a FITS file.

        Args:

            hdu_source (str or astropy.io.fits.HDUList) : HDUList or name of the
            file which contains a HDUList.
        """
        from os.path import exists

        if isinstance(hdu_source, str):

            if not exists(hdu_source):
                raise (IOError, '%s file not found.' % hdu_source)

            hdu_source = _pyfits.open(hdu_source)

        h0 = hdu_source[0].header
        h1 = hdu_source[1].header

        h0.append('UNITS')
        h0.set('UNITS', value='ADU', comment='Pixel intensity units.')

        # Save the CCD binning in the main header
        h0['CCDSUM'] = h1['CCDSUM']
        h0['DETSEC'] = h1['DETSEC']

        # Save the area that corresponds to each amplifier
        bin_size = _np.array(h0['CCDSUM'].split(' '), dtype=int)

        dx, dy = slices.iraf2python(h0['DETSEC'])
        dx, dy = dx // bin_size[0], dy // bin_size[1]

        h0['AMP_SEC1'] = slices.python2iraf(
            dx[0], dx[1], dy[0], dy[1])

        h0['AMP_SEC2'] = slices.python2iraf(
            dx[0] + dx[1], dx[1] + dx[1], dy[0], dy[1])

        h0['AMP_SEC3'] = slices.python2iraf(
            dx[0], dx[1], dy[0] + dy[1], dy[1] + dy[1])

        h0['AMP_SEC4'] = slices.python2iraf(
            dx[0] + dx[1], dx[1] + dx[1], dy[0] + dy[1], dy[1] + dy[1])

        return h0

    def get_prefix(self):
        """
        Return a prefix to be added to the file deppending on the data
        reduction steps.

        Returns
        -------
            prefix : (str)
                The prefix that can be used.
                    m = merged amplifiers.
                    z = zero subtracted.
                    f = flat corrected.
        """

        prefix = 'm_'

        if self.zero_file:
            prefix = 'z' + prefix

        if self.dark_file:
            prefix = 'd' + prefix

        if self.flat_file:
            prefix = 'f' + prefix

        return prefix

    def merge(self,  hdul):
        """
        Open a FITS image and try to join its extensions in a single array.

        Args:

            hdul (astropy.io.fits.HDUList) : an HDUList that contains one
            PrimaryHDU and four ImageHDU

        """
        w, h = slices.iraf2python(hdul[1].header['DETSIZE'])

        if len(hdul) is 1:
            logger.warning('%s file contains a single extension. ' % hdul +
                           'Not doing anything')
            return hdul[0].data

        # Correct for binning
        bin_size = _np.array(hdul[1].header['CCDSUM'].split(' '),
                             dtype=int)
        bw, bh = w[1] // bin_size[0], h[1] // bin_size[1]

        # Create empty full frame
        new_data = _np.empty((bh, bw), dtype=float)

        # Process each extension
        for i in range(1, 5):
            tx, ty = slices.iraf2python(hdul[i].header['TRIMSEC'])
            bx, by = slices.iraf2python(hdul[i].header['BIASSEC'])

            data = hdul[i].data
            trim = data[ty[0]:ty[1], tx[0]:tx[1]]
            bias = data[by[0]:by[1], bx[0]:bx[1]]

            # Collapse the bias columns to a single column.
            bias = _np.median(bias, axis=1)

            # Fit and remove OVERSCAN
            x = _np.arange(bias.size) + 1
            bias_fit_pars = _np.polyfit(x, bias, 2)  # Last par = inf
            bias_fit = _np.polyval(bias_fit_pars, x)
            bias_fit = bias_fit.reshape((bias_fit.size, 1))
            bias_fit = _np.repeat(bias_fit, trim.shape[1], axis=1)

            trim = trim - bias_fit
            dx, dy = slices.iraf2python(hdul[i].header['DETSEC'])
            dx, dy = dx // bin_size[0], dy // bin_size[1]
            new_data[dy[0]:dy[1], dx[0]:dx[1]] = trim

        header = self.get_header(hdul)

        return new_data, header, "m_"

    @staticmethod
    def remove_cosmic_rays(data, header, prefix, cosmic_rays):
        """
        Use LACosmic to remove cosmic rays.

        Args:

            data (numpy.ndarray) : A 2D numpy array that contains the data.

            header (astropy.io.fits.Header) : A header that will be updated.

            prefix : str
                Filename prefix to flag images that were clean.

            cosmic_rays : bool
                Flag to indicate if cosmic rays removal should be performed.
        """
        if cosmic_rays:

            d = data
            d, _ = _cosmicray_lacosmic(
                d, gain=2.6, readnoise=10.0, sigclip=2.5, sigfrac=0.3,
                objlim=5.0)
            d /= 2.6

            h = header
            h.add_history(
                'Cosmic rays and hot pixels removed using LACosmic')

            data = d
            header = h

        return data, header, prefix

    @staticmethod
    def remove_wcs(header):

        return header

class SamiReducer(Reducer):

    gain = [2.1, 2.0537, 2.1, 2.0823]
    read_noise = [10., 10., 10., 10.]

    def reduce(self, hdu_list, prefix=""):

        # If the number of extensions is just 1, then the file is already
        # processed.
        if len(hdu_list) == 1:
            return hdu_list, ''

        # Merge file
        data, header, prefix = self.merge(hdu_list)

        # Removing bad column and line
        data, header, prefix = self.remove_central_bad_columns(
            data, header, prefix,
        )

        # Correct ZERO
        data, header, prefix = self.correct_zero(
            data, header, prefix, self.zero_file
        )

        # Correct DARK
        data, header, prefix = self.correct_dark(
            data, header, prefix, self.dark_file
        )

        # Remove cosmic rays and hot pixels
        data, header, prefix = self.remove_cosmic_rays(
            data, header, prefix, self.cosmic_rays
        )

        # Remove lateral glows
        data, header, prefix = self.correct_lateral_glow(
            data, header, prefix, self.glow_file
        )

        # Correct FLAT
        data, header, prefix = self.correct_flat(
            data, header, prefix, self.flat_file
        )

        # Normalize by the EXPOSURE TIME
        data, header, prefix = self.divide_by_exposuretime(
            data, header, prefix, self.time
        )

        # Clean known bad columns and lines
        data, header, prefix = self.clean_hot_columns_and_lines(
            data, header, prefix, self.clean
        )

        # Add WCS
        data, header = self.create_wcs(
            data, header
        )

        return data, header, prefix

    def clean_columns(self, data, header):
        """
        Clean the known bad columns that exists in most of SAMI's, SOI's or
        SIFS's data. This method is meant to be overwritten via inheritance.

        Args:

            data (numpy.ndarray) : A 2D numpy array that contains the data.

            header (astropy.io.fits.Header) : A header that will be updated.

        See also:

            Reducer.clean_column
            Reducer.clean_line
            Reducer.clean_lines
        """
        binning = header['CCDSUM'].split(' ')[0]
        binning = int(binning.strip())

        if binning == 4:
            bad_columns = [
                [167, 0, 513],
                [213, 513, 1023],
                [304, 0, 513],
                [309, 1, 512],
                [386, 0, 513],
                [476, 0, 513],
                [602, 0, 513],
                [671, 0, 513],
                [673, 475, 513],
                [678, 0, 513],
                [741, 0, 513],
                [810, 0, 513],
                [919, 0, 513],
                [212, 513, 1023],
                [680, 513, 1023],
                [725, 513, 1023],
                [848, 513, 1023],
                [948, 0, 512],
                [949, 0, 512]
                ]
        else:
            []

        for column in bad_columns:
            x0 = column[0]
            y0 = column[1]
            yf = column[2]
            data = self.clean_column(data, x0, y0, yf)

        return data

    def clean_lines(self, data, header):
        """
        Clean the known bad lines that exists in most of SAMI's, SOI's or
        SIFS's data. This method is meant to be overwritten via inheritance.

        Args:

            data (numpy.ndarray) : A 2D numpy array that contains the data.

            header (astropy.io.fits.Header) : A header that will be updated.

        See also:

            Reducer.clean_column
            Reducer.clean_line
            Reducer.clean_lines
        """
        binning = header['CCDSUM'].split(' ')[0]
        binning = int(binning.strip())

        if binning == 4:
            bad_lines = [
                [166, 206, 282],
                [212, 258, 689],
                [214, 239, 688],
                [304, 345, 291],
                [386, 422, 454],
                [398, 422, 38],
                [477, 516, 490],
                [387, 429, 455],
                [574, 603, 494],
                [574, 603, 493],
                [640, 672, 388],
                [604, 671, 388],
                [698, 746, 198],
                [706, 634, 634],
                [772, 812, 354],
                [900, 938, 426],
                [904, 920, 396]
            ]

        else:
            bad_lines = []

        for line in bad_lines:
            x0 = line[0]
            xf = line[1]
            y = line[2]
            data = self.clean_line(data, x0, xf, y)

        return data

    @staticmethod
    def remove_central_bad_columns(data, header, prefix):
        """
        Remove central bad columns at the interface of the four extensions.

        Parameter
        ---------
            data : numpy.ndarray
                2D Array containing the data.
        """
        n_rows, n_columns = data.shape

        # Copy the central bad columns to a temp array
        temp_column = data[:, n_columns // 2 - 1:n_columns // 2 + 1]

        # Shift the whole image by two columns
        data[:, n_columns // 2 - 1:-2] = data[:, n_columns // 2 + 1:]

        # Copy the bad array in the end (right) of the image).
        data[:, -2:] = temp_column

        return data, header, prefix

class SifsReducer(SamiReducer):
    pass

class SoiReducer(Reducer):
    """
    SoiReducer

    This class holds all the methods used to join the extensions within a
    FITS file obtained with SOI.

    Parameters
    ----------
        zero_file : str
            The filename of the master zero that will be used in subtraction.

        clean : bool
            Clean bad collumns by taking the _median value of the pixels around
            them.

        cosmic_rays : bool
            Clean cosmic rays using LACosmic package. See noted bellow for
            reference.

        dark_file : str
            Master Dark's filename to be used for dark subtraction.

        debug : bool
            Turn on debug mode with lots of printing.

        flat_file : str
            Master Flat filename to be used for normalization.

        glow_file : str
            Master file that contains the lateral glowings sometimes present in
            SAMI's data.

        time : bool
            Divide each pixel's values by the exposure time and update header.

        verbose : bool
            Turn on verbose mode (not so talktive as debug mode).

    See also
    --------
        LACosmic - http://www.astro.yale.edu/dokkum/lacosmic/
    """

    @staticmethod
    def add_gap(data, header, interpolation_factor=10):
        """
        SOI has two detectors which are separated by 7.8 arcsec (or 102
        unbinned pixels). This method reads an merged array and adds the gap
        based on the detector's binning.

        Parameters
        ----------
            data : numpy.ndarray
                2D array with the data merged.

            header : astropy.io.fits.Header
                a header that contains the binning information on the 'CCDSUM'
                key.
        """
        if header['OBSTYPE'] == 'OBJECT':

            binning = header['CCDSUM']
            binning = int(binning.split()[0])

            gap_size = 7.8  # arcseconds
            pixel_scale = 0.0767  # arcsecond / pixel
            gap_pixel = int(round(gap_size / pixel_scale / binning, 0))

            nrow, ncol = data.shape

            data = _np.append(data, _np.zeros((nrow, gap_pixel)), axis=1)
            data[:, ncol // 2 + gap_pixel:] = data[:, ncol // 2:- gap_pixel]
            data[:, ncol // 2:ncol // 2 + gap_pixel] = 0

        return data, header

    def clean_columns(self, _data, _header):
        """
        Clean the known bad columns that exists in most of SAMI's data.

        Parameters
        ----------
            _data : numpy.ndarray
                A 2D numpy array that contains the data.

            _header : astropy.io.fits.Header
                a header that contains the binning information on the 'CCDSUM'
                key.

        See also
        --------
            SoiMerger.clean_column
            SoiMerger.clean_line
            SoiMerger.clean_lines
        """
        if not isinstance(_data, _np.ndarray):
            raise (TypeError, 'Please, use a np.array as input')
        if _data.ndim is not 2:
            raise (TypeError, 'Data contains %d dimensions while it was '
                              'expected 2 dimensions.')

        b = int(_header['CCDSUM'].strip().split(' ')[0])

        if b == 1:
            bad_columns = []
        elif b == 2:
            bad_columns = [
                [855, 0, 2047],
            ]
        elif b == 4:
            bad_columns = [
                [427, 0, 1023]
            ]
        else:
            logger.warning(
                'Skipping clean_columns for binning {} x {}'.format(b, b))
            bad_columns = []

        for column in bad_columns:
            x0 = column[0]
            y0 = column[1]
            yf = column[2]
            _data = self.clean_column(_data, x0, y0, yf)

        return _data

    def clean_lines(self, hdu_list):
        """
        Clean the known bad lines that exists in most of SAMI's, SOI's or
        SIFS's data. This method is meant to be overwritten via inheritance.

        Args:

            hdu_list (astropy.io.fits.HDUList)

        See also:

            Reducer.clean_column
            Reducer.clean_line
            Reducer.clean_lines
        """
        if not isinstance(hdu_list, _pyfits.HDUList):
            raise TypeError('Please, use a HDUList as input')

        if len(hdu_list) != 5:
            raise ValueError(
                "HDUList is expected to have 1 + 4 elements. Found {}".format(
                    len(hdu_list)))

        for i in range(1, len(hdu_list)):

            _data = hdu_list[i].data
            _hdr = hdu_list[i].header

            bad_lines = [
                # [166, 206, 282],
                # [212, 258, 689],
                # [214, 239, 688],
                # [304, 345, 291],
                # [386, 422, 454],
                # [398, 422, 38],
                # [477, 516, 490],
                # [387, 429, 455],
                # [574, 603, 494],
                # [574, 603, 493],
                # [640, 672, 388],
                # [604, 671, 388],
                # [698, 746, 198],
                # [706, 634, 634],
                # [772, 812, 354],
                # [900, 938, 426],
                # [904, 920, 396]
            ]

            for line in bad_lines:
                x0 = line[0]
                xf = line[1]
                y = line[2]
                _data = self.clean_line(_data, x0, xf, y)

            hdu_list[i].data = _data

        return hdu_list

def _normalize_data(data):
    """
    This method is intended to normalize flat data before it is applied to the
    images that are being reduced. A total of 1000 random points are used to
    estimate the _median level that will be used for normalization.

    Args:

        data (numpy.ndarray) : Data that will be normalized

    Returns:
        norm_data (numpy.ndarray) : Normalized data.
    """
    sample = _np.random.randint(0, high=data.size - 1, size=1000)
    mode = stats.mode(data.ravel()[sample])[0]

    return data / mode


# SVMs for Food Experiment Data
import matplotlib
import numpy as np
import matplotlib.pyplot as pp
import optparse
import unittest
import random
import itertools
from sklearn import decomposition
from sklearn import svm
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import StratifiedKFold
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import LeaveOneOut
from sklearn.model_selection import KFold
from sklearn.metrics import classification_report
from sklearn.metrics import confusion_matrix
from scipy.interpolate import InterpolatedUnivariateSpline
from scipy.ndimage.filters import gaussian_filter1d
import sys
import lib_get_subset as lib_get_subset
import matplotlib.pyplot as plt

TRIAL_ORDER = [[1, 4, 6, 2, 3, 7, 0, 5, 8],
               [0, 3, 6, 2, 5, 7, 1, 4, 8],
               [2, 3, 7, 1, 4, 6, 0, 5, 8],
               [2, 4, 6, 0, 5, 8, 1, 3, 7],
               [0, 5, 6, 1, 4, 7, 2, 3, 8],
               [1, 3, 8, 0, 4, 6, 2, 5, 7],
               [0, 4, 7, 1, 3, 8, 2, 5, 6],
               [2, 3, 8, 0, 4, 6, 1, 5, 7],
               [0, 3, 8, 2, 4, 6, 1, 5, 7],
               [1, 3, 6, 0, 5, 7, 2, 4, 8],
               [0, 5, 7, 1, 3, 8, 2, 4, 6],
               [1, 3, 8, 0, 4, 7, 2, 5, 6],
               [1, 5, 6, 0, 4, 7, 2, 3, 8],
               [1, 5, 8, 2, 4, 6, 0, 3, 7],
               [1, 5, 8, 2, 4, 6, 0, 3, 7],
               [1, 4, 6, 2, 5, 7, 0, 3, 8],
               [0, 4, 6, 2, 5, 8, 1, 3, 7],
               [2, 4, 7, 0, 5, 6, 1, 3, 8],
               [1, 5, 7, 0, 3, 6, 2, 4, 8],
               [2, 5, 7, 0, 4, 8, 1, 3, 6],
               [2, 5, 7, 1, 4, 6, 0, 3, 8],
               [0, 5, 7, 2, 3, 6, 1, 4, 8],
               [2, 4, 8, 0, 5, 6, 1, 3, 7],
               [2, 5, 7, 0, 4, 8, 1, 3, 6],
               [1, 3, 7, 2, 5, 8, 0, 4, 6],
               [2, 3, 8, 1, 5, 7, 0, 4, 6],
               [2, 3, 7, 0, 5, 8, 1, 4, 6],
               [1, 4, 7, 0, 5, 6, 2, 3, 8],
               [2, 5, 8, 1, 3, 6, 0, 4, 7],
               [1, 4, 8, 0, 5, 7, 2, 3, 6]]

NUM_TRIAL_SETS = len(TRIAL_ORDER)

def get_most_overlapping(x_full, data_dict, num_pairs):

    matching_criteria = []

    for trial_set in range(len(TRIAL_ORDER)):
        optim_set = data_dict['trial_set_' + str(trial_set + 1) + '_X']
        print("trial set 1 shape", optim_set.shape)

        for block in [0, 1, 2]:
            present_set = optim_set[block, :]
            for other_trial_set in range(len(TRIAL_ORDER)):
                if other_trial_set != trial_set:
                    for CW_block in [6, 7, 8]:
                        check_set = data_dict['trial_set_' + str(other_trial_set + 1) + '_X'][CW_block, :]
                        #print(present_set)
                        #print(check_set)
                        matching_criteria.append([trial_set, block, other_trial_set, CW_block, np.sum(np.square(present_set - check_set))])

            #break

        #break

    matching_criteria = np.array(matching_criteria)

    matching_criteria = matching_criteria[matching_criteria[:,4].argsort()]
    #matching_criteria = np.flipud(matching_criteria)

    #for criteria in matching_criteria:
    #    print(criteria)

    #matching_criteria = matching_criteria[0:num_pairs, :]

    #here make sure there arne't any duplicates in the list
    reduced_matching_criteria = []
    iterator = 0
    while len(reduced_matching_criteria) < num_pairs:
        next_criteria = matching_criteria[iterator, :]
        should_append = True
        print(next_criteria)
        for item in reduced_matching_criteria:
            if item[0] == next_criteria[0] and item[1] == next_criteria[1]:
                should_append = False
            if item[2] == next_criteria[2] and item[3] == next_criteria[3]:
                should_append = False
        if should_append == True:
            reduced_matching_criteria.append(next_criteria)
        iterator += 1

    matching_criteria = np.array(reduced_matching_criteria).astype(int)

    print(matching_criteria.astype(int))

    X_data = []
    y_data = []
    #Use matching criteria indexing to create data subsets that we can use for LOOCV
    for ct in range(num_pairs):
        trial_set = int(matching_criteria[ct, 0])
        block = int(matching_criteria[ct, 1])
        X_data.append(data_dict['trial_set_' + str(trial_set + 1) + '_X'][block, :])
        y_data.append(0)
    for ct in range(num_pairs):
        trial_set = int(matching_criteria[ct, 2])
        block = int(matching_criteria[ct, 3])
        X_data.append(data_dict['trial_set_' + str(trial_set + 1) + '_X'][block, :])
        y_data.append(1)
    X_data = np.array(X_data)
    y_data = np.array(y_data)

    # print(len(skf))
    # for train_index, test_index in skf:

    cm_total = np.zeros((2,2))

    #loo = LeaveOneOut(num_pairs)
    #for train_index, test_index in loo:

    for i in range(1000):
        kfold = KFold(n_splits=3, shuffle = True)
        kfold.get_n_splits(np.arange(num_pairs))
        for train_index, test_index in kfold.split(np.arange(num_pairs)):
            print("TRAIN:", train_index, "TEST:", test_index)
            train_index_second = list(np.array(train_index) + num_pairs)
            train_index = np.array(list(train_index) + train_index_second)
            test_index_second = list(np.array(test_index) + num_pairs)
            test_index = np.array(list(test_index) + test_index_second)
            print("TRAIN:", train_index, "TEST:", test_index)

            print(np.shape(X_data))

            X_train, X_test = X_data[train_index], X_data[test_index]
            print(np.shape(X_train))

            y_train, y_test = y_data[train_index], y_data[test_index]
            # print(X_test)
            # print(y_test)

            svc = svm.SVC(kernel='linear')
            # svc = svm.SVC(kernel='rbf')
            # svc = svm.SVC(kernel='poly', degree=3)
            clf = svc.fit(X_train, y_train)
            preds = svc.predict(X_test)

            print(y_test, preds)

            cm = confusion_matrix(y_test, preds, labels=[0, 1])

            scores = clf.score(X_test, y_test)

            #print(test_index, y_test, np.array(preds))
            cm_total += cm
            print(cm)
            print(cm_total)
    print(cm_total / 1000)

    '''
    ax=[]
    fig = plt.figure()
    print(np.shape(X_data))
    for trial_set in range(num_pairs):
        ax.append(trial_set)
        print("trial set num: ", trial_set)

        ax[-1] = fig.add_subplot(4, 5, trial_set+1)

        ax[-1].plot(np.arange(1000)/200., X_data[trial_set+num_pairs, :], 'r-', linewidth=0.5)
        ax[-1].plot(np.arange(1000)/200., X_data[trial_set, :], 'b-', linewidth=0.5)

        ax[-1].set_ylim(26.0, 30.0)
        ax[-1].set_xlim(0, 5)
        ax[-1].grid()

    plt.tight_layout()
    fig = plt.gcf()
    plt.show()
    '''

def get_meanstd_overlapping2(data_dict):

    print("GOT TO OVERLAPPING2")

    X_data_all_M = []
    y_data_all_M = []
    X_data_all_CW = []
    y_data_all_CW = []
    X_data_all_W = []
    y_data_all_W = []

    for trial_set in range(len(TRIAL_ORDER)):
        optim_set = data_dict['trial_set_' + str(trial_set + 1) + '_X']
        print("trial set 1 shape", optim_set.shape)

        for block in [0, 1, 2]:
            X_data_all_M.append(optim_set[block, :])
            y_data_all_M.append(0)

        for block in [3, 4, 5]:
            X_data_all_W.append(optim_set[block, :])
            y_data_all_W.append(1)

        for block in [6, 7, 8]:
            X_data_all_CW.append(optim_set[block, :])
            y_data_all_CW.append(1)

    X_data_all_M = np.array(X_data_all_M)
    X_data_all_CW = np.array(X_data_all_CW)
    X_data_all_W = np.array(X_data_all_W)
    print(X_data_all_M.shape)
    X_data_all_M_mean = np.mean(X_data_all_M, axis = 0)
    X_data_all_M_std = np.std(X_data_all_M, axis = 0)
    X_data_all_CW_mean = np.mean(X_data_all_CW, axis = 0)
    X_data_all_CW_std = np.std(X_data_all_CW, axis = 0)
    X_data_all_W_mean = np.mean(X_data_all_W, axis = 0)
    X_data_all_W_std = np.std(X_data_all_W, axis = 0)
    #print(X_data_all_M_mean.shape)
    #print(X_data_all_M_mean)
    #print(X_data_all_M_std)

    #print(X_data_all_M_mean - X_data_all_M_std)

    X_upper_bound = np.amin([X_data_all_CW_mean + X_data_all_CW_std, X_data_all_M_mean + X_data_all_M_std], axis = 0)
    X_lower_bound = np.amax([X_data_all_CW_mean - X_data_all_CW_std, X_data_all_M_mean - X_data_all_M_std], axis = 0)
    #print("BLAH")
    #print(X_upper_bound)
    #print(X_lower_bound)
    #print("BLAH2")

    fig, ax = plt.subplots()

    #ax.plot(np.arange(0, 1000)/200., X_data_all_M_mean[0:1000], color = '#d95f02', label="mean M, active")
    #ax.plot(np.arange(0, 1000)/200., X_data_all_CW_mean[0:1000], '--', color = '#1b9e77', label="mean CW, active")
    #ax.plot(np.arange(0, 1000)/200., X_data_all_W_mean[0:1000],  '-.', color = '#7570b3', label="mean W, active")
    #ax.fill_between(np.arange(0, 1000)/200.,
    #                X_data_all_M_mean[0:1000] - X_data_all_M_std[0:1000],
    #                X_data_all_M_mean[0:1000] + X_data_all_M_std[0:1000], facecolor='#d95f02', alpha = 0.3)
    #ax.fill_between(np.arange(0, 1000)/200.,
    #                X_data_all_CW_mean[0:1000] - X_data_all_CW_std[0:1000],
    #                X_data_all_CW_mean[0:1000] + X_data_all_CW_std[0:1000], facecolor='#1b9e77', alpha = 0.3)
    #ax.fill_between(np.arange(0, 1000)/200.,
    #                X_data_all_W_mean[0:1000] - X_data_all_W_std[0:1000],
    #                X_data_all_W_mean[0:1000] + X_data_all_W_std[0:1000], facecolor='#7570b3', alpha = 0.3)
    #ax.plot(np.arange(0, 1000)/200., X_lower_bound, 'c-')
    #ax.plot(np.arange(0, 1000)/200., X_upper_bound, 'c-')

    ax.plot(np.arange(0, 1000)/200., X_data_all_M_mean[1000:2000], color = '#d95f02', label="mean M, passive")
    ax.plot(np.arange(0, 1000)/200., X_data_all_CW_mean[1000:2000], '--', color = '#1b9e77', label="mean CW, passive")
    ax.plot(np.arange(0, 1000)/200., X_data_all_W_mean[1000:2000],  '-.', color = '#7570b3', label="mean W, passive")

    ax.fill_between(np.arange(0, 1000)/200.,
                    X_data_all_M_mean[1000:2000] - X_data_all_M_std[1000:2000],
                    X_data_all_M_mean[1000:2000] + X_data_all_M_std[1000:2000], facecolor='#d95f02', alpha = 0.3)
    ax.fill_between(np.arange(0, 1000)/200.,
                    X_data_all_CW_mean[1000:2000] - X_data_all_CW_std[1000:2000],
                    X_data_all_CW_mean[1000:2000] + X_data_all_CW_std[1000:2000], facecolor='#1b9e77', alpha = 0.3)
    ax.fill_between(np.arange(0, 1000)/200.,
                    X_data_all_W_mean[1000:2000] - X_data_all_W_std[1000:2000],
                    X_data_all_W_mean[1000:2000] + X_data_all_W_std[1000:2000], facecolor='#7570b3', alpha = 0.3)
    #ax.plot(np.arange(0, 1000)/200., X_lower_bound, 'c-')
    #ax.plot(np.arange(0, 1000)/200., X_upper_bound, 'c-')

    ax.legend(loc=3, fontsize=10)
    ax.set_title("Mean Temperature and \n Standard Deviation Bounds", fontsize=16)
    ax.set_xlabel("Time Elapsed (seconds)", fontsize=14)
    ax.set_ylabel("Temperature, Celsius", fontsize=14)
    plt.tight_layout()
    fig = plt.gcf()
    plt.show()

    X_data_all_M = list(X_data_all_M)
    X_data_all_CW = list(X_data_all_CW)

    X_data = []
    y_data = []

    #ax.plot(np.arange(0, 1000) / 200., X_data_all_mean + X_data_all_std, 'k-', label="mean")
    for ct in range(len(X_data_all_M)):

        if np.min(X_upper_bound - X_data_all_M[ct]) > 0 and\
                np.min(X_data_all_M[ct] - X_lower_bound) > 0:
            #ax.plot(np.arange(0, 1000) / 200., X_data_all[ct], 'r-', linewidth = 0.5)
            X_data.append(X_data_all_M[ct])
            y_data.append(y_data_all_M[ct])

    for ct in range(len(X_data_all_CW)):

        if np.min(X_upper_bound - X_data_all_CW[ct]) > 0 and\
                np.min(X_data_all_CW[ct] - X_lower_bound) > 0:
            #ax.plot(np.arange(0, 1000) / 200., X_data_all[ct], 'r-', linewidth = 0.5)
            X_data.append(X_data_all_CW[ct])
            y_data.append(y_data_all_CW[ct])

    print(np.shape(X_data))
    print(np.shape(y_data))
    print(np.sum(y_data))

    #ax.legend()
    #plt.tight_layout()
    #fig = plt.gcf()
    #plt.show()

    X_data = np.array(X_data)
    y_data = np.array(y_data)

    from time import sleep
    sleep(4)

    cm_total = np.zeros((2,2))

    for i in range(1000):

        skf = StratifiedKFold(y_data, n_folds=3, shuffle=True)
        for train_index, test_index in skf:
            print("TRAIN:", train_index, "TEST:", test_index)

            print(np.shape(X_data))

            X_train = X_data[train_index]
            X_test = X_data[test_index]
            print(np.shape(X_train))

            y_train, y_test = y_data[train_index], y_data[test_index]
            # print(X_test)
            # print(y_test)

            svc = svm.SVC(kernel='linear')
            # svc = svm.SVC(kernel='rbf')
            # svc = svm.SVC(kernel='poly', degree=3)
            clf = svc.fit(X_train, y_train)
            preds = svc.predict(X_test)

            print(y_test, preds)

            cm = confusion_matrix(y_test, preds, labels=[0, 1])

            scores = clf.score(X_test, y_test)

            #print(test_index, y_test, np.array(preds))
            cm_total += cm
            print(cm)
            print(cm_total)
    print(cm_total/1000)

    '''
    ax=[]
    fig = plt.figure()
    print(np.shape(X_data))
    for trial_set in range(num_pairs):
        ax.append(trial_set)
        print("trial set num: ", trial_set)

        ax[-1] = fig.add_subplot(4, 5, trial_set+1)

        ax[-1].plot(np.arange(1000)/200., X_data[trial_set+num_pairs, :], 'r-', linewidth=0.5)
        ax[-1].plot(np.arange(1000)/200., X_data[trial_set, :], 'b-', linewidth=0.5)

        ax[-1].set_ylim(26.0, 30.0)
        ax[-1].set_xlim(0, 5)
        ax[-1].grid()

    plt.tight_layout()
    fig = plt.gcf()
    plt.show()
    '''
def get_meanstd_overlapping(x_full, data_dict):

    X_data_all = []
    y_data_all = []

    for trial_set in range(len(TRIAL_ORDER)):
        optim_set = data_dict['trial_set_' + str(trial_set + 1) + '_X']
        print("trial set 1 shape", optim_set.shape)

        for block in [0, 1, 2, 6, 7, 8]:
            X_data_all.append(optim_set[block, :])
            if block in [0, 1, 2]:
                y_data_all.append(0)
            else:
                y_data_all.append(1)

    X_data_all = np.array(X_data_all)
    print(X_data_all.shape)
    X_data_all_mean = np.mean(X_data_all, axis = 0)
    X_data_all_std = np.std(X_data_all, axis = 0)
    print(X_data_all_mean.shape)
    print(X_data_all_mean)
    print(X_data_all_std)

    print(X_data_all_mean - X_data_all_std)

    fig, ax = plt.subplots()

    ax.plot(np.arange(0, 1000)/200., X_data_all_mean, 'k-', label="mean")
    ax.fill_between(np.arange(0, 1000)/200.,
                    X_data_all_mean - X_data_all_std,
                    X_data_all_mean + X_data_all_std)
    ax.legend()
    plt.tight_layout()
    fig = plt.gcf()
    plt.show()

    X_data_all = list(X_data_all)

    X_data = []
    y_data = []

    ax.plot(np.arange(0, 1000) / 200., X_data_all_mean + X_data_all_std, 'k-', label="mean")
    for ct in range(len(X_data_all)):

        print(ct, np.shape(X_data_all_mean + X_data_all_std - X_data_all[ct]),
              np.min(X_data_all_mean + X_data_all_std - X_data_all[ct]),
              np.min(X_data_all[ct] + X_data_all_std - X_data_all_mean))

        if np.min(X_data_all_mean + X_data_all_std - X_data_all[ct]) > 0 and\
                np.min(X_data_all[ct] - X_data_all_mean + X_data_all_std) > 0:
            #ax.plot(np.arange(0, 1000) / 200., X_data_all[ct], 'r-', linewidth = 0.5)
            X_data.append(X_data_all[ct])
            y_data.append(y_data_all[ct])

    print(np.shape(X_data))
    print(np.shape(y_data))
    print(np.sum(y_data))

    #ax.legend()
    #plt.tight_layout()
    #fig = plt.gcf()
    #plt.show()

    X_data = np.array(X_data)
    y_data = np.array(y_data)

    cm_total = np.zeros((2,2))

    for i in range(1000):

        skf = StratifiedKFold(y_data, n_folds=3, shuffle=True)
        for train_index, test_index in skf:
            print("TRAIN:", train_index, "TEST:", test_index)

            print(np.shape(X_data))

            X_train = X_data[train_index]
            X_test = X_data[test_index]
            print(np.shape(X_train))

            y_train, y_test = y_data[train_index], y_data[test_index]
            # print(X_test)
            # print(y_test)

            svc = svm.SVC(kernel='linear')
            # svc = svm.SVC(kernel='rbf')
            # svc = svm.SVC(kernel='poly', degree=3)
            clf = svc.fit(X_train, y_train)
            preds = svc.predict(X_test)

            print(y_test, preds)

            cm = confusion_matrix(y_test, preds, labels=[0, 1])

            scores = clf.score(X_test, y_test)

            #print(test_index, y_test, np.array(preds))
            cm_total += cm
            print(cm)
            print(cm_total)
    print(cm_total/1000)

    '''
    ax=[]
    fig = plt.figure()
    print(np.shape(X_data))
    for trial_set in range(num_pairs):
        ax.append(trial_set)
        print("trial set num: ", trial_set)

        ax[-1] = fig.add_subplot(4, 5, trial_set+1)

        ax[-1].plot(np.arange(1000)/200., X_data[trial_set+num_pairs, :], 'r-', linewidth=0.5)
        ax[-1].plot(np.arange(1000)/200., X_data[trial_set, :], 'b-', linewidth=0.5)

        ax[-1].set_ylim(26.0, 30.0)
        ax[-1].set_xlim(0, 5)
        ax[-1].grid()

    plt.tight_layout()
    fig = plt.gcf()
    plt.show()
    '''

def get_overlapping_reference(x_full, data_dict, num_referenced):

    matching_criteria = []

    #get set of all possible pairs of RMSE scores
    for trial_set in range(len(TRIAL_ORDER)):
        optim_set = data_dict['trial_set_' + str(trial_set + 1) + '_X']
        print("trial set 1 shape", optim_set.shape)
        for block in [0, 1, 2]:
            present_set = optim_set[block, :]
            for other_trial_set in range(len(TRIAL_ORDER)):
                if other_trial_set != trial_set:
                    for CW_block in [6, 7, 8]:
                        check_set = data_dict['trial_set_' + str(other_trial_set + 1) + '_X'][CW_block, :]
                        #print(present_set)
                        #print(check_set)
                        matching_criteria.append([trial_set, block, other_trial_set, CW_block, np.sum(np.square(present_set - check_set))])

    matching_criteria = np.array(matching_criteria)
    matching_criteria = matching_criteria[matching_criteria[:,4].argsort()]
    #matching_criteria = np.flipud(matching_criteria)

    referenced_pair = matching_criteria[0, :]
    print(referenced_pair)

    #get set of nearest metals to the metal
    matching_M_criteria = []
    M_trial_set = int(referenced_pair[0])
    M_block = int(referenced_pair[1])
    present_set = data_dict['trial_set_' + str(M_trial_set + 1) + '_X'][M_block, :]
    for other_trial_set in range(len(TRIAL_ORDER)):
        for other_block in [0, 1, 2]:
            if other_trial_set == M_trial_set and other_block == M_block: pass
            else:
                check_set = data_dict['trial_set_' + str(other_trial_set + 1) + '_X'][other_block, :]
                matching_M_criteria.append([other_trial_set, other_block, np.sum(np.square(present_set - check_set))])
    matching_M_criteria = np.array(matching_M_criteria)
    matching_M_criteria = matching_M_criteria[matching_M_criteria[:,2].argsort()]

    print(matching_M_criteria[0:num_referenced-1, :])

    #get set of nearest cold woods to the cold wood
    matching_CW_criteria = []
    CW_trial_set = int(referenced_pair[2])
    CW_block = int(referenced_pair[3])
    present_set = data_dict['trial_set_' + str(CW_trial_set + 1) + '_X'][CW_block, :]
    for other_trial_set in range(len(TRIAL_ORDER)):
        for other_block in [6, 7, 8]:
            if other_trial_set == CW_trial_set and other_block == CW_block: pass
            else:
                check_set = data_dict['trial_set_' + str(other_trial_set + 1) + '_X'][other_block, :]
                matching_CW_criteria.append([other_trial_set, other_block, np.sum(np.square(present_set - check_set))])
    matching_CW_criteria = np.array(matching_CW_criteria)
    matching_CW_criteria = matching_CW_criteria[matching_CW_criteria[:,2].argsort()]

    print(matching_CW_criteria[0:num_referenced-1, :])

    print(matching_criteria.astype(int))

    #Use matching criteria indexing to create data subsets that we can use for LOOCV
    X_data = []
    y_data = []
    trial_set = int(referenced_pair[0])
    block = int(referenced_pair[1])
    X_data.append(data_dict['trial_set_' + str(trial_set + 1) + '_X'][block, :])
    y_data.append(0)
    for ct in range(num_referenced-1):
        trial_set = int(matching_M_criteria[ct, 0])
        block = int(matching_M_criteria[ct, 1])
        X_data.append(data_dict['trial_set_' + str(trial_set + 1) + '_X'][block, :])
        y_data.append(0)

    trial_set = int(referenced_pair[2])
    block = int(referenced_pair[3])
    X_data.append(data_dict['trial_set_' + str(trial_set + 1) + '_X'][block, :])
    y_data.append(1)
    for ct in range(num_referenced-1):
        trial_set = int(matching_CW_criteria[ct, 0])
        block = int(matching_CW_criteria[ct, 1])
        X_data.append(data_dict['trial_set_' + str(trial_set + 1) + '_X'][block, :])
        y_data.append(1)

    X_data = np.array(X_data)
    y_data = np.array(y_data)

    print(y_data)

    # print(len(skf)
    # for train_index, test_index in skf:

    cm_total = np.zeros((2,2))

    #loo = LeaveOneOut(num_pairs)
    #for train_index, test_index in loo:

    for i in range(1000):

        skf = StratifiedKFold(y_data, n_folds=3, shuffle=True)
        for train_index, test_index in skf:
            print("TRAIN:", train_index, "TEST:", test_index)

            print(np.shape(X_data))

            X_train, X_test = X_data[train_index], X_data[test_index]
            print(np.shape(X_train))

            y_train, y_test = y_data[train_index], y_data[test_index]
            # print(X_test)
            # print(y_test)

            svc = svm.SVC(kernel='linear')
            # svc = svm.SVC(kernel='rbf')
            # svc = svm.SVC(kernel='poly', degree=3)
            clf = svc.fit(X_train, y_train)
            preds = svc.predict(X_test)

            print(y_test, preds)

            cm = confusion_matrix(y_test, preds, labels=[0, 1])

            scores = clf.score(X_test, y_test)

            #print(test_index, y_test, np.array(preds))
            cm_total += cm
            print(cm)
            print(cm_total)
    print(cm_total/1000)
    #[[6.234 3.766]
    # [3.221 6.779]]

    '''
    ax=[]
    fig = plt.figure()
    print(np.shape(X_data))
    for trial_set in range(num_pairs):
        ax.append(trial_set)
        print("trial set num: ", trial_set)

        ax[-1] = fig.add_subplot(4, 5, trial_set+1)

        ax[-1].plot(np.arange(1000)/200., X_data[trial_set+num_pairs, :], 'r-', linewidth=0.5)
        ax[-1].plot(np.arange(1000)/200., X_data[trial_set, :], 'b-', linewidth=0.5)

        ax[-1].set_ylim(26.0, 30.0)
        ax[-1].set_xlim(0, 5)
        ax[-1].grid()

    plt.tight_layout()
    fig = plt.gcf()
    plt.show()
    '''