#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""hyperscreen is an improved background rejection algorithm
for """

from __future__ import division
from __future__ import print_function

import colorama
from tqdm import tqdm as progressbar
import sys
import os

import warnings

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm

from astropy.io import fits
from astropy.table import Table

import skimage.data as data
import skimage.segmentation as seg
import skimage.filters as filters
import skimage.draw as draw
import skimage.color as color

import pandas as pd
import numpy as np
np.seterr(divide='ignore')

colorama.init()

class HRCevt1:
    """This is a conceptual class representation of a Chandra High Resolution Camera (HRC) Level 1 Event File

    :return: HRCevt1 object
    :rtype: pandas.DataFrame or astropy.table.table.Table
    """

    def __init__(self, evt1file, verbose=False, as_astropy_table=False):
        """The constructor method for the HRCevt1 class

        :param evt1file: A .fits (or fits.gz) file containing the level 1 event list. If downloaded from the Chandra database, this file always has a *evt1.fits extension. This event list includes all events telemetered.
        :type evt1file: .fits or .fits.gz
        :param verbose: Set verbose=True to make the constructor chatty on the command line, defaults to False
        :type verbose: bool, optional
        :param as_astropy_table: Set as_astropy_table to True in order to have the HRCevt1 constructor method return an Astropy Table object, rather than a Pandas DataFrame. Defaults to False.
        :type as_astropy_table: bool, optional
        """

        # Define how chatty to be
        self.verbose = verbose

        if self.verbose is True:
            print(colorama.Fore.BLUE + '\nParsing HRC EVT1 file...', end=" ")
        # Do a standard read in of the EVT1 fits table
        self.filename = evt1file
        self.hdulist = fits.open(evt1file)
        self.data = Table(self.hdulist[1].data)
        self.header = self.hdulist[1].header
        self.gti = self.hdulist[2].data
        self.hdulist.close()  # Don't forget to close your fits file!

        # Make sure the user isn't running this on an ACIS observation!
        if self.header["DETNAM"][:4] == 'ACIS':
            raise Exception(
                "ERROR: HRCevt1 objects can only be initialized for Chandra/HRC observations. This is a Chandra/ACIS observation.")

        # Populate the fp, fb values for ever event
        if self.verbose is True:
            print(colorama.Fore.BLUE + 'Calculating fp, fb values...', end=" ")
        fp_u, fb_u, fp_v, fb_v = self.calculate_fp_fb()

        # Populate the fp, fb values for ever event
        if self.verbose is True:
            print(colorama.Fore.BLUE + 'Applying GTI mask... ', end=" ")
        self.gti.starts = self.gti['START']
        self.gti.stops = self.gti['STOP']

        self.gtimask = (self.data["time"] > self.gti.starts[0]) & (
            self.data["time"] < self.gti.stops[-1])

        # Populate the fp, fb values for every event
        if self.verbose is True:
            print(colorama.Fore.BLUE + 'Populating metadata columns...', end=" ")
        self.data["fp_u"] = fp_u
        self.data["fb_u"] = fb_u
        self.data["fp_v"] = fp_v
        self.data["fb_v"] = fb_v

        # Make individual status bit columns with legible names
        self.data["AV3 corrected for ringing"] = self.data["status"][:, 0]
        self.data["AU3 corrected for ringing"] = self.data["status"][:, 1]
        self.data["Event impacted by prior event (piled up)"] = self.data["status"][:, 2]
        # Bit 4 (Python 3) is spare
        self.data["Shifted event time"] = self.data["status"][:, 4]
        self.data["Event telemetered in NIL mode"] = self.data["status"][:, 5]
        self.data["V axis not triggered"] = self.data["status"][:, 6]
        self.data["U axis not triggered"] = self.data["status"][:, 7]
        self.data["V axis center blank event"] = self.data["status"][:, 8]
        self.data["U axis center blank event"] = self.data["status"][:, 9]
        self.data["V axis width exceeded"] = self.data["status"][:, 10]
        self.data["U axis width exceeded"] = self.data["status"][:, 11]
        self.data["Shield PMT active"] = self.data["status"][:, 12]
        # Bit 14 (Python 13) is hardware spare
        self.data["Upper level discriminator not exceeded"] = self.data["status"][:, 14]
        self.data["Lower level discriminator not exceeded"] = self.data["status"][:, 15]
        self.data["Event in bad region"] = self.data["status"][:, 16]
        self.data["Amp total on V or U = 0"] = self.data["status"][:, 17]
        self.data["Incorrect V center"] = self.data["status"][:, 18]
        self.data["Incorrect U center"] = self.data["status"][:, 19]
        self.data["PHA ratio test failed"] = self.data["status"][:, 20]
        self.data["Sum of 6 taps = 0"] = self.data["status"][:, 21]
        self.data["Grid ratio test failed"] = self.data["status"][:, 22]
        self.data["ADC sum on V or U = 0"] = self.data["status"][:, 23]
        self.data["PI exceeding 255"] = self.data["status"][:, 24]
        self.data["Event time tag is out of sequence"] = self.data["status"][:, 25]
        self.data["V amp flatness test failed"] = self.data["status"][:, 26]
        self.data["U amp flatness test failed"] = self.data["status"][:, 27]
        self.data["V amp saturation test failed"] = self.data["status"][:, 28]
        self.data["U amp saturation test failed"] = self.data["status"][:, 29]
        self.data["V hyperbolic test failed"] = self.data["status"][:, 30]
        self.data["U hyperbolic test failed"] = self.data["status"][:, 31]
        self.data["Hyperbola test passed"] = np.logical_not(np.logical_or(
            self.data['U hyperbolic test failed'], self.data['V hyperbolic test failed']))
        self.data["Hyperbola test failed"] = np.logical_or(
            self.data['U hyperbolic test failed'], self.data['V hyperbolic test failed'])

        self.obsid = self.header["OBS_ID"]
        self.obs_date = self.header["DATE"]
        self.target = self.header["OBJECT"]
        self.detector = self.header["DETNAM"]
        self.grating = self.header["GRATING"]
        self.exptime = self.header["EXPOSURE"]

        self.numevents = len(self.data["time"])
        self.goodtimeevents = len(self.data["time"][self.gtimask])
        self.badtimeevents = self.numevents - self.goodtimeevents

        self.hyperbola_passes = np.sum(np.logical_or(
            self.data['U hyperbolic test failed'], self.data['V hyperbolic test failed']))
        self.hyperbola_failures = np.sum(np.logical_not(np.logical_or(
            self.data['U hyperbolic test failed'], self.data['V hyperbolic test failed'])))

        if self.hyperbola_passes + self.hyperbola_failures != self.numevents:
            warnings.warn("Number of Hyperbola Test Failures and Passes ({}) does not equal total number of events ({}).".format(
                self.hyperbola_passes + self.hyperbola_failures, self.numevents))

        if self.verbose is True:
            print(colorama.Fore.GREEN + 'Done')

        if self.verbose is True:
            if as_astropy_table is True:
                read_type = "Astropy Table"
            else:
                read_type = "Pandas DataFrame"
            print(colorama.Fore.CYAN + 'Observation Details: ')
            print(colorama.Fore.CYAN + 'ObsID {}  |    {}    |    {}      |    {} ksec     |      {:,} counts (level 1 events)'.format(self.obsid,
                                                                                self.target, self.detector, np.round(self.exptime/1000, 2), self.numevents))
            print(colorama.Fore.RED + '\nConverting EVT1 file to {}...'.format(read_type), end=" ")

        if as_astropy_table is False:
            # Multidimensional columns don't grok with Pandas
            self.data.remove_column('status')
            self.data = self.data.to_pandas()

        if self.verbose is True:
            print(colorama.Fore.GREEN + 'Done')

    def __str__(self):
        """This method returns the string representation of the HRCevt1 object. It is called when the print() or str() function is invoked on an HRCevt1 object.

        :return: A string describing the HRCevt1 object
        :rtype: str
        """
        return "HRC EVT1 object with {} events. Data is packaged as a Pandas Dataframe (or an Astropy Table if as_astropy_table=True on initialization.)".format(self.numevents)

    def calculate_fp_fb(self):
        """Method to calculate the Fine Position (f_p) and normalized central tap amplitude (fb) for the HRC U- and V- axes.

        :return: fp_u, fb_u, fp_v, fb_v; the calculated fine positions and normalized central tap amplitudes, respectively, for the HRC U- and V- axes of the I or S detector
        :rtype: float
        """

        a_u = self.data["au1"]  # otherwise known as "a1"
        b_u = self.data["au2"]  # "a2"
        c_u = self.data["au3"]  # "a3"

        a_v = self.data["av1"]
        b_v = self.data["av2"]
        c_v = self.data["av3"]

        with np.errstate(invalid='ignore'):
            # Do the U axis
            fp_u = ((c_u - a_u) / (a_u + b_u + c_u))
            fb_u = b_u / (a_u + b_u + c_u)

            # Do the V axis
            fp_v = ((c_v - a_v) / (a_v + b_v + c_v))
            fb_v = b_v / (a_v + b_v + c_v)

        return fp_u, fb_u, fp_v, fb_v

    def threshold(self, img, bins, softening=None):
        """HyperScreen (a) separates events by both axis and tap, (b) creates an
        image (a 2D histogram, with bin sizes dependent on the total number of counts
        in the observation), then (c) performs efficient image segmentation on that image
        by applying Otsu's Method (https://en.wikipedia.org/wiki/Otsu%27s_method) to
        create a bespoke threshold for each tap image. This efficiently separates the 'boomerang'
        locus for 'real' events with the 'everything else' region for background events. This function
        performs operation (c).

        Arguments:
            img {[type]} -- [description]
            bins {[type]} -- [description]
        """

        # You don't want to be verbose in this function; it's called many times

        thresh_img = img.copy()
        thresh_img[img == 0] = np.nan

        if softening is None:
            thresh = filters.threshold_otsu(img)
        elif isinstance(softening, float):
            otsu_thresh = filters.threshold_otsu(img)
            thresh = otsu_thresh - (otsu_thresh * softening)

        # "If you don't ignore the warning, you'll get a warning" ~~ G. Tremblay
        with np.errstate(invalid='ignore'):
            thresh_img[thresh_img < thresh] = np.nan
        thresh_img[:int(bins[1] / 2), :] = np.nan
    #     thresh_img[:,int(bins[1]-5):] = np.nan

        return thresh_img

    def hyperscreen(self, softening=1.0):
        """[summary]

        Returns:
            [type] -- [description]
        """

        data = self.data[self.data['Hyperbola test passed']]

        # taprange = range(data['crsu'].min(), data['crsu'].max() + 1)
        taprange_u = range(data['crsu'].min() - 1, data['crsu'].max() + 1)
        taprange_v = range(data['crsv'].min() - 1, data['crsv'].max() + 1)

        if self.numevents < 100000:
            bins = [50, 50]  # number of bins
        else:
            bins = [200, 200]

        # Instantiate these empty dictionaries to hold our results
        u_axis_survivals = {}
        v_axis_survivals = {}

        if self.verbose is False:
            progressbar_disable = True
        elif self.verbose is True:
            progressbar_disable = False

        if self.verbose is True:
            print(colorama.Fore.YELLOW + "\nApplying Otsu's Method to every Tap-specific boomerang across U-axis taps {} through {}".format(taprange_u[0] + 1, taprange_u[-1] + 1))

        skiptaps_u = []
        skiptaps_v = []

        for tap in progressbar(taprange_u, disable=progressbar_disable, ascii=False):
            # Do the U axis
            tapmask_u = data[data['crsu'] == tap].index.values
            if len(tapmask_u) < 20:
                skiptaps_u.append((tap + 1, len(tapmask_u)))
                continue
            keep_u = np.isfinite(data['fb_u'][tapmask_u])

            hist_u, xbounds_u, ybounds_u = np.histogram2d(
                data['fb_u'][tapmask_u][keep_u], data['fp_u'][tapmask_u][keep_u], bins=bins)
            thresh_hist_u = self.threshold(
                hist_u, bins=bins, softening=softening)

            posx_u = np.digitize(data['fb_u'][tapmask_u], xbounds_u)
            posy_u = np.digitize(data['fp_u'][tapmask_u], ybounds_u)
            hist_mask_u = (posx_u > 0) & (posx_u <= bins[0]) & (
                posy_u > -1) & (posy_u <= bins[1])

            # Values of the histogram where the points are
            hhsub_u = thresh_hist_u[posx_u[hist_mask_u] -
                                    1, posy_u[hist_mask_u] - 1]
            pass_fb_u = data['fb_u'][tapmask_u][hist_mask_u][np.isfinite(
                hhsub_u)]

            u_axis_survivals["U Axis Tap {:02d}".format(
                tap)] = pass_fb_u.index.values

        if self.verbose is True:
            print("\nThe following {} U-axis taps were skipped due to a (very) low number of counts: ".format(len(skiptaps_u)))
            for skipped_tap in skiptaps_u:
                tapnum, counts = skipped_tap
                print("Skipped U-axis Tap {}, which had {} count(s)".format(tapnum, counts))
            print(colorama.Fore.MAGENTA + "\n... doing the same for the V axis taps {} through {}".format(taprange_v[0] + 1, taprange_v[-1] + 1))

        for tap in progressbar(taprange_v, disable=progressbar_disable, ascii=False):
            # Now do the V axis:
            tapmask_v = data[data['crsv'] == tap].index.values
            if len(tapmask_v) < 20:
                skiptaps_v.append((tap + 1, len(tapmask_v)))
                continue
            keep_v = np.isfinite(data['fb_v'][tapmask_v])

            hist_v, xbounds_v, ybounds_v = np.histogram2d(
                data['fb_v'][tapmask_v][keep_v], data['fp_v'][tapmask_v][keep_v], bins=bins)
            thresh_hist_v = self.threshold(
                hist_v, bins=bins, softening=softening)

            posx_v = np.digitize(data['fb_v'][tapmask_v], xbounds_v)
            posy_v = np.digitize(data['fp_v'][tapmask_v], ybounds_v)
            hist_mask_v = (posx_v > 0) & (posx_v <= bins[0]) & (
                posy_v > -1) & (posy_v <= bins[1])

            # Values of the histogram where the points are
            hhsub_v = thresh_hist_v[posx_v[hist_mask_v] -
                                    1, posy_v[hist_mask_v] - 1]
            pass_fb_v = data['fb_v'][tapmask_v][hist_mask_v][np.isfinite(
                hhsub_v)]

            v_axis_survivals["V Axis Tap {:02d}".format(
                tap)] = pass_fb_v.index.values

        if self.verbose is True:
            print("\nThe following {} V-axis taps were skipped due to a (very) low number of counts: ".format(len(skiptaps_v)))
            for skipped_tap in skiptaps_v:
                tapnum, counts = skipped_tap
                print("Skipped V-axis Tap {}, which had {} count(s)".format(tapnum, counts))

        # Done looping over taps

        if self.verbose is True:
            print(colorama.Fore.BLUE + "\nCollecting events that pass both U- and V-axis HyperScreen tests...", end=" ")

        u_all_survivals = np.concatenate(
            [x for x in u_axis_survivals.values()])
        v_all_survivals = np.concatenate(
            [x for x in v_axis_survivals.values()])

        # If the event passes both U- and V-axis tests, it survives
        all_survivals = np.intersect1d(u_all_survivals, v_all_survivals)
        survival_mask = np.isin(self.data.index.values, all_survivals)
        failure_mask = np.logical_not(survival_mask)

        num_survivals = sum(survival_mask)
        num_failures = sum(failure_mask)

        percent_hyperscreen_rejected = round(
            ((num_failures / self.numevents) * 100), 2)

        # Do a sanity check to look for lost events. Shouldn't be any.
        if num_survivals + num_failures != self.numevents:
            print("WARNING: Total Number of survivals and failures does \
            not equal total events in the EVT1 file. Something is wrong!")

        legacy_hyperbola_test_failures = sum(
            self.data['Hyperbola test failed'])
        percent_legacy_hyperbola_test_rejected = round(
            ((legacy_hyperbola_test_failures / self.numevents) * 100), 2)

        percent_improvement_over_legacy_test = round(
            (percent_hyperscreen_rejected - percent_legacy_hyperbola_test_rejected), 2)

        if self.verbose is True:
            print("Done")
            print(colorama.Fore.GREEN + "HyperScreen rejected" + colorama.Fore.YELLOW + " {}% of all events ({:,} bad events / {:,} total events)".format(percent_hyperscreen_rejected, sum(failure_mask), self.numevents) + colorama.Fore.GREEN +
                  "\nThe Murray+ algorithm rejects" + colorama.Fore.MAGENTA + " {}% of all events ({:,} bad events / {:,} total events)".format(percent_legacy_hyperbola_test_rejected, legacy_hyperbola_test_failures, self.numevents))

            print(colorama.Fore.GREEN + "As long as the results pass sanity checks, this is a POTENTIAL improvement of \n" +
                  colorama.Fore.BLUE + "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ POTENTIAL Improvement ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n" +
                  colorama.Fore.WHITE + "                                      {}%\n".format(percent_improvement_over_legacy_test) +
                  colorama.Fore.BLUE + "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n")

        hyperscreen_results_dict = {"ObsID": self.obsid,
                                    "Target": self.target,
                                    "Exposure Time": self.exptime,
                                    "Detector": self.detector,
                                    "Number of Events": self.numevents,
                                    "Number of Good Time Events": self.goodtimeevents,
                                    "U Axis Survivals by Tap": u_axis_survivals,
                                    "V Axis Survivals by Tap": v_axis_survivals,
                                    "U Axis All Survivals": u_all_survivals,
                                    "V Axis All Survivals": v_all_survivals,
                                    "All Survivals (event indices)": all_survivals,
                                    "All Survivals (boolean mask)": survival_mask,
                                    "All Failures (boolean mask)": failure_mask,
                                    "Percent rejected by Tapscreen": percent_hyperscreen_rejected,
                                    "Percent rejected by Hyperbola": percent_legacy_hyperbola_test_rejected,
                                    "Percent improvement": percent_improvement_over_legacy_test
                                    }

        return hyperscreen_results_dict

    def hyperbola(self, fb, a, b, h):
        """Given the normalized central tap amplitude, a, b, and h,
        return an array of length len(fb) that gives a hyperbola.

        Arguments:
            fb {[type]} -- [description]
            a {[type]} -- [description]
            b {[type]} -- [description]
            h {[type]} -- [description]

        Returns:
            [type] -- [description]
        """

        hyperbola = b * np.sqrt(((fb - h)**2 / a**2) - 1)

        return hyperbola

    def legacy_hyperbola_test(self, tolerance=0.035):
        """[summary]

        Keyword Arguments:
            tolerance {float} -- [description] (default: {0.035})

        Returns:
            [type] -- [description]
        """

        # Remind the user what tolerance they're using
        # print("{0: <25}| Using tolerance = {1}".format(" ", tolerance))

        # Set hyperbolic coefficients, depending on whether this is HRC-I or -S
        if self.detector == "HRC-I":
            a_u = 0.3110
            b_u = 0.3030
            h_u = 1.0580

            a_v = 0.3050
            b_v = 0.2730
            h_v = 1.1
            # print("{0: <25}| Using HRC-I hyperbolic coefficients: ".format(" "))
            # print("{0: <25}|    Au={1}, Bu={2}, Hu={3}".format(" ", a_u, b_u, h_u))
            # print("{0: <25}|    Av={1}, Bv={2}, Hv={3}".format(" ", a_v, b_v, h_v))

        if self.detector == "HRC-S":
            a_u = 0.2706
            b_u = 0.2620
            h_u = 1.0180

            a_v = 0.2706
            b_v = 0.2480
            h_v = 1.0710
            # print("{0: <25}| Using HRC-S hyperbolic coefficients: ".format(" "))
            # print("{0: <25}|    Au={1}, Bu={2}, Hu={3}".format(" ", a_u, b_u, h_u))
            # print("{0: <25}|    Av={1}, Bv={2}, Hv={3}".format(" ", a_v, b_v, h_v))

        # Set the tolerance boundary ("width" of the hyperbolic region)

        h_u_lowerbound = h_u * (1 + tolerance)
        h_u_upperbound = h_u * (1 - tolerance)

        h_v_lowerbound = h_v * (1 + tolerance)
        h_v_upperbound = h_v * (1 - tolerance)

        # Compute the Hyperbolae
        with np.errstate(invalid='ignore'):
            zone_u_fit = self.hyperbola(self.data["fb_u"], a_u, b_u, h_u)
            zone_u_lowerbound = self.hyperbola(
                self.data["fb_u"], a_u, b_u, h_u_lowerbound)
            zone_u_upperbound = self.hyperbola(
                self.data["fb_u"], a_u, b_u, h_u_upperbound)

            zone_v_fit = self.hyperbola(self.data["fb_v"], a_v, b_v, h_v)
            zone_v_lowerbound = self.hyperbola(
                self.data["fb_v"], a_v, b_v, h_v_lowerbound)
            zone_v_upperbound = self.hyperbola(
                self.data["fb_v"], a_v, b_v, h_v_upperbound)

        zone_u = [zone_u_lowerbound, zone_u_upperbound]
        zone_v = [zone_v_lowerbound, zone_v_upperbound]

        # Apply the masks
        # print("{0: <25}| Hyperbolic masks for U and V axes computed".format(""))

        with np.errstate(invalid='ignore'):
            # print("{0: <25}| Creating U-axis mask".format(""), end=" |")
            between_u = np.logical_not(np.logical_and(
                self.data["fp_u"] < zone_u[1], self.data["fp_u"] > -1 * zone_u[1]))
            not_beyond_u = np.logical_and(
                self.data["fp_u"] < zone_u[0], self.data["fp_u"] > -1 * zone_u[0])
            condition_u_final = np.logical_and(between_u, not_beyond_u)

            # print(" Creating V-axis mask")
            between_v = np.logical_not(np.logical_and(
                self.data["fp_v"] < zone_v[1], self.data["fp_v"] > -1 * zone_v[1]))
            not_beyond_v = np.logical_and(
                self.data["fp_v"] < zone_v[0], self.data["fp_v"] > -1 * zone_v[0])
            condition_v_final = np.logical_and(between_v, not_beyond_v)

        mask_u = condition_u_final
        mask_v = condition_v_final

        hyperzones = {"zone_u_fit": zone_u_fit,
                      "zone_u_lowerbound": zone_u_lowerbound,
                      "zone_u_upperbound": zone_u_upperbound,
                      "zone_v_fit": zone_v_fit,
                      "zone_v_lowerbound": zone_v_lowerbound,
                      "zone_v_upperbound": zone_v_upperbound}

        hypermasks = {"mask_u": mask_u, "mask_v": mask_v}

        # print("{0: <25}| Hyperbolic masks created".format(""))
        # print("{0: <25}| ".format(""))
        return hyperzones, hypermasks

    def boomerang(self, mask=None, show=True, plot_legacy_zone=True, title=None, cmap=None, savepath=None, create_subplot=False, ax=None, rasterized=True):
        """[summary]
        """
        # You can plot the image on axes of a subplot by passing
        # that axis to this function. Here are some switches to enable that.

        if create_subplot is False:
            self.fig, self.ax = plt.subplots(figsize=(12, 8))
        elif create_subplot is True:
            if ax is None:
                self.ax = plt.gca()
            else:
                self.ax = ax

        if cmap is None:
            cmap = 'plasma'

        if mask is not None:
            self.ax.scatter(self.data['fb_u'], self.data['fp_u'],
                            c=self.data['sumamps'], cmap='bone', s=0.3, alpha=0.8, rasterized=rasterized, label='All Events')

            frame = self.ax.scatter(self.data['fb_u'][mask], self.data['fp_u'][mask],
                                    c=self.data['sumamps'][mask], cmap=cmap, s=0.5, rasterized=rasterized, label='HyperScreen Accepted Events')

        else:
            frame = self.ax.scatter(self.data['fb_u'], self.data['fp_u'],
                                    c=self.data['sumamps'], cmap=cmap, s=0.5, rasterized=rasterized)

        if plot_legacy_zone is True:
            hyperzones, hypermasks = self.legacy_hyperbola_test(
                tolerance=0.035)
            self.ax.plot(self.data["fb_u"], hyperzones["zone_u_lowerbound"],
                         'o', markersize=0.3, color='black', alpha=0.8, rasterized=rasterized, label='Murray Exclusion Hyperbola')
            self.ax.plot(self.data["fb_u"], -1 * hyperzones["zone_u_lowerbound"],
                         'o', markersize=0.3, color='black', alpha=0.8, rasterized=rasterized)

            self.ax.plot(self.data["fb_u"], hyperzones["zone_u_upperbound"],
                         'o', markersize=0.3, color='black', alpha=0.8, rasterized=rasterized)
            self.ax.plot(self.data["fb_u"], -1 * hyperzones["zone_u_upperbound"],
                         'o', markersize=0.3, color='black', alpha=0.8, rasterized=rasterized)

        self.ax.grid(False)

        if title is None:
            self.ax.set_title('{} | {} | ObsID {} | {} ksec | {:,} counts'.format(
                self.target, self.detector, self.obsid, round(self.exptime / 1000, 1), self.numevents))
        else:
            self.ax.set_title(title)

        self.ax.set_ylim(-1.1, 1.1)
        self.ax.set_xlim(-0.1, 1.1)

        self.ax.set_ylabel(r'Fine Position $f_p$ $(C-A)/(A + B + C)$')
        self.ax.set_xlabel(
            r'Normalized Central Tap Amplitude $f_b$ $B / (A+B+C)$')

        if create_subplot is False:
            self.cbar = plt.colorbar(frame, pad=-0.005)
            self.cbar.set_label("SUMAMPS")

        if show is True:
            plt.show()

        if savepath is not None:
            plt.savefig(savepath, dpi=150, bbox_inches='tight')

        plt.close()

    def image(self, masked_x=None, masked_y=None, xlim=None, ylim=None, detcoords=False, title=None, cmap=None, show=True, rasterized=True, savepath=None, create_subplot=False, ax=None, nbins=(400, 400)):
        """Create a quicklook image, in detector or sky coordinates, of the
        observation. The image will be binned to 400x400 by default.

        Keyword Arguments:
            masked_x {[type]} -- [description] (default: {None})
            masked_y {[type]} -- [description] (default: {None})
            xlim {[type]} -- [description] (default: {None})
            ylim {[type]} -- [description] (default: {None})
            detcoords {bool} -- [description] (default: {False})
            title {[type]} -- [description] (default: {None})
            cmap {[type]} -- [description] (default: {None})
            show {bool} -- [description] (default: {True})
            rasterized {bool} -- [description] (default: {True})
            savepath {[type]} -- [description] (default: {None})
            create_subplot {bool} -- [description] (default: {False})
            ax {[type]} -- [description] (default: {None})
        """

        if masked_x is not None and masked_y is not None:
            x = masked_x
            y = masked_y
            img_data, yedges, xedges = np.histogram2d(y, x, nbins)
        else:
            if detcoords is False:
                x = self.data['x'][self.gtimask]
                y = self.data['y'][self.gtimask]
            elif detcoords is True:
                x = self.data['detx'][self.gtimask]
                y = self.data['dety'][self.gtimask]
            img_data, yedges, xedges = np.histogram2d(y, x, nbins)

        extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]

        # Create the Figure
        styleplots()

        # You can plot the image on axes of a subplot by passing
        # that axis to this function. Here are some switches to enable that.
        if create_subplot is False:
            self.fig, self.ax = plt.subplots()
        elif create_subplot is True:
            if ax is None:
                self.ax = plt.gca()
            else:
                self.ax = ax

        self.ax.grid(False)

        if cmap is None:
            cmap = 'viridis'

        self.ax.imshow(img_data, extent=extent, norm=LogNorm(),
                       interpolation=None, rasterized=rasterized, cmap=cmap, origin='lower')

        if title is None:
            self.ax.set_title("ObsID {} | {} | {} | {:,} events".format(
                self.obsid, self.target, self.detector, self.goodtimeevents))
        else:
            self.ax.set_title("{}".format(title))
        if detcoords is False:
            self.ax.set_xlabel("Sky X")
            self.ax.set_ylabel("Sky Y")
        elif detcoords is True:
            self.ax.set_xlabel("Detector X")
            self.ax.set_ylabel("Detector Y")

        if xlim is not None:
            self.ax.set_xlim(xlim)
        if ylim is not None:
            self.ax.set_ylim(ylim)

        if show is True:
            plt.show(block=True)

        if savepath is not None:
            plt.savefig('{}'.format(savepath))
            print("Saved image to {}".format(savepath))

        plt.close()

def styleplots():  # pragma: no cover
    """Make the plots pretty.
    """

    mpl.rcParams['agg.path.chunksize'] = 10000

    # Make things pretty
    plt.style.use('ggplot')

    labelsizes = 10

    plt.rcParams['font.size'] = labelsizes
    plt.rcParams['axes.titlesize'] = 12
    plt.rcParams['axes.labelsize'] = labelsizes
    plt.rcParams['xtick.labelsize'] = labelsizes
    plt.rcParams['ytick.labelsize'] = labelsizes


"""Photometric redshift analysis. Includes a wrapper to LEPHARE and BPZ.

- LEPHARE: http://www.cfht.hawaii.edu/~arnouts/LEPHARE/lephare.html
- BPZ: http://www.stsci.edu/~dcoe/BPZ
"""

from __future__ import print_function
import os
import sys
import subprocess
import numpy as N
import pylab as P
from scipy.optimize import curve_fit
from astropy.io import ascii
from astropy.table import Table, hstack
from astropy.coordinates import SkyCoord
from . import utils as cutils

class LEPHARE(object):

    """Wrapper to the LEPHARE photometric redshift code.

    http://www.cfht.hawaii.edu/~arnouts/LEPHARE/lephare.html
    """

    def __init__(self, magnitudes, errors, zpara=None, spectro_file=None, **kwargs):
        """
        Run the LEPHARE progam (zphota).

        :param list magnitudes: Magnitudes. A list of list
        :param list errors: Error on magnitudes. Same shape as the magnitude list.
        :param string zpara: default is $LEPHAREDIR/config/zphot_megacam.para
        :param dictionnary files: A dictioannry conating the names of the

        Kwargs include the following list of possible parameters

        :param string basename: Base name for all created files
        :param string cname: Name of the studied cluster. If basename if not given,
         it will be used as base name for all created file
        :param string filters: filter list
        :param list ra: list of ra
        :param list dec: list of dec
        :param list id: list of ID
        """
        self.data = {'mag': magnitudes, 'err': errors}
        self.kwargs = kwargs
        self.config = os.environ["LEPHAREDIR"] + \
            "/config/zphot_megacam.para" if zpara is None else zpara

        self.spectro_file = spectro_file
        # Name of created files?
        prefix = ""
        if 'basename' in kwargs:
            prefix = kwargs['basename'] + "_"
        elif 'cname' in kwargs:
            prefix = kwargs['cname'] + "_"
        self.files = {}
        self.files['input'] = prefix + "zphot.in"
        self.files['output'] = prefix + "zphot.out"
        self.files['pdz_output'] = prefix + "zphot"
        self.files['all_input'] = self.files['input'].replace('.in', '.all')

        # Initialize lephare output variables
        self.lephare_out = None
        self.data_out = None

        self.write_input()

    def write_input(self):
        """
        Create and write files needed to run LEPHARE.

        - the input data file for LEPHARE
        - a similar file containing the sources ID along with their RA DEC.
        """
        f = open(self.files['input'], 'w')
        if self.spectro_file is None:
            # No spectroscopic redshift file provided in config.yaml
            # --> only need the SHORT format for LePhare input file
            if 'filters' in self.kwargs:
                f.write("# id " + " ".join(["mag_%s" % filt for filt in self.kwargs['filters']]) +
                        " " + " ".join(["err_mag_%s" % filt for filt in self.kwargs['filters']]) +
                        "\n")
                for i, mags in enumerate(N.concatenate([self.data['mag'], self.data['err']]).T):
                    f.write("%i %s\n" %
                            (i, " ".join(["%.3f" % m for m in mags])))

                f.close()
                print("INFO: Input data saved in", self.files['input'])
        else:
            # Spectroscopic redshift file provided in config.yaml
            # --> Need to write LePhare input file in the LONG format,
            # i.e with 'context' and 'spectroz' of matching galaxies
            zspec = ZSPEC(self.spectro_file, names=[
                          'object', 'ra', 'dec', 'zspec'])
            zspec.skycoords = SkyCoord(
                zspec.data['ra'], zspec.data['dec'], unit='deg')
            skycoords_cat = SkyCoord(
                self.kwargs['ra'], self.kwargs['dec'], unit='deg')
            idx, d2d, d3d = skycoords_cat.match_to_catalog_sky(zspec.skycoords)
            zp = zspec.data['zspec'][idx]
            # identify galaxies with bad match, i.e. dist > 300 mas
            bad = N.where(d2d.mas > 300)
            zp[bad] = -99
            print("INFO: Using " + str(len(idx) - N.size(bad)) +
                  " galaxies for spectroz training")
            if 'filters' in self.kwargs:
                f.write("# id " + " ".join(["mag_%s" % filt for filt in self.kwargs['filters']]) +
                        " " + " ".join(["err_mag_%s" % filt for filt in self.kwargs['filters']]) +
                        " context" + " zspec" + "\n")
                # tells LePhare to run using the u, g, r, i and z bands.
                context = 31
                for i, mags in enumerate(N.concatenate([self.data['mag'], self.data['err']]).T):
                    f.write("%i %s %s\n" % (i, " ".join(["%.3f" % m for m in mags]),
                                            " ".join(("%i" % context, "%.3f" % zp[i]))))
                f.close()
        if 'ra' in self.kwargs:
            f = open(self.files['all_input'], 'w')
            if 'filters' in self.kwargs is not None:
                f.write("# id ID RA DEC " +
                        " ".join(["mag_%s" % filt for filt in self.kwargs['filters']]) + " " +
                        " ".join(["err_mag_%s" % filt for filt in self.kwargs['filters']]) +
                        "\n")
            for i, mags in enumerate(N.concatenate([self.data['mag'], self.data['err']]).T):
                f.write("%i %i %f %f %s\n" % (i, self.kwargs['id'][i],
                                              self.kwargs['ra'][i], self.kwargs['dec'][i],
                                              " ".join(["%.3f" % m for m in mags])))
            f.close()
            print("INFO: All data saved in", self.files['all_input'])

    def check_config(self, config=None):
        """
        Check that the SED and filters requested for the LePhare run do exist.

        If not: explains where the problem is and aborts.
        """
        if config is not None:
            if os.path.exists(config):
                self.config = config
            else:
                raise ValueError("%s does not exist" % config)
        with open(self.config) as f:
            for line in f:
                if "ZPHOTLIB" in line:
                    libs = line.split()[1].split(",")
        path_to_lib = os.environ["LEPHAREWORK"] + "/lib_mag/"

        for lib in libs:
            lib_tmp = path_to_lib + lib + ".bin"
            if not os.path.exists(lib_tmp):
                print("\nERROR: Requested library %s does not exist " % lib_tmp)
                print("INFO: Available SED libraries are:\n ")
                for f in os.listdir(path_to_lib):
                    if f.endswith(".bin"):
                        print(f)
                        print("--> Correct the ZPHOTLIB variable in %s" % self.config +
                              "or generate the missing LePhare SED libraries")
                sys.exit()

        for counter, lib in enumerate(libs):
            lib_tmp = path_to_lib + lib + ".doc"
            with open(lib_tmp) as f:
                for line in f:
                    if "FILTER_FILE" in line:
                        filt_tmp = line.split()[1]
                        if counter == 0:
                            filt_ref = filt_tmp
                        if filt_tmp != filt_ref:
                            print("\nERROR: Requested SED libraries %s" % libs +
                                  " have been built using different filters. Either change the " +
                                  "requested libraries or re-generate them accordingly.")
                            sys.exit()

        path_to_filt = os.environ["LEPHAREWORK"] + "/filt/"
        if not os.path.exists(path_to_filt + filt_ref):
            print("\nERROR: The FILTER_FILE %s used by the SED libraries %s does not exists." %
                  (path_to_filt + filt_ref, libs))
            sys.exit()

    def run(self, config=None):
        """
        Run LEPHARE.

        Default config file is $LEPHAREDIR/config/zphot_megacam.para.
        Can be overwritten with the config argument
        """
        if config is not None:
            if os.path.exists(config):
                self.config = config
            else:
                raise ValueError("%s does not exist" % config)

        # build command line
        cmd = "zphota"
        cmd += " -c " + self.config
        cmd += " -CAT_IN " + self.files['input']
        cmd += " -CAT_OUT " + self.files['output']
        cmd += " -PDZ_OUT " + self.files['pdz_output']
        print("INFO: Will run '%s'" % cmd)

        self.lephare_out = subprocess.check_output(
            cmd, stderr=subprocess.STDOUT, shell=True)
        print("INFO: LEPHARE output summary (full output in self.lephare_out)")
        print(
            "\n".join(["   " + zo for zo in self.lephare_out.split("\n")[-6:]]))

        self.data_out = ZPHOTO(self.files['output'], self.files['pdz_output'], zcode_name='lephare',
                               all_input=self.files['all_input'], **self.kwargs)

class BPZ(object):

    """
    Wrapper to the BPZ photometric redshift code.

    http://www.stsci.edu/~dcoe/BPZ
    """

    def __init__(self, magnitudes, errors, zpara=None, spectro_file=None, **kwargs):
        """
        Run the BPZ progam (zphota).

        :param list magnitudes: Magnitudes. A list of list
        :param list errors: Error on magnitudes. Same shape as the magnitude list.

        Kwargs include the following list of possible parameters

        :param string basename: Base name for all created files
        :param string cname: Name of the studied cluster. If basename if not given,
                             it will be used as base name for all created file
        :param string filters: filter list
        :param list ra: list of ra
        :param list dec: list of dec
        :param list id: list of ID
        """
        self.data = {'mag': magnitudes, 'err': errors}
        self.kwargs = kwargs
        self.config = zpara
        self.spectro_file = spectro_file

        # Name of created files?
        prefix = ""
        if 'basename' in kwargs:
            prefix = kwargs['basename'] + "_"
        elif 'cname' in kwargs:
            prefix = kwargs['cname'] + "_"
        self.files = {}
        self.files['input'] = prefix + "bpz.in"
        self.files['flux_comparison'] = prefix + "bpz.flux_comparison"
        self.files['output'] = prefix + "bpz.bpz"
        self.files['pdz_output'] = prefix + "bpz.probs"
        self.files['columns'] = prefix + "bpz.columns"
        self.files['all_input'] = self.files['input'].replace('.in', '.all')

        # Initialize BPZ output variables
        self.bpz_out = None
        self.data_out = None

        self.write_input()
        self.build_columns_file()

    def write_input(self):
        """
        Create and write files needed to run BPZ.

        - the input data file for BPZ
        - a similar file containing the sources ID along with their RA DEC.
        """
        f = open(self.files['input'], 'w')
        if 'filters' in self.kwargs:
            f.write("# id " + " ".join(["mag_%s" % filt for filt in self.kwargs['filters']]) +
                    " " + " ".join(["err_mag_%s" % filt for filt in self.kwargs['filters']]) +
                    "\n")
            for i, mags in enumerate(N.concatenate([self.data['mag'], self.data['err']]).T):
                f.write("%i %s\n" % (i, " ".join(["%.3f" % m for m in mags])))

            f.close()
            print("INFO: Input data saved in", self.files['input'])
        if 'ra' in self.kwargs:
            f = open(self.files['all_input'], 'w')
            if 'filters' in self.kwargs is not None:
                f.write("# id ID RA DEC " +
                        " ".join(["mag_%s" % filt for filt in self.kwargs['filters']]) + " " +
                        " ".join(["err_mag_%s" % filt for filt in self.kwargs['filters']]) +
                        "\n")
            for i, mags in enumerate(N.concatenate([self.data['mag'], self.data['err']]).T):
                f.write("%i %i %f %f %s\n" % (i, self.kwargs['id'][i],
                                              self.kwargs['ra'][i], self.kwargs['dec'][i],
                                              " ".join(["%.3f" % m for m in mags])))
            f.close()
            print("INFO: All data saved in", self.files['all_input'])

#    def build_columns_file(self, prefix='HSC-', sufix='',
#                           filters=None, ref='i', z_s=False):
    def build_columns_file(self, prefix='CFHT_megacam_', sufix='p',
                           filters=None, ref='i', z_s=False):
        """Build and write the 'columns' file.

        Hardcoded for test purpose.
        """
        if filters is None:  # set the default values as the CFHT ones
            filters = ['u', 'g', 'r', 'i', 'z']
#            filters = ['g', 'r', 'i', 'z', 'Y']
        f = open(self.files['columns'], 'w')
        f.write("# Filter  columns  AB/Vega  zp_error  zp_offset\n")
        for i, filt in enumerate(filters):
            f.write("%s%s%s     %i, %i   AB        0.01      0.00\n" %
                    (prefix, filt, sufix, i + 2, i + len(filters) + 2))
        if z_s:
            f.write("Z_S                  %i\n" % (len(filters) * 2 + 2))
        f.write("M_0                 %i\n"
                    % [j + 2 for j, filt in enumerate(filters) if filt == ref][0])
        f.write("ID                    1\n")
        f.close()

    def run(self):
        """
        Run BPZ.

        Configuration file must exist in the current directory.

        .. todo:: Build the configuration file on the fly (the .columns)
        """
        if not os.getenv('BPZPATH'):
            mess = "BPZPATH if not defined. Please install BPZ "
            mess += "(see http://www.stsci.edu/~dcoe/BPZ/install.html)."
            raise IOError(mess)

        if not os.path.exists(self.files['columns']):
            raise IOError("%s does not exist" % self.files['columns'])

        # build command line options from param file
        if self.config is not None:
            opt_arr = N.genfromtxt(self.config, dtype=None)
            option = ['-' + opt_arr[i, 0].decode() + ' ' + opt_arr[i, 1].decode()
                      for i in N.arange(len(opt_arr))]
            options = ' '.join(e for e in option)
        else:
            options = ''

        # build command line
#        cmd = "python $BPZPATH/bpz.py %s " % self.files['input'] + options
        cmd = "bpz_run.py %s " % self.files['input'] + options
        print("INFO: Will run '%s'" % cmd)

        self.bpz_out = subprocess.check_output(
            cmd, stderr=subprocess.STDOUT, shell=True)
        print("INFO: BPZ output summary (full output in self.bpz_out)")
#        print("\n".join(["   " + zo for zo in self.bpz_out.split("\n")[:20]]))

        self.data_out = ZPHOTO(self.files['output'], self.files['pdz_output'],
                               zcode_name='bpz', all_input=self.files['all_input'],
                               **self.kwargs)

class ZPHOTO(object):

    """Read photoz code (LePhare, BPZ) output file and creates/saves astropy tables."""

    def __init__(self, zphot_output, zphot_pdz_output, zcode_name=None, all_input=None, **kwargs):
        """Read the photoz progam (LePhare, BPZ, ...) output."""
        self.files = {}
        self.files['output'] = zphot_output
        self.files['pdz_output'] = zphot_pdz_output
        self.kwargs = kwargs
        if all_input is not None:
            self.files['input'] = all_input
            self.read_input()
        self.code = zcode_name
        self.read()

    def read(self):
        """Read the output."""
        f = open(self.files['output'], 'r')
        self.data_array = N.loadtxt(self.files['output'], unpack=True)

        if self.code == 'lephare':
            self.header = [l for l in f if l.startswith('#')]
            f.close()
            self.variables = N.loadtxt(os.getenv('LEPHAREDIR') +
                                       "/config/zphot_output.para", dtype='string')
            self.data_dict = {v: a for v, a in zip(
                self.variables, self.data_array)}
            self.nsources = len(self.data_dict['Z_BEST'])
            self.pdz_zbins = N.loadtxt(
                self.files['pdz_output'] + '.zph', unpack=True)
            self.pdz_val = N.loadtxt(
                self.files['pdz_output'] + '.pdz', unpack=True)

        elif self.code == 'bpz':
            self.header = [l for l in f if l.startswith('##')]
            f.seek(0)
            self.variables = [l[4:].replace(' ', '').split(
                '\n')[0] for l in f if l.startswith('# ')]
            f.seek(0)
            # BPZ does not provide a zbins file.
            # Needs to create it from zmin, zmax and dz specified in output file
            zmin = float([line.split('=')[1]
                          for line in f if 'ZMIN' in line][0])
            f.seek(0)
            zmax = float([line.split('=')[1]
                          for line in f if 'ZMAX' in line][0])
            f.seek(0)
            dz = float([line.split('=')[1] for line in f if 'DZ' in line][0])
            f.close()

            self.data_dict = {v: a for v, a in zip(
                self.variables, self.data_array)}
            self.nsources = len(self.data_dict['Z_B'])
            self.pdz_zbins = N.arange(zmin, zmax + dz, dz)
            self.pdz_val = N.loadtxt(self.files['pdz_output'], unpack=True,
                                     usecols=N.arange(1, len(self.pdz_zbins) + 1))

    def save_zphot(self, file_out, path_output, overwrite=False):
        """Save the output of photoz code (z_best, chi^2, pdz) into astropy table."""
        # Duplicates the zbins vector for each object.
        # It is redundant information but astropy tables need each field to have the
        # same size. Or maybe I'm missing something.
        zbins = N.tile(self.pdz_zbins, (len(self.kwargs['id']), 1))

        # Converts LePhare or BPZ likelihood to actual probability density
        for i in N.arange(len(self.pdz_val.T)):
            norm = N.trapz(self.pdz_val[:, i], self.pdz_zbins)
            new_pdz_val = self.pdz_val[:, i] / norm
            self.pdz_val[:, i] = new_pdz_val

        # Creates astropy table to be saved in path_output of file_out
        new_tab = hstack([Table([self.kwargs['id']]), Table([self.kwargs['ra']]),
                          Table([self.kwargs['dec']]), Table(self.data_dict),
                          Table([zbins], names=['zbins']),
                          Table([self.pdz_val.T], names=['pdz'])],
                         join_type='inner')

        # Rename BPZ Z_B to Z_BEST to match LePhare
        if 'Z_B' in new_tab.keys():
            new_tab.rename_column('Z_B', 'Z_BEST')

        # overwrite keyword of data.write(file,path) does not only overwrites
        # the data in path, but the whole file, i.e. (we lose all other
        # paths in the process) --> overwrite_or_append (see above)

        cutils.overwrite_or_append(file_out, path_output, new_tab, overwrite=overwrite)

        print("INFO: ", self.code, "data saved in", file_out, "as", path_output)

    def read_input(self):
        """Read the input."""
        data = N.loadtxt(self.files['input'], unpack=True)
        f = open(self.files['input'], 'r')
        l = f.readlines()[0]
        self.input_data = {k: d for k, d in zip(l[2:-1].split(), data)}
        f.close()

    def hist(self, param, **kwargs):
        """Plot histograms.

        Possible kwargs

        :params float minv: Lower value of the histogram
        :params float maxv: Upper value of the histogram
        :params int nbins: Number of bins. Default is 10.
        :params string xlabel: An xlbal for the figure
        :params string title: A title for the figure
        :params float zclust: Redshift of the studies cluster
        """
        # Get the data and apply some cut if asked
        pval = self.data_dict[param]
        filt = N.array([1] * len(pval), dtype='bool')
        if 'minv' in kwargs:
            filt &= (pval >= kwargs['minv'])
        if 'maxv' in kwargs:
            filt &= (pval <= kwargs['maxv'])
        pval = pval[filt]

        # Plot the histogram
        fig = P.figure()
        ax = fig.add_subplot(111, ylabel='#')
        ax.hist(pval, bins=kwargs['nbins'] if 'nbins' in kwargs else 10)
        xlabel = kwargs['xlabel'] if 'xlabel' in kwargs else param
        ax.set_xlabel(xlabel)
        if 'title' in kwargs:
            ax.set_title(kwargs['title'])
        if 'zclust' in kwargs:
            ax.axvline(kwargs['zclust'], color='r',
                       label='Cluster redshift (%.4f)' % kwargs['zclust'])
            ax.legend(loc='best')

        # Save the figure
        fig.savefig(self.files['output'].replace(
            '.out', '') + "_" + xlabel + "_zphot_hist.png")

    def plot(self, px, py, **kwargs):
        """
        Plot x vs. y.

        Possible kwargs are:
        :params float minx: lower limit of the x axis
        :params float maxx: upper limit of the x axis
        :params float miny: lower limit of the y axis
        :params float maxy: upper limit of the y axis
        :params string xlabel: label of the x axis
        :params string ylabel: label of the y axis
        :params string title: title of the figure
        """
        pvalx = self.data_dict[px]
        pvaly = self.data_dict[py]
        filt = N.array([1] * len(pvalx), dtype='bool')
        if 'minx' in kwargs and kwargs['minx'] is not None:
            filt &= (pvalx >= kwargs['minx'])
        if 'maxx' in kwargs and kwargs['maxx'] is not None:
            filt &= (pvalx <= kwargs['maxx'])
        if 'miny' in kwargs and kwargs['miny'] is not None:
            filt &= (pvaly >= kwargs['miny'])
        if 'maxy' in kwargs and kwargs['maxy'] is not None:
            filt &= (pvaly <= kwargs['maxy'])
        pvalx, pvaly = pvalx[filt], pvaly[filt]
        fig = P.figure()
        ax = fig.add_subplot(111)
        ax.scatter(pvalx, pvaly)
        if 'xlabel' in kwargs and kwargs['xlabel'] is not None:
            px = kwargs['xlabel']
        if 'ylabel' in kwargs and kwargs['ylabel'] is not None:
            py = kwargs['ylabel']
        ax.set_xlabel(px)
        ax.set_ylabel(py)
        if 'title' in kwargs and kwargs['title'] is not None:
            ax.set_title(kwargs['title'])

        if 'figname' in kwargs and kwargs['figname'] is not None:
            fig.savefig(self.files['output'].replace(
                '.out', '') + "_%s_vs_%s_zphot.png" % (py, px))
        else:
            fig.savefig("%s_vs_%s_zphot.png" % (py, px))

    def plot_map(self, title=None, zmin=0, zmax=999):
        """Plot the redshift sky-map."""
        if not hasattr(self, 'input_data'):
            print("WARNING: No input data given. Cannot plot the redshift map.")
            return

        ra, dec, redshift = self.input_data['RA'], self.input_data['DEC'], self.data_dict['Z_BEST']

        # redshift has to be >= 0
        filt = (redshift >= zmin) & (redshift < zmax)
        ra, dec, redshift = ra[filt], dec[filt], redshift[filt]

        fig = P.figure()
        ax = fig.add_subplot(111, xlabel='RA (deg)', ylabel='DEC (deg)')
        scat = ax.scatter(ra, dec, c=redshift, cmap=(P.cm.jet))
        cb = fig.colorbar(scat)
        cb.set_label('Photometric redshift')
        if title is not None:
            ax.set_title(title)
        ax.set_xlim(xmin=min(ra) - 0.001, xmax=max(ra) + 0.001)
        ax.set_ylim(ymin=min(dec) - 0.001, ymax=max(dec) + 0.001)
        fig.savefig(self.files['output'].replace(
            '.out', '') + "_redshift_map.png")

def dict_to_array(d, filters='ugriz'):
    """Transform a dictionnary into a list of arrays."""
    return N.array([N.array(d[f]) for f in filters])

class ZSPEC(object):

    """Compare spectroscopic and photometric redshifts."""

    def __init__(self, sfile, names, unit='deg'):
        """Read the input data file and organize the data.

        :param str sfile: File containing the spectroscopic redshift
        :param list names: A list of names for the columns. You must use 'ra' and 'dec'
         for the coordinates.
        :param str unit: Unit of the given coordinates ('deg', 'rad'). Should be understandable by
         astropy.coordinates.SkyCoord

        The spectroscopic redshift must be named 'zspec'.
        """
        self.sfile = sfile
        self.data = ascii.read(sfile, names=names)
        if 'ra' not in self.data.keys():
            raise IOError("RA coordinate must be called 'ra'")
        elif 'dec' not in self.data.keys():
            raise IOError("DEC coordinate must be called 'dec'")
        elif 'zspec' not in self.data.keys():
            raise IOError("Spectroscopic reshift must be 'zspec'")

        # Check whether duplicate galaxies exist in the spectroz sample
        # and remove them. Ideally would average out the various occurances.
        # (To be done later)
        ra = ['{:.20}'.format(x) for x in self.data['ra']]
        dec = ['{:.20}'.format(x) for x in self.data['dec']]
        radec = N.core.defchararray.add(ra, dec)
        unique_radec, good = N.unique(radec, return_index=True)
        if len(unique_radec) < len(radec):
            print("INFO: There are " + str(len(radec) - len(unique_radec)) +
                  " duplicate galaxies in spectroz sample. They are removed.")
        bad = N.delete(N.arange(len(self.data)), good)
        self.data.remove_rows(bad)

        self.skycoords = SkyCoord(self.data['ra'], self.data['dec'], unit=unit)
        self.zphot = self.skycoords_phot = self.match = None

    def load_zphot(self, ra, dec, zphot, unit='deg'):
        """Load the photometric informations and match them to the spectro ones.

        :param list ra: List of RA coordinates
        :param list dec: List of DEC coordinates
        :param list zphot: List of photometric redshift
        :param list unit: List of RA coordinates

        All lists must have the same length.
        """
        assert len(ra) == len(dec) == len(zphot)
        self.zphot = Table([ra, dec, zphot], names=['ra', 'dec', 'zphot'])
        self.skycoords_phot = SkyCoord(ra, dec, unit=unit)
        idx, d2d, d3d = self.skycoords.match_to_catalog_sky(
            self.skycoords_phot)
        self.match = Table([idx, d2d.marcsec, d3d], names=['idx', 'd2d', 'd3d'],
                           meta={'description': ['Indices into first catalog.',
                                                 'On-sky separation between the '
                                                 'closest match for each element',
                                                 '3D distance between the closest'
                                                 'match for each element'],
                                 'unit': ['', 'marcsec', '']})

    def plot(self, cut=300, path_to_png=None):
        """Plot a sky-map of the matches."""
        if self.match is None:
            raise IOError(
                "ERROR: You must load the photometric data first (load_zphot).")

        zspec = self.data['zspec']
        zphot = self.zphot['zphot'][self.match['idx']]
        sdist = N.array(self.match['d2d'])
        filt = sdist < cut
        zspec, zphot, sdist = [x[filt] for x in [zspec, zphot, sdist]]

        fig = P.figure(figsize=(15, 8))

        # z-phot as a function of z-spec
        ax = fig.add_subplot(121, xlabel='Z-spec', ylabel='Z-phot')
        ax.set_xlim([0, 1.5])
        ax.set_ylim([0, 3.5])
        scat = ax.scatter(zspec, zphot, c=sdist, cmap=(P.cm.copper))
        ax.plot(N.arange(int(max(zspec + 1))),
                N.arange(int(max(zspec + 1))), c='red')
        cb = fig.colorbar(scat)
        cb.set_label('On-sky distance (marcsec)')
        ax.set_title("%i galaxies" % len(self.match[filt]))

        # z-phot - zspec as a function of sources on-sky distances
        ax = fig.add_subplot(122, xlabel='On-sky distance',
                             ylabel='(Z-phot - Z-spec)')
        scat = ax.scatter(sdist, zphot - zspec, color='k')
        ax.set_title("%i galaxies" % len(self.match[filt]))
        ax.set_xscale('log')
        ax.set_xlim([1., 1.e6])
        ax.set_ylim([-1., 3.5])

        if path_to_png is not None:
            fig.savefig(path_to_png)

        fig = P.figure()

        # radec scatter plot of all catalogue and zspec galaxies
        ax = fig.add_subplot(121, xlabel='ra', ylabel='dec')
        ax.scatter(self.skycoords_phot.ra, self.skycoords_phot.dec,
                   color='k', label='Photo-z', s=15)
        ax.scatter(self.skycoords.ra, self.skycoords.dec,
                   color='r', label='Spectro-z', s=12)

        # radec scatter plot of matched catalogue and zspec galaxies, within the cut criterion
        ax = fig.add_subplot(122, xlabel='ra', ylabel='dec')
        ax.scatter(self.skycoords_phot.ra[self.match['idx'][filt]],
                   self.skycoords_phot.dec[self.match['idx'][filt]],
                   color='k', label='Photo-z', s=5)
        ax.scatter(self.skycoords.ra[filt], self.skycoords.dec[filt],
                   color='r', label='Spectro-z', s=5)

        if path_to_png is not None:
            fig.savefig(path_to_png.replace('.png', '_map.png'))

        P.show()

    def scatter(self, zclust, cluster=None, cut=0.1, stability=False):
        """Redshift scatter in the cluster.

        Plot the spectroscopic redshift distribution and apply a gaussian fit.
        """
        # Apply selection
        zspec = self.data['zspec'][N.abs(self.data['zspec'] - zclust) < cut]

        # The gaussian fit
        hist, bin_edges = N.histogram(zspec, bins=len(zspec) / 4)
        bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2

        # Use initial guess for the fitting coefficients (A, mu and sigma above)
        coeff = curve_fit(gauss, bin_centers, hist, p0=[
                          N.max(hist[0]), zclust, 0.3])[0]

        # Plot
        fig = P.figure()
        ax = fig.add_subplot(111, xlabel='Z-spec', ylabel='#',
                             title='' if cluster is None else cluster +
                             ', %i object included (cut=%.2f)' % (len(zspec), cut))
        ax.hist(zspec, bins=len(zspec) / 4, histtype='stepfilled')
        ax.axvline(zclust, color='r', label='Z-cluster (%.3f)' % zclust, lw=2)
        ax.plot(bin_centers, gauss(bin_centers, *coeff),
                label='Gaussian fit (mean, std)=(%.3f, %.3f)' %
                (coeff[1], N.abs(coeff[2])), lw=2, color='k')
        ax.legend(loc='best')

        if stability:
            cuts = [0.5, 0.1, 0.05, 0.03]
            z, ze = zip(*[self.scatter(zclust, cluster=cluster, cut=c)
                          for c in cuts])
            fig = P.figure()
            ax = fig.add_subplot(111, xlabel='Selection cut around cluste rredshift',
                                 ylabel='Estimated redshift',
                                 title='' if cluster is None else cluster +
                                 ', %i object included (cut=%.2f)' % (len(zspec), cut))
            ax.axhline(N.mean(z), label='Average value', color='r', lw=2)
            ax.errorbar(cuts, z, yerr=ze, label='Individual fits', color='k',
                        capsize=20, elinewidth=3)
            ax.legend(loc='best')
            print("INFO: Stability over the followed range of selection cuts:")
            print("       ", cuts)
            print("INFO: Input redshift: %.4f" % zclust)
            print("INFO: Average redshift: %.4f =/- %.4f" %
                  (N.mean(z), N.sqrt(N.std(z)**2 + N.mean(ze)**2)))
            P.show()
            return N.mean(z), N.sqrt(N.std(z)**2 + N.mean(ze)**2)
        P.show()
        return coeff[1], N.abs(coeff[2])

def gauss(x, *p):
    """Model function to be used to fit a gaussian distribution."""
    A, mu, sigma = p
    return A * N.exp(- (x - mu) ** 2 / (2. * sigma ** 2))


import cv2
import numpy as np
import sympy as sp

#linePoint_obj = [(200, 200), (300, 200)]
#areaPoint_obj = [(500, 520), (400, 650), (600, 650), (650, 520)]

def is_in_banned_area(center, bboxArea, areaPoint, factor=0.1, img='NULL'):
    """
        判断越线
        车辆中心坐标，车辆BBox面积，区域，压线检测圆域系数，是否DEBUG
    """
    if len(areaPoint) > 2: #多边形区域
        areaMinX = min(list([k[0] for k in areaPoint]))
        areaMaxX = max(list([k[0] for k in areaPoint]))
        areaMinY = min(list([k[1] for k in areaPoint]))
        areaMaxY = max(list([k[1] for k in areaPoint]))
        if not (areaMinX - 5 <= center[0] <= areaMaxX + 5 and
                areaMinY - 5 <= center[1] <= areaMaxY + 5):  # 在区域外
            return False

        base_point = (10, 10)  # 初始化射线
        base_x = [base_point[0], center[0]]
        base_y = [base_point[1], center[1]]
        base_f = np.polyfit(base_y, base_x, 1) #直线拟合
        base_fy = np.poly1d(base_f)
        base_fy_k = base_fy[1]
        base_fy_b = base_fy[0]

        y = sp.symbols('y')  # 未知数Y
        x = sp.symbols('x')  # 未知数X

        crash = 0 #交点个数，奇为内，偶为外
        for count in range(len(areaPoint) - 1):  # 遍历区域边
            point1 = areaPoint[count]
            point2 = areaPoint[count + 1]
            if point1[1] == point2[1]:  # 平行于X轴的线段处理
                Jy = point1[1]
                Jx = base_fy(Jy)
                xx = [point1[0], point2[0]]
                minx, maxx = min(xx), max(xx)

                if not img == 'NULL':
                    cv2.circle(img, (int(Jx), int(Jy)), 5, (255, 255, 255), 10)  # 展示所有平行点

                if minx - 5 <= Jx <= maxx + 5 and \
                        base_point[1] - 5 <= Jy <= center[1] + 5: #交点在两条线段之内
                    crash += 1

                    if not img == 'NULL':  # 交点可视化
                        cv2.circle(img, (int(Jx), int(Jy)), 2, (0, 255, 0), 8)
                if not img == 'NULL':  # 区域可视化
                    cv2.line(img, point1, point2, (0, 0, 255), 2, 1)

                continue
            else:
                xx = [point1[0], point2[0]]
                yy = [point1[1], point2[1]]
                f = np.polyfit(yy, xx, 1) # 拟合一次函数
                fy = np.poly1d(f)
                fy_k = fy[1]
                fy_b = fy[0]

                fy1 = fy_k * y + fy_b - x  # 求解方程
                fy2 = base_fy_k * y + base_fy_b - x
                JXY = sp.solve([fy1, fy2], [x, y])  # 求交点坐标
                Jx = JXY[x]
                Jy = JXY[y]

                if not img == 'NULL':
                    cv2.circle(img, (int(Jx), int(Jy)), 5, (155, 155, 255), 10)  # 展示所有点

                miny, maxy = min(yy), max(yy)
                minx, maxx = min(xx), max(xx)
                if miny - 5 <= Jy <= maxy + 5 and minx - 5 <= Jx <= maxx + 5 and \
                        base_point[1] - 5 <= Jy <= center[1] + 5 and \
                        base_point[0] - 5 <= Jx <= center[0] + 5: #交点在两条线段之内
                    crash += 1

                    if not img == 'NULL':  # 交点可视化
                        cv2.circle(img, (int(Jx), int(Jy)), 2, (0, 255, 0), 8)
                if not img == 'NULL':  # 车道可视化
                    cv2.line(img, point1, point2, (0, 0, 255), 2, 1)

        point1 = areaPoint[0]  # 最后一个区域边
        point2 = areaPoint[-1]

        if point1[1] == point2[1]:  # 平行于X轴的线段处理
            Jy = (point1[1] + point2[1]) // 2
            Jx = base_fy(Jy)
            xx = [point1[0], point2[0]]
            yy = [point1[1], point2[1]]
            minx, maxx = min(xx), max(xx)

            if not img == 'NULL':
                cv2.circle(img, (int(Jx), int(Jy)), 5, (255, 255, 255), 10)  # 展示所有平行点

            if minx - 5 <= Jx <= maxx + 5 and base_point[1] - 5 <= Jy <= center[1] + 5:  # 交点在线段内
                crash += 1

                if not img == 'NULL':  # 交点可视化
                    cv2.circle(img, (int(Jx), int(Jy)), 2, (0, 255, 0), 4)
            if not img == 'NULL':  # 车道可视化
                cv2.line(img, point1, point2, (0, 0, 255), 2, 1)

        else:
            xx = [point1[0], point2[0]]
            yy = [point1[1], point2[1]]
            f = np.polyfit(yy, xx, 1)
            fy = np.poly1d(f)
            fy_k = fy[1]
            fy_b = fy[0]
            fy1 = fy_k * y + fy_b - x
            fy2 = base_fy_k * y + base_fy_b - x
            JXY = sp.solve([fy1, fy2], [x, y])  # 求交点坐标
            Jx = JXY[x]
            Jy = JXY[y]

            if not img == 'NULL':
                cv2.circle(img, (int(Jx), int(Jy)), 5, (155, 155, 255), 10)  # 展示所有平行点

            miny, maxy = min(yy), max(yy)
            minx, maxx = min(xx), max(xx)
            if miny - 5 <= Jy <= maxy + 5 and minx - 5 <= Jx <= maxx + 5 and \
                    base_point[1] - 5 <= Jy <= center[1] + 5 and \
                    base_point[0] - 5 <= Jx <= center[0] + 5: #交点在线段内
                crash += 1

                if not img == 'NULL':  # 交点可视化
                    cv2.circle(img, (int(Jx), int(Jy)), 2, (0, 255, 0), 4)
            if not img == 'NULL':  # 车道可视化
                cv2.line(img, point1, point2, (0, 0, 255), 2, 1)

        if crash % 2 == 1:  # Area判定

            if not img == 'NULL':
                zi = 'YaArea'  # 可视化
                cv2.putText(img, zi, center, 4, .8, (0, 0, 255), 1)

            return True
        else:
            return False
    else: #线区域
        circle_r = bboxArea * factor  # 检测圆域半径

        areaMinX = min(list([k[0] for k in areaPoint]))
        areaMaxX = max(list([k[0] for k in areaPoint]))
        areaMinY = min(list([k[1] for k in areaPoint]))
        areaMaxY = max(list([k[1] for k in areaPoint]))
        if not (areaMinX - circle_r <= center[0] <= areaMaxX + circle_r and
                areaMinY - circle_r <= center[1] <= areaMaxY + circle_r):  # 在区域外
            return False

        if not img == 'NULL':
            cv2.circle(img, center, int(circle_r), (0, 255, 0), 4)
            cv2.line(img, areaPoint[0], areaPoint[1], (0, 0, 255), 2, 1)

        if areaPoint[0][1] != areaPoint[1][1]: #斜率存在
            xx = [areaPoint[0][0], areaPoint[1][0]]
            yy = [areaPoint[0][1], areaPoint[1][1]]
            f = np.polyfit(yy, xx, 1)  # 拟合一次函数
            fy = np.poly1d(f)
            fy_k = fy[1]
            fy_j = fy[0] #截距
            fy_a = fy_k
            fy_b = -1
            fy_c = fy_j
            d = abs(fy_a * center[1] + fy_b * center[0] + fy_c) / (fy_a**2 * fy_b**2)**0.5
            if d <= circle_r: #压线

                if not img == 'NULL':
                    zi = 'YaLine'  # 可视化
                    cv2.putText(img, zi, center, 4, .8, (0, 0, 255), 1)

                return True
            else:
                return False
        else: #斜率不存在
            d = abs(areaPoint[0][1] - center[1])
            if d <= circle_r: #压线

                if not img == 'NULL':
                    zi = 'YaLine'  # 可视化
                    cv2.putText(img, zi, center, 4, .8, (0, 0, 255), 1)

                return True
            else:
                return False

if __name__ == '__main__':
    pass


try:
    import galsim
except ImportError:
    class GalSimWCS(object):
        def __init__(self, *args, **kwargs):
            raise NotImplementedError(
                'Unable to import galsim. The GalSimWCS interface is not available.')
else:

    from .decaminfo import DECamInfo
    from .PixelMapCollection import PixelMapCollection
    from .DESMaps import DESMaps
    import os
    import astropy.coordinates
    import numpy as np
    from . import files

    class GalSimWCS(galsim.wcs.CelestialWCS):
        """A wrapper of the `PixelMapCollection` class that can be used as a galsim `WCS` type.

        This can be constructed using either:
        * an existing `PixelMapCollection` object (supplied as `pmc`), or
        * a generic serialized PixelMapCollection in a YAML file,
          (supplied by `yaml_file`), or
        * The `DESMaps` class, which accesses files in the format of
          the DES Y6A1_ASTROMETRY release, which one selects via
          `use_DESMaps=True`.  In this case one can supply names for
          the `exposure_file,guts_file,resids_file,affine_file` that
          the `DESMaps` needs, or any of these will use Y6A1 default
          names if they are supplied as `None` (This is default
          behavior).

        Exactly one of these must be true: `pmc is not None`;
        `yaml_file is not None`; `use_DESMaps`.

        The WCS within the collection is selected by name.  One can
        either specify the name directly with `wcs_name`; or a name
        can be constructed using DES conventions from an exposure
        number `exp` and `ccdnum`.

        :param pmc:         An existing pixmappy.PixelMapCollection instance [default: None]
        :param yaml_file: The yaml file with the PixelMapCollection
                            description. [default: None]
        :param use_DESMaps:  If `True`, use `DESMaps` to construct WCS. [default: False]

        :param dir:         Optional directory to prepend to all filename
                            arguments. [default: None.  Note that if use_DESMaps is
                            True, then you can place a search path for setup files into
                            the environment variable CAL_PATH]

        :param wcs_name: The name of the WCS within the
                            PixelMapCollection to use.  [default:
                            None; either `wcs_name` or (`exp` and
                            `ccdnum`) is required.  DESMaps require
                            the exp+ccdnum.]
        :param exp:         The exposure number of the desired WCS. [default: None]
        :param ccdnum:      The CCD number of the desired WCS. [default: None]

        :param exposure_file: FITS file holding binary table of DES
                            per-exposure info for DESMaps.  [default:
                            None; if `use_DESMaps` then the file in the
                            Y6A1_ASTRONOMY release will be used in
                            this default case.]
        :param guts_file:   YAML file holding static DECam distortions
                            for DESMaps.  [default: None; (same
                            behavior as above).
        :param resids_file: FITS file holding 2d residual adjustment
                            maps for DESMaps [default: None; (same
                            behavior as above).
        :param affine_file: FITS file holding time-dependent DECam CCD
                            affine tweaks for DESMaps [default: None;
                            (same behavior as above).

        :param origin:      Optional origin position for the image coordinate system.
                            If provided, it should be a PositionD or
                            PositionI.  [default: None]
        :param cache:       Cache this file's PixelMapCollection in the GalSimWCS.cache dict?
                            [default: True]
        :param default_color: The default color to use if this WCS involves color terms and
                            `wcs.toWorld` or similar methods do not pass in a color value.
                            [default: None, which means an exception will be raised if no
                            color term is in the map and no color value is provided]
        """
        _opt_params = { "origin" : galsim.PositionD, "ccdnum": int,
                        "dir" : str, "guts_file" : str,
                        "exposure_file" : str, "resids_file" : str, "affine_file" : str,
                        "default_color" : float}
        _single_params = [ { "wcs_name" : str, "exp" : int },
                               {"yaml_file" : str, "use_DESMaps":bool}]
        _takes_rng = False

        info = DECamInfo()
        cache = dict()

        def __init__(self, pmc=None, yaml_file=None, use_DESMaps=False, dir=None,
                     wcs_name=None, exp=None, ccdnum=None,
                     exposure_file=None, guts_file=None, resids_file=None, affine_file=None,
                     origin=None, cache=True, default_color=None):
            self._color = default_color

            # Make sure only one method is in use:
            count = int(pmc is not None) + int(yaml_file is not None) + int(use_DESMaps)
            if count!=1:
                raise TypeError("Must provide exactly one of yaml_file, pmc, or use_DESMaps")

            if pmc is not None:
                self._pmc = pmc
                self._tag = 'pmc=%r'%pmc  # Used in __repr__

            if yaml_file is not None:
                if dir is not None:
                    yaml_file = os.path.join(dir,yaml_file)
                if yaml_file in self.cache:
                    pmc = self.cache[yaml_file]
                else:
                    pmc = PixelMapCollection(yaml_file)
                    if cache:
                        self.cache[yaml_file] = pmc
                self._tag = 'yaml_file=%r'%yaml_file
                self._pmc = pmc

            if use_DESMaps:
                if exp is None or ccdnum is None:
                    raise TypeError("exp and ccdnum must be provided when using DESMaps")

                self._tag = 'use_DESMaps=True'

                if exposure_file is None:
                    exposure_file = files.default_exposure_file
                else:
                    self._tag = self._tag + ', exposure_file=%s'%exposure_file

                if guts_file is None:
                    guts_file = files.default_guts_file
                else:
                    self._tag = self._tag + ', guts_file=%s'%guts_file

                if resids_file is None:
                    resids_file = files.default_resids_file
                else:
                    self._tag = self._tag + ', resids_file=%s'%resids_file

                if affine_file is None:
                    affine_file = files.default_affine_file
                else:
                    self._tag = self._tag + ', affine_file=%s'%affine_file

                if dir is not None:
                    exposure_file = os.path.join(dir,exposure_file)
                    guts_file = os.path.join(dir,guts_file)
                    resids_file = os.path.join(dir,resids_file)
                    affine_file = os.path.join(dir,affine_file)
                    self._tag = self._tag + ', dir=%s'%dir

                # We'll cache the DESMaps object by the exposure_file name
                if exposure_file in self.cache:
                    pmc = self.cache[exposure_file]
                else:
                    pmc = DESMaps(guts_file=guts_file,
                                  exposure_file=exposure_file,
                                  resids_file=resids_file,
                                  affine_file=affine_file)
                    if cache:
                        self.cache[exposure_file] = pmc
                self._pmc = pmc

            # Now extract the desired WCS from our PixelMapCollection or DESMaps
            if use_DESMaps:
                if exp is None or ccdnum is None:
                    raise TypeError("DESMaps require exp,ccdnum")
                ccdname = self.info.ccddict[ccdnum]
                self._wcs_name = 'D%s/%s'%(exp, ccdname)  #Used by __eq__
                self._wcs = pmc.getDESWCS(exp, ccdname)
                self._tag = self._tag + ', exp=%r, ccdnum=%r'%(exp,ccdnum)

            else:
                if wcs_name is None:
                    if exp is None or ccdnum is None:
                        raise TypeError("Must provide either wcs_name or (exp,ccdnum)")
                    ccdname = self.info.ccddict[ccdnum]
                    self._wcs_name = 'D%s/%s'%(exp, ccdname)
                elif exp is not None or ccdnum is not None:
                    raise TypeError("Cannot provide both wcs_name and (exp,ccdnum)")
                else:
                    self._wcs_name = wcs_name

                self._wcs = pmc.getWCS(self._wcs_name)
                self._tag = self._tag + ', wcs_name=%r'%self._wcs_name

            # Set origin, if any
            if origin is None:
                self._origin = galsim.PositionD(0,0)
            else:
                if isinstance(origin, galsim.PositionI):
                    origin = galsim.PositionD(origin.x, origin.y)
                elif not isinstance(origin, galsim.PositionD):
                    raise TypeError("origin must be a PositionD or PositionI argument")
                self._origin = origin

        @property
        def pmc(self): return self._pmc

        @property
        def wcs_name(self):
            # Note: older versions of pixmappy didn't set _wcs_name if it was None.
            # This is fixed above, but to accommodate reading in older serialized pixmappy
            # objects, we check for the attribute existing here.
            return self._wcs_name if hasattr(self,'_wcs_name') else None

        @property
        def origin(self): return self._origin

        @classmethod
        def clear_cache(cls):
            """Clear the cache of PixelMapCollections.

            The PixelMapCollection objects that are read in from a file are often rather large,
            and a typical use case involves getting many wcs objects from the same file.
            So the GalSimWCS class caches them avoid needing to read in the file repeatedly for
            each WCS you want to extract from it.

            However, if you are done with a particular input file, you might not want to keep
            it around anymore.  So ``pixmappy.GalSimWCS.clear_cache()`` will release the memory
            currently being used by the cache.

            You can also modify the cache yourself if you want (say to remove a particular element
            rather than all objects).  It is a dict indexed by the the yaml filename or the exposure
            filename for DESMaps.
            """
            cls.cache.clear()

        def _radec(self, x, y, c=None):
            ra, dec = self._wcs.toSky(x, y, c=c )
            ra *= galsim.degrees / galsim.radians
            dec *= galsim.degrees / galsim.radians
            return ra, dec

        def _xy(self, ra, dec, c=None):
            ra *= galsim.radians / galsim.degrees
            dec *= galsim.radians / galsim.degrees
            x, y = self._wcs.toPix(ra, dec, c=c)
            return x, y

        def _newOrigin(self, origin):
            ret = self.copy()
            ret._origin = origin
            return ret

        def _writeHeader(self, header, bounds):
            raise NotImplementedError("Cannot write PixelMap to a fits header")

        @staticmethod
        def _readHeader(header):
            raise NotImplementedError("Cannot read PixelMap from a fits header")

        def copy(self):
            # The copy module version of copying the dict works fine here.
            import copy
            return copy.copy(self)

        def __eq__(self, other):
            return (isinstance(other, GalSimWCS) and
                    self._tag == other._tag and
                    self.wcs_name == other.wcs_name and
                    self.origin == other.origin )

        def __repr__(self):
            # Should eval back into itself
            s = "pixmappy.GalSimWCS(%s, origin=%r"%(self._tag,  self.origin)
            if self._color is not None:
                s += ', default_color=%r'%self._color
            s += ')'
            return s

        def __hash__(self): return hash(repr(self))

        def __getstate__(self):
            # The naive pickling works, but it includes _pmc, which is huge, and not actually
            # necessary for the functioning of the object.  (It's just for information purposes
            # really.)  So remove it from the dict to be pickled.
            d = self.__dict__.copy()
            d['_pmc'] = None
            return d


# -- coding: utf-8 --
import os
from time import sleep
import pandas as pd
import numpy as np

switch_on = 0
state_cur = []
# pull requests via jmeter
for i in range(0, 20):
    if switch_on == 0:
        os.system("jmeter -n -t newplan.jmx -l lowtest_ver/ver_%s.csv -JthreadNum=1650000" % i)
    elif switch_on == 1:
        os.system("jmeter -n -t newplan.jmx -l lowtest_ver/ver_%s.csv -JthreadNum=825000" % i)
    jmeter = pd.read_csv('lowtest_ver/ver_%s.csv' % i)
    elapsed = np.array(jmeter['elapsed'])
    timeStamp = jmeter['timeStamp']
    success = np.array(jmeter['success'].astype('int'))
    latency = np.array(jmeter['Latency'])
    connect = np.array(jmeter['Connect'])
    lenth = len(timeStamp)
    timeStamp_len = timeStamp[lenth - 1] - timeStamp[0]
    # QoS: version 3
    state_cur_val = (np.mean(elapsed) * 0.05 + np.mean(latency) * 0.1 + np.mean(connect) * 0.05 + (
            1 - np.mean(success)) * 2 + (timeStamp_len - 300000) / 300000) / 8
    if state_cur_val > 4:
        state_cur_val = 4
    state_cur.append(int(round(state_cur_val)))
    csv = pd.DataFrame(data=state_cur)
    names = ['state']
    csv.columns = names
    csv.to_csv('lowtest_ver.csv', mode='w')
    # run requests_server.py
    if i > 4:
        with open('lowtest_hor.py', 'r') as f:
            exec(f.read())
    os.system("pkill -9 java")
    sleep(90)