import inspect
import numpy as np
import os
import copy

from astropy import stats
from abc import ABCMeta

from pypeit import ginga, utils, msgs, processimages, specobjs
from pypeit.core import skysub, extract, trace_slits, pixels, wave

from pypeit import debugger

class Reduce(object):
    """
     This class will organize and run actions related to
     a Science or Standard star exposure

     Args:
         file_list : list
           List of raw files to produce the flat field
         spectrograph : str
         settings : dict-like
         tslits_dict : dict
           dict from TraceSlits class
         tilts : ndarray
           tilts from WaveTilts class
           used for sky subtraction and object finding
         det : int
         setup : str
         datasec_img : ndarray
           Identifies pixels to amplifiers
         bpm : ndarray
           Bad pixel mask
         maskslits : ndarray (bool)
           Specifies masked out slits
         pixlocn : ndarray
         objtype : str
           'science'
           'standard'
         scidx : int
           Row in the fitstbl corresponding to the exposure

     Attributes:
         frametype : str
           Set to 'science'
         sciframe : ndarray
           Processed 2D frame
         rawvarframe : ndarray
           Variance generated without a sky (or object) model
         modelvarframe : ndarray
           Variance generated with a sky model
         finalvar : ndarray
           Final variance frame
         global_sky : ndarray
           Sky model across the slit/order
         skycorr_box : ndarray
           Local corrections to the sky model
         final_sky : ndarray
           Final sky model; may include 'local' corrections
         obj_model : ndarray
           Model of the object flux
         trcmask : ndarray
           Masks of objects for sky subtraction
         tracelist : list
           List of traces for objects in slits
         inst_name : str
           Short name of the spectrograph, e.g. KASTb
         target_name : str
           Parsed from the Header
         basename : str
           Combination of camera, target, and time
           e.g. J1217p3905_KASTb_2015May20T045733.56
         time : Time
           time object
         specobjs : list
           List of specobjs
         bm: ScienceImageBitMask
           Object used to select bits of a given type
     """

    __metaclass__ = ABCMeta

    def __init__(self, spectrograph, tslits_dict, mask, par, ir_redux=False, det=1,
                 objtype='science', binning=None, setup=None, maskslits=None):

        # Setup the parameters sets for this object. NOTE: This uses objtype, not frametype!
        self.objtype = objtype
        self.par = par
        # TODO -- I find it very confusing to break apart the main parset
        self.proc_par = self.par['scienceframe']['process']
        # TODO Rename the scienceimage arset to reduce.
        self.redux_par = self.par['scienceimage']
        self.wave_par = self.par['calibrations']['wavelengths']
        self.flex_par = self.par['flexure']

        # Instantiation attributes for this object
        self.spectrograph = spectrograph
        self.tslits_dict = tslits_dict
        self.mask = mask
        self.slitmask = pixels.tslits2mask(self.tslits_dict)
        # Now add the slitmask to the mask (i.e. post CR rejection in proc)
        self.mask = processimages.ProcessImages.update_mask_slitmask(self.mask, self.slitmask)
        self.maskslits=None
        self.maskslits = self._get_goodslits(maskslits)
        self.ir_redux = ir_redux
        self.det = det
        self.binning = binning
        self.setup = setup
        self.pypeline = spectrograph.pypeline

        self.steps = []

        # Other attributes that will be set later during object finding,
        # sky-subtraction, and extraction
        self.tilts = None  # used by extract
        self.waveimage = None  # used by extract

        # Key outputs images for extraction
        self.sciimg = None
        self.sciivar = None
        self.ivarmodel = None
        self.objimage = None
        self.skyimage = None
        self.global_sky = None
        self.skymask = None
        self.outmask = None
        self.extractmask = None
        # SpecObjs object
        self.sobjs_obj = None  # Only object finding but no extraction
        self.sobjs = None  # Final extracted object list with trace corrections applied

        # Other bookeeping internals
        self.crmask = None

    def _chk_objs(self, items):
        """

        Args:
            items:

        Returns:

        """
        for obj in items:
            if getattr(self, obj) is None:
                msgs.warn('You need to generate {:s} prior to this step..'.format(obj))
                if obj in ['sciimg', 'sciivar', 'rn2_img']:
                    msgs.warn('Run the process() method')
                elif obj in ['sobjs_obj']:
                    msgs.warn('Run the find_objects() method')
                elif obj in['global_sky']:
                    msgs.warn('Run the global_skysub() method')
                elif obj in ['tilts', 'tslits_dict'] :
                    msgs.warn('Calibrations missing: these were required to run find_objects() '
                              'and global_skysub()')
                elif obj in ['waveimg']:
                    msgs.warn('Calibrations missing: waveimg must be input as a parameter. Try '
                              'running calibrations')
                return False
        return True

    def parse_manual_dict(self, manual_dict, neg=False):
        """
        Parse the manual dict
        This method is here mainly to deal with negative images

        Args:
            manual_dict (dict or None):
            neg (bool, optional):

        Returns:
            None or dict:  None if no matches; dict if there are for manual extraction

        """
        if manual_dict is None:
            return None
        #
        dets = manual_dict['hand_extract_det']
        # Grab the ones we want
        gd_det = dets > 0
        if neg:
            gd_det = np.invert(gd_det)
        # Any?
        if not np.any(gd_det):
            return None
        # Fill
        manual_extract_dict = {}
        for key in manual_dict.keys():
            sgn = 1
            if key == 'hand_extract_det':
                sgn = -1
            manual_extract_dict[key] = sgn*manual_dict[key][gd_det]
        # Return
        return manual_extract_dict

    def find_objects(self, image, ivar, std=False, ir_redux=False, std_trace=None, maskslits=None,
                          show_peaks=False, show_fits=False, show_trace=False, show=False,
                     manual_extract_dict=None):

        # Positive image
        parse_manual = self.parse_manual_dict(manual_extract_dict, neg=False)
        sobjs_obj_init, nobj_init, skymask_pos = \
            self.find_objects_pypeline(image, ivar, std=std, std_trace=std_trace, maskslits=maskslits,
                                   show_peaks = show_peaks, show_fits = show_fits, show_trace = show_trace,
                                       manual_extract_dict=parse_manual)

        # For nobj we take only the positive objects
        if ir_redux:
            # Parses
            parse_manual = self.parse_manual_dict(manual_extract_dict, neg=True)
            sobjs_obj_init_neg, nobj_init_neg, skymask_neg = \
                self.find_objects_pypeline(-image, ivar, std=std, std_trace=std_trace, maskslits=maskslits,
                show_peaks=show_peaks, show_fits=show_fits, show_trace=show_trace,
                                           manual_extract_dict=parse_manual)
            skymask = skymask_pos & skymask_neg
            sobjs_obj_init.append_neg(sobjs_obj_init_neg)
        else:
            skymask = skymask_pos

        if show:
            self.show('image', image=image*(self.mask == 0), chname='objfind',sobjs=sobjs_obj_init, slits=True)

        # For nobj we take only the positive objects
        return sobjs_obj_init, nobj_init, skymask

    def find_objects_pypeline(self, image, ivar, std=False, std_trace=None, maskslits=None,
                              show_peaks=False, show_fits=False, show_trace=False, show=False, debug=False,
                              manual_extract_dict=None):

        """
         Dummy method for object finding. Overloaded by class specific object finding.

         Returns:

         """
        return None, None, None

    def global_skysub(self, sciimg, sciivar, tilts, std=False, skymask=None, update_crmask=True, maskslits=None, show_fit=False,
                      show=False, show_objs=False):
        """
        Perform global sky subtraction, slit by slit

        Wrapper to skysub.global_skysub

        Parameters
        ----------
        tslits_dict: dict
           Dictionary containing information on the slits traced for this image

        Optional Parameters
        -------------------
        bspline_spaceing: (float):
           Break-point spacing for bspline

        use_skymask: (bool, optional):
           Mask objects using self.skymask if object finding has been run
           (This requires they were found previously, i.e. that find_objects was already run)

        Returns:
            global_sky: (numpy.ndarray) image of the the global sky model
        """

        # Prep
        self.global_sky = np.zeros_like(sciimg)
        if (std and not self.redux_par['global_sky_std']):
            msgs.info('Skipping global sky-subtraction for standard star.')
            return self.global_sky

        if std:
            sigrej = 7.0
            update_crmask = False
        else:
            sigrej = 3.0

        self.sciimg = sciimg
        self.sciivar = sciivar
        self.tilts = tilts
        self.maskslits = self.maskslits if maskslits is None else maskslits
        gdslits = np.where(np.invert(self.maskslits))[0]

        # Mask objects using the skymask? If skymask has been set by objfinding, and masking is requested, then do so
        skymask_now = skymask if (skymask is not None) else np.ones_like(self.sciimg, dtype=bool)
        # Loop on slits
        for slit in gdslits:
            msgs.info("Global sky subtraction for slit: {:d}".format(slit))
            thismask = (self.slitmask == slit)
            inmask = (self.mask == 0) & thismask & skymask_now
            # Find sky
            self.global_sky[thismask] = skysub.global_skysub(self.sciimg, self.sciivar,
                                                             self.tilts, thismask,
                                                             self.tslits_dict['slit_left'][:,slit],
                                                             self.tslits_dict['slit_righ'][:,slit],
                                                             inmask=inmask,
                                                             sigrej=sigrej,
                                                             bsp=self.redux_par['bspline_spacing'],
                                                             no_poly=self.redux_par['no_poly'],
                                                             pos_mask = (not self.ir_redux),
                                                             show_fit=show_fit)
            # Mask if something went wrong
            if np.sum(self.global_sky[thismask]) == 0.:
                self.maskslits[slit] = True

        if update_crmask:
            self.crmask = processimages.ProcessImages.build_crmask(self.sciimg - self.global_sky, self.proc_par,
                                                                   self.det, self.spectrograph, ivar = self.sciivar,
                                                                   binning=self.binning)
            # Rebuild the mask with this new crmask
            self.mask = processimages.ProcessImages.update_mask_cr(self.mask, self.crmask)

        # Step
        self.steps.append(inspect.stack()[0][3])

        if show:
            sobjs_show = None if show_objs else self.sobjs_obj
            # Global skysub is the first step in a new extraction so clear the channels here
            self.show('global', slits=True, sobjs =sobjs_show, clear=False)

        # Return
        return self.global_sky

    def local_skysub_extract(self, sciimg, sciivar, tilts, waveimg, global_sky, rn2img, sobjs,
                             maskslits=None, model_noise=True, std=False,
                             show_profile=False, show_resids=False, show=False):

        """
         Dummy method for locak skysubtraction and extraction. Overloaded by class specific skysub and extraction.

         Returns:

         """

        return None, None, None, None, None

    def flexure_correct(self, sobjs, basename):
        """ Correct for flexure

        Spectra are modified in place (wavelengths are shifted)

        Args:
            sobjs: SpecObjs object
            maskslits: ndarray

        """

        if self.flex_par['method'] != 'skip':
            flex_list = wave.flexure_obj(sobjs, self.maskslits, self.flex_par['method'],
                                         self.flex_par['spectrum'],
                                         mxshft=self.flex_par['maxshift'])
            # QA
            wave.flexure_qa(sobjs, self.maskslits, basename, self.det, flex_list,out_dir=self.par['rdx']['redux_path'])
        else:
            msgs.info('Skipping flexure correction.')

    def helio_correct(self, sobjs, radec, obstime):
        """ Perform a heliocentric correction """
        # Helio, correct Earth's motion
        if (self.wave_par['frame'] in ['heliocentric', 'barycentric']) \
                and (self.wave_par['reference'] != 'pixel'):
            # TODO change this keyword to refframe instead of frame
            msgs.info("Performing a {0} correction".format(self.wave_par['frame']))
            vel, vel_corr = wave.geomotion_correct(sobjs, radec, obstime, self.maskslits,
                                                   self.spectrograph.telescope['longitude'],
                                                   self.spectrograph.telescope['latitude'],
                                                   self.spectrograph.telescope['elevation'],
                                                   self.wave_par['frame'])
        else:
            msgs.info('A wavelength reference-frame correction will not be performed.')
            vel_corr = None

        return vel_corr

    def _get_goodslits(self, maskslits):
        """
        Return the slits to be reduce by going through the maskslits
        logic below. If the input maskslits is None it uses previously
        assigned maskslits

        Returns
        -------
        gdslits
            numpy array of slit numbers to be reduced
        """

        # Identify the slits that we want to consider.
        if maskslits is not None:
            # If maskslits was passed in use it, and update self
            self.maskslits = maskslits
        elif (self.maskslits is None):
            # If maskslits was not passed, and it does not exist in self, reduce all slits
            self.maskslits = np.zeros(self.tslits_dict['slit_left'].shape[1], dtype=bool)
        else: # Otherwise, if self.maskslits exists, use the previously set maskslits
            pass
        return self.maskslits

    def show(self, attr, image=None, showmask=False, sobjs=None, chname=None, slits=False,clear=False):
        """
        Show one of the internal images

        .. todo::
            Should probably put some of these in ProcessImages

        Parameters
        ----------
        attr : str
          global -- Sky model (global)
          sci -- Processed science image
          rawvar -- Raw variance image
          modelvar -- Model variance image
          crmasked -- Science image with CRs set to 0
          skysub -- Science image with global sky subtracted
          image -- Input image
        display : str, optional
        image : ndarray, optional
          User supplied image to display

        Returns
        -------

        """

        if showmask:
            mask_in = self.mask
            bitmask_in = self.bitmask
        else:
            mask_in = None
            bitmask_in = None

        if attr == 'global':
            # global sky subtraction
            if self.sciimg is not None and self.global_sky is not None and self.mask is not None:
                # sky subtracted image
                image = (self.sciimg - self.global_sky)*(self.mask == 0)
                mean, med, sigma = stats.sigma_clipped_stats(image[self.mask == 0], sigma_lower=5.0,
                                                       sigma_upper=5.0)
                cut_min = mean - 1.0 * sigma
                cut_max = mean + 4.0 * sigma
                ch_name = chname if chname is not None else 'global_sky_{}'.format(self.det)
                viewer, ch = ginga.show_image(image, chname=ch_name, bitmask=bitmask_in,
                                              mask=mask_in, clear=clear, wcs_match=True)
                                              #, cuts=(cut_min, cut_max))
        elif attr == 'local':
            # local sky subtraction
            if self.sciimg is not None and self.skymodel is not None and self.mask is not None:
                # sky subtracted image
                image = (self.sciimg - self.skymodel)*(self.mask == 0)
                mean, med, sigma = stats.sigma_clipped_stats(image[self.mask == 0], sigma_lower=5.0,
                                                       sigma_upper=5.0)
                cut_min = mean - 1.0 * sigma
                cut_max = mean + 4.0 * sigma
                ch_name = chname if chname is not None else 'local_sky_{}'.format(self.det)
                viewer, ch = ginga.show_image(image, chname=ch_name, bitmask=bitmask_in,
                                              mask=mask_in, clear=clear, wcs_match=True)
                                              #, cuts=(cut_min, cut_max))
        elif attr == 'sky_resid':
            # sky residual map with object included
            if self.sciimg is not None and self.skymodel is not None \
                    and self.objmodel is not None and self.ivarmodel is not None \
                    and self.mask is not None:
                image = (self.sciimg - self.skymodel) * np.sqrt(self.ivarmodel)
                image *= (self.mask == 0)
                ch_name = chname if chname is not None else 'sky_resid_{}'.format(self.det)
                viewer, ch = ginga.show_image(image, chname=ch_name, cuts=(-5.0, 5.0),
                                              bitmask=bitmask_in, mask=mask_in, clear=clear,
                                              wcs_match=True)
        elif attr == 'resid':
            # full residual map with object model subtractede
            if self.sciimg is not None and self.skymodel is not None \
                    and self.objmodel is not None and self.ivarmodel is not None \
                    and self.mask is not None:
                # full model residual map
                image = (self.sciimg - self.skymodel - self.objmodel) * np.sqrt(self.ivarmodel)
                image *= (self.mask == 0)
                ch_name = chname if chname is not None else 'resid_{}'.format(self.det)
                viewer, ch = ginga.show_image(image, chname=ch_name, cuts=(-5.0, 5.0),
                                              bitmask=bitmask_in, mask=mask_in, clear=clear,
                                              wcs_match=True)
        elif attr == 'image':
            ch_name = chname if chname is not None else 'image'
            viewer, ch = ginga.show_image(image, chname=ch_name, clear=clear, wcs_match=True)
        else:
            msgs.warn("Not an option for show")

        if sobjs is not None:
            for spec in sobjs:
                color = 'magenta' if spec.hand_extract_flag else 'orange'
                ginga.show_trace(viewer, ch, spec.trace_spat, spec.idx, color=color)

        if slits:
            if self.tslits_dict is not None:
                slit_ids = [trace_slits.get_slitid(
                    self.mask.shape, self.tslits_dict['slit_left'],
                    self.tslits_dict['slit_righ'], ii)[0] for ii in range(self.tslits_dict['slit_left'].shape[1])]

                ginga.show_slits(viewer, ch, self.tslits_dict['slit_left'], self.tslits_dict['slit_righ'],
                                 slit_ids)  # , args.det)

    def __repr__(self):
        txt = '<{:s}: nimg={:d}'.format(self.__class__.__name__,
                                        self.nsci)
        if len(self.steps) > 0:
            txt+= ' steps: ['
            for step in self.steps:
                txt += '{:s}, '.format(step)
            txt = txt[:-2]+']'  # Trim the trailing comma
        txt += '>'
        return txt

class MultiSlit(Reduce):
    """
    Child of Reduce for Multislit and Longslit reductions

    """
    def __init__(self, spectrograph, tslits_dict, mask, par, **kwargs):
        super(MultiSlit, self).__init__(spectrograph, tslits_dict, mask, par, **kwargs)

    def find_objects_pypeline(self, image, ivar, std=False, std_trace = None, maskslits=None,
                              manual_extract_dict=None,
                              show_peaks=False, show_fits=False, show_trace=False, show=False, debug=False):
        """
        Find objects in the slits. This is currently setup only for ARMS

        Wrapper to extract.objfind

        Parameters
        ----------
        tslits_dict: dict
           Dictionary containing information on the slits traced for this image

        Optional Parameters
        -------------------
        SHOW_PEAKS:  bool
          Generate QA showing peaks identified by object finding

        SHOW_FITS:  bool
          Generate QA  showing fits to traces

        SHOW_TRACE:  bool
          Generate QA  showing traces identified. Requires an open ginga RC modules window

        Returns
        -------
        specobjs : Specobjs object
            Container holding Specobj objects
        nobj:
            Number of objects identified
        self.skymask : ndarray
                Boolean image indicating which pixels are useful for global sky subtraction

        """

        self.maskslits = self.maskslits if maskslits is None else maskslits
        gdslits = np.where(np.invert(self.maskslits))[0]

        # create the ouptut image for skymask
        skymask = np.zeros_like(image, dtype=bool)
        # Instantiate the specobjs container
        sobjs = specobjs.SpecObjs()

        # Loop on slits
        for slit in gdslits:
            qa_title ="Finding objects on slit # {:d}".format(slit)
            msgs.info(qa_title)
            thismask = (self.slitmask == slit)
            inmask = (self.mask == 0) & thismask
            # Find objects
            specobj_dict = {'setup': self.setup, 'slitid': slit, 'orderindx': 999,
                            'det': self.det, 'objtype': self.objtype, 'pypeline': self.pypeline}

            # TODO we need to add QA paths and QA hooks. QA should be
            # done through objfind where all the relevant information
            # is. This will be a png file(s) per slit.

            sig_thresh = 30.0 if std else self.redux_par['sig_thresh']
            #
            sobjs_slit, skymask[thismask] = \
                extract.objfind(image, thismask, self.tslits_dict['slit_left'][:,slit],self.tslits_dict['slit_righ'][:,slit],
                inmask=inmask, ncoeff=self.redux_par['trace_npoly'],
                std_trace=std_trace, sig_thresh=sig_thresh, hand_extract_dict=manual_extract_dict, #self.redux_par['manual'],
                specobj_dict=specobj_dict, show_peaks=show_peaks,show_fits=show_fits, show_trace=show_trace,
                qa_title=qa_title, nperslit=self.redux_par['maxnumber'])
            sobjs.add_sobj(sobjs_slit)

        # Steps
        self.steps.append(inspect.stack()[0][3])
        if show:
            self.show('image', image=image*(self.mask == 0), chname = 'objfind',
                      sobjs=sobjs, slits=True)

        # Return
        return sobjs, len(sobjs), skymask

    # JFH TODO Should we reduce the number of iterations for standards or near-IR redux where the noise model is not
    # being updated?
    def local_skysub_extract(self, sciimg, sciivar, tilts, waveimg, global_sky, rn2img, sobjs,
                             spat_pix=None, maskslits=None, model_noise=True, std = False,
                             show_profile=False, show=False):
        """
        Perform local sky subtraction, profile fitting, and optimal extraction slit by slit

        Wrapper to skysub.local_skysub_extract

        Parameters
        ----------
        sobjs: object
           Specobjs object containing Specobj objects containing information about objects found.
        waveimg: ndarray, shape (nspec, nspat)
           Wavelength map

        Optional Parameters
        -------------------

        Returns:
            global_sky: (numpy.ndarray) image of the the global sky model
        """

        self.sciimg = sciimg
        self.sciivar = sciivar
        self.tilts = tilts
        self.waveimg = waveimg
        self.global_sky = global_sky
        self.rn2img = rn2img

        # get the good slits and assign self.maskslits
        self.maskslits = self.maskslits if maskslits is None else maskslits
        gdslits = np.where(np.invert(self.maskslits))[0]

        # Allocate the images that are needed
        # Initialize to mask in case no objects were found
        self.outmask = np.copy(self.mask)
        # Initialize to input mask in case no objects were found
        self.extractmask = (self.mask == 0)
        # Initialize to zero in case no objects were found
        self.objmodel = np.zeros_like(self.sciimg)
        # Set initially to global sky in case no objects were found
        self.skymodel  = np.copy(self.global_sky)
        # Set initially to sciivar in case no obects were found.
        self.ivarmodel = np.copy(self.sciivar)

        # Could actually create a model anyway here, but probably
        # overkill since nothing is extracted

        self.sobjs = sobjs.copy()
        # Loop on slits
        for slit in gdslits:
            msgs.info("Local sky subtraction and extraction for slit: {:d}".format(slit))
            thisobj = (self.sobjs.slitid == slit) # indices of objects for this slit
            if np.any(thisobj):
                thismask = (self.slitmask == slit) # pixels for this slit
                # True  = Good, False = Bad for inmask
                inmask = (self.mask == 0) & thismask
                # Local sky subtraction and extraction
                self.skymodel[thismask], self.objmodel[thismask], self.ivarmodel[thismask], \
                    self.extractmask[thismask] = skysub.local_skysub_extract(
                    self.sciimg, self.sciivar, self.tilts, self.waveimg, self.global_sky, self.rn2img,
                    thismask, self.tslits_dict['slit_left'][:,slit], self.tslits_dict['slit_righ'][:, slit],
                    self.sobjs[thisobj], spat_pix=spat_pix, model_full_slit=self.redux_par['model_full_slit'],
                    box_rad=self.redux_par['boxcar_radius']/self.spectrograph.detector[self.det-1]['platescale'],
                    sigrej=self.redux_par['sky_sigrej'],
                    model_noise=model_noise, std=std, bsp=self.redux_par['bspline_spacing'],
                    sn_gauss=self.redux_par['sn_gauss'], inmask=inmask, show_profile=show_profile)

        # Set the bit for pixels which were masked by the extraction.
        # For extractmask, True = Good, False = Bad
        iextract = (self.mask == 0) & (self.extractmask == False)
        self.outmask[iextract] = processimages.ProcessImages.bitmask.turn_on(self.outmask[iextract], 'EXTRACT')

        # Step
        self.steps.append(inspect.stack()[0][3])

        if show:
            self.show('local', sobjs = self.sobjs, slits= True)
            self.show('resid', sobjs = self.sobjs, slits= True)

        # Return
        return self.skymodel, self.objmodel, self.ivarmodel, self.outmask, self.sobjs

class Echelle(Reduce):
    """
    Child of Reduce for Echelle reductions

    """
    def __init__(self, spectrograph, tslits_dict, mask, par, **kwargs):
        super(Echelle, self).__init__(spectrograph, tslits_dict, mask, par, **kwargs)

    def find_objects_pypeline(self, image, ivar, std=False, std_trace = None, maskslits=None,
                              show=False, show_peaks=False, show_fits=False, show_trace = False, debug=False,
                              manual_extract_dict=None):

        # create the ouptut image for skymask
        skymask = np.zeros_like(image, dtype=bool)

        plate_scale = self.spectrograph.order_platescale(binning=self.binning)
        inmask = self.mask == 0
        # Find objects
        specobj_dict = {'setup': self.setup, 'slitid': 999, 'orderindx': 999,
                        'det': self.det, 'objtype': self.objtype, 'pypeline': self.pypeline}
        # ToDO implement parsets here!
        sig_thresh = 30.0 if std else self.redux_par['sig_thresh']
        sobjs_ech, skymask[self.slitmask > -1] = \
            extract.ech_objfind(image, ivar, self.slitmask, self.tslits_dict['slit_left'], self.tslits_dict['slit_righ'],
                                inmask=inmask, ncoeff=self.redux_par['trace_npoly'],
                                hand_extract_dict=manual_extract_dict,
                                plate_scale=plate_scale, std_trace=std_trace,
                                specobj_dict=specobj_dict,sig_thresh=sig_thresh,
                                show_peaks=show_peaks, show_fits=show_fits,
                                show_trace=show_trace, debug=debug)

        # Steps
        self.steps.append(inspect.stack()[0][3])
        if show:
            self.show('image', image=image*(self.mask == 0), chname = 'ech_objfind',sobjs=sobjs_ech, slits=False)

        return sobjs_ech, len(sobjs_ech), skymask

    # JFH TODO Should we reduce the number of iterations for standards or near-IR redux where the noise model is not
    # being updated?
    def local_skysub_extract(self, sciimg, sciivar, tilts, waveimg, global_sky, rn2img, sobjs,
                             spat_pix=None, model_noise=True, min_snr=2.0, std = False, fit_fwhm=False,
                             maskslits=None, show_profile=False, show_resids=False, show_fwhm=False, show=False):
        """
        Perform local sky subtraction, profile fitting, and optimal extraction slit by slit

        Wrapper to skysub.local_skysub_extract

        Parameters
        ----------
        sobjs: object
           Specobjs object containing Specobj objects containing information about objects found.
        waveimg: ndarray, shape (nspec, nspat)
           Wavelength map

        Optional Parameters
        -------------------

        Returns:
            global_sky: (numpy.ndarray) image of the the global sky model
        """

        self.sciimg = sciimg
        self.sciivar = sciivar
        self.tilts = tilts
        self.waveimg = waveimg
        self.global_sky = global_sky
        self.rn2img = rn2img
        order_vec = self.spectrograph.order_vec()
        plate_scale = self.spectrograph.order_platescale(binning=self.binning)
        self.skymodel, self.objmodel, self.ivarmodel, self.outmask, self.sobjs = skysub.ech_local_skysub_extract(
            self.sciimg, self.sciivar, self.mask, self.tilts, self.waveimg, self.global_sky,
            self.rn2img, self.tslits_dict, sobjs, order_vec, spat_pix=spat_pix,
            std=std, fit_fwhm=fit_fwhm, min_snr=min_snr, bsp=self.redux_par['bspline_spacing'],
            box_rad_order=self.redux_par['boxcar_radius']/plate_scale,
            sigrej=self.redux_par['sky_sigrej'],
            sn_gauss=self.redux_par['sn_gauss'], model_full_slit=self.redux_par['model_full_slit'],
            model_noise=model_noise, show_profile=show_profile, show_resids=show_resids, show_fwhm=show_fwhm)

        # Step
        self.steps.append(inspect.stack()[0][3])

        if show:
            self.show('local', sobjs = self.sobjs, slits= True, chname='ech_local')
            self.show('resid', sobjs = self.sobjs, slits= True, chname='ech_resid')

        return self.skymodel, self.objmodel, self.ivarmodel, self.outmask, self.sobjs

def instantiate_me(spectrograph, tslits_dict, mask, par, **kwargs):
    """
    Instantiate the Reduce subclass appropriate for the provided
    spectrograph.

    The class must be subclassed from Reduce.  See :class:`Reduce` for
    the description of the valid keyword arguments.

    Args:
        spectrograph
            (:class:`pypeit.spectrographs.spectrograph.Spectrograph`):
            The instrument used to collect the data to be reduced.

        tslits_dict: dict
            dictionary containing slit/order boundary information

    Returns:
        :class:`PypeIt`: One of the classes with :class:`PypeIt` as its
        base.
    """
    indx = [ c.__name__ == spectrograph.pypeline for c in Reduce.__subclasses__() ]
    if not np.any(indx):
        msgs.error('Pipeline {0} is not defined!'.format(spectrograph.pypeline))
    return Reduce.__subclasses__()[np.where(indx)[0][0]](spectrograph, tslits_dict, mask, par, **kwargs)


#!/usr/bin/env python

from __future__ import division, print_function

import sys, os
import copy
import numpy as np
import simtk.openmm as mm
import simtk.openmm.app as app
import simtk.unit as u
import parmed as pmd
from parmed.openmm.reporters import NetCDFReporter
from pymbar import mbar
from openmm_surface_affinities_lib import *
import pytraj as pt

#Given a topology and structure file, this script sets up and runs umbrella-sampling simulations
#For the structure file given, the solute should be far from the surface so it can be pulled towards it
#This will generate starting points for the umbrella simulations

def main(args):
  #Get the structure and topology files from the command line
  #ParmEd accepts a wide range of file types (Amber, GROMACS, CHARMM, OpenMM... but not LAMMPS)
  try:
    topFile = args[0]
    strucFile = args[1]
  except IndexError:
    print("Specify topology and structure files from the command line.")
    sys.exit(2)

  print("Using topology file: %s" % topFile)
  print("Using structure file: %s" % strucFile)

  print("\nSetting up system...")

  #Load in the files for initial simulations
  top = pmd.load_file(topFile)
  struc = pmd.load_file(strucFile)

  #Transfer unit cell information to topology object
  top.box = struc.box[:]

  #Set up some global features to use in all simulations
  temperature = 298.15*u.kelvin

  #Define the platform (i.e. hardware and drivers) to use for running the simulation
  #This can be CUDA, OpenCL, CPU, or Reference
  #CUDA is for NVIDIA GPUs
  #OpenCL is for CPUs or GPUs, but must be used for old CPUs (not SSE4.1 compatible)
  #CPU only allows single precision (CUDA and OpenCL allow single, mixed, or double)
  #Reference is a clear, stable reference for other code development and is very slow, using double precision by default
  platform = mm.Platform.getPlatformByName('CUDA')
  prop = {#'Threads': '2', #number of threads for CPU - all definitions must be strings (I think)
          'Precision': 'mixed', #for CUDA or OpenCL, select the precision (single, mixed, or double)
          'DeviceIndex': '0', #selects which GPUs to use - set this to zero if using CUDA_VISIBLE_DEVICES
          'DeterministicForces': 'True' #Makes sure forces with CUDA and PME are deterministic
         }

  #Create the OpenMM system that can be used as a reference
  systemRef = top.createSystem(
                               nonbondedMethod=app.PME, #Uses PME for long-range electrostatics, simple cut-off for LJ
                               nonbondedCutoff=12.0*u.angstroms, #Defines cut-off for non-bonded interactions
                               rigidWater=True, #Use rigid water molecules
                               constraints=app.HBonds, #Constrains all bonds involving hydrogens
                               flexibleConstraints=False, #Whether to include energies for constrained DOFs
                               removeCMMotion=False, #Whether or not to remove COM motion (don't want to if part of system frozen)
  )

  #Set up the integrator to use as a reference
  integratorRef = mm.LangevinIntegrator(
                                        temperature, #Temperature for Langevin
                                        1.0/u.picoseconds, #Friction coefficient
                                        2.0*u.femtoseconds, #Integration timestep
  )
  integratorRef.setConstraintTolerance(1.0E-08)

  #To freeze atoms, set mass to zero (does not apply to virtual sites, termed "extra particles" in OpenMM)
  #Here assume (correctly, I think) that the topology indices for atoms correspond to those in the system
  for i, atom in enumerate(top.atoms):
    if atom.type in ('SU'): #, 'CU', 'CUO'):
      systemRef.setParticleMass(i, 0*u.dalton)

  #Track non-bonded force, mainly to turn off dispersion correction
  NBForce = None
  for frc in systemRef.getForces():
    if (isinstance(frc, mm.NonbondedForce)):
      NBForce = frc

  #Turn off dispersion correction since have interface
  NBForce.setUseDispersionCorrection(False)

  #Get solute atoms and solute heavy atoms separately
  soluteIndices = []
  heavyIndices = []
  for res in top.residues:
    if res.name not in ['OTM', 'CTM', 'STM', 'NTM', 'SOL']:
      for atom in res.atoms:
        soluteIndices.append(atom.idx)
        if 'H' not in atom.name[0]:
          heavyIndices.append(atom.idx)

  #Also get surface SU atoms
  surfIndices = []
  for atom in top.atoms:
    if atom.type == 'SU':
      surfIndices.append(atom.idx)

  startPos = np.array(struc.positions.value_in_unit(u.nanometer))

  print("\nSolute indices: %s" % str(soluteIndices))
  print("Solute heavy atom indices: %s" % str(heavyIndices))
  print("Surface SU atom indices: %s" % str(surfIndices))

  #Solute should already be placed far from the surface
  #If this is not done right, or it is too close to half the periodic box distance, will have issues
  #Either way, set this as the starting reference z distance
  initRefZ = np.average(startPos[heavyIndices, 2]) - np.average(startPos[surfIndices, 2])

  print(initRefZ)

  #Will now add a custom bonded force between solute heavy atoms and surface SU atoms
  #Should be in units of kJ/mol*nm^2, but should check this
  #For expanded ensemble, fine to assume solute is less than half the box distance from the surface
  #But for umbrella sampling, want to apply pulling towards surface regardless of whether pull from above or below
  #Also allows us to get further from the surface with our umbrellas without worrying about PBCs
  restraintExpression = '0.5*k*(abs(z2 - z1) - refZ)^2'
  restraintForce = mm.CustomCentroidBondForce(2, restraintExpression)
  restraintForce.addPerBondParameter('k')
  restraintForce.addGlobalParameter('refZ', initRefZ) #Make global so can modify during simulation
  restraintForce.addGroup(surfIndices, np.ones(len(surfIndices))) #Don't weight with masses
  restraintForce.addGroup(heavyIndices, np.ones(len(heavyIndices)))
  restraintForce.addBond([0, 1], [1000.0])
  restraintForce.setUsesPeriodicBoundaryConditions(True) #Only when doing umbrella sampling
  systemRef.addForce(restraintForce)

  forceLabelsRef = getForceLabels(systemRef)

  decompEnergy(systemRef, struc.positions, labels=forceLabelsRef, verbose=False)

  #Do NVT simulation
  stateFileNVT1, stateNVT1 = doSimNVT(top, systemRef, integratorRef, platform, prop, temperature, pos=struc.positions)

  #Do pulling simulation - really I'm just slowly changing the equilibrium bond distance between the surface and the solute
  stateFilePull, statePull = doSimPull(top, systemRef, integratorRef, platform, prop, temperature, state=stateFileNVT1)

  decompEnergy(systemRef, statePull, labels=forceLabelsRef, verbose=False)

  #Load in pulling restraint data and pulling trajectory to identify starting structures for each umbrella
  pullData = np.loadtxt('pull_restraint.txt')
  trajtop = copy.deepcopy(top)
  trajtop.rb_torsions = pmd.TrackedList([])
  trajtop = pt.load_parmed(trajtop, traj=False)
  pulltraj = pt.iterload('pull.nc', trajtop)
  frameTimes = np.array([frame.time for frame in pulltraj])

  #Define some umbrella distances based on a fixed spacing between initial and final pulling coordinate
  #Actually for final pulling coordinate, close to the surface, using average of actual coordinate and reference
  zSpace = 0.1 #nm
  zUmbs = np.arange(0.5*(pullData[-1,1] + pullData[-1,2]), pullData[0,2], zSpace)

  print("\nUsing following umbrellas:")
  print(zUmbs)

  #Then loop over umbrellas and run simulation for each
  for i, zRefDist in enumerate(zUmbs):

    os.mkdir("umbrella%i"%i)
    os.chdir("umbrella%i"%i)

    #Find where in the pulling trajectory the solute came closest to this umbrella
    pullDatInd = np.argmin(abs(pullData[:,2] - zRefDist))
    frameInd = np.argmin(abs(frameTimes - pullData[pullDatInd,0]))

    #Get starting coordinates
    #Making sure to assign the units that will be returned by pytraj
    thisCoords = np.array(pulltraj[frameInd].xyz) * u.angstrom

    #Set reference value for harmonic force
    restraintForce.setGlobalParameterDefaultValue(0, zRefDist)

    print("\nUmbrella %i:"%i)
    print("\tReference distance: %f"%zRefDist)
    print("\tFrame chosen from trajectory: %i (%f ps)"%(frameInd, frameTimes[frameInd]))

    #And run simulations, first equilibrating in NVT, then NPT, then production in NPT
    stateFileNVT, stateNVT = doSimNVT(top, systemRef, integratorRef, platform, prop, temperature, pos=thisCoords)

    stateFileNPT, stateNPT = doSimNPT(top, systemRef, integratorRef, platform, prop, temperature, state=stateFileNVT)

    stateFileProd, stateProd = doSimUmbrella(top, systemRef, integratorRef, platform, prop, temperature, state=stateFileNPT)

    os.chdir("../")

if __name__ == "__main__":
  main(sys.argv[1:])


import numpy as np
import cv2
#import matplotlib.pyplot as plt
face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
eye_cascade = cv2.CascadeClassifier('haarcascade_eye.xml')
cap = cv2.VideoCapture(0)
#to get height and width of the camera capture screen
if cap.isOpened():
    height=cap.get(3)
    width=cap.get(4)
    fheight=height/4
    fwidth=width/4
    a='fheight'
    print("height: ",fheight," width: ",fwidth)
#height,width,chanel=cap.shape()

while True:
    ret, img = cap.read()

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.3, 5)
    cv2.rectangle(img,(int(fheight),int(fwidth)),(int(height-fheight),int(width-fwidth)),(0,0,255),2)
    for (x,y,w,h) in faces:
        cv2.putText(img,(str(x)+","+str(y)),(x,y),cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,255,255),1,cv2.LINE_AA)
        print(x,"  ",y,"  ",w,"  ",h)
        cv2.rectangle(img,(x,y),(x+w,y+h),(255,0,0),2)
        roi_gray = gray[y:y+h, x:x+w]
        roi_color = img[y:y+h, x:x+w]

        eyes = eye_cascade.detectMultiScale(roi_gray)
        for (ex,ey,ew,eh) in eyes:
            cv2.rectangle(roi_color,(ex,ey),(ex+ew,ey+eh),(0,255,0),2)

    cv2.imshow('img',img)
    k = cv2.waitKey(30) & 0xff
    if k == 27:
        break

cap.release()
cv2.destroyAllWindows()


import torch
import torch.nn as nn
from torchvision.models import vgg16
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np

from model.utils.roipooling import RoIPool
from model.utils.loss import delta_loss
from model.utils.bbox_tools import delta2bbox
from model.utils.nms_cpu import py_cpu_nms

def get_vgg16_extractor_and_head(n_class, roip_size=7, vgg_pretrained=False):
    vgg16_net = vgg16(pretrained=True)
    features = list(vgg16_net.features)[0:30]

    for layer in features[0:10]:    # freeze top 4 conv2d layers
        for p in layer.parameters():
            p.requires_grad = False
    extractor = nn.Sequential(*features)
    output_feature_channel = 512

    classifier = list(vgg16_net.classifier)
    del(classifier[6])  # delete last fc layer
    classifier = nn.Sequential(*classifier)     # classifier : (N,25088) -> (N,4096); 25088 = 512*7*7 = C*H*W
    if torch.cuda.is_available():
        classifier = classifier.cuda()
    head = _VGG16Head(n_class_bg=n_class+1, roip_size=roip_size, classifier=classifier)
    if torch.cuda.is_available():
        extractor, head = extractor.cuda(), head.cuda()
    return extractor, head, output_feature_channel

class _VGG16Head(nn.Module):
    def __init__(self, n_class_bg, roip_size, classifier):
        """n_class_bg: n_class plus background = n_class + 1"""
        super(_VGG16Head, self).__init__()
        self.n_class_bg = n_class_bg
        self.roip_size = roip_size

        self.roip = RoIPool(roip_size, roip_size)
        self.classifier = classifier
        self.delta = nn.Linear(in_features=4096, out_features=n_class_bg*4)    # Note: predice a delta for each class
        self.score = nn.Linear(in_features=4096, out_features=n_class_bg)

        self._normal_init(self.delta, 0, 0.001)
        self._normal_init(self.score, 0, 0.01)

    def forward(self, feature_map, rois, image_size):
        """
        Args:
            feature_map: (N1=1,C,H,W)
            rois : (N2,4)
        """
        #---------- debug
        assert isinstance(feature_map, Variable)
        assert isinstance(rois, np.ndarray)
        assert len(feature_map.shape) == 4 and feature_map.shape[0] == 1    # batch size should be 1
        assert len(rois.shape) == 2 and rois.shape[1] == 4
        #---------- debug

        # this is important because rois are in image scale, we need to pass this ratio
        # to roipooing layer to map roi into feature_map scale
        feature_image_scale = feature_map.shape[2] / image_size[0]

        # meet roi_pooling's input requirement
        temp = np.zeros((rois.shape[0], 1), dtype=rois.dtype)
        rois = np.concatenate([temp, rois], axis=1)

        rois = Variable(torch.FloatTensor(rois))
        if torch.cuda.is_available():
            rois = rois.cuda()

        roipool_out = self.roip(feature_map, rois, spatial_scale=feature_image_scale)

        roipool_out = roipool_out.view(roipool_out.size(0), -1) # (N, 25088)
        if torch.cuda.is_available():
            roipool_out = roipool_out.cuda()

        mid_output = self.classifier(roipool_out)   # (N, 4096)
        delta_per_class = self.delta(mid_output)    # (N, n_class_bg*4)
        score = self.score(mid_output)      # (N, n_class_bg)
        #---------- debug
        assert isinstance(delta_per_class, Variable) and isinstance(score, Variable)
        assert delta_per_class.shape[0] == score.shape[0] == rois.shape[0]
        assert delta_per_class.shape[1] == score.shape[1] * 4 == self.n_class_bg * 4
        assert len(delta_per_class.shape) == len(score.shape) == 2
        #---------- debug
        return delta_per_class, score

    def loss(self, score, delta_per_class, target_delta_for_sample_roi, bbox_bg_label_for_sample_roi):
        """
        Args:
            score: (N, 2)
            delta_per_class: (N, 4*n_class_bg)
            target_delta_for_sample_roi: (N, 4)
            bbox_bg_label_for_sample_roi: (N,)
        """
        #---------- debug
        assert isinstance(score, Variable)
        assert isinstance(delta_per_class, Variable)
        assert isinstance(target_delta_for_sample_roi, np.ndarray)
        assert isinstance(bbox_bg_label_for_sample_roi, np.ndarray)
        #---------- debug
        target_delta_for_sample_roi = Variable(torch.FloatTensor(target_delta_for_sample_roi))
        bbox_bg_label_for_sample_roi = Variable(torch.LongTensor(bbox_bg_label_for_sample_roi))
        if torch.cuda.is_available():
            target_delta_for_sample_roi = target_delta_for_sample_roi.cuda()
            bbox_bg_label_for_sample_roi = bbox_bg_label_for_sample_roi.cuda()

        n_sample = score.shape[0]
        delta_per_class = delta_per_class.view(n_sample, -1, 4)

        # get delta for roi w.r.t its corresponding bbox label
        index = torch.arange(0, n_sample).long()
        if torch.cuda.is_available():
            index = index.cuda()
        delta = delta_per_class[index, bbox_bg_label_for_sample_roi.data]

        head_delta_loss = delta_loss(delta, target_delta_for_sample_roi, bbox_bg_label_for_sample_roi, 1)
        head_class_loss = F.cross_entropy(score, bbox_bg_label_for_sample_roi)

        return head_delta_loss + head_class_loss

    def predict(self, roi, delta_per_class, score, image_size, prob_threshold=0.5):
        """
        Args:
            roi: (N, 4)
            delta_per_class: (N, 4*n_class_bg)
            score: (N, n_class_bg)
        """
        #---------- debug
        assert isinstance(roi, np.ndarray)
        assert isinstance(delta_per_class, Variable)
        assert isinstance(score, Variable)
        #---------- debug
        roi = torch.FloatTensor(roi)
        if torch.cuda.is_available():
            roi = roi.cuda()
        delta_per_class = delta_per_class.data
        prob = F.softmax(score, dim=1).data

        delta_per_class = delta_per_class.view(-1, self.n_class_bg, 4)

        #!!!!!
        delta_per_class = delta_per_class * torch.cuda.FloatTensor([0.1, 0.1, 0.2, 0.2]) + torch.cuda.FloatTensor([0., 0., 0., 0.])

        roi = roi.view(-1,1,4).expand_as(delta_per_class)
        bbox_per_class = delta2bbox(roi.cpu().numpy().reshape(-1,4), delta_per_class.cpu().numpy().reshape(-1,4))
        bbox_per_class = torch.FloatTensor(bbox_per_class)

        bbox_per_class[:,0::2] = bbox_per_class[:,0::2].clamp(min=0, max=image_size[0])
        bbox_per_class[:,1::2] = bbox_per_class[:,1::2].clamp(min=0, max=image_size[1])

        bbox_per_class = bbox_per_class.numpy().reshape(-1,self.n_class_bg,4)
        prob = prob.cpu().numpy()
        #---------- debug
        assert bbox_per_class.shape[0] == prob.shape[0]
        assert bbox_per_class.shape[2] == 4
        assert bbox_per_class.shape[1] == prob.shape[1] == self.n_class_bg
        #---------- debug

        # suppress:
        bbox_out = []
        class_out = []
        prob_out = []
        # skip class_id = 0 because it is the background class
        for t in range(1, self.n_class_bg):
            bbox_for_class_t = bbox_per_class[:,t,:]    #(N, 4)
            prob_for_class_t = prob[:,t]                #(N,)
            mask = prob_for_class_t > prob_threshold    #(N,)
            # debug:
            # print("mask", mask.sum())
            left_bbox_for_class_t = bbox_for_class_t[mask]  #(N2,4)
            left_prob_for_class_t = prob_for_class_t[mask]  #(N2,)
            keep = py_cpu_nms(left_bbox_for_class_t, score=left_prob_for_class_t)
            bbox_out.append(left_bbox_for_class_t[keep])
            prob_out.append(left_prob_for_class_t[keep])
            class_out.append((t-1)*np.ones(len(keep)))

        bbox_out = np.concatenate(bbox_out, axis=0).astype(np.float32)
        prob_out = np.concatenate(prob_out, axis=0).astype(np.float32)
        class_out = np.concatenate(class_out, axis=0).astype(np.int32)
        #---------- debug
        assert isinstance(bbox_out, np.ndarray)
        assert isinstance(prob_out, np.ndarray)
        assert isinstance(class_out, np.ndarray)
        #---------- debug
        return bbox_out, class_out, prob_out

    def _normal_init(self, m, mean, stddev, truncated=False):
        """
        weight initalizer: truncated normal and random normal.
        """
        if truncated:
            m.weight.data.normal_().fmod_(2).mul_(stddev).add_(mean)  # not a perfect approximation
        else:
            m.weight.data.normal_(mean, stddev)
            m.bias.data.zero_()

def _get_resnet50_extractor_and_head():
    pass

if __name__ == '__main__':
    from model.utils.proposal_target_creator import ProposalTargetCreator
    extractor, head, output_feature_channel = get_vgg16_extractor_and_head(20, 7)
    features = Variable(torch.randn(1,512,50,50))
    if torch.cuda.is_available():
        extractor, head, features = extractor.cuda(), head.cuda(), features.cuda()

    rois = (np.random.rand(2000,4)+[0,0,1,1])*240
    gt_bbox = (np.random.rand(10,4) + [0,0,1,1])*240
    gt_bbox_label = np.random.randint(0,20,size=10)

    proposal_target_creator = ProposalTargetCreator()
    sample_roi, target_delta_for_sample_roi, bbox_bg_label_for_sample_roi = proposal_target_creator.make_proposal_target(rois, gt_bbox, gt_bbox_label)

    delta_per_class, score = head.forward(features, sample_roi, image_size=(500,500))
    loss = head.loss(score, delta_per_class,target_delta_for_sample_roi,bbox_bg_label_for_sample_roi)
    print(loss)
    loss.backward()

    rois = (np.random.rand(300,4)+[0,0,1,1])*240
    delta_per_class, score = head.forward(features, rois, image_size=(500,500))
    bbox_out, class_out, prob_out = head.predict(rois, delta_per_class, score, image_size=(500,500),prob_threshold=0.2)
    print(bbox_out.shape)
    print(class_out.shape)
    print(prob_out.shape)


"""
Fine-tune Faster R-CNN on HICO-DET

Fred Zhang <frederic.zhang@anu.edu.au>

The Australian National University
Australian Centre for Robotic Vision
"""

import os
import math
import json
import copy
import time
import torch
import bisect
import argparse
import torchvision
import numpy as np

from tqdm import tqdm
from itertools import repeat, chain
from collections import defaultdict
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.sampler import Sampler, BatchSampler

import pocket
from pocket.data import HICODet, DatasetConcat
from pocket.ops import RandomHorizontalFlip, to_tensor

class DetectorEngine(pocket.core.LearningEngine):
    def __init__(self, net, train_loader, val_loader, **kwargs):
        super().__init__(net, None, train_loader, **kwargs)
        self._val_loader = val_loader
        self.timer = pocket.utils.HandyTimer(1)

    def _on_each_iteration(self):
        self._state.output = self._state.net(*self._state.inputs, targets=self._state.targets)
        self._state.loss = sum(loss for loss in self._state.output.values())
        self._state.optimizer.zero_grad()
        self._state.loss.backward()
        self._state.optimizer.step()

    def _on_end_epoch(self):
        with self.timer:
            ap, max_rec = self.validate()
        print("\n=> Validation (+{:.2f})\n"
            "Epoch: {} | mAP: {:.4f}, mRec: {:.4f} | Time: {:.2f}s\n".format(
                time.time() - self._dawn, self._state.epoch,
                ap.mean().item(), max_rec.mean().item(), self.timer[0]
            ))
        super()._on_end_epoch()

    @torch.no_grad()
    def validate(self, min_iou=0.5, nms_thresh=0.5):
        num_gt = torch.zeros(80)
        associate = pocket.utils.BoxAssociation(min_iou=min_iou)
        meter = pocket.utils.DetectionAPMeter(
            80, algorithm='INT', nproc=10
        )
        self._state.net.eval()
        for batch in tqdm(self._val_loader):
            inputs = pocket.ops.relocate_to_cuda(batch[0])
            output = self._state.net(inputs)
            assert len(output) == 1, "The batch size should be one"
            # Relocate back to cpu
            output = pocket.ops.relocate_to_cpu(output[0])
            target = batch[1][0]
            # Do NMS on ground truth boxes
            # NOTE This is because certain objects appear multiple times in
            # different pairs and different interactions
            keep_gt_idx = torchvision.ops.boxes.batched_nms(
                target['boxes'], torch.ones_like(target['labels']).float(),
                target['labels'], nms_thresh
            )
            gt_boxes = target['boxes'][keep_gt_idx].view(-1, 4)
            gt_classes = target['labels'][keep_gt_idx].view(-1)
            # Update the number of ground truth instances
            # Convert the object index to zero based
            for c in gt_classes:
                num_gt[c - 1] += 1
            # Associate detections with ground truth
            binary_labels = torch.zeros_like(output['scores'])
            unique_obj = output['labels'].unique()
            for obj_idx in unique_obj:
                det_idx = torch.nonzero(output['labels'] == obj_idx).squeeze(1)
                gt_idx = torch.nonzero(gt_classes == obj_idx).squeeze(1)
                if len(gt_idx) == 0:
                    continue
                binary_labels[det_idx] = associate(
                    gt_boxes[gt_idx].view(-1, 4),
                    output['boxes'][det_idx].view(-1, 4),
                    output['scores'][det_idx].view(-1)
                )

            meter.append(output['scores'], output['labels'] - 1, binary_labels)

        meter.num_gt = num_gt.tolist()
        ap = meter.eval()
        return ap, meter.max_rec

class HICODetObject(Dataset):
    def __init__(self, dataset, data_root, nms_thresh=0.5, random_flip=False):
        self.dataset = dataset
        self.nms_thresh = nms_thresh
        with open(os.path.join(data_root, 'coco80tohico80.json'), 'r') as f:
            corr = json.load(f)
        self.hico2coco91 = dict(zip(corr.values(), corr.keys()))
        self.transform = RandomHorizontalFlip() if random_flip else None
    def __len__(self):
        return len(self.dataset)
    def __getitem__(self, idx):
        image, target = self.dataset[idx]
        boxes = torch.cat([
            target['boxes_h'],
            target['boxes_o']
        ])
        # Convert ground truth boxes to zero-based index and the
        # representation from pixel indices to coordinates
        boxes[:, :2] -= 1
        labels = torch.cat([
            49 * torch.ones_like(target['object']),
            target['object']
        ])
        # Convert HICODet object indices to COCO indices
        converted_labels = torch.tensor([int(self.hico2coco91[i.item()]) for i in labels])
        # Apply transform
        if self.transform is not None:
            image, boxes = self.transform(image, boxes)
        image = to_tensor(image, input_format='pil')
        return [image], [dict(boxes=boxes, labels=converted_labels)]

def collate_fn(batch):
    images = []
    targets = []
    for im, tar in batch:
        images += im
        targets += tar
    return images, targets

"""
Batch sampler that groups images by aspect ratio
https://github.com/pytorch/vision/blob/master/references/detection/group_by_aspect_ratio.py
"""

def _repeat_to_at_least(iterable, n):
    repeat_times = math.ceil(n / len(iterable))
    repeated = chain.from_iterable(repeat(iterable, repeat_times))
    return list(repeated)

class GroupedBatchSampler(BatchSampler):
    """
    Wraps another sampler to yield a mini-batch of indices.
    It enforces that the batch only contain elements from the same group.
    It also tries to provide mini-batches which follows an ordering which is
    as close as possible to the ordering from the original sampler.
    Arguments:
        sampler (Sampler): Base sampler.
        group_ids (list[int]): If the sampler produces indices in range [0, N),
            `group_ids` must be a list of `N` ints which contains the group id of each sample.
            The group ids must be a continuous set of integers starting from
            0, i.e. they must be in the range [0, num_groups).
        batch_size (int): Size of mini-batch.
    """
    def __init__(self, sampler, group_ids, batch_size):
        if not isinstance(sampler, Sampler):
            raise ValueError(
                "sampler should be an instance of "
                "torch.utils.data.Sampler, but got sampler={}".format(sampler)
            )
        self.sampler = sampler
        self.group_ids = group_ids
        self.batch_size = batch_size

    def __iter__(self):
        buffer_per_group = defaultdict(list)
        samples_per_group = defaultdict(list)

        num_batches = 0
        for idx in self.sampler:
            group_id = self.group_ids[idx]
            buffer_per_group[group_id].append(idx)
            samples_per_group[group_id].append(idx)
            if len(buffer_per_group[group_id]) == self.batch_size:
                yield buffer_per_group[group_id]
                num_batches += 1
                del buffer_per_group[group_id]
            assert len(buffer_per_group[group_id]) < self.batch_size

        # now we have run out of elements that satisfy
        # the group criteria, let's return the remaining
        # elements so that the size of the sampler is
        # deterministic
        expected_num_batches = len(self)
        num_remaining = expected_num_batches - num_batches
        if num_remaining > 0:
            # for the remaining batches, take first the buffers with largest number
            # of elements
            for group_id, _ in sorted(buffer_per_group.items(),
                                      key=lambda x: len(x[1]), reverse=True):
                remaining = self.batch_size - len(buffer_per_group[group_id])
                samples_from_group_id = _repeat_to_at_least(samples_per_group[group_id], remaining)
                buffer_per_group[group_id].extend(samples_from_group_id[:remaining])
                assert len(buffer_per_group[group_id]) == self.batch_size
                yield buffer_per_group[group_id]
                num_remaining -= 1
                if num_remaining == 0:
                    break
        assert num_remaining == 0

    def __len__(self):
        return len(self.sampler) // self.batch_size

def _quantize(x, bins):
    bins = copy.deepcopy(bins)
    bins = sorted(bins)
    quantized = list(map(lambda y: bisect.bisect_right(bins, y), x))
    return quantized

def create_aspect_ratio_groups(aspect_ratios, k=0):
    bins = (2 ** np.linspace(-1, 1, 2 * k + 1)).tolist() if k > 0 else [1.0]
    groups = _quantize(aspect_ratios, bins)
    # count number of elements per group
    counts = np.unique(groups, return_counts=True)[1]
    fbins = [0] + bins + [np.inf]
    print("Using {} as bins for aspect ratio quantization".format(fbins))
    print("Count of instances per bin: {}".format(counts))
    return groups

def main(args):

    torch.cuda.set_device(0)
    torch.manual_seed(args.random_seed)

    train2015 = HICODetObject(HICODet(
        root=os.path.join(args.data_root, "hico_20160224_det/images/train2015"),
        anno_file=os.path.join(args.data_root, "instances_train2015.json"),
        target_transform=pocket.ops.ToTensor(input_format='dict')
    ), data_root=args.data_root, random_flip=True)
    test2015 = HICODetObject(HICODet(
        root=os.path.join(args.data_root, "hico_20160224_det/images/test2015"),
        anno_file=os.path.join(args.data_root, "instances_test2015.json"),
        target_transform=pocket.ops.ToTensor(input_format='dict')
    ), data_root=args.data_root)

    def div(a, b):
        return a / b
    use_train2015 = 'train2015' in args.training_data
    use_test2015 = 'test2015' in args.training_data
    if len(args.training_data) == 1 and use_train2015:
        trainset = train2015
        aspect_ratios = [div(*train2015.dataset.image_size(i)) for i in range(len(train2015))]
    elif len(args.training_data) == 1 and use_test2015:
        trainset = test2015
        aspect_ratios = [div(*test2015.dataset.image_size(i)) for i in range(len(test2015))]
    elif len(args.training_data) == 2 and use_train2015 and use_train2015:
        trainset = DatasetConcat(train2015, test2015)
        aspect_ratios = [
            div(*train2015.dataset.image_size(i)) for i in range(len(train2015))
        ] + [div(*test2015.dataset.image_size(i)) for i in range(len(test2015))]
    else:
        raise ValueError("Unknown dataset partition in ", args.training_data)

    sampler = torch.utils.data.RandomSampler(trainset)
    group_ids = create_aspect_ratio_groups(aspect_ratios, k=args.aspect_ratio_group_factor)
    batch_sampler = GroupedBatchSampler(sampler, group_ids, args.batch_size)
    train_loader = DataLoader(
        dataset=trainset, batch_sampler=batch_sampler,
        num_workers=4, collate_fn=collate_fn,
    )

    val_loader = DataLoader(
        dataset=test2015, batch_size=1, shuffle=False,
        num_workers=4, collate_fn=collate_fn
    )

    net = pocket.models.fasterrcnn_resnet_fpn('resnet50', pretrained=True)
    net.cuda()

    engine = DetectorEngine(
        net, train_loader, val_loader,
        print_interval=args.print_interval,
        cache_dir=args.cache_dir,
        optim_params=dict(
            lr=args.learning_rate,
            momentum=args.momentum,
            weight_decay=args.weight_decay
        ),
        lr_scheduler=True,
        lr_sched_params=dict(
            milestones=args.milestones,
            gamma=args.lr_decay
        )
    )

    engine(args.num_epochs)

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description="Fine-tune Faster R-CNN on HICO-DET")
    parser.add_argument('--data-root', type=str, default='../')
    parser.add_argument('--training-data', nargs='+', default=['train2015',], type=str)
    parser.add_argument('--num-epochs', default=15, type=int)
    parser.add_argument('--random-seed', default=1, type=int)
    parser.add_argument('--learning-rate', default=0.00025, type=float)
    parser.add_argument('--momentum', default=0.9, type=float)
    parser.add_argument('--weight-decay', default=1e-4, type=float)
    parser.add_argument('--batch-size', default=2, type=int)
    parser.add_argument('--milestones', nargs='+', default=[8, 12], type=int)
    parser.add_argument('--lr-decay', default=0.1, type=float)
    parser.add_argument('--aspect-ratio-group-factor', default=3, type=int)
    parser.add_argument('--print-interval', default=2000, type=int)
    parser.add_argument('--cache-dir', type=str, default='./checkpoints')

    args = parser.parse_args()
    print(args)

    main(args)


#! /usr/bin/env python

'''
reload of photostim trial event times from file
'''

import os
import logging
import math

from pathlib import Path
from itertools import chain
from collections import namedtuple
from datetime import datetime

import numpy as np
import scipy.io as spio

from pipeline import lab
from pipeline import experiment
from pipeline.fixes import fix_history

from pipeline.ingest import behavior as behavior_ingest

log = logging.getLogger(__name__)

# TODO: were we inserting full rig-subpaths paths before pathlib?
#       or just names.. this would not be necessary with rig-subpath stored
#       also, could be separate issue from which is better.

def find_path(path, fname):
    log.info('finding', fname, 'in', path)

    for root, dirs, files in os.walk(path):
        log.info('RigDataFile.make(): entering {r}'.format(r=root))
        for f in files:
            if f == fname:
                log.info('found', root, f)
                return Path(root) / Path(f)

    return None

def fix_session(session_key):
    paths = behavior_ingest.RigDataPath.fetch(as_dict=True)
    files = (behavior_ingest.BehaviorIngest
             * behavior_ingest.BehaviorIngest.BehaviorFile
             & session_key).fetch(as_dict=True, order_by='behavior_file asc')

    filelist = []
    for pf in [(p, f) for f in files for p in paths]:
        p, f = pf
        found = find_path(p['rig_data_path'], f['behavior_file'])
        if found:
            filelist.append(found)

    if len(filelist) != len(files):
        log.warning("session {} has behavior files missing. skipping")
        return

    #
    # Prepare PhotoStim
    #
    photosti_duration = 0.5  # (s) Hard-coded here
    photostims = {4: {'photo_stim': 4, 'photostim_device': 'OBIS470',
                      'brain_location_name': 'left_alm', 'duration': photosti_duration},
                  5: {'photo_stim': 5, 'photostim_device': 'OBIS470',
                      'brain_location_name': 'right_alm', 'duration': photosti_duration},
                  6: {'photo_stim': 6, 'photostim_device': 'OBIS470',
                      'brain_location_name': 'both_alm', 'duration': photosti_duration}}

    #
    # Load all files & create combined list of per-trial data
    #

    trials = zip()

    trial = namedtuple(  # simple structure to track per-trial vars
        'trial', ('ttype', 'stim', 'settings', 'state_times', 'state_names',
                  'state_data', 'event_data', 'event_times'))

    for f in filelist:

        if os.stat(f).st_size/1024 < 1000:
            log.info('skipping file {f} - too small'.format(f=f))
            continue

        log.debug('loading file {}'.format(f))

        mat = spio.loadmat(f, squeeze_me=True)

        SessionData = mat['SessionData'].flatten()

        AllTrialTypes = SessionData['TrialTypes'][0]
        AllTrialSettings = SessionData['TrialSettings'][0]

        RawData = SessionData['RawData'][0].flatten()
        AllStateNames = RawData['OriginalStateNamesByNumber'][0]
        AllStateData = RawData['OriginalStateData'][0]
        AllEventData = RawData['OriginalEventData'][0]
        AllStateTimestamps = RawData['OriginalStateTimestamps'][0]
        AllEventTimestamps = RawData['OriginalEventTimestamps'][0]

        # verify trial-related data arrays are all same length
        assert(all((x.shape[0] == AllStateTimestamps.shape[0] for x in
                    (AllTrialTypes, AllTrialSettings,
                     AllStateNames, AllStateData, AllEventData,
                     AllEventTimestamps))))

        if 'StimTrials' in SessionData.dtype.fields:
            log.debug('StimTrials detected in session - will include')
            AllStimTrials = SessionData['StimTrials'][0]
            assert(AllStimTrials.shape[0] == AllStateTimestamps.shape[0])
        else:
            log.debug('StimTrials not detected in session - will skip')
            AllStimTrials = np.array([
                None for i in enumerate(range(AllStateTimestamps.shape[0]))])

        z = zip(AllTrialTypes, AllStimTrials, AllTrialSettings,
                AllStateTimestamps, AllStateNames, AllStateData,
                AllEventData, AllEventTimestamps)

        trials = chain(trials, z)

    trials = list(trials)

    # all files were internally invalid or size < 100k
    if not trials:
        log.warning('skipping .. no valid files')

    #
    # Trial data seems valid; synthesize session id & add session record
    # XXX: note - later breaks can result in Sessions without valid trials
    #

    key = session_key
    skey = (experiment.Session & key).fetch1()

    #
    # Actually load the per-trial data
    #
    log.info('BehaviorIngest.make(): trial parsing phase')

    # lists of various records for batch-insert
    rows = {k: list() for k in ('trial', 'behavior_trial', 'trial_note',
                                'trial_event', 'corrected_trial_event',
                                'action_event', 'photostim',
                                'photostim_location', 'photostim_trial',
                                'photostim_trial_event')}

    i = -1
    for t in trials:

        #
        # Misc
        #

        t = trial(*t)  # convert list of items to a 'trial' structure
        i += 1  # increment trial counter

        log.debug('BehaviorIngest.make(): parsing trial {i}'.format(i=i))

        # covert state data names into a lookup dictionary
        #
        # names (seem to be? are?):
        #
        # Trigtrialstart
        # PreSamplePeriod
        # SamplePeriod
        # DelayPeriod
        # EarlyLickDelay
        # EarlyLickSample
        # ResponseCue
        # GiveLeftDrop
        # GiveRightDrop
        # GiveLeftDropShort
        # GiveRightDropShort
        # AnswerPeriod
        # Reward
        # RewardConsumption
        # NoResponse
        # TimeOut
        # StopLicking
        # StopLickingReturn
        # TrialEnd

        states = {k: (v+1) for v, k in enumerate(t.state_names)}
        required_states = ('PreSamplePeriod', 'SamplePeriod',
                           'DelayPeriod', 'ResponseCue', 'StopLicking',
                           'TrialEnd')

        missing = list(k for k in required_states if k not in states)

        if len(missing):
            log.warning('skipping trial {i}; missing {m}'
                        .format(i=i, m=missing))
            continue

        gui = t.settings['GUI'].flatten()

        # ProtocolType - only ingest protocol >= 3
        #
        # 1 Water-Valve-Calibration 2 Licking 3 Autoassist
        # 4 No autoassist 5 DelayEnforce 6 SampleEnforce 7 Fixed
        #

        if 'ProtocolType' not in gui.dtype.names:
            log.warning('skipping trial {i}; protocol undefined'
                        .format(i=i))
            continue

        protocol_type = gui['ProtocolType'][0]
        if gui['ProtocolType'][0] < 3:
            log.warning('skipping trial {i}; protocol {n} < 3'
                        .format(i=i, n=gui['ProtocolType'][0]))
            continue

        #
        # Top-level 'Trial' record
        #

        tkey = dict(skey)
        startindex = np.where(t.state_data == states['PreSamplePeriod'])[0]

        # should be only end of 1st StopLicking;
        # rest of data is irrelevant w/r/t separately ingested ephys
        endindex = np.where(t.state_data == states['StopLicking'])[0]

        log.debug('states\n' + str(states))
        log.debug('state_data\n' + str(t.state_data))
        log.debug('startindex\n' + str(startindex))
        log.debug('endindex\n' + str(endindex))

        if not(len(startindex) and len(endindex)):
            log.warning('skipping trial {i}: start/end index error: {s}/{e}'.format(i=i,s=str(startindex), e=str(endindex)))
            continue

        try:
            tkey['trial'] = i
            tkey['trial_uid'] = i  # Arseny has unique id to identify some trials
            tkey['start_time'] = t.state_times[startindex][0]
            tkey['stop_time'] = t.state_times[endindex][0]
        except IndexError:
            log.warning('skipping trial {i}: error indexing {s}/{e} into {t}'.format(i=i, s=str(startindex), e=str(endindex), t=str(t.state_times)))
            continue

        log.debug('BehaviorIngest.make(): Trial().insert1')  # TODO msg
        log.debug('tkey' + str(tkey))
        rows['trial'].append(tkey)

        #
        # Specific BehaviorTrial information for this trial
        #

        bkey = dict(tkey)
        bkey['task'] = 'audio delay'  # hard-coded here
        bkey['task_protocol'] = 1     # hard-coded here

        # determine trial instruction
        trial_instruction = 'left'    # hard-coded here

        if gui['Reversal'][0] == 1:
            if t.ttype == 1:
                trial_instruction = 'left'
            elif t.ttype == 0:
                trial_instruction = 'right'
        elif gui['Reversal'][0] == 2:
            if t.ttype == 1:
                trial_instruction = 'right'
            elif t.ttype == 0:
                trial_instruction = 'left'

        bkey['trial_instruction'] = trial_instruction

        # determine early lick
        early_lick = 'no early'

        if (protocol_type >= 5
            and 'EarlyLickDelay' in states
                and np.any(t.state_data == states['EarlyLickDelay'])):
                early_lick = 'early'
        if (protocol_type > 5
            and ('EarlyLickSample' in states
                 and np.any(t.state_data == states['EarlyLickSample']))):
                early_lick = 'early'

        bkey['early_lick'] = early_lick

        # determine outcome
        outcome = 'ignore'

        if ('Reward' in states
                and np.any(t.state_data == states['Reward'])):
            outcome = 'hit'
        elif ('TimeOut' in states
                and np.any(t.state_data == states['TimeOut'])):
            outcome = 'miss'
        elif ('NoResponse' in states
                and np.any(t.state_data == states['NoResponse'])):
            outcome = 'ignore'

        bkey['outcome'] = outcome
        rows['behavior_trial'].append(bkey)

        #
        # Add 'protocol' note
        #
        nkey = dict(tkey)
        nkey['trial_note_type'] = 'protocol #'
        nkey['trial_note'] = str(protocol_type)
        rows['trial_note'].append(nkey)

        #
        # Add 'autolearn' note
        #
        nkey = dict(tkey)
        nkey['trial_note_type'] = 'autolearn'
        nkey['trial_note'] = str(gui['Autolearn'][0])
        rows['trial_note'].append(nkey)

        #
        # Add 'bitcode' note
        #
        if 'randomID' in gui.dtype.names:
            nkey = dict(tkey)
            nkey['trial_note_type'] = 'bitcode'
            nkey['trial_note'] = str(gui['randomID'][0])
            rows['trial_note'].append(nkey)

        #
        # Add presample event
        #
        log.debug('BehaviorIngest.make(): presample')

        ekey = dict(tkey)
        sampleindex = np.where(t.state_data == states['SamplePeriod'])[0]

        ekey['trial_event_id'] = len(rows['trial_event'])
        ekey['trial_event_type'] = 'presample'
        ekey['trial_event_time'] = t.state_times[startindex][0]
        ekey['duration'] = (t.state_times[sampleindex[0]]
                            - t.state_times[startindex])[0]

        if math.isnan(ekey['duration']):
            log.debug('BehaviorIngest.make(): fixing presample duration')
            ekey['duration'] = 0.0  # FIXDUR: lookup from previous trial

        rows['trial_event'].append(ekey)

        #
        # Add 'go' event
        #
        log.debug('BehaviorIngest.make(): go')

        ekey = dict(tkey)
        responseindex = np.where(t.state_data == states['ResponseCue'])[0]

        ekey['trial_event_id'] = len(rows['trial_event'])
        ekey['trial_event_type'] = 'go'
        ekey['trial_event_time'] = t.state_times[responseindex][0]
        ekey['duration'] = gui['AnswerPeriod'][0]

        if math.isnan(ekey['duration']):
            log.debug('BehaviorIngest.make(): fixing go duration')
            ekey['duration'] = 0.0  # FIXDUR: lookup from previous trials
            rows['corrected_trial_event'].append(ekey)

        rows['trial_event'].append(ekey)

        #
        # Add other 'sample' events
        #

        log.debug('BehaviorIngest.make(): sample events')

        last_dur = None

        for s in sampleindex:  # in protocol > 6 ~-> n>1
            # todo: batch events
            ekey = dict(tkey)
            ekey['trial_event_id'] = len(rows['trial_event'])
            ekey['trial_event_type'] = 'sample'
            ekey['trial_event_time'] = t.state_times[s]
            ekey['duration'] = gui['SamplePeriod'][0]

            if math.isnan(ekey['duration']) and last_dur is None:
                log.warning('... trial {} bad duration, no last_edur'
                            .format(i, last_dur))
                ekey['duration'] = 0.0  # FIXDUR: cross-trial check
                rows['corrected_trial_event'].append(ekey)

            elif math.isnan(ekey['duration']) and last_dur is not None:
                log.warning('... trial {} duration using last_edur {}'
                            .format(i, last_dur))
                ekey['duration'] = last_dur
                rows['corrected_trial_event'].append(ekey)

            else:
                last_dur = ekey['duration']  # only track 'good' values.

            rows['trial_event'].append(ekey)

        #
        # Add 'delay' events
        #

        log.debug('BehaviorIngest.make(): delay events')

        last_dur = None
        delayindex = np.where(t.state_data == states['DelayPeriod'])[0]

        for d in delayindex:  # protocol > 6 ~-> n>1
            ekey = dict(tkey)
            ekey['trial_event_id'] = len(rows['trial_event'])
            ekey['trial_event_type'] = 'delay'
            ekey['trial_event_time'] = t.state_times[d]
            ekey['duration'] = gui['DelayPeriod'][0]

            if math.isnan(ekey['duration']) and last_dur is None:
                log.warning('... {} bad duration, no last_edur'
                            .format(i, last_dur))
                ekey['duration'] = 0.0  # FIXDUR: cross-trial check
                rows['corrected_trial_event'].append(ekey)

            elif math.isnan(ekey['duration']) and last_dur is not None:
                log.warning('... {} duration using last_edur {}'
                            .format(i, last_dur))
                ekey['duration'] = last_dur
                rows['corrected_trial_event'].append(ekey)

            else:
                last_dur = ekey['duration']  # only track 'good' values.

            log.debug('delay event duration: {}'.format(ekey['duration']))
            rows['trial_event'].append(ekey)

        #
        # Add 'trialEnd' events
        #

        log.debug('BehaviorIngest.make(): trialend events')

        last_dur = None
        trialendindex = np.where(t.state_data == states['TrialEnd'])[0]

        ekey['trial_event_id'] = len(rows['trial_event'])
        ekey['trial_event_type'] = 'trialend'
        ekey['trial_event_time'] = t.state_times[trialendindex][0]
        ekey['duration'] = 0.0

        rows['trial_event'].append(ekey)

        #
        # Add lick events
        #

        lickleft = np.where(t.event_data == 69)[0]
        log.debug('... lickleft: {r}'.format(r=str(lickleft)))

        action_event_count = len(rows['action_event'])
        if len(lickleft):
            [rows['action_event'].append(
                dict(tkey, action_event_id=action_event_count+idx, action_event_type='left lick',
                     action_event_time=t.event_times[l]))
             for idx, l in enumerate(lickleft)]

        lickright = np.where(t.event_data == 70)[0]
        log.debug('... lickright: {r}'.format(r=str(lickright)))

        action_event_count = len(rows['action_event'])
        if len(lickright):
            [rows['action_event'].append(
                dict(tkey, action_event_id=action_event_count+idx, action_event_type='right lick',
                     action_event_time=t.event_times[r]))
                for idx, r in enumerate(lickright)]

        # Photostim Events
        #
        # TODO:
        #
        # - base stimulation parameters:
        #
        #   - should be loaded elsewhere - where
        #   - actual ccf locations - cannot be known apriori apparently?
        #   - Photostim.Profile: what is? fix/add
        #
        # - stim data
        #
        #   - how retrieve power from file (didn't see) or should
        #     be statically coded here?
        #   - how encode stim type 6?
        #     - we have hemisphere as boolean or
        #     - but adding an event 4 and event 5 means querying
        #       is less straightforwrard (e.g. sessions with 5 & 6)

        if t.stim:
            log.info('BehaviorIngest.make(): t.stim == {}'.format(t.stim))
            rows['photostim_trial'].append(tkey)
            delay_period_idx = np.where(t.state_data == states['DelayPeriod'])[0][0]
            rows['photostim_trial_event'].append(dict(
                tkey, **photostims[t.stim], photostim_event_id=len(rows['photostim_trial_event']),
                photostim_event_time=t.state_times[delay_period_idx],
                power=5.5))

        # end of trial loop.

    pe = experiment.PhotostimEvent
    for pev in rows['photostim_trial_event']:

        log.debug('updating photostim_event {}'.format(pev))

        pe_k = {k: v for k, v in pev.items() if k in pe.primary_key}

        try:
            (pe & pe_k)._update('photostim_event_time',
                                pev['photostim_event_time'])

            (pe & pe_k)._update('power', pev['power'])
        except Exception as e:
            log.error('exception with pe_k : {}, pev: {} - {}'.format(
                pe_k, pev, repr(e)))
            raise

def verify_session(session_key):
    if len(experiment.PhotostimEvent & session_key & {'power': 0.0}) != 0:
        log.warning('NOK: session {} session_key still has power==0.0 events')
    else:
        log.info('OK: session {} has no power==0.0 events')

def fix_0001_photostim():

    fix_history.schema.connection.start_transaction()

    fh = {'fix_name': 'fix_0001_photostim', 'fix_timestamp': datetime.now()}

    fix_history.FixHistory.insert1(fh)

    q = (experiment.Session & behavior_ingest.BehaviorIngest)

    for s in q.fetch('KEY'):
        fix_session(s)
        verify_session(s)
        # "FixHistorySession" table is obsolete now (by fix_0007), but this fix is no longer needed
        fix_history.FixHistory.FixHistorySession.insert1({**fh, **s})

    fix_history.schema.connection.commit_transaction()

if __name__ == '__main__':
    fix_0001_photostim()


#!/search/odin/liyaozong/tools/python3/bin/python3
# coding: utf8

import random
import numpy as np
from collections import defaultdict, deque
from game import Board, Game
from mcts_pure import MCTSPlayer as MCTS_Pure
from mcts_alphaZero import MCTSPlayer
import multiprocessing
from multiprocessing import Manager, Pool
import time
import logging
from tensorflow.python.platform import gfile
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

model_dir='./multi_3'
log_name = model_dir + '/logs'
current_model_name = model_dir + '/current_policy_model'
temp_current_model_name = model_dir + '/temp_current_policy_model'
best_model_name = model_dir + '/best_policy_model'

class TrainPipeline():

    def __init__(self, init_model=None):
        # params of the board and the game
        self.board_width = 8
        self.board_height = 8
        self.n_in_row = 5
        # training params
        self.learn_rate = 2e-3
        # self.lr_multiplier = 1.0
        self.temp = 1.0  # the temperature param
        self.n_playout = 400  # num of simulations for each move
        self.c_puct = 5
        self.batch_size = 512  # mini-batch size for training
        self.buffer_num = self.batch_size * 10
        self.epochs = 5  # num of train_steps for each update
        self.kl_targ = 0.02
        self.check_freq = 60
        self.game_batch_num = 1000000000
        self.process_num = 12
        self.summary_record_freq = 5
        # self.global_update_step = 0

    def collect_selfplay_data_thread(self, thread_id, shared_queue, net_lock, data_lock):
        logging.info('selfpaly process start: {}'.format(thread_id))
        os.environ["CUDA_VISIBLE_DEVICES"] = str(thread_id % 6 + 2)
        from policy_value_net_tensorflow import PolicyValueNet
        # 读取模型文件，加锁
        with net_lock:
            current_policy = PolicyValueNet(self.board_width, self.board_height, model_file=current_model_name)
        local_board = Board(width=self.board_width,
                           height=self.board_height,
                           n_in_row=self.n_in_row)
        local_game = Game(local_board)
        local_mcts_player = MCTSPlayer(current_policy.policy_value_fn,
                                      c_puct=self.c_puct,
                                      n_playout=self.n_playout,
                                      is_selfplay=1)
        winner, play_data = local_game.start_self_play(local_mcts_player,
                                                       temp=self.temp)
        play_data = list(play_data)
        play_data = self.get_equi_data(play_data)
        # 添加对弈数据，加锁
        with data_lock:
            shared_queue.extend(play_data)
            while len(shared_queue) > self.buffer_num:
                shared_queue.pop(0)
        logging.info('selfpaly process finished: {}'.format(thread_id))

    def get_equi_data(self, play_data):
        """augment the data set by rotation and flipping
        play_data: [(state, mcts_prob, winner_z), ..., ...]
        """
        extend_data = []
        for state, mcts_porb, winner in play_data:
            for i in [1, 2, 3, 4]:
                # rotate counterclockwise
                equi_state = np.array([np.rot90(s, i) for s in state])
                equi_mcts_prob = np.rot90(np.flipud(
                    mcts_porb.reshape(self.board_height, self.board_width)), i)
                extend_data.append((equi_state,
                                    np.flipud(equi_mcts_prob).flatten(),
                                    winner))
                # flip horizontally
                equi_state = np.array([np.fliplr(s) for s in equi_state])
                equi_mcts_prob = np.fliplr(equi_mcts_prob)
                extend_data.append((equi_state,
                                    np.flipud(equi_mcts_prob).flatten(),
                                    winner))
        return extend_data

    def policy_update(self, current_policy_value_net, shared_queue, net_lock, global_update_step, lr_multiplier):
        """update the policy-value net"""
        random_index = list(range(len(shared_queue)))
        random.shuffle(random_index)
        mini_batch = []
        for i in range(self.batch_size):
            mini_batch.append(shared_queue[random_index[i]])
        state_batch = [data[0] for data in mini_batch]
        mcts_probs_batch = [data[1] for data in mini_batch]
        winner_batch = [data[2] for data in mini_batch]
        old_probs, old_v = current_policy_value_net.policy_value(state_batch)
        loss, entropy = current_policy_value_net.train_step(
                state_batch,
                mcts_probs_batch,
                winner_batch,
                self.learn_rate * lr_multiplier.value)
        new_probs, new_v = current_policy_value_net.policy_value(state_batch)
        kl = np.mean(np.sum(old_probs * (
                np.log(old_probs + 1e-10) - np.log(new_probs + 1e-10)),
                axis=1)
        )
        # adaptively adjust the learning rate
        if kl > self.kl_targ * 2 and lr_multiplier.value > 0.1:
            lr_multiplier.value /= 1.5
        elif kl < self.kl_targ / 2 and lr_multiplier.value < 10:
            lr_multiplier.value *= 1.5

        explained_var_old = (1 -
                             np.var(np.array(winner_batch) - old_v.flatten()) /
                             np.var(np.array(winner_batch)))
        explained_var_new = (1 -
                             np.var(np.array(winner_batch) - new_v.flatten()) /
                             np.var(np.array(winner_batch)))
        logging.info("update current model process kl:{:.5f},lr_multiplier:{:.3f},loss:{},entropy:{},explained_var_old:{:.3f},explained_var_new:{:.3f}".format(
            kl,
            lr_multiplier.value,
            loss,
            entropy,
            explained_var_old,
            explained_var_new))
        # summary for tensorboard
        if global_update_step.value % self.summary_record_freq == 0:
            current_policy_value_net.summary_record(
                    state_batch,
                    mcts_probs_batch,
                    winner_batch,
                    global_update_step.value,
                    )

    def update_net(self, shared_queue, net_lock, update_best_model, global_update_step, lr_multiplier, stop_update_process, update_or_selfplay):
        os.environ["CUDA_VISIBLE_DEVICES"] = "1"
        from policy_value_net_tensorflow import PolicyValueNet
        current_policy_value_net = PolicyValueNet(self.board_width, self.board_height, model_dir)
        current_policy_value_net.save_model(current_model_name)
        current_policy_value_net.save_model(best_model_name)
        while global_update_step.value <= self.game_batch_num:
            if update_or_selfplay.value == 0:
                if len(shared_queue) >= self.batch_size:
                    for _ in range(self.epochs):
                        global_update_step.value += 1
                        logging.info('update current model process start self train: {}'.format(global_update_step.value))
                        self.policy_update(current_policy_value_net, shared_queue, net_lock, global_update_step, lr_multiplier)
                        if (global_update_step.value) % self.check_freq == 0:
                            update_best_model.value = 1
                    # 这里更新最新模型文件
                    with net_lock:
                        logging.info('update process update current model')
                        current_policy_value_net.save_model(current_model_name)
                update_or_selfplay.value = 1
            else:
                time.sleep(1)
        stop_update_process.value = 1

    def update_best_model_thread(self, current_model_name, best_model_name, net_lock, update_best_model, stop_update_process):
        os.environ["CUDA_VISIBLE_DEVICES"] = "1"
        def update_best_model_local_process_func(name1, name2):
            from policy_value_net_tensorflow import PolicyValueNet
            PolicyValueNet(self.board_width, self.board_height, model_file=name1).save_model(name2)
        update_best_model_global_num = 0
        update_best_model_time = 0
        logging.info('update best model process start!')
        while stop_update_process.value == 0:
            time.sleep(1)
            if update_best_model.value == 1:
                update_best_model.value = 0
                update_best_model_global_num += 1
                logging.info('update best model process global_num:{}'.format(update_best_model_global_num))
                win_num = 0
                start_player = 0
                with net_lock:
                    p = multiprocessing.Process(target=update_best_model_local_process_func, args=(current_model_name, temp_current_model_name))
                    p.start()
                    p.join()
                for i in range(30):
                    board = Board(width=self.board_width, height=self.board_height, n_in_row=self.n_in_row)
                    game = Game(board)
                    win_player = game.two_net_play(temp_current_model_name, best_model_name, net_lock, start_player=start_player, is_shown=0)
                    start_player = 1 - start_player  # play first in turn
                    if win_player == 0:
                        win_num += 1
                logging.info('update best model process global_num:{} finished, current model win total {}'.format(update_best_model_global_num, win_num))
                if win_num >= 16:
                    update_best_model_time += 1
                    logging.info('update best model process get new best model:{}'.format(update_best_model_time))
                    p = multiprocessing.Process(target=update_best_model_local_process_func, args=(temp_current_model_name, best_model_name))
                    p.start()
                    p.join()

    def run(self):
        """run the training pipeline"""
        try:
            # 必须在一个线程中引入tensorflow，否则会造成其他线程由于错误阻塞。
            # ERROR: could not retrieve CUDA device count: CUDA_ERROR_NOT_INITIALIZED
            m = Manager()
            shared_queue = m.list()
            net_lock = m.Lock()
            data_lock = m.Lock()
            stop_update_process = multiprocessing.Value('i', 0)
            update_best_model = multiprocessing.Value('i', 0)
            global_update_step = multiprocessing.Value('i', 0)
            lr_multiplier = multiprocessing.Value('d', 1.0)
            update_or_selfplay = multiprocessing.Value('i', 0)
            update_best_model_process = multiprocessing.Process(target=self.update_best_model_thread,
                    args=(current_model_name, best_model_name, net_lock, update_best_model, stop_update_process))
            update_best_model_process.start()
            update_current_model_process = multiprocessing.Process(target=self.update_net,
                    args=(shared_queue, net_lock, update_best_model, global_update_step, lr_multiplier, stop_update_process, update_or_selfplay))
            update_current_model_process.start()
            while stop_update_process.value == 0:
                if update_or_selfplay.value == 1:
                    # 产生训练数据
                    pro_list = []
                    for i in range(self.process_num):
                        pro = multiprocessing.Process(target=self.collect_selfplay_data_thread, args=(i, shared_queue, net_lock, data_lock))
                        pro_list.append(pro)
                        pro.start()
                    for pro in pro_list:
                        pro.join()
                    update_or_selfplay.value = 0
                else:
                    time.sleep(0.1)
            # 等待更新最有模型的进程结束，再结束主程序
            stop_update_process.value = 1
            while update_best_model_process.is_alive():
                time.sleep(5)
        except Exception as e:
            logging.error(e)

if __name__ == '__main__':
    if gfile.Exists(model_dir):
        gfile.DeleteRecursively(model_dir)
    gfile.MakeDirs(model_dir)
    # log 配置
    logging.basicConfig(filename=log_name, level=logging.INFO, format="[%(levelname)s]\t%(asctime)s\tLINENO:%(lineno)d\t%(message)s", datefmt="%Y-%m-%d %H:%M:%S")
    training_pipeline = TrainPipeline()
    logging.info('start training')
    training_pipeline.run()
    logging.info('all finished')