<commit_before># This file is part of beets.
# Copyright 2010, Adrian Sampson.
#
# 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.

"""Facilities for automatically determining files' correct metadata.
"""

import os
from collections import defaultdict
from beets.autotag import mb
import re
from munkres import Munkres
from beets import library, mediafile
import logging

# Try 5 releases. In the future, this should be more dynamic: let the
# probability of continuing to the next release be inversely
# proportional to how good our current best is and how long we've
# already taken.
MAX_CANDIDATES = 5

# Distance parameters.
# Text distance weights: proportions on the normalized intuitive edit
# distance.
ARTIST_WEIGHT = 3.0
ALBUM_WEIGHT = 3.0
# The weight of the entire distance calculated for a given track.
TRACK_WEIGHT = 1.0
# These distances are components of the track distance (that is, they
# compete against each other but not ARTIST_WEIGHT and ALBUM_WEIGHT;
# the overall TRACK_WEIGHT does that).
TRACK_TITLE_WEIGHT = 3.0
# Added when the indices of tracks don't match.
TRACK_INDEX_WEIGHT = 0.0
# Track length weights: no penalty before GRACE, maximum (WEIGHT)
# penalty at GRACE+MAX discrepancy.
TRACK_LENGTH_GRACE = 15
TRACK_LENGTH_MAX = 30
TRACK_LENGTH_WEIGHT = 1.0
# MusicBrainz track ID matches.
TRACK_ID_WEIGHT = 5.0

# Recommendation constants.
RECOMMEND_STRONG = 'RECOMMEND_STRONG'
RECOMMEND_MEDIUM = 'RECOMMEND_MEDIUM'
RECOMMEND_NONE = 'RECOMMEND_NONE'
# Thresholds for recommendations.
STRONG_REC_THRESH = 0.03
MEDIUM_REC_THRESH = 0.2
REC_GAP_THRESH = 0.25

# Autotagging exceptions.
class AutotagError(Exception):
    pass
class InsufficientMetadataError(AutotagError):
    pass

# Global logger.
log = logging.getLogger('beets')

def _first_n(it, n):
    """Takes an iterator and returns another iterator, truncated to
    yield only the first n elements.
    """
    for i, v in enumerate(it):
        if i >= n:
            break
        yield v

def _sorted_walk(path):
    """Like os.walk, but yields things in sorted, breadth-first
    order.
    """
    # Get all the directories and files at this level.
    dirs = []
    files = []
    for base in os.listdir(path):
        cur = os.path.join(path, base)
        if os.path.isdir(cur):
            dirs.append(base)
        else:
            files.append(base)

    # Sort lists and yield the current level.
    dirs.sort()
    files.sort()
    yield (path, dirs, files)

    # Recurse into directories.
    for base in dirs:
        cur = os.path.join(path, base)
        # yield from _sorted_walk(cur)
        for res in _sorted_walk(cur):
            yield res

def albums_in_dir(path):
    """Recursively searches the given directory and returns an iterable
    of lists of items where each list is probably an album.
    Specifically, any folder containing any media files is an album.
    """
    path = library._unicode_path(path)
    for root, dirs, files in _sorted_walk(path):
        # Get a list of items in the directory.
        items = []
        for filename in files:
            try:
                i = library.Item.from_path(os.path.join(root, filename))
            except mediafile.FileTypeError:
                pass
            except mediafile.UnreadableFileError:
                log.warn('unreadable file: ' + filename)
            else:
                items.append(i)

        # If it's nonempty, yield it.
        if items:
            yield items

def _ie_dist(str1, str2):
    """Gives an "intuitive" edit distance between two strings. This is
    an edit distance, normalized by the string length, ignoring case
    and nonalphanumeric characters.
    """
    str1 = re.sub(r'[^a-z0-9]', '', str1.lower())
    str2 = re.sub(r'[^a-z0-9]', '', str2.lower())

    # Avoid divide-by-zero. Two emptry strings are identical.
    if not str1 and not str2:
        return 0

    # Here's a nice DP edit distance implementation from Wikibooks:
    # http://en.wikibooks.org/wiki/Algorithm_implementation/Strings/
    # Levenshtein_distance#Python
    # This should probably be written in a C module.
    def levenshtein(s1, s2):
        if len(s1) < len(s2):
            return levenshtein(s2, s1)
        if not s1:
            return len(s2)

        previous_row = xrange(len(s2) + 1)
        for i, c1 in enumerate(s1):
            current_row = [i + 1]
            for j, c2 in enumerate(s2):
                insertions = previous_row[j + 1] + 1
                deletions = current_row[j] + 1
                substitutions = previous_row[j] + (c1 != c2)
                current_row.append(min(insertions, deletions, substitutions))
            previous_row = current_row

        return previous_row[-1]

    return levenshtein(str1, str2) / float(max(len(str1), len(str2)))

def current_metadata(items):
    """Returns the most likely artist and album for a set of Items.
    Each is determined by tag reflected by the plurality of the Items.
    """
    # The tags we'll try to determine.
    keys = 'artist', 'album'

    # Make dictionaries in which to count the freqencies of different
    # artist and album tags. We'll use this to find the most likely
    # artist and album. Defaultdicts let the frequency default to zero.
    freqs = {}
    for key in keys:
        freqs[key] = defaultdict(int)

    # Count the frequencies.
    for item in items:
        for key in keys:
            value = getattr(item, key)
            if value: # Don't count empty tags.
                freqs[key][value] += 1

    # Find max-frequency tags.
    likelies = {}
    for key in keys:
        max_freq = 0
        likelies[key] = None
        for tag, freq in freqs[key].items():
            if freq > max_freq:
                max_freq = freq
                likelies[key] = tag

    return likelies['artist'], likelies['album']

def order_items(items, trackinfo):
    """Orders the items based on how they match some canonical track
    information. This always produces a result if the numbers of tracks
    match.
    """
    # Construct the cost matrix.
    costs = []
    for cur_item in items:
        row = []
        for i, canon_item in enumerate(trackinfo):
            row.append(track_distance(cur_item, canon_item, i+1))
        costs.append(row)

    # Find a minimum-cost bipartite matching.
    matching = Munkres().compute(costs)

    # Order items based on the matching.
    ordered_items = [None]*len(items)
    for cur_idx, canon_idx in matching:
        ordered_items[canon_idx] = items[cur_idx]
    return ordered_items

def track_distance(item, track_data, track_index=None):
    """Determines the significance of a track metadata change. Returns
    a float in [0.0,1.0]. `track_index` is the track number of the
    `track_data` metadata set. If `track_index` is provided and
    item.track is set, then these indices are used as a component of
    the distance calculation.
    """
    # Distance and normalization accumulators.
    dist, dist_max = 0.0, 0.0

    # Check track length.
    if 'length' not in track_data:
        # If there's no length to check, assume the worst.
        dist += TRACK_LENGTH_WEIGHT
    else:
        diff = abs(item.length - track_data['length'])
        diff = max(diff - TRACK_LENGTH_GRACE, 0.0)
        diff = min(diff, TRACK_LENGTH_MAX)
        dist += (diff / TRACK_LENGTH_MAX) * TRACK_LENGTH_WEIGHT
    dist_max += TRACK_LENGTH_WEIGHT

    # Track title.
    dist += _ie_dist(item.title, track_data['title']) * TRACK_TITLE_WEIGHT
    dist_max += TRACK_TITLE_WEIGHT

    # Track index.
    if track_index and item.track:
        if track_index != item.track:
            dist += TRACK_INDEX_WEIGHT
        dist_max += TRACK_INDEX_WEIGHT

    # MusicBrainz track ID.
    if item.mb_trackid:
        if item.mb_trackid != track_data['id'].rsplit('/', 1)[1]:
            dist += TRACK_ID_WEIGHT
        dist_max += TRACK_ID_WEIGHT

    return dist / dist_max

def distance(items, info):
    """Determines how "significant" an album metadata change would be.
    Returns a float in [0.0,1.0]. The list of items must be ordered.
    """
    cur_artist, cur_album = current_metadata(items)
    cur_artist = cur_artist or ''
    cur_album = cur_album or ''

    # These accumulate the possible distance components. The final
    # distance will be dist/dist_max.
    dist = 0.0
    dist_max = 0.0

    # Artist/album metadata.
    dist += _ie_dist(cur_artist, info['artist']) * ARTIST_WEIGHT
    dist_max += ARTIST_WEIGHT
    dist += _ie_dist(cur_album,  info['album']) * ALBUM_WEIGHT
    dist_max += ALBUM_WEIGHT

    # Track distances.
    for i, (item, track_data) in enumerate(zip(items, info['tracks'])):
        dist += track_distance(item, track_data, i+1) * TRACK_WEIGHT
        dist_max += TRACK_WEIGHT

    # Normalize distance, avoiding divide-by-zero.
    if dist_max == 0.0:
        return 0.0
    else:
        return dist/dist_max

def apply_metadata(items, info):
    """Set the items' metadata to match the data given in info. The
    list of items must be ordered.
    """
    # Global MusicBrainz IDs (album and artist).
    mb_albumid = info['album_id'].rsplit('/', 1)[1]
    mb_artistid = info['artist_id'].rsplit('/', 1)[1]

    for index, (item, track_data) in enumerate(zip(items,  info['tracks'])):
        # Album, artist, track count.
        item.artist = info['artist']
        item.album = info['album']
        item.tracktotal = len(items)

        # Release date.
        if 'year' in info:
            item.year = info['year']
        if 'month' in info:
            item.month = info['month']
        if 'day' in info:
            item.day = info['day']

        # Title and track index.
        item.title = track_data['title']
        item.track = index + 1

        # MusicBrainz track ID.
        item.mb_trackid = track_data['id'].rsplit('/', 1)[1]
        # Album and artist IDs.
        item.mb_albumid = mb_albumid
        item.mb_artistid = mb_artistid

def match_by_id(items):
    """If the items are tagged with a MusicBrainz album ID, returns an
    info dict for the corresponding album. Otherwise, returns None.
    """
    # Is there a consensus on the MB album ID?
    albumids = [item.mb_albumid for item in items if item.mb_albumid]
    if not albumids:
        return None

    # If all album IDs are equal, look up the album.
    if bool(reduce(lambda x,y: x if x==y else (), albumids)):
        albumid = albumids[0]
        return mb.album_for_id(albumid)
    else:
        return None

    #fixme In the future, at the expense of performance, we could use
    # other IDs (i.e., track and artist) in case the album tag isn't
    # present, but that event seems very unlikely.

def tag_album(items, search_artist=None, search_album=None):
    """Bundles together the functionality used to infer tags for a
    set of items comprised by an album. Returns everything relevant:
        - The current artist.
        - The current album.
        - A list of (distance, items, info) tuples where info is a
          dictionary containing the inferred tags and items is a
          reordered version of the input items list. The candidates are
          sorted by distance (i.e., best match first).
        - A recommendation, one of RECOMMEND_STRONG, RECOMMEND_MEDIUM,
          or RECOMMEND_NONE; indicating that the first candidate is
          very likely, it is somewhat likely, or no conclusion could
          be reached.
    If search_artist and search_album are provided, then they are used
    as search terms in place of the current metadata.
    May raise an AutotagError if existing metadata is insufficient.
    """
    # Get current metadata.
    cur_artist, cur_album = current_metadata(items)

    # The output list of result tuples:
    dist_ordered_cands = []

    # Try to find album indicated by MusicBrainz iDs.
    id_info = match_by_id(items)
    if id_info:
        ordered = order_items(items, id_info['tracks'])
        if ordered:
            dist = distance(ordered, id_info)
            dist_ordered_cands.append((dist, ordered, id_info))
            # If we have a very good MBID match, don't bother
            # continuing.
            if dist < STRONG_REC_THRESH:
                return (cur_artist, cur_album, dist_ordered_cands,
                        RECOMMEND_STRONG)
            # Otherwise, this match will compete against metadata-based
            # matches.

    # Search terms.
    if not (search_artist and search_album):
        # No explicit search terms -- use current metadata.
        search_artist, search_album = cur_artist, cur_album

    # Get candidate metadata.
    if not search_artist or not search_album:
        raise InsufficientMetadataError()
    candidates = mb.match_album(search_artist, search_album, len(items))

    # Get the distance to each candidate.
    for info in _first_n(candidates, MAX_CANDIDATES):
        # Make sure the album has the correct number of tracks.
        if len(items) != len(info['tracks']):
            continue

        # Put items in order.
        ordered = order_items(items, info['tracks'])
        if not ordered:
            continue

        # Get the change distance.
        dist = distance(ordered, info)

        dist_ordered_cands.append((dist, ordered, info))

    # Sort by distance.
    dist_ordered_cands.sort()

    # Make a recommendation.
    if not dist_ordered_cands:
        # No candidates: no recommendation.
        rec = RECOMMEND_NONE
    else:
        min_dist = dist_ordered_cands[0][0]
        if min_dist < STRONG_REC_THRESH:
            # Strong recommendation level.
            rec = RECOMMEND_STRONG
        elif len(dist_ordered_cands) == 1:
            # Only a single candidate. Medium recommendation.
            rec = RECOMMEND_MEDIUM
        elif min_dist <= MEDIUM_REC_THRESH:
            # Medium recommendation level.
            rec = RECOMMEND_MEDIUM
        elif dist_ordered_cands[1][0] - min_dist >= REC_GAP_THRESH:
            # Gap between first two candidates is large.
            rec = RECOMMEND_MEDIUM
        else:
            # No conclusion.
            rec = RECOMMEND_NONE

    return cur_artist, cur_album, dist_ordered_cands, rec
<commit_msg>fixed situation where the MusicBrainz match duplicated a metadata match<commit_after># This file is part of beets.
# Copyright 2010, Adrian Sampson.
#
# 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.

"""Facilities for automatically determining files' correct metadata.
"""

import os
from collections import defaultdict
from beets.autotag import mb
import re
from munkres import Munkres
from beets import library, mediafile
import logging

# Try 5 releases. In the future, this should be more dynamic: let the
# probability of continuing to the next release be inversely
# proportional to how good our current best is and how long we've
# already taken.
MAX_CANDIDATES = 5

# Distance parameters.
# Text distance weights: proportions on the normalized intuitive edit
# distance.
ARTIST_WEIGHT = 3.0
ALBUM_WEIGHT = 3.0
# The weight of the entire distance calculated for a given track.
TRACK_WEIGHT = 1.0
# These distances are components of the track distance (that is, they
# compete against each other but not ARTIST_WEIGHT and ALBUM_WEIGHT;
# the overall TRACK_WEIGHT does that).
TRACK_TITLE_WEIGHT = 3.0
# Added when the indices of tracks don't match.
TRACK_INDEX_WEIGHT = 0.0
# Track length weights: no penalty before GRACE, maximum (WEIGHT)
# penalty at GRACE+MAX discrepancy.
TRACK_LENGTH_GRACE = 15
TRACK_LENGTH_MAX = 30
TRACK_LENGTH_WEIGHT = 1.0
# MusicBrainz track ID matches.
TRACK_ID_WEIGHT = 5.0

# Recommendation constants.
RECOMMEND_STRONG = 'RECOMMEND_STRONG'
RECOMMEND_MEDIUM = 'RECOMMEND_MEDIUM'
RECOMMEND_NONE = 'RECOMMEND_NONE'
# Thresholds for recommendations.
STRONG_REC_THRESH = 0.03
MEDIUM_REC_THRESH = 0.2
REC_GAP_THRESH = 0.25

# Autotagging exceptions.
class AutotagError(Exception):
    pass
class InsufficientMetadataError(AutotagError):
    pass

# Global logger.
log = logging.getLogger('beets')

def _first_n(it, n):
    """Takes an iterator and returns another iterator, truncated to
    yield only the first n elements.
    """
    for i, v in enumerate(it):
        if i >= n:
            break
        yield v

def _sorted_walk(path):
    """Like os.walk, but yields things in sorted, breadth-first
    order.
    """
    # Get all the directories and files at this level.
    dirs = []
    files = []
    for base in os.listdir(path):
        cur = os.path.join(path, base)
        if os.path.isdir(cur):
            dirs.append(base)
        else:
            files.append(base)

    # Sort lists and yield the current level.
    dirs.sort()
    files.sort()
    yield (path, dirs, files)

    # Recurse into directories.
    for base in dirs:
        cur = os.path.join(path, base)
        # yield from _sorted_walk(cur)
        for res in _sorted_walk(cur):
            yield res

def albums_in_dir(path):
    """Recursively searches the given directory and returns an iterable
    of lists of items where each list is probably an album.
    Specifically, any folder containing any media files is an album.
    """
    path = library._unicode_path(path)
    for root, dirs, files in _sorted_walk(path):
        # Get a list of items in the directory.
        items = []
        for filename in files:
            try:
                i = library.Item.from_path(os.path.join(root, filename))
            except mediafile.FileTypeError:
                pass
            except mediafile.UnreadableFileError:
                log.warn('unreadable file: ' + filename)
            else:
                items.append(i)

        # If it's nonempty, yield it.
        if items:
            yield items

def _ie_dist(str1, str2):
    """Gives an "intuitive" edit distance between two strings. This is
    an edit distance, normalized by the string length, ignoring case
    and nonalphanumeric characters.
    """
    str1 = re.sub(r'[^a-z0-9]', '', str1.lower())
    str2 = re.sub(r'[^a-z0-9]', '', str2.lower())

    # Avoid divide-by-zero. Two emptry strings are identical.
    if not str1 and not str2:
        return 0

    # Here's a nice DP edit distance implementation from Wikibooks:
    # http://en.wikibooks.org/wiki/Algorithm_implementation/Strings/
    # Levenshtein_distance#Python
    # This should probably be written in a C module.
    def levenshtein(s1, s2):
        if len(s1) < len(s2):
            return levenshtein(s2, s1)
        if not s1:
            return len(s2)

        previous_row = xrange(len(s2) + 1)
        for i, c1 in enumerate(s1):
            current_row = [i + 1]
            for j, c2 in enumerate(s2):
                insertions = previous_row[j + 1] + 1
                deletions = current_row[j] + 1
                substitutions = previous_row[j] + (c1 != c2)
                current_row.append(min(insertions, deletions, substitutions))
            previous_row = current_row

        return previous_row[-1]

    return levenshtein(str1, str2) / float(max(len(str1), len(str2)))

def current_metadata(items):
    """Returns the most likely artist and album for a set of Items.
    Each is determined by tag reflected by the plurality of the Items.
    """
    # The tags we'll try to determine.
    keys = 'artist', 'album'

    # Make dictionaries in which to count the freqencies of different
    # artist and album tags. We'll use this to find the most likely
    # artist and album. Defaultdicts let the frequency default to zero.
    freqs = {}
    for key in keys:
        freqs[key] = defaultdict(int)

    # Count the frequencies.
    for item in items:
        for key in keys:
            value = getattr(item, key)
            if value: # Don't count empty tags.
                freqs[key][value] += 1

    # Find max-frequency tags.
    likelies = {}
    for key in keys:
        max_freq = 0
        likelies[key] = None
        for tag, freq in freqs[key].items():
            if freq > max_freq:
                max_freq = freq
                likelies[key] = tag

    return likelies['artist'], likelies['album']

def order_items(items, trackinfo):
    """Orders the items based on how they match some canonical track
    information. This always produces a result if the numbers of tracks
    match.
    """
    # Construct the cost matrix.
    costs = []
    for cur_item in items:
        row = []
        for i, canon_item in enumerate(trackinfo):
            row.append(track_distance(cur_item, canon_item, i+1))
        costs.append(row)

    # Find a minimum-cost bipartite matching.
    matching = Munkres().compute(costs)

    # Order items based on the matching.
    ordered_items = [None]*len(items)
    for cur_idx, canon_idx in matching:
        ordered_items[canon_idx] = items[cur_idx]
    return ordered_items

def track_distance(item, track_data, track_index=None):
    """Determines the significance of a track metadata change. Returns
    a float in [0.0,1.0]. `track_index` is the track number of the
    `track_data` metadata set. If `track_index` is provided and
    item.track is set, then these indices are used as a component of
    the distance calculation.
    """
    # Distance and normalization accumulators.
    dist, dist_max = 0.0, 0.0

    # Check track length.
    if 'length' not in track_data:
        # If there's no length to check, assume the worst.
        dist += TRACK_LENGTH_WEIGHT
    else:
        diff = abs(item.length - track_data['length'])
        diff = max(diff - TRACK_LENGTH_GRACE, 0.0)
        diff = min(diff, TRACK_LENGTH_MAX)
        dist += (diff / TRACK_LENGTH_MAX) * TRACK_LENGTH_WEIGHT
    dist_max += TRACK_LENGTH_WEIGHT

    # Track title.
    dist += _ie_dist(item.title, track_data['title']) * TRACK_TITLE_WEIGHT
    dist_max += TRACK_TITLE_WEIGHT

    # Track index.
    if track_index and item.track:
        if track_index != item.track:
            dist += TRACK_INDEX_WEIGHT
        dist_max += TRACK_INDEX_WEIGHT

    # MusicBrainz track ID.
    if item.mb_trackid:
        if item.mb_trackid != track_data['id'].rsplit('/', 1)[1]:
            dist += TRACK_ID_WEIGHT
        dist_max += TRACK_ID_WEIGHT

    return dist / dist_max

def distance(items, info):
    """Determines how "significant" an album metadata change would be.
    Returns a float in [0.0,1.0]. The list of items must be ordered.
    """
    cur_artist, cur_album = current_metadata(items)
    cur_artist = cur_artist or ''
    cur_album = cur_album or ''

    # These accumulate the possible distance components. The final
    # distance will be dist/dist_max.
    dist = 0.0
    dist_max = 0.0

    # Artist/album metadata.
    dist += _ie_dist(cur_artist, info['artist']) * ARTIST_WEIGHT
    dist_max += ARTIST_WEIGHT
    dist += _ie_dist(cur_album,  info['album']) * ALBUM_WEIGHT
    dist_max += ALBUM_WEIGHT

    # Track distances.
    for i, (item, track_data) in enumerate(zip(items, info['tracks'])):
        dist += track_distance(item, track_data, i+1) * TRACK_WEIGHT
        dist_max += TRACK_WEIGHT

    # Normalize distance, avoiding divide-by-zero.
    if dist_max == 0.0:
        return 0.0
    else:
        return dist/dist_max

def apply_metadata(items, info):
    """Set the items' metadata to match the data given in info. The
    list of items must be ordered.
    """
    # Global MusicBrainz IDs (album and artist).
    mb_albumid = info['album_id'].rsplit('/', 1)[1]
    mb_artistid = info['artist_id'].rsplit('/', 1)[1]

    for index, (item, track_data) in enumerate(zip(items,  info['tracks'])):
        # Album, artist, track count.
        item.artist = info['artist']
        item.album = info['album']
        item.tracktotal = len(items)

        # Release date.
        if 'year' in info:
            item.year = info['year']
        if 'month' in info:
            item.month = info['month']
        if 'day' in info:
            item.day = info['day']

        # Title and track index.
        item.title = track_data['title']
        item.track = index + 1

        # MusicBrainz track ID.
        item.mb_trackid = track_data['id'].rsplit('/', 1)[1]
        # Album and artist IDs.
        item.mb_albumid = mb_albumid
        item.mb_artistid = mb_artistid

def match_by_id(items):
    """If the items are tagged with a MusicBrainz album ID, returns an
    info dict for the corresponding album. Otherwise, returns None.
    """
    # Is there a consensus on the MB album ID?
    albumids = [item.mb_albumid for item in items if item.mb_albumid]
    if not albumids:
        return None

    # If all album IDs are equal, look up the album.
    if bool(reduce(lambda x,y: x if x==y else (), albumids)):
        albumid = albumids[0]
        return mb.album_for_id(albumid)
    else:
        return None

    #fixme In the future, at the expense of performance, we could use
    # other IDs (i.e., track and artist) in case the album tag isn't
    # present, but that event seems very unlikely.

def tag_album(items, search_artist=None, search_album=None):
    """Bundles together the functionality used to infer tags for a
    set of items comprised by an album. Returns everything relevant:
        - The current artist.
        - The current album.
        - A list of (distance, items, info) tuples where info is a
          dictionary containing the inferred tags and items is a
          reordered version of the input items list. The candidates are
          sorted by distance (i.e., best match first).
        - A recommendation, one of RECOMMEND_STRONG, RECOMMEND_MEDIUM,
          or RECOMMEND_NONE; indicating that the first candidate is
          very likely, it is somewhat likely, or no conclusion could
          be reached.
    If search_artist and search_album are provided, then they are used
    as search terms in place of the current metadata.
    May raise an AutotagError if existing metadata is insufficient.
    """
    # Get current metadata.
    cur_artist, cur_album = current_metadata(items)

    # The output list of result tuples:
    dist_ordered_cands = []

    # Try to find album indicated by MusicBrainz IDs.
    id_info = match_by_id(items)
    if id_info:
        ordered = order_items(items, id_info['tracks'])
        if ordered:
            dist = distance(ordered, id_info)
            dist_ordered_cands.append((dist, ordered, id_info))
            # If we have a very good MBID match, don't bother
            # continuing.
            if dist < STRONG_REC_THRESH:
                return (cur_artist, cur_album, dist_ordered_cands,
                        RECOMMEND_STRONG)
            # Otherwise, this match will compete against metadata-based
            # matches.

    # Search terms.
    if not (search_artist and search_album):
        # No explicit search terms -- use current metadata.
        search_artist, search_album = cur_artist, cur_album

    # Get candidate metadata.
    if not search_artist or not search_album:
        raise InsufficientMetadataError()
    candidates = mb.match_album(search_artist, search_album, len(items))

    # Get the distance to each candidate.
    for info in _first_n(candidates, MAX_CANDIDATES):
        # Don't duplicate the MusicBrainz match.
        if id_info and id_info['album_id'] == info['album_id']:
            continue

        # Make sure the album has the correct number of tracks.
        if len(items) != len(info['tracks']):
            continue

        # Put items in order.
        ordered = order_items(items, info['tracks'])
        if not ordered:
            continue

        # Get the change distance.
        dist = distance(ordered, info)

        dist_ordered_cands.append((dist, ordered, info))

    # Sort by distance.
    dist_ordered_cands.sort()

    # Make a recommendation.
    if not dist_ordered_cands:
        # No candidates: no recommendation.
        rec = RECOMMEND_NONE
    else:
        min_dist = dist_ordered_cands[0][0]
        if min_dist < STRONG_REC_THRESH:
            # Strong recommendation level.
            rec = RECOMMEND_STRONG
        elif len(dist_ordered_cands) == 1:
            # Only a single candidate. Medium recommendation.
            rec = RECOMMEND_MEDIUM
        elif min_dist <= MEDIUM_REC_THRESH:
            # Medium recommendation level.
            rec = RECOMMEND_MEDIUM
        elif dist_ordered_cands[1][0] - min_dist >= REC_GAP_THRESH:
            # Gap between first two candidates is large.
            rec = RECOMMEND_MEDIUM
        else:
            # No conclusion.
            rec = RECOMMEND_NONE

    return cur_artist, cur_album, dist_ordered_cands, rec
<|endoftext|>


<commit_before>from __future__ import division

import logging
import sys
import os
import math
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
import xml.etree.ElementTree as ET
from collections import OrderedDict

import numpy as NP
from scipy.constants import mu_0
from scipy.io import savemat
from scipy.interpolate import CubicSpline

from ..util.signal import nextpow2
from ..mag.iaga2002 import parse
from ..util.nan import nan_interp
from ..util.date import toJ2000

def calc_e_3d(Bx,
              By,
              Zxx_function,
              Zxy_function,
              Zyx_function,
              Zyy_function,
              interval):
    """
    Calculate the electric field induced by the magnetic field (*Bx*,
    eastward oriented magnetic field intensity, and *By*, northward
    oriented magnetic field intensity). *Zxx_function*,
    *Z_xy_function*, *Z_yx_function*, and *Z_yy_function* are the
    surface impedance functions (input is angular frequency [rad] and
    output is impedance [Ohms]). Return the tuple *Ex*, *Ey* (same
    orientation as the magnetic field).

    Units:
       (input)  Bx, By: [T]
       (input)  Zxx_function: [rad] -> [Ohm]
       (input)  Zxy_function: [rad] -> [Ohm]
       (input)  Zyx_function: [rad] -> [Ohm]
       (input)  Zyy_function: [rad] -> [Ohm]
       (output) Ex, Ey: [V/m]
    """
    # This function is the reimplementation of Prof. Zhu's matlab code
    # which is based on the algorithm detailed in the NERC Application
    # Guide, Dec. 2013. Note that the NERC Guide contains a sign error
    # that has been corrected in the code below.

    # function global variables
    lp = 60  # length at beginning and end of time series to fit
             # detrend parameters
    p  = 60  # length at beginning and end of time series for
             # frequency domain window
    # get the FFT size
    assert len(Bx) == len(By)
    N = len(Bx)
    Nfft = nextpow2(N)
    # zero-mean and remove linear trend from the magnetic time series
    # data
    ax = NP.mean(Bx[:lp])
    bx = NP.mean(Bx[-lp-1:])
    N_range = NP.arange(1, N+1)
    cx = ax * (N - N_range) / N + bx * N_range / N
    Bx0 = Bx - cx

    ay = NP.mean(By[:lp])
    by = NP.mean(By[-lp-1:])
    cy = ay * (N - N_range) / N + by * N_range / N
    By0 = By - cy
    # window the magnetic data time series
    wp = NP.hstack((0.5 * (1 - NP.cos(2 * math.pi * NP.arange(0, p/2) / p)),
                    NP.ones(N - p),
                    0.5 * (1 - NP.cos(2 * math.pi * NP.arange(p/2 - 1, -1, -1) / p))))
    Bx1 = Bx0 * wp
    By1 = By0 * wp
    # compute FFT
    Sbx = NP.fft.fft(Bx1, n=Nfft) / N
    Sby = NP.fft.fft(By1, n=Nfft) / N
    freq = NP.arange(0, Nfft) / Nfft / interval
    omega = 2 * math.pi * freq

    Zxx_positive = Zxx_function(omega[1:])
    Zxx = NP.hstack(([0], Zxx_positive))

    Zxx2 = NP.hstack((Zxx[:(int(Nfft/2)+1)],
                      NP.conj(Zxx[1:int(Nfft/2)])[::-1]))

    Zxy_positive = Zxy_function(omega[1:])
    Zxy = NP.hstack(([0], Zxy_positive))

    Zxy2 = NP.hstack((Zxy[:(int(Nfft/2)+1)],
                      NP.conj(Zxy[1:int(Nfft/2)])[::-1]))

    Zyx_positive = Zyx_function(omega[1:])
    Zyx = NP.hstack(([0], Zyx_positive))

    Zyx2 = NP.hstack((Zyx[:(int(Nfft/2)+1)],
                      NP.conj(Zyx[1:int(Nfft/2)])[::-1]))

    Zyy_positive = Zyy_function(omega[1:])
    Zyy = NP.hstack(([0], Zyy_positive))

    Zyy2 = NP.hstack((Zyy[:(int(Nfft/2)+1)],
                      NP.conj(Zyy[1:int(Nfft/2)])[::-1]))

    Se_x = []
    Se_y = []
    for Zxx_i, Zxy_i, Zyx_i, Zyy_i, Sbx_i, Sby_i in zip(Zxx2, Zxy2, Zyx2, Zyy2, Sbx, Sby):
        Z = NP.array([[Zxx_i, Zxy_i], [Zyx_i, Zyy_i]])
        Sb_i = NP.array([[Sbx_i / mu_0], [Sby_i / mu_0]])
        Se_i = NP.dot(Z, Sb_i)
        Se_x.append(Se_i[0, 0])
        Se_y.append(Se_i[1, 0])

    Ex = NP.real(NP.fft.ifft(Se_x, Nfft) * N)
    Ey = NP.real(NP.fft.ifft(Se_y, Nfft) * N)

    return (Ex[:N],
            Ey[:N])

def parse_xml(xml_fname):
    """
    Parse the E-M transfer function file *xml_fname* returning a
    mapping between period keys (in [s]) and the 2x2 matrices
    containing the associated Zxx, Zxy, Zyx, and Zyy parameters (in
    [mV / km] / [nT]).

    These files are available at http://ds.iris.edu/spud/emtf.
    """
    tree = ET.parse(xml_fname)
    root = tree.getroot()

    data_list = root.findall('Data')
    assert len(data_list) == 1
    data = data_list[0]

    Z_map = OrderedDict()
    for period in data.findall('Period'):
        Z_list = period.findall('Z')
        assert len(Z_list) == 1
        Z = Z_list[0]
        values = []
        for value, name in zip(Z, ['Zxx', 'Zxy', 'Zyx', 'Zyy']):
            assert value.attrib['name'] == name
            values.append(complex(*map(float, value.text.split())))
        Z_array = NP.array(values)
        Z_array.shape = 2, 2
        Z_map[float(period.attrib['value'])] = Z_array
    return Z_map

def parse_xml_header(xml_fname):
    """
    Parse the E-M transfer function file *xml_fname* returning a
    mapping of site specific information, such as location and data
    quality.
    """
    header_map = {}

    tree = ET.parse(xml_fname)
    root = tree.getroot()

    copyright_list = root.findall('Copyright')
    assert len(copyright_list) == 1
    copyright_ = copyright_list[0]

    for copyright_node in copyright_:
        if copyright_node.tag == 'Acknowledgement':
            header_map['acknowledgement'] = copyright_node.text

    site_list = root.findall('Site')
    assert len(site_list) == 1
    site = site_list[0]

    for site_node in site:
        if site_node.tag == 'Id':
            header_map['id'] = site_node.text
        elif site_node.tag == 'Location':
            for child in site_node:
                if child.tag == 'Latitude':
                    header_map['lat'] = float(child.text)
                elif child.tag == 'Longitude':
                    header_map['lon'] = float(child.text)
                elif child.tag == 'Elevation':
                    assert child.attrib['units'] == 'meters'
                    header_map['el'] = float(child.text)
        elif site_node.tag == 'DataQualityNotes':
            for child in site_node:
                if child.tag == 'Rating':
                    header_map['rating'] = int(child.text)
                elif child.tag == 'GoodFromPeriod':
                    header_map['good_from'] = float(child.text)
                    header_map['good_to_mHz'] = 1/float(child.text) * 1e3
                elif child.tag == 'GoodToPeriod':
                    header_map['good_to'] = float(child.text)
                    header_map['good_from_mHz'] = 1/float(child.text) * 1e3
                elif child.tag == 'Comments':
                    header_map['data_quality_comments'] = child.text
        elif site_node.tag == 'DataQualityWarnings':
            for child in site_node:
                if child.tag == 'Flag':
                    header_map['data_quality_flag'] = int(child.text)
                elif child.tag == 'Comments':
                    header_map['data_quality_warning_comments'] = child.text
        elif site_node.tag == 'Acknowledgment':
            header_map['acknowledgment'] = site_node.text

    return header_map

class Zw_interpolator(object):
    def __init__(self, Z_map):
        """
        Construct a cubic-spline 3-D E-M transfer function interpolater
        using the information in *Z_map* returned from
        :func:`parse_xml` as the function samples.
        """
        self.Z_map = Z_map
        periods = Z_map.keys()
        self.f = NP.array([1/x for x in periods[::-1]])
        self.omega = 2 * math.pi * self.f
        self.Zxx_interp = CubicSpline(self.omega, [x[0, 0] for x in Z_map.values()[::-1]])
        self.Zxy_interp = CubicSpline(self.omega, [x[0, 1] for x in Z_map.values()[::-1]])
        self.Zyx_interp = CubicSpline(self.omega, [x[1, 0] for x in Z_map.values()[::-1]])
        self.Zyy_interp = CubicSpline(self.omega, [x[1, 1] for x in Z_map.values()[::-1]])
        self.key_map = {'xx': self.Zxx_interp,
                        'xy': self.Zxy_interp,
                        'yx': self.Zyx_interp,
                        'yy': self.Zyy_interp}

    def __call__(self, omega, key):
        """
        Return the interpolated value of the transfer function (*key* is
        xx, xy, yx, or yy) at angular frequency *omega*. The units of
        *omega* and the output are the same as *Z_map* ([rad] and
        [mV/km]/[nT] by default).
        """
        return self.key_map[key](omega)

def apply_transfer_function(Bx, By, interval, xml_fname):
    """
    Filter *Bx* and *By* (in [T]) with 3-D transfer function where
    *interval* is the sample period (in [s]). Uses the 3-D transfer
    function model given in *xml_fname*. Return Ex and Ey (in [V/m]).
    """
    # setup surface impedance function
    Z_map = parse_xml(xml_fname)
    interp = Zw_interpolator(Z_map)
    # mu_0 * 1e3 converts from [mv / km] / [nT] to [Ohm]
    Zxx_function = lambda omega: interp(omega, 'xx') * mu_0 * 1e3
    Zxy_function = lambda omega: interp(omega, 'xy') * mu_0 * 1e3
    Zyx_function = lambda omega: interp(omega, 'yx') * mu_0 * 1e3
    Zyy_function = lambda omega: interp(omega, 'yy') * mu_0 * 1e3
    # calculate E field
    Ex, Ey = calc_e_3d(nan_interp(Bx),
                       nan_interp(By),
                       Zxx_function,
                       Zxy_function,
                       Zyx_function,
                       Zyy_function,
                       interval)
    return Ex, Ey

def process(output_mat_fname,
            input_iaga2002_fname,
            xml_fname,
            save_B=False):
    """
    End-to-end processing of an IAGA2002 magnetometer data record
    *input_iaga2002_fname* to the output file *output_mat_fname*
    containing the calculated E-field (units are [V/m]). Use the 3-D
    transfer function model given in *xml_fname.* Also save the
    B-field information if *save_B*.
    """
    # gather Bx and By magnetometer measurements
    _, data_map = parse(iaga2002_fname)
    interval = int((data_map.keys()[1] - data_map.keys()[0]).total_seconds())
    Bx = nan_interp([record.x * 1e-9 for record in data_map.itervalues()])
    By = nan_interp([record.y * 1e-9 for record in data_map.itervalues()])
    # filter with transfer function
    Ex, Ey = apply_transfer_function(Bx,
                                     By,
                                     interval,
                                     xml_fname)
    # save E field
    stn_name = os.path.basename(iaga2002_fname)[:3]
    j2000 = map(toJ2000, data_map.iterkeys())
    mdict = {'Ex': Ex,
             'Ey': Ey,
             'j2000': j2000,
             'input_fname': os.path.abspath(input_iaga2002_fname),
             'xml_fname': os.path.abspath(xml_fname)}
    if save_B:
        mdict['Bx'] = Bx
        mdict['By'] = By
        mdict['Bx_raw'] = [getattr(x, stn_name.upper() + 'X') * 1e-9 for x in data_map.itervalues()]
        mdict['By_raw'] = [getattr(x, stn_name.upper() + 'Y') * 1e-9 for x in data_map.itervalues()]
    savemat(output_mat_fname, mdict)
    return output_mat_fname

def main(argv=None):
    if argv is None:
        argv = sys.argv

    parser = ArgumentParser('Compute E-field from B-field using 3-D transfer function model.',
                            formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument('output_mat_fname',
                        type=str,
                        help='output, containing modeled E-field, in .mat format')
    parser.add_argument('input_iaga2002_fname',
                        type=str,
                        help='input IAGA2002 magnetometer data file')
    parser.add_argument('xml_fname',
                        type=str,
                        help='EM transfer function XML file')
    args = parser.parse_args(argv[1:])

    process(args.output_mat_fname,
            args.input_iaga2002_fname,
            args.xml_fname)

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    sys.exit(main())
<commit_msg>Added divide by zero work around.<commit_after>from __future__ import division

import logging
import sys
import os
import math
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
import xml.etree.ElementTree as ET
from collections import OrderedDict

import numpy as NP
from scipy.constants import mu_0
from scipy.io import savemat
from scipy.interpolate import CubicSpline

from ..util.signal import nextpow2
from ..mag.iaga2002 import parse
from ..util.nan import nan_interp
from ..util.date import toJ2000

def calc_e_3d(Bx,
              By,
              Zxx_function,
              Zxy_function,
              Zyx_function,
              Zyy_function,
              interval):
    """
    Calculate the electric field induced by the magnetic field (*Bx*,
    eastward oriented magnetic field intensity, and *By*, northward
    oriented magnetic field intensity). *Zxx_function*,
    *Z_xy_function*, *Z_yx_function*, and *Z_yy_function* are the
    surface impedance functions (input is angular frequency [rad] and
    output is impedance [Ohms]). Return the tuple *Ex*, *Ey* (same
    orientation as the magnetic field).

    Units:
       (input)  Bx, By: [T]
       (input)  Zxx_function: [rad] -> [Ohm]
       (input)  Zxy_function: [rad] -> [Ohm]
       (input)  Zyx_function: [rad] -> [Ohm]
       (input)  Zyy_function: [rad] -> [Ohm]
       (output) Ex, Ey: [V/m]
    """
    # This function is the reimplementation of Prof. Zhu's matlab code
    # which is based on the algorithm detailed in the NERC Application
    # Guide, Dec. 2013. Note that the NERC Guide contains a sign error
    # that has been corrected in the code below.

    # function global variables
    lp = 60  # length at beginning and end of time series to fit
             # detrend parameters
    p  = 60  # length at beginning and end of time series for
             # frequency domain window
    # get the FFT size
    assert len(Bx) == len(By)
    N = len(Bx)
    Nfft = nextpow2(N)
    # zero-mean and remove linear trend from the magnetic time series
    # data
    ax = NP.mean(Bx[:lp])
    bx = NP.mean(Bx[-lp-1:])
    N_range = NP.arange(1, N+1)
    cx = ax * (N - N_range) / N + bx * N_range / N
    Bx0 = Bx - cx

    ay = NP.mean(By[:lp])
    by = NP.mean(By[-lp-1:])
    cy = ay * (N - N_range) / N + by * N_range / N
    By0 = By - cy
    # window the magnetic data time series
    wp = NP.hstack((0.5 * (1 - NP.cos(2 * math.pi * NP.arange(0, p/2) / p)),
                    NP.ones(N - p),
                    0.5 * (1 - NP.cos(2 * math.pi * NP.arange(p/2 - 1, -1, -1) / p))))
    Bx1 = Bx0 * wp
    By1 = By0 * wp
    # compute FFT
    Sbx = NP.fft.fft(Bx1, n=Nfft) / N
    Sby = NP.fft.fft(By1, n=Nfft) / N
    freq = NP.arange(0, Nfft) / Nfft / interval
    omega = 2 * math.pi * freq

    Zxx_positive = Zxx_function(omega[1:])
    Zxx = NP.hstack(([0], Zxx_positive))

    Zxx2 = NP.hstack((Zxx[:(int(Nfft/2)+1)],
                      NP.conj(Zxx[1:int(Nfft/2)])[::-1]))

    Zxy_positive = Zxy_function(omega[1:])
    Zxy = NP.hstack(([0], Zxy_positive))

    Zxy2 = NP.hstack((Zxy[:(int(Nfft/2)+1)],
                      NP.conj(Zxy[1:int(Nfft/2)])[::-1]))

    Zyx_positive = Zyx_function(omega[1:])
    Zyx = NP.hstack(([0], Zyx_positive))

    Zyx2 = NP.hstack((Zyx[:(int(Nfft/2)+1)],
                      NP.conj(Zyx[1:int(Nfft/2)])[::-1]))

    Zyy_positive = Zyy_function(omega[1:])
    Zyy = NP.hstack(([0], Zyy_positive))

    Zyy2 = NP.hstack((Zyy[:(int(Nfft/2)+1)],
                      NP.conj(Zyy[1:int(Nfft/2)])[::-1]))

    Se_x = []
    Se_y = []
    for Zxx_i, Zxy_i, Zyx_i, Zyy_i, Sbx_i, Sby_i in zip(Zxx2, Zxy2, Zyx2, Zyy2, Sbx, Sby):
        Z = NP.array([[Zxx_i, Zxy_i], [Zyx_i, Zyy_i]])
        Sb_i = NP.array([[Sbx_i / mu_0], [Sby_i / mu_0]])
        Se_i = NP.dot(Z, Sb_i)
        Se_x.append(Se_i[0, 0])
        Se_y.append(Se_i[1, 0])

    Ex = NP.real(NP.fft.ifft(Se_x, Nfft) * N)
    Ey = NP.real(NP.fft.ifft(Se_y, Nfft) * N)

    return (Ex[:N],
            Ey[:N])

def parse_xml(xml_fname):
    """
    Parse the E-M transfer function file *xml_fname* returning a
    mapping between period keys (in [s]) and the 2x2 matrices
    containing the associated Zxx, Zxy, Zyx, and Zyy parameters (in
    [mV / km] / [nT]).

    These files are available at http://ds.iris.edu/spud/emtf.
    """
    tree = ET.parse(xml_fname)
    root = tree.getroot()

    data_list = root.findall('Data')
    assert len(data_list) == 1
    data = data_list[0]

    Z_map = OrderedDict()
    for period in data.findall('Period'):
        Z_list = period.findall('Z')
        assert len(Z_list) == 1
        Z = Z_list[0]
        values = []
        for value, name in zip(Z, ['Zxx', 'Zxy', 'Zyx', 'Zyy']):
            assert value.attrib['name'] == name
            values.append(complex(*map(float, value.text.split())))
        Z_array = NP.array(values)
        Z_array.shape = 2, 2
        Z_map[float(period.attrib['value'])] = Z_array
    return Z_map

def parse_xml_header(xml_fname):
    """
    Parse the E-M transfer function file *xml_fname* returning a
    mapping of site specific information, such as location and data
    quality.
    """
    header_map = {}

    tree = ET.parse(xml_fname)
    root = tree.getroot()

    copyright_list = root.findall('Copyright')
    assert len(copyright_list) == 1
    copyright_ = copyright_list[0]

    for copyright_node in copyright_:
        if copyright_node.tag == 'Acknowledgement':
            header_map['acknowledgement'] = copyright_node.text

    site_list = root.findall('Site')
    assert len(site_list) == 1
    site = site_list[0]

    for site_node in site:
        if site_node.tag == 'Id':
            header_map['id'] = site_node.text
        elif site_node.tag == 'Location':
            for child in site_node:
                if child.tag == 'Latitude':
                    header_map['lat'] = float(child.text)
                elif child.tag == 'Longitude':
                    header_map['lon'] = float(child.text)
                elif child.tag == 'Elevation':
                    assert child.attrib['units'] == 'meters'
                    header_map['el'] = float(child.text)
        elif site_node.tag == 'DataQualityNotes':
            for child in site_node:
                if child.tag == 'Rating':
                    header_map['rating'] = int(child.text)
                elif child.tag == 'GoodFromPeriod':
                    header_map['good_from'] = float(child.text)
                    try:
                        header_map['good_to_mHz'] = 1/float(child.text) * 1e3
                    except ZeroDivisionError:
                        header_map['good_to_mHz'] = float('inf')
                elif child.tag == 'GoodToPeriod':
                    header_map['good_to'] = float(child.text)
                    try:
                        header_map['good_from_mHz'] = 1/float(child.text) * 1e3
                    except ZeroDivisionError:
                        header_map['good_from_mHz'] = float('inf')
                elif child.tag == 'Comments':
                    header_map['data_quality_comments'] = child.text
        elif site_node.tag == 'DataQualityWarnings':
            for child in site_node:
                if child.tag == 'Flag':
                    header_map['data_quality_flag'] = int(child.text)
                elif child.tag == 'Comments':
                    header_map['data_quality_warning_comments'] = child.text
        elif site_node.tag == 'Acknowledgment':
            header_map['acknowledgment'] = site_node.text

    return header_map

class Zw_interpolator(object):
    def __init__(self, Z_map):
        """
        Construct a cubic-spline 3-D E-M transfer function interpolater
        using the information in *Z_map* returned from
        :func:`parse_xml` as the function samples.
        """
        self.Z_map = Z_map
        periods = Z_map.keys()
        self.f = NP.array([1/x for x in periods[::-1]])
        self.omega = 2 * math.pi * self.f
        self.Zxx_interp = CubicSpline(self.omega, [x[0, 0] for x in Z_map.values()[::-1]])
        self.Zxy_interp = CubicSpline(self.omega, [x[0, 1] for x in Z_map.values()[::-1]])
        self.Zyx_interp = CubicSpline(self.omega, [x[1, 0] for x in Z_map.values()[::-1]])
        self.Zyy_interp = CubicSpline(self.omega, [x[1, 1] for x in Z_map.values()[::-1]])
        self.key_map = {'xx': self.Zxx_interp,
                        'xy': self.Zxy_interp,
                        'yx': self.Zyx_interp,
                        'yy': self.Zyy_interp}

    def __call__(self, omega, key):
        """
        Return the interpolated value of the transfer function (*key* is
        xx, xy, yx, or yy) at angular frequency *omega*. The units of
        *omega* and the output are the same as *Z_map* ([rad] and
        [mV/km]/[nT] by default).
        """
        return self.key_map[key](omega)

def apply_transfer_function(Bx, By, interval, xml_fname):
    """
    Filter *Bx* and *By* (in [T]) with 3-D transfer function where
    *interval* is the sample period (in [s]). Uses the 3-D transfer
    function model given in *xml_fname*. Return Ex and Ey (in [V/m]).
    """
    # setup surface impedance function
    Z_map = parse_xml(xml_fname)
    interp = Zw_interpolator(Z_map)
    # mu_0 * 1e3 converts from [mv / km] / [nT] to [Ohm]
    Zxx_function = lambda omega: interp(omega, 'xx') * mu_0 * 1e3
    Zxy_function = lambda omega: interp(omega, 'xy') * mu_0 * 1e3
    Zyx_function = lambda omega: interp(omega, 'yx') * mu_0 * 1e3
    Zyy_function = lambda omega: interp(omega, 'yy') * mu_0 * 1e3
    # calculate E field
    Ex, Ey = calc_e_3d(nan_interp(Bx),
                       nan_interp(By),
                       Zxx_function,
                       Zxy_function,
                       Zyx_function,
                       Zyy_function,
                       interval)
    return Ex, Ey

def process(output_mat_fname,
            input_iaga2002_fname,
            xml_fname,
            save_B=False):
    """
    End-to-end processing of an IAGA2002 magnetometer data record
    *input_iaga2002_fname* to the output file *output_mat_fname*
    containing the calculated E-field (units are [V/m]). Use the 3-D
    transfer function model given in *xml_fname.* Also save the
    B-field information if *save_B*.
    """
    # gather Bx and By magnetometer measurements
    _, data_map = parse(iaga2002_fname)
    interval = int((data_map.keys()[1] - data_map.keys()[0]).total_seconds())
    Bx = nan_interp([record.x * 1e-9 for record in data_map.itervalues()])
    By = nan_interp([record.y * 1e-9 for record in data_map.itervalues()])
    # filter with transfer function
    Ex, Ey = apply_transfer_function(Bx,
                                     By,
                                     interval,
                                     xml_fname)
    # save E field
    stn_name = os.path.basename(iaga2002_fname)[:3]
    j2000 = map(toJ2000, data_map.iterkeys())
    mdict = {'Ex': Ex,
             'Ey': Ey,
             'j2000': j2000,
             'input_fname': os.path.abspath(input_iaga2002_fname),
             'xml_fname': os.path.abspath(xml_fname)}
    if save_B:
        mdict['Bx'] = Bx
        mdict['By'] = By
        mdict['Bx_raw'] = [getattr(x, stn_name.upper() + 'X') * 1e-9 for x in data_map.itervalues()]
        mdict['By_raw'] = [getattr(x, stn_name.upper() + 'Y') * 1e-9 for x in data_map.itervalues()]
    savemat(output_mat_fname, mdict)
    return output_mat_fname

def main(argv=None):
    if argv is None:
        argv = sys.argv

    parser = ArgumentParser('Compute E-field from B-field using 3-D transfer function model.',
                            formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument('output_mat_fname',
                        type=str,
                        help='output, containing modeled E-field, in .mat format')
    parser.add_argument('input_iaga2002_fname',
                        type=str,
                        help='input IAGA2002 magnetometer data file')
    parser.add_argument('xml_fname',
                        type=str,
                        help='EM transfer function XML file')
    args = parser.parse_args(argv[1:])

    process(args.output_mat_fname,
            args.input_iaga2002_fname,
            args.xml_fname)

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    sys.exit(main())
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# This file is part of beets.
# Copyright 2016, Jan-Erik Dahlin
#
# 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.

"""If the title is empty, try to extract track and title from the
filename.
"""
from __future__ import division, absolute_import, print_function

from beets import plugins
from beets.util import displayable_path
import os
import re
import six

# Filename field extraction patterns.
PATTERNS = [
    # Useful patterns.
    r'^(?P<artist>.+)[\-_](?P<title>.+)[\-_](?P<tag>.*)$',
    r'^(?P<track>\d+)[\s.\-_]+(?P<artist>.+)[\-_](?P<title>.+)[\-_](?P<tag>.*)$',
    r'^(?P<artist>.+)[\-_](?P<title>.+)$',
    r'^(?P<track>\d+)[\s.\-_]+(?P<artist>.+)[\-_](?P<title>.+)$',
    r'^(?P<title>.+)$',
    r'^(?P<track>\d+)[\s.\-_]+(?P<title>.+)$',
    r'^(?P<track>\d+)\s+(?P<title>.+)$',
    r'^(?P<title>.+) by (?P<artist>.+)$',
    r'^(?P<track>\d+).*$',
]

# Titles considered "empty" and in need of replacement.
BAD_TITLE_PATTERNS = [
    r'^$',
]

def equal(seq):
    """Determine whether a sequence holds identical elements.
    """
    return len(set(seq)) <= 1

def equal_fields(matchdict, field):
    """Do all items in `matchdict`, whose values are dictionaries, have
    the same value for `field`? (If they do, the field is probably not
    the title.)
    """
    return equal(m[field] for m in matchdict.values())

def all_matches(names, pattern):
    """If all the filenames in the item/filename mapping match the
    pattern, return a dictionary mapping the items to dictionaries
    giving the value for each named subpattern in the match. Otherwise,
    return None.
    """
    matches = {}
    for item, name in names.items():
        m = re.match(pattern, name, re.IGNORECASE)
        if m and m.groupdict():
            # Only yield a match when the regex applies *and* has
            # capture groups. Otherwise, no information can be extracted
            # from the filename.
            matches[item] = m.groupdict()
        else:
            return None
    return matches

def bad_title(title):
    """Determine whether a given title is "bad" (empty or otherwise
    meaningless) and in need of replacement.
    """
    for pat in BAD_TITLE_PATTERNS:
        if re.match(pat, title, re.IGNORECASE):
            return True
    return False

def apply_matches(d):
    """Given a mapping from items to field dicts, apply the fields to
    the objects.
    """
    some_map = list(d.values())[0]
    keys = some_map.keys()

    # Only proceed if the "tag" field is equal across all filenames.
    if 'tag' in keys and not equal_fields(d, 'tag'):
        return

    # Given both an "artist" and "title" field, assume that one is
    # *actually* the artist, which must be uniform, and use the other
    # for the title. This, of course, won't work for VA albums.
    if 'artist' in keys:
        if equal_fields(d, 'artist'):
            artist = some_map['artist']
            title_field = 'title'
        elif equal_fields(d, 'title'):
            artist = some_map['title']
            title_field = 'artist'
        else:
            # Both vary. Abort.
            return

        for item in d:
            if not item.artist:
                item.artist = artist

    # No artist field: remaining field is the title.
    else:
        title_field = 'title'

    # Apply the title and track.
    for item in d:
        if bad_title(item.title):
            item.title = six.text_type(d[item][title_field])
        if 'track' in d[item] and item.track == 0:
            item.track = int(d[item]['track'])

# Plugin structure and hook into import process.

class FromFilenamePlugin(plugins.BeetsPlugin):
    def __init__(self):
        super(FromFilenamePlugin, self).__init__()
        self.register_listener('import_task_start', filename_task)

def filename_task(task, session):
    """Examine each item in the task to see if we can extract a title
    from the filename. Try to match all filenames to a number of
    regexps, starting with the most complex patterns and successively
    trying less complex patterns. As soon as all filenames match the
    same regex we can make an educated guess of which part of the
    regex that contains the title.
    """
    items = task.items if task.is_album else [task.item]

    # Look for suspicious (empty or meaningless) titles.
    missing_titles = sum(bad_title(i.title) for i in items)

    if missing_titles:
        # Get the base filenames (no path or extension).
        names = {}
        for item in items:
            path = displayable_path(item.path)
            name, _ = os.path.splitext(os.path.basename(path))
            names[item] = name

        # Look for useful information in the filenames.
        for pattern in PATTERNS:
            d = all_matches(names, pattern)
            if d:
                apply_matches(d)
<commit_msg>Fix "line too long"  Travis CI error<commit_after># -*- coding: utf-8 -*-
# This file is part of beets.
# Copyright 2016, Jan-Erik Dahlin
#
# 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.

"""If the title is empty, try to extract track and title from the
filename.
"""
from __future__ import division, absolute_import, print_function

from beets import plugins
from beets.util import displayable_path
import os
import re
import six

# Filename field extraction patterns.
PATTERNS = [
  # Useful patterns.
  r'^(?P<artist>.+)[\-_](?P<title>.+)[\-_](?P<tag>.*)$',
  r'^(?P<track>\d+)[\s.\-_]+(?P<artist>.+)[\-_](?P<title>.+)[\-_](?P<tag>.*)$',
  r'^(?P<artist>.+)[\-_](?P<title>.+)$',
  r'^(?P<track>\d+)[\s.\-_]+(?P<artist>.+)[\-_](?P<title>.+)$',
  r'^(?P<title>.+)$',
  r'^(?P<track>\d+)[\s.\-_]+(?P<title>.+)$',
  r'^(?P<track>\d+)\s+(?P<title>.+)$',
  r'^(?P<title>.+) by (?P<artist>.+)$',
  r'^(?P<track>\d+).*$',
]

# Titles considered "empty" and in need of replacement.
BAD_TITLE_PATTERNS = [
    r'^$',
]

def equal(seq):
    """Determine whether a sequence holds identical elements.
    """
    return len(set(seq)) <= 1

def equal_fields(matchdict, field):
    """Do all items in `matchdict`, whose values are dictionaries, have
    the same value for `field`? (If they do, the field is probably not
    the title.)
    """
    return equal(m[field] for m in matchdict.values())

def all_matches(names, pattern):
    """If all the filenames in the item/filename mapping match the
    pattern, return a dictionary mapping the items to dictionaries
    giving the value for each named subpattern in the match. Otherwise,
    return None.
    """
    matches = {}
    for item, name in names.items():
        m = re.match(pattern, name, re.IGNORECASE)
        if m and m.groupdict():
            # Only yield a match when the regex applies *and* has
            # capture groups. Otherwise, no information can be extracted
            # from the filename.
            matches[item] = m.groupdict()
        else:
            return None
    return matches

def bad_title(title):
    """Determine whether a given title is "bad" (empty or otherwise
    meaningless) and in need of replacement.
    """
    for pat in BAD_TITLE_PATTERNS:
        if re.match(pat, title, re.IGNORECASE):
            return True
    return False

def apply_matches(d):
    """Given a mapping from items to field dicts, apply the fields to
    the objects.
    """
    some_map = list(d.values())[0]
    keys = some_map.keys()

    # Only proceed if the "tag" field is equal across all filenames.
    if 'tag' in keys and not equal_fields(d, 'tag'):
        return

    # Given both an "artist" and "title" field, assume that one is
    # *actually* the artist, which must be uniform, and use the other
    # for the title. This, of course, won't work for VA albums.
    if 'artist' in keys:
        if equal_fields(d, 'artist'):
            artist = some_map['artist']
            title_field = 'title'
        elif equal_fields(d, 'title'):
            artist = some_map['title']
            title_field = 'artist'
        else:
            # Both vary. Abort.
            return

        for item in d:
            if not item.artist:
                item.artist = artist

    # No artist field: remaining field is the title.
    else:
        title_field = 'title'

    # Apply the title and track.
    for item in d:
        if bad_title(item.title):
            item.title = six.text_type(d[item][title_field])
        if 'track' in d[item] and item.track == 0:
            item.track = int(d[item]['track'])

# Plugin structure and hook into import process.

class FromFilenamePlugin(plugins.BeetsPlugin):
    def __init__(self):
        super(FromFilenamePlugin, self).__init__()
        self.register_listener('import_task_start', filename_task)

def filename_task(task, session):
    """Examine each item in the task to see if we can extract a title
    from the filename. Try to match all filenames to a number of
    regexps, starting with the most complex patterns and successively
    trying less complex patterns. As soon as all filenames match the
    same regex we can make an educated guess of which part of the
    regex that contains the title.
    """
    items = task.items if task.is_album else [task.item]

    # Look for suspicious (empty or meaningless) titles.
    missing_titles = sum(bad_title(i.title) for i in items)

    if missing_titles:
        # Get the base filenames (no path or extension).
        names = {}
        for item in items:
            path = displayable_path(item.path)
            name, _ = os.path.splitext(os.path.basename(path))
            names[item] = name

        # Look for useful information in the filenames.
        for pattern in PATTERNS:
            d = all_matches(names, pattern)
            if d:
                apply_matches(d)
<|endoftext|>


<commit_before>import kivy
kivy.require('1.9.1')

from kivy.app import App
from kivy.uix.boxlayout import BoxLayout
from kivy.uix.gridlayout import GridLayout
from kivy.uix.widget import Widget
from kivy.uix.label import Label
from kivy.uix.button import Button
from kivy.uix.dropdown import DropDown
from kivy.graphics import Color, Rectangle, Line

from kivy.core.window import Window

Window.clearcolor = (0.85, 0.85, 0.85, 0.5)
Window.borderless = True

class Net:
    def __init__(self):
        pass

    def init(self):
        self.active_slots = [[], [], []]  # free, black, white
    '''
        active_slots[0] = new
        Vector. < Slot > (); // free
        active_slots[1] = new
        Vector. < Slot > (); // black
        active_slots[2] = new
        Vector. < Slot > (); // white

        empty_points = new
        Vector. < Point > ();
        all_slots = new
        Vector. < Slot > ();

        for (var i: int = 0; i < 225; i + +) {
            var
        p:Point = new
        Point(this, Math.floor(i / 15) - 7, i % 15 - 7);
        // trace(i, p.x, p.y);
        all_points[i] = p;
        empty_points.push(p);

        for (var j: int = 0; j < 4; j + +) {
        if (p.is_valid_scp(j))
        {
            var
        s: Slot = new
        Slot(this, p, j);
        all_slots.push(s);
        active_slots[0].push(s);
        }
        }
        }

        for each(var item: Slot in all_slots) {
            item.init();
        }
        '''

class Game:
    def __init__(self, app):
        self.is_play = False
        self.is_run = False
        self.is_busy = False

        self.n_step = 0
        self.mes = ""

        self.app = app
        self.app.status.text = "Press 'New'"

    def play(self):
        self.is_play = True
        self.is_run = False
        self.is_busy = False

        if self.app.mode() == 0:
            self.app.action_step.disabled = False
        else:
            self.app.action_step.disabled = True
            self.app.action_back.disabled = True

        self.app.action_run.disabled = False
        self.app.action_mode.disabled = True

        self.app.desk.draw_init()
        self.app.desk.draw_grid()

        self.app.net.init()

        self.mes = "Start"
        self.n_step = 1
        self.app.net.step(0, 0, 1)
        self.app.status.text = "New game"

        if self.app.mode() == 1:
            self.run(False)

class Desk(Widget):
    def draw_init(self):
        self.canvas.clear()

        self.h = 31
        self.s = 16 * self.h
        self.pady = (Window.height - self.s - self.h * 2) / 2
        self.padx = (Window.width - self.s) / 2
        self.size = (self.s, self.s)
        self.pos = (self.padx, self.h + self.pady)
        self.d = self.h
        self.cx = [self.pos[0] + self.d + i * self.d for i in range(0, 15)]
        self.cy = [self.pos[1] + self.d + i * self.d for i in range(0, 15)]

    def draw_grid(self):

        self.canvas.add(Color(1., 1., 1.))
        self.canvas.add(Rectangle(size=self.size, pos=self.pos))
        self.canvas.add(Color(0., 0., 0.))
        for c in self.cx:
            self.canvas.add(Line(points = [c, self.pos[1] + self.d, c, self.pos[1] + self.size[1] - self.d]))
        for c in self.cy:
            self.canvas.add(Line(points=[self.pos[0] + self.d, c, self.pos[0] + self.size[0] - self.d, c]))

class Gomoku(App):

    def build(self):
        main = GridLayout(cols = 1, rows = 1)

        actions = BoxLayout(orientation='horizontal', pos_hint = {'x': 0}, size_hint_x = 1)
        self.action_new = Button(text = "New",
                            pos_hint = {'top': 1}, size = (70, 30), size_hint = (None, None),
                            on_press = lambda x: self.game.play())

        self.action_mode = Button(text="Mode:Black",
                            pos_hint = {'top': 1}, size = (100, 30), size_hint = (None, None),
                            on_release = lambda x: drop_mode.open(x))
        drop_mode = DropDown()
        drop_mode.bind(on_select=lambda instance, x: setattr(self.action_mode, 'text', x))
        mode_on_release = lambda btn: drop_mode.select("Mode:" + btn.text)

        black_mode = Button(text = "Black", size = (100, 30), size_hint = (None, None), on_release=mode_on_release)
        white_mode = Button(text = "White", size=(100, 30), size_hint=(None, None), on_release=mode_on_release)
        manual_mode = Button(text = "Manual", size=(100, 30), size_hint=(None, None), on_release=mode_on_release)
        drop_mode.add_widget(black_mode)
        drop_mode.add_widget(white_mode)
        drop_mode.add_widget(manual_mode)

        self.action_step = Button(text = "Step", pos_hint = {'top': 1}, size = (70, 30), size_hint = (None, None))
        self.action_back = Button(text = "Back", pos_hint = {'top': 1}, size = (70, 30), size_hint = (None, None))
        self.action_run = Button(text = "Run", pos_hint = {'top': 1}, size = (70, 30), size_hint = (None, None))
        self.action_close = Button(text="Close",
                            pos_hint={'top': 1}, size=(70, 30), size_hint=(None, None),
                            on_press = lambda x: Window.close())

        actions.add_widget(self.action_new)
        actions.add_widget(self.action_mode)
        actions.add_widget(self.action_step)
        actions.add_widget(self.action_back)
        actions.add_widget(self.action_run)
        actions.add_widget(self.action_close)

        self.desk = Desk()

        self.status = Label(text = "Ready", pos_hint = {'bottom': 1}, size = (200, 30), size_hint = (None, None), text_size=(180, None))

        vbox = BoxLayout(orientation='vertical', size = main.size, pos_hint = {'x': 0, 'y': 0}, size_hint_x = 1, size_hint_y = 1)
        vbox.add_widget(actions)
        vbox.add_widget(self.desk)
        vbox.add_widget(self.status)
        main.add_widget(vbox)

        self.action_step.disabled = True
        self.action_back.disabled = True
        self.action_run.disabled = True

        self.desk.draw_init()
        self.desk.draw_grid()

        self.game = Game(self)
        self.colors = ([0, 0, 0, 0.5], [1, 1, 1, 0.5])
        self.qsteps = 0
        self.steps = [None] * 225
        self.mode = lambda : {"Mode:Manual":0, "Mode:Black":1, "Mode:White":2}[self.action_mode.text]

        print self.mode()

        self.net = Net()

        return main

if __name__ == '__main__':
    Gomoku().run()

<commit_msg>In progress<commit_after>import kivy
kivy.require('1.9.1')

from kivy.app import App
from kivy.uix.boxlayout import BoxLayout
from kivy.uix.gridlayout import GridLayout
from kivy.uix.widget import Widget
from kivy.uix.label import Label
from kivy.uix.button import Button
from kivy.uix.dropdown import DropDown
from kivy.graphics import Color, Rectangle, Line, Ellipse

from kivy.core.window import Window

import random

Window.clearcolor = (0.85, 0.85, 0.85, 0.5)
Window.borderless = True

class Point:
    def __init__(self, net, x, y):
        self.net = net
        self.x = x
        self.y = y
        self.slots = []
        self.r = [0, 0, 0]
        self.s = 0

    def is_valid_scp(self, d):
        x = self.x - 7
        y = self.y - 7
    # 0 - vert, 1 - horiz, 2 - up, 3 - down
        if d == 0 and y > -6 and y < 6:
            return True
        if d == 1 and x > -6 and x < 6:
            return True
        if d == 2 and (x > -6 and y < 6) and (x < 6 and y > -6):
            return True
        if d == 3 and (x > -6 and y > -6) and (x < 6 and y < 6):
            return True
        return False

    def add_slot(self, s):
        self.slots.append(s)
        self.r[s.s] += 1

    def to_string(self):
        return "x:{0},y:{1};".format(self.x, self.y)

class Slot:
    def __init__(self, net, scp, d):
        self.net = net
        self.scp = scp
        self.d = d
        self.points = [None] * 5
        self.r = 0
        self.s = 0

    def init(self):
        #print "Init Slot: d: {0}, scp: {1}".format(self.d, self.scp.to_string())
        self.points[2] = self.net.get_point(self.scp.x, self.scp.y)

        if self.d == 0:
            self.points[0] = self.net.get_point(self.scp.x, self.scp.y - 2)
            self.points[1] = self.net.get_point(self.scp.x, self.scp.y - 1)
            self.points[3] = self.net.get_point(self.scp.x, self.scp.y + 1)
            self.points[4] = self.net.get_point(self.scp.x, self.scp.y + 2)
        elif self.d == 1:
            self.points[0] = self.net.get_point(self.scp.x - 2, self.scp.y)
            self.points[1] = self.net.get_point(self.scp.x - 1, self.scp.y)
            self.points[3] = self.net.get_point(self.scp.x + 1, self.scp.y)
            self.points[4] = self.net.get_point(self.scp.x + 2, self.scp.y)
        elif self.d == 2:
            self.points[0] = self.net.get_point(self.scp.x - 2, self.scp.y - 2)
            self.points[1] = self.net.get_point(self.scp.x - 1, self.scp.y - 1)
            self.points[3] = self.net.get_point(self.scp.x + 1, self.scp.y + 1)
            self.points[4] = self.net.get_point(self.scp.x + 2, self.scp.y + 2)
        elif self.d == 3:
            self.points[0] = self.net.get_point(self.scp.x - 2, self.scp.y + 2)
            self.points[1] = self.net.get_point(self.scp.x - 1, self.scp.y + 1)
            self.points[3] = self.net.get_point(self.scp.x + 1, self.scp.y - 1)
            self.points[4] = self.net.get_point(self.scp.x + 2, self.scp.y - 2)

        for i in range(0, 5):
            self.points[i].add_slot(self)

class Net:
    def __init__(self):
        pass

    def init(self):
        self.all_slots = []
        self.active_slots = [[], [], []]  # free, black, white

        self.all_points = [Point(self, int(i / 15), i % 15 ) for i in range(0, 225)]
        self.empty_points = self.all_points[:]

        for p in self.all_points:
            for d in range(0, 4):
                if p.is_valid_scp(d):
                    s = Slot(self, p, d)
                    self.all_slots.append(s)

        self.active_slots[0] = self.all_slots[:]
        for s in self.all_slots:
            s.init()

    def step(self, x, y, c):
        p = self.get_point(x, y)
        p.s = c
        self.empty_points.remove(p)

        for s in p.slots:
            if s.s == 0:
                p.r[0] -= 1
                p.r[c] += 1
                s.s = c
                s.r = 1
                self.active_slots[0].remove(s)
                self.active_slots[c].append(s)
            elif s.s == c:
                p.r[c] += 1
                s.r += 1
            elif s.s != 3:
                p.r[c] -= 1
                self.active_slots[s.s].remove(s)
                s.s = 3

    def get_point(self, x, y):
        #print (x, y)
        return self.all_points[x * 15 + y]

class Game:
    def __init__(self, app):
        self.is_play = False
        self.is_run = False
        self.is_busy = False

        self.n_step = 0
        self.mes = ""

        self.app = app
        self.name_c = ["", "Black", "White"]

    def play(self):
        self.app.desk.draw_init()
        self.app.desk.draw_grid()

        self.is_play = True
        self.is_run = False
        self.is_busy = False

        if self.app.mode() == 0:
            self.app.action_step.disabled = False
        else:
            self.app.action_step.disabled = True
            self.app.action_back.disabled = True

        self.app.action_run.disabled = False
        self.app.action_mode.disabled = True

        self.app.net.init()

        self.app.qsteps = 0
        self.app.add_step(7, 7, 0, "Start")

        self.mes = "Start"
        self.n_step = 1
        self.app.net.step(7, 7, 1)
        self.app.status.text = "New game"

        if self.app.mode() == 1:
            self.run(False)

    def run(self, r):
        if self.is_play:
            self.app.status.text = "Thinking..."
            self.is_busy = True
            self.is_run = r
            self.go(True, 0, 0)
            self.is_busy = False

    def go(self, auto, x, y):
        ret = 0
        if auto:
            ret = self.next_step()
        else:
            ret = self.manual_step(x, y)

        self.app.status.text = "Finish! -> {0}".format(self.mes) if ret < 0 else "Step {0} -> {1}".format(ret, self.mes)

        if self.app.mode == 0:
            self.app.action_back.disabled = False

        if ret < 0 or ret > 224:
            self.app.action_run.disabled = True
            self.app.action_step.disabled = True
            self.app.action_mode.disabled = False

            self.is_run = False
            self.is_play = False

        elif not auto and self.app.mode() > 0:
            self.go(True, 7, 7)
        elif self.is_run:
            self.go(True, 7, 7)

    def next_step(self):
        self.n_step += 1

        if self.check_win(3 - (2 - self.n_step % 2)) or self.check_draw():
            return -1
        else:
            p = self.calc_point(2 - self.n_step % 2)
            self.app.net.step(p.x, p.y, 2 - self.n_step % 2)
            self.app.add_step(p.x, p.y, 1 - self.n_step % 2, self.mes)
            return self.n_step

    def manual_step(self, x, y):
        self.n_step += 1
        if self.check_win(3 - (2 - self.n_step % 2)) or self.check_draw():
            return -1
        else:
            self.app.net.step(x, y, 2 - self.n_step % 2)
            self.mes = self.name_c[2 - self.n_step % 2] + " :: manual (" + x + ":" + y + ")"
            self.app.add_step(x, y, 1 - self.n_step % 2, self.mes)
            return self.n_step

    def check_win(self, c):
        for s in self.app.net.active_slots[c]:
            if s.r == 5:
                mes = self.name_c[c] + " :: win!!!"
                return True
        return False

    def check_draw(self):
        if len(self.app.net.active_slots[0]) == 0 and \
                len(self.app.net.active_slots[1]) == 0 and \
                len(self.app.net.active_slots[2]) == 0:
            self.mes = " draw :("
            return True
        else:
            return False

    def calc_point(self, c):
        ret = []
        self.mes = self.name_c[c] + " :: auto :: "

        ret = self.find_slot_4(c)
        if len(ret) == 0:
            ret = self.find_slot_4(3 - c)
        if len(ret) == 0:

            ret = self.find_point_x(c, 2, 1)
        if len(ret) == 0:
            ret = self.find_point_x(3 - c, 2, 1)
        if len(ret) == 0:

            ret = self.find_point_x(c, 1, 5)
        if len(ret) == 0:
            ret = self.find_point_x(3 - c, 1, 5)
        if len(ret) == 0:

            ret = self.find_point_x(c, 1, 4)
        if len(ret) == 0:
            ret = self.find_point_x(3 - c, 1, 4)

        if len(ret) == 0:
            ret = self.find_point_x(c, 1, 3)
        if len(ret) == 0:
            ret = self.find_point_x(3 - c, 1, 3)

        if len(ret) == 0:
            ret = self.find_point_x(c, 1, 2)
        if len(ret) == 0:
            ret = self.find_point_x(3 - c, 1, 2)

        if len(ret) == 0:
            ret = self.find_point_x(c, 1, 1)
        if len(ret) == 0:
            ret = self.find_point_x(3 - c, 1, 1)

        if len(ret) == 0:
            ret = self.find_point_x(c, 0, 10)
        if len(ret) == 0:
            ret = self.find_point_x(3 - c, 0, 10)

        if len(ret) == 0:
            ret = self.find_point_x(c, 0, 9)
        if len(ret) == 0:
            ret = self.find_point_x(3 - c, 0, 9)

        if len(ret) == 0:
            ret = self.find_point_x(c, 0, 8)
        if len(ret) == 0:
            ret = self.find_point_x(3 - c, 0, 8)

        if len(ret) == 0:
            ret = self.find_point_x(c, 0, 7)
        if len(ret) == 0:
            ret = self.find_point_x(3 - c, 0, 7)

        if len(ret) == 0:
            ret = self.calc_point_max_rate(c)

        #mes = ret[0].m;
		#return ret[0].p;
		#i = int(Math.random() * ret.length)
        o = random.choice(ret)
        self.mes = o["m"]
        return o["p"]

    def find_slot_4(self, c):
        ret = []
        m = 0
        for s in self.app.net.active_slots[c]:
            if s.r == 4:
                for p in s.points:
                    if p.s == 0:
                        m = "{0} {1} :: find_slot_4 -> -> ({2},{3})".format(self.mes, self.name_c[c], p.x, p.y)
                        ret.append({"p": p, "m": m})
        return ret

    def find_point_x(self, c, r, b):
        ret = []
        m = ""
        for p in self.app.net.empty_points:
            i = 0
            for s in p.slots:
                if s.s == c and s.r > r:
                    i += 1
            if i > b:
                m = "{0} {1} :: point_max_rate({2},{3}) -> ({4},{5})".format(self.mes, self.name_c[c], r, b, p.x, p.y)
                ret.append({"p": p, "m": m})

        return ret

    def calc_point_max_rate(self, c):
        ret = []
        m = ""
        r = -1
        d = 0
        i = 0

        for p in self.app.net.empty_points:
            d = 0
            for s in p.slots:
                if s.s == 0:
                    d += 1
                elif s.s == c:
                    d += (1 + s.r) * (1 + s.r)
                elif s.s != 3:
                    d += (1 + s.r) * (1 + s.r)

            if d > r:
                i = 1
                r = d
                ret = []
                m = "{0} {1} :: point_max_rate({2},{3}) -> ({4},{5})".format(self.mes, self.name_c[c], i, r, p.x, p.y)
                ret.append({"p": p, "m": m})
            elif d == r:
                i += 1
                m = "{0} {1} :: point_max_rate({2},{3}) -> ({4},{5})".format(self.mes, self.name_c[c], i, r, p.x, p.y)
                ret.append({"p": p, "m": m})

        return ret

class Desk(Widget):
    def draw_init(self):
        self.canvas.clear()

        self.h = 31
        self.s = 16 * self.h
        self.pady = (Window.height - self.s - self.h * 2) / 2
        self.padx = (Window.width - self.s) / 2
        self.size = (self.s, self.s)
        self.pos = (self.padx, self.h + self.pady)
        self.d = self.h
        self.cx = [self.pos[0] + self.d + i * self.d for i in range(0, 15)]
        self.cy = [self.pos[1] + self.d + i * self.d for i in range(0, 15)]

    def draw_grid(self):

        self.canvas.add(Color(1., 1., 1.))
        self.canvas.add(Rectangle(size=self.size, pos=self.pos))
        self.canvas.add(Color(0., 0., 0.))
        for c in self.cx:
            self.canvas.add(Line(points = [c, self.pos[1] + self.d, c, self.pos[1] + self.size[1] - self.d]))
        for c in self.cy:
            self.canvas.add(Line(points=[self.pos[0] + self.d, c, self.pos[0] + self.size[0] - self.d, c]))

    def draw_step(self, x, y, c):
        self.canvas.add(Color(c[0], c[1], c[2]))
        self.canvas.add(Ellipse(pos = [self.cx[x] - self.d/2, self.cy[y] - self.d/2], size = [self.d, self.d]))
        self.canvas.add(Color(0., 0., 0.))
        self.canvas.add(Line(circle = (self.cx[x], self.cy[y], self.d/2)))

class Gomoku(App):

    def build(self):
        self.game = Game(self)
        main = GridLayout(cols = 1, rows = 1)

        actions = BoxLayout(orientation='horizontal', pos_hint = {'x': 0}, size_hint_x = 1)
        self.action_new = Button(text = "New",
                            pos_hint = {'top': 1}, size = (70, 30), size_hint = (None, None),
                            on_press = lambda x: self.game.play())

        self.action_mode = Button(text="Mode:Black",
                            pos_hint = {'top': 1}, size = (100, 30), size_hint = (None, None),
                            on_release = lambda x: drop_mode.open(x))
        drop_mode = DropDown()
        drop_mode.bind(on_select=lambda instance, x: setattr(self.action_mode, 'text', x))
        mode_on_release = lambda btn: drop_mode.select("Mode:" + btn.text)

        black_mode = Button(text = "Black", size = (100, 30), size_hint = (None, None), on_release=mode_on_release)
        white_mode = Button(text = "White", size=(100, 30), size_hint=(None, None), on_release=mode_on_release)
        manual_mode = Button(text = "Manual", size=(100, 30), size_hint=(None, None), on_release=mode_on_release)
        drop_mode.add_widget(black_mode)
        drop_mode.add_widget(white_mode)
        drop_mode.add_widget(manual_mode)

        self.action_step = Button(text = "Step",
                            pos_hint = {'top': 1}, size = (70, 30), size_hint = (None, None),
                            on_press = lambda x: self.game.run(False))

        self.action_back = Button(text = "Back", pos_hint = {'top': 1}, size = (70, 30), size_hint = (None, None))

        self.action_run = Button(text = "Run",
                            pos_hint = {'top': 1}, size = (70, 30), size_hint = (None, None),
                            on_press=lambda x: self.game.run(True))
        self.action_close = Button(text="Close",
                            pos_hint={'top': 1}, size=(70, 30), size_hint=(None, None),
                            on_press = lambda x: Window.close())

        actions.add_widget(self.action_new)
        actions.add_widget(self.action_mode)
        actions.add_widget(self.action_step)
        actions.add_widget(self.action_back)
        actions.add_widget(self.action_run)
        actions.add_widget(self.action_close)

        self.desk = Desk()

        self.status = Label(text = "Ready", pos_hint = {'bottom': 1}, size = (500, 30), size_hint = (None, None), text_size=(500, None))

        vbox = BoxLayout(orientation='vertical', size = main.size, pos_hint = {'x': 0, 'y': 0}, size_hint_x = 1, size_hint_y = 1)
        vbox.add_widget(actions)
        vbox.add_widget(self.desk)
        vbox.add_widget(self.status)
        main.add_widget(vbox)

        self.action_step.disabled = True
        self.action_back.disabled = True
        self.action_run.disabled = True

        self.desk.draw_init()
        self.desk.draw_grid()

        self.colors = ([0, 0, 0, 0.5], [1, 1, 1, 0.5])
        self.qsteps = 0
        self.steps = [None] * 225
        self.mode = lambda : {"Mode:Manual":0, "Mode:Black":1, "Mode:White":2}[self.action_mode.text]

        print self.mode()

        self.net = Net()
        self.status.text = "Press 'New'"
        return main

    def add_step(self, x, y, c, mes):
        self.steps[self.qsteps] = {"x": x, "y": y, "mes": mes}
        self.qsteps += 1
        self.desk.draw_step(x, y, self.colors[c])

if __name__ == '__main__':
    Gomoku().run()

<|endoftext|>


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

# Copyright (c) 2016 Jérémie DECOCK (http://www.jdhp.org)

# 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.

# See also: http://effbot.org/tkinterbook/checkbutton.htm

import tkinter as tk

root = tk.Tk()

test_var = tk.IntVar()
#test_var.set(1)

def callback():
    print("var = ", test_var.get())

checkbutton = tk.Checkbutton(root, text="Test", variable=test_var, command=callback)

checkbutton.pack()

root.mainloop()

<commit_msg>Misc update.<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright (c) 2016 Jérémie DECOCK (http://www.jdhp.org)

# 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.

# See also: http://effbot.org/tkinterbook/checkbutton.htm

import tkinter as tk

root = tk.Tk()

test_var = tk.IntVar()

# Initialize
#test_var.set(1)

def callback():
    print("var = ", test_var.get())

checkbutton = tk.Checkbutton(root, text="Test", variable=test_var, command=callback)

checkbutton.pack()

root.mainloop()

<|endoftext|>


<commit_before>from django.contrib.auth.decorators import login_required
from django.core.cache              import cache
from django.core.cache.backends     import locmem
from django.http                    import HttpResponse, Http404
from django.shortcuts               import get_object_or_404, render, redirect
from django.template                import RequestContext, loader
from django.views.decorators.csrf   import csrf_exempt
from moveon.models                  import Company
import json
import osmlineadapters.settings as settings

@csrf_exempt
@login_required(login_url='moveon_login')
def newline(request, company_id):
    #Looking for the
    if request.method == 'POST':
        json_request = json.loads(request.body.decode("utf-8"))
        osm_line_id = json_request['osmline']['id']

        adapter_path = "osmlineadapters.adapters.{0}.osm_line".format(company_id)
        mod = __import__(adapter_path, fromlist=['OSMLine'])
        class_ = getattr(mod, 'OSMLine')

        instance = class_(osm_line_id)

        cache_id = osm_line_id
        cache_simplified_id = str(osm_line_id) + "_simple"
        cache.set(cache_id, instance.to_json(), 600)
        cache.set(cache_simplified_id, instance.to_simplified_json(), 600)
        return HttpResponse(request.body)

    return redirect('copany', company_id=company_id)

@login_required(login_url='moveon_login')
def newlinedetail(request, company_id, osm_line_id):

    cache_simplified_id = str(osm_line_id) + "_simple"
    if request.method == "GET":
        line = json.loads(cache.get(cache_simplified_id))
        if line:
            context = { 'line': line,
                        'company_id': company_id
                      }
            return render(request, 'new_line_detail.html', context)
        else:
            return HttpResponse("Line not retrieved yet")
    elif request.method == "POST":
        json_request = json.loads(request.body.decode("utf-8"))
        agree = bool(json_request['osmline']['accept'])
        if agree:
            line = json.loads(cache.get(osm_line_id))
            osmlinemanager = _get_osmlinemanager(line)
            osmlinemanager.save()
            status_code = 201
        else:
            status_code = 200

        cache.delete(osm_line_id)
        cache.delete(cache_simplified_id)
        return HttpResponse(status=status_code)

def _get_osmlinemanager(line):
    manager_module = settings.OSMLINEMANAGERMODULE
    manager_class_name = settings.OSMLINEMANAGERCLASS
    mod = __import__(manager_module, fromlist=[manager_class_name])
    class_ = getattr(mod, manager_class_name)

    instance = class_(line)
    return instance
<commit_msg>Removes unused imports<commit_after>from django.contrib.auth.decorators import login_required
from django.core.cache              import cache
from django.http                    import HttpResponse
from django.shortcuts               import render, redirect
from django.views.decorators.csrf   import csrf_exempt
import json
import osmlineadapters.settings as settings
from django.core.urlresolvers import reverse

@csrf_exempt
@login_required(login_url='moveon_login')
def newline(request, company_id):
    #Looking for the
    if request.method == 'POST':
        json_request = json.loads(request.body.decode("utf-8"))
        osm_line_id = json_request['osmline']['id']

        adapter_path = "osmlineadapters.adapters.{0}.osm_line".format(company_id)
        mod = __import__(adapter_path, fromlist=['OSMLine'])
        class_ = getattr(mod, 'OSMLine')

        instance = class_(osm_line_id)

        cache_id = osm_line_id
        cache_simplified_id = str(osm_line_id) + "_simple"
        cache.set(cache_id, instance.to_json(), 600)
        cache.set(cache_simplified_id, instance.to_simplified_json(), 600)
        return HttpResponse(request.body)

    return redirect('copany', company_id=company_id)

@login_required(login_url='moveon_login')
def newlinedetail(request, company_id, osm_line_id):

    cache_simplified_id = str(osm_line_id) + "_simple"
    if request.method == "GET":
        line = json.loads(cache.get(cache_simplified_id))
        if line:
            context = { 'line': line,
                        'company_id': company_id
                      }
            return render(request, 'new_line_detail.html', context)
        else:
            return HttpResponse("Line not retrieved yet")
    elif request.method == "POST":
        json_request = json.loads(request.body.decode("utf-8"))
        agree = bool(json_request['osmline']['accept'])
        if agree:
            line = json.loads(cache.get(osm_line_id))
            osmlinemanager = _get_osmlinemanager(line)
            osmlinemanager.save()
            status_code = 201
        else:
            status_code = 200

        cache.delete(osm_line_id)
        cache.delete(cache_simplified_id)
        return HttpResponse(status=status_code)

def _get_osmlinemanager(line):
    manager_module = settings.OSMLINEMANAGERMODULE
    manager_class_name = settings.OSMLINEMANAGERCLASS
    mod = __import__(manager_module, fromlist=[manager_class_name])
    class_ = getattr(mod, manager_class_name)

    instance = class_(line)
    return instance
<|endoftext|>