<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''Feature inversion'''

import numpy as np
import scipy.fftpack

from ..util.exceptions import ParameterError
from ..core.spectrum import griffinlim
from ..core.spectrum import db_to_power
from .. import filters
from ..util import nnls

__all__ = ['mel_to_stft', 'mel_to_audio',
           'mfcc_to_mel', 'mfcc_to_audio']

def mel_to_stft(M, sr=22050, n_fft=2048, power=2.0, **kwargs):
    '''Approximate STFT magnitude from a Mel power spectrogram.

    Parameters
    ----------
    M : np.ndarray [shape=(n_mels, n), non-negative]
        The spectrogram as produced by `feature.melspectrogram`

    sr : number > 0 [scalar]
        sampling rate of the underlying signal

    n_fft : int > 0 [scalar]
        number of FFT components in the resulting STFT

    power : float > 0 [scalar]
        Exponent for the magnitude melspectrogram

    kwargs : additional keyword arguments
        Mel filter bank parameters.
        See `librosa.filters.mel` for details

    Returns
    -------
    S : np.ndarray [shape=(n_fft, t), non-negative]
        An approximate linear magnitude spectrogram

    See Also
    --------
    feature.melspectrogram
    core.stft
    filters.mel
    util.nnls

    Examples
    --------
    >>> y, sr = librosa.load(librosa.util.example_audio_file(), duration=5, offset=10)
    >>> S = np.abs(librosa.stft(y))
    >>> mel_spec = librosa.feature.melspectrogram(S=S, sr=sr)
    >>> S_inv = librosa.feature.inverse.mel_to_stft(mel_spec, sr=sr)

    Compare the results visually

    >>> import matplotlib.pyplot as plt
    >>> plt.figure()
    >>> plt.subplot(2,1,1)
    >>> librosa.display.specshow(librosa.amplitude_to_db(S, ref=np.max, top_db=None),
    ...                          y_axis='log', x_axis='time')
    >>> plt.colorbar()
    >>> plt.title('Original STFT')
    >>> plt.subplot(2,1,2)
    >>> librosa.display.specshow(librosa.amplitude_to_db(np.abs(S_inv - S),
    ...                                                  ref=S.max(), top_db=None),
    ...                          vmax=0, y_axis='log', x_axis='time', cmap='magma')
    >>> plt.title('Residual error (dB)')
    >>> plt.colorbar()
    >>> plt.tight_layout()
    >>> plt.show()
    '''

    # Construct a mel basis with dtype matching the input data
    mel_basis = filters.mel(sr, n_fft, n_mels=M.shape[0],
                            dtype=M.dtype,
                            **kwargs)

    # Find the non-negative least squares solution, and apply
    # the inverse exponent.
    # We'll do the exponentiation in-place.
    inverse = nnls(mel_basis, M)
    return np.power(inverse, 1. / power, out=inverse)

def mel_to_audio(M, sr=22050, n_fft=2048, hop_length=512, win_length=None,
                 window='hann', center=True, pad_mode='reflect', power=2.0, n_iter=32,
                 length=None, dtype=np.float32, **kwargs):
    """Invert a mel power spectrogram to audio using Griffin-Lim.

    This is primarily a convenience wrapper for:

        >>> S = librosa.feature.inverse.mel_to_stft(M)
        >>> y = librosa.griffinlim(S)

    Parameters
    ----------
    M : np.ndarray [shape=(n_mels, n), non-negative]
        The spectrogram as produced by `feature.melspectrogram`

    sr : number > 0 [scalar]
        sampling rate of the underlying signal

    n_fft : int > 0 [scalar]
        number of FFT components in the resulting STFT

    hop_length : None or int > 0
        The hop length of the STFT.  If not provided, it will default to `n_fft // 4`

    win_length : None or int > 0
        The window length of the STFT.  By default, it will equal `n_fft`

    window : string, tuple, number, function, or np.ndarray [shape=(n_fft,)]
        A window specification as supported by `stft` or `istft`

    center : boolean
        If `True`, the STFT is assumed to use centered frames.
        If `False`, the STFT is assumed to use left-aligned frames.

    pad_mode : string
        If `center=True`, the padding mode to use at the edges of the signal.
        By default, STFT uses reflection padding.

    power : float > 0 [scalar]
        Exponent for the magnitude melspectrogram

    n_iter : int > 0
        The number of iterations for Griffin-Lim

    length : None or int > 0
        If provided, the output `y` is zero-padded or clipped to exactly `length`
        samples.

    dtype : np.dtype
        Real numeric type for the time-domain signal.  Default is 32-bit float.

    kwargs : additional keyword arguments
        Mel filter bank parameters

    Returns
    -------
    y : np.ndarray [shape(n,)]
        time-domain signal reconstructed from `M`

    See Also
    --------
    core.griffinlim
    feature.melspectrogram
    filters.mel
    feature.inverse.mel_to_stft
    """

    stft = mel_to_stft(M, sr=sr, n_fft=n_fft, power=power, **kwargs)

    return griffinlim(stft, n_iter=n_iter, hop_length=hop_length, win_length=win_length,
                      window=window, center=center, dtype=dtype, length=length,
                      pad_mode=pad_mode)

def mfcc_to_mel(mfcc, n_mels=128, dct_type=2, norm='ortho', ref=1.0, lifter=0):
    '''Invert Mel-frequency cepstral coefficients to approximate a Mel power
    spectrogram.

    This inversion proceeds in two steps:

        1. The inverse DCT is applied to the MFCCs
        2. `core.db_to_power` is applied to map the dB-scaled result to a power spectrogram

    Parameters
    ----------
    mfcc : np.ndarray [shape=(n_mfcc, n)]
        The Mel-frequency cepstral coefficients

    n_mels : int > 0
        The number of Mel frequencies

    dct_type : None or {1, 2, 3}
        Discrete cosine transform (DCT) type
        By default, DCT type-2 is used.

    norm : None or 'ortho'
        If `dct_type` is `2 or 3`, setting `norm='ortho'` uses an orthonormal
        DCT basis.

        Normalization is not supported for `dct_type=1`.

    ref : number or callable
        Reference power for (inverse) decibel calculation

    lifter : number >= 0
        If `lifter>0`, apply inverse liftering (inverse cepstral filtering):
        `M[n, :] <- M[n, :] / (1 + sin(pi * (n + 1) / lifter)) * lifter / 2`

    Returns
    -------
    M : np.ndarray [shape=(n_mels, n)]
        An approximate Mel power spectrum recovered from `mfcc`

    Warns
    --------
    RuntimeWarning
        overflow encountered in power or
        divide by zero encountered during inverse liftering.

    See Also
    --------
    mfcc
    melspectrogram
    scipy.fftpack.dct
    '''
    if lifter > 0:
        n_mfcc = mfcc.shape[0]
        idx = np.arange(1, 1 + n_mfcc, dtype=mfcc.dtype)
        lifter_sine = 1 + (lifter / 2) * np.sin(np.pi * idx / lifter)[:, np.newaxis]
        lifter_sine = np.where(lifter_sine == 0, np.finfo(float).eps, lifter_sine)
        mfcc = np.where(mfcc == 0, np.finfo(float).eps, mfcc)
        liftered_mfcc = mfcc / lifter_sine
        mfcc = liftered_mfcc

    elif lifter != 0:
        raise ParameterError('MFCC to mel lifter={} must be a non-negative number'.format(lifter))

    logmel = scipy.fftpack.idct(mfcc, axis=0, type=dct_type, norm=norm, n=n_mels)
    return db_to_power(logmel, ref=ref)

def mfcc_to_audio(mfcc, n_mels=128, dct_type=2, norm='ortho', ref=1.0, lifter=0, **kwargs):
    '''Convert Mel-frequency cepstral coefficients to a time-domain audio signal

    This function is primarily a convenience wrapper for the following steps:

        1. Convert mfcc to Mel power spectrum (`mfcc_to_mel`)
        2. Convert Mel power spectrum to time-domain audio (`mel_to_audio`)

    Parameters
    ----------
    mfcc : np.ndarray [shape=(n_mfcc, n)]
        The Mel-frequency cepstral coefficients

    n_mels : int > 0
        The number of Mel frequencies

    dct_type : None or {1, 2, 3}
        Discrete cosine transform (DCT) type
        By default, DCT type-2 is used.

    norm : None or 'ortho'
        If `dct_type` is `2 or 3`, setting `norm='ortho'` uses an orthonormal
        DCT basis.

        Normalization is not supported for `dct_type=1`.

    ref : number or callable
        Reference power for (inverse) decibel calculation

    lifter : number >= 0
        If `lifter>0`, apply inverse liftering (inverse cepstral filtering):
        `M[n, :] <- M[n, :] / (1 + sin(pi * (n + 1) / lifter)) * lifter / 2`

    kwargs : additional keyword arguments
        Parameters to pass through to `mel_to_audio`

    Returns
    -------
    y : np.ndarray [shape=(n)]
        A time-domain signal reconstructed from `mfcc`

    See Also
    --------
    mfcc_to_mel
    mel_to_audio
    feature.mfcc
    core.griffinlim
    scipy.fftpack.dct
    '''
    mel_spec = mfcc_to_mel(mfcc, n_mels=n_mels, dct_type=dct_type, norm=norm,
                           ref=ref, lifter=lifter)

    return mel_to_audio(mel_spec, **kwargs)
<commit_msg>adjust code according to bmcfee review (from 24 Nov 2019)<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''Feature inversion'''

import numpy as np
import scipy.fftpack

from ..util.exceptions import ParameterError
from ..core.spectrum import griffinlim
from ..core.spectrum import db_to_power
from .. import filters
from ..util import nnls

__all__ = ['mel_to_stft', 'mel_to_audio',
           'mfcc_to_mel', 'mfcc_to_audio']

def mel_to_stft(M, sr=22050, n_fft=2048, power=2.0, **kwargs):
    '''Approximate STFT magnitude from a Mel power spectrogram.

    Parameters
    ----------
    M : np.ndarray [shape=(n_mels, n), non-negative]
        The spectrogram as produced by `feature.melspectrogram`

    sr : number > 0 [scalar]
        sampling rate of the underlying signal

    n_fft : int > 0 [scalar]
        number of FFT components in the resulting STFT

    power : float > 0 [scalar]
        Exponent for the magnitude melspectrogram

    kwargs : additional keyword arguments
        Mel filter bank parameters.
        See `librosa.filters.mel` for details

    Returns
    -------
    S : np.ndarray [shape=(n_fft, t), non-negative]
        An approximate linear magnitude spectrogram

    See Also
    --------
    feature.melspectrogram
    core.stft
    filters.mel
    util.nnls

    Examples
    --------
    >>> y, sr = librosa.load(librosa.util.example_audio_file(), duration=5, offset=10)
    >>> S = np.abs(librosa.stft(y))
    >>> mel_spec = librosa.feature.melspectrogram(S=S, sr=sr)
    >>> S_inv = librosa.feature.inverse.mel_to_stft(mel_spec, sr=sr)

    Compare the results visually

    >>> import matplotlib.pyplot as plt
    >>> plt.figure()
    >>> plt.subplot(2,1,1)
    >>> librosa.display.specshow(librosa.amplitude_to_db(S, ref=np.max, top_db=None),
    ...                          y_axis='log', x_axis='time')
    >>> plt.colorbar()
    >>> plt.title('Original STFT')
    >>> plt.subplot(2,1,2)
    >>> librosa.display.specshow(librosa.amplitude_to_db(np.abs(S_inv - S),
    ...                                                  ref=S.max(), top_db=None),
    ...                          vmax=0, y_axis='log', x_axis='time', cmap='magma')
    >>> plt.title('Residual error (dB)')
    >>> plt.colorbar()
    >>> plt.tight_layout()
    >>> plt.show()
    '''

    # Construct a mel basis with dtype matching the input data
    mel_basis = filters.mel(sr, n_fft, n_mels=M.shape[0],
                            dtype=M.dtype,
                            **kwargs)

    # Find the non-negative least squares solution, and apply
    # the inverse exponent.
    # We'll do the exponentiation in-place.
    inverse = nnls(mel_basis, M)
    return np.power(inverse, 1. / power, out=inverse)

def mel_to_audio(M, sr=22050, n_fft=2048, hop_length=512, win_length=None,
                 window='hann', center=True, pad_mode='reflect', power=2.0, n_iter=32,
                 length=None, dtype=np.float32, **kwargs):
    """Invert a mel power spectrogram to audio using Griffin-Lim.

    This is primarily a convenience wrapper for:

        >>> S = librosa.feature.inverse.mel_to_stft(M)
        >>> y = librosa.griffinlim(S)

    Parameters
    ----------
    M : np.ndarray [shape=(n_mels, n), non-negative]
        The spectrogram as produced by `feature.melspectrogram`

    sr : number > 0 [scalar]
        sampling rate of the underlying signal

    n_fft : int > 0 [scalar]
        number of FFT components in the resulting STFT

    hop_length : None or int > 0
        The hop length of the STFT.  If not provided, it will default to `n_fft // 4`

    win_length : None or int > 0
        The window length of the STFT.  By default, it will equal `n_fft`

    window : string, tuple, number, function, or np.ndarray [shape=(n_fft,)]
        A window specification as supported by `stft` or `istft`

    center : boolean
        If `True`, the STFT is assumed to use centered frames.
        If `False`, the STFT is assumed to use left-aligned frames.

    pad_mode : string
        If `center=True`, the padding mode to use at the edges of the signal.
        By default, STFT uses reflection padding.

    power : float > 0 [scalar]
        Exponent for the magnitude melspectrogram

    n_iter : int > 0
        The number of iterations for Griffin-Lim

    length : None or int > 0
        If provided, the output `y` is zero-padded or clipped to exactly `length`
        samples.

    dtype : np.dtype
        Real numeric type for the time-domain signal.  Default is 32-bit float.

    kwargs : additional keyword arguments
        Mel filter bank parameters

    Returns
    -------
    y : np.ndarray [shape(n,)]
        time-domain signal reconstructed from `M`

    See Also
    --------
    core.griffinlim
    feature.melspectrogram
    filters.mel
    feature.inverse.mel_to_stft
    """

    stft = mel_to_stft(M, sr=sr, n_fft=n_fft, power=power, **kwargs)

    return griffinlim(stft, n_iter=n_iter, hop_length=hop_length, win_length=win_length,
                      window=window, center=center, dtype=dtype, length=length,
                      pad_mode=pad_mode)

def mfcc_to_mel(mfcc, n_mels=128, dct_type=2, norm='ortho', ref=1.0, lifter=0):
    '''Invert Mel-frequency cepstral coefficients to approximate a Mel power
    spectrogram.

    This inversion proceeds in two steps:

        1. The inverse DCT is applied to the MFCCs
        2. `core.db_to_power` is applied to map the dB-scaled result to a power spectrogram

    Parameters
    ----------
    mfcc : np.ndarray [shape=(n_mfcc, n)]
        The Mel-frequency cepstral coefficients

    n_mels : int > 0
        The number of Mel frequencies

    dct_type : None or {1, 2, 3}
        Discrete cosine transform (DCT) type
        By default, DCT type-2 is used.

    norm : None or 'ortho'
        If `dct_type` is `2 or 3`, setting `norm='ortho'` uses an orthonormal
        DCT basis.

        Normalization is not supported for `dct_type=1`.

    ref : number or callable
        Reference power for (inverse) decibel calculation

    lifter : number >= 0
        If `lifter>0`, apply inverse liftering (inverse cepstral filtering):
        `M[n, :] <- M[n, :] / (1 + sin(pi * (n + 1) / lifter)) * lifter / 2`

    Returns
    -------
    M : np.ndarray [shape=(n_mels, n)]
        An approximate Mel power spectrum recovered from `mfcc`

    Warns
    --------
    RuntimeWarning
        overflow encountered in power or
        divide by zero encountered during inverse liftering.

    See Also
    --------
    mfcc
    melspectrogram
    scipy.fftpack.dct
    '''
    if lifter > 0:
        n_mfcc = mfcc.shape[0]
        idx = np.arange(1, 1 + n_mfcc, dtype=mfcc.dtype)
        lifter_sine = 1 + (lifter / 2) * np.sin(np.pi * idx / lifter)[:, np.newaxis]

        # check lifter array for critical values
        if (lifter_sine == 0).any():
            warnings.warn("Warning: lifter array includes critial values that will likely to invoke underflow.")

        # lifter mfcc values
        mfcc = mfcc / lifter_sine

    elif lifter != 0:
        raise ParameterError('MFCC to mel lifter={} must be a non-negative number'.format(lifter))

    logmel = scipy.fftpack.idct(mfcc, axis=0, type=dct_type, norm=norm, n=n_mels)
    return db_to_power(logmel, ref=ref)

def mfcc_to_audio(mfcc, n_mels=128, dct_type=2, norm='ortho', ref=1.0, lifter=0, **kwargs):
    '''Convert Mel-frequency cepstral coefficients to a time-domain audio signal

    This function is primarily a convenience wrapper for the following steps:

        1. Convert mfcc to Mel power spectrum (`mfcc_to_mel`)
        2. Convert Mel power spectrum to time-domain audio (`mel_to_audio`)

    Parameters
    ----------
    mfcc : np.ndarray [shape=(n_mfcc, n)]
        The Mel-frequency cepstral coefficients

    n_mels : int > 0
        The number of Mel frequencies

    dct_type : None or {1, 2, 3}
        Discrete cosine transform (DCT) type
        By default, DCT type-2 is used.

    norm : None or 'ortho'
        If `dct_type` is `2 or 3`, setting `norm='ortho'` uses an orthonormal
        DCT basis.

        Normalization is not supported for `dct_type=1`.

    ref : number or callable
        Reference power for (inverse) decibel calculation

    lifter : number >= 0
        If `lifter>0`, apply inverse liftering (inverse cepstral filtering):
        `M[n, :] <- M[n, :] / (1 + sin(pi * (n + 1) / lifter)) * lifter / 2`

    kwargs : additional keyword arguments
        Parameters to pass through to `mel_to_audio`

    Returns
    -------
    y : np.ndarray [shape=(n)]
        A time-domain signal reconstructed from `mfcc`

    See Also
    --------
    mfcc_to_mel
    mel_to_audio
    feature.mfcc
    core.griffinlim
    scipy.fftpack.dct
    '''
    mel_spec = mfcc_to_mel(mfcc, n_mels=n_mels, dct_type=dct_type, norm=norm,
                           ref=ref, lifter=lifter)

    return mel_to_audio(mel_spec, **kwargs)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
pgo-evoleval
Copyright (c) 2016 nipil <https://github.com/nipil>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE
OR OTHER DEALINGS IN THE SOFTWARE.

Author: nipil <https://github.com/nipil>
"""

import os
import sys
import json
import codecs
import hashlib
import argparse

from operator import itemgetter, attrgetter, methodcaller

class Pokemon():

	def __init__(self, data):
		self.id = data[u"pokemon_id"]
		try:
			self.iv_a = data[u"individual_attack"] + 0
		except KeyError:
			self.iv_a = 0
		try:
			self.iv_d = data[u"individual_defense"] + 0
		except KeyError:
			self.iv_d = 0
		try:
			self.iv_s = data[u"individual_stamina"] + 0
		except KeyError:
			self.iv_s = 0
		self.iv_p = float(self.iv_a + self.iv_d + self.iv_s) / 45.0
		self.cur_cp = data[u"cp"]

	def __str__(self):
		return "#{0} A({1})+D({2})+S({3})=P({4}) CP({5})".format(self.id, self.iv_a, self.iv_d, self.iv_s, self.iv_p, self.cur_cp)

class Evoleval():

	@staticmethod
	def load_json(path):
		with open(path) as json_data:
			d = json.load(json_data)
		return d

	def __init__(self, inventory_file, args):
		# init datastores
		self.pokemon_bag = {}
		self.names = {}
		self.families = {}
		self.evolutions = {}
		self.candies = {}
		self.planning = {}
		self.actions = {}
		self.locale = args.locale
		self.egg_time = args.egg_time
		self.evolve_time = args.evolve_time
		# paths
		self.sp_inv = os.path.normpath(os.path.expandvars(os.path.expanduser(inventory_file)))
		self.sp_ref = os.path.normpath(os.path.join(os.path.dirname(self.sp_inv), '../data/pokemon.json'))
		if self.locale is None:
			self.sp_loc = None
		else:
			self.sp_loc = os.path.normpath(os.path.join(os.path.dirname(self.sp_inv), '../data/locales/{0}.json'.format(self.locale)))

	def add_candy(self, data):
		family_id = int(data[u'family_id'])
		family_candy = int(data[u'candy'])
		self.candies[family_id] = family_candy

	def add_pokemon(self, poke):
		if poke.id not in self.pokemon_bag:
			self.pokemon_bag[poke.id] = []
		self.pokemon_bag[poke.id].append(poke)

	def get_pokemon_count(self, p_id):
		try:
			return len(self.pokemon_bag[p_id])
		except KeyError:
			return 0

	def get_family_candies(self, f_id):
		try:
			return self.candies[f_id]
		except KeyError:
			return 0

	def get_evolution_requirement(self, p_id):
		try:
			return self.evolutions[p_id]["next_req"]
		except KeyError:
			return 0

	def load_locale(self):
		if self.sp_loc is None:
			self.locale = {}
		else:
			self.locale = self.load_json(self.sp_loc)

	def localize(self, value):
		if self.locale is None:
			return value
		if value not in self.locale:
			return value
		return self.locale[value]

	def load_reference(self):
		# load reference data
		ref = self.load_json(self.sp_ref)
		for poke in ref:
			info = {}
			p_id = int(poke[u"Number"])
			name = self.localize(poke[u'Name'])
			# ensure that localization is defined
			self.locale[name] = name
			self.locale[p_id] = name
			# family links
			if u"Next Evolution Requirements" in poke:
				# store evolution
				info["family"] = poke[u"Next Evolution Requirements"][u"Family"]
				info["next_req"] = poke[u"Next Evolution Requirements"][u"Amount"]
				info["next_id"] = int(poke[u"Next evolution(s)"][0][u"Number"])
				# add to family
				if info["family"] not in self.families:
					self.families[info["family"]] = []
				self.families[info["family"]].append(p_id)
			else:
				if u"Previous evolution(s)" in poke:
					prev_id = int(poke[u"Previous evolution(s)"][0][u"Number"])
					info["family"] = self.evolutions[prev_id]["family"]
					self.families[info["family"]].append(p_id)
			# store
			self.evolutions[p_id] = info
		# sort families
		for i in sorted(self.families.keys()):
			self.families[i] = sorted(self.families[i])

	def load_inventory(self):
		# load inventory items
		inv = self.load_json(self.sp_inv)
		for k in inv:
			for iid_t, inv_item_data in k.items():
				for iid_st, data in inv_item_data.items():
					if iid_st == "candy":
						self.add_candy(data)
					if iid_st == "pokemon_data":
						if u"is_egg" in data and data[u"is_egg"]:
							continue
						poke = Pokemon(data)
						self.add_pokemon(poke)
		# sort pokemons according to perfection
		for id, pokes in self.pokemon_bag.items():
			self.pokemon_bag[id] = sorted(pokes, key=attrgetter('iv_p', 'cur_cp'), reverse=True)

	def plan_evolution(self):

		for f_id in sorted(self.families.keys()):

			# setup
			candies = self.get_family_candies(f_id)
			p_id_base = self.families[f_id][0]
			pokes = self.get_pokemon_count(p_id_base)
			req = self.get_evolution_requirement(p_id_base)

			# evolve/transfer loop
			transfered = 0
			evolve = 0
			while True:
				# print "candies", candies, "pokes", pokes, "total", candies + pokes, "req", req, "evolve", evolve, "transfered", transfered
				if candies >= req and pokes >= 1:
					# print "evolve"
					candies -= req
					pokes -= 1
					evolve += 1
				elif candies < req and pokes > 1:
					transfer = min(pokes - 1, req - candies)
					# print "transfer", transfer
					candies += transfer
					pokes -= transfer
					transfered += transfer
				else:
					# print "nothing more to do"
					break

			# store action
			self.actions[f_id] = {
				"pid": p_id_base,
				"possible": evolve,
				"missing": candies // req,
				"transfer": transfered,
			}

	def print_pokemons(self):
		print "\n== Pokemon details\n"
		for p_id in sorted(self.pokemon_bag.keys()):
			pokes = self.pokemon_bag[p_id]
			print u"* {0}".format(self.localize(p_id))
			for poke in pokes:
				print "** ({3:.2f}) A/D/S={0}/{1}/{2} CP={4}".format(
					poke.iv_a, poke.iv_d, poke.iv_s, poke.iv_p, poke.cur_cp)

	def print_actions(self):
		print "\n== Evolution actions\n"

		n = sum([ action["possible"] for i, action in self.actions.items() ])

		dt = self.evolve_time
		t = n * dt
		e = float(t / (self.egg_time * 60))
		print "{0} evolutions available, {1} seconds per, {2} seconds total, use {3:.2f} eggs".format(n, dt, t, e)

		print "\n|==="
		print "|Pokemon|Transfer|Evolve|Missing"
		for f_id in sorted(self.actions.keys()):
			action = self.actions[f_id]
			p_id = action["pid"]
			try:
				pokes = self.pokemon_bag[p_id]
			except KeyError:
				pass
			print ""

			# only display if there is something interesting
			if action["transfer"] + action["possible"] + action["missing"] == 0:
				continue

			# display pokemon amount, name and candies
			print u"a|{0} {1} (#{2})\n".format(
				self.get_pokemon_count(action["pid"]),
				self.localize(action["pid"]),
				action["pid"])
			print "{0} candies".format(self.get_family_candies(f_id))

			# display transfer summary and details
			print "a|{0} transfer".format(action["transfer"])
			if action["transfer"] > 0:
				print ""
				for poke in pokes[-action["transfer"]:]:
					print "* ({0:.2f}) CP={1}".format(poke.iv_p, poke.cur_cp)
				print ""

			# display evolution summary and details
			print "a|{0} evolution\n".format(action["possible"])
			if action["possible"] > 0:
				print ""
				for poke in pokes[:action["possible"]]:
					print "* ({0:.2f}) CP={1}".format(poke.iv_p, poke.cur_cp)
				print ""

			# notify of missing pokemons
			print "|{0} missing".format(action["missing"])
		print "|==="

	def get_inv_hash(self, hasher=hashlib.sha256(), blocksize=65536):
		with open(p_inv, 'rb') as inv_file:
			buf = inv_file.read(blocksize)
			while len(buf) > 0:
				hasher.update(buf)
				buf = inv_file.read(blocksize)
			return hasher.hexdigest()

	def output_report(self):
		# compute hash of input file
		inv_hash = self.get_inv_hash()

		# build target file name
		base = os.path.splitext(os.path.split(self.sp_inv)[1])[0]
		target = "{0}_{1}.adoc".format(base, inv_hash)

		# display progress
		print "Generating report {0}".format(target)

		# redirect stdout to file
		stdout_backup = sys.stdout
		sys.stdout = codecs.open(target, 'w', 'utf-8')

		# actual output
		print "= Evoleval report for {0}".format(base)
		self.print_actions()
		self.print_pokemons()

		# restore default stdout
		sys.stdout = stdout_backup

	def run(self):
		self.load_locale()
		self.load_reference()
		self.load_inventory()
		self.plan_evolution()
		self.output_report()

if __name__ == '__main__':

	def verify_strictly_positive(value):
		ivalue = int(value)
		if ivalue <= 0:
			raise argparse.ArgumentTypeError("%s must be a strictly positive integer" % value)
		return ivalue

	parser = argparse.ArgumentParser(description='Evaluate PokemonGo evolutions and inventory')
	parser.add_argument(
		'pgob_dir',
		help='PokemonGo-bot install directory')
	parser.add_argument(
		'--evolve-time',
		action='store',
		default=26,
		metavar='Nsec',
		type=verify_strictly_positive,
		help='duration of a FULL evolution (animation + menu)')
	parser.add_argument(
		'--egg-time',
		action='store',
		default=30,
		metavar='Nmin',
		type=verify_strictly_positive,
		help='duration of a lucky-egg (in minutes)')
	parser.add_argument(
		'--locale',
		action='store',
		default=None,
		metavar='LANG',
		choices=['de','fr','ja','pt_br','ru','zh_cn','zh_hk','zh_tw'],
		help='select language other than ENGLISH for pokemon names')
	args = parser.parse_args()

	pgob_web = os.path.join(os.path.normpath(os.path.expandvars(os.path.expanduser(args.pgob_dir))), "web")

	for file in os.listdir(pgob_web):
		p = os.path.splitext(file)
		if p[1] == '.json' and p[0].startswith('inventory-'):
			p_inv = os.path.join(pgob_web, file)
			ee = Evoleval(p_inv, args)
			ee.run()

	print "Now convert the ascidoctor reports (*.adoc) to HTML using the following command: asciidoctor *.adoc"

	sys.exit(0)
<commit_msg>fix egg calculation type<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
pgo-evoleval
Copyright (c) 2016 nipil <https://github.com/nipil>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE
OR OTHER DEALINGS IN THE SOFTWARE.

Author: nipil <https://github.com/nipil>
"""

import os
import sys
import json
import codecs
import hashlib
import argparse

from operator import itemgetter, attrgetter, methodcaller

class Pokemon():

	def __init__(self, data):
		self.id = data[u"pokemon_id"]
		try:
			self.iv_a = data[u"individual_attack"] + 0
		except KeyError:
			self.iv_a = 0
		try:
			self.iv_d = data[u"individual_defense"] + 0
		except KeyError:
			self.iv_d = 0
		try:
			self.iv_s = data[u"individual_stamina"] + 0
		except KeyError:
			self.iv_s = 0
		self.iv_p = float(self.iv_a + self.iv_d + self.iv_s) / 45.0
		self.cur_cp = data[u"cp"]

	def __str__(self):
		return "#{0} A({1})+D({2})+S({3})=P({4}) CP({5})".format(self.id, self.iv_a, self.iv_d, self.iv_s, self.iv_p, self.cur_cp)

class Evoleval():

	@staticmethod
	def load_json(path):
		with open(path) as json_data:
			d = json.load(json_data)
		return d

	def __init__(self, inventory_file, args):
		# init datastores
		self.pokemon_bag = {}
		self.names = {}
		self.families = {}
		self.evolutions = {}
		self.candies = {}
		self.planning = {}
		self.actions = {}
		self.locale = args.locale
		self.egg_time = args.egg_time
		self.evolve_time = args.evolve_time
		# paths
		self.sp_inv = os.path.normpath(os.path.expandvars(os.path.expanduser(inventory_file)))
		self.sp_ref = os.path.normpath(os.path.join(os.path.dirname(self.sp_inv), '../data/pokemon.json'))
		if self.locale is None:
			self.sp_loc = None
		else:
			self.sp_loc = os.path.normpath(os.path.join(os.path.dirname(self.sp_inv), '../data/locales/{0}.json'.format(self.locale)))

	def add_candy(self, data):
		family_id = int(data[u'family_id'])
		family_candy = int(data[u'candy'])
		self.candies[family_id] = family_candy

	def add_pokemon(self, poke):
		if poke.id not in self.pokemon_bag:
			self.pokemon_bag[poke.id] = []
		self.pokemon_bag[poke.id].append(poke)

	def get_pokemon_count(self, p_id):
		try:
			return len(self.pokemon_bag[p_id])
		except KeyError:
			return 0

	def get_family_candies(self, f_id):
		try:
			return self.candies[f_id]
		except KeyError:
			return 0

	def get_evolution_requirement(self, p_id):
		try:
			return self.evolutions[p_id]["next_req"]
		except KeyError:
			return 0

	def load_locale(self):
		if self.sp_loc is None:
			self.locale = {}
		else:
			self.locale = self.load_json(self.sp_loc)

	def localize(self, value):
		if self.locale is None:
			return value
		if value not in self.locale:
			return value
		return self.locale[value]

	def load_reference(self):
		# load reference data
		ref = self.load_json(self.sp_ref)
		for poke in ref:
			info = {}
			p_id = int(poke[u"Number"])
			name = self.localize(poke[u'Name'])
			# ensure that localization is defined
			self.locale[name] = name
			self.locale[p_id] = name
			# family links
			if u"Next Evolution Requirements" in poke:
				# store evolution
				info["family"] = poke[u"Next Evolution Requirements"][u"Family"]
				info["next_req"] = poke[u"Next Evolution Requirements"][u"Amount"]
				info["next_id"] = int(poke[u"Next evolution(s)"][0][u"Number"])
				# add to family
				if info["family"] not in self.families:
					self.families[info["family"]] = []
				self.families[info["family"]].append(p_id)
			else:
				if u"Previous evolution(s)" in poke:
					prev_id = int(poke[u"Previous evolution(s)"][0][u"Number"])
					info["family"] = self.evolutions[prev_id]["family"]
					self.families[info["family"]].append(p_id)
			# store
			self.evolutions[p_id] = info
		# sort families
		for i in sorted(self.families.keys()):
			self.families[i] = sorted(self.families[i])

	def load_inventory(self):
		# load inventory items
		inv = self.load_json(self.sp_inv)
		for k in inv:
			for iid_t, inv_item_data in k.items():
				for iid_st, data in inv_item_data.items():
					if iid_st == "candy":
						self.add_candy(data)
					if iid_st == "pokemon_data":
						if u"is_egg" in data and data[u"is_egg"]:
							continue
						poke = Pokemon(data)
						self.add_pokemon(poke)
		# sort pokemons according to perfection
		for id, pokes in self.pokemon_bag.items():
			self.pokemon_bag[id] = sorted(pokes, key=attrgetter('iv_p', 'cur_cp'), reverse=True)

	def plan_evolution(self):

		for f_id in sorted(self.families.keys()):

			# setup
			candies = self.get_family_candies(f_id)
			p_id_base = self.families[f_id][0]
			pokes = self.get_pokemon_count(p_id_base)
			req = self.get_evolution_requirement(p_id_base)

			# evolve/transfer loop
			transfered = 0
			evolve = 0
			while True:
				# print "candies", candies, "pokes", pokes, "total", candies + pokes, "req", req, "evolve", evolve, "transfered", transfered
				if candies >= req and pokes >= 1:
					# print "evolve"
					candies -= req
					pokes -= 1
					evolve += 1
				elif candies < req and pokes > 1:
					transfer = min(pokes - 1, req - candies)
					# print "transfer", transfer
					candies += transfer
					pokes -= transfer
					transfered += transfer
				else:
					# print "nothing more to do"
					break

			# store action
			self.actions[f_id] = {
				"pid": p_id_base,
				"possible": evolve,
				"missing": candies // req,
				"transfer": transfered,
			}

	def print_pokemons(self):
		print "\n== Pokemon details\n"
		for p_id in sorted(self.pokemon_bag.keys()):
			pokes = self.pokemon_bag[p_id]
			print u"* {0}".format(self.localize(p_id))
			for poke in pokes:
				print "** ({3:.2f}) A/D/S={0}/{1}/{2} CP={4}".format(
					poke.iv_a, poke.iv_d, poke.iv_s, poke.iv_p, poke.cur_cp)

	def print_actions(self):
		print "\n== Evolution actions\n"

		n = sum([ action["possible"] for i, action in self.actions.items() ])

		dt = self.evolve_time
		t = n * dt
		e = float(t / float(self.egg_time * 60.0))
		print "{0} evolutions available, {1} seconds per, {2} seconds total, use {3:.2f} eggs".format(n, dt, t, e)

		print "\n|==="
		print "|Pokemon|Transfer|Evolve|Missing"
		for f_id in sorted(self.actions.keys()):
			action = self.actions[f_id]
			p_id = action["pid"]
			try:
				pokes = self.pokemon_bag[p_id]
			except KeyError:
				pass
			print ""

			# only display if there is something interesting
			if action["transfer"] + action["possible"] + action["missing"] == 0:
				continue

			# display pokemon amount, name and candies
			print u"a|{0} {1} (#{2})\n".format(
				self.get_pokemon_count(action["pid"]),
				self.localize(action["pid"]),
				action["pid"])
			print "{0} candies".format(self.get_family_candies(f_id))

			# display transfer summary and details
			print "a|{0} transfer".format(action["transfer"])
			if action["transfer"] > 0:
				print ""
				for poke in pokes[-action["transfer"]:]:
					print "* ({0:.2f}) CP={1}".format(poke.iv_p, poke.cur_cp)
				print ""

			# display evolution summary and details
			print "a|{0} evolution\n".format(action["possible"])
			if action["possible"] > 0:
				print ""
				for poke in pokes[:action["possible"]]:
					print "* ({0:.2f}) CP={1}".format(poke.iv_p, poke.cur_cp)
				print ""

			# notify of missing pokemons
			print "|{0} missing".format(action["missing"])
		print "|==="

	def get_inv_hash(self, hasher=hashlib.sha256(), blocksize=65536):
		with open(p_inv, 'rb') as inv_file:
			buf = inv_file.read(blocksize)
			while len(buf) > 0:
				hasher.update(buf)
				buf = inv_file.read(blocksize)
			return hasher.hexdigest()

	def output_report(self):
		# compute hash of input file
		inv_hash = self.get_inv_hash()

		# build target file name
		base = os.path.splitext(os.path.split(self.sp_inv)[1])[0]
		target = "{0}_{1}.adoc".format(base, inv_hash)

		# display progress
		print "Generating report {0}".format(target)

		# redirect stdout to file
		stdout_backup = sys.stdout
		sys.stdout = codecs.open(target, 'w', 'utf-8')

		# actual output
		print "= Evoleval report for {0}".format(base)
		self.print_actions()
		self.print_pokemons()

		# restore default stdout
		sys.stdout = stdout_backup

	def run(self):
		self.load_locale()
		self.load_reference()
		self.load_inventory()
		self.plan_evolution()
		self.output_report()

if __name__ == '__main__':

	def verify_strictly_positive(value):
		ivalue = int(value)
		if ivalue <= 0:
			raise argparse.ArgumentTypeError("%s must be a strictly positive integer" % value)
		return ivalue

	parser = argparse.ArgumentParser(description='Evaluate PokemonGo evolutions and inventory')
	parser.add_argument(
		'pgob_dir',
		help='PokemonGo-bot install directory')
	parser.add_argument(
		'--evolve-time',
		action='store',
		default=26,
		metavar='Nsec',
		type=verify_strictly_positive,
		help='duration of a FULL evolution (animation + menu)')
	parser.add_argument(
		'--egg-time',
		action='store',
		default=30,
		metavar='Nmin',
		type=verify_strictly_positive,
		help='duration of a lucky-egg (in minutes)')
	parser.add_argument(
		'--locale',
		action='store',
		default=None,
		metavar='LANG',
		choices=['de','fr','ja','pt_br','ru','zh_cn','zh_hk','zh_tw'],
		help='select language other than ENGLISH for pokemon names')
	args = parser.parse_args()

	pgob_web = os.path.join(os.path.normpath(os.path.expandvars(os.path.expanduser(args.pgob_dir))), "web")

	for file in os.listdir(pgob_web):
		p = os.path.splitext(file)
		if p[1] == '.json' and p[0].startswith('inventory-'):
			p_inv = os.path.join(pgob_web, file)
			ee = Evoleval(p_inv, args)
			ee.run()

	print "Now convert the ascidoctor reports (*.adoc) to HTML using the following command: asciidoctor *.adoc"

	sys.exit(0)
<|endoftext|>


<commit_before># Copyright (c) 2012 Azavea, Inc.
# See LICENSE in the project root for copying permission

from django.contrib.gis.db import models
from portal.templatetags.portal_extras import nicephone
from django.template.defaultfilters import title
from django.utils.translation import ugettext_lazy as _
from django.utils.translation import pgettext_lazy

class Neighborhood(models.Model):
    """Model for Neighborhoods
    Model for a Neighborhood, which is simply a geographic area with a name
    boundary -- MultiPolygon for the neighborhood
    primary_name -- Display name for the neighborhood
    secondary_name -- Detailed/Unformatted name for the neighborhood
    """
    boundary = models.MultiPolygonField()
    primary_name = models.CharField(max_length=100)
    secondary_name = models.CharField(max_length=100)
    # GeoManager is required so Django knows it can do geospatial queries to the db
    objects = models.GeoManager()

    def save(self, *args, **kwargs):
        """ Override for Model.save()
        Updates neighborhood relation for all Locations that intersect the new Neighborhood boundaries
        """
        super(Neighborhood, self).save(*args, **kwargs)

        # update location AFTER updating neighborhood polygon
        locations = Location.objects.filter(geom__intersects=self.boundary)
        for location in locations:
            # update if no neighborhood or the neighborhood is different
            if location.neighborhood is None or location.neighborhood.id != self.id:
                location.save()

    def __unicode__(self):
        return self.primary_name

class Location(models.Model):
    """Model for school locations in Chicago
    """
    site_name = models.CharField(_('Site Name'), max_length=100)
    address = models.CharField(pgettext_lazy(u'field name', u'Address'), max_length=75)
    city = models.CharField(_('City'), max_length=75)
    state = models.CharField(_('State'), max_length=2)
    zip = models.CharField(_('Zip Code'), max_length=10)
    neighborhood = models.ForeignKey('Neighborhood', null=True)
    phone = models.CharField(_('Phone Number'), max_length=20, blank=True)
    q_rating = models.CharField(_('Quality Rating'), max_length=10, blank=True)
    url = models.CharField(_('Website'), max_length=256, blank=True)
    q_stmt = models.TextField(_('Quality Statement'), blank=True)
    accred = models.CharField(_('Accreditation'), max_length=100, blank=True)
    prg_hours = models.CharField(_('Program Hours'), max_length=50, blank=True)
    is_full_day = models.NullBooleanField(_('Full Day'))
    is_part_day = models.NullBooleanField(_('Part Day'))
    is_full_week = models.NullBooleanField(_('Full Week'))
    is_part_week = models.NullBooleanField(_('Part Week'))
    is_school_year = models.NullBooleanField(_('School Year'))
    is_full_year = models.NullBooleanField(_('Full Year'))
    ages = models.CharField(_('Ages Served'), max_length=50, blank=True)
    is_age_lt_3 = models.NullBooleanField(_('Ages 0-3'))
    is_age_gt_3 = models.NullBooleanField(_('Ages 3-5'))
    language_1 = models.CharField(_('Language 1 (other than English)'), max_length=50, blank=True)
    language_2 = models.CharField(_('Language 2 (other than English)'), max_length=50, blank=True)
    language_3 = models.CharField(_('Language 3 (other than English)'), max_length=50, blank=True)
    is_community_based = models.NullBooleanField(_('Community Based'))
    is_cps_based = models.NullBooleanField(_('CPS Based'))
    is_home_visiting = models.NullBooleanField(_('Home Visiting'))
    accept_ccap = models.NullBooleanField(_('Accepts CCAP'))
    is_hs = models.NullBooleanField(_('Head Start'))
    is_ehs = models.NullBooleanField(_('Early Head Start'))
    geom = models.PointField(_('Geometry'), srid=4326, null=True)
    objects = models.GeoManager()

    # List of simple/boolean fields that should be displayed by Location renderers/views
    display_include = set(['ages', 'prg_hours', 'accred', 'is_home_visiting', 'accept_ccap'])

    def __unicode__(self):
        return self.site_name

    def verbose_name(self, field):
        """
        Given the name of field, returns the verbose_name property for it
        """
        return self._meta.get_field_by_name(field)[0].verbose_name

    @staticmethod
    def get_filter_fields():
        """
        Returns all boolean fields that should be used for filtering Location objects

        This method does not use a static list of names, but rather inspects
        the meta class attached to the model to introspect on the field types.
        This also has the benefit of providing the verbose name of the field.
        """

        # TODO: This relies on fields being ordered in the fields object in the
        # same order they're defined above.  This is probably a sketchy assumption.
        exclude = ['is_montessori', 'is_special_ed']

        fields = []
        for field in Location._meta.fields:
            if (field.get_internal_type() == 'NullBooleanField' and
                    not field.get_attname() in exclude):
                fields.append((field.get_attname(), field.verbose_name,))

        return fields

    def is_true_bool_field(self, field):
        """
        Returns true if field is a boolean field and self.field is True
        """
        fname = field.get_attname()
        return (field.get_internal_type() == 'NullBooleanField' and getattr(self, fname))

    def is_simple_field(self, field):
        """
        Returns true if field is of type CharField or TextField
        """
        ftype = field.get_internal_type()
        return ((ftype == 'CharField' or ftype == 'TextField'))

    def get_context_dict(self):
        # Fix some ugly data
        if self.site_name.isupper():
            self.site_name = title(self.site_name)
        if self.address.isupper():
            self.address = title(self.address)
        if self.city.isupper():
            self.city = title(self.city)

        item = {'address': self.address,
                'city': self.city,
                'site_name': self.site_name,
                'zip': self.zip,
                'url': self.url,
                'state': self.state,
                'key': self.pk}

        # simple fields to present -- these are the attributes that have text content
        sfields = []

        # boolean fields to present -- these are the attributes that are set to True
        bfields = []

        for field in Location._meta.fields:
            fname = field.get_attname()
            if not fname in Location.display_include:
                continue

            if self.is_true_bool_field(field):
                bfields.append(field.verbose_name)
            elif self.is_simple_field(field):
        # Affiliations
        affiliation_fields = [(self.is_community_based, 'is_community_based'),
                              (self.is_cps_based, 'is_cps_based'),
                              (self.is_hs, 'is_hs'),
                              (self.is_ehs, 'is_ehs')]
        affiliation_values = [self.verbose_name(aff[1]).encode() for aff in affiliation_fields if aff[0]]
        sfields.append({'fieldname': u'%s' % _('Affiliations'), 'value': u', '.join(affiliation_values) if affiliation_values else 'None'})

                kv = (field.verbose_name, getattr(self, fname))
                sfields.append(kv)

        # Combine Languages
        lang_list = [lang for lang in self.language_1, self.language_2, self.language_3 if lang]
        languages = ", ".join(lang_list)
        if languages != '':
            sfields.append((_('Languages (other than English)'), languages))

        # Program Duration
        sfields.append((_('Program Duration'), 'Full Year' if self.is_full_year else 'School Year'))

        # Week Duration
        sfields.append((_('Weekday Availability'), 'Full Week' if self.is_full_week else 'Partial Week'))

        bfields.sort()
        sfields.sort(key=lambda a: a[0])
        return { 'item': self, 'sfields': sfields, 'bfields': bfields }

    def val_or_empty(self, field, f=(lambda x: x)):
        """
        Returns a pretty string representing the value of a field if it's present, otherwise an
        empty string.  The pretty string ends in a newline.
        field:  string with the name of a field on this model
        f:      optional rendering function.  It's given the value of the field, modifies it, and
                returns the result.  Default returns its input
        """
        val = self.__dict__[field]
        return ("%s: %s\n" % (self.verbose_name(field), f(val))) if val else ""

    def save(self, *args, **kwargs):
        """ Override for Model.save()
        Overrides Location.save(). Provides the additional functionality of updating the Neighborhood
        of the Location before the save.
        """
        if self.geom is not None:
            neighborhoods = Neighborhood.objects.filter(boundary__intersects=self.geom)
            if len(neighborhoods):
                self.neighborhood = neighborhoods[0]

        super(Location, self).save(*args, **kwargs)

<commit_msg>Remove translation stuff, set up field names and values to return json<commit_after># Copyright (c) 2012 Azavea, Inc.
# See LICENSE in the project root for copying permission

from django.contrib.gis.db import models
from portal.templatetags.portal_extras import nicephone
from django.template.defaultfilters import title

class Neighborhood(models.Model):
    """Model for Neighborhoods
    Model for a Neighborhood, which is simply a geographic area with a name
    boundary -- MultiPolygon for the neighborhood
    primary_name -- Display name for the neighborhood
    secondary_name -- Detailed/Unformatted name for the neighborhood
    """
    boundary = models.MultiPolygonField()
    primary_name = models.CharField(max_length=100)
    secondary_name = models.CharField(max_length=100)
    # GeoManager is required so Django knows it can do geospatial queries to the db
    objects = models.GeoManager()

    def save(self, *args, **kwargs):
        """ Override for Model.save()
        Updates neighborhood relation for all Locations that intersect the new Neighborhood boundaries
        """
        super(Neighborhood, self).save(*args, **kwargs)

        # update location AFTER updating neighborhood polygon
        locations = Location.objects.filter(geom__intersects=self.boundary)
        for location in locations:
            # update if no neighborhood or the neighborhood is different
            if location.neighborhood is None or location.neighborhood.id != self.id:
                location.save()

    def __unicode__(self):
        return self.primary_name

class Location(models.Model):
    """Model for school locations in Chicago
    """
    site_name = models.CharField('Site Name', max_length=100)
    address = models.CharField('Address', max_length=75)
    city = models.CharField('City', max_length=75)
    state = models.CharField('State', max_length=2)
    zip = models.CharField('Zip Code', max_length=10)
    neighborhood = models.ForeignKey('Neighborhood', null=True)
    phone = models.CharField('Phone Number', max_length=20, blank=True)
    q_rating = models.CharField('Quality Rating', max_length=10, blank=True)
    url = models.CharField('Website', max_length=256, blank=True)
    q_stmt = models.TextField('Quality Statement', blank=True)
    accred = models.CharField('Accreditation', max_length=100, blank=True)
    prg_hours = models.CharField('Program Hours', max_length=50, blank=True)
    is_full_day = models.NullBooleanField('Full Day')
    is_part_day = models.NullBooleanField('Part Day')
    is_full_week = models.NullBooleanField('Full Week')
    is_part_week = models.NullBooleanField('Part Week')
    is_school_year = models.NullBooleanField('School Year')
    is_full_year = models.NullBooleanField('Full Year')
    ages = models.CharField('Ages Served', max_length=50, blank=True)
    is_age_lt_3 = models.NullBooleanField('Ages 0-3')
    is_age_gt_3 = models.NullBooleanField('Ages 3-5')
    language_1 = models.CharField('Language 1 (other than English)', max_length=50, blank=True)
    language_2 = models.CharField('Language 2 (other than English)', max_length=50, blank=True)
    language_3 = models.CharField('Language 3 (other than English)', max_length=50, blank=True)
    is_community_based = models.NullBooleanField('Community Based')
    is_cps_based = models.NullBooleanField('CPS Based')
    is_home_visiting = models.NullBooleanField('Home Visiting')
    accept_ccap = models.NullBooleanField('Accepts CCAP')
    is_hs = models.NullBooleanField('Head Start')
    is_ehs = models.NullBooleanField('Early Head Start')
    geom = models.PointField('Geometry', srid=4326, null=True)
    objects = models.GeoManager()

    # List of simple/boolean fields that should be displayed by Location renderers/views
    display_include = set(['ages', 'prg_hours', 'accred', 'is_home_visiting', 'accept_ccap'])

    def __unicode__(self):
        return self.site_name

    def verbose_name(self, field):
        """
        Given the name of field, returns the verbose_name property for it
        """
        return self._meta.get_field_by_name(field)[0].verbose_name

    @staticmethod
    def get_filter_fields():
        """
        Returns all boolean fields that should be used for filtering Location objects

        This method does not use a static list of names, but rather inspects
        the meta class attached to the model to introspect on the field types.
        This also has the benefit of providing the verbose name of the field.
        """

        # TODO: This relies on fields being ordered in the fields object in the
        # same order they're defined above.  This is probably a sketchy assumption.
        exclude = ['is_montessori', 'is_special_ed']

        fields = []
        for field in Location._meta.fields:
            if (field.get_internal_type() == 'NullBooleanField' and
                    not field.get_attname() in exclude):
                fields.append((field.get_attname(), field.verbose_name,))

        return fields

    def is_true_bool_field(self, field):
        """
        Returns true if field is a boolean field and self.field is True
        """
        fname = field.get_attname()
        return (field.get_internal_type() == 'NullBooleanField' and getattr(self, fname))

    def is_simple_field(self, field):
        """
        Returns true if field is of type CharField or TextField
        """
        ftype = field.get_internal_type()
        return ((ftype == 'CharField' or ftype == 'TextField'))

    def get_context_dict(self):
        # Fix some ugly data
        if self.site_name.isupper():
            self.site_name = title(self.site_name)
        if self.address.isupper():
            self.address = title(self.address)
        if self.city.isupper():
            self.city = title(self.city)

        item = {'address': self.address,
                'city': self.city,
                'site_name': self.site_name,
                'zip': self.zip,
                'url': self.url,
                'state': self.state,
                'key': self.pk}

        # simple fields to present -- these are the attributes that have text content
        sfields = []

        # boolean fields to present -- these are the attributes that are set to True
        bfields = []

        for field in Location._meta.fields:
            fname = field.get_attname()
            if not fname in Location.display_include:
                continue

            if self.is_true_bool_field(field):
                bfields.append(field.verbose_name)
            elif self.is_simple_field(field):
                kv = {'fieldname': field.verbose_name, 'value': getattr(self, fname)}
                sfields.append(kv)

        # Affiliations
        affiliation_fields = [(self.is_community_based, 'is_community_based'),
                              (self.is_cps_based, 'is_cps_based'),
                              (self.is_hs, 'is_hs'),
                              (self.is_ehs, 'is_ehs')]
        affiliation_values = [self.verbose_name(aff[1]).encode() for aff in affiliation_fields if aff[0]]
        sfields.append({'fieldname': 'Affiliations', 'value': ', '.join(affiliation_values) if affiliation_values else 'None'})

        # Combine Languages
        lang_list = [lang for lang in self.language_1, self.language_2, self.language_3 if lang]
        languages = ", ".join(lang_list)
        if languages != '':
            sfields.append({'fieldname': 'Languages (other than English)', 'value': languages})

        # Program Duration
        sfields.append({'fieldname': 'Program Duration', 'value': 'Full Year' if self.is_full_year else 'School Year'})

        # Week Duration
        sfields.append({'fieldname': 'Weekday Availability', 'value': 'Full Week' if self.is_full_week else 'Partial Week'})

        # Phone
        phone = {'fieldname': self.verbose_name('phone'), 'number': nicephone(self.phone)}

        # Position
        position = {'lng': self.geom[0], 'lat': self.geom[1]}
        bfields.sort()
        sfields.sort(key=lambda a: a['fieldname'])
        return {'item': item, 'phone': phone, 'sfields': sfields, 'bfields': bfields, 'position': position}

    def val_or_empty(self, field, f=(lambda x: x)):
        """
        Returns a pretty string representing the value of a field if it's present, otherwise an
        empty string.  The pretty string ends in a newline.
        field:  string with the name of a field on this model
        f:      optional rendering function.  It's given the value of the field, modifies it, and
                returns the result.  Default returns its input
        """
        val = self.__dict__[field]
        return ("%s: %s\n" % (self.verbose_name(field), f(val))) if val else ""

    def save(self, *args, **kwargs):
        """ Override for Model.save()
        Overrides Location.save(). Provides the additional functionality of updating the Neighborhood
        of the Location before the save.
        """
        if self.geom is not None:
            neighborhoods = Neighborhood.objects.filter(boundary__intersects=self.geom)
            if len(neighborhoods):
                self.neighborhood = neighborhoods[0]

        super(Location, self).save(*args, **kwargs)

<|endoftext|>


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

from subprocess import call
import os

import argparse
parser = argparse.ArgumentParser()

parser.add_argument("-i", "--input", help="GVCF.GZ file (can be the location on S3)")

args = parser.parse_args()
vcf_file = args.input

print(vcf_file)

input_folder = "/home/ubuntu/projects/input/vcf"

if vcf_file.startswith('s3://'):
    #download file to input folder
    command = "s3cmd get %s %s/" % (vcf_file, input_folder)
    output = call(command, shell=True)
    print(output)

base=os.path.basename(vcf_file)
# base_name = os.path.splitext(base)[0]
base_name = base.split('.')[0]
print(base_name)

die()

#create one folder per sample
output_folder = "/home/ubuntu/projects/output/reports/vcf/%s" % (base_name)
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

output_base = "%s/%s" % (output_folder, base_name)

memory_use = "15g"
gvcftools_path = "/home/ubuntu/projects/programs/gvcftools-0.16/bin"
vcftools_path = "/home/ubuntu/projects/programs/vcftools/vcftools-0.1.14/src"
bcftools_path = "/home/ubuntu/projects/programs/bcftools/bcftools-1.3.1"
snpeff_path = "/home/ubuntu/projects/programs/snpeff/snpEff"
#/home/ubuntu/projects/programs/vcftools/vcftools-0.1.14/src/cpp/
#/home/ubuntu/projects/programs/vcftools/vcftools-0.1.14/src/perl/

#extract vcf from gvcf
print('extract vcf from gvcf')
#gzip -dc ../../input/WGC081270U.g.vcf.gz | ../../programs/gvcftools-0.16/bin/extract_variants | bgzip -c > WGC081270U.vcf.gz
command = """gzip -dc %s | %s/extract_variants | bgzip -c > %s.vcf.gz
""" % (vcf_file, gvcftools_path, output_base)
output = call(command, shell=True)
print(output)

#index with tabix
command = """tabix -p vcf %s.vcf.gz""" % (output_base)
output = call(command, shell=True)
print(output)

print('vcftools stats')
#vcftools metrics
command = """%s/perl/vcf-stats %s.vcf.gz > %s.vcftools.stats.txt
""" % (vcftools_path, output_base, output_base)
output = call(command, shell=True)
print(output)

#*coverage
#bcftools metrics
print('bcftools stats')
command = "%s/bcftools stats %s.vcf.gz > %s.bcftools.stats.txt" % (bcftools_path, output_base, output_base)
output = call(command, shell=True)
print(output)

#snpeff
#extract vcf
print('extract vcf')
command = "bgzip -d -c %s.vcf.gz > %s.vcf" % (output_base, output_base)
output = call(command, shell=True)
print(output)

print('snpeff')
command = "java -Xmx5g -jar %s/snpEff.jar eff -stats %s.snpeff.full.html -i vcf GRCh37.75 %s.vcf > %s.snpeff.full.vcf" % (snpeff_path, output_base, output_base, output_base)
output = call(command, shell=True)
print(output)

# print('filtering')
# #filtering VCF
# command = "%s/bcftools filter -T %s %s.vcf > %s.filtered.exons.vcf" % (bcftools_path, target_file, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# #filtered.exons.q50.dp50.vcf
# command = "%s/bcftools filter -T %s -i'QUAL>50 && FMT/DP>50' %s.vcf > %s.filtered.exons.q50.dp50.vcf" % (bcftools_path, target_file, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# #WGC081270U.filtered.exons.q100.dp100.vcf
# command = "%s/bcftools filter -T %s -i'QUAL>100 && FMT/DP>100' %s.vcf > %s.filtered.exons.q100.dp100.vcf" % (bcftools_path, target_file, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# command = "java -Xmx5g -jar %s/snpEff.jar eff -stats %s.snpeff.exons.html -i vcf GRCh37.75 %s.filtered.exons.vcf > %s.snpeff.exons.vcf" % (snpeff_path, output_base, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# print('snpeff')
# command = "java -Xmx5g -jar %s/snpEff.jar eff -stats %s.snpeff.exons.q100.dp100.html -i vcf GRCh37.75 %s.filtered.exons.q100.dp100.vcf > %s.snpeff.exons.q100.dp100.vcf" % (snpeff_path, output_base, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# print('bcftools stats')
# command = "%s/bcftools stats %s.filtered.exons.q100.dp100.vcf > %s.bcftools.stats.exons.q100.dp100.txt" % (bcftools_path, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# print('plot vcf-stats')
# command = "%s/plot-vcfstats %s.bcftools.stats.exons.q100.dp100.txt -p %s" % (bcftools_path, output_base, output_base)
# output = call(command, shell=True)
# print(output)
<commit_msg>vcf genome stats<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from subprocess import call
import os

import argparse
parser = argparse.ArgumentParser()

parser.add_argument("-i", "--input", help="GVCF.GZ file (can be the location on S3)")

args = parser.parse_args()
vcf_file = args.input

print(vcf_file)

input_folder = "/home/ubuntu/projects/input/vcf"

if vcf_file.startswith('s3://'):
    #download file to input folder
    command = "s3cmd get %s %s/" % (vcf_file, input_folder)
    output = call(command, shell=True)
    print(output)

base=os.path.basename(vcf_file)
# base_name = os.path.splitext(base)[0]
base_name = base.split('.')[0]
print(base_name)

#create one folder per sample
output_folder = "/home/ubuntu/projects/output/reports/vcf/%s" % (base_name)
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

output_base = "%s/%s" % (output_folder, base_name)

memory_use = "15g"
gvcftools_path = "/home/ubuntu/projects/programs/gvcftools-0.16/bin"
vcftools_path = "/home/ubuntu/projects/programs/vcftools/vcftools-0.1.14/src"
bcftools_path = "/home/ubuntu/projects/programs/bcftools/bcftools-1.3.1"
snpeff_path = "/home/ubuntu/projects/programs/snpeff/snpEff"
#/home/ubuntu/projects/programs/vcftools/vcftools-0.1.14/src/cpp/
#/home/ubuntu/projects/programs/vcftools/vcftools-0.1.14/src/perl/

#extract vcf from gvcf
print('extract vcf from gvcf')
#gzip -dc ../../input/WGC081270U.g.vcf.gz | ../../programs/gvcftools-0.16/bin/extract_variants | bgzip -c > WGC081270U.vcf.gz
command = """gzip -dc %s | %s/extract_variants | bgzip -c > %s.vcf.gz
""" % (vcf_file, gvcftools_path, output_base)
output = call(command, shell=True)
print(output)

#index with tabix
command = """tabix -p vcf %s.vcf.gz""" % (output_base)
output = call(command, shell=True)
print(output)

print('vcftools stats')
#vcftools metrics
command = """%s/perl/vcf-stats %s.vcf.gz > %s.vcftools.stats.txt
""" % (vcftools_path, output_base, output_base)
output = call(command, shell=True)
print(output)

#*coverage
#bcftools metrics
print('bcftools stats')
command = "%s/bcftools stats %s.vcf.gz > %s.bcftools.stats.txt" % (bcftools_path, output_base, output_base)
output = call(command, shell=True)
print(output)

#snpeff
#extract vcf
# print('extract vcf')
# command = "bgzip -d -c %s.vcf.gz > %s.vcf" % (output_base, output_base)
# output = call(command, shell=True)
# print(output)

print('snpeff')
command = "java -Xmx5g -jar %s/snpEff.jar eff -stats %s.snpeff.full.html -i vcf GRCh37.75 %s.vcf | bgzip -c > %s.snpeff.full.vcf.gz" % (snpeff_path, output_base, output_base, output_base)
output = call(command, shell=True)
print(output)

# print('filtering')
# #filtering VCF
# command = "%s/bcftools filter -T %s %s.vcf > %s.filtered.exons.vcf" % (bcftools_path, target_file, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# #filtered.exons.q50.dp50.vcf
# command = "%s/bcftools filter -T %s -i'QUAL>50 && FMT/DP>50' %s.vcf > %s.filtered.exons.q50.dp50.vcf" % (bcftools_path, target_file, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# #WGC081270U.filtered.exons.q100.dp100.vcf
# command = "%s/bcftools filter -T %s -i'QUAL>100 && FMT/DP>100' %s.vcf > %s.filtered.exons.q100.dp100.vcf" % (bcftools_path, target_file, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# command = "java -Xmx5g -jar %s/snpEff.jar eff -stats %s.snpeff.exons.html -i vcf GRCh37.75 %s.filtered.exons.vcf > %s.snpeff.exons.vcf" % (snpeff_path, output_base, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# print('snpeff')
# command = "java -Xmx5g -jar %s/snpEff.jar eff -stats %s.snpeff.exons.q100.dp100.html -i vcf GRCh37.75 %s.filtered.exons.q100.dp100.vcf > %s.snpeff.exons.q100.dp100.vcf" % (snpeff_path, output_base, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# print('bcftools stats')
# command = "%s/bcftools stats %s.filtered.exons.q100.dp100.vcf > %s.bcftools.stats.exons.q100.dp100.txt" % (bcftools_path, output_base, output_base)
# output = call(command, shell=True)
# print(output)

# print('plot vcf-stats')
# command = "%s/plot-vcfstats %s.bcftools.stats.exons.q100.dp100.txt -p %s" % (bcftools_path, output_base, output_base)
# output = call(command, shell=True)
# print(output)
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
GeoBase. Rebooted.

    >>> b = NeoBase()
    >>> b.get('ORY', 'city_code_list')
    ['PAR']
    >>> b.get('ORY', 'city_name_list')
    ['Paris']
    >>> b.get('ORY', 'country_code')
    'FR'
    >>> b.distance('ORY', 'CDG')
    34.87...
    >>> b.get_location('ORY')
    LatLng(lat=48.72..., lng=2.35...)
"""

from __future__ import with_statement, print_function, division

import os.path as op
import csv
import heapq
from collections import namedtuple
from math import pi, cos, sin, asin, sqrt

__all__ = ['NeoBase', 'LatLng']

_DIR = op.dirname(__file__)
_FILE_PATH = op.join(_DIR, 'optd_por_public.csv')
_DEFAULT_RADIUS = 50

_FIELDS = (
    ('iata_code', 0),
    ('name', 6),
    ('lat', 8),
    ('lng', 9),
    ('country_code', 16),
    ('continent_name', 19),
    ('timezone', 31),
    ('city_code_list', 36),
    ('city_name_list', 37),
    ('location_type', 41),
)

_SPLITS = (
    ('city_code_list', lambda s: s.split(',')),
    ('city_name_list', lambda s: s.split('=')),
    ('location_type', list),
)

_KEY = 0 # iata_code

LatLng = namedtuple('LatLng', ['lat', 'lng'])

class NeoBase(object):
    """Main structure, a wrapper around a dict, with dict-like behavior.
    """
    def __init__(self, rows=None):
        if rows is None:
            with open(_FILE_PATH) as f:
                self._data = self.load(f)
        else:
            self._data = self.load(rows)

    @staticmethod
    def _empty_value():
        return {'__dup__': set()}

    @classmethod
    def load(cls, f):
        """Building a dictionary of geographical data from optd_por.

        >>> path = op.join(_DIR, 'optd_por_public.csv')
        >>> with open(path) as f:
        ...     b = NeoBase.load(f)
        >>> b['ORY']['city_code_list']
        ['PAR']
        """
        data = {}
        next(f) # skipping first line

        for row in csv.reader(f, delimiter='^', quotechar='"'):
            # Comments and empty lines
            if not row or row[0].startswith('#'):
                continue

            key = row[_KEY]
            if key not in data:
                data[key] = d = cls._empty_value()
            else:
                prev_d = data[key]
                new_key = '{0}@{1}'.format(key, 1 + len(prev_d['__dup__']))
                data[new_key] = d = cls._empty_value()
                # Exchanging duplicata information
                d['__dup__'] = prev_d['__dup__'] | set([key])
                prev_d['__dup__'].add(new_key)

            for field, c in _FIELDS:
                d[field] = row[c]

            for field, splitter in _SPLITS:
                d[field] = splitter(d[field])

        return data

    def __iter__(self):
        """Returns iterator of all keys in the base.

        :returns: the iterator of all keys

        >>> b = NeoBase()
        >>> sorted(b)
        ['AAA', 'AAA@1', 'AAB', 'AAC', 'AAC@1', ...
        """
        return iter(self._data)

    def __contains__(self, key):
        """Test if a key is in the base.

        :param key: the key of to be tested
        :returns:   a boolean

        >>> b = NeoBase()
        >>> 'AN' in b
        False
        >>> 'AGN' in b
        True
        """
        return key in self._data

    def __nonzero__(self):
        """Testing structure emptiness.

        :returns: a boolean

        >>> b = NeoBase()
        >>> if b:
        ...     print('not empty')
        not empty
        """
        return bool(self._data)

    def __len__(self):
        """Testing structure size.

        :returns: a integer

        >>> b = NeoBase()
        >>> 18000 < len(b) < 20000
        True
        """
        return len(self._data)

    def keys(self):
        """Returns iterator of all keys in the base.

        :returns: the iterator of all keys

        >>> b = NeoBase()
        >>> sorted(b.keys())
        ['AAA', 'AAA@1', 'AAB', 'AAC', 'AAC@1', ...
        """
        return iter(self)

    def set(self, key, **kwargs):
        """Set information.

        >>> b = NeoBase()
        >>> b.get('ORY', 'name')
        'Paris Orly Airport'
        >>> b.set('ORY', name='test')
        >>> b.get('ORY', 'name')
        'test'
        >>> b.set('Wow!', name='test')
        >>> b.get('Wow!', 'name')
        'test'
        """
        if key not in self:
            self._data[key] = self._empty_value()
        self._data[key].update(kwargs)

    def get(self, key, field=None, **kwargs):
        """Get data from structure.

        >>> b = NeoBase()
        >>> b.get('OR', 'city_code_list', default=None)
        >>> b.get('ORY', 'city_code_list')
        ['PAR']
        """
        try:
            d = self._data[key]
        except KeyError:
            # Unless default is set, we raise an Exception
            if 'default' in kwargs:
                return kwargs['default']
            else:
                raise KeyError("Key not found: {0}".format(key))

        if field is None:
            return d # we return the whole dictionary

        try:
            res = d[field]
        except KeyError:
            raise KeyError("Field '{0}' (for key '{1}') not in {2}".format(
                field, key, list(d)))
        else:
            return res

    def get_location(self, key, **kwargs):
        """Get None or the geocode.

        >>> b = NeoBase()
        >>> b.get_location('ORY')
        LatLng(lat=48.72..., lng=2.35...)
        """
        if key not in self:
            # Unless default is set, we raise an Exception
            if 'default' in kwargs:
                return kwargs['default']
            else:
                raise KeyError("Key not found: {0}".format(key))

        try:
            loc = LatLng(float(self.get(key, 'lat')),
                         float(self.get(key, 'lng')))

        except (ValueError, TypeError, KeyError):
            # Decode geocode, if error, returns None
            # TypeError : input type is not a string, probably None
            # ValueError: could not convert to float
            # KeyError  : could not find lat or lng 'fields'
            return None
        else:
            return loc

    @staticmethod
    def distance_between_locations(l0, l1):
        """Great circle distance

        :param l0: the LatLng tuple of the first location
        :param l1: the LatLng tuple of the second location
        :returns:  the distance in kilometers

        >>> NeoBase.distance_between_locations((48.84, 2.367), (43.70, 7.26)) # Paris -> Nice
        683.85...

        Case of unknown location.

        >>> NeoBase.distance_between_locations(None, (43.70, 7.26)) # returns None
        """
        if l0 is None or l1 is None:
            return None

        l0_lat = l0[0] / 180 * pi
        l0_lng = l0[1] / 180 * pi
        l1_lat = l1[0] / 180 * pi
        l1_lng = l1[1] / 180 * pi

        # Haversine formula (6371 is Earth radius)
        return 2 * 6371.0 * asin(sqrt(
            sin(0.5 * (l0_lat - l1_lat)) ** 2 +
            sin(0.5 * (l0_lng - l1_lng)) ** 2 *
            cos(l0_lat) * cos(l1_lat)
        ))

    def distance(self, key_0, key_1):
        """Compute distance between two elements.

        This is just a wrapper between the original haversine
        function, but it is probably the most used feature :)

        :param key_0: the first key
        :param key_1: the second key
        :returns:    the distance (km)

        >>> b = NeoBase()
        >>> b.distance('ORY', 'CDG')
        34.87...
        """
        return self.distance_between_locations(self.get_location(key_0),
                                               self.get_location(key_1))

    def _build_distances(self, lat_lng_ref, keys):
        """
        Compute the iterable of (dist, keys) of a reference
        lat_lng and a list of keys. Keys which have not valid
        geocodes will not appear in the results.

        >>> b = NeoBase()
        >>> list(b._build_distances((0,0), ['ORY', 'CDG']))
        [(5422.74..., 'ORY'), (5455.45..., 'CDG')]
        """
        if lat_lng_ref is None:
            return

        for key in keys:
            if key in self:
                lat_lng = self.get_location(key)
                if lat_lng is not None:
                    yield self.distance_between_locations(lat_lng_ref, lat_lng), key

    def find_near_location(self, lat_lng, radius=_DEFAULT_RADIUS, from_keys=None):
        """
        Returns a list of nearby keys from a location (given
        latidude and longitude), and a radius for the search.
        Note that the haversine function, which compute distance
        at the surface of a sphere, here returns kilometers,
        so the radius should be in kms.

        :param lat_lng: the lat_lng of the location
        :param radius:  the radius of the search (kilometers)
        :param from_keys: if None, it takes all keys in consideration, else takes from_keys \
            iterable of keys to perform search.
        :returns:       an iterable of keys (like ['ORY', 'CDG'])

        >>> b = NeoBase()
        >>> # Paris, airports <= 50km
        >>> [b.get(k, 'iata_code') for d, k in sorted(b.find_near_location((48.84, 2.367), 10))]
        ['PAR', 'XGB', 'XHP', 'XPG', 'XEX', 'XTT', 'QAF', 'JDP', 'XBT', 'QNL', 'QBH', 'QFC', 'QEV']
        """
        if from_keys is None:
            from_keys = iter(self)

        for dist, key in self._build_distances(lat_lng, from_keys):
            if dist <= radius:
                yield dist, key

    def find_near(self, key, radius=_DEFAULT_RADIUS, from_keys=None):
        """
        Same as find_near_location, except the location is given
        not by a lat/lng, but with its key, like ORY or SFO.
        We just look up in the base to retrieve lat/lng, and
        call find_near_location.

        :param key:     the key of the location
        :param radius:  the radius of the search (kilometers)
        :param from_keys: if None, it takes all keys in consideration, else takes from_keys \
            iterable of keys to perform search.
        :returns:       a list of keys (like ['ORY', 'CDG'])

        >>> b = NeoBase()
        >>> sorted(b.find_near('ORY', 10)) # Orly, por <= 10km
        [(0.0, 'ORY'), (6.94..., 'XJY'), (9.96..., 'QFC')]
        """
        if from_keys is None:
            from_keys = iter(self)

        if key not in self:
            return

        for dist, key in self.find_near_location(self.get_location(key),
                                                 radius=radius,
                                                 from_keys=from_keys):
            yield dist, key

    def find_closest_from_location(self, lat_lng, N=1, from_keys=None):
        """
        Concept close to find_near_location, but here we do not
        look for the keys radius-close to a location,
        we look for the closest key from this location, given by
        latitude/longitude.

        Note that a similar implementation is done in
        the LocalHelper, to find efficiently N closest location
        in a graph, from a location (using heaps).

        :param lat_lng:   the lat_lng of the location
        :param N:         the N closest results wanted
        :param from_keys: if None, it takes all keys in consideration, else takes from_keys \
            iterable of keys to perform find_closest_from_location. This is useful to combine \
            searches
        :returns:   one key (like 'SFO'), or a list if approximate is not None

        >>> b = NeoBase()
        >>> list(b.find_closest_from_location((43.70, 7.26))) # Nice
        [(0.60..., 'NCE@1')]
        >>> list(b.find_closest_from_location((43.70, 7.26), N=3)) # Nice
        [(0.60..., 'NCE@1'), (5.82..., 'NCE'), (5.89..., 'XBM')]
        """
        if from_keys is None:
            from_keys = iter(self)

        iterable = self._build_distances(lat_lng, from_keys)

        for dist, key in heapq.nsmallest(N, iterable):
            yield dist, key

    def find_with(self, conditions, from_keys=None, reverse=False):
        """Get iterator of all keys with particular field.

        For example, if you want to know all airports in Paris.

        :param conditions: a list of (field, value) conditions
        :param reverse:    we look keys where the field is *not* the particular value
        :returns:          an iterator of matching keys

        Testing several conditions.

        >>> b = NeoBase()
        >>> c0 = [('city_code_list', ['PAR'])]
        >>> c1 = [('location_type', ['H'])]
        >>> len(list(b.find_with(c0)))
        17
        >>> len(list(b.find_with(c1)))
        146
        >>> len(list(b.find_with(c0 + c1)))
        2
        """
        if from_keys is None:
            from_keys = iter(self)

        if not reverse:
            match = lambda a, b: a == b
        else:
            match = lambda a, b: a != b

        for key in from_keys:
            if key in self:
                if all(match(self.get(key, f), v) for f, v in conditions):
                    yield key

def main():
    import argparse
    from pprint import pprint

    parser = argparse.ArgumentParser()
    parser.add_argument('keys', nargs='+', help="List of IATA codes")
    args = parser.parse_args()
    b = NeoBase()

    for key in args.keys:
        if key in b:
            pprint(b.get(key))
        else:
            print('{0} not in data.'.format(key))

if __name__ == '__main__':
    main()

<commit_msg>Add proper sentinel implementation<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
GeoBase. Rebooted.

    >>> b = NeoBase()
    >>> b.get('ORY', 'city_code_list')
    ['PAR']
    >>> b.get('ORY', 'city_name_list')
    ['Paris']
    >>> b.get('ORY', 'country_code')
    'FR'
    >>> b.distance('ORY', 'CDG')
    34.87...
    >>> b.get_location('ORY')
    LatLng(lat=48.72..., lng=2.35...)
"""

from __future__ import with_statement, print_function, division

import os.path as op
import csv
import heapq
from collections import namedtuple
from math import pi, cos, sin, asin, sqrt

__all__ = ['NeoBase', 'LatLng']

_DIR = op.dirname(__file__)
_FILE_PATH = op.join(_DIR, 'optd_por_public.csv')
_DEFAULT_RADIUS = 50

_FIELDS = (
    ('iata_code', 0),
    ('name', 6),
    ('lat', 8),
    ('lng', 9),
    ('country_code', 16),
    ('continent_name', 19),
    ('timezone', 31),
    ('city_code_list', 36),
    ('city_name_list', 37),
    ('location_type', 41),
)

_SPLITS = (
    ('city_code_list', lambda s: s.split(',')),
    ('city_name_list', lambda s: s.split('=')),
    ('location_type', list),
)

_KEY = 0 # iata_code

LatLng = namedtuple('LatLng', ['lat', 'lng'])

# Sentinel value for signatures
_sentinel = object()

class NeoBase(object):
    """Main structure, a wrapper around a dict, with dict-like behavior.
    """
    def __init__(self, rows=None):
        if rows is None:
            with open(_FILE_PATH) as f:
                self._data = self.load(f)
        else:
            self._data = self.load(rows)

    @staticmethod
    def _empty_value():
        return {'__dup__': set()}

    @classmethod
    def load(cls, f):
        """Building a dictionary of geographical data from optd_por.

        >>> path = op.join(_DIR, 'optd_por_public.csv')
        >>> with open(path) as f:
        ...     b = NeoBase.load(f)
        >>> b['ORY']['city_code_list']
        ['PAR']
        """
        data = {}
        next(f) # skipping first line

        for row in csv.reader(f, delimiter='^', quotechar='"'):
            # Comments and empty lines
            if not row or row[0].startswith('#'):
                continue

            key = row[_KEY]
            if key not in data:
                data[key] = d = cls._empty_value()
            else:
                prev_d = data[key]
                new_key = '{0}@{1}'.format(key, 1 + len(prev_d['__dup__']))
                data[new_key] = d = cls._empty_value()
                # Exchanging duplicata information
                d['__dup__'] = prev_d['__dup__'] | set([key])
                prev_d['__dup__'].add(new_key)

            for field, c in _FIELDS:
                d[field] = row[c]

            for field, splitter in _SPLITS:
                d[field] = splitter(d[field])

        return data

    def __iter__(self):
        """Returns iterator of all keys in the base.

        :returns: the iterator of all keys

        >>> b = NeoBase()
        >>> sorted(b)
        ['AAA', 'AAA@1', 'AAB', 'AAC', 'AAC@1', ...
        """
        return iter(self._data)

    def __contains__(self, key):
        """Test if a key is in the base.

        :param key: the key of to be tested
        :returns:   a boolean

        >>> b = NeoBase()
        >>> 'AN' in b
        False
        >>> 'AGN' in b
        True
        """
        return key in self._data

    def __nonzero__(self):
        """Testing structure emptiness.

        :returns: a boolean

        >>> b = NeoBase()
        >>> if b:
        ...     print('not empty')
        not empty
        """
        return bool(self._data)

    def __len__(self):
        """Testing structure size.

        :returns: a integer

        >>> b = NeoBase()
        >>> 18000 < len(b) < 20000
        True
        """
        return len(self._data)

    def keys(self):
        """Returns iterator of all keys in the base.

        :returns: the iterator of all keys

        >>> b = NeoBase()
        >>> sorted(b.keys())
        ['AAA', 'AAA@1', 'AAB', 'AAC', 'AAC@1', ...
        """
        return iter(self)

    def set(self, key, **data):
        """Set information.

        >>> b = NeoBase()
        >>> b.get('ORY', 'name')
        'Paris Orly Airport'
        >>> b.set('ORY', name='test')
        >>> b.get('ORY', 'name')
        'test'
        >>> b.set('Wow!', name='test')
        >>> b.get('Wow!', 'name')
        'test'
        """
        if key not in self:
            self._data[key] = self._empty_value()
        self._data[key].update(data)

    def get(self, key, field=None, default=_sentinel):
        """Get data from structure.

        >>> b = NeoBase()
        >>> b.get('OR', 'city_code_list', default=None)
        >>> b.get('ORY', 'city_code_list')
        ['PAR']
        """
        try:
            d = self._data[key]
        except KeyError:
            # Unless default is set, we raise an Exception
            if default is _sentinel:
                raise KeyError("Key not found: {0}".format(key))
            return default

        if field is None:
            return d # we return the whole dictionary

        try:
            res = d[field]
        except KeyError:
            raise KeyError("Field '{0}' (for key '{1}') not in {2}".format(
                field, key, list(d)))
        else:
            return res

    def get_location(self, key, default=_sentinel):
        """Get None or the geocode.

        >>> b = NeoBase()
        >>> b.get_location('ORY')
        LatLng(lat=48.72..., lng=2.35...)
        """
        if key not in self:
            # Unless default is set, we raise an Exception
            if default is _sentinel:
                raise KeyError("Key not found: {0}".format(key))
            return default

        try:
            loc = LatLng(float(self.get(key, 'lat')),
                         float(self.get(key, 'lng')))

        except (ValueError, TypeError, KeyError):
            # Decode geocode, if error, returns None
            # TypeError : input type is not a string, probably None
            # ValueError: could not convert to float
            # KeyError  : could not find lat or lng 'fields'
            return None
        else:
            return loc

    @staticmethod
    def distance_between_locations(l0, l1):
        """Great circle distance

        :param l0: the LatLng tuple of the first location
        :param l1: the LatLng tuple of the second location
        :returns:  the distance in kilometers

        >>> NeoBase.distance_between_locations((48.84, 2.367), (43.70, 7.26)) # Paris -> Nice
        683.85...

        Case of unknown location.

        >>> NeoBase.distance_between_locations(None, (43.70, 7.26)) # returns None
        """
        if l0 is None or l1 is None:
            return None

        l0_lat = l0[0] / 180 * pi
        l0_lng = l0[1] / 180 * pi
        l1_lat = l1[0] / 180 * pi
        l1_lng = l1[1] / 180 * pi

        # Haversine formula (6371 is Earth radius)
        return 2 * 6371.0 * asin(sqrt(
            sin(0.5 * (l0_lat - l1_lat)) ** 2 +
            sin(0.5 * (l0_lng - l1_lng)) ** 2 *
            cos(l0_lat) * cos(l1_lat)
        ))

    def distance(self, key_0, key_1):
        """Compute distance between two elements.

        This is just a wrapper between the original haversine
        function, but it is probably the most used feature :)

        :param key_0: the first key
        :param key_1: the second key
        :returns:    the distance (km)

        >>> b = NeoBase()
        >>> b.distance('ORY', 'CDG')
        34.87...
        """
        return self.distance_between_locations(self.get_location(key_0),
                                               self.get_location(key_1))

    def _build_distances(self, lat_lng_ref, keys):
        """
        Compute the iterable of (dist, keys) of a reference
        lat_lng and a list of keys. Keys which have not valid
        geocodes will not appear in the results.

        >>> b = NeoBase()
        >>> list(b._build_distances((0,0), ['ORY', 'CDG']))
        [(5422.74..., 'ORY'), (5455.45..., 'CDG')]
        """
        if lat_lng_ref is None:
            return

        for key in keys:
            if key in self:
                lat_lng = self.get_location(key)
                if lat_lng is not None:
                    yield self.distance_between_locations(lat_lng_ref, lat_lng), key

    def find_near_location(self, lat_lng, radius=_DEFAULT_RADIUS, from_keys=None):
        """
        Returns a list of nearby keys from a location (given
        latidude and longitude), and a radius for the search.
        Note that the haversine function, which compute distance
        at the surface of a sphere, here returns kilometers,
        so the radius should be in kms.

        :param lat_lng: the lat_lng of the location
        :param radius:  the radius of the search (kilometers)
        :param from_keys: if None, it takes all keys in consideration, else takes from_keys \
            iterable of keys to perform search.
        :returns:       an iterable of keys (like ['ORY', 'CDG'])

        >>> b = NeoBase()
        >>> # Paris, airports <= 50km
        >>> [b.get(k, 'iata_code') for d, k in sorted(b.find_near_location((48.84, 2.367), 10))]
        ['PAR', 'XGB', 'XHP', 'XPG', 'XEX', 'XTT', 'QAF', 'JDP', 'XBT', 'QNL', 'QBH', 'QFC', 'QEV']
        """
        if from_keys is None:
            from_keys = iter(self)

        for dist, key in self._build_distances(lat_lng, from_keys):
            if dist <= radius:
                yield dist, key

    def find_near(self, key, radius=_DEFAULT_RADIUS, from_keys=None):
        """
        Same as find_near_location, except the location is given
        not by a lat/lng, but with its key, like ORY or SFO.
        We just look up in the base to retrieve lat/lng, and
        call find_near_location.

        :param key:     the key of the location
        :param radius:  the radius of the search (kilometers)
        :param from_keys: if None, it takes all keys in consideration, else takes from_keys \
            iterable of keys to perform search.
        :returns:       a list of keys (like ['ORY', 'CDG'])

        >>> b = NeoBase()
        >>> sorted(b.find_near('ORY', 10)) # Orly, por <= 10km
        [(0.0, 'ORY'), (6.94..., 'XJY'), (9.96..., 'QFC')]
        """
        if from_keys is None:
            from_keys = iter(self)

        if key not in self:
            return

        for dist, key in self.find_near_location(self.get_location(key),
                                                 radius=radius,
                                                 from_keys=from_keys):
            yield dist, key

    def find_closest_from_location(self, lat_lng, N=1, from_keys=None):
        """
        Concept close to find_near_location, but here we do not
        look for the keys radius-close to a location,
        we look for the closest key from this location, given by
        latitude/longitude.

        Note that a similar implementation is done in
        the LocalHelper, to find efficiently N closest location
        in a graph, from a location (using heaps).

        :param lat_lng:   the lat_lng of the location
        :param N:         the N closest results wanted
        :param from_keys: if None, it takes all keys in consideration, else takes from_keys \
            iterable of keys to perform find_closest_from_location. This is useful to combine \
            searches
        :returns:   one key (like 'SFO'), or a list if approximate is not None

        >>> b = NeoBase()
        >>> list(b.find_closest_from_location((43.70, 7.26))) # Nice
        [(0.60..., 'NCE@1')]
        >>> list(b.find_closest_from_location((43.70, 7.26), N=3)) # Nice
        [(0.60..., 'NCE@1'), (5.82..., 'NCE'), (5.89..., 'XBM')]
        """
        if from_keys is None:
            from_keys = iter(self)

        iterable = self._build_distances(lat_lng, from_keys)

        for dist, key in heapq.nsmallest(N, iterable):
            yield dist, key

    def find_with(self, conditions, from_keys=None, reverse=False):
        """Get iterator of all keys with particular field.

        For example, if you want to know all airports in Paris.

        :param conditions: a list of (field, value) conditions
        :param reverse:    we look keys where the field is *not* the particular value
        :returns:          an iterator of matching keys

        Testing several conditions.

        >>> b = NeoBase()
        >>> c0 = [('city_code_list', ['PAR'])]
        >>> c1 = [('location_type', ['H'])]
        >>> len(list(b.find_with(c0)))
        17
        >>> len(list(b.find_with(c1)))
        146
        >>> len(list(b.find_with(c0 + c1)))
        2
        """
        if from_keys is None:
            from_keys = iter(self)

        if not reverse:
            match = lambda a, b: a == b
        else:
            match = lambda a, b: a != b

        for key in from_keys:
            if key in self:
                if all(match(self.get(key, f), v) for f, v in conditions):
                    yield key

def main():
    import argparse
    from pprint import pprint

    parser = argparse.ArgumentParser()
    parser.add_argument('keys', nargs='+', help="List of IATA codes")
    args = parser.parse_args()
    b = NeoBase()

    for key in args.keys:
        if key in b:
            pprint(b.get(key))
        else:
            print('{0} not in data.'.format(key))

if __name__ == '__main__':
    main()

<|endoftext|>