#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#    Project: Azimuthal integration
#             https://github.com/silx-kit/pyFAI
#
#    Copyright (C) 2014-2018 European Synchrotron Radiation Facility, Grenoble, France
#
#    Principal author:       Jérôme Kieffer (Jerome.Kieffer@ESRF.eu)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

__author__ = "Jérôme Kieffer"
__contact__ = "Jerome.Kieffer@ESRF.eu"
__license__ = "MIT"
__copyright__ = "European Synchrotron Radiation Facility, Grenoble, France"
__date__ = "17/10/2018"
__status__ = "production"

import sys
import os
import threading
from math import ceil, sqrt
import logging
logger = logging.getLogger(__name__)
import numpy
import fabio
from scipy.ndimage import label, distance_transform_edt
from scipy.ndimage.filters import median_filter
from .utils.decorators import deprecated
from .ext.bilinear import Bilinear
from .utils import gaussian_filter, binning, unbinning, is_far_from_group
from .third_party import six

if os.name != "nt":
    WindowsError = RuntimeError

class Massif(object):
    """
    A massif is defined as an area around a peak, it is used to find neighboring peaks
    """
    TARGET_SIZE = 1024

    def __init__(self, data=None, mask=None):
        """Constructor of the class...

        :param data: 2D array or filename (discouraged)
        :param mask: array with non zero for invalid data
        """
        if isinstance(data, six.string_types) and os.path.isfile(data):
            self.data = fabio.open(data).data.astype("float32")
        elif isinstance(data, fabio.fabioimage.fabioimage):
            self.data = data.data.astype("float32")
        else:
            try:
                self.data = data.astype("float32")
            except Exception as error:
                logger.error("Unable to understand this type of data %s: %s", data, error)
        self.mask = mask
        self._cleaned_data = None
        self._bilin = Bilinear(self.data)
        self._blurred_data = None
        self._median_data = None
        self._labeled_massif = None
        self._number_massif = None
        self._valley_size = None
        self._binned_data = None
        self._reconstruct_used = None
        self.binning = None  # Binning is 2-list usually
        self._sem = threading.Semaphore()
        self._sem_label = threading.Semaphore()
        self._sem_binning = threading.Semaphore()
        self._sem_median = threading.Semaphore()

    def nearest_peak(self, x):
        """
        :param x: coordinates of the peak
        :returns: the coordinates of the nearest peak
        """
        out = self._bilin.local_maxi(x)
        if isinstance(out, tuple):
            res = out
        elif isinstance(out, numpy.ndarray):
            res = tuple(out)
        else:
            res = [int(i) for idx, i in enumerate(out) if 0 <= i < self.data.shape[idx]]
        if (len(res) != 2) or not((0 <= out[0] < self.data.shape[0]) and (0 <= res[1] < self.data.shape[1])):
            logger.warning("in nearest_peak %s -> %s", x, out)
            return
        elif (self.mask is not None) and self.mask[int(res[0]), int(res[1])]:
            logger.info("Masked pixel %s -> %s", x, out)
            return
        else:
            return res

    def calculate_massif(self, x):
        """
        defines a map of the massif around x and returns the mask
        """
        labeled = self.get_labeled_massif()
        if labeled[x[0], x[1]] > 0:  # without relabeled the background is 0 labeled.max():
            return (labeled == labeled[x[0], x[1]])

    def find_peaks(self, x, nmax=200, annotate=None, massif_contour=None, stdout=sys.stdout):
        """
        All in one function that finds a maximum from the given seed (x)
        then calculates the region extension and extract position of the neighboring peaks.
        :param x: coordinates of the peak, seed for the calculation
        :type x: tuple of integer
        :param nmax: maximum number of peak per region
        :param annotate: call back method taking number of points + coordinate as input.
        :param massif_contour: callback to show the contour of a massif with the given index.
        :param stdout: this is the file where output is written by default.
        :return: list of peaks
        """
        listpeaks = []
        region = self.calculate_massif(x)
        if region is None:
            logger.error("You picked a background point at %s", x)
            return listpeaks
        xinit = self.nearest_peak(x)
        if xinit is None:
            logger.error("Unable to find peak in the vinicy of %s", x)
            return listpeaks
        else:
            if not region[int(xinit[0] + 0.5), int(xinit[1] + 0.5)]:
                logger.error("Nearest peak %s is not in the same region  %s", xinit, x)
                return listpeaks

            if annotate is not None:
                try:
                    annotate(xinit, x)
                except Exception as error:
                    logger.error("Error in annotate %i: %i %i. %s", len(listpeaks), xinit[0], xinit[1], error)

        listpeaks.append(xinit)
        cleaned_data = self.cleaned_data
        mean = cleaned_data[region].mean(dtype=numpy.float64)
        region2 = region * (cleaned_data > mean)
        idx = numpy.vstack(numpy.where(region2)).T
        numpy.random.shuffle(idx)
        nmax = min(nmax, int(ceil(sqrt(idx.shape[0]))))
        if massif_contour is not None:
            try:
                massif_contour(region)
            except (WindowsError, MemoryError) as error:
                logger.error("Error in plotting region: %s", error)
        nbFailure = 0
        for j in idx:
            xopt = self.nearest_peak(j)
            if xopt is None:
                nbFailure += 1
                continue
            if (region2[int(xopt[0] + 0.5), int(xopt[1] + 0.5)]) and not (xopt in listpeaks):
                stdout.write("[ %4i, %4i ] --> [ %5.1f, %5.1f ] after %3i iterations %s" % (tuple(j) + tuple(xopt) + (nbFailure, os.linesep)))
                listpeaks.append(xopt)
                nbFailure = 0
            else:
                nbFailure += 1
            if (len(listpeaks) > nmax) or (nbFailure > 2 * nmax):
                break
        return listpeaks

    def peaks_from_area(self, mask, Imin=numpy.finfo(numpy.float64).min,
                        keep=1000, dmin=0.0, seed=None, **kwarg):
        """
        Return the list of peaks within an area

        :param mask: 2d array with mask.
        :param Imin: minimum of intensity above the background to keep the point
        :param keep: maximum number of points to keep
        :param kwarg: ignored parameters
        :param dmin: minimum distance to another peak
        :param seed: list of good guesses to start with
        :return: list of peaks [y,x], [y,x], ...]
        """
        all_points = numpy.vstack(numpy.where(mask)).T
        res = []
        cnt = 0
        dmin2 = dmin * dmin
        if len(all_points) > 0:
            numpy.random.shuffle(all_points)
        if seed:
            seeds = numpy.array(list(seed))
            if len(seeds) > 0:
                numpy.random.shuffle(seeds)
            all_points = numpy.concatenate((seeds, all_points))
        for idx in all_points:
            out = self.nearest_peak(idx)
            if out is not None:
                msg = "[ %3i, %3i ] -> [ %.1f, %.1f ]"
                logger.debug(msg, idx[1], idx[0], out[1], out[0])
                p0, p1 = int(round(out[0])), int(round(out[1]))
                if mask[p0, p1]:
                    if (self.data[p0, p1] > Imin) and is_far_from_group(out, res, dmin2):
                        res.append(out)
                        cnt = 0
            if len(res) >= keep or cnt > keep:
                break
            else:
                cnt += 1
        return res

    def init_valley_size(self):
        if self._valley_size is None:
            self.valley_size = max(5., max(self.data.shape) / 50.)

    @property
    def valley_size(self):
        "Defines the minimum distance between two massifs"
        if self._valley_size is None:
            self.init_valley_size()
        return self._valley_size

    @valley_size.setter
    def valley_size(self, size):
        new_size = float(size)
        if self._valley_size != new_size:
            self._valley_size = new_size
            t = threading.Thread(target=self.get_labeled_massif)
            t.start()

    @valley_size.deleter
    def valley_size(self):
        self._valley_size = None
        self._blurred_data = None

    @property
    def cleaned_data(self):
        if self.mask is None:
            return self.data
        else:
            if self._cleaned_data is None:
                idx = distance_transform_edt(self.mask,
                                             return_distances=False,
                                             return_indices=True)
                self._cleaned_data = self.data[tuple(idx)]
            return self._cleaned_data

    def get_binned_data(self):
        """
        :return: binned data
        """
        if self._binned_data is None:
            with self._sem_binning:
                if self._binned_data is None:
                    logger.info("Image size is %s", self.data.shape)
                    self.binning = []
                    for i in self.data.shape:
                        if i % self.TARGET_SIZE == 0:
                            self.binning.append(max(1, i // self.TARGET_SIZE))
                        else:
                            for j in range(i // self.TARGET_SIZE - 1, 0, -1):
                                if i % j == 0:
                                    self.binning.append(max(1, j))
                                    break
                            else:
                                self.binning.append(1)
#                    self.binning = max([max(1, i // self.TARGET_SIZE) for i in self.data.shape])
                    logger.info("Binning size is %s", self.binning)
                    self._binned_data = binning(self.cleaned_data, self.binning)
        return self._binned_data

    def get_median_data(self):
        """
        :return: a spatial median filtered image 3x3
        """
        if self._median_data is None:
            with self._sem_median:
                if self._median_data is None:
                    self._median_data = median_filter(self.cleaned_data, 3)
        return self._median_data

    def get_blurred_data(self):
        """
        :return: a blurred image
        """
        if self._blurred_data is None:
            with self._sem:
                if self._blurred_data is None:
                    logger.debug("Blurring image with kernel size: %s", self.valley_size)
                    self._blurred_data = gaussian_filter(self.get_binned_data(),
                                                         [self.valley_size / i for i in self.binning],
                                                         mode="reflect")
        return self._blurred_data

    def get_labeled_massif(self, pattern=None, reconstruct=True):
        """
        :param pattern: 3x3 matrix
        :param reconstruct: if False, split massif at masked position, else reconstruct missing part.
        :return: an image composed of int with a different value for each massif
        """
        if self._labeled_massif is None:
            with self._sem_label:
                if self._labeled_massif is None:
                    if pattern is None:
                        pattern = numpy.ones((3, 3), dtype=numpy.int8)
                    logger.debug("Labeling all massifs. This takes some time !!!")
                    massif_binarization = (self.get_binned_data() > self.get_blurred_data())
                    if (self.mask is not None) and (not reconstruct):
                            binned_mask = binning(self.mask.astype(int), self.binning, norm=False)
                            massif_binarization = numpy.logical_and(massif_binarization, binned_mask == 0)
                    self._reconstruct_used = reconstruct
                    labeled_massif, self._number_massif = label(massif_binarization,
                                                                pattern)
                    # TODO: investigate why relabel fails
                    # relabeled = relabel(labeled_massif, self.get_binned_data(), self.get_blurred_data())
                    relabeled = labeled_massif
                    self._labeled_massif = unbinning(relabeled, self.binning, False)
                    logger.info("Labeling found %s massifs.", self._number_massif)
        return self._labeled_massif

    @deprecated(reason="switch to pep8 style", replacement="init_valley_size", since_version="0.16.0")
    def initValleySize(self):
        self.init_valley_size()

    @deprecated(reason="switch to PEP8 style", replacement="get_median_data", since_version="0.16.0")
    def getMedianData(self):
        return self.get_median_data()

    @deprecated(reason="switch to PEP8 style", replacement="get_binned_data", since_version="0.16.0")
    def getBinnedData(self):
        return self.get_binned_data()

    @deprecated(reason="switch to PEP8 style", replacement="get_blurred_data", since_version="0.16.0")
    def getBluredData(self):
        return self.get_blurred_data()

    @deprecated(reason="switch to PEP8 style", replacement="get_labeled_massif", since_version="0.16.0")
    def getLabeledMassif(self, pattern=None):
        return self.get_labeled_massif(pattern)


import numpy as np
from astropy.io import fits
from astropy.table import Table, vstack
from astropy.wcs import WCS
import os
import argparse
import logging, traceback
import pandas as pd
from copy import copy

from bkg_rate_estimation import get_avg_lin_cub_rate_quad_obs
from config import quad_dicts, EBINS0, EBINS1,\
                solid_angle_dpi_fname, bright_source_table_fname
from sqlite_funcs import write_rate_fits_from_obj, get_conn
from dbread_funcs import get_info_tab, guess_dbfname, get_files_tab
from event2dpi_funcs import filter_evdata
from models import Bkg_Model_wFlatA, CompoundModel, Point_Source_Model_Binned_Rates
from LLH import LLH_webins
from minimizers import NLLH_ScipyMinimize, NLLH_ScipyMinimize_Wjacob
from ray_trace_funcs import RayTraces
from coord_conv_funcs import convert_radec2imxy
from gti_funcs import add_bti2gti, bti2gti, gti2bti, union_gtis
from wcs_funcs import world2val

def cli():
    parser = argparse.ArgumentParser()
    parser.add_argument('--evfname', type=str,\
            help="Event data file",
            default=None)
    parser.add_argument('--dmask', type=str,\
            help="Detmask fname",
            default=None)
    parser.add_argument('--dbfname', type=str,\
            help="Name to save the database to",\
            default=None)
    parser.add_argument('--job_id', type=int,\
            help="Job ID",\
            default=0)
    parser.add_argument('--Njobs', type=int,\
            help="Number of jobs",\
            default=1)
    parser.add_argument('--twind', type=float,\
            help="Number of seconds to go +/- from the trigtime",\
            default=20)
    parser.add_argument('--bkg_dur', type=float,\
            help="bkg duration",\
            default=40.0)
    parser.add_argument('--bkg_nopost',\
            help="Don't use time after signal window for bkg",\
            action='store_true')
    parser.add_argument('--bkg_nopre',\
            help="Don't use time before signal window for bkg",\
            action='store_true')
    parser.add_argument('--pcfname', type=str,\
            help="partial coding file name",\
            default='pc_2.img')

    args = parser.parse_args()
    return args

def ang_sep(ra0, dec0, ra1, dec1):

    dcos = np.cos(np.radians(np.abs(ra0 - ra1)))
    angsep = np.arccos(np.cos(np.radians(90-dec0))*np.cos(np.radians(90-dec1)) +\
                        np.sin(np.radians(90-dec0))*np.sin(np.radians(90-dec1))*dcos)
    return np.rad2deg(angsep)

def im_dist(imx0, imy0, imx1, imy1):
    return np.hypot((imx1 - imx0), (imy1 - imy0))

def add_imxy2src_tab(src_tab, attfile, t0):

    att_ind = np.argmin(np.abs(attfile['TIME'] - t0))
    att_quat = attfile['QPARAM'][att_ind]
    pnt_ra, pnt_dec = attfile['POINTING'][att_ind,:2]
    imxs = np.zeros(len(src_tab))
    imys = np.zeros(len(src_tab))
    src_tab['PntSep'] = ang_sep(pnt_ra, pnt_dec, src_tab['RAJ2000'], src_tab['DEJ2000'])
    for i in range(len(imxs)):
        if src_tab['PntSep'][i] > 80.0:
            imxs[i], imys[i] = np.nan, np.nan
            continue
        imxs[i], imys[i] = convert_radec2imxy(src_tab['RAJ2000'][i],\
                                              src_tab['DEJ2000'][i],\
                                              att_quat)
    src_tab['imx'] = imxs
    src_tab['imy'] = imys
    return src_tab

def get_srcs_infov(attfile, t0, pcfname=None, pcmin=5e-2):

    brt_src_tab = Table.read(bright_source_table_fname)
    add_imxy2src_tab(brt_src_tab, attfile, t0)
    bl_infov = (np.abs(brt_src_tab['imy'])<.95)&(np.abs(brt_src_tab['imx'])<1.75)
    if pcfname is not None:
        PC = fits.open(pcfname)[0]
        pc = PC.data
        w_t = WCS(PC.header, key='T')
        pcvals = world2val(w_t, pc, brt_src_tab['imx'],\
                            brt_src_tab['imy'])
        bl_infov = bl_infov&(pcvals>=pcmin)
    N_infov = np.sum(bl_infov)
    return brt_src_tab[bl_infov]

def bkg_withPS_fit(PS_tab, model, llh_obj, t0s, t1s,\
                   dimxy=2e-3, im_steps=5, test_null=False):

    Nps = len(PS_tab)
    imax = np.linspace(-dimxy, dimxy, im_steps)
    if im_steps == 3:
        imax = np.linspace(-dimxy/2., dimxy/2., im_steps)
    elif im_steps == 2:
        imax = np.linspace(-dimxy/2., dimxy/2., im_steps)
    elif im_steps == 1:
        imax = np.array([0.0])

    imlist = []
    for i in range(Nps):
        imlist += [imax,imax]
    imgs = np.meshgrid(*imlist)
    Npnts = imgs[0].size

    bkg_miner = NLLH_ScipyMinimize_Wjacob('')
    bkg_miner.set_llh(llh_obj)
    llh_obj.set_time(t0s,t1s)

    bf_params_list = []
    bkg_nllhs = np.zeros(Npnts)

    nebins = model.nebins

    for i in range(Npnts):
        bf_params = {}
        im_names = []
        for j in range(Nps):
            row = PS_tab[j]
            psname = row['Name']
            bf_params[psname+'_imx'] = imgs[2*j].ravel()[i] + row['imx']
            bf_params[psname+'_imy'] = imgs[2*j+1].ravel()[i] + row['imy']
            im_names = [psname+'_imx', psname+'_imy']
        im_vals = [bf_params[nm] for nm in im_names]

        for e0 in range(nebins):
            bkg_miner.set_fixed_params(bkg_miner.param_names)
            bkg_miner.set_fixed_params(im_names, im_vals)
            e0_pnames = []
            for pname in bkg_miner.param_names:
                try:
                    if int(pname[-1])==e0:
                        e0_pnames.append(pname)
                except:
                    pass
            bkg_miner.set_fixed_params(e0_pnames, fixed=False)
            llh_obj.set_ebin(e0)

            bf_vals, bkg_nllh, res = bkg_miner.minimize()
            bkg_nllhs[i] += bkg_nllh[0]
            for ii, pname in enumerate(e0_pnames):
                bf_params[pname] = bf_vals[0][ii]
        bf_params_list.append(bf_params)

    bf_ind = np.argmin(bkg_nllhs)
    bf_params = bf_params_list[bf_ind]
    bf_nllh = bkg_nllhs[bf_ind]

    if test_null:
        TS_nulls = {}
        for i in range(Nps):
            params_ = copy(bf_params)
            row = PS_tab[i]
            psname = row['Name']
            for j in range(nebins):
                params_[psname+'_rate_'+str(j)] = 0.0
            llh_obj.set_ebin(-1)
            nllh_null = -llh_obj.get_logprob(params_)
            TS_nulls[psname] = np.sqrt(2.*(nllh_null - bf_nllh))
            if np.isnan(TS_nulls[psname]):
                TS_nulls[psname] = 0.0
        return bf_nllh, bf_params, TS_nulls

    return bf_nllh, bf_params

def do_init_bkg_wPSs(bkg_mod, llh_obj, src_tab, rt_obj, GTI, sig_twind):

    gti_bkg = add_bti2gti(sig_twind, GTI)
    bkg_t0s = gti_bkg['START']
    bkg_t1s = gti_bkg['STOP']

    Nsrcs = len(src_tab)
    nebins = bkg_mod.nebins

    for ii in range(Nsrcs):
        mod_list = [bkg_mod]
        im_steps = 5
        if Nsrcs >= 3:
            im_steps = 3
        if Nsrcs >= 5:
            im_steps = 2
        if Nsrcs >= 7:
            im_steps = 1

        ps_mods = []
        for i in range(Nsrcs):
            row = src_tab[i]
            mod = Point_Source_Model_Binned_Rates(row['imx'], row['imy'], 0.1,\
                                                  [llh_obj.ebins0,llh_obj.ebins1],\
                                                  rt_obj, llh_obj.bl_dmask,\
                                                  use_deriv=True, name=row['Name'])
            ps_mods.append(mod)

        mod_list += ps_mods
        comp_mod = CompoundModel(mod_list)

        llh_obj.set_model(comp_mod)

        bf_nllh, bf_params, TS_nulls = bkg_withPS_fit(src_tab, comp_mod,\
                                      llh_obj, bkg_t0s, bkg_t1s,\
                                      test_null=True, im_steps=im_steps)

        logging.debug("TS_nulls: ")
        logging.debug(TS_nulls)

        bkg_rates = np.array([bf_params['Background'+'_bkg_rate_'+str(j)]\
                            for j in range(nebins)])
        min_rate = 1e-1*bkg_rates
        logging.debug("min_rate: ")
        logging.debug(min_rate)
        PSs2keep = []
        for name, TS in TS_nulls.items():
            if TS < 8.0:
                ps_rates = np.array([bf_params[name+'_rate_'+str(j)] for j in range(nebins)])
                logging.debug(name + " rates: ")
                logging.debug(ps_rates)
                # print ps_rates
                if np.all(ps_rates<min_rate):
                    continue
            PSs2keep.append(name)

        if len(PSs2keep) == len(src_tab):
            break
        if len(PSs2keep) == 0:
            Nsrcs = 0
            src_tab = src_tab[np.zeros(len(src_tab), dtype=np.bool)]
            break
        bl = np.array([src_tab['Name'][i] in PSs2keep for i in range(Nsrcs)])
        src_tab = src_tab[bl]
        Nsrcs = len(src_tab)
        logging.debug("src_tab: ")
        logging.debug(src_tab)

    return bf_params, src_tab

def get_info_mat_around_min(llh_obj, mod, params_, ebin):

    params = copy(params_)
    dt = llh_obj.dt
    mod_cnts = llh_obj.model.get_rate_dpi(params, ebin)*dt
    data_cnts = llh_obj.data_dpis[ebin]

    dR_dparams = mod.get_dr_dp(params, ebin)

    cov_ndim = len(dR_dparams)

    info_mat = np.zeros((cov_ndim,cov_ndim))

    for i in range(cov_ndim):
        for j in range(cov_ndim):
            info_mat[i,j] = np.sum(((dR_dparams[j]*dt)*(dR_dparams[i]*dt)*\
                                    data_cnts)/np.square(mod_cnts))
#             cov_mat[i,j] = np.sum(np.square(mod_cnts)/\
#                         ((dR_dparams[j]*dt)*(dR_dparams[i]*dt)*data_cnts))

    return info_mat

def get_errs_corrs(llh_obj, model, params, e0, pnames2skip=[]):

    imat = get_info_mat_around_min(llh_obj, model, copy(params), e0)
    cov_mat = np.linalg.inv(imat)
    e0_pnames = []
    for pname in model.param_names:
        try:
            if int(pname[-1])==e0 and not (pname in pnames2skip):
                e0_pnames.append(pname)
        except:
            pass

    err_dict = {}
    corr_dict = {}
    errs = np.sqrt(np.diag(cov_mat))
    for i, pname in enumerate(e0_pnames):
        k = 'err_' + pname
        err_dict[k] = errs[i]
    Npars = len(e0_pnames)
    for i in range(Npars-1):
        pname0 = e0_pnames[i]
        for j in range(i+1,Npars):
            pname1 = e0_pnames[j]
            k = 'corr_' + pname0 + '_' + pname1
            corr_dict[k] = cov_mat[i,j]/(errs[i]*errs[j])

    return err_dict, corr_dict

def bkg_withPS_fit_fiximxy(PS_tab, model, llh_obj, t0s, t1s, params_,\
                           fixed_pnames=None):

    Nps = len(PS_tab)
    params = copy(params_)

    llh_obj.set_model(model)
    bkg_miner = NLLH_ScipyMinimize_Wjacob('')
    bkg_miner.set_llh(llh_obj)
    llh_obj.set_time(t0s,t1s)

    nllh = 0.0
#     bf_params = {fixed_pars[i]:fixed_vals[i] for i in range(len(fixed_pars))}
    bf_params = copy(params)
    fixed_vals = [bf_params[pname] for pname in fixed_pnames]
    errs_dict = {}
    corrs_dict = {}

    for e0 in range(llh_obj.nebins):
        bkg_miner.set_fixed_params(bkg_miner.param_names)
        if fixed_pnames is not None:
            bkg_miner.set_fixed_params(fixed_pnames, values=fixed_vals)
        e0_pnames = []
        for pname in bkg_miner.param_names:
            try:
                if int(pname[-1])==e0 and not (pname in fixed_pnames):
                    e0_pnames.append(pname)
            except:
                pass
        bkg_miner.set_fixed_params(e0_pnames, fixed=False)
        llh_obj.set_ebin(e0)

        bf_vals, bkg_nllh, res = bkg_miner.minimize()
        nllh += bkg_nllh[0]
        for ii, pname in enumerate(e0_pnames):
            bf_params[pname] = bf_vals[0][ii]

        print(bf_params)

        err_dict, corr_dict = get_errs_corrs(llh_obj, model, copy(bf_params), e0, pnames2skip=fixed_pnames)
        for k, val in err_dict.items():
            errs_dict[k] = val
        for k, val in corr_dict.items():
            corrs_dict[k] = val

    return nllh, bf_params, errs_dict, corrs_dict

def main(args):

    logging.basicConfig(filename='bkg_rate_estimation_wPSs.log', level=logging.DEBUG,\
                    format='%(asctime)s-' '%(levelname)s- %(message)s')

    if args.bkg_nopost and args.bkg_nopre:
        raise Exception("Can't have no pre and no post")

    if args.dbfname is None:
        db_fname = guess_dbfname()
        if isinstance(db_fname, list):
            db_fname = db_fname[0]
    else:
        db_fname = args.dbfname

    logging.info('Connecting to DB')
    conn = get_conn(db_fname)

    info_tab = get_info_tab(conn)
    logging.info('Got info table')

    files_tab = get_files_tab(conn)
    logging.info('Got files table')

    trigtime = info_tab['trigtimeMET'][0]
    tstart = trigtime - args.twind
    tstop = trigtime + args.twind

    evfname = files_tab['evfname'][0]
    dmfname = files_tab['detmask'][0]
    attfname = files_tab['attfname'][0]
    ev_data = fits.open(evfname)[1].data
    try:
        GTI = Table.read(evfname, hdu='GTI_POINTING')
    except:
        GTI = Table.read(evfname, hdu='GTI')
    dmask = fits.open(dmfname)[0].data
    bl_dmask = (dmask==0.)
    attfile = fits.open(attfname)[1].data
    logging.debug('Opened up event, detmask, and att files')

    ebins0 = np.array(EBINS0)
    ebins1 = np.array(EBINS1)
    nebins = len(ebins0)
    logging.debug("ebins0")
    logging.debug(ebins0)
    logging.debug("ebins1")
    logging.debug(ebins1)

    solid_angle_dpi = np.load(solid_angle_dpi_fname)

    src_tab = get_srcs_infov(attfile, trigtime, pcfname=args.pcfname)
    Nsrcs = len(src_tab)
    logging.info("src_tab: ")
    logging.info(src_tab)

    bkg_mod = Bkg_Model_wFlatA(bl_dmask, solid_angle_dpi, nebins, use_deriv=True)

    llh_obj = LLH_webins(ev_data, ebins0, ebins1, bl_dmask)

    # add in stuff later for if there's no srcs

    if Nsrcs > 0:

        rt_dir = files_tab['rtDir'][0]
        rt_obj = RayTraces(rt_dir)

        sig_dtwind = (-10*1.024, 20*1.024)
        sig_twind = (trigtime + sig_dtwind[0], trigtime + sig_dtwind[1])

        init_bf_params, src_tab = do_init_bkg_wPSs(bkg_mod, llh_obj, src_tab, rt_obj, GTI, sig_twind)

        Nsrcs = len(src_tab)

        logging.debug("Final src_tab:")
        logging.debug(src_tab)

        # Now need to do each time, with these PSs and these imxys

    else:
        init_bf_params = {k:bkg_mod.param_dict[k]['val'] for k in bkg_mod.param_names}

    if Nsrcs > 0:
        fixed_pars = [pname for pname in list(init_bf_params.keys()) if '_flat_' in pname\
                    or '_imx' in pname or '_imy' in pname]

        mod_list = [bkg_mod]
        ps_mods = []
        for i in range(Nsrcs):
            row = src_tab[i]
            mod = Point_Source_Model_Binned_Rates(row['imx'], row['imy'], 0.1,\
                                                  [ebins0,ebins1], rt_obj, bl_dmask,\
                                                  use_deriv=True, name=row['Name'])
            ps_mods.append(mod)

        mod_list += ps_mods
        mod = CompoundModel(mod_list)

    else:
        init_bf_params = {k:bkg_mod.param_dict[k]['val'] for k in bkg_mod.param_names}
        mod = bkg_mod
        fixed_pars = []

    bkg_dur = args.bkg_dur*1.024
    twind = args.twind*1.024
    bkg_tstep = 2*1.024
    dt_ax = np.arange(-twind, twind+1, bkg_tstep)
    t_ax = dt_ax + trigtime
    Ntpnts = len(dt_ax)
    logging.info('Ntpnts: %d'%(Ntpnts))
    logging.info('min(dt_ax): %.3f'%(np.min(dt_ax)))
    logging.info('max(dt_ax): %.3f'%(np.max(dt_ax)))
    sig_wind = 10*1.024
    # sig_twind = (trigger_time + sig_dtwind[0], trigger_time + sig_dtwind[1])

    bkg_bf_dicts = []

    for i in range(Ntpnts):

        tmid = t_ax[i]
        sig_twind = (-sig_wind/2. + tmid, sig_wind/2. + tmid)
        gti_ = add_bti2gti(sig_twind, GTI)
        bkg_t0 = tmid - sig_wind/2. - bkg_dur/2.
        bkg_t1 = tmid + sig_wind/2. + bkg_dur/2.
        bkg_bti = Table(data=([-np.inf, bkg_t1], [bkg_t0, np.inf]), names=('START', 'STOP'))
        gti_ = add_bti2gti(bkg_bti, gti_)
        print(tmid - trigtime)
        print(gti_)
        t0s = gti_['START']
        t1s = gti_['STOP']

        nllh, params, errs_dict, corrs_dict = bkg_withPS_fit_fiximxy(src_tab, mod, llh_obj, t0s, t1s,\
                                                                     copy(init_bf_params),\
                                                                     fixed_pnames=fixed_pars)

        params.update(errs_dict)
        params.update(corrs_dict)
        params['nllh'] = nllh
        params['time'] = tmid
        params['dt'] = tmid - trigtime
        bkg_bf_dicts.append(params)
        params['exp'] = llh_obj.dt

    bkg_df = pd.DataFrame(bkg_bf_dicts)

    save_fname = 'bkg_estimation.csv'
    logging.info("Saving results in a DataFrame to file: ")
    logging.info(save_fname)
    bkg_df.to_csv(save_fname, index=False)

if __name__ == "__main__":

    args = cli()

    main(args)


from ase_interface import ANI
import pyNeuroChem as pync
import hdnntools as hdt
import pymolfrag as pmf
import numpy as np

# ASE Stuff
from ase.io import read, write
from ase.md.langevin import Langevin
from ase.optimize import BFGS, LBFGS
from ase import units

#wkdir = '/home/jujuman/Dropbox/ChemSciencePaper.AER/networks/ANI-c08f-ntwk-cv/'
wkdir = '/home/jujuman/Research/DataReductionMethods/model6r/model-gdb_r06_comb08/cv5/'
#wkdir = '/home/jujuman/Scratch/Research/WaterData/CV2/'
cnstfile = wkdir + 'rHCNO-4.6A_16-3.1A_a4-8.params'
saefile = wkdir + 'sae_6-31gd.dat'

stdir = '/home/jujuman/Research/WaterData/cv_md_2/'

T = 10.0
dt = 0.25

mol = read('/home/jujuman/Dropbox/ChemSciencePaper.AER/TestCases/water.pdb')
Nnt = 5

# Construct pyNeuroChem classes
print('Constructing CV network list...')
ncl =  [pync.molecule(cnstfile, saefile, wkdir + 'train' + str(l) + '/networks/', 0, False) for l in range(Nnt)]
print('Complete.')

L1 = 16.761
L2 = 16.592
L3 = 16.566
mol.set_cell(([[L1,0,0],[0,L2,0],[0,0,L3]]))
mol.set_pbc((True, True, True))

# Set the calculator
nc = pync.molecule(cnstfile, saefile, wkdir + 'train' + str(0) + '/networks/', 0, False)
mol.set_calculator(ANI(False))
mol.calc.setnc(nc)

dyn = LBFGS(mol)
dyn.run(fmax=1.0)

# We want to run MD with constant energy using the Langevin algorithm
# with a time step of 5 fs, the temperature T and the friction
# coefficient to 0.02 atomic units.
dyn = Langevin(mol, dt * units.fs, T * units.kB, 0.01)

mdcrd = open(stdir + "mdcrd.xyz",'w')
temp = open(stdir + "temp.dat",'w')

dyn.get_time()
def printenergy(a=mol,b=mdcrd,d=dyn,t=temp):  # store a reference to atoms in the
    """Function to print the potential, kinetic and total energy."""
    epot = a.get_potential_energy() / len(a)
    ekin = a.get_kinetic_energy() / len(a)
    print('Step %i - Energy per atom: Epot = %.3feV  Ekin = %.3feV (T=%3.0fK)  '
          'Etot = %.3feV' % (d.get_number_of_steps(),epot, ekin, ekin / (1.5 * units.kB), epot + ekin))
    t.write(str(d.get_number_of_steps()) + ' ' + str(d.get_time()) + ' ' + str(ekin / (1.5 * units.kB)) + ' ' + str(epot) + ' ' +  str(ekin) + ' ' + str(epot + ekin) + '\n')
    b.write('\n' + str(len(a)) + '\n')
    c=a.get_positions(wrap=True)
    for j,i in zip(a,c):
        b.write(str(j.symbol) + ' ' + str(i[0]) + ' ' + str(i[1]) + ' ' + str(i[2]) + '\n')

dyn.attach(printenergy, interval=10)
dyn.set_temperature(T * units.kB)

data = dict()
#dyn.run(4000)  # Do 100 steps of MD
for i in range(100):

    spc = mol.get_chemical_symbols().copy()
    xyz = np.array(mol.get_positions(wrap=False).copy(), dtype=np.float32).reshape(len(spc), 3)

    mf = pmf.molfrag(xyz,spc)

    ofrags1 = mf.get_all_frags('O', 2.7, 3.8)
    ofrags2 = mf.get_all_frags('O', 4.5, 5.6)

    ofrags = ofrags1 + ofrags2
    Nb = 0
    Nf = 0
    Ne = 0
    Ftot = 0.0
    Etot = 0.0
    for i,of in enumerate(ofrags):
        spc_l = of[0]
        xyz_l = of[1]

        energies = np.zeros((Nnt), dtype=np.float64)
        forces = []
        N = 0
        for comp in ncl:
            comp.setMolecule(coords=xyz_l, types=list(spc_l))
            energies[N] = comp.energy()[0]
            forces.append(comp.force().reshape(1, len(list(spc_l)), 3))
            N = N + 1

        Ftot = Ftot + np.mean(np.std(np.concatenate(forces),axis=0))

        energies = hdt.hatokcal * energies

        sigma = np.std(energies) / float(len(spc_l))
        Etot = Etot + sigma

        No = len([s for s in spc_l if s == 'O'])
        Nh = len([s for s in spc_l if s == 'H'])

        Nf = Nf + 1
        if Nh == 2 * No:

            key = str(No).zfill(3) + str(Nh).zfill(3)
            if sigma > 0.1:
                Nb = Nb + 1
                if key not in data.keys():
                    data[key] = (spc_l,[xyz_l])
                else:
                    #print(type())
                    p_xyz = data[key][1]
                    p_xyz.append(xyz_l)
                    data[key] = (spc_l, p_xyz)

                '''
                print('H2O_frag' + str(i).zfill(3) + '.xyz : O=' + str(No).zfill(3) + ' : H=' + str(Nh).zfill(3) + ' : '
                            + "{:13.3f}".format(energies[0]) +
                        ' ' + "{:13.3f}".format(energies[1]) +
                        ' ' + "{:13.3f}".format(energies[2]) +
                        ' ' + "{:13.3f}".format(energies[3]) +
                        ' ' + "{:13.3f}".format(energies[4]) +
                        ' ' + "{:.3f}".format(sigma))
                '''
        else:
            Ne = Ne + 1
            print('Warning: fragmentation failed! No:',No,'Nh:',Nh)
    print('Bad Frags:', Nb, 'of',Nf,'Error:',Ne,'Estd:',Etot/float(Nf), ' Avg. Frc. Std.: ', Ftot/float(Nf))
    dyn.run(40)  # Do 100 steps of MD

for key in data:
    xyz = np.array(data[key][1])
    spc = data[key][0]
    print(key,xyz.shape[0])
    hdt.writexyzfile(stdir + 'frag' + key + '.xyz', xyz, spc)


# -*- coding: utf-8 -*-
# --------------------------------------
# @Time    : 2020/11/01
# @Author  : Oscar Chen
# @Email   : 530824679@qq.com
# @File    : network.py
# Description :YOLO v5 network architecture
# --------------------------------------

import numpy as np
import tensorflow as tf
from cfg.config import *
from model.basenet import *
from utils.data_utils import *

class Network(object):
    def __init__(self, is_train):
        self.is_train = is_train
        self.anchors = read_anchors(path_params['anchor_file'])
        self.classes = read_class_names(path_params['class_file'])
        self.class_num = len(self.classes)
        self.strides = np.array(model_params['strides'])
        self.anchor_per_scale = model_params['anchor_per_scale']

    def forward(self, inputs):
        try:
            conv_lbbox, conv_mbbox, conv_sbbox = self.build_network(inputs)
        except:
            raise NotImplementedError("Can not build up yolov5 network!")

        with tf.variable_scope('pred_sbbox'):
            pred_sbbox = self.reorg_layer(conv_sbbox, self.anchors[0], self.strides[0])

        with tf.variable_scope('pred_mbbox'):
            pred_mbbox = self.reorg_layer(conv_mbbox, self.anchors[1], self.strides[1])

        with tf.variable_scope('pred_lbbox'):
            pred_lbbox = self.reorg_layer(conv_lbbox, self.anchors[2], self.strides[2])

        logits = [conv_sbbox, conv_mbbox, conv_lbbox]
        preds = [pred_sbbox, pred_mbbox, pred_lbbox]
        return logits, preds

    def build_network(self, inputs):

        # backbone
        focus_0 = focus(inputs, 64, 3, 'model/0')
        conv_1 = convBnLeakly(focus_0, 128, 3, 2, "model/1")
        bottleneck_csp_2 = bottleneckCSP(conv_1, 128, 128, 3, True, 0.5, "model/2")
        conv_3 = convBnLeakly(bottleneck_csp_2, 256, 3, 2, 'model/3')
        bottleneck_csp_4 = bottleneckCSP(conv_3, 256, 256, 9, True, 0.5, 'model/4')
        conv_5 = convBnLeakly(bottleneck_csp_4, 512, 3, 2, 'model/5')
        bottleneck_csp_6 = bottleneckCSP(conv_5, 512, 512, 9, True, 0.5, 'model/6')
        conv_7 = convBnLeakly(bottleneck_csp_6, 1024, 3, 2, 'model/7')
        spp_8 = spp(conv_7, 1024, 1024, 5, 9, 13, 'model/8')

        # neck
        bottleneck_csp_9 = bottleneckCSP(spp_8, 1024, 1024, 3, False, 0.5, 'model/9')
        conv_10 = convBnLeakly(bottleneck_csp_9, 512, 1, 1, 'model/10')

        shape = [conv_10.shape[1].value * 2, conv_10.shape[2].value * 2]
        # 0：双线性差值。1：最近邻居法。2：双三次插值法。3：面积插值法
        deconv_11 = tf.image.resize_images(conv_10, shape, method=1)

        concat_12 = tf.concat((deconv_11, bottleneck_csp_6), -1)
        bottleneck_csp_13 = bottleneckCSP(concat_12, 1024, 512, 3, False, 0.5, 'model/13')
        conv_14 = convBnLeakly(bottleneck_csp_13, 256, 1, 1, 'model/14')

        shape = [conv_14.shape[1].value * 2, conv_14.shape[2].value * 2]
        deconv_15 = tf.image.resize_images(conv_14, shape, method=1)

        concat_16 = tf.concat((deconv_15, bottleneck_csp_4), -1)
        bottleneck_csp_17 = bottleneckCSP(concat_16, 512, 256, 3, False, 0.5, 'model/17')
        conv_18 = convBnLeakly(bottleneck_csp_17, 256, 3, 2, 'model/18')

        concat_19 = tf.concat((conv_18, conv_14), -1)
        bottleneck_csp_20 = bottleneckCSP(concat_19, 512, 512, 3, False, 0.5, 'model/20')
        conv_21 = convBnLeakly(bottleneck_csp_20, 512, 3, 2, 'model/21')

        concat_22 = tf.concat((conv_21, conv_10), -1)
        bottleneck_csp_23 = bottleneckCSP(concat_22, 1024, 1024, 3, False, 0.5, 'model/23')

        # head
        conv_24_m0 = conv(bottleneck_csp_17, 3 * (self.class_num + 5), 1, 1, 'model/24/m/0', add_bias=True)
        conv_24_m1 = conv(bottleneck_csp_20, 3 * (self.class_num + 5), 1, 1, 'model/24/m/1', add_bias=True)
        conv_24_m2 = conv(bottleneck_csp_23, 3 * (self.class_num + 5), 1, 1, 'model/24/m/2', add_bias=True)

        return conv_24_m0, conv_24_m1, conv_24_m2

    def reorg_layer(self, feature_maps, anchors, strides):
        """
        解码网络输出的特征图
        :param feature_maps:网络输出的特征图
        :param anchors:当前层使用的anchor尺度
        :param stride:特征图相比原图的缩放比例
        :return: 预测层最终的输出 shape=[batch_size, feature_size, feature_size, anchor_per_scale, 5 + class_num]
        """
        feature_shape = tf.shape(feature_maps)[1:3]
        batch_size = tf.shape(feature_maps)[0]
        anchor_per_scale = len(anchors)

        predict = tf.reshape(feature_maps, (batch_size, feature_shape[0], feature_shape[1], anchor_per_scale, 5 + self.class_num))
        conv_raw_xy = predict[:, :, :, :, 0:2]
        conv_raw_wh = predict[:, :, :, :, 2:4]
        conv_raw_conf = predict[:, :, :, :, 4:5]
        conv_raw_prob = predict[:, :, :, :, 5:]

        y = tf.tile(tf.range(feature_shape[0], dtype=tf.int32)[:, tf.newaxis], [1, feature_shape[0]])
        x = tf.tile(tf.range(feature_shape[1], dtype=tf.int32)[tf.newaxis, :], [feature_shape[1], 1])

        xy_cell = tf.concat([x[:, :, tf.newaxis], y[:, :, tf.newaxis]], axis=-1)
        xy_cell = tf.tile(xy_cell[tf.newaxis, :, :, tf.newaxis, :], [batch_size, 1, 1, anchor_per_scale, 1])
        xy_cell = tf.cast(xy_cell, tf.float32)

        bboxes_xy = (tf.sigmoid(conv_raw_xy) + xy_cell) * strides
        bboxes_wh = (tf.sigmoid(conv_raw_wh) * anchors) * strides

        pred_xywh = tf.concat([bboxes_xy, bboxes_wh], axis=-1)
        pred_box_confidence = tf.sigmoid(conv_raw_conf)
        pred_box_class_prob = tf.sigmoid(conv_raw_prob)
        return tf.concat([pred_xywh, pred_box_confidence, pred_box_class_prob], axis=-1)


import cv2
import pycuda.autoinit
from multiprocessing import Process,Queue

from utils.yolo_with_plugins import TrtYOLO
from utils.yolo_classes import get_cls_dict
from utils.camera import add_camera_args, Camera
from utils.display import open_window, set_display, show_fps
from utils.visualization import BBoxVisualization

import pyaudio
import os
import struct
import numpy as np
from scipy.fftpack import fft
import time

Qa=Queue()

def audio():
    C = 1024 * 2
    F = pyaudio.paInt16
    CH = 1
    R = 44100

    p = pyaudio.PyAudio()

    stream = p.open(
    format=F,
    channels=CH,
    rate=R,
    input=True,
    output=True,
    frames_per_buffer=C
    )
    xf = np.linspace(0, R, C)
    while True:
        data = stream.read(C)
        data_int = struct.unpack(str(2*C) + 'B', data)
        #data_np = np.array(data_int, dtype='b')[::2] + 128
        yf = fft(data_int)
        ydata=np.abs(yf[0:C]) / (128 * C)
        for x,y in zip(xf,ydata):
            if x>600 and x<1300 and y>0.4:
                Qa.put(1)

if __name__ == '__main__':
    trt_yolo = TrtYOLO('yolov4-tiny-detektor-416', (416, 416), 3)
    cap = cv2.VideoCapture(0)  # USB WebCam 0
    p1=Process(target=audio)
    p1.start()
    while True:
        _, img = cap.read()
        if img is None:  break
        middleImg=img.shape[1]/2
        Det=trt_yolo.detect(img)
        if Det[0].size==0 and Qa.get()==1:
            print("pojazd na sygnale")
        for i,j,k in zip(Det[0],Det[2],Det[1]):
            if Qa.get()==1 and i.size!=0 and k>0.5:
                middleDet = i[0] + (i[2]-i[0])/2
                if middleDet>middleImg+middleImg/10 and j==0:
                    print("karetka z prawej strony")
                if middleDet<middleImg-middleImg/10 and j==0:
                    print("karetka z lewej strony")
                if middleDet>middleImg-middleImg/10 and middleDet<middleImg+middleImg/10 and j==0:
                    print("karetka po środku")
                if middleDet>middleImg+middleImg/10 and j==1:
                    print("radiowóz z prawej strony")
                if middleDet<middleImg-middleImg/10 and j==1:
                    print("radiowóz z lewej strony")
                if middleDet>middleImg-middleImg/10 and middleDet<middleImg+middleImg/10 and j==1:
                    print("radiowóz po środku")
                if middleDet>middleImg+middleImg/10 and j==2:
                    print("wóz strażacki z prawej strony")
                if middleDet<middleImg-middleImg/10 and j==2:
                    print("wóz strażacki z lewej strony")
                if middleDet>middleImg-middleImg/10 and middleDet<middleImg+middleImg/10 and j==2:
                    print("wóz strażacki po środku")

    p1.join()


from keras.models import Sequential
from keras.layers import Dense,Embedding, Dropout, Flatten, Conv1D,GlobalMaxPooling1D
import numpy as np

class CNN:
    def __init__(self,x_train,x_test,y_train,y_test=None,vocab_size=None,typ=None,max_length=None,embed_matrix=None):
        self.model = Sequential()
        self.x_train = x_train
        self.y_train = y_train
        self.x_test = x_test
        self.y_test = y_test
        self.vocab_size = vocab_size
        self.typ = typ
        self.max_length = max_length
        self.embed_matrix = embed_matrix

    def evaluate(self):
        print ("Evaluating Training Data")
        self.model.add(Embedding(self.vocab_size,200,input_length=self.max_length,weights=[self.embed_matrix]))
        self.model.add(Conv1D(filters=100, kernel_size=2, padding='valid', activation='relu', strides=1))
        self.model.add(GlobalMaxPooling1D())
        self.model.add(Dense(128, activation='relu'))
        self.model.add(Dropout(0.5))
        self.model.add(Dense(128, activation='relu'))
        self.model.add(Dense(3,activation="softmax"))
        self.model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['categorical_accuracy'])
        self.model.fit(self.x_train, self.y_train, epochs=25, batch_size=32, verbose=0)
        scores = self.model.predict(self.x_test, batch_size=32, verbose=0)
        prediction = np.subtract(np.argmax(scores, axis=1), 1)
        print("End Evaluating Test Data")
        return prediction


import tensorflow as tf
import numpy as np
import tensorflow_hub as hub
import json

#Return loaded model
def load_model(modelpath):
    model = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
    return model

# return prediction based on loaded model (from the step above) and an input payload
def predict(model, payload):
    try:
        if(type(payload) == str):
            data = [payload]
        else:
            data = [payload.decode()]# Multi model endpoints -> [payload[0]['body'].decode()]

        out = np.asarray(model(data)).tolist()
    except Exception as e:
        out = str(e)
    return [json.dumps({'output':[out],'tfeager': tf.executing_eagerly()})]