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

"""This stimulus class defines a field of dots with an update rule that
determines how they change on every call to the .draw() method.
"""

# Part of the PsychoPy library
# Copyright (C) 2002-2018 Jonathan Peirce (C) 2019-2021 Open Science Tools Ltd.
# Distributed under the terms of the GNU General Public License (GPL).

# Bugfix by Andrew Schofield.
# Replaces out of bounds but still live dots at opposite edge of aperture
# instead of randomly within the field. This stops the concentration of dots at
# one side of field when lifetime is long.
# Update the dot direction immediately for 'walk' as otherwise when the
# coherence varies some signal dots will inherit the random directions of
# previous walking dots.
# Provide a visible wrapper function to refresh all the dot locations so that
# the whole field can be more easily refreshed between trials.

from __future__ import absolute_import, division, print_function

from builtins import str
from builtins import range

# Ensure setting pyglet.options['debug_gl'] to False is done prior to any
# other calls to pyglet or pyglet submodules, otherwise it may not get picked
# up by the pyglet GL engine and have no effect.
# Shaders will work but require OpenGL2.0 drivers AND PyOpenGL3.0+
import pyglet
pyglet.options['debug_gl'] = False
import ctypes
GL = pyglet.gl

import psychopy  # so we can get the __path__
from psychopy import logging

# tools must only be imported *after* event or MovieStim breaks on win32
# (JWP has no idea why!)
from psychopy.tools.attributetools import attributeSetter, setAttribute
from psychopy.tools.arraytools import val2array
from psychopy.visual.basevisual import (BaseVisualStim, ColorMixin,
                                        ContainerMixin)

import numpy as np

# some constants
_piOver2 = np.pi / 2.
_piOver180 = np.pi / 180.
_2pi = 2 * np.pi

class DotStim(BaseVisualStim, ColorMixin, ContainerMixin):
    """This stimulus class defines a field of dots with an update rule that
    determines how they change on every call to the .draw() method.

    This single class can be used to generate a wide variety of dot motion
    types. For a review of possible types and their pros and cons see Scase,
    Braddick & Raymond (1996). All six possible motions they describe can be
    generated with appropriate choices of the `signalDots` (which determines
    whether signal dots are the 'same' or 'different' on each frame),
    `noiseDots` (which determines the locations of the noise dots on each frame)
    and the `dotLife` (which determines for how many frames the dot will
    continue before being regenerated).

    The default settings (as of v1.70.00) is for the noise dots to have
    identical velocity but random direction and signal dots remain the 'same'
    (once a signal dot, always a signal dot).

    For further detail about the different configurations see :ref:`dots` in the
    Builder Components section of the documentation.

    If further customisation is required, then the DotStim should be subclassed
    and its _update_dotsXY and _newDotsXY methods overridden.

    The maximum number of dots that can be drawn is limited by system
    performance.

    Attributes
    ----------
    fieldShape : str
        *'sqr'* or 'circle'. Defines the envelope used to present the dots. If
        changed while drawing, dots outside new envelope will be respawned.
    dotSize : float
        Dot size specified in pixels (overridden if `element` is specified).
        :ref:`operations <attrib-operations>` are supported.
    dotLife : int
        Number of frames each dot lives for (-1=infinite). Dot lives are
        initiated randomly from a uniform distribution from 0 to dotLife. If
        changed while drawing, the lives of all dots will be randomly initiated
        again.
    signalDots : str
        If 'same' then the signal and noise dots are constant. If 'different'
        then the choice of which is signal and which is noise gets randomised on
        each frame. This corresponds to Scase et al's (1996) categories of RDK.
    noiseDots : str
        Determines the behaviour of the noise dots, taken directly from Scase et
        al's (1996) categories. For 'position', noise dots take a random
        position every frame. For 'direction' noise dots follow a random, but
        constant direction. For 'walk' noise dots vary their direction every
        frame, but keep a constant speed.
    element : object
        This can be any object that has a ``.draw()`` method and a
        ``.setPos([x,y])`` method (e.g. a GratingStim, TextStim...)!! DotStim
        assumes that the element uses pixels as units. ``None`` defaults to
        dots.
    fieldPos : array_like
        Specifying the location of the centre of the stimulus using a
        :ref:`x,y-pair <attrib-xy>`. See e.g. :class:`.ShapeStim` for more
        documentation/examples on how to set position.
        :ref:`operations <attrib-operations>` are supported.
    fieldSize : array_like
        Specifying the size of the field of dots using a
        :ref:`x,y-pair <attrib-xy>`. See e.g. :class:`.ShapeStim` for more
        documentation/examples on how to set position.
        :ref:`operations <attrib-operations>` are supported.
    coherence : float
        Change the coherence (%) of the DotStim. This will be rounded according
        to the number of dots in the stimulus.
    dir : float
        Direction of the coherent dots in degrees. :ref:`operations
        <attrib-operations>` are supported.
    speed : float
        Speed of the dots (in *units*/frame). :ref:`operations
        <attrib-operations>` are supported.

    """
    def __init__(self,
                 win,
                 units='',
                 nDots=1,
                 coherence=0.5,
                 fieldPos=(0.0, 0.0),
                 fieldSize=(1.0, 1.0),
                 fieldShape='sqr',
                 dotSize=2.0,
                 dotLife=3,
                 dir=0.0,
                 speed=0.5,
                 rgb=None,
                 color=(1.0, 1.0, 1.0),
                 colorSpace='rgb',
                 opacity=None,
                 contrast=1.0,
                 depth=0,
                 element=None,
                 signalDots='same',
                 noiseDots='direction',
                 name=None,
                 autoLog=None):
        """
        Parameters
        ----------
        win : window.Window
            Window this stimulus is associated with.
        units : str
            Units to use.
        nDots : int
            Number of dots to present in the field.
        coherence : float
            Proportion of dots which are coherent. This value can be set using
            the `coherence` property after initialization.
        fieldPos : array_like
            (x,y) or [x,y] position of the field. This value can be set using
            the `fieldPos` property after initialization.
        fieldSize : array_like, int or float
            (x,y) or [x,y] or single value (applied to both dimensions). Sizes
            can be negative and can extend beyond the window. This value can be
            set using the `fieldSize` property after initialization.
        fieldShape : str
            Defines the envelope used to present the dots. If changed while
            drawing by setting the `fieldShape` property, dots outside new
            envelope will be respawned., valid values are 'square', 'sqr' or
            'circle'.
        dotSize : array_like or float
            Size of the dots. If given an array, the sizes of individual dots
            will be set. The array must have length `nDots`. If a single value
            is given, all dots will be set to the same size.
        dotLife : int
            Lifetime of a dot in frames. Dot lives are initiated randomly from a
            uniform distribution from 0 to dotLife. If changed while drawing,
            the lives of all dots will be randomly initiated again. A value of
            -1 results in dots having an infinite lifetime. This value can be
            set using the `dotLife` property after initialization.
        dir : float
            Direction of the coherent dots in degrees. At 0 degrees, coherent
            dots will move from left to right. Increasing the angle will rotate
            the direction counter-clockwise. This value can be set using the
            `dir` property after initialization.
        speed : float
            Speed of the dots (in *units* per frame). This value can be set
            using the `speed` property after initialization.
        rgb : array_like, optional
            Color of the dots in form (r, g, b) or [r, g, b]. **Deprecated**,
            use `color` instead.
        color : array_like or str
            Color of the dots in form (r, g, b) or [r, g, b].
        colorSpace : str
            Colorspace to use.
        opacity : float
            Opacity of the dots from 0.0 to 1.0.
        contrast : float
            Contrast of the dots 0.0 to 1.0. This value is simply multiplied by
            the `color` value.
        depth : float
            **Deprecated**, depth is now controlled simply by drawing order.
        element : object
            This can be any object that has a ``.draw()`` method and a
            ``.setPos([x,y])`` method (e.g. a GratingStim, TextStim...)!!
            DotStim assumes that the element uses pixels as units.
            ``None`` defaults to dots.
        signalDots : str
            If 'same' then the signal and noise dots are constant. If different
            then the choice of which is signal and which is noise gets
            randomised on each frame. This corresponds to Scase et al's (1996)
            categories of RDK. This value can be set using the `signalDots`
            property after initialization.
        noiseDots : str
            Determines the behaviour of the noise dots, taken directly from
            Scase et al's (1996) categories. For 'position', noise dots take a
            random position every frame. For 'direction' noise dots follow a
            random, but constant direction. For 'walk' noise dots vary their
            direction every frame, but keep a constant speed. This value can be
            set using the `noiseDots` property after initialization.
        name : str, optional
            Optional name to use for logging.
        autoLog : bool
            Enable automatic logging.

        """
        # what local vars are defined (these are the init params) for use by
        # __repr__
        self._initParams = __builtins__['dir']()
        self._initParams.remove('self')

        super(DotStim, self).__init__(win, units=units, name=name,
                                      autoLog=False)  # set at end of init

        self.nDots = nDots
        # pos and size are ambiguous for dots so DotStim explicitly has
        # fieldPos = pos, fieldSize=size and then dotSize as additional param
        self.fieldPos = fieldPos  # self.pos is also set here
        self.fieldSize = val2array(fieldSize, False)  # self.size is also set
        if type(dotSize) in (tuple, list):
            self.dotSize = np.array(dotSize)
        else:
            self.dotSize = dotSize
        if self.win.useRetina:
            self.dotSize *= 2  # double dot size to make up for 1/2-size pixels
        self.fieldShape = fieldShape
        self.__dict__['dir'] = dir
        self.speed = speed
        self.element = element
        self.dotLife = dotLife
        self.signalDots = signalDots

        self.useShaders = False  # not needed for dots?
        if rgb != None:
            logging.warning("Use of rgb arguments to stimuli are deprecated."
                            " Please use color and colorSpace args instead")
            self.colorSpace = 'rgba'
            self.color = rgb
        else:
            self.colorSpace = colorSpace
            self.color = color
        self.opacity = opacity
        self.contrast = float(contrast)
        self.depth = depth

        # initialise the dots themselves - give them all random dir and then
        # fix the first n in the array to have the direction specified
        self.coherence = coherence  # using the attributeSetter
        self.noiseDots = noiseDots

        # initialise a random array of X,Y
        self._verticesBase = self._dotsXY = self._newDotsXY(self.nDots)
        # all dots have the same speed
        self._dotsSpeed = np.ones(self.nDots, dtype=float) * self.speed
        # abs() means we can ignore the -1 case (no life)
        self._dotsLife = np.abs(dotLife) * np.random.rand(self.nDots)
        # pre-allocate array for flagging dead dots
        self._deadDots = np.zeros(self.nDots, dtype=bool)
        # set directions (only used when self.noiseDots='direction')
        self._dotsDir = np.random.rand(self.nDots) * _2pi
        self._dotsDir[self._signalDots] = self.dir * _piOver180

        self._update_dotsXY()

        # set autoLog now that params have been initialised
        wantLog = autoLog is None and self.win.autoLog
        self.__dict__['autoLog'] = autoLog or wantLog
        if self.autoLog:
            logging.exp("Created %s = %s" % (self.name, str(self)))

    def set(self, attrib, val, op='', log=None):
        """DEPRECATED: DotStim.set() is obsolete and may not be supported
        in future versions of PsychoPy. Use the specific method for each
        parameter instead (e.g. setFieldPos(), setCoherence()...).
        """
        self._set(attrib, val, op, log=log)

    @attributeSetter
    def fieldShape(self, fieldShape):
        """*'sqr'* or 'circle'. Defines the envelope used to present the dots.
        If changed while drawing, dots outside new envelope will be respawned.
        """
        self.__dict__['fieldShape'] = fieldShape

    @attributeSetter
    def dotSize(self, dotSize):
        """Float specified in pixels (overridden if `element` is specified).
        :ref:`operations <attrib-operations>` are supported."""
        self.__dict__['dotSize'] = dotSize

    @attributeSetter
    def dotLife(self, dotLife):
        """Int. Number of frames each dot lives for (-1=infinite).
        Dot lives are initiated randomly from a uniform distribution
        from 0 to dotLife. If changed while drawing, the lives of all
        dots will be randomly initiated again.

        :ref:`operations <attrib-operations>` are supported.
        """
        self.__dict__['dotLife'] = dotLife
        self._dotsLife = abs(self.dotLife) * np.random.rand(self.nDots)

    @attributeSetter
    def signalDots(self, signalDots):
        """str - 'same' or *'different'*
        If 'same' then the signal and noise dots are constant. If different
        then the choice of which is signal and which is noise gets
        randomised on each frame. This corresponds to Scase et al's (1996)
        categories of RDK.
        """
        self.__dict__['signalDots'] = signalDots

    @attributeSetter
    def noiseDots(self, noiseDots):
        """str - *'direction'*, 'position' or 'walk'
        Determines the behaviour of the noise dots, taken directly from
        Scase et al's (1996) categories. For 'position', noise dots take a
        random position every frame. For 'direction' noise dots follow a
        random, but constant direction. For 'walk' noise dots vary their
        direction every frame, but keep a constant speed.
        """
        self.__dict__['noiseDots'] = noiseDots
        self.coherence = self.coherence  # update using attributeSetter

    @attributeSetter
    def element(self, element):
        """*None* or a visual stimulus object
        This can be any object that has a ``.draw()`` method and a
        ``.setPos([x,y])`` method (e.g. a GratingStim, TextStim...)!!
        DotStim assumes that the element uses pixels as units.
        ``None`` defaults to dots.

        See `ElementArrayStim` for a faster implementation of this idea.
        """
        self.__dict__['element'] = element

    @attributeSetter
    def fieldPos(self, pos):
        """Specifying the location of the centre of the stimulus
        using a :ref:`x,y-pair <attrib-xy>`.
        See e.g. :class:`.ShapeStim` for more documentation / examples
        on how to set position.

        :ref:`operations <attrib-operations>` are supported.
        """
        # Isn't there a way to use BaseVisualStim.pos.__doc__ as docstring
        # here?
        self.pos = pos  # using BaseVisualStim. we'll store this as both
        self.__dict__['fieldPos'] = self.pos

    def setFieldPos(self, val, op='', log=None):
        """Usually you can use 'stim.attribute = value' syntax instead, but use
        this method if you need to suppress the log message.
        """
        setAttribute(self, 'fieldPos', val, log, op)  # calls attributeSetter

    def setPos(self, newPos=None, operation='', units=None, log=None):
        """Obsolete - users should use setFieldPos instead of setPos
        """
        logging.error("User called DotStim.setPos(pos). "
                      "Use DotStim.SetFieldPos(pos) instead.")

    def setFieldSize(self, val, op='', log=None):
        """Usually you can use 'stim.attribute = value' syntax instead, but use
        this method if you need to suppress the log message.
        """
        setAttribute(self, 'fieldSize', val, log, op)  # calls attributeSetter

    @attributeSetter
    def fieldSize(self, size):
        """Specifying the size of the field of dots using a
        :ref:`x,y-pair <attrib-xy>`. See e.g. :class:`.ShapeStim` for more
        documentation/examples on how to set position.

        :ref:`operations <attrib-operations>` are supported.
        """
        # Isn't there a way to use BaseVisualStim.pos.__doc__ as docstring
        # here?
        self.size = size  # using BaseVisualStim. we'll store this as both
        self.__dict__['fieldSize'] = self.size

    @attributeSetter
    def coherence(self, coherence):
        """Scalar between 0 and 1.

        Change the coherence (%) of the DotStim. This will be rounded according
        to the number of dots in the stimulus.

        :ref:`operations <attrib-operations>` are supported.
        """
        if not 0 <= coherence <= 1:
            raise ValueError('DotStim.coherence must be between 0 and 1')

        _cohDots = coherence * self.nDots

        self.__dict__['coherence'] = round(_cohDots) /self.nDots
        self._signalDots = np.zeros(self.nDots, dtype=bool)
        self._signalDots[0:int(self.coherence * self.nDots)] = True
        # for 'direction' method we need to update the direction of the number
        # of signal dots immediately, but for other methods it will be done
        # during updateXY

        # NB - AJS Actually you need to do this for 'walk' also
        # otherwise would be signal dots adopt random directions when the become
        # sinal dots in later trails
        if self.noiseDots in ('direction', 'position', 'walk'):
            self._dotsDir = np.random.rand(self.nDots) * _2pi
            self._dotsDir[self._signalDots] = self.dir * _piOver180

    def setFieldCoherence(self, val, op='', log=None):
        """Usually you can use 'stim.attribute = value' syntax instead, but use
        this method if you need to suppress the log message.
        """
        setAttribute(self, 'coherence', val, log, op)  # calls attributeSetter

    @attributeSetter
    def dir(self, dir):
        """float (degrees). direction of the coherent dots. :ref:`operations
        <attrib-operations>` are supported.
        """
        # check which dots are signal before setting new dir
        signalDots = self._dotsDir == (self.dir * _piOver180)
        self.__dict__['dir'] = dir

        # dots currently moving in the signal direction also need to update
        # their direction
        self._dotsDir[signalDots] = self.dir * _piOver180

    def setDir(self, val, op='', log=None):
        """Usually you can use 'stim.attribute = value' syntax instead, but use
        this method if you need to suppress the log message.
        """
        setAttribute(self, 'dir', val, log, op)

    @attributeSetter
    def speed(self, speed):
        """float. speed of the dots (in *units*/frame). :ref:`operations
        <attrib-operations>` are supported.
        """
        self.__dict__['speed'] = speed

    def setSpeed(self, val, op='', log=None):
        """Usually you can use 'stim.attribute = value' syntax instead, but use
        this method if you need to suppress the log message.

        """
        setAttribute(self, 'speed', val, log, op)

    def draw(self, win=None):
        """Draw the stimulus in its relevant window. You must call this method
        after every MyWin.flip() if you want the stimulus to appear on that
        frame and then update the screen again.

        Parameters
        ----------
        win : window.Window, optional
            Window to draw dots to. If `None`, dots will be drawn to the parent
            window.

        """
        if win is None:
            win = self.win
        self._selectWindow(win)

        self._update_dotsXY()

        GL.glPushMatrix()  # push before drawing, pop after

        # draw the dots
        if self.element is None:
            win.setScale('pix')
            GL.glPointSize(self.dotSize)

            # load Null textures into multitexteureARB - they modulate with
            # glColor
            GL.glActiveTexture(GL.GL_TEXTURE0)
            GL.glEnable(GL.GL_TEXTURE_2D)
            GL.glBindTexture(GL.GL_TEXTURE_2D, 0)
            GL.glActiveTexture(GL.GL_TEXTURE1)
            GL.glEnable(GL.GL_TEXTURE_2D)
            GL.glBindTexture(GL.GL_TEXTURE_2D, 0)

            CPCD = ctypes.POINTER(ctypes.c_double)
            GL.glVertexPointer(2, GL.GL_DOUBLE, 0,
                               self.verticesPix.ctypes.data_as(CPCD))
            GL.glColor4f(*self._foreColor.render('rgba1'))
            GL.glEnableClientState(GL.GL_VERTEX_ARRAY)
            GL.glDrawArrays(GL.GL_POINTS, 0, self.nDots)
            GL.glDisableClientState(GL.GL_VERTEX_ARRAY)
        else:
            # we don't want to do the screen scaling twice so for each dot
            # subtract the screen centre
            initialDepth = self.element.depth
            for pointN in range(0, self.nDots):
                _p = self.verticesPix[pointN, :] + self.fieldPos
                self.element.setPos(_p)
                self.element.draw()
            # reset depth before going to next frame
            self.element.setDepth(initialDepth)
        GL.glPopMatrix()

    def _newDotsXY(self, nDots):
        """Returns a uniform spread of dots, according to the `fieldShape` and
        `fieldSize`.

        Parameters
        ----------
        nDots : int
            Number of dots to sample.

        Returns
        -------
        ndarray
            Nx2 array of X and Y positions of dots.

        Examples
        --------
        Create a new array of dot positions::

            dots = self._newDots(nDots)

        """
        if self.fieldShape == 'circle':
            length = np.sqrt(np.random.uniform(0, 1, (nDots,)))
            angle = np.random.uniform(0., _2pi, (nDots,))

            newDots = np.zeros((nDots, 2))
            newDots[:, 0] = length * np.cos(angle)
            newDots[:, 1] = length * np.sin(angle)

            newDots *= self.fieldSize * .5
        else:
            newDots = np.random.uniform(-0.5, 0.5, size = (nDots, 2)) * self.fieldSize

        return newDots

    def refreshDots(self):
        """Callable user function to choose a new set of dots."""
        self._verticesBase = self._dotsXY = self._newDotsXY(self.nDots)

        # Don't allocate another array if the new number of dots is equal to
        # the last.
        if self.nDots != len(self._deadDots):
            self._deadDots = np.zeros(self.nDots, dtype=bool)

    def _update_dotsXY(self):
        """The user shouldn't call this - its gets done within draw().
        """
        # Find dead dots, update positions, get new positions for
        # dead and out-of-bounds
        # renew dead dots
        if self.dotLife > 0:  # if less than zero ignore it
            # decrement. Then dots to be reborn will be negative
            self._dotsLife -= 1
            self._deadDots[:] = (self._dotsLife <= 0)
            self._dotsLife[self._deadDots] = self.dotLife
        else:
            self._deadDots[:] = False

        # update XY based on speed and dir
        # NB self._dotsDir is in radians, but self.dir is in degs
        # update which are the noise/signal dots
        if self.signalDots == 'different':
            #  **up to version 1.70.00 this was the other way around,
            # not in keeping with Scase et al**
            # noise and signal dots change identity constantly
            np.random.shuffle(self._dotsDir)
            # and then update _signalDots from that
            self._signalDots = (self._dotsDir == (self.dir * _piOver180))

        # update the locations of signal and noise; 0 radians=East!
        reshape = np.reshape
        if self.noiseDots == 'walk':
            # noise dots are ~self._signalDots
            sig = np.random.rand(np.sum(~self._signalDots))
            self._dotsDir[~self._signalDots] = sig * _2pi
            # then update all positions from dir*speed
            cosDots = reshape(np.cos(self._dotsDir), (self.nDots,))
            sinDots = reshape(np.sin(self._dotsDir), (self.nDots,))
            self._verticesBase[:, 0] += self.speed * cosDots
            self._verticesBase[:, 1] += self.speed * sinDots
        elif self.noiseDots == 'direction':
            # simply use the stored directions to update position
            cosDots = reshape(np.cos(self._dotsDir), (self.nDots,))
            sinDots = reshape(np.sin(self._dotsDir), (self.nDots,))
            self._verticesBase[:, 0] += self.speed * cosDots
            self._verticesBase[:, 1] += self.speed * sinDots
        elif self.noiseDots == 'position':
            # update signal dots
            sd = self._signalDots
            sdSum = self._signalDots.sum()
            cosDots = reshape(np.cos(self._dotsDir[sd]), (sdSum,))
            sinDots = reshape(np.sin(self._dotsDir[sd]), (sdSum,))
            self._verticesBase[sd, 0] += self.speed * cosDots
            self._verticesBase[sd, 1] += self.speed * sinDots
            # update noise dots
            self._deadDots[:] = self._deadDots + (~self._signalDots)

        # handle boundaries of the field
        if self.fieldShape in (None, 'square', 'sqr'):
            out0 = (np.abs(self._verticesBase[:, 0]) > .5 * self.fieldSize[0])
            out1 = (np.abs(self._verticesBase[:, 1]) > .5 * self.fieldSize[1])
            outofbounds = out0 + out1
        else:
            # transform to a normalised circle (radius = 1 all around)
            # then to polar coords to check
            # the normalised XY position (where radius should be < 1)
            normXY = self._verticesBase / .5 / self.fieldSize
            # add out-of-bounds to those that need replacing
            outofbounds = np.hypot(normXY[:, 0], normXY[:, 1]) > 1.

        # update any dead dots
        nDead = self._deadDots.sum()
        if nDead:
            self._verticesBase[self._deadDots, :] = self._newDotsXY(nDead)

        # Reposition any dots that have gone out of bounds. Net effect is to
        # place dot one step inside the boundary on the other side of the
        # aperture.
        nOutOfBounds = outofbounds.sum()
        if nOutOfBounds:
            self._verticesBase[outofbounds, :] = self._newDotsXY(nOutOfBounds)

        # update the pixel XY coordinates in pixels (using _BaseVisual class)
        self._updateVertices()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Functions to deal with HSC filters."""

import os
import json

import numpy as np

from astropy.io import fits
from astropy.table import Table
from astropy import constants as const

import matplotlib.pyplot as plt
from matplotlib import rc
from matplotlib import rcParams

import unagi

__all__ = ['filters_to_kcorrect', 'hsc_filters', 'HscFilter', 'SolarSpectrum']

plt.rc('text', usetex=True)
rcParams.update({'xtick.major.pad': '7.0'})
rcParams.update({'xtick.major.size': '7.5'})
rcParams.update({'xtick.major.width': '1.5'})
rcParams.update({'xtick.minor.pad': '7.0'})
rcParams.update({'xtick.minor.size': '3.5'})
rcParams.update({'xtick.minor.width': '1.0'})
rcParams.update({'ytick.major.pad': '7.0'})
rcParams.update({'ytick.major.size': '7.5'})
rcParams.update({'ytick.major.width': '1.5'})
rcParams.update({'ytick.minor.pad': '7.0'})
rcParams.update({'ytick.minor.size': '3.5'})
rcParams.update({'ytick.minor.width': '1.0'})
rcParams.update({'axes.titlepad': '15.0'})
rcParams.update({'font.size': 26})

# Directory that save the transmission curves of the filter
FILTER_DIR = os.path.join(os.path.dirname(unagi.__file__), 'data', 'filters')
# Directory that keeps the solar spectra
SOLAR_DIR = os.path.join(os.path.dirname(unagi.__file__), 'data', 'solar')

# Formal names and nick names of the filters
FILTER_LIST = ['HSC-G', 'HSC-R', 'HSC-R2', 'HSC-I', 'HSC-I2', 'HSC-Z', 'HSC-Y',
               'HSC-NB387', 'HSC-NB816', 'HSC-NB921']
FILTER_SHORT = ['g', 'r', 'r2', 'i', 'i2', 'z', 'y', 'nb387', 'nb816', 'nb921']

# Colors to show the filter
FILTER_COLORS = ['#2ca02c', '#ff7f0e', '#ff7f0e', '#d62728', '#d62728',
                 '#8c564b', '#7f7f7f', 'c', 'm', 'purple']

# AB reference flux in unit of erg/s/cm^2/Hz
AB_FLUX = 3.631e-20
# Speed of light in unit of AA/s
C_AA_PER_SEC = const.c.to('AA/s').value

class Filter(object):
    """Class for organizing HSC filters.

    Please see the following page for more details about HSC filters:
    https://www.subarutelescope.org/Observing/Instruments/HSC/sensitivity.html

    This class is based on the Filter() class from Ben Johnson's sedpy
    https://github.com/bd-j/sedpy/blob/master/sedpy/observate.py
    """
    def __init__(self, band, origin=False, center=False):
        """Read in the HSC filter transmission curves.
        """
        if band.strip().upper() in FILTER_LIST:
            self.index = FILTER_LIST.index(band.strip().upper())
        elif band.strip().lower() in FILTER_SHORT:
            self.index = FILTER_SHORT.index(band.strip().lower())
        else:
            raise NameError("# Wrong filter name!")

        # Name of the filter
        self.filter = FILTER_LIST[self.index]
        self.short = FILTER_SHORT[self.index]

        # Color to show the filter
        self.color = FILTER_COLORS[self.index]

        # Old or new version of r- and i-band filters
        if self.short == 'r' or self.short == 'i':
            self.version = 'old'
        else:
            self.version = 'new'

        # Whether the transmission curve is the total one or the just-filter one
        if origin:
            # This is the "origin" filter response curve from the HSC webpage:
            # https://www.subarutelescope.org/Observing/Instruments/HSC/sensitivity.html
            # Type of the transmission curve: area weighted one or for center
            filter_dir = os.path.join(FILTER_DIR, 'origin')
            self.type = 'just_filter'
            self.center = False
            if center:
                # This is for the center of the camera
                self.center = True
                self.name = '{}.txt'.format(self.filter)
            else:
                # This is the area-weighted mean one
                self.center = False
                self.name = 'w{}.txt'.format(self.filter)
        else:
            # This is the total transmission curves from Kawanomoto et al. 2018
            # Downloaded from here:
            # https://hsc-release.mtk.nao.ac.jp/doc/wp-content/uploads/2019/04/hsc_responses_all_rev3.tar.gz
            # Assume airmass=1.2 and PWV=1.5
            filter_dir = os.path.join(FILTER_DIR, 'total')
            self.type = 'total'
            self.name = 'hsc_{}_v2018.dat'.format(self.short)

        # Find the file and read in the transmission curve
        self.filename = os.path.join(filter_dir, self.name)
        self.wave, self.trans = self._load_filter()
        self.npts = len(self.wave)

        # Basic properties of the filter
        self.wave_effective = None
        self.wave_pivot = None
        self.wave_mean = None
        self.wave_average = None
        self.rectangular_width = None
        self.gauss_width = None
        self.effective_width = None
        self._basic_properties()

        # AB Zeropoint in counts
        self.ab_zero_counts = self._counts(
            self.wave, AB_FLUX * C_AA_PER_SEC / (self.wave ** 2))

        # AB absolute magnitude of Sun in this band
        self.solar_ab_mag = self._solar_ab_mag(kind='Willmer2018')

        # Convert to Kcorrect filter format
        self.to_kfilter()

        # Save it as a JSON file
        self.to_json()

    def _load_filter(self):
        """Load and process the transimission curve."""
        if not os.path.isfile(self.filename):
            raise IOError("# Cannot find the response curve file {}".format(self.filename))

        # Read in the .txt response curve
        wave, trans = np.genfromtxt(self.filename, usecols=(0, 1), unpack=True)

        use = np.isfinite(trans) & (trans >= 0)
        order = wave[use].argsort()
        wave = wave[use][order]
        trans = trans[use][order]

        return wave, trans

    def print(self):
        """Print out basic properties of the filter."""
        print("# Filter: {}".format(self.filename))
        print("# Filter type: {}".format(self.type))
        print("# Version: {}".format(self.version))
        print("# Effective wavelength   : {:8.3f} Angstrom".format(self.wave_effective))
        print("# Pivot wavelength       : {:8.3f} Angstrom".format(self.wave_pivot))
        print("# Mean wavelength        : {:8.3f} Angstrom".format(self.wave_mean))
        print("# Effective width        : {:8.3f} Angstrom".format(self.effective_width))
        print("# Rectangular width      : {:8.3f} Angstrom".format(self.rectangular_width))
        print("# AB absolute magnitude of the sun : {:6.3f} mag".format(self.solar_ab_mag))

    def plot(self):
        """Plot the transmission curves of the filter."""
        fig = plt.figure(figsize=(8, 5))
        ax1 = fig.add_subplot(111)
        ax1.grid(linestyle='--', linewidth=2, alpha=0.8)
        ax1.axhline(0.0, linewidth=2, color='k', alpha=0.8)
        # Filled the transmission curve
        ax1.fill_between(self.wave, 0.0, self.trans, edgecolor='k', alpha=0.4,
                         linewidth=2.5, facecolor=FILTER_COLORS[self.index])
        ax1.text(0.9, 0.85, r'$\rm {}$'.format(self.short), transform=ax1.transAxes)
        # Highlight the effective and pivot wavelength
        ax1.axvline(self.wave_effective, linewidth=2.0, linestyle='--',
                    label=r'$\lambda_{\rm eff}$')
        ax1.axvline(self.wave_pivot, linewidth=2.0, linestyle='-.',
                    label=r'$\lambda_{\rm piv}$')
        _ = ax1.set_xlabel(r'$\mathrm{Wavelength}\ [\AA]$')
        _ = ax1.set_ylabel(r'$\mathrm{Transmission}$')
        ax1.legend(loc='lower right', fontsize=18)

        return fig

    def _basic_properties(self):
        """Get the basic properties of the filter.

        These properties include several 'effective' wavelength definitions and
        several width definitions.

        See Fukugita et al. (1996) AJ 111, 1748 for discussion and definition
        of many of these quantities.
        """
        # Calculate some useful integrals
        i0 = np.trapz(self.trans * np.log(self.wave), np.log(self.wave))
        i1 = np.trapz(self.trans, np.log(self.wave))
        i2 = np.trapz(self.trans * self.wave, self.wave)
        i3 = np.trapz(self.trans, self.wave)

        # Effective wavelength
        self.wave_effective = np.exp(i0 / i1)
        # Pivot wavelength
        self.wave_pivot = np.sqrt(i2 / i1)
        # Mean wavelength
        self.wave_mean = self.wave_effective
        self.wave_average = i2 / i3
        # Rectangular width of the filter
        self.rectangular_width = i3 / self.trans.max()

        i4 = np.trapz(self.trans * (np.log(self.wave / self.wave_effective)) ** 2.0,
                      np.log(self.wave))

        # Gaussian width of the filter
        self.gauss_width = (i4 / i1) ** (0.5)
        # Effectove width of the filter
        self.effective_width = (2.0 * np.sqrt(2. * np.log(2.)) *
                                self.gauss_width * self.wave_effective)

    def _counts(self, obj_wave, obj_flux):
        """Project source spectrum onto filter and return the detector signal.
        """
        # Interpolate filter transmission to source spectrum
        newtrans = np.interp(obj_wave, self.wave, self.trans, left=0., right=0.)

        # Integrate lambda * f_lambda * R
        if True in (newtrans > 0.):
            positive = np.where(newtrans > 0.)[0]
            ind = slice(max(positive.min() - 1, 0),
                        min(positive.max() + 2, len(obj_wave)))
            counts = np.trapz(obj_wave[ind] * newtrans[ind] *
                              obj_flux[..., ind], obj_wave[ind], axis=-1)
            return np.squeeze(counts)

        return float('NaN')

    def _solar_ab_mag(self, kind='Willmer2018'):
        """Get the absolute magnitude of Sun in this band.
        """
        # Get the solar spectrum
        solar_spectrum = SolarSpectrum(kind=kind)
        counts = self._counts(solar_spectrum.wave, solar_spectrum.flux)
        return -2.5 * np.log10(counts / self.ab_zero_counts)

    def to_kfilter(self):
        """Convert the filter transmission curve into Kcorrect format.
        """
        _ = filters_to_kcorrect(self.filename)

    def as_dict(self):
        """Convert it into a dict."""
        filter_dict = self.__dict__
        filter_dict['wave'] = list(filter_dict['wave'])
        filter_dict['trans'] = list(filter_dict['trans'])
        return filter_dict

    def to_json(self):
        """Save as a JSON file."""
        pre, _ = os.path.splitext(self.filename)
        json_out = pre + '.json'
        with open(json_out, 'w') as jf:
            json.dump(self.as_dict(), jf)

def hsc_filters(origin=False, center=False, use_saved=True):
    """Get the tabel that summarizes information about HSC filters."""
    if origin:
        if center:
            tab_id = 'origin_center'
            xml_table = os.path.join(FILTER_DIR, 'hsc_filters_origin_center.xml')
        else:
            tab_id = 'origin_weighted'
            xml_table = os.path.join(FILTER_DIR, 'hsc_filters_origin_weighted.xml')
    else:
        tab_id = 'total'
        xml_table = os.path.join(FILTER_DIR, 'hsc_filters_total.xml')

    # If already created, just read it in
    if os.path.isfile(xml_table) and use_saved:
        return Table.read(xml_table, format='votable', table_id=tab_id)
    else:
        # Summarize the filters
        # Total transmission curves
        f_table = Table(
            [Filter(f, origin=origin, center=center).as_dict() for f in FILTER_LIST])
        f_table.write(xml_table, format='votable', table_id=tab_id, overwrite=True)
        return f_table

def filters_to_kcorrect(curve_file, verbose=False):
    """
    Convert a filter response curve to the Kcorrect format.

    This is used by Kcorrect and iSEDFit.
    """
    if not os.path.isfile(curve_file):
        raise IOError("# Cannot find the response curve file {}".format(curve_file))

    # Read in the .txt response curve
    wave, response = np.genfromtxt(curve_file, usecols=(0, 1), unpack=True)

    # Output file name
    prefix, _ = os.path.splitext(curve_file)
    output_par = prefix + '.par'

    if os.path.isfile(output_par):
        if verbose:
            print("# Curve {0} is already available".format(output_par))
    else:
        assert len(wave) == len(response), '''
            Wavelength and response curve should have the same size'''

        par = open(output_par, 'w')
        par.write(
            "# %s\n typedef struct {\n  double lambda;\n  double pass;\n } KFILTER;\n\n")
        for w, r in zip(wave, response):
            par.write("KFILTER %10.4f %11.7f\n" % (w, r))
        par.close()

    return wave, response

class SolarSpectrum(object):
    """Class to hold solar spectrum.

    On deriving Solar absolute magnitude:
    ------------------------------------
    To derive the Sun's absolute magnitudes, the IAU 2012 definitions of the astronomical unit
    (au) (Prša et al. 2016) and parsec were used, giving a distance modulus for the
    Sun of −31.5721 mag.

    To rationalize the use of solar constants, the IAU in 2015 adopted a nominal value for the Sun's
    luminosity L⊙ = 3.828 × 10^8 W (Prša et al. 2016), which corresponds to an average
    TSI of 1361 W m−2 at 1 au and an absolute bolometric magnitude of MBol = 4.74.
    """
    def __init__(self, kind='Willmer2018'):
        """Read in the solar spectrum and process it."""
        # Conversion to d=10 pc from 1 AU
        au_to_10pc = (1.0 / (3600 * 180 / np.pi * 10)) ** 2.0

        if kind.strip() == 'Willmer2018':
            # This is based on the work by Willmer 2018:
            #   https://iopscience.iop.org/article/10.3847/1538-4365/aabfdf
            # More information can be found on this webpage:
            #   http://mips.as.arizona.edu/~cnaw/sun.html
            # About the spectrum:
            #   The solar SED used here also combines observations with model spectra.
            #   The observed spectrum is a composite calculated by Haberreiter et al. (2017)
            #   using data from over 20 space-based instruments for an arbitrary date
            #   (2008 December 19, JDN = 2454820) during the solar minimum.
            #   The absolute calibration is set using the ${ATLAS}\,3$ composite spectrum
            #   of Thuillier et al. (2004), and constraining the Total Solar Irradiance (TSI)
            #   to the value measured for each day by Dudok de Wit et al. (2017).
            #   The observed composite ends at ~2.0 μm, and to extend the SED into the infrared;
            #   the model spectra of Fontenla et al. (2011) and Kurucz (2011) are used.
            self.file = os.path.join(SOLAR_DIR, 'sun_composite.fits')
            if os.path.isfile(self.file):
                solar_fits = fits.open(self.file)[1].data
                # Wavelength
                self.wave = solar_fits['WAVE']
                self.wave_unit = 'Angstrom'
                self.flux = solar_fits['FLUX'] * au_to_10pc
                self.flux_unit = 'erg/s/cm^2/AA'
            else:
                raise IOError("# Cannot find the solar spectrum {}".format(self.file))
        elif kind.strip() == 'Kurucz1993':
            # This is the theoretical spectrum of the Sun from Kurucz
            # from: ftp://ftp.stsci.edu/cdbs/grid/k93models/standards/sun_kurucz93.fits
            # The theoretical spectrum is scaled to match the observed spectrum
            # from 1.5 - 2.5 microns, and then it is used where the observed spectrum ends.
            # The theoretical model of the Sun from Kurucz93 atlas using the following
            # parameters when the Sun is at 1 au.
            #   log_Z T_eff log_g V_{Johnson} +0.0 5777 +4.44 -26.75
            self.file = os.path.join(SOLAR_DIR, 'sun_kurucz93.fits')

            # This file should be in AA and erg/s/cm^2/AA at 1AU
            if os.path.isfile(self.file):
                solar_fits = fits.open(self.file)[1].data
                # Wavelength
                self.wave = solar_fits['WAVELENGTH']
                self.wave_unit = 'Angstrom'
                self.flux = solar_fits['FLUX'] * au_to_10pc
                self.flux_unit = 'erg/s/cm^2/AA'
            else:
                raise IOError("# Cannot find the solar spectrum {}".format(self.file))
        else:
            raise NameError("# Wrong type of Solar spectrum: [Willmer2018|Kurucz1993]")

        self.wave_min, self.wave_max = self.wave.min(), self.wave.max()
        self.solar = np.column_stack((self.wave, self.flux))


import numpy
import os
import sys
from main import train

if __name__ == '__main__':
    model_name = os.path.basename(os.path.dirname(os.path.realpath(__file__)))
    train(
    saveto           = '../../models/{}.npz'.format(model_name),
    reload_          = False,
    dim_word         = 300,
    dim              = 300,
    patience         = 7,
    n_words          = 110497,
    n_words_lemma    = 100360,
    decay_c          = 0.,
    clip_c           = 10.,
    lrate            = 0.0004,
    optimizer        = 'adam',
    maxlen           = 100,
    batch_size       = 32,
    valid_batch_size = 32,
    dispFreq         = 100,
    validFreq        = int(549367/32+1),
    saveFreq         = int(549367/32+1),
    use_dropout      = True,
    verbose          = False,
    datasets         = ['../../data/Infotabs/premise_infotabs_train_token.txt',
                        '../../data/Infotabs/hypothesis_infotabs_train_token.txt',
                        '../../data/Infotabs/premise_infotabs_train_lemma.txt',
                        '../../data/Infotabs/hypothesis_infotabs_train_lemma.txt',
                        '../../data/Infotabs/label_infotabs_train.txt'],
    valid_datasets   = ['../../data/Infotabs/premise_infotabs_dev_token.txt',
                        '../../data/Infotabs/hypothesis_infotabs_dev_token.txt',
                        '../../data/Infotabs/premise_infotabs_dev_lemma.txt',
                        '../../data/Infotabs/hypothesis_infotabs_dev_lemma.txt',
                        '../../data/Infotabs/label_infotabs_dev.txt'],
    test_datasets    = ['../../data/Infotabs/premise_infotabs_test_token.txt',
                        '../../data/Infotabs/hypothesis_infotabs_test_token.txt',
                        '../../data/Infotabs/premise_infotabs_test_lemma.txt',
                        '../../data/Infotabs/hypothesis_infotabs_test_lemma.txt',
                        '../../data/Infotabs/label_infotabs_test.txt'],
    dictionary       = ['../../data/Infotabs/vocab_cased.pkl',
                        '../../data/Infotabs/vocab_cased_lemma.pkl'],
    kb_dicts         = ['../../data/Infotabs/pair_features.pkl'],
    embedding        = '../../data/glove/glove.840B.300d.txt',
    dim_kb           = 5,
    kb_inference     = True,
    kb_composition   = False,
    attention_lambda = 0,
    )


#===================================
#        Chengwen Liu              #
#      liuchw2010@gmail.com        #
#   University of Texas at Austin  #
#===================================

import os
import sys
import time
import argparse
import subprocess
import numpy as np
import concurrent.futures
from datetime import datetime

def readxyz():
  lines = open(xyz).readlines()
  natoms = int(lines[0].split()[0])
  nlines = len(lines)
  delta = nlines - natoms
  if delta == 1:
    ftype = "TXYZ"
  elif delta == 2:
    ncols = len(lines[2].split())
    if ncols == 4:
      ftype = "XYZ"
    else:
      ftype = "TXYZ"
  else:
    sys.exit("Could not guess the file type")

  pbc = []
  if ftype == "TXYZ":
    atoms = np.loadtxt(xyz, usecols=(1,), dtype='str', unpack=True, skiprows=delta)
    xs,ys,zs = np.loadtxt(xyz, usecols=(2,3,4), dtype='float', unpack=True, skiprows=delta)
    if delta == 2:
      pbc = [float(x) for x in lines[1].split()]
  else:
    atoms = np.loadtxt(xyz, usecols=(0,), dtype='str', unpack=True, skiprows=delta)
    xs,ys,zs = np.loadtxt(xyz, usecols=(1,2,3), dtype='float', unpack=True, skiprows=delta)
  return atoms, xs, ys, zs, pbc

def handlepdb(pdb):
  atoms, xs, ys, zs, pbc = readxyz()
  natoms = len(atoms)
  lines = open(pdb).readlines()
  nlines = 0
  for line in lines:
    if ("TER" not in line) and ("END" not in line) and ("CRYST1" not in line):
      nlines += 1
  if natoms != nlines:
    sys.exit(f"Unequal # of atoms: {pdb}:{nlines} and {xyz}:{natoms}")
  else:
    with open("result.pdb", 'w') as f:
      idx = 0
      if pbc != []:
        f.write("CRYST1%9.3f%9.3f%9.3f%7.2f%7.2f%7.2f P 1           1\n" % (pbc[0], pbc[1], pbc[2], pbc[3], pbc[4], pbc[5]))
      for i in range(len(lines)):
        if ("TER" not in lines[i]) and ("END" not in lines[i]) and ("CRYST1" not in lines[i]):
          d = lines[i]
          atom = d[-3:].strip()
          x = float(d[30:38])
          y = float(d[38:46])
          z = float(d[46:54])
          s = len(atom)
          if atom.lower() != atoms[i-idx].strip()[0:s].lower():
            sys.exit(f"Atom does not match {atom} --- {atoms[i-idx][0:s]}")
          else:
            newline = d[:30] + "%8.3f%8.3f%8.3f"%(xs[i-idx], ys[i-idx], zs[i-idx]) + d[54:]
            f.write(newline)
        if ("TER" in lines[i]) or ("END" in lines[i]):
          f.write(lines[i])
          idx += 1
  return

if __name__ == "__main__":
  global xyz
  xyz = sys.argv[1]
  pdb = sys.argv[2]
  ## Tinker recognizes certain names
  ions_in_xyzs = {"POT":"K  ", "CLA":"Cl "}
  for ion, ion_ in ions_in_xyzs.items():
    cmd = f"sed 's/{ion}/{ion_}/g' -i {xyz}"
    os.system(cmd)
  ress_in_pdbs = {"CYT":"DC ", "GUA":"DG ", "ADE":"DA ", "THY":"DT "}
  for res, res_ in ress_in_pdbs.items():
    cmd = f"sed 's/{res}/{res_}/g' -i {pdb}"
    os.system(cmd)
  handlepdb(pdb)


from .base_stage import PipelineStage
from .data_types import HDFFile
import numpy as np

class TXExposureInfo(PipelineStage):
    """
    """
    name='TXExposureInfo'

    inputs = [
    ]
    outputs = [
        ('exposures', HDFFile),
    ]
    config_options = {
        'dc2_name': '1.2p',
        'opsim_db': "/global/projecta/projectdirs/lsst/groups/SSim/DC2/minion_1016_desc_dithered_v4.db",
        'propId': 54,
        }

    def run(self):
        from astropy.io import fits
        import sqlite3

        # find butler by name using this tool.
        # change later to general repo path.
        from desc_dc2_dm_data import get_butler
        from lsst.daf.persistence import NoResults

        run = self.config['dc2_name']
        propId = self.config['propId']

        print(f"Getting butler for repo {run}")
        butler = get_butler(run)
        print("Butler loaded")

        # central detector in whole focal plane, just as example
        refs = butler.subset('calexp', raftName='R22', detectorName='S11')
        n = len(refs)
        print(f"Found {n} exposure centers.  Reading exposure info.")

        matching_visits = self.find_matching_opsim_visits()
        nmatch = len(matching_visits)
        print(f"Found list of {nmatch} visits with propId=={propId}")

        float_params =[
            "mjd-obs",
            "bore-ra",
            "bore-dec",
            "bore-az",
            "bore-alt",
            "bore-rotang",
            "bore-airmass",
            "humidity",
            "bgmean",
            "bgvar",
            "magzero_rms",
            "colorterm1",
            "colorterm2",
            "colorterm3",
            "exptime",
            "darktime",

        ]

        int_params = [
            "expid",
            "magzero_nobj",
            "ap_corr_map_id",
            "psf_id",
            "skywcs_id",
        ]

        str_params  = [
            "date-avg",
            "timesys",
            "rottype",
            "filter",
            "testtype",
            "obstype",
        ]
        # We can add WCS information here if needed, there is lots
        # of it.

        params = float_params + int_params + str_params
        # Spaces for output columns
        data = {p:list() for p in params}

        num_params = float_params + int_params
        # Loop through the images and get their metadata
        for i,ref in enumerate(refs):
            # Progress update
            if i%100==0:
                print(f'Reading metadata for exposure {i+1} / {n}')

            # Read the metadata for this exposure reference
            try:
                metadata = butler.get('calexp_md', dataId=ref.dataId).toDict()
            except NoResults:
                continue
            obsid = int(str(metadata['EXPID'])[:-3])
            if obsid not in matching_visits:
                continue

            # columns that we want to save to file, from the metadata
            for p in params:
                data[p].append(metadata[p.upper()])

        m = len(data['bore-ra'])
        f = 100. * m / n
        print(f"{m} / {n} visits match propId={propId} ({f:.2f}%)")

        # Save output
        f = self.open_output('exposures')
        g = f.create_group('exposures')

        for name in num_params:
            g.create_dataset(name, data=data[name])

        for name in str_params:
            # H5PY cannot deal with fixed-length unicode arrays (the numpy default in py3)
            # so convert to ASCII
            g.create_dataset(name, data=np.array(data[name], dtype="S"))

        f.close()

    def find_matching_opsim_visits(self):
        import sqlite3
        db = self.config['opsim_db']
        propId = self.config['propId']
        connection = sqlite3.connect(db)
        cursor = connection.cursor()
        cursor.execute('select obsHistID from summary where propId=:propId', {'propId':propId})
        obsHistID = {Id[0] for Id in cursor.fetchall()}
        return obsHistID


import numpy as np
import torch
import random
import logging
import time
import pickle
import os
from retro_star.common import args, prepare_starting_molecules, prepare_mlp, \
    prepare_molstar_planner, smiles_to_fp
from retro_star.model import ValueMLP
from retro_star.api import RSPlanner

import rdkit.Chem as Chem
from rdkit.Chem import AllChem
import pickle
import pandas as pd

from database import collection_templates, collection_example_reactions_smilesonly
TRANSFORM_DB = collection_templates()
REACTION_DB = collection_example_reactions_smilesonly()
reaction_smiles_field = 'reaction_smiles'

if __name__ == '__main__':
    n_reactions = REACTION_DB.count()
    N = 1000000
    i = -1
    N = min([N, n_reactions])
    target_mols = []

    for example_doc in REACTION_DB.find(no_cursor_timeout = True):
        i+=1
        if i == N:
            i-=1
            break
        reaction_smiles = str(example_doc[reaction_smiles_field])
        reactants, agents, products = [x for x in [mols.split('.') for mols in reaction_smiles.split('>')]]
        if len(products)==1:
            mol = Chem.MolFromSmiles(products[0])
            if mol == None:
                continue
            for atom in mol.GetAtoms():
                if atom.HasProp('molAtomMapNumber'):
                    atom.ClearProp('molAtomMapNumber')
            product = Chem.MolToSmiles(mol)
            if product not in target_mols:
                target_mols.append(product)

    planner = RSPlanner(
    gpu=-1,
    use_value_fn=True,
    iterations=200,
    expansion_topk=50
    )

    file = open("synthetic_routes.txt", "w")

    for i,target_mol in enumerate(target_mols):
        try:
            result = planner.plan(target_mol)
            if len(result) != 0:
                file.write(" ".join(result) + "\n")
                print(i, target_mol, len(result))
        except:
            print("Error occurred")
    file.close()


import os
import glob
import numpy as np

def launch_jobs_ccin2p3(rep_inout='/pbs/home/l/leget/sps_lsst/HSC/gp_output/',
                        filt_telescop = ['z', 'g', 'r', 'r2', 'i', 'i2', 'y'],
                        NBIN=21, MAX = 17.*60.,
                        P0=[3000., 0., 0.],
                        queue = 'lsst',
                        cpu_time = '04:00:00',
                        kernel = "15**2 * AnisotropicVonKarman(invLam=np.array([[1./3000.**2,0],[0,1./3000.**2]]))",
                        rep_mean = None):
    J = 0
    stop = False
    for f in filt_telescop:
        visits = os.path.join(rep_inout, '*%s'%f)
        visits = glob.glob(visits)
        for i in range(len(visits)):
            print(visits[i])

            yaml_name = os.path.join(visits[i], 'config_file.yaml')
            config_yaml = open(yaml_name,'w')
            config_yaml.write('rep: %s\n'%visits[i])
            config_yaml.write('\n')
            config_yaml.write('NBIN: %i\n'%NBIN)
            config_yaml.write('\n')
            config_yaml.write('MAX: %f\n'%MAX)
            config_yaml.write('\n')
            config_yaml.write('P0: {0}\n'.format(P0))
            config_yaml.write('\n')
            config_yaml.write('kernel: %s\n'%kernel)
            if rep_mean is not None:
                config_yaml.write('\n')
                config_yaml.write('rep_mean: %s'%rep_mean)
            config_yaml.close()

            fichier=open('machine_gun_jobs_%i.sh'%(J),'w')
            fichier.write('#!/bin/bash \n')
            fichier.write('\n')
            fichier.write('gastrify %s'%(yaml_name))

            fichier.close()

            o_log = os.path.join(visits[i], "output_o.log")
            e_log = os.path.join(visits[i], "output_e.log")
            command = 'qsub -P P_%s -pe multicores 8 -q mc_long'%(queue)
            command += ' -o %s -e %s -l h_cpu=%s -l sps=1 machine_gun_jobs_%i.sh'%((o_log, e_log, cpu_time, J))
            os.system(command)
            os.system('rm machine_gun_jobs_%i.sh*'%(J))
            J += 1
       #     if J > 10:
       #         stop = True
       #         break
       # if stop:
       #     break


import os
import requests
import base64
import logging
import datetime
import hashlib
from pathlib import Path
import tempfile
import traceback

from synthesizer.inference import Synthesizer
from encoder import inference as encoder
from vocoder import inference as vocoder
import numpy as np
import librosa

logging.basicConfig(level=10, format="%(asctime)s - [%(levelname)8s] - %(name)s - %(message)s")
log = logging.getLogger("voice_cloning")

def clone(audio=None, audio_url=None, sentence=""):
    try:
        if not 10 <= len(sentence.split(" ")) <= 30:
            return {"audio": b"Sentence is invalid! (length must be 10 to 30 words)"}
        audio_data = audio
        if audio_url:
            # Link
            if "http://" in audio_url or "https://" in audio_url:
                header = {'User-Agent': 'Mozilla/5.0 (Windows NT x.y; Win64; x64; rv:9.0) Gecko/20100101 Firefox/10.0'}
                # Check if audio file has less than 5Mb
                r = requests.head(audio_url, headers=header, allow_redirects=True)
                size = r.headers.get('content-length', 0)
                size = int(size) / float(1 << 20)
                log.info("File size: {:.2f} Mb".format(size))
                if size > 5:
                    return {"audio": b"Input audio file is too large! (max 5Mb)"}
                r = requests.get(audio_url, headers=header, allow_redirects=True)
                audio_data = r.content
            # Base64
            elif len(audio_url) > 500:
                audio_data = base64.b64decode(audio_url)

        audio_path = generate_uid() + ".audio"
        with open(audio_path, "wb") as f:
            f.write(audio_data)

        # Load the models one by one.
        log.info("Preparing the encoder, the synthesizer and the vocoder...")
        encoder.load_model(Path("rtvc/encoder/saved_models/pretrained.pt"))
        synthesizer = Synthesizer(Path("rtvc/synthesizer/saved_models/logs-pretrained/taco_pretrained"))
        vocoder.load_model(Path("rtvc/vocoder/saved_models/pretrained/pretrained.pt"))

        # Computing the embedding
        original_wav, sampling_rate = librosa.load(audio_path)
        preprocessed_wav = encoder.preprocess_wav(original_wav, sampling_rate)
        log.info("Loaded file successfully")

        if os.path.exists(audio_path):
            os.remove(audio_path)

        embed = encoder.embed_utterance(preprocessed_wav)
        log.info("Created the embedding")

        specs = synthesizer.synthesize_spectrograms([sentence], [embed])
        spec = np.concatenate(specs, axis=1)
        # spec = specs[0]
        log.info("Created the mel spectrogram")

        # Generating the waveform
        log.info("Synthesizing the waveform:")
        generated_wav = vocoder.infer_waveform(spec, progress_callback=lambda *args: None)

        # Post-generation
        # There's a bug with sounddevice that makes the audio cut one second earlier, so we
        # pad it.
        generated_wav = np.pad(generated_wav,
                               (0, synthesizer.sample_rate),
                               mode="constant")

        # Save it on the disk
        fp = tempfile.TemporaryFile()
        librosa.output.write_wav(fp, generated_wav.astype(np.float32), synthesizer.sample_rate)
        return {"audio": fp.read()}

    except Exception as e:
        log.error(e)
        traceback.print_exc()
        return {"audio": b"Fail"}

def generate_uid():
    # Setting a hash accordingly to the timestamp
    seed = "{}".format(datetime.datetime.now())
    m = hashlib.sha256()
    m.update(seed.encode("utf-8"))
    m = m.hexdigest()
    # Returns only the first and the last 10 hex
    return m[:10] + m[-10:]


from __future__ import print_function
from flask import Flask, request
from structs import *
from copy import deepcopy
import json
import numpy as np
import astar

app = Flask(__name__)
niv = 0
costs = [0, 15000, 50000, 100000, 250000, 500000]

def create_action(action_type, target):
    actionContent = ActionContent(action_type, target.__dict__)
    return json.dumps(actionContent.__dict__)

def create_move_action(target):
    return create_action("MoveAction", target)

def create_attack_action(target):
    return create_action("AttackAction", target)

def create_collect_action(target):
    return create_action("CollectAction", target)

def create_steal_action(target):
    return create_action("StealAction", target)

def create_heal_action():
    return create_action("HealAction", "")

def create_purchase_action(item):
    return create_action("PurchaseAction", item)

def create_upgrade_action(lvl):
    return create_action("UpgradeAction", lvl)

def deserialize_map(serialized_map):
    """
    Fonction utilitaire pour comprendre la map
    """
    serialized_map = serialized_map[1:]
    rows = serialized_map.split('[')
    column = rows[0].split('{')
    deserialized_map = [[Tile() for x in range(20)] for y in range(20)]
    for i in range(len(rows) - 1):
        column = rows[i + 1].split('{')

        for j in range(len(column) - 1):
            infos = column[j + 1].split(',')
            end_index = infos[2].find('}')
            content = int(infos[0])
            x = int(infos[1])
            y = int(infos[2][:end_index])
            deserialized_map[i][j] = Tile(content, x, y)

    return deserialized_map

def bot():
    """
    Main de votre bot.
    """
    print('debut du main')
    map_json = request.form["map"]

    # Player info
    encoded_map = map_json.encode()
    map_json = json.loads(encoded_map)
    p = map_json["Player"]
    pos = p["Position"]
    x = pos["X"]
    y = pos["Y"]
    house = p["HouseLocation"]
    player = Player(p["Health"], p["MaxHealth"], Point(x,y),
                    Point(house["X"], house["Y"]), p["Score"],
                    p["CarriedResources"], p["CarryingCapacity"])

    print('Player info')
    print('Position',pos)
    print('house pos',house)
    print()
    print('Player')
    print('Pos', 'x:', x, 'y:', y)
    print('health', p["Health"])
    print('maxhealth', p["MaxHealth"])
    print('ressources', p["CarriedResources"])
    print('capacity', p["CarryingCapacity"])
    print('Score', p["Score"])
    # Map
    serialized_map = map_json["CustomSerializedMap"]
    deserialized_map = deserialize_map(serialized_map)

    print(serialized_map.replace('],','],\n'))

    #list of tiles
    wall_tiles = []
    house_tiles = []
    lava_tiles = []
    ressource_tiles = []
    shop_tiles = []

    astar_array_res = np.array([[0 for k in range(20)] for l in range(20)])
    astar_array_mais = np.array([[0 for k in range(20)] for l in range(20)])

    for i, row in enumerate(deserialized_map[:20]):
        for j, item in enumerate(row[:20]):
            if item.Content in [1,3,5]:
                astar_array_res[i][j] = 1
                astar_array_mais[i][j] = 1

            if item.Content == 4:
                ressource_tiles.append((i,j))
                astar_array_mais[i][j] = 1

            if item.Content == 2:
                house_tiles.append((i,j))

    print(*astar_array_mais,sep='\n')

    pos_top_corner = (x - 10,y - 10)
    path = False
    r = 0

    a = create_move_action(Point(x,y))

    if x == p["HouseLocation"]["X"] and y == p["HouseLocation"]["Y"]:
        if costs[niv] < p["Score"]:
            a = create_upgrade_action(n)

    if len(ressource_tiles) > 0 and p["CarriedResources"] < p["CarryingCapacity"]: #miner
        while not path and r < len(ressource_tiles):
            pos_cible_rel = ressource_tiles[r]
            pos_cible_abs = (pos_cible_rel[0] + pos_top_corner[0], pos_cible_rel[1] + pos_top_corner[1])
            print("CIBLE : ", pos_cible_abs)

            path = astar.astar(astar_array_res,pos_cible_rel,(10,10))
            r += 1

        if path:
            print("il y a un chemin")
            if len(path) > 1:
                print("PATH: ", path)
                prochain_dep = Point(path[1][0] - 10, path[1][1] - 10)
                print("prochain depl ", prochain_dep)
                a = create_move_action(Point(x + prochain_dep.X, y + prochain_dep.Y))
            else:
                print("LONG ", len(path))
                print("collect")
                a = create_collect_action(Point(pos_cible_abs[0], pos_cible_abs[1]))
        else:
            print("pas de chemin vers les ressources")
            a = create_move_action(Point(x-1,y))

    else: #rentrer chez nous
        cible_x = x
        cible_y = y

        pos_cible_abs = p["HouseLocation"]
        print("CIBLE : ", pos_cible_abs)

        if pos_cible_abs["X"] <= x + 10 and pos_cible_abs["X"] >= x - 10:
            cible_x = pos_cible_abs["X"]

        elif pos_cible_abs["X"] > x + 10:
            cible_x = x + 10

        else:
            cible_x = x - 10

        if pos_cible_abs["Y"] <= y + 10 and pos_cible_abs["Y"] >= y - 10:
            cible_y = pos_cible_abs["Y"]

        elif pos_cible_abs["Y"] > y + 10:
            cible_y = y + 10

        else:
            cible_y = y - 10

        path = astar.astar(astar_array_mais,(cible_x - x + 10,cible_y - y + 10),(10,10))
        if path: #si maison dans champ de vision
            if len(path) > 1:
                prochain_dep = Point(path[1][0] - 10, path[1][1] - 10)
                a = create_move_action(Point(x + prochain_dep.X, y + prochain_dep.Y))
            else:
                print("LONG ", len(path))
                print("maison!")
                a = create_move_action(Point(pos_cible_abs["X"], pos_cible_abs["Y"]))

    print(path)

    #find other players Contient des bugs
    """
    otherPlayers = []

    for player_dict in map_json["OtherPlayers"]:
        for player_name in player_dict.keys():
            player_info = player_dict[player_name]
            p_pos = player_info["Position"]
            player_info = PlayerInfo(player_info["Health"],
                                     player_info["MaxHealth"],
                                     Point(p_pos["X"], p_pos["Y"]))

            otherPlayers.append({player_name: player_info })
    """

    print('Wall')
    print(wall_tiles)
    print('House')
    print(house_tiles)
    print('Lava')
    print(lava_tiles)
    print('Ressource')
    print(ressource_tiles)
    print('Shop')
    print(shop_tiles)

    # return decision
    print('Fin du Main\n')
    print(a)
    return a

@app.route("/", methods=["POST"])
def reponse():
    """
    Point d'entree appelle par le GameServer
    """
    return bot()

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=3000)


import os
import time
from datetime import datetime
import numpy as np
import tensorflow as tf
from tqdm import trange
from inception.slim import slim
from shutil import copyfile
import importlib.util
import subprocess
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

import libs.cylib as cylib
import eval_helper
from train_helper import *
from datasets.cityscapes.cityscapes import CityscapesDataset
from datasets.cityscapes.cityscapes_info import class_info, class_color_map
import datasets.reader as reader

#g_model_path = './models/vgg_16s.py'
g_model_path = './models/vgg_16s_baseline.py'
spec = importlib.util.spec_from_file_location("model", g_model_path)
model = importlib.util.module_from_spec(spec)
spec.loader.exec_module(model)
g_save_dir = '/home/kivan/source/results/semseg/tf/nets/' + get_time_string()
tf.gfile.MakeDirs(g_save_dir)
copyfile(g_model_path, g_save_dir + 'model.py')
print('\ntensorboard --logdir=' + g_save_dir + '\n')

#subprocess.Popen(['tensorboard', '--logdir=' + g_save_dir])
#subprocess.Popen('tensorboard --logdir=' + g_save_dir + ' > /dev/null', shell=True)

FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_string('train_dir', g_save_dir,
    #'/home/kivan/source/deep-learning/semantic_segmentation/tensorflow/results/',
    """Directory where to write event logs """
    """and checkpoint.""")
tf.app.flags.DEFINE_integer('img_width', 1024, '')
tf.app.flags.DEFINE_integer('img_height', 432, '')
tf.app.flags.DEFINE_integer('img_depth', 3, '')

tf.app.flags.DEFINE_string('dataset_dir',
    '/home/kivan/datasets/Cityscapes/tensorflow/records/1024x432/', '')
#tf.app.flags.DEFINE_integer('max_steps', 100000,
#                            """Number of batches to run.""")
tf.app.flags.DEFINE_integer('max_epochs', 20, 'Number of epochs to run.')
tf.app.flags.DEFINE_integer('batch_size', 1, '')
tf.app.flags.DEFINE_integer('num_classes', 19, '')
tf.app.flags.DEFINE_boolean('log_device_placement', False, 'Whether to log device placement.')
#tf.app.flags.DEFINE_boolean('is_training', True, '')

#tf.app.flags.DEFINE_float('initial_learning_rate', 1e-5,
tf.app.flags.DEFINE_float('initial_learning_rate', 1e-4,
#tf.app.flags.DEFINE_float('initial_learning_rate', 1e-3,
#tf.app.flags.DEFINE_float('initial_learning_rate', 1e-4,
                          """Initial learning rate.""")
tf.app.flags.DEFINE_float('num_epochs_per_decay', 3.0,
#tf.app.flags.DEFINE_float('num_epochs_per_decay', 30.0,
                          """Epochs after which learning rate decays.""")
tf.app.flags.DEFINE_float('learning_rate_decay_factor', 0.5,
#tf.app.flags.DEFINE_float('learning_rate_decay_factor', 0.2,
                          """Learning rate decay factor.""")
tf.app.flags.DEFINE_float('moving_average_decay', 0.9999, '')

def evaluate_epoch(dataset, checkpoint_path):
  with tf.Graph().as_default(), tf.Session() as sess:
    image, labels, _, _, img_name, queue = reader.inputs_for_inference(sess, dataset)
    # Restores from checkpoint
    logits = model.inference(image, is_training=False)

    #saver = tf.train.Saver()
    variables_to_restore = tf.get_collection(slim.variables.VARIABLES_TO_RESTORE)
    saver = tf.train.Saver(variables_to_restore)
    saver.restore(sess, checkpoint_path)

    # Start the queue runners.
    coord = tf.train.Coordinator()
    data_threads = tf.train.start_queue_runners(sess=sess, coord=coord)
    #pixel_acc, iou_acc, recall, precision, num_pixels = evaluate_model(sess, dataset, logits, labels)
    conf_mat = np.ascontiguousarray(np.zeros((FLAGS.num_classes, FLAGS.num_classes), dtype=np.uint64))
    for i in trange(dataset.num_examples()):
      out_logits, yt = sess.run([logits, labels])
      y = out_logits[0].argmax(2).astype(np.int32, copy=False)
      yt = yt.astype(np.int32, copy=False)
      cylib.collect_confusion_matrix(y.reshape(-1), yt, conf_mat)
      #print(q_size)
    pixel_acc, iou_acc, recall, precision, num_pixels = eval_helper.compute_errors(conf_mat,
                                                        'Validation', class_info, verbose=True)

    coord.request_stop()
    coord.join(data_threads)
  return pixel_acc, iou_acc, recall, precision

def train_epoch(dataset, resume_path, epoch_num):
  with tf.Graph().as_default():
    #sess = tf.Session(config=tf.ConfigProto(log_device_placement=FLAGS.log_device_placement))
    sess = tf.Session()
    # Create a variable to count the number of train() calls. This equals the
    # number of batches processed * FLAGS.num_gpus.
    global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0),
                                  trainable=False)

    # Calculate the learning rate schedule.
    num_batches_per_epoch = (dataset.num_examples() /
                             FLAGS.batch_size)
    decay_steps = int(num_batches_per_epoch * FLAGS.num_epochs_per_decay)

    # Decay the learning rate exponentially based on the number of steps.
    lr = tf.train.exponential_decay(FLAGS.initial_learning_rate,
                                    global_step,
                                    decay_steps,
                                    FLAGS.learning_rate_decay_factor,
                                    staircase=True)

    # Create an optimizer that performs gradient descent.
    #opt = tf.train.RMSPropOptimizer(lr, RMSPROP_DECAY,
    #                                momentum=RMSPROP_MOMENTUM,
    #                                epsilon=RMSPROP_EPSILON)

    # Get images and labels.
    #image, labels, weights, num_labels, img_name, filename_queue, enqueue_op, \
    image, labels, weights, num_labels, img_name = reader.inputs_single_epoch(sess, dataset)

    # Build a Graph that computes the logits predictions from the
    # inference model.
    logits = model.inference(image)

    # Calculate loss.
    loss = model.loss(logits, labels, weights, num_labels)

    # Build a Graph that trains the model with one batch of examples and
    # updates the model parameters.
    #train_op = model.train(loss, global_step)
      # Variables that affect learning rate.

    # Add a summary to track the learning rate.
    tf.scalar_summary('learning_rate', lr)
    #tf.scalar_summary('learning_rate', tf.mul(lr, tf.constant(1 / FLAGS.initial_learning_rate)))

    # Generate moving averages of all losses and associated summaries.
    #loss_averages_op = add_loss_summaries(loss)

    # Compute gradients.
    #with tf.control_dependencies([loss_averages_op]):
    #opt = tf.train.GradientDescentOptimizer(lr)
    opt = tf.train.AdamOptimizer(lr)
    grads = opt.compute_gradients(loss)

    # Apply gradients.
    apply_gradient_op = opt.apply_gradients(grads, global_step=global_step)

    # Add histograms for trainable variables.
    for var in tf.trainable_variables():
      tf.histogram_summary(var.op.name, var)
    # Add histograms for gradients.
    for grad, var in grads:
      if grad is not None:
        tf.histogram_summary(var.op.name + '/gradients', grad)

    # Track the moving averages of all trainable variables.
    variable_averages = tf.train.ExponentialMovingAverage(FLAGS.moving_average_decay, global_step)
    variables_averages_op = variable_averages.apply(tf.trainable_variables())

    with tf.control_dependencies([apply_gradient_op, variables_averages_op]):
      train_op = tf.no_op(name='train')

    # Create a saver.
    #saver = tf.train.Saver(tf.all_variables(), max_to_keep=FLAGS.max_epochs)
    saver = tf.train.Saver(tf.all_variables())
    if resume_path == None:
      # Build an initialization operation to run below.
      init = tf.initialize_all_variables()
      sess.run(init)
    else:
      assert tf.gfile.Exists(resume_path)
      saver.restore(sess, resume_path)
      #variables_to_restore = tf.get_collection(
      #    slim.variables.VARIABLES_TO_RESTORE)
      #restorer = tf.train.Saver(variables_to_restore)
      #restorer.restore(sess, resume_path)

    # Build the summary operation based on the TF collection of Summaries.
    summary_op = tf.merge_all_summaries()

    # Start running operations on the Graph.
    #dataset.enqueue(sess, enqueue_op, enqueue_placeholder)
    #filename_queue.close()

    # Start the queue runners.
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    #graph_def = sess.graph.as_graph_def(add_shapes=True)
    #summary_writer = tf.train.SummaryWriter(FLAGS.train_dir, graph_def=graph_def)
    summary_writer = tf.train.SummaryWriter(FLAGS.train_dir, graph=sess.graph)

    num_examples = FLAGS.batch_size * FLAGS.img_height * FLAGS.img_width
    ex_start_time = time.time()
    conf_mat = np.ascontiguousarray(np.zeros((FLAGS.num_classes, FLAGS.num_classes), dtype=np.uint64))
    for step in range(dataset.num_examples()):
      start_time = time.time()
      #_, loss_value = sess.run([train_op, loss])
      if step % 100 == 0:
        _, loss_value, scores, yt, clr, filename, global_step_val, summary_str = sess.run(
            [train_op, loss, logits, labels, lr, img_name, global_step, summary_op])
        summary_writer.add_summary(summary_str, global_step_val)
      else:
        _, loss_value, scores, yt, clr, filename, global_step_val = sess.run(
            [train_op, loss, logits, labels, lr, img_name, global_step])
      #loss_val, out_logits, yt = sess.run([loss, logits, labels])
      duration = time.time() - start_time

      assert not np.isnan(loss_value), 'Model diverged with loss = NaN'

      y = scores[0].argmax(2).astype(np.int32, copy=False)
      cylib.collect_confusion_matrix(y.reshape(-1), yt, conf_mat)

      if step % 10 == 0:
        num_examples_per_step = FLAGS.batch_size
        examples_per_sec = num_examples_per_step / duration
        sec_per_batch = float(duration)

        format_str = '%s: epoch %d, step %d, loss = %.2f (%.1f examples/sec; %.3f sec/batch)'
        #print('lr = ', clr)
        print(format_str % (get_expired_time(ex_start_time), epoch_num, step, loss_value,
                            examples_per_sec, sec_per_batch))

    # Save the model checkpoint periodically.
    eval_helper.compute_errors(conf_mat, 'train', class_info)
    checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt')
    saver.save(sess, checkpoint_path, global_step=epoch_num)

    #print('stop')
    coord.request_stop()
    #print('join')
    coord.join(threads)
    sess.close()

  return checkpoint_path + '-' + str(epoch_num)

def train(train_dataset, val_dataset):
  iou_fig = plt.figure()
  pixel_fig = plt.figure()
  iou_acc_data = []
  pixel_acc_data = []
  checkpoint_path = None
  for epoch_num in range(1, FLAGS.max_epochs):
    checkpoint_path = train_epoch(train_dataset, checkpoint_path, epoch_num)
    pixel_acc, iou, recall, precision = evaluate_epoch(val_dataset, checkpoint_path)
    #time.sleep(1)
    iou_acc_data += [iou]
    pixel_acc_data += [pixel_acc]
    #iou_acc_data += [epoch_num]
    #pixel_acc_data += [epoch_num]
    if epoch_num > 1:
      eval_helper.plot_accuracy(iou_fig, iou_acc_data)
      eval_helper.plot_accuracy(pixel_fig, pixel_acc_data)
  input('Press ENTER to exit...')

def main(argv=None):  # pylint: disable=unused-argument
  #if tf.gfile.Exists(FLAGS.train_dir):
  #  tf.gfile.DeleteRecursively(FLAGS.train_dir)
  #tf.gfile.MakeDirs(FLAGS.train_dir)
  print('Experiment dir: ' + FLAGS.train_dir)
  print('Dataset dir: ' + FLAGS.dataset_dir)
  train_dataset = CityscapesDataset(FLAGS.dataset_dir, 'train')
  val_dataset = CityscapesDataset(FLAGS.dataset_dir, 'val')
  train(train_dataset, val_dataset)

if __name__ == '__main__':
  tf.app.run()


from __future__ import division

import math
import os
import subprocess

import numpy as np
from bw_plex import LOG
from bw_plex.misc import sec_to_hh_mm_ss
from bw_plex.video import video_frame_by_frame

try:
    import cv2
except ImportError:
    cv2 = None
    LOG.warning('Scanning for credits is not supported. '
                'Install the package with pip install bw_plex[all] or bw_plex[video]')

try:
    import pytesseract
except ImportError:
    pytesseract = None
    LOG.warning('Extracting text from images is not supported. '
                'Install the package with pip install bw_plex[all] or bw_plex[video]')

try:
    import Image
except ImportError:
    from PIL import Image

# b g, r

# (255, 0, 0)
color = {'yellow': (0, 255, 255),
         'red': (0, 0, 255),
         'blue': (255, 0, 0),
         'lime': (0, 255, 0),
         'white': (255, 255, 255),
         'fuchsia': (255, 0, 255),
         'black': (0, 0, 0)
         }

NET = None

EAST_MODEL = os.path.join(os.path.dirname(__file__), 'models', 'frozen_east_text_detection.pb')

class DEBUG_STOP(Exception):
    pass

def check_stop_in_credits(value, cutoff=1500):
    """Helper to check if the credits are consecutive."""
    for i, v in enumerate(np.diff(value, n=1).tolist()):
        if v > cutoff:
            return i, v

    return None, None

def crop_img(i, edge=0):
    """ crop the image edge % pr side."""
    new_img = i.copy()
    height = new_img.shape[0]
    width = new_img.shape[1]
    sh = int(height / 100 * edge)
    sw = int(width / 100 * edge)

    return new_img[sh:height - sh, sw:width - sw]

def decode(scores, geometry, scoreThresh=0.9999):
    # Stolen from https://github.com/opencv/opencv/blob/master/samples/dnn/text_detection.py
    # scoreTresh is set insanely high as we dont want false positives.
    detections = []
    confidences = []

    # CHECK DIMENSIONS AND SHAPES OF geometry AND scores #
    assert len(scores.shape) == 4, "Incorrect dimensions of scores"
    assert len(geometry.shape) == 4, "Incorrect dimensions of geometry"
    assert scores.shape[0] == 1, "Invalid dimensions of scores"
    assert geometry.shape[0] == 1, "Invalid dimensions of geometry"
    assert scores.shape[1] == 1, "Invalid dimensions of scores"
    assert geometry.shape[1] == 5, "Invalid dimensions of geometry"
    assert scores.shape[2] == geometry.shape[2], "Invalid dimensions of scores and geometry"
    assert scores.shape[3] == geometry.shape[3], "Invalid dimensions of scores and geometry"
    height = scores.shape[2]
    width = scores.shape[3]
    for y in range(0, height):

        # Extract data from scores
        scoresData = scores[0][0][y]
        x0_data = geometry[0][0][y]
        x1_data = geometry[0][1][y]
        x2_data = geometry[0][2][y]
        x3_data = geometry[0][3][y]
        anglesData = geometry[0][4][y]
        for x in range(0, width):
            score = scoresData[x]

            # If score is lower than threshold score, move to next x
            if(score < scoreThresh):
                continue

            # Calculate offset
            offsetX = x * 4.0
            offsetY = y * 4.0
            angle = anglesData[x]

            # Calculate cos and sin of angle
            cosA = math.cos(angle)
            sinA = math.sin(angle)
            h = x0_data[x] + x2_data[x]
            w = x1_data[x] + x3_data[x]

            # Calculate offset
            offset = ([offsetX + cosA * x1_data[x] + sinA * x2_data[x], offsetY - sinA * x1_data[x] + cosA * x2_data[x]])

            # Find points for rectangle
            p1 = (-sinA * h + offset[0], -cosA * h + offset[1])
            p3 = (-cosA * w + offset[0], sinA * w + offset[1])
            center = (0.5 * (p1[0] + p3[0]), 0.5 * (p1[1] + p3[1]))
            detections.append((center, (w, h), -1 * angle * 180.0 / math.pi))
            # This should be the format for non rotation nms boxes
            #detections.append([int(center[0]), int(center[1]), int(w), int(h)])
            confidences.append(float(score))

    #print(detections)
    #print(confidences)

    # Return detections and confidences
    return [detections, confidences]

def make_imgz(afile, start=600, dest=None, fps=1):  # pragma: no cover
    """Helper to generate images."""

    dest_path = dest + '\out%d.jpg'
    fps = 'fps=%s' % fps
    t = sec_to_hh_mm_ss(start)

    cmd = [
        'ffmpeg', '-ss', t, '-i',
        afile, '-vf', fps, dest_path
    ]

    # fix me
    subprocess.call(cmd)
    return dest

def extract_text(img, lang='eng', encoding='utf-8'):
    """Very simple way to find the text in a image, it don't work work well for
       natural scene images but it good enoght for clean frames like credits.

       Ideally we should prop set some roi, but cba for now.

    """
    if pytesseract is None:
        return

    if isinstance(img, str):
        img = Image.open(img)

    return pytesseract.image_to_string(img, lang=lang).encode(encoding, 'ignore')

def calc_success(rectangles, img_height, img_width, success=0.9):  # pragma: no cover
    """Helper to check the n percentage of the image is covered in text."""
    t = sum([i[2] * i[3] for i in rectangles if i])
    p = 100 * float(t) / float(img_height * img_width)
    return p > success

def locate_text_east(image, debug=False, width=320, height=320, confedence_tresh=0.5, nms_treshhold=0):
    import cv2

    global NET

    if NET is None:
        NET = cv2.dnn.readNet(EAST_MODEL)

    features = ['feature_fusion/Conv_7/Sigmoid', 'feature_fusion/concat_3']

    if isinstance(image, str) and os.path.isfile(image):
        image = cv2.imread(image)

    frame = image

    # crop the image as we dont want
    # the logo and burnt in subs.
    # Maybe this edge needs to be a config option # TODO
    frame = crop_img(frame, edge=15)

    # Get frame height and width
    height_ = frame.shape[0]
    width_ = frame.shape[1]
    rW = width_ / float(width)
    rH = height_ / float(height)

    # Create a 4D blob from frame.
    blob = cv2.dnn.blobFromImage(frame, 1.0, (width, height), (123.68, 116.78, 103.94), swapRB=True, crop=False)

    # Run the model
    NET.setInput(blob)
    kWinName = "EAST: An Efficient and Accurate Scene Text Detector"
    # Get scores and geometry
    scores, geometry = NET.forward(features)
    t, _ = NET.getPerfProfile()
    label = 'Inference time: %.2f ms' % (t * 1000.0 / cv2.getTickFrequency())
    boxes, confidences = decode(scores, geometry)
    # print(confidences)
    indices = cv2.dnn.NMSBoxesRotated(boxes, confidences, confedence_tresh, nms_treshhold)
    # Why the fuck dont you work??
    # https://github.com/opencv/opencv/issues/12299
    # indices = cv2.dnn.NMSBoxes(boxes, confidences, confedence_tresh, nms_treshhold)
    locs = []

    if debug is False:
        if isinstance(indices, tuple):
            return []
        else:
            return indices

    for i in indices:
        vertices = cv2.boxPoints(boxes[i[0]])
        # scale the bounding box coordinates based on the respective ratios
        for j in range(4):
            vertices[j][0] *= rW
            vertices[j][1] *= rH

        box = np.int0(vertices)
        cv2.drawContours(frame, [box], 0, color['blue'], 2)
        locs.append([box])
        # print(t)
        # rects will do for now.
        # for j in range(4):
        #     p1 = (vertices[j][0], vertices[j][1])
        #     p2 = (vertices[(j + 1) % 4][0], vertices[(j + 1) % 4][1])
        #     cv2.line(frame, p1, p2, (0, 255, 0), 2)

    # Display the frame
    if debug:
        # Put efficiency information
        cv2.putText(frame, label, (0, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0))
        cv2.imshow(kWinName, frame)

        k = cv2.waitKey(0) & 0xff
        if k == 27:
            raise DEBUG_STOP

    if isinstance(indices, tuple):
        return []
    else:

        return locs

def check_movement(path, debug=True):  # pragma: no cover
    """Nothing usefull atm. TODO"""

    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    #k = np.zeros((3,3),np.uint8)
    #fgbg = cv2.createBackgroundSubtractorGMG()
    #fgbg = cv2.createBackgroundSubtractorMOG2(history=1, varThreshold=2, detectShadows=False)
    #fgbg = cv2.createBackgroundSubtractorMOG()
    # (int history=500, double dist2Threshold=400.0, bool detectShadows=true
    fgbg = cv2.createBackgroundSubtractorKNN(1, 200, False)
    frame = None
    r_size = (640, 480)

    for _, (frame, millisec) in enumerate(video_frame_by_frame(path, offset=0,
                                                               step=0, frame_range=False)):
        if frame is not None:
            fgmask = fgbg.apply(frame)
            fgmask = cv2.erode(fgmask, kernel, iterations=20)
            fgmask = cv2.morphologyEx(fgmask, cv2.MORPH_OPEN, kernel)
            fgmask = cv2.morphologyEx(fgmask, cv2.MORPH_CLOSE, kernel)

            if debug:
                # Need to add a extra channel
                m = cv2.cvtColor(fgmask, cv2.COLOR_GRAY2BGR)
                # Resize so it easier to see them side by side.
                m = cv2.resize(m, r_size)
                f = cv2.resize(frame.copy(), r_size)
                vis = np.concatenate((m, f), axis=1)

                cv2.imshow('frame', vis)

            k = cv2.waitKey(0) & 0xff
            if k == 27:
                break

    # cv2.destroyAllWindows()

def locate_text(image, debug=False):
    """Locate where and if there are text in the images.

       Args:
            image(numpy.ndarray, str): str would be path to image
            debug(bool): Show each of the images using open cv.

       Returns:
            list of rectangles

    """
    # Mostly ripped from https://github.com/hurdlea/Movie-Credits-Detect
    # Thanks!
    import cv2

    # Compat so we can use a frame and img file..
    if isinstance(image, str) and os.path.isfile(image):
        image = cv2.imread(image)

    height, width, _ = image.shape
    mser = cv2.MSER_create(4, 10, 8000, 0.8, 0.2, 200, 1.01, 0.003, 5)
    # Convert to gray.
    grey = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Pull out grahically overlayed text from a video image
    blur = cv2.GaussianBlur(grey, (3, 3), 0)
    adapt_threshold = cv2.adaptiveThreshold(blur, 255, cv2.ADAPTIVE_THRESH_MEAN_C,
                                            cv2.THRESH_BINARY, 5, -25)

    contours, _ = mser.detectRegions(adapt_threshold)

    # for each contour get a bounding box and remove
    rects = []
    for contour in contours:
        # get rectangle bounding contour
        [x, y, w, h] = cv2.boundingRect(contour)

        # Remove small rects
        if w < 5 or h < 5:  # 2
            continue

        # Throw away rectangles which don't match a character aspect ratio
        if (float(w * h) / (width * height)) > 0.005 or float(w) / h > 1:
            continue

        rects.append(cv2.boundingRect(contour))

    # Mask of original image
    mask = np.zeros((height, width, 1), np.uint8)
    # To expand rectangles, i.e. increase sensitivity to nearby rectangles
    # Add knobs?
    # lets scale this alot so we get mostly one big square
    # todo when/if detect motion.
    xscaleFactor = 14  # 14
    yscaleFactor = 4  # 4
    for box in rects:
        [x, y, w, h] = box
        # Draw filled bounding boxes on mask
        cv2.rectangle(mask, (x - xscaleFactor, y - yscaleFactor),
                      (x + w + xscaleFactor, y + h + yscaleFactor),
                      color['white'], cv2.FILLED)

    # Find contours in mask if bounding boxes overlap,
    # they will be joined by this function call
    rectangles = []
    contours = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    if cv2.__version__.startswith('4'):
        contours = contours[0]
    else:
        contours = contours[1]

    for contour in contours:
        # This is disabled since we are not after the text but the text area.
        # Only preserve "squarish" features
        # peri = cv2.arcLength(contour, True)
        # approx = cv2.approxPolyDP(contour, 0.01 * peri, True)

        # the contour is 'bad' if it is not a rectangluarish
        # This doesnt have to be bad, since we match more then one char.
        # if len(approx) > 8:
        #    cv2.drawContours(image, [contour], -1, color['lime'])
        #    if debug:
        #        cv2.imshow("bad Rectangles check lime", image)
        #    continue

        rect = cv2.boundingRect(contour)

        x, y, w, h = rect
        cv2.rectangle(image, (x, y), (x + w, y + h), color['blue'], 2)

        # Remove small areas and areas that don't have text like features
        # such as a long width.
        if ((float(w * h) / (width * height)) < 0.006):
            # remove small areas
            if float(w * h) / (width * height) < 0.0018:
                continue

            # remove areas that aren't long
            if (float(w) / h < 2.5):
                continue

        else:
            pass
            # This is disabled as we want to cache the large area of text
            # to backup this shit for movement detection
            # General catch for larger identified areas that they have
            # a text width profile
            # and it does not fit for jap letters.

            # if float(w) / h < 1.8:
            #    continue

        rectangles.append(rect)

        cv2.rectangle(image, (x, y), (x + w, y + h), color['fuchsia'], 2)

    if debug:
        cv2.imshow("Final image", image)
        cv2.waitKey(0)

    return rectangles

def find_credits(path, offset=0, fps=None, duration=None,
                 check=7, step=1, frame_range=True, debug=False, method='east'):
    """Find the start/end of the credits and end in a videofile.
       This only check frames so if there is any silence in the video this is simply skipped as
       opencv only handles videofiles.

       use frame_range to so we only check frames every 1 sec.

       # TODO just ffmepg to check for silence so we calculate the correct time? :(

       Args:
            path (str): path to the videofile
            offset(int): If given we should start from this one.
            fps(float?): fps of the video file
            duration(None, int): Duration of the vfile in seconds.
            check(int): Stop after n frames with text, set a insane high number to check all.
                        end is not correct without this!
            step(int): only use every n frame
            frame_range(bool). default true, precalc the frames and only check thous frames.
            debug(bool): Disable the images.
            method(str): east is better but slower.

       Returns:
            1, 2

    """
    # LOG.debug('%r %r %r %r %r %r %r', path, offset, fps, duration, check, step, frame_range)
    if cv2 is None:
        return
    frames = []
    start = -1
    end = -1
    LOG.debug('Trying to find the credits for %s', path)

    if method == 'east':
        func = locate_text_east
    else:
        func = locate_text

    try:
        if fps is None:
            # we can just grab the fps from plex.
            cap = cv2.VideoCapture(path)
            fps = cap.get(cv2.CAP_PROP_FPS)
            cap.release()

        for _, (frame, millisec) in enumerate(video_frame_by_frame(path, offset=offset,
                                                                   step=step, frame_range=frame_range)):

            try:
                # LOG.debug('progress %s', millisec / 1000)
                if frame is not None:
                    # recs = locate_text(frame, debug=True)
                    recs = func(frame, debug=debug)
                    len_recs = len(recs)

                    # If we get 1 match we should verify.
                    # now this is pretty harsh but we really
                    # don't want false positives.
                    if len_recs == 0:
                        continue
                    elif len_recs == 1:
                        t = extract_text(frame)
                        if t:
                            frames.append(millisec)
                    else:
                        frames.append(millisec)

                    # check for motion here?

                    if check != -1 and len(frames) >= check:
                        break

            except DEBUG_STOP:
                break
                if hasattr(cv2, 'destroyAllWindows'):
                    cv2.destroyAllWindows()

        if frames:
            start = min(frames) / 1000
            end = max(frames) / 1000

        LOG.debug('credits_start %s, credits_end %s', start, end)

    except:  # pragma: no cover
        # We just want to log the exception not halt the entire process to db.
        LOG.exception('There was a error in find_credits')

    return start, end

def fill_rects(image, rects):  # pragma: no cover
    """This is used to fill the rects (location of credits)

       The idea if to mask the credits so we can check if there is any background
       movement while the credits are running. Like the movie cars etc.
       See if we can grab something usefull from
       https://gist.github.com/luipillmann/d76eb4f4eea0320bb35dcd1b2a4575ee
    """
    for rect in rects:
        try:
            x, y, w, h = rect
            cv2.rectangle(image, (x, y), (x + w, y + h), color['black'], cv2.FILLED)
        except ValueError:
            cv2.drawContours(image, rect, 0, color['black'], cv2.FILLED)

    return image

if __name__ == '__main__':
    def test():
        # import logging
        # logging.basicConfig(level=logging.DEBUG)
        # import cv2
        # i = r"C:\Users\steff\Documents\GitHub\bw_plex\tests\test_data\blacktext_whitebg_2.png"
        # i = r'C:\Users\alexa\.config\bw_plex\third_images\out165.jpg'

        # img = cv2.imread(i)
        # ffs = img.copy()
        # rects = locate_text(ffs, debug=True)
        # locate_text2(img, debug=True, width=320, height=320, confedence_tresh=0.8, nms_treshhold=0.1)
        #out = r'C:\stuff\GUNDAM BUILD FIGHTERS TRY-Episode 1 - The Boy Who Calls The Wind (ENG sub)-M7fLOQXlPmE.mkv' # 21*60
        # out = r'C:\Users\steff\Documents\GitHub\bw_plex\tests\test_data\part.mkv'
        out = r'C:\Users\steff\Documents\GitHub\bw_plex\tests\test_data\out.mkv'
        t = find_credits(out, offset=0, fps=None, duration=None, check=600, step=1, frame_range=True, debug=True)
        print(t)
        # print(out)
        # for z in check_movement(out):
        #     print()

        # f = fill_rects(img, rects)
        # cv2.imshow('ass', f)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()

    test()