# ******************************************************************************
# Copyright 2017-2018 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ******************************************************************************

import logging
import pickle

import numpy as np

from nlp_architect.common.cdc.mention_data import MentionDataLight

logger = logging.getLogger(__name__)

class GloveEmbedding(object):
    def __init__(self, glove_file):
        logger.info("Loading Glove Online Embedding module, This my take a while...")
        self.word_to_ix, self.word_embeddings = self.load_glove_for_vocab(glove_file)
        logger.info("Glove Offline Embedding module lead successfully")

    @staticmethod
    def load_glove_for_vocab(glove_filename):
        vocab = []
        embd = []
        with open(glove_filename) as glove_file:
            for line in glove_file:
                row = line.strip().split(" ")
                word = row[0]
                vocab.append(word)
                embd.append(row[1:])

        embeddings = np.asarray(embd, dtype=float)
        word_to_ix = {word: i for i, word in enumerate(vocab)}
        return word_to_ix, embeddings

class GloveEmbeddingOffline(object):
    def __init__(self, embed_resources):
        logger.info("Loading Glove Offline Embedding module")
        with open(embed_resources, "rb") as out:
            self.word_to_ix, self.word_embeddings = pickle.load(out, encoding="latin1")
        logger.info("Glove Offline Embedding module lead successfully")

    def get_feature_vector(self, mention: MentionDataLight):
        embed = None
        head = mention.mention_head
        lemma = mention.mention_head_lemma
        if head in self.word_to_ix:
            embed = self.word_embeddings[self.word_to_ix[head]]
        elif lemma in self.word_to_ix:
            embed = self.word_embeddings[self.word_to_ix[lemma]]

        return embed

    def get_avrg_feature_vector(self, tokens_str):
        embed = np.zeros(300, dtype=np.float64)
        mention_size = 0
        for token in tokens_str.split():
            if token in self.word_to_ix:
                token_embed = self.word_embeddings[self.word_to_ix[token]]
                embed = np.add(embed, token_embed)
                mention_size += 1

        if mention_size == 0:
            mention_size = 1

        return np.true_divide(embed, mention_size)


# Licensed under a 3-clause BSD style license - see LICENSE.rst

from warnings import warn
import collections
import socket
import json
import urllib.request
import urllib.error
import urllib.parse

import numpy as np
import erfa

from astropy import units as u
from astropy import constants as consts
from astropy.units.quantity import QuantityInfoBase
from astropy.utils.exceptions import AstropyUserWarning
from .angles import Angle, Longitude, Latitude
from .representation import CartesianRepresentation, CartesianDifferential
from .errors import UnknownSiteException
from astropy.utils import data

__all__ = ['EarthLocation']

GeodeticLocation = collections.namedtuple('GeodeticLocation', ['lon', 'lat', 'height'])

# Available ellipsoids (defined in erfam.h, with numbers exposed in erfa).
ELLIPSOIDS = ('WGS84', 'GRS80', 'WGS72')

OMEGA_EARTH = u.Quantity(7.292115855306589e-5, 1./u.s)
"""
Rotational velocity of Earth. In UT1 seconds, this would be 2 pi / (24 * 3600),
but we need the value in SI seconds.
See Explanatory Supplement to the Astronomical Almanac, ed. P. Kenneth Seidelmann (1992),
University Science Books.
"""

def _check_ellipsoid(ellipsoid=None, default='WGS84'):
    if ellipsoid is None:
        ellipsoid = default
    if ellipsoid not in ELLIPSOIDS:
        raise ValueError(f'Ellipsoid {ellipsoid} not among known ones ({ELLIPSOIDS})')
    return ellipsoid

def _get_json_result(url, err_str, use_google):

    # need to do this here to prevent a series of complicated circular imports
    from .name_resolve import NameResolveError
    try:
        # Retrieve JSON response from Google maps API
        resp = urllib.request.urlopen(url, timeout=data.conf.remote_timeout)
        resp_data = json.loads(resp.read().decode('utf8'))

    except urllib.error.URLError as e:
        # This catches a timeout error, see:
        #   http://stackoverflow.com/questions/2712524/handling-urllib2s-timeout-python
        if isinstance(e.reason, socket.timeout):
            raise NameResolveError(err_str.format(msg="connection timed out"))
        else:
            raise NameResolveError(err_str.format(msg=e.reason))

    except socket.timeout:
        # There are some cases where urllib2 does not catch socket.timeout
        # especially while receiving response data on an already previously
        # working request
        raise NameResolveError(err_str.format(msg="connection timed out"))

    if use_google:
        results = resp_data.get('results', [])

        if resp_data.get('status', None) != 'OK':
            raise NameResolveError(err_str.format(msg="unknown failure with "
                                                  "Google API"))

    else:  # OpenStreetMap returns a list
        results = resp_data

    if not results:
        raise NameResolveError(err_str.format(msg="no results returned"))

    return results

class EarthLocationInfo(QuantityInfoBase):
    """
    Container for meta information like name, description, format.  This is
    required when the object is used as a mixin column within a table, but can
    be used as a general way to store meta information.
    """
    _represent_as_dict_attrs = ('x', 'y', 'z', 'ellipsoid')

    def _construct_from_dict(self, map):
        # Need to pop ellipsoid off and update post-instantiation.  This is
        # on the to-fix list in #4261.
        ellipsoid = map.pop('ellipsoid')
        out = self._parent_cls(**map)
        out.ellipsoid = ellipsoid
        return out

    def new_like(self, cols, length, metadata_conflicts='warn', name=None):
        """
        Return a new EarthLocation instance which is consistent with the
        input ``cols`` and has ``length`` rows.

        This is intended for creating an empty column object whose elements can
        be set in-place for table operations like join or vstack.

        Parameters
        ----------
        cols : list
            List of input columns
        length : int
            Length of the output column object
        metadata_conflicts : str ('warn'|'error'|'silent')
            How to handle metadata conflicts
        name : str
            Output column name

        Returns
        -------
        col : EarthLocation (or subclass)
            Empty instance of this class consistent with ``cols``
        """
        # Very similar to QuantityInfo.new_like, but the creation of the
        # map is different enough that this needs its own rouinte.
        # Get merged info attributes shape, dtype, format, description.
        attrs = self.merge_cols_attributes(cols, metadata_conflicts, name,
                                           ('meta', 'format', 'description'))
        # The above raises an error if the dtypes do not match, but returns
        # just the string representation, which is not useful, so remove.
        attrs.pop('dtype')
        # Make empty EarthLocation using the dtype and unit of the last column.
        # Use zeros so we do not get problems for possible conversion to
        # geodetic coordinates.
        shape = (length,) + attrs.pop('shape')
        data = u.Quantity(np.zeros(shape=shape, dtype=cols[0].dtype),
                          unit=cols[0].unit, copy=False)
        # Get arguments needed to reconstruct class
        map = {key: (data[key] if key in 'xyz' else getattr(cols[-1], key))
               for key in self._represent_as_dict_attrs}
        out = self._construct_from_dict(map)
        # Set remaining info attributes
        for attr, value in attrs.items():
            setattr(out.info, attr, value)

        return out

class EarthLocation(u.Quantity):
    """
    Location on the Earth.

    Initialization is first attempted assuming geocentric (x, y, z) coordinates
    are given; if that fails, another attempt is made assuming geodetic
    coordinates (longitude, latitude, height above a reference ellipsoid).
    When using the geodetic forms, Longitudes are measured increasing to the
    east, so west longitudes are negative. Internally, the coordinates are
    stored as geocentric.

    To ensure a specific type of coordinates is used, use the corresponding
    class methods (`from_geocentric` and `from_geodetic`) or initialize the
    arguments with names (``x``, ``y``, ``z`` for geocentric; ``lon``, ``lat``,
    ``height`` for geodetic).  See the class methods for details.

    Notes
    -----
    This class fits into the coordinates transformation framework in that it
    encodes a position on the `~astropy.coordinates.ITRS` frame.  To get a
    proper `~astropy.coordinates.ITRS` object from this object, use the ``itrs``
    property.
    """

    _ellipsoid = 'WGS84'
    _location_dtype = np.dtype({'names': ['x', 'y', 'z'],
                                'formats': [np.float64]*3})
    _array_dtype = np.dtype((np.float64, (3,)))

    info = EarthLocationInfo()

    def __new__(cls, *args, **kwargs):
        # TODO: needs copy argument and better dealing with inputs.
        if (len(args) == 1 and len(kwargs) == 0 and
                isinstance(args[0], EarthLocation)):
            return args[0].copy()
        try:
            self = cls.from_geocentric(*args, **kwargs)
        except (u.UnitsError, TypeError) as exc_geocentric:
            try:
                self = cls.from_geodetic(*args, **kwargs)
            except Exception as exc_geodetic:
                raise TypeError('Coordinates could not be parsed as either '
                                'geocentric or geodetic, with respective '
                                'exceptions "{}" and "{}"'
                                .format(exc_geocentric, exc_geodetic))
        return self

    @classmethod
    def from_geocentric(cls, x, y, z, unit=None):
        """
        Location on Earth, initialized from geocentric coordinates.

        Parameters
        ----------
        x, y, z : `~astropy.units.Quantity` or array_like
            Cartesian coordinates.  If not quantities, ``unit`` should be given.
        unit : `~astropy.units.UnitBase` object or None
            Physical unit of the coordinate values.  If ``x``, ``y``, and/or
            ``z`` are quantities, they will be converted to this unit.

        Raises
        ------
        astropy.units.UnitsError
            If the units on ``x``, ``y``, and ``z`` do not match or an invalid
            unit is given.
        ValueError
            If the shapes of ``x``, ``y``, and ``z`` do not match.
        TypeError
            If ``x`` is not a `~astropy.units.Quantity` and no unit is given.
        """
        if unit is None:
            try:
                unit = x.unit
            except AttributeError:
                raise TypeError("Geocentric coordinates should be Quantities "
                                "unless an explicit unit is given.")
        else:
            unit = u.Unit(unit)

        if unit.physical_type != 'length':
            raise u.UnitsError("Geocentric coordinates should be in "
                               "units of length.")

        try:
            x = u.Quantity(x, unit, copy=False)
            y = u.Quantity(y, unit, copy=False)
            z = u.Quantity(z, unit, copy=False)
        except u.UnitsError:
            raise u.UnitsError("Geocentric coordinate units should all be "
                               "consistent.")

        x, y, z = np.broadcast_arrays(x, y, z)
        struc = np.empty(x.shape, cls._location_dtype)
        struc['x'], struc['y'], struc['z'] = x, y, z
        return super().__new__(cls, struc, unit, copy=False)

    @classmethod
    def from_geodetic(cls, lon, lat, height=0., ellipsoid=None):
        """
        Location on Earth, initialized from geodetic coordinates.

        Parameters
        ----------
        lon : `~astropy.coordinates.Longitude` or float
            Earth East longitude.  Can be anything that initialises an
            `~astropy.coordinates.Angle` object (if float, in degrees).
        lat : `~astropy.coordinates.Latitude` or float
            Earth latitude.  Can be anything that initialises an
            `~astropy.coordinates.Latitude` object (if float, in degrees).
        height : `~astropy.units.Quantity` or float, optional
            Height above reference ellipsoid (if float, in meters; default: 0).
        ellipsoid : str, optional
            Name of the reference ellipsoid to use (default: 'WGS84').
            Available ellipsoids are:  'WGS84', 'GRS80', 'WGS72'.

        Raises
        ------
        astropy.units.UnitsError
            If the units on ``lon`` and ``lat`` are inconsistent with angular
            ones, or that on ``height`` with a length.
        ValueError
            If ``lon``, ``lat``, and ``height`` do not have the same shape, or
            if ``ellipsoid`` is not recognized as among the ones implemented.

        Notes
        -----
        For the conversion to geocentric coordinates, the ERFA routine
        ``gd2gc`` is used.  See https://github.com/liberfa/erfa
        """
        ellipsoid = _check_ellipsoid(ellipsoid, default=cls._ellipsoid)
        # We use Angle here since there is no need to wrap the longitude -
        # gd2gc will just take cos/sin anyway.  And wrapping might fail
        # on readonly input.
        lon = Angle(lon, u.degree, copy=False)
        lat = Latitude(lat, u.degree, copy=False)
        # don't convert to m by default, so we can use the height unit below.
        if not isinstance(height, u.Quantity):
            height = u.Quantity(height, u.m, copy=False)
        # get geocentric coordinates. Have to give one-dimensional array.
        xyz = erfa.gd2gc(getattr(erfa, ellipsoid),
                         lon.to_value(u.radian),
                         lat.to_value(u.radian),
                         height.to_value(u.m))
        self = xyz.ravel().view(cls._location_dtype,
                                cls).reshape(xyz.shape[:-1])
        self._unit = u.meter
        self._ellipsoid = ellipsoid
        return self.to(height.unit)

    @classmethod
    def of_site(cls, site_name):
        """
        Return an object of this class for a known observatory/site by name.

        This is intended as a quick convenience function to get basic site
        information, not a fully-featured exhaustive registry of observatories
        and all their properties.

        Additional information about the site is stored in the ``.info.meta``
        dictionary of sites obtained using this method (see the examples below).

        .. note::
            When this function is called, it will attempt to download site
            information from the astropy data server. If you would like a site
            to be added, issue a pull request to the
            `astropy-data repository <https://github.com/astropy/astropy-data>`_ .
            If a site cannot be found in the registry (i.e., an internet
            connection is not available), it will fall back on a built-in list,
            In the future, this bundled list might include a version-controlled
            list of canonical observatories extracted from the online version,
            but it currently only contains the Greenwich Royal Observatory as an
            example case.

        Parameters
        ----------
        site_name : str
            Name of the observatory (case-insensitive).

        Returns
        -------
        site : This class (a `~astropy.coordinates.EarthLocation` or subclass)
            The location of the observatory.

        Examples
        --------

        >>> from astropy.coordinates import EarthLocation
        >>> keck = EarthLocation.of_site('Keck Observatory')  # doctest: +REMOTE_DATA
        >>> keck.geodetic  # doctest: +REMOTE_DATA +FLOAT_CMP
        GeodeticLocation(lon=<Longitude -155.47833333 deg>, lat=<Latitude 19.82833333 deg>, height=<Quantity 4160. m>)
        >>> keck.info  # doctest: +REMOTE_DATA
        name = W. M. Keck Observatory
        dtype = void192
        unit = m
        class = EarthLocation
        n_bad = 0
        >>> keck.info.meta  # doctest: +REMOTE_DATA
        {'source': 'IRAF Observatory Database', 'timezone': 'US/Aleutian'}

        See Also
        --------
        get_site_names : the list of sites that this function can access
        """  # noqa
        registry = cls._get_site_registry()
        try:
            el = registry[site_name]
        except UnknownSiteException as e:
            raise UnknownSiteException(e.site, 'EarthLocation.get_site_names',
                                       close_names=e.close_names)

        if cls is el.__class__:
            return el
        else:
            newel = cls.from_geodetic(*el.to_geodetic())
            newel.info.name = el.info.name
            return newel

    @classmethod
    def of_address(cls, address, get_height=False, google_api_key=None):
        """
        Return an object of this class for a given address by querying either
        the OpenStreetMap Nominatim tool [1]_ (default) or the Google geocoding
        API [2]_, which requires a specified API key.

        This is intended as a quick convenience function to get easy access to
        locations. If you need to specify a precise location, you should use the
        initializer directly and pass in a longitude, latitude, and elevation.

        In the background, this just issues a web query to either of
        the APIs noted above. This is not meant to be abused! Both
        OpenStreetMap and Google use IP-based query limiting and will ban your
        IP if you send more than a few thousand queries per hour [2]_.

        .. warning::
            If the query returns more than one location (e.g., searching on
            ``address='springfield'``), this function will use the **first**
            returned location.

        Parameters
        ----------
        address : str
            The address to get the location for. As per the Google maps API,
            this can be a fully specified street address (e.g., 123 Main St.,
            New York, NY) or a city name (e.g., Danbury, CT), or etc.
        get_height : bool, optional
            This only works when using the Google API! See the ``google_api_key``
            block below. Use the retrieved location to perform a second query to
            the Google maps elevation API to retrieve the height of the input
            address [3]_.
        google_api_key : str, optional
            A Google API key with the Geocoding API and (optionally) the
            elevation API enabled. See [4]_ for more information.

        Returns
        -------
        location : This class (a `~astropy.coordinates.EarthLocation` or subclass)
            The location of the input address.

        References
        ----------
        .. [1] https://nominatim.openstreetmap.org/
        .. [2] https://developers.google.com/maps/documentation/geocoding/start
        .. [3] https://developers.google.com/maps/documentation/elevation/start
        .. [4] https://developers.google.com/maps/documentation/geocoding/get-api-key

        """

        use_google = google_api_key is not None

        # Fail fast if invalid options are passed:
        if not use_google and get_height:
            raise ValueError(
                'Currently, `get_height` only works when using '
                'the Google geocoding API, which requires passing '
                'a Google API key with `google_api_key`. See: '
                'https://developers.google.com/maps/documentation/geocoding/get-api-key '
                'for information on obtaining an API key.')

        if use_google:  # Google
            pars = urllib.parse.urlencode({'address': address,
                                           'key': google_api_key})
            geo_url = f"https://maps.googleapis.com/maps/api/geocode/json?{pars}"

        else:  # OpenStreetMap
            pars = urllib.parse.urlencode({'q': address,
                                           'format': 'json'})
            geo_url = f"https://nominatim.openstreetmap.org/search?{pars}"

        # get longitude and latitude location
        err_str = f"Unable to retrieve coordinates for address '{address}'; {{msg}}"
        geo_result = _get_json_result(geo_url, err_str=err_str,
                                      use_google=use_google)

        if use_google:
            loc = geo_result[0]['geometry']['location']
            lat = loc['lat']
            lon = loc['lng']

        else:
            loc = geo_result[0]
            lat = float(loc['lat'])  # strings are returned by OpenStreetMap
            lon = float(loc['lon'])

        if get_height:
            pars = {'locations': f'{lat:.8f},{lon:.8f}',
                    'key': google_api_key}
            pars = urllib.parse.urlencode(pars)
            ele_url = f"https://maps.googleapis.com/maps/api/elevation/json?{pars}"

            err_str = f"Unable to retrieve elevation for address '{address}'; {{msg}}"
            ele_result = _get_json_result(ele_url, err_str=err_str,
                                          use_google=use_google)
            height = ele_result[0]['elevation']*u.meter

        else:
            height = 0.

        return cls.from_geodetic(lon=lon*u.deg, lat=lat*u.deg, height=height)

    @classmethod
    def get_site_names(cls):
        """
        Get list of names of observatories for use with
        `~astropy.coordinates.EarthLocation.of_site`.

        .. note::
            When this function is called, it will first attempt to
            download site information from the astropy data server.  If it
            cannot (i.e., an internet connection is not available), it will fall
            back on the list included with astropy (which is a limited and dated
            set of sites).  If you think a site should be added, issue a pull
            request to the
            `astropy-data repository <https://github.com/astropy/astropy-data>`_ .

        Returns
        -------
        names : list of str
            List of valid observatory names

        See Also
        --------
        of_site : Gets the actual location object for one of the sites names
                  this returns.
        """
        return cls._get_site_registry().names

    @classmethod
    def _get_site_registry(cls, force_download=False, force_builtin=False):
        """
        Gets the site registry.  The first time this either downloads or loads
        from the data file packaged with astropy.  Subsequent calls will use the
        cached version unless explicitly overridden.

        Parameters
        ----------
        force_download : bool or str
            If not False, force replacement of the cached registry with a
            downloaded version. If a str, that will be used as the URL to
            download from (if just True, the default URL will be used).
        force_builtin : bool
            If True, load from the data file bundled with astropy and set the
            cache to that.

        returns
        -------
        reg : astropy.coordinates.sites.SiteRegistry
        """
        # need to do this here at the bottom to avoid circular dependencies
        from .sites import get_builtin_sites, get_downloaded_sites

        if force_builtin and force_download:
            raise ValueError('Cannot have both force_builtin and force_download True')

        if force_builtin:
            reg = cls._site_registry = get_builtin_sites()
        else:
            reg = getattr(cls, '_site_registry', None)
            if force_download or not reg:
                try:
                    if isinstance(force_download, str):
                        reg = get_downloaded_sites(force_download)
                    else:
                        reg = get_downloaded_sites()
                except OSError:
                    if force_download:
                        raise
                    msg = ('Could not access the online site list. Falling '
                           'back on the built-in version, which is rather '
                           'limited. If you want to retry the download, do '
                           '{0}._get_site_registry(force_download=True)')
                    warn(AstropyUserWarning(msg.format(cls.__name__)))
                    reg = get_builtin_sites()
                cls._site_registry = reg

        return reg

    @property
    def ellipsoid(self):
        """The default ellipsoid used to convert to geodetic coordinates."""
        return self._ellipsoid

    @ellipsoid.setter
    def ellipsoid(self, ellipsoid):
        self._ellipsoid = _check_ellipsoid(ellipsoid)

    @property
    def geodetic(self):
        """Convert to geodetic coordinates for the default ellipsoid."""
        return self.to_geodetic()

    def to_geodetic(self, ellipsoid=None):
        """Convert to geodetic coordinates.

        Parameters
        ----------
        ellipsoid : str, optional
            Reference ellipsoid to use.  Default is the one the coordinates
            were initialized with.  Available are: 'WGS84', 'GRS80', 'WGS72'

        Returns
        -------
        (lon, lat, height) : tuple
            The tuple contains instances of `~astropy.coordinates.Longitude`,
            `~astropy.coordinates.Latitude`, and `~astropy.units.Quantity`

        Raises
        ------
        ValueError
            if ``ellipsoid`` is not recognized as among the ones implemented.

        Notes
        -----
        For the conversion to geodetic coordinates, the ERFA routine
        ``gc2gd`` is used.  See https://github.com/liberfa/erfa
        """
        ellipsoid = _check_ellipsoid(ellipsoid, default=self.ellipsoid)
        self_array = self.to(u.meter).view(self._array_dtype, np.ndarray)
        lon, lat, height = erfa.gc2gd(getattr(erfa, ellipsoid), self_array)
        return GeodeticLocation(
            Longitude(lon * u.radian, u.degree,
                      wrap_angle=180.*u.degree, copy=False),
            Latitude(lat * u.radian, u.degree, copy=False),
            u.Quantity(height * u.meter, self.unit, copy=False))

    @property
    def lon(self):
        """Longitude of the location, for the default ellipsoid."""
        return self.geodetic[0]

    @property
    def lat(self):
        """Longitude of the location, for the default ellipsoid."""
        return self.geodetic[1]

    @property
    def height(self):
        """Height of the location, for the default ellipsoid."""
        return self.geodetic[2]

    # mostly for symmetry with geodetic and to_geodetic.
    @property
    def geocentric(self):
        """Convert to a tuple with X, Y, and Z as quantities"""
        return self.to_geocentric()

    def to_geocentric(self):
        """Convert to a tuple with X, Y, and Z as quantities"""
        return (self.x, self.y, self.z)

    def get_itrs(self, obstime=None):
        """
        Generates an `~astropy.coordinates.ITRS` object with the location of
        this object at the requested ``obstime``.

        Parameters
        ----------
        obstime : `~astropy.time.Time` or None
            The ``obstime`` to apply to the new `~astropy.coordinates.ITRS`, or
            if None, the default ``obstime`` will be used.

        Returns
        -------
        itrs : `~astropy.coordinates.ITRS`
            The new object in the ITRS frame
        """
        # Broadcast for a single position at multiple times, but don't attempt
        # to be more general here.
        if obstime and self.size == 1 and obstime.shape:
            self = np.broadcast_to(self, obstime.shape, subok=True)

        # do this here to prevent a series of complicated circular imports
        from .builtin_frames import ITRS
        return ITRS(x=self.x, y=self.y, z=self.z, obstime=obstime)

    itrs = property(get_itrs, doc="""An `~astropy.coordinates.ITRS` object  with
                                     for the location of this object at the
                                     default ``obstime``.""")

    def get_gcrs(self, obstime):
        """GCRS position with velocity at ``obstime`` as a GCRS coordinate.

        Parameters
        ----------
        obstime : `~astropy.time.Time`
            The ``obstime`` to calculate the GCRS position/velocity at.

        Returns
        --------
        gcrs : `~astropy.coordinates.GCRS` instance
            With velocity included.
        """
        # do this here to prevent a series of complicated circular imports
        from .builtin_frames import GCRS

        itrs = self.get_itrs(obstime)
        # Assume the observatory itself is fixed on the ground.
        # We do a direct assignment rather than an update to avoid validation
        # and creation of a new object.
        zeros = np.broadcast_to(0. * u.km / u.s, (3,) + itrs.shape, subok=True)
        itrs.data.differentials['s'] = CartesianDifferential(zeros)
        return itrs.transform_to(GCRS(obstime=obstime))

    def get_gcrs_posvel(self, obstime):
        """
        Calculate the GCRS position and velocity of this object at the
        requested ``obstime``.

        Parameters
        ----------
        obstime : `~astropy.time.Time`
            The ``obstime`` to calculate the GCRS position/velocity at.

        Returns
        --------
        obsgeoloc : `~astropy.coordinates.CartesianRepresentation`
            The GCRS position of the object
        obsgeovel : `~astropy.coordinates.CartesianRepresentation`
            The GCRS velocity of the object
        """
        # GCRS position
        gcrs_data = self.get_gcrs(obstime).data
        obsgeopos = gcrs_data.without_differentials()
        obsgeovel = gcrs_data.differentials['s'].to_cartesian()
        return obsgeopos, obsgeovel

    def gravitational_redshift(self, obstime,
                               bodies=['sun', 'jupiter', 'moon'],
                               masses={}):
        """Return the gravitational redshift at this EarthLocation.

        Calculates the gravitational redshift, of order 3 m/s, due to the
        requested solar system bodies.

        Parameters
        ----------
        obstime : `~astropy.time.Time`
            The ``obstime`` to calculate the redshift at.

        bodies : iterable, optional
            The bodies (other than the Earth) to include in the redshift
            calculation.  List elements should be any body name
            `get_body_barycentric` accepts.  Defaults to Jupiter, the Sun, and
            the Moon.  Earth is always included (because the class represents
            an *Earth* location).

        masses : dict of str to Quantity, optional
            The mass or gravitational parameters (G * mass) to assume for the
            bodies requested in ``bodies``. Can be used to override the
            defaults for the Sun, Jupiter, the Moon, and the Earth, or to
            pass in masses for other bodies.

        Returns
        --------
        redshift :  `~astropy.units.Quantity`
            Gravitational redshift in velocity units at given obstime.
        """
        # needs to be here to avoid circular imports
        from .solar_system import get_body_barycentric

        bodies = list(bodies)
        # Ensure earth is included and last in the list.
        if 'earth' in bodies:
            bodies.remove('earth')
        bodies.append('earth')
        _masses = {'sun': consts.GM_sun,
                   'jupiter': consts.GM_jup,
                   'moon': consts.G * 7.34767309e22*u.kg,
                   'earth': consts.GM_earth}
        _masses.update(masses)
        GMs = []
        M_GM_equivalency = (u.kg, u.Unit(consts.G * u.kg))
        for body in bodies:
            try:
                GMs.append(_masses[body].to(u.m**3/u.s**2, [M_GM_equivalency]))
            except KeyError:
                raise KeyError(f'body "{body}" does not have a mass!')
            except u.UnitsError as exc:
                exc.args += ('"masses" argument values must be masses or '
                             'gravitational parameters',)
                raise

        positions = [get_body_barycentric(name, obstime) for name in bodies]
        # Calculate distances to objects other than earth.
        distances = [(pos - positions[-1]).norm() for pos in positions[:-1]]
        # Append distance from Earth's center for Earth's contribution.
        distances.append(CartesianRepresentation(self.geocentric).norm())
        # Get redshifts due to all objects.
        redshifts = [-GM / consts.c / distance for (GM, distance) in
                     zip(GMs, distances)]
        # Reverse order of summing, to go from small to big, and to get
        # "earth" first, which gives m/s as unit.
        return sum(redshifts[::-1])

    @property
    def x(self):
        """The X component of the geocentric coordinates."""
        return self['x']

    @property
    def y(self):
        """The Y component of the geocentric coordinates."""
        return self['y']

    @property
    def z(self):
        """The Z component of the geocentric coordinates."""
        return self['z']

    def __getitem__(self, item):
        result = super().__getitem__(item)
        if result.dtype is self.dtype:
            return result.view(self.__class__)
        else:
            return result.view(u.Quantity)

    def __array_finalize__(self, obj):
        super().__array_finalize__(obj)
        if hasattr(obj, '_ellipsoid'):
            self._ellipsoid = obj._ellipsoid

    def __len__(self):
        if self.shape == ():
            raise IndexError('0-d EarthLocation arrays cannot be indexed')
        else:
            return super().__len__()

    def _to_value(self, unit, equivalencies=[]):
        """Helper method for to and to_value."""
        # Conversion to another unit in both ``to`` and ``to_value`` goes
        # via this routine. To make the regular quantity routines work, we
        # temporarily turn the structured array into a regular one.
        array_view = self.view(self._array_dtype, np.ndarray)
        if equivalencies == []:
            equivalencies = self._equivalencies
        new_array = self.unit.to(unit, array_view, equivalencies=equivalencies)
        return new_array.view(self.dtype).reshape(self.shape)


import operator
import tflite_runtime.interpreter as tflite
from pathlib import Path
import platform
import json
from PIL import Image
import sys
import numpy as np
import collections
import time

WEIGHTS_DIR = Path(__file__).parent / "weights"
EDGETPU_SHARED_LIBRARY = {
    'Linux': 'libedgetpu.so.1',
    'Darwin': 'libedgetpu.1.dylib',
    'Windows': 'edgetpu.dll'
}[platform.system()]

Class = collections.namedtuple('Class', ['id', 'score'])

def preprocess_img(img_path, batch=True):
    img = Image.open(img_path).convert('RGB')
    img = img.resize((480, 480))
    img_arr = np.asarray(img)/255
    if batch:
        img_arr = img_arr[None, :, :, :3]
    # img_arr = 2.0 * (img_arr - img_arr.min()) / img_arr.ptp() -1
    return img_arr.astype('float32')

class EdgeTPUClassifier:
    def __init__(self, weights_dir: Path):
        self._weights_dir = weights_dir
        self.labels = self.load_labels()
        self.interpreter = self.make_interpreter()
        self.interpreter.allocate_tensors()

    def get_latest_weights(self) -> Path:
        latest_epoch = None
        for entity in self._weights_dir.iterdir():
            if not entity.is_file() or entity.suffix != '.tflite':
                continue
            filename = entity.stem
            if filename.endswith("_edgetpu") and filename.startswith("8bit_epoch"):
                epoch = int(filename.split("_")[1][5:])
                if latest_epoch is None or epoch > latest_epoch:
                    latest_epoch = epoch
        return self._weights_dir / f"8bit_epoch{latest_epoch}_edgetpu.tflite"

    def load_labels(self):
        if not (self._weights_dir / "labels.json").is_file():
            raise Exception("No labels file found in the weights directory")
        with (self._weights_dir / "labels.json").open() as f:
            return json.load(f)

    def make_interpreter(self) -> tflite.Interpreter:
        return tflite.Interpreter(
            model_path=str(self.get_latest_weights()),
            experimental_delegates=[
                tflite.load_delegate(
                    EDGETPU_SHARED_LIBRARY,
                    {}
                )
            ]
        )

    def input_details(self, key):
        return self.interpreter.get_input_details()[0][key]

    def input_size(self):
        _, height, width, _ = self.input_details('shape')
        return width, height

    def input_tensor(self):
        tensor_index = self.input_details('index')
        return self.interpreter.tensor(tensor_index)()[0]

    def output_tensor(self, dequantize=True):
        output_details = self.interpreter.get_output_details()[0]
        output_data = np.squeeze(self.interpreter.tensor(output_details['index'])())
        if dequantize and np.issubdtype(output_details['dtype'], np.integer):
            scale, zero_point = output_details['quantization']
            return scale * (output_data - zero_point)
        return output_data

    def set_input(self, data):
        self.input_tensor()[:, :] = data

    def get_output(self, top_k=1, score_threshold=0.0):
        scores = self.output_tensor()
        classes = [
            Class(i, scores[i])
            for i in np.argpartition(scores, -top_k)[-top_k:]
            if scores[i] >= score_threshold
        ]
        return sorted(classes, key=operator.itemgetter(1), reverse=True)

    def infer(self, image, top_k=1, score_threshold=0.0):
        self.set_input(image)
        self.interpreter.invoke()
        return self.get_output(top_k, score_threshold)

if __name__ == "__main__":
    classifier = EdgeTPUClassifier(WEIGHTS_DIR)
    image = preprocess_img(Path(sys.argv[1]), False)
    result = classifier.infer(image, 3)

    benchmark_iterations = 50

    start = time.perf_counter()
    for _ in range(benchmark_iterations):
        classifier.infer(image)
    end = time.perf_counter()
    duration = end - start
    print(f"Took {duration * 1000:.5f} ms to run inference {benchmark_iterations} times")
    print(f"Averaged {duration * 1000 / benchmark_iterations:.5f} ms per inference")
    print(f"Or {benchmark_iterations / duration:.5f} inferences per second")

    print("----- RESULTS -----")
    for klass in result:
        print(f"{classifier.labels[klass.id]}: {klass.score:.5f}")


from __future__ import (absolute_import, print_function, division)

from astropy.table import Table
import numpy as np
from astropy.io import fits

# fix needed till astropy.io get's updated - need to check this
fits.column.ASCII2NUMPY['E'] = fits.column.ASCII2NUMPY['F'] = 'f8'

import matplotlib.pyplot as plt
import string as st
import subprocess
from scipy import interpolate
from scipy.integrate import simps
from pydirtygrid.PhotDG import PhotDG
import h5py

__all__ = ['SpecDG']

class SpecDG:
    def __init__(self):
        """
        Read in the mapping file for all spectra

        Returns
        -------
        (self.)plan:   Table
            table containing the identification of each spectrum,
            as well as its parameters

            (self.)seds:   list
                empty list to save the spectra if you want to
        """
        mapfile = '/astro/dust_kg3/klaw/cloudy2/dirtygrid_db/django/param_table6_combined.tsvx'
        self.plan = Table.read(mapfile, format='ascii')
        self.seds = []
        #self.waves = []

    def findFile(self,file_id):
       """
       Returns the full filename for a given GID from the spectrum mapping.

       Parameter
       ---------
       file_id:   string
          identification number of the spectrum to be read;
          'gid' column of the self.plan Table

       Returns
       -------
       filepath+filename:   string
          the absolute path and name of the .fits file
       """

       alphab = []
       for i in range(26): alphab.append(st.uppercase[i])
       alphab = np.array(alphab)
       ind_id = []
       for i in range(len(file_id)):
           tmp_i = np.where(file_id[i] == alphab)
           if len(tmp_i[0] == 1):
               ind_id.append(tmp_i[0][0])
           else:
               break
       prefix = ''.join(alphab[ind_id])
       filepath = '/astro/dust_kg3/klaw/cloudy2/nasa_fits/'+prefix \
           +'/'+str(file_id[-4:-2])+'/'
       files = []
       proc = subprocess.Popen(['ls', filepath], stdout=subprocess.PIPE)
       for line in proc.stdout.readlines(): files.append(line.rstrip())
       filename = [s for s in files if file_id in s]
       if len(filename) == 0:
           return 'Void'
       else:
           filename = filename[0]
           return filepath+filename

    def findGidFromParam(self, grain, geom, sf_type, metal, age, sfr, tau):
        """
        Returns the GID given a set of parameters
        (If you want to plot or something)

        Parameters
        ----------
        grain:   string
        type of grain

        geom:   string
        geometry

        sf_type: string
        star formation type

        metal:   float
        metallicity

        age:   float
        age of the stellar population

        sfr:   float
        star formation rate

        tau:    float
        optical depth

        Returns
        -------
        file_gid:   string
        identification number of the spectrum with the given set
        of parameters
        """
        # ToDo: A better way to do this?
        ind = np.where( (self.plan['grain'] == grain)
                        & (self.plan['geom'] == geom)
                        & (self.plan['sf_type'] == sf_type)
                        & (self.plan['metal'] == metal)
                        & (self.plan['age'] == age)
                        & (self.plan['sfr'] == sfr)
                        & (self.plan['tau'] == tau) )
        # Not all spectra available
        #  -- Exit if non existant - Until interpolation?
        if len(ind[0]) == 0:
            return
        # If valid, then go on
        file_gid = self.plan['gid'][ind[0][0]]
        return file_gid

    def findParamsFromGid(self, thisgid):
        """
        Returns the parameter values for a given GID
        (Used later to update the cube)

        Parameters
        ----------
        thisgid:   string
        identification number of the spectrum

        Returns
        -------
        thisgt, thisgm, thissf, thismt, thissa, thissr, thista:      integers
        indices of the parameter values for the given spectrum,
        in the multidimensional photometry cube
        """
        ind_gid = np.where(self.plan['gid'] == thisgid)
        thisgt = self.plan['grain'][ind_gid[0][0]]      # Grain type
        thisgm = self.plan['geom'][ind_gid[0][0]]      # Geometry
        thissf = self.plan['sf_type'][ind_gid[0][0]]   # Star formation type
        thismt = self.plan['metal'][ind_gid[0][0]]      # Metal
        thissa = self.plan['age'][ind_gid[0][0]]      # Stellar age
        thissr = self.plan['sfr'][ind_gid[0][0]]      # Star formation rate
        thista = self.plan['tau'][ind_gid[0][0]]      # Optical depth
        return thisgt, thisgm, thissf, thismt, thissa, thissr, thista

    def specGet(self, filename):
        """
        Save a spectrum from the filename

        Parameters
        ----------
        filename:   string
        absolute path to the spectrum you want to read

        Returns
        -------
        (self.)seds:   list
        update the list to add the SED
        (self.)waves:   array
        the wavelengths to the SED
        """
        # Read the fits file
        hdulist = fits.open(filename)
        scidata = hdulist[1].data
        self.seds.append(scidata['Flux'])
        self.waves = scidata['Wavelength']

    def specPlot(self, ind=-1):
        """
        Plot SEDs, either giving a specific index of which if more that
        one saved, or all of them

        Parameters
        ----------
        ind:   integer(s)   (Optional)
        indices of the SED you wish to plot
        """
        plt.xscale('log')
        plt.yscale('log')
        plt.ylabel('Luminosity  [$erg$ $s^{-1}$ $\mu m^{-1}$]', fontsize=15)
        plt.xlabel('Wavelength  [$\mu m$]', fontsize=15)
        if ind == -1:
            print('Plotting all SEDs in object')
            for i in range(len(self.seds)):
                plt.plot(self.waves, self.seds[i])
        else:
            for i in ind: plt.plot(self.waves, np.array(self.seds[ind]))
        # ToDo: need a legend
        # ToDo: specify an figure/subplot if desired?
        plt.show()

    def spec2Phot(self, trans_curve, trans_waves, wave0, energy=1):
        """
        Compute the new photometry

        Parameters
        ----------
        trans_curve:   array
        transmission curve of the new filter /!\ Alread read, as an array

        trans_waves:   array
        wavelengths of the new filter

        wave0:   float
        center/effective wavelength of the new filter

        energy:   integer   (Optional)
        specification for integration with energy(=1, default) or photons

        Returns
        -------
        newcube:   array(3, 6, 2, 5, 50, 29, 25) (float)
        new photometry cube
        """
        # Create an instance of the PhotDG class to call functions
        phtemp = PhotDG(silent=1)
        light_speed = 2.998e14      # microns/s
        n_files = len(self.plan['gid'])
        # Create the array to fill
        newcube = np.empty((3,6,2,5,50,29,25))
        # Read in each spectrum
        for i in range(n_files):
            if (i % 1000) == 0:
                print(i)
            filename = self.findFile(self.plan['gid'][i])
            if filename == 'Void':
                continue
            hdu = fits.open(filename)
            scidata = hdu[1].data
            spec = scidata['Flux']                  # SED
            # Wavelength (same for all files)
            if i == 0:
                dgwave = scidata['Wavelength']
            # Interpolate the spectral response on DG wavelength grid
            # Do it in log space
            func = interpolate.interp1d(trans_waves, np.log10(trans_curve),
                                        kind='slinear')
            tr_ind = np.where( (dgwave > np.min(trans_waves))
                               & (dgwave < np.max(trans_waves)))
            newtran = func(dgwave[tr_ind])
            newtran = 10**newtran
            # Frequencies
            nu0 = light_speed/wave0#[indf]
            nu = light_speed/dgwave[tr_ind]
            # Total response curve, for integration      (Used simps)
            if energy == 1:
                tot_resp_curve = simps(newtran*nu0/nu, dgwave[tr_ind])
                tmp_phot = simps(spec[tr_ind]*newtran, dgwave[tr_ind])
                newphot = tmp_phot / tot_resp_curve
            else:      # In photons. Do the Blackbody too?
                tot_resp_curve = simps(newtran*nu0/nu**2, dgwave[tr_ind])
                tmp_phot = simps(spec[tr_ind]*newtran/nu, dgwave[tr_ind])
                newphot = tmp_phot / tot_resp_curve
            # Fill the array
            curgt, curgm, cursf, curmt, cursa, cursr, curta = \
                self.findParamsFromGid(self.plan['gid'][i])
            indgt, indgm, indsf, indmt, indsa, indsr, indta = \
                phtemp.findIndexFromParams(curgt, curgm, cursf, curmt,
                                           cursa, cursr, curta)
            newcube[indgt, indgm, indsf, indmt, indsa, indsr, indta] = newphot

        return newcube
        #ToDo: maybe save sometimes to avoid all recompute if fails somewhere?

##########


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

# Copyright 2020 Patrick Lumban Tobing (Nagoya University)
#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)

from __future__ import print_function

import argparse
import multiprocessing as mp
import logging
import os

import numpy as np
from sklearn.preprocessing import StandardScaler

from utils import check_hdf5
from utils import read_hdf5
from utils import read_txt
from utils import write_hdf5

from multiprocessing import Array

def main():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--feats", default=None, required=True,
        help="name of the list of hdf5 files")
    parser.add_argument(
        "--stats", default=None, required=True,
        help="filename of stats for hdf5 format")
    parser.add_argument("--expdir", required=True,
        type=str, help="directory to save the log")
    parser.add_argument("--mcep_dim", default=50,
        type=int, help="dimension of mel-cepstrum")
    parser.add_argument(
        "--n_jobs", default=10,
        type=int, help="number of parallel jobs")
    parser.add_argument(
        "--verbose", default=1,
        type=int, help="log message level")

    args = parser.parse_args()

    # set log level
    if args.verbose == 1:
        logging.basicConfig(level=logging.INFO,
                            format='%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s',
                            datefmt='%m/%d/%Y %I:%M:%S',
                            filename=args.expdir + "/calc_stats.log")
        logging.getLogger().addHandler(logging.StreamHandler())
    elif args.verbose > 1:
        logging.basicConfig(level=logging.DEBUG,
                            format='%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s',
                            datefmt='%m/%d/%Y %I:%M:%S',
                            filename=args.expdir + "/calc_stats.log")
        logging.getLogger().addHandler(logging.StreamHandler())
    else:
        logging.basicConfig(level=logging.WARN,
                            format='%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s',
                            datefmt='%m/%d/%Y %I:%M:%S',
                            filename=args.expdir + "/calc_stats.log")
        logging.getLogger().addHandler(logging.StreamHandler())
        logging.warn("logging is disabled.")

    # read list and define scaler
    filenames = read_txt(args.feats)
    logging.info("number of training utterances = "+str(len(filenames)))

    def calc_stats(filenames, cpu, feat_mceplf0cap_list, feat_orglf0_list, varmcep_list, f0_list, melsp_list, varmelsp_list):
        feat_mceplf0cap_arr = None
        feat_orglf0_arr = None
        varmcep_arr = None
        f0_arr = None
        melsp_arr = None
        varmelsp_arr = None
        count = 0
        # process over all of data
        for filename in filenames:
            logging.info(filename)
            feat_mceplf0cap = read_hdf5(filename, "/feat_mceplf0cap")
            logging.info(feat_mceplf0cap.shape)
            feat_orglf0 = read_hdf5(filename, "/feat_org_lf0")
            logging.info(feat_orglf0.shape)
            melsp = read_hdf5(filename, "/log_1pmelmagsp")
            logging.info(melsp.shape)
            if feat_mceplf0cap_arr is not None:
                feat_mceplf0cap_arr = np.r_[feat_mceplf0cap_arr, feat_mceplf0cap]
            else:
                feat_mceplf0cap_arr = feat_mceplf0cap
            if feat_orglf0_arr is not None:
                feat_orglf0_arr = np.r_[feat_orglf0_arr, feat_orglf0]
            else:
                feat_orglf0_arr = feat_orglf0
            logging.info('feat')
            logging.info(feat_mceplf0cap_arr.shape)
            logging.info(feat_orglf0_arr.shape)
            if varmcep_arr is not None:
                varmcep_arr = np.r_[varmcep_arr, np.var(feat_mceplf0cap[:,-args.mcep_dim:], axis=0, keepdims=True)]
            else:
                varmcep_arr = np.var(feat_mceplf0cap[:,-args.mcep_dim:], axis=0, keepdims=True)
            logging.info('var')
            logging.info(varmcep_arr.shape)
            logging.info('f0')
            f0 = read_hdf5(filename, "/f0_range")
            logging.info(f0.shape)
            logging.info('f0 > 0')
            f0 = f0[np.nonzero(f0)]
            logging.info(f0.shape)
            if f0_arr is not None:
                f0_arr = np.r_[f0_arr, f0]
            else:
                f0_arr = f0
            logging.info(f0_arr.shape)
            if melsp_arr is not None:
                melsp_arr = np.r_[melsp_arr, melsp]
            else:
                melsp_arr = melsp
            logging.info(melsp_arr.shape)
            if varmelsp_arr is not None:
                varmelsp_arr = np.r_[varmelsp_arr, np.var((np.exp(melsp)-1)/10000, axis=0, keepdims=True)]
            else:
                varmelsp_arr = np.var((np.exp(melsp)-1)/10000, axis=0, keepdims=True)
            logging.info('var')
            logging.info(varmelsp_arr.shape)
            count += 1
            logging.info("cpu %d %d %d %d %d %d %d %d" % (cpu, count, len(feat_mceplf0cap_arr), len(feat_orglf0_arr), len(varmcep_arr), len(f0_arr), len(melsp_arr), len(varmelsp_arr)))
            #if count >= 5:
            #    break

        feat_mceplf0cap_list.append(feat_mceplf0cap_arr)
        feat_orglf0_list.append(feat_orglf0_arr)
        varmcep_list.append(varmcep_arr)
        f0_list.append(f0_arr)
        melsp_list.append(melsp_arr)
        varmelsp_list.append(varmelsp_arr)

    # divie list
    feat_lists = np.array_split(filenames, args.n_jobs)
    feat_lists = [f_list.tolist() for f_list in feat_lists]

    for i in range(len(feat_lists)):
        logging.info("%d %d" % (i+1, len(feat_lists[i])))

    # multi processing
    with mp.Manager() as manager:
        processes = []
        feat_mceplf0cap_list = manager.list()
        feat_orglf0_list = manager.list()
        varmcep_list = manager.list()
        f0_list = manager.list()
        melsp_list = manager.list()
        varmelsp_list = manager.list()
        for i, feat_list in enumerate(feat_lists):
            p = mp.Process(target=calc_stats, args=(feat_list, i+1, feat_mceplf0cap_list, feat_orglf0_list, varmcep_list, f0_list, melsp_list, varmelsp_list,))
            p.start()
            processes.append(p)

        # wait for all process
        for p in processes:
            p.join()

        feat_mceplf0cap = None
        for i in range(len(feat_mceplf0cap_list)):
            if feat_mceplf0cap_list[i] is not None:
                logging.info(i)
                logging.info(feat_mceplf0cap_list[i].shape)
                if feat_mceplf0cap is not None:
                    feat_mceplf0cap = np.r_[feat_mceplf0cap, feat_mceplf0cap_list[i]]
                else:
                    feat_mceplf0cap = feat_mceplf0cap_list[i]
        logging.info('feat mceplf0cap: %d' % (len(feat_mceplf0cap)))
        logging.info(feat_mceplf0cap.shape)

        feat_orglf0 = None
        for i in range(len(feat_orglf0_list)):
            if feat_orglf0_list[i] is not None:
                logging.info(i)
                logging.info(feat_orglf0_list[i].shape)
                if feat_orglf0 is not None:
                    feat_orglf0 = np.r_[feat_orglf0, feat_orglf0_list[i]]
                else:
                    feat_orglf0 = feat_orglf0_list[i]
        logging.info('feat orglf0: %d' % (len(feat_orglf0)))
        logging.info(feat_orglf0.shape)

        var_range = None
        for i in range(len(varmcep_list)):
            if varmcep_list[i] is not None:
                logging.info(i)
                logging.info(varmcep_list[i].shape)
                if var_range is not None:
                    var_range = np.r_[var_range, varmcep_list[i]]
                else:
                    var_range = varmcep_list[i]
        logging.info('var mcep: %d' % (len(var_range)))
        logging.info(var_range.shape)

        f0s_range = None
        for i in range(len(f0_list)):
            if f0_list[i] is not None:
                logging.info(i)
                logging.info(f0_list[i].shape)
                if f0s_range is not None:
                    f0s_range = np.r_[f0s_range, f0_list[i]]
                else:
                    f0s_range = f0_list[i]
        logging.info('f0: %d' % (len(f0s_range)))
        logging.info(f0s_range.shape)

        melsp = None
        for i in range(len(melsp_list)):
            if melsp_list[i] is not None:
                logging.info(i)
                logging.info(melsp_list[i].shape)
                if melsp is not None:
                    melsp = np.r_[melsp, melsp_list[i]]
                else:
                    melsp = melsp_list[i]
        logging.info('melsp: %d' % (len(melsp)))
        logging.info(melsp.shape)

        var_melsp = None
        for i in range(len(varmelsp_list)):
            if varmelsp_list[i] is not None:
                logging.info(i)
                logging.info(varmelsp_list[i].shape)
                if var_melsp is not None:
                    var_melsp = np.r_[var_melsp, varmelsp_list[i]]
                else:
                    var_melsp = varmelsp_list[i]
        logging.info('var melsp: %d' % (len(var_melsp)))
        logging.info(var_melsp.shape)

        scaler_feat_mceplf0cap = StandardScaler()
        scaler_feat_orglf0 = StandardScaler()

        scaler_feat_mceplf0cap.partial_fit(feat_mceplf0cap)
        scaler_feat_orglf0.partial_fit(feat_orglf0)

        scaler_melsp = StandardScaler()
        scaler_melsp.partial_fit(melsp)

        mean_feat_mceplf0cap = scaler_feat_mceplf0cap.mean_
        scale_feat_mceplf0cap = scaler_feat_mceplf0cap.scale_
        mean_feat_orglf0 = scaler_feat_orglf0.mean_
        scale_feat_orglf0 = scaler_feat_orglf0.scale_
        gv_range_mean = np.mean(np.array(var_range), axis=0)
        gv_range_var = np.var(np.array(var_range), axis=0)
        logging.info(gv_range_mean)
        logging.info(gv_range_var)
        f0_range_mean = np.mean(f0s_range)
        f0_range_std = np.std(f0s_range)
        logging.info(f0_range_mean)
        logging.info(f0_range_std)
        lf0_range_mean = np.mean(np.log(f0s_range))
        lf0_range_std = np.std(np.log(f0s_range))
        logging.info(lf0_range_mean)
        logging.info(lf0_range_std)

        logging.info(mean_feat_mceplf0cap)
        logging.info(scale_feat_mceplf0cap)
        write_hdf5(args.stats, "/mean_feat_mceplf0cap", mean_feat_mceplf0cap)
        write_hdf5(args.stats, "/scale_feat_mceplf0cap", scale_feat_mceplf0cap)
        logging.info(mean_feat_orglf0)
        logging.info(scale_feat_orglf0)
        write_hdf5(args.stats, "/mean_feat_org_lf0", mean_feat_orglf0)
        write_hdf5(args.stats, "/scale_feat_org_lf0", scale_feat_orglf0)
        write_hdf5(args.stats, "/gv_range_mean", gv_range_mean)
        write_hdf5(args.stats, "/gv_range_var", gv_range_var)
        write_hdf5(args.stats, "/f0_range_mean", f0_range_mean)
        write_hdf5(args.stats, "/f0_range_std", f0_range_std)
        write_hdf5(args.stats, "/lf0_range_mean", lf0_range_mean)
        write_hdf5(args.stats, "/lf0_range_std", lf0_range_std)

        mean_melsp = scaler_melsp.mean_
        scale_melsp = scaler_melsp.scale_
        gv_melsp_mean = np.mean(np.array(var_melsp), axis=0)
        gv_melsp_var = np.var(np.array(var_melsp), axis=0)
        logging.info(gv_melsp_mean)
        logging.info(gv_melsp_var)
        logging.info(mean_melsp)
        logging.info(scale_melsp)
        write_hdf5(args.stats, "/mean_melsp", mean_melsp)
        write_hdf5(args.stats, "/scale_melsp", scale_melsp)
        write_hdf5(args.stats, "/gv_melsp_mean", gv_melsp_mean)
        write_hdf5(args.stats, "/gv_melsp_var", gv_melsp_var)

if __name__ == "__main__":
    main()


#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
Run a YOLO_v3 style detection model on test images.
"""

import colorsys
import os
import glob
from timeit import default_timer as timer

import numpy as np
from keras import backend as K
from keras.models import load_model
from keras.layers import Input
from PIL import Image, ImageFont, ImageDraw

from yolo3.model import yolo_eval, yolo_body, tiny_yolo_body
from yolo3.utils import letterbox_image

'''
before run, run prepare_test_set.py
Only change model_path, anchors_path and score!!!!!!!!!!!!!!!!!!!!!!!!!!
'''

class YOLO(object):
    def __init__(self):
        self.model_path = 'E:/Teade checkpoints/000/finalmodel_stage_20.h5' # only thing to change
        self.anchors_path = 'model_data/my_anchors4.txt' # only thing to change
        self.classes_path = 'model_data/teade_classes.txt'
        self.score = 0.2
        self.iou = 0.5
        self.class_names = self._get_class()
        self.anchors = self._get_anchors()
        self.sess = K.get_session()
        self.model_image_size = (416, 416) # fixed size or (None, None), hw
        self.boxes, self.scores, self.classes = self.generate()

    def _get_class(self):
        classes_path = os.path.expanduser(self.classes_path)
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    def _get_anchors(self):
        anchors_path = os.path.expanduser(self.anchors_path)
        with open(anchors_path) as f:
            anchors = f.readline()
        anchors = [float(x) for x in anchors.split(',')]
        return np.array(anchors).reshape(-1, 2)

    def generate(self):
        model_path = os.path.expanduser(self.model_path)
        assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'

        # Load model, or construct model and load weights.
        num_anchors = len(self.anchors)
        num_classes = len(self.class_names)
        is_tiny_version = num_anchors==6 # default setting
        try:
            self.yolo_model = load_model(model_path, compile=False)
        except:
            self.yolo_model = tiny_yolo_body(Input(shape=(None,None,3)), num_anchors//2, num_classes) \
                if is_tiny_version else yolo_body(Input(shape=(None,None,3)), num_anchors//3, num_classes)
            self.yolo_model.load_weights(self.model_path) # make sure model, anchors and classes match
        else:
            assert self.yolo_model.layers[-1].output_shape[-1] == \
                num_anchors/len(self.yolo_model.output) * (num_classes + 5), \
                'Mismatch between model and given anchor and class sizes'

        print('{} model, anchors, and classes loaded.'.format(model_path))

        # Generate output tensor targets for filtered bounding boxes.
        self.input_image_shape = K.placeholder(shape=(2, ))
        boxes, scores, classes = yolo_eval(self.yolo_model.output, self.anchors,
                len(self.class_names), self.input_image_shape,
                score_threshold=self.score, iou_threshold=self.iou)
        return boxes, scores, classes

    def detect_image(self, image,filename):
        if self.model_image_size != (None, None):
            assert self.model_image_size[0]%32 == 0, 'Multiples of 32 required'
            assert self.model_image_size[1]%32 == 0, 'Multiples of 32 required'
            boxed_image = letterbox_image(image, tuple(reversed(self.model_image_size)))
        else:
            new_image_size = (image.width - (image.width % 32),
                              image.height - (image.height % 32))
            boxed_image = letterbox_image(image, new_image_size)
        image_data = np.array(boxed_image, dtype='float32')

        #print(image_data.shape)
        image_data /= 255.
        image_data = np.expand_dims(image_data, 0)  # Add batch dimension.

        out_boxes, out_scores, out_classes = self.sess.run(
            [self.boxes, self.scores, self.classes],
            feed_dict={
                self.yolo_model.input: image_data,
                self.input_image_shape: [image.size[1], image.size[0]],
                K.learning_phase(): 0
            })

        #print('Found {} boxes for {}'.format(len(out_boxes), 'img'))

        pred_prefix = 'C:/Users/CTK_CAD/PycharmProjects/mAP/predicted/'
        text_file = open(pred_prefix + filename.split('.')[0] + '.txt', 'a')
        text_file.close()

        for i, c in reversed(list(enumerate(out_classes))):
            predicted_class = self.class_names[c]
            box = out_boxes[i]
            score = out_scores[i]

            label = '{} {:.2f}'.format(predicted_class, score)

            top, left, bottom, right = box
            top = max(0, np.floor(top + 0.5).astype('int32'))
            left = max(0, np.floor(left + 0.5).astype('int32'))
            bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
            right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
            #print(label, (left, top), (right, bottom))

            text_file = open(pred_prefix + filename.split('.')[0]+'.txt','a')
            text_file.write(label + ' ' +str(left)+' '+str(top)+' '+str(right)+' '+str(bottom) + '\n')
            text_file.close()

    def close_session(self):
        self.sess.close()

def detect_img(yolo):
    prefix = 'C:\\Users\\CTK_CAD\\Chalmers Teknologkonsulter AB\\Bird Classification - Images\\Bird Detection - Images\\Original images\\'
    anno_file = 'test_anno.txt'

    pred_dir = 'C:/Users/CTK_CAD/PycharmProjects/mAP/predicted'
    for subdir, dirs, files in os.walk(pred_dir):
        for f in files:
            os.remove(os.path.join(subdir,f))

    with open(prefix+anno_file) as f:
        lines = f.readlines()

    for line in lines:
        filename = line.split()[0]
        if filename.endswith(".jpg") or filename.endswith(".png"):
            image = Image.open(prefix + filename)
            yolo.detect_image(image, filename)
    yolo.close_session()

if __name__ == '__main__':
    detect_img(YOLO())