<commit_before>"""
Implementations of standard library functions, because it's not possible to
understand them with Jedi.

To add a new implementation, create a function and add it to the
``_implemented`` dict at the bottom of this module.

"""
import collections
from jedi._compatibility import unicode
from jedi.evaluate import compiled
from jedi.evaluate import representation as er
from jedi.evaluate import iterable
from jedi.evaluate.helpers import FakeArray, FakeStatement
from jedi.parser import Parser
from jedi.parser import representation as pr
from jedi import debug

class NotInStdLib(LookupError):
    pass

def execute(evaluator, obj, params):
    try:
        obj_name = str(obj.name)
    except AttributeError:
        pass
    else:
        if obj.parent == compiled.builtin:
            module_name = 'builtins'
        elif isinstance(obj.parent, pr.Module):
            module_name = str(obj.parent.name)
        else:
            module_name = ''

        # for now we just support builtin functions.
        try:
            return _implemented[module_name][obj_name](evaluator, obj, params)
        except KeyError:
            pass
    raise NotInStdLib()

def _follow_param(evaluator, params, index):
    try:
        stmt = params[index]
    except IndexError:
        return []
    else:
        if isinstance(stmt, pr.Statement):
            return evaluator.eval_statement(stmt)
        else:
            return [stmt]  # just some arbitrary object

def builtins_getattr(evaluator, obj, params):
    stmts = []
    # follow the first param
    objects = _follow_param(evaluator, params, 0)
    names = _follow_param(evaluator, params, 1)
    for obj in objects:
        if not isinstance(obj, (er.Instance, er.Class, pr.Module, compiled.CompiledObject)):
            debug.warning('getattr called without instance')
            continue

        for name in names:
            s = unicode, str
            if isinstance(name, compiled.CompiledObject) and isinstance(name.obj, s):
                stmts += evaluator.follow_path(iter([name.obj]), [obj], obj)
            else:
                debug.warning('getattr called without str')
                continue
    return stmts

def builtins_type(evaluator, obj, params):
    if len(params) == 1:
        # otherwise it would be a metaclass... maybe someday...
        objects = _follow_param(evaluator, params, 0)
        return [o.base for o in objects if isinstance(o, er.Instance)]
    return []

def builtins_super(evaluator, obj, params):
    # TODO make this able to detect multiple inheritance super
    accept = (pr.Function,)
    func = params.get_parent_until(accept)
    if func.isinstance(*accept):
        cls = func.get_parent_until(accept + (pr.Class,),
                                    include_current=False)
        if isinstance(cls, pr.Class):
            cls = er.Class(evaluator, cls)
            su = cls.get_super_classes()
            if su:
                return evaluator.execute(su[0])
    return []

def builtins_reversed(evaluator, obj, params):
    objects = tuple(_follow_param(evaluator, params, 0))
    if objects:
        # unpack the iterator values
        objects = tuple(iterable.get_iterator_types(objects))
        if objects:
            rev = reversed(objects)
            # Repack iterator values and then run it the normal way. This is
            # necessary, because `reversed` is a function and autocompletion
            # would fail in certain cases like `reversed(x).__iter__` if we
            # just returned the result directly.
            stmts = [FakeStatement([r]) for r in rev]
            objects = (iterable.Array(evaluator, FakeArray(stmts, objects[0].parent)),)
    return [er.Instance(evaluator, obj, objects)]

def collections_namedtuple(evaluator, obj, params):
    """
    Implementation of the namedtuple function.

    This has to be done by processing the namedtuple class template and
    evaluating the result.

    .. note:: |jedi| only supports namedtuples on Python >2.6.

    """
    # Namedtuples are not supported on Python 2.6
    if not hasattr(collections, '_class_template'):
        return []

    # Process arguments
    name = _follow_param(evaluator, params, 0)[0].obj
    _fields = _follow_param(evaluator, params, 1)[0]
    if isinstance(_fields, compiled.CompiledObject):
        fields = _fields.obj.replace(',', ' ').split()
    elif isinstance(_fields, iterable.Array):
        try:
            fields = [v.obj for v in _fields.values()]
        except AttributeError:
            return []
    else:
        return []

    # Build source
    source = collections._class_template.format(
        typename=name,
        field_names=fields,
        num_fields=len(fields),
        arg_list=', '.join(fields),
        repr_fmt=', '.join(collections._repr_template.format(name=name) for name in fields),
        field_defs='\n'.join(collections._field_template.format(index=index, name=name)
                             for index, name in enumerate(fields))
    )

    # Parse source
    generated_class = Parser(unicode(source)).module.subscopes[0]
    return [er.Class(evaluator, generated_class)]

def _return_first_param(evaluator, obj, params):
    if len(params) == 1:
        return _follow_param(evaluator, params, 0)
    return []

_implemented = {
    'builtins': {
        'getattr': builtins_getattr,
        'type': builtins_type,
        'super': builtins_super,
        'reversed': builtins_reversed,
    },
    'copy': {
        'copy': _return_first_param,
        'deepcopy': _return_first_param,
    },
    'json': {
        'load': lambda *args: [],
        'loads': lambda *args: [],
    },
    'collections': {
        'namedtuple': collections_namedtuple,
    },
}
<commit_msg>temporary SuperInstance class to eventually handle super. But need something like mro() resolution first.<commit_after>"""
Implementations of standard library functions, because it's not possible to
understand them with Jedi.

To add a new implementation, create a function and add it to the
``_implemented`` dict at the bottom of this module.

"""
import collections
from jedi._compatibility import unicode
from jedi.evaluate import compiled
from jedi.evaluate import representation as er
from jedi.evaluate import iterable
from jedi.evaluate.helpers import FakeArray, FakeStatement
from jedi.parser import Parser
from jedi.parser import representation as pr
from jedi import debug

class NotInStdLib(LookupError):
    pass

def execute(evaluator, obj, params):
    try:
        obj_name = str(obj.name)
    except AttributeError:
        pass
    else:
        if obj.parent == compiled.builtin:
            module_name = 'builtins'
        elif isinstance(obj.parent, pr.Module):
            module_name = str(obj.parent.name)
        else:
            module_name = ''

        # for now we just support builtin functions.
        try:
            return _implemented[module_name][obj_name](evaluator, obj, params)
        except KeyError:
            pass
    raise NotInStdLib()

def _follow_param(evaluator, params, index):
    try:
        stmt = params[index]
    except IndexError:
        return []
    else:
        if isinstance(stmt, pr.Statement):
            return evaluator.eval_statement(stmt)
        else:
            return [stmt]  # just some arbitrary object

def builtins_getattr(evaluator, obj, params):
    stmts = []
    # follow the first param
    objects = _follow_param(evaluator, params, 0)
    names = _follow_param(evaluator, params, 1)
    for obj in objects:
        if not isinstance(obj, (er.Instance, er.Class, pr.Module, compiled.CompiledObject)):
            debug.warning('getattr called without instance')
            continue

        for name in names:
            s = unicode, str
            if isinstance(name, compiled.CompiledObject) and isinstance(name.obj, s):
                stmts += evaluator.follow_path(iter([name.obj]), [obj], obj)
            else:
                debug.warning('getattr called without str')
                continue
    return stmts

def builtins_type(evaluator, obj, params):
    if len(params) == 1:
        # otherwise it would be a metaclass... maybe someday...
        objects = _follow_param(evaluator, params, 0)
        return [o.base for o in objects if isinstance(o, er.Instance)]
    return []

class SuperInstance(er.Instance):
    """To be used like the object ``super`` returns."""
    def __init__(self, evaluator, cls):
        su = cls.get_super_classes()
        super().__init__(evaluator, su and su[0] or self)

def builtins_super(evaluator, obj, params):
    # TODO make this able to detect multiple inheritance super
    accept = (pr.Function,)
    func = params.get_parent_until(accept)
    if func.isinstance(*accept):
        cls = func.get_parent_until(accept + (pr.Class,),
                                    include_current=False)
        if isinstance(cls, pr.Class):
            cls = er.Class(evaluator, cls)
            su = cls.get_super_classes()
            if su:
                return evaluator.execute(su[0])
    return []

def builtins_reversed(evaluator, obj, params):
    objects = tuple(_follow_param(evaluator, params, 0))
    if objects:
        # unpack the iterator values
        objects = tuple(iterable.get_iterator_types(objects))
        if objects:
            rev = reversed(objects)
            # Repack iterator values and then run it the normal way. This is
            # necessary, because `reversed` is a function and autocompletion
            # would fail in certain cases like `reversed(x).__iter__` if we
            # just returned the result directly.
            stmts = [FakeStatement([r]) for r in rev]
            objects = (iterable.Array(evaluator, FakeArray(stmts, objects[0].parent)),)
    return [er.Instance(evaluator, obj, objects)]

def collections_namedtuple(evaluator, obj, params):
    """
    Implementation of the namedtuple function.

    This has to be done by processing the namedtuple class template and
    evaluating the result.

    .. note:: |jedi| only supports namedtuples on Python >2.6.

    """
    # Namedtuples are not supported on Python 2.6
    if not hasattr(collections, '_class_template'):
        return []

    # Process arguments
    name = _follow_param(evaluator, params, 0)[0].obj
    _fields = _follow_param(evaluator, params, 1)[0]
    if isinstance(_fields, compiled.CompiledObject):
        fields = _fields.obj.replace(',', ' ').split()
    elif isinstance(_fields, iterable.Array):
        try:
            fields = [v.obj for v in _fields.values()]
        except AttributeError:
            return []
    else:
        return []

    # Build source
    source = collections._class_template.format(
        typename=name,
        field_names=fields,
        num_fields=len(fields),
        arg_list=', '.join(fields),
        repr_fmt=', '.join(collections._repr_template.format(name=name) for name in fields),
        field_defs='\n'.join(collections._field_template.format(index=index, name=name)
                             for index, name in enumerate(fields))
    )

    # Parse source
    generated_class = Parser(unicode(source)).module.subscopes[0]
    return [er.Class(evaluator, generated_class)]

def _return_first_param(evaluator, obj, params):
    if len(params) == 1:
        return _follow_param(evaluator, params, 0)
    return []

_implemented = {
    'builtins': {
        'getattr': builtins_getattr,
        'type': builtins_type,
        'super': builtins_super,
        'reversed': builtins_reversed,
    },
    'copy': {
        'copy': _return_first_param,
        'deepcopy': _return_first_param,
    },
    'json': {
        'load': lambda *args: [],
        'loads': lambda *args: [],
    },
    'collections': {
        'namedtuple': collections_namedtuple,
    },
}
<|endoftext|>


<commit_before>import pandas as pd
import csv

from context import campaignadvisor

"""
Data cleanup utility

May eventually be used to store cleaned tables instead.
"""

# Zip code file-specific constants
# Replacement text for when there is no county name
NO_COUNTY_NAME = "NO_COUNTY_NAME"
# Replacement text for when there is no state name
NO_STATE_NAME = "NO_STATE_NAME"
# Replacement text for when there is no fips code
NO_FIPS_CODE = "0"
# Extraneous words in the zip code file
ZIP_STOP_WORDS = ['County', 'City']
# Length of zip codes provided in zip code file
SHORT_ZIP_CODE_LENGTH = 5
# Length of zip codes that can occur in contributions
LONG_ZIP_CODE_LENGTH = 9

# References points for field names in datasets
ZIP_CODES_FIELD_NAMES = ["zip", "type", "primary_city", "acceptable_cities", "unacceptable_cities", "state", "county",
                         "timezone", "area_codes", "latitude", "longitude", "world_region", "country", "decommissioned",
                         "estimated_population", "notes"]

CONTRIBUTIONS_FIELD_NAMES = ["cmte_id", "cand_id", "cand_nm", "contbr_nm", "contbr_city", "contbr_st", "contbr_zip",
                             "contbr_employer", "contbr_occupation", "contb_receipt_amt", "contb_receipt_dt",
                             "receipt_desc", "memo_cd", "memo_text", "form_tp", "file_num", "tran_id", "election_tp"]

_STATE_ABBREVIATIONS = states = {
    'AK': 'Alaska',
    'AL': 'Alabama',
    'AR': 'Arkansas',
    'AS': 'American Samoa',
    'AZ': 'Arizona',
    'CA': 'California',
    'CO': 'Colorado',
    'CT': 'Connecticut',
    'DC': 'District of Columbia',
    'DE': 'Delaware',
    'FL': 'Florida',
    'GA': 'Georgia',
    'GU': 'Guam',
    'HI': 'Hawaii',
    'IA': 'Iowa',
    'ID': 'Idaho',
    'IL': 'Illinois',
    'IN': 'Indiana',
    'KS': 'Kansas',
    'KY': 'Kentucky',
    'LA': 'Louisiana',
    'MA': 'Massachusetts',
    'MD': 'Maryland',
    'ME': 'Maine',
    'MI': 'Michigan',
    'MN': 'Minnesota',
    'MO': 'Missouri',
    'MP': 'Northern Mariana Islands',
    'MS': 'Mississippi',
    'MT': 'Montana',
    'NA': 'National',
    'NC': 'North Carolina',
    'ND': 'North Dakota',
    'NE': 'Nebraska',
    'NH': 'New Hampshire',
    'NJ': 'New Jersey',
    'NM': 'New Mexico',
    'NV': 'Nevada',
    'NY': 'New York',
    'OH': 'Ohio',
    'OK': 'Oklahoma',
    'OR': 'Oregon',
    'PA': 'Pennsylvania',
    'PR': 'Puerto Rico',
    'RI': 'Rhode Island',
    'SC': 'South Carolina',
    'SD': 'South Dakota',
    'TN': 'Tennessee',
    'TX': 'Texas',
    # NOTE: There are 9 minor outlying islands, collisions may occur
    'UM': 'United States Minor Outlying Islands',
    'UT': 'Utah',
    'VA': 'Virginia',
    'VI': 'Virgin Islands',
    'VT': 'Vermont',
    'WA': 'Washington',
    'WI': 'Wisconsin',
    'WV': 'West Virginia',
    'WY': 'Wyoming'
}

def _get_state_name(state_abbreviation):
    if state_abbreviation in _STATE_ABBREVIATIONS:
        return _STATE_ABBREVIATIONS[state_abbreviation]
    else:
        return NO_STATE_NAME

class FIPSMapper:
    def __init__(self, fips_to_state_county, fips_to_state, state_county_to_fips, state_to_fips):
        self.fips_to_state_county = fips_to_state_county
        self.fips_to_state = fips_to_state
        self.state_county_to_fips = state_county_to_fips
        self.state_to_fips = state_to_fips

# Global variables (strictly for optimization purposes)
the_zip_code_map = None
the_fips_mapper = None

def _get_the_zip_code_map():
    global the_zip_code_map
    if the_zip_code_map is None:
        the_zip_code_map = _get_zip_codes_map()
    return the_zip_code_map

def _get_the_fips_mapper():
    global the_fips_mapper
    if the_fips_mapper is None:
        the_fips_mapper = _get_fips_mapper()
    return the_fips_mapper

def _get_zip_codes_map():
    # Mapping of zipcode to county name
    zip_codes = {}
    zip_codes_resource = campaignadvisor.resources.get_resource("zip_code_database.csv")
    with zip_codes_resource.get_file() as csvfile:
        reader = csv.DictReader(csvfile, ZIP_CODES_FIELD_NAMES)
        for row in reader:
            # Look out for types, they're different everywhere and need to be consistent
            zipcode = str(row['zip'])
            state_abbreviation = str(row['state'])
            state_name = _get_state_name(state_abbreviation)
            county = str(row['county'])
            # Do not want empty strings in the map
            # Cannot split no length strings
            if len(county) == 0:
                continue
            """
            split = county.split()
            # TODO: will cause errors when city is valid.
            # TODO: add city whitelist, perhaps don't use it at all
            # Throw out stop words at the end of the county name
            if split[-1] in ZIP_STOP_WORDS:
                county = " ".join(split[:-1]).strip()
            """
            state_county = state_name, county
            zip_codes[zipcode] = state_county
    return zip_codes

def get_state_county_from_zip_code(zip_code):
    """
    Given a zip code, return its corresponding (state, county) tuple
    :param zip_code: A zip code of length 5 or 9
    :return: The zip code's corresponding (state's name, county's name)
    """
    zip_codes_map = _get_the_zip_code_map()
    # Accept any type of zip code representation (numpy.float64, int, str...etc)
    # Throws away trailing decimals that are unwanted
    try:
        zip_code = str(int(zip_code))
    except ValueError:
        return NO_STATE_NAME, NO_COUNTY_NAME

    # Convert long zip codes to short
    if len(zip_code) == LONG_ZIP_CODE_LENGTH:
        zip_code = zip_code[:SHORT_ZIP_CODE_LENGTH]
    if zip_code in zip_codes_map:
        return zip_codes_map[zip_code]
    else:
        return NO_STATE_NAME, NO_COUNTY_NAME

FIPS_FIELD_NAMES = ["state_abbreviation", "state_fips", "county_fips", "county_name", "fips_class"]

# TODO: standardize what county names will be
def _get_fips_mapper():
    national_counties_resource = campaignadvisor.resources.get_resource("national_counties.csv")
    fips_to_state_county = {}
    fips_to_state = {}
    state_county_to_fips = {}
    state_to_fips = {}
    with national_counties_resource.get_file() as csvfile:
        reader = csv.DictReader(csvfile, FIPS_FIELD_NAMES)
        for row in reader:
            # Look out for types, they're different everywhere and need to be consistent
            state_abbreviation = str(row["state_abbreviation"])
            state_name = _get_state_name(state_abbreviation)
            state_fips = str(row["state_fips"])
            county_fips = state_fips + str(row["county_fips"])
            county_name = str(row["county_name"])

            state_county = (state_name, county_name)

            # TODO:  For all below dictionaries, add error handling for collisions. Should never happen.
            # Build fips_to_state_county
            if county_fips not in fips_to_state_county:
                # TODO: ensure consistent naming convention for counties. This presently lists Fairfax as Fairfax County
                fips_to_state_county[county_fips] = state_county

            # Build fips_to_state
            if state_fips not in fips_to_state:
                fips_to_state[state_fips] = state_name

            # Build state_county_to_fips
            if state_county not in state_county_to_fips:
                state_county_to_fips[state_county] = county_fips

                # Build state_to_fips:
                if state_name not in state_to_fips:
                    state_to_fips[state_name] = state_fips
        return FIPSMapper(fips_to_state_county, fips_to_state, state_county_to_fips, state_to_fips)

def get_state_county_from_fips(fips_code):
    """
    Given a county's fips code , return a tuple of its state's properly capitalized name followed by
    its properly capitalized name.

    NOTE: The formatting of the name returned is currently "Alger County" rather than "Alger"
    This is subject to change.

    :param fips_code: A county's fips code
    :return: The tuple: (properly capitalized state name, properly capitalized county name)
    """
    fips_to_state_county = _get_the_fips_mapper().fips_to_state_county
    if fips_code in fips_to_state_county:
        return fips_to_state_county[fips_code]
    else:
        return NO_STATE_NAME, NO_COUNTY_NAME

def get_state_from_fips(fips_code):
    """
    Given a state's fips code , return its properly capitalized name
    :param fips_code: A state's fips code
    :return: A state's properly capitalized name
    """
    fips_to_state = _get_the_fips_mapper().fips_to_state
    if fips_code in fips_to_state:
        return fips_to_state[fips_code]
    else:
        return NO_STATE_NAME

# TODO: Standardize how which will be the county name "Alger County" or "Alger"
# TODO: Perhaps add case insensitivity?

def get_fips_from_state_county(state_county):
    """
    Get the fips code of a county in a given state

    NOTE: This function expects input in the form "COUNTY_NAME + County"
    This implementation may vary

    NOTE: If a state were to have multiple counties with the same name, the return value is not guaranteed to be correct
    NOTE: Some states have municipalities with the same name as the county (Fairfax vs Fairfax County)
    This likely will resolve correctly to the county name, but this is not guaranteed

    :param state_county: A tuple containing the properly capitalized (state, county_ ie:  ("Maryland", "Fairfax County")
    :return: The fips code of the desired city
    """
    # Convert to tuple in case was passed in as a list etc
    state_county = state_county[0], state_county[1]
    state_county_to_fips = _get_the_fips_mapper().state_county_to_fips
    if state_county in state_county_to_fips:
        return state_county_to_fips[state_county]
    else:
        return NO_FIPS_CODE

def get_fips_from_state(state):
    """
    Returns the fips code of a given state.
    :param state: The properly capitalized name of a state
    :return: The state's fips code
    """
    state_to_fips = _get_the_fips_mapper().state_to_fips
    if state in state_to_fips:
        return state_to_fips[state]
    else:
        return NO_FIPS_CODE

def example_zip_to_county():
    # Rename CONTRIBUTIONS_PATH if you would like to use a smaller test file
    # CONTRIBUTIONS_PATH = os.path.join(DATA_DIR_PATH, "VA.csv")
    contributions_resource = campaignadvisor.resources.get_resource("contributions.csv")
    contributions = pd.read_csv(contributions_resource.get_local_path())
    # NOTE: keyword argument is not optional. The apply method is just a little hanky
    # Full blown code smell, probably a better way.
    print contributions["contbr_zip"].apply(get_state_county_from_zip_code)

def debug():
    starting_state_county = ("Virginia", "Fairfax County")
    fips_code = get_fips_from_state_county(starting_state_county)
    state_county = get_state_county_from_fips(fips_code)
    state = state_county[0]
    state_fips = get_fips_from_state(state)

    print "Start:", starting_state_county
    print "County fips:", fips_code
    print "State County found:", state_county
    print "State fips:", state_fips

    # Baltimore city
    print get_state_county_from_fips("24510")

if __name__ == "__main__":
    debug()
    example_zip_to_county()<commit_msg>Clean all occurences of state or county<commit_after>import pandas as pd
import csv

from context import campaignadvisor

"""
Data cleanup utility

May eventually be used to store cleaned tables instead.
"""

# Zip code file-specific constants
# Replacement text for when there is no county name
NO_COUNTY_NAME = "NO_COUNTY_NAME"
# Replacement text for when there is no state name
NO_STATE_NAME = "NO_STATE_NAME"
# Replacement text for when there is no fips code
NO_FIPS_CODE = "0"
# Extraneous words to be removed from county names
# NOTE: stop words that are substrings of other stop words must appear AFTER them in this list.
COUNTY_STOP_WORDS = [
    'city and municipality of',
    'city of',
    'municipality of',
    'census area',
    'borough',
    'city',
    'county',
]
# Length of zip codes provided in zip code file
SHORT_ZIP_CODE_LENGTH = 5
# Length of zip codes that can occur in contributions
LONG_ZIP_CODE_LENGTH = 9

# References points for field names in datasets
ZIP_CODES_FIELD_NAMES = ["zip", "type", "primary_city", "acceptable_cities", "unacceptable_cities", "state", "county",
                         "timezone", "area_codes", "latitude", "longitude", "world_region", "country", "decommissioned",
                         "estimated_population", "notes"]

CONTRIBUTIONS_FIELD_NAMES = ["cmte_id", "cand_id", "cand_nm", "contbr_nm", "contbr_city", "contbr_st", "contbr_zip",
                             "contbr_employer", "contbr_occupation", "contb_receipt_amt", "contb_receipt_dt",
                             "receipt_desc", "memo_cd", "memo_text", "form_tp", "file_num", "tran_id", "election_tp"]

_STATE_ABBREVIATIONS = states = {
    'AK': 'Alaska',
    'AL': 'Alabama',
    'AR': 'Arkansas',
    'AS': 'American Samoa',
    'AZ': 'Arizona',
    'CA': 'California',
    'CO': 'Colorado',
    'CT': 'Connecticut',
    'DC': 'District of Columbia',
    'DE': 'Delaware',
    'FL': 'Florida',
    'GA': 'Georgia',
    'GU': 'Guam',
    'HI': 'Hawaii',
    'IA': 'Iowa',
    'ID': 'Idaho',
    'IL': 'Illinois',
    'IN': 'Indiana',
    'KS': 'Kansas',
    'KY': 'Kentucky',
    'LA': 'Louisiana',
    'MA': 'Massachusetts',
    'MD': 'Maryland',
    'ME': 'Maine',
    'MI': 'Michigan',
    'MN': 'Minnesota',
    'MO': 'Missouri',
    'MP': 'Northern Mariana Islands',
    'MS': 'Mississippi',
    'MT': 'Montana',
    'NA': 'National',
    'NC': 'North Carolina',
    'ND': 'North Dakota',
    'NE': 'Nebraska',
    'NH': 'New Hampshire',
    'NJ': 'New Jersey',
    'NM': 'New Mexico',
    'NV': 'Nevada',
    'NY': 'New York',
    'OH': 'Ohio',
    'OK': 'Oklahoma',
    'OR': 'Oregon',
    'PA': 'Pennsylvania',
    'PR': 'Puerto Rico',
    'RI': 'Rhode Island',
    'SC': 'South Carolina',
    'SD': 'South Dakota',
    'TN': 'Tennessee',
    'TX': 'Texas',
    # NOTE: There are 9 minor outlying islands, collisions may occur
    'UM': 'United States Minor Outlying Islands',
    'UT': 'Utah',
    'VA': 'Virginia',
    'VI': 'Virgin Islands',
    'VT': 'Vermont',
    'WA': 'Washington',
    'WI': 'Wisconsin',
    'WV': 'West Virginia',
    'WY': 'Wyoming'
}

def _get_state_name(state_abbreviation):
    if state_abbreviation in _STATE_ABBREVIATIONS:
        return _STATE_ABBREVIATIONS[state_abbreviation]
    else:
        return NO_STATE_NAME

# TODO: add whitelist when should actually read Kansas City (if that were a real county)
def _get_clean_county_name(dirty_county_name):
    clean_county_name = dirty_county_name.lower()
    for stop_word in COUNTY_STOP_WORDS:
        if stop_word in clean_county_name:
            clean_county_name = clean_county_name.replace(stop_word, "").strip()
    return clean_county_name

def _get_clean_state_name(dirty_state_name):
    return dirty_state_name.lower()

class FIPSMapper:
    def __init__(self, fips_to_state_county, fips_to_state, state_county_to_fips, state_to_fips):
        self.fips_to_state_county = fips_to_state_county
        self.fips_to_state = fips_to_state
        self.state_county_to_fips = state_county_to_fips
        self.state_to_fips = state_to_fips

# Global variables (strictly for optimization purposes)
the_zip_code_map = None
the_fips_mapper = None

def _get_the_zip_code_map():
    global the_zip_code_map
    if the_zip_code_map is None:
        the_zip_code_map = _get_zip_codes_map()
    return the_zip_code_map

def _get_the_fips_mapper():
    global the_fips_mapper
    if the_fips_mapper is None:
        the_fips_mapper = _get_fips_mapper()
    return the_fips_mapper

def _get_zip_codes_map():
    # Mapping of zipcode to county name
    zip_codes = {}
    zip_codes_resource = campaignadvisor.resources.get_resource("zip_code_database.csv")
    with zip_codes_resource.get_file() as csvfile:
        reader = csv.DictReader(csvfile, ZIP_CODES_FIELD_NAMES)
        for row in reader:
            # Look out for types, they're different everywhere and need to be consistent
            zipcode = str(row['zip'])
            state_abbreviation = str(row['state'])
            state_name = _get_state_name(state_abbreviation)
            # Put to lower case for consistency
            state_name = _get_clean_state_name(state_name)
            # Throw away city, borough, county, etc change to lower case
            county = _get_clean_county_name(str(row['county']))
            # Do not want empty strings in the map
            if len(county) == 0:
                continue

            state_county = state_name, county
            zip_codes[zipcode] = state_county
    return zip_codes

def get_state_county_from_zip_code(zip_code):
    """
    Given a zip code, return its corresponding (state, county) tuple
    :param zip_code: A zip code of length 5 or 9
    :return: The zip code's corresponding (state's name, county's name)
    """
    zip_codes_map = _get_the_zip_code_map()
    # Accept any type of zip code representation (numpy.float64, int, str...etc)
    # Throws away trailing decimals that are unwanted
    try:
        zip_code = str(int(zip_code))
    except ValueError:
        return NO_STATE_NAME, NO_COUNTY_NAME

    # Convert long zip codes to short
    if len(zip_code) == LONG_ZIP_CODE_LENGTH:
        zip_code = zip_code[:SHORT_ZIP_CODE_LENGTH]
    if zip_code in zip_codes_map:
        return zip_codes_map[zip_code]
    else:
        return NO_STATE_NAME, NO_COUNTY_NAME

FIPS_FIELD_NAMES = ["state_abbreviation", "state_fips", "county_fips", "county_name", "fips_class"]

# TODO: standardize what county names will be
def _get_fips_mapper():
    national_counties_resource = campaignadvisor.resources.get_resource("national_counties.csv")
    fips_to_state_county = {}
    fips_to_state = {}
    state_county_to_fips = {}
    state_to_fips = {}
    with national_counties_resource.get_file() as csvfile:
        reader = csv.DictReader(csvfile, FIPS_FIELD_NAMES)
        for row in reader:
            # Look out for types, they're different everywhere and need to be consistent
            state_abbreviation = str(row["state_abbreviation"])
            state_name = _get_state_name(state_abbreviation)
            state_name = _get_clean_state_name(state_name)
            state_fips = str(row["state_fips"])
            county_fips = state_fips + str(row["county_fips"])
            county_name = str(row["county_name"])
            county_name = _get_clean_county_name(county_name)

            state_county = (state_name, county_name)

            # TODO:  For all below dictionaries, add error handling for collisions. Should never happen.
            # Build fips_to_state_county
            if county_fips not in fips_to_state_county:
                # TODO: ensure consistent naming convention for counties. This presently lists Fairfax as Fairfax County
                fips_to_state_county[county_fips] = state_county

            # Build fips_to_state
            if state_fips not in fips_to_state:
                fips_to_state[state_fips] = state_name

            # Build state_county_to_fips
            if state_county not in state_county_to_fips:
                state_county_to_fips[state_county] = county_fips

                # Build state_to_fips:
                if state_name not in state_to_fips:
                    state_to_fips[state_name] = state_fips
        return FIPSMapper(fips_to_state_county, fips_to_state, state_county_to_fips, state_to_fips)

def get_state_county_from_fips(fips_code):
    """
    Given a county's fips code , return a tuple of its state's properly capitalized name followed by
    its properly capitalized name.

    NOTE: The formatting of the name returned is currently "Alger County" rather than "Alger"
    This is subject to change.

    :param fips_code: A county's fips code
    :return: The tuple: (properly capitalized state name, properly capitalized county name)
    """
    fips_to_state_county = _get_the_fips_mapper().fips_to_state_county
    if fips_code in fips_to_state_county:
        return fips_to_state_county[fips_code]
    else:
        return NO_STATE_NAME, NO_COUNTY_NAME

def get_state_from_fips(fips_code):
    """
    Given a state's fips code , return its properly capitalized name
    :param fips_code: A state's fips code
    :return: A state's properly capitalized name
    """
    fips_to_state = _get_the_fips_mapper().fips_to_state
    if fips_code in fips_to_state:
        return fips_to_state[fips_code]
    else:
        return NO_STATE_NAME

# TODO: Standardize how which will be the county name "Alger County" or "Alger"
# TODO: Perhaps add case insensitivity?

def get_fips_from_state_county(state_county):
    """
    Get the fips code of a county in a given state

    NOTE: This function expects input in the form "COUNTY_NAME + County"
    This implementation may vary

    NOTE: If a state were to have multiple counties with the same name, the return value is not guaranteed to be correct
    NOTE: Some states have municipalities with the same name as the county (Fairfax vs Fairfax County)
    This likely will resolve correctly to the county name, but this is not guaranteed

    :param state_county: A tuple containing the properly capitalized (state, county_ ie:  ("Maryland", "Fairfax County")
    :return: The fips code of the desired city
    """
    # Convert to tuple in case was passed in as a list etc
    # Convert to lowercase to support lookup
    state_county = _get_clean_state_name(state_county[0]), _get_clean_county_name(state_county[1])
    state_county_to_fips = _get_the_fips_mapper().state_county_to_fips
    if state_county in state_county_to_fips:
        return state_county_to_fips[state_county]
    else:
        return NO_FIPS_CODE

def get_fips_from_state(state):
    """
    Returns the fips code of a given state.
    :param state: The properly capitalized name of a state
    :return: The state's fips code
    """
    # Lowercase for lookup
    state = _get_clean_state_name(state)
    state_to_fips = _get_the_fips_mapper().state_to_fips
    if state in state_to_fips:
        return state_to_fips[state]
    else:
        return NO_FIPS_CODE

def example_zip_to_county():
    # Rename CONTRIBUTIONS_PATH if you would like to use a smaller test file
    # CONTRIBUTIONS_PATH = os.path.join(DATA_DIR_PATH, "VA.csv")
    contributions_resource = campaignadvisor.resources.get_resource("contributions.csv")
    contributions = pd.read_csv(contributions_resource.get_local_path())
    # NOTE: keyword argument is not optional. The apply method is just a little hanky
    # Full blown code smell, probably a better way.
    print contributions["contbr_zip"].apply(get_state_county_from_zip_code)

def debug():
    starting_state_county = ("Virginia", "Fairfax County")
    fips_code = get_fips_from_state_county(starting_state_county)
    state_county = get_state_county_from_fips(fips_code)
    state = state_county[0]
    state_fips = get_fips_from_state(state)

    print "Start:", starting_state_county
    print "County fips:", fips_code
    print "State County found:", state_county
    print "State fips:", state_fips

    # Baltimore city
    print get_state_county_from_fips("24510")

if __name__ == "__main__":
    debug()
    example_zip_to_county()<|endoftext|>


<commit_before>import logging
from enum import Enum

from django.contrib.auth.models import User
from django.contrib.gis.db import models
from django.utils import timezone
from django.urls import reverse
from django.template.defaulttags import register

from emstrack.latlon import calculate_orientation, calculate_distance, stationary_radius
from emstrack.models import AddressModel, UpdatedByModel, defaults

logger = logging.getLogger(__name__)

# filters

@register.filter
def get_ambulance_status(key):
    return AmbulanceStatus[key].value

@register.filter
def get_ambulance_capability(key):
    return AmbulanceCapability[key].value

@register.filter
def get_location_type(key):
    return LocationType[key].value

@register.filter
def get_call_status(key):
    return CallStatus[key].value

@register.filter
def get_call_priority(key):
    return CallPriority[key].value

# Ambulance location models

# Ambulance model

class AmbulanceStatus(Enum):
    UK = 'Unknown'
    AV = 'Available'
    OS = 'Out of service'
    PB = 'Patient bound'
    AP = 'At patient'
    HB = 'Hospital bound'
    AH = 'At hospital'

class AmbulanceCapability(Enum):
    B = 'Basic'
    A = 'Advanced'
    R = 'Rescue'

class Ambulance(UpdatedByModel):

    # TODO: Should we consider denormalizing Ambulance to avoid duplication with AmbulanceUpdate?

    # ambulance properties
    identifier = models.CharField(max_length=50, unique=True)

    # TODO: Should we add an active flag?

    AMBULANCE_CAPABILITY_CHOICES = \
        [(m.name, m.value) for m in AmbulanceCapability]
    capability = models.CharField(max_length=1,
                                  choices=AMBULANCE_CAPABILITY_CHOICES)

    # status
    AMBULANCE_STATUS_CHOICES = \
        [(m.name, m.value) for m in AmbulanceStatus]
    status = models.CharField(max_length=2,
                              choices=AMBULANCE_STATUS_CHOICES,
                              default=AmbulanceStatus.UK.name)

    # location
    orientation = models.FloatField(default=0.0)
    location = models.PointField(srid=4326, default=defaults['location'])

    # timestamp
    timestamp = models.DateTimeField(default=timezone.now)

    # location client
    location_client = models.ForeignKey('login.Client',
                                        on_delete=models.CASCADE,
                                        blank=True, null=True,
                                        related_name='location_client_set')

    # default value for _loaded_values
    _loaded_values = None

    @classmethod
    def from_db(cls, db, field_names, values):

        # call super
        instance = super(Ambulance, cls).from_db(db, field_names, values)

        # store the original field values on the instance
        instance._loaded_values = dict(zip(field_names, values))

        # return instance
        return instance

    def save(self, *args, **kwargs):

        # creation?
        created = self.pk is None

        # loaded_values?
        loaded_values = self._loaded_values is not None

        # has location changed?
        has_moved = False
        if (not loaded_values) or \
                calculate_distance(self._loaded_values['location'], self.location) > stationary_radius:
            has_moved = True

        # calculate orientation only if location has changed and orientation has not changed
        if has_moved and loaded_values and self._loaded_values['orientation'] == self.orientation:
            # TODO: should we allow for a small radius before updating direction?
            self.orientation = calculate_orientation(self._loaded_values['location'], self.location)
            logger.debug('< {} - {} = {}'.format(self._loaded_values['location'],
                                                 self.location,
                                                 self.orientation))

        logger.debug('loaded_values: {}'.format(loaded_values))
        logger.debug('_loaded_values: {}'.format(self._loaded_values))
        logger.debug('self.location_client: {}'.format(self.location_client))

        # location_client changed?
        if self.location_client is None:
            location_client_id = None
        else:
            location_client_id = self.location_client.id
        location_client_changed = False
        if loaded_values and location_client_id != self._loaded_values['location_client_id']:
            location_client_changed = True

        logger.debug('location_client_changed: {}'.format(location_client_changed))
        # TODO: Check if client is logged with ambulance if setting location_client

        # if comment, status or location changed
        model_changed = False
        if has_moved or \
                self._loaded_values['status'] != self.status or \
                self._loaded_values['comment'] != self.comment:

            # save to Ambulance
            super().save(*args, **kwargs)

            logger.debug('SAVED')

            # save to AmbulanceUpdate
            data = {k: getattr(self, k)
                    for k in ('status', 'orientation',
                              'location', 'timestamp',
                              'comment', 'updated_by', 'updated_on')}
            data['ambulance'] = self
            obj = AmbulanceUpdate(**data)
            obj.save()

            logger.debug('UPDATE SAVED')

            # model changed
            model_changed = True

        # if identifier or capability changed
        # NOTE: self._loaded_values is NEVER None because has_moved is True
        elif (location_client_changed or
              self._loaded_values['identifier'] != self.identifier or
              self._loaded_values['capability'] != self.capability):

            # save only to Ambulance
            super().save(*args, **kwargs)

            logger.debug('SAVED')

            # model changed
            model_changed = True

        # Did the model change?
        if model_changed:

            # publish to mqtt
            from mqtt.publish import SingletonPublishClient
            SingletonPublishClient().publish_ambulance(self)

            logger.debug('PUBLISHED ON MQTT')

        # just created?
        if created:
            # invalidate permissions cache
            from login.permissions import cache_clear
            cache_clear()

    def delete(self, *args, **kwargs):

        # remove from mqtt
        from mqtt.publish import SingletonPublishClient
        SingletonPublishClient().remove_ambulance(self)

        # invalidate permissions cache
        from login.permissions import cache_clear
        cache_clear()

        # delete from Ambulance
        super().delete(*args, **kwargs)

    def get_absolute_url(self):
        return reverse('ambulance:detail', kwargs={'pk': self.id})

    def __str__(self):
        return ('Ambulance {}(id={}) ({}) [{}]:\n' +
                '    Status: {}\n' +
                '  Location: {} @ {}\n' +
                ' LocClient: {}\n' +
                '   Updated: {} by {}').format(self.identifier,
                                               self.id,
                                               AmbulanceCapability[self.capability].value,
                                               self.comment,
                                               AmbulanceStatus[self.status].value,
                                               self.location,
                                               self.timestamp,
                                               self.location_client,
                                               self.updated_by,
                                               self.updated_on)

# Call related models

class CallPriority(Enum):
    A = 'Resucitation'
    B = 'Emergent'
    C = 'Urgent'
    D = 'Less urgent'
    E = 'Not urgent'
    O = 'Omega'

class CallStatus(Enum):
    P = 'Pending'
    S = 'Started'
    E = 'Ended'

class CallPublishMixin:

    def save(self, *args, **kwargs):

        # publish?
        publish = kwargs.pop('publish', True)

        # save to Call
        super().save(*args, **kwargs)

        if publish:
            self.publish()

class Call(CallPublishMixin,
           AddressModel,
           UpdatedByModel):

    # status
    CALL_STATUS_CHOICES = \
        [(m.name, m.value) for m in CallStatus]
    status = models.CharField(max_length=1,
                              choices=CALL_STATUS_CHOICES,
                              default=CallStatus.P.name)

    # details
    details = models.CharField(max_length=500, default="")

    # call priority
    CALL_PRIORITY_CHOICES = \
        [(m.name, m.value) for m in CallPriority]
    priority = models.CharField(max_length=1,
                                choices=CALL_PRIORITY_CHOICES,
                                default=CallPriority.E.name)

    # timestamps
    pending_at = models.DateTimeField(null=True, blank=True)
    started_at = models.DateTimeField(null=True, blank=True)
    ended_at = models.DateTimeField(null=True, blank=True)

    # created at
    created_at = models.DateTimeField(auto_now_add=True)

    def save(self, *args, **kwargs):

        # publish?
        publish = kwargs.pop('publish', True)

        if self.status == CallStatus.E.name:
            # remove topic from mqtt server
            from mqtt.publish import SingletonPublishClient
            SingletonPublishClient().remove_call(self)

            # prevent publication
            publish = False

        # call super
        super().save(*args, **kwargs, publish=publish)

    def publish(self):

        # publish to mqtt
        from mqtt.publish import SingletonPublishClient
        SingletonPublishClient().publish_call(self)

    def __str__(self):
        return "{} ({})".format(self.location, self.priority)

class AmbulanceCallStatus(Enum):
    R = 'Requested'
    O = 'Ongoing'
    I = 'Interrupted'
    C = 'Completed'

class AmbulanceCall(CallPublishMixin,
                    models.Model):

    # status
    AMBULANCE_CALL_STATUS_CHOICES = \
        [(m.name, m.value) for m in AmbulanceCallStatus]
    status = models.CharField(max_length=1,
                              choices=AMBULANCE_CALL_STATUS_CHOICES,
                              default=AmbulanceCallStatus.R.name)

    # call
    call = models.ForeignKey(Call,
                             on_delete=models.CASCADE)

    # ambulance
    ambulance = models.ForeignKey(Ambulance,
                                  on_delete=models.CASCADE)

    # created at
    created_at = models.DateTimeField(auto_now_add=True)

    def save(self, *args, **kwargs):

        # publish?
        publish = kwargs.pop('publish', True)

        # changed to ongoing?
        if self.status == AmbulanceCallStatus.O.name:

            # retrieve call
            call = self.call

            if call.status != CallStatus.S.name:

                # change call status to started
                call.status = CallStatus.S.name
                call.save()

        # changed to complete?
        elif self.status == AmbulanceCallStatus.C.name:

            # retrieve call
            call = self.call

            # retrieve all ongoing ambulances
            ongoing_ambulancecalls = call.ambulancecall_set.exclude(status=AmbulanceCallStatus.C.name)

            set_size = len(ongoing_ambulancecalls)
            if (set_size == 0 or
                (set_size == 1 and ongoing_ambulancecalls[0].ambulance is not self)):

                # change call status to finished
                call.status = CallStatus.E.name
                call.save()

                # prevent publication
                publish = False

        # call super
        super().save(*args, **kwargs, publish=publish)

    def publish(self):

        # publish to mqtt
        from mqtt.publish import SingletonPublishClient
        SingletonPublishClient().publish_call_status(self)

    class Meta:
        unique_together = ('call', 'ambulance')

class AmbulanceUpdate(models.Model):

    # ambulance
    ambulance = models.ForeignKey(Ambulance,
                                  on_delete=models.CASCADE)

    # ambulance call
    # TODO: Is it possible to enforce that ambulance_call.ambulance == ambulance?
    ambulance_call = models.ForeignKey(AmbulanceCall,
                                       on_delete=models.SET_NULL,
                                       null=True)

    # ambulance status
    AMBULANCE_STATUS_CHOICES = \
        [(m.name, m.value) for m in AmbulanceStatus]
    status = models.CharField(max_length=2,
                              choices=AMBULANCE_STATUS_CHOICES,
                              default=AmbulanceStatus.UK.name)

    # location
    orientation = models.FloatField(default=0.0)
    location = models.PointField(srid=4326, default=defaults['location'])

    # timestamp, indexed
    timestamp = models.DateTimeField(db_index=True, default=timezone.now)

    # comment
    comment = models.CharField(max_length=254, null=True, blank=True)

    # updated by
    updated_by = models.ForeignKey(User,
                                   on_delete=models.CASCADE)
    updated_on = models.DateTimeField(default=timezone.now)

    class Meta:
        indexes = [
            models.Index(
                fields=['ambulance', 'timestamp'],
                name='ambulance_timestamp_idx',
            ),
        ]

# Patient might be expanded in the future

class Patient(CallPublishMixin,
              models.Model):
    """
    A model that provides patient fields.
    """

    call = models.ForeignKey(Call,
                             on_delete=models.CASCADE)

    name = models.CharField(max_length=254, default="")
    age = models.IntegerField(null=True)

    def publish(self):

        # publish to mqtt
        from mqtt.publish import SingletonPublishClient
        SingletonPublishClient().publish_call(self.call)

# Location related models

# noinspection PyPep8
class LocationType(Enum):
    b = 'Base'
    a = 'AED'
    o = 'Other'

class Location(AddressModel, UpdatedByModel):
    # location name
    name = models.CharField(max_length=254, unique=True)

    # location type
    LOCATION_TYPE_CHOICES = \
        [(m.name, m.value) for m in LocationType]
    type = models.CharField(max_length=1,
                            choices=LOCATION_TYPE_CHOICES,
                            default=LocationType.o.name)

    # location
    location = models.PointField(srid=4326, null=True)

    def get_absolute_url(self):
        return reverse('ambulance:location_detail', kwargs={'pk': self.id})

    def __str__(self):
        return "{} @{} ({})".format(self.name, self.location, self.comment)

# THOSE NEED REVIEWING

class Region(models.Model):
    name = models.CharField(max_length=254, unique=True)
    center = models.PointField(srid=4326, null=True)

    def __str__(self):
        return self.name
<commit_msg>Making sure ambulance updates are from authorized clients.<commit_after>import logging
from enum import Enum

from django.contrib.auth.models import User
from django.contrib.gis.db import models
from django.utils import timezone
from django.urls import reverse
from django.template.defaulttags import register

from emstrack.latlon import calculate_orientation, calculate_distance, stationary_radius
from emstrack.models import AddressModel, UpdatedByModel, defaults

logger = logging.getLogger(__name__)

# filters

@register.filter
def get_ambulance_status(key):
    return AmbulanceStatus[key].value

@register.filter
def get_ambulance_capability(key):
    return AmbulanceCapability[key].value

@register.filter
def get_location_type(key):
    return LocationType[key].value

@register.filter
def get_call_status(key):
    return CallStatus[key].value

@register.filter
def get_call_priority(key):
    return CallPriority[key].value

# Ambulance location models

# Ambulance model

class AmbulanceStatus(Enum):
    UK = 'Unknown'
    AV = 'Available'
    OS = 'Out of service'
    PB = 'Patient bound'
    AP = 'At patient'
    HB = 'Hospital bound'
    AH = 'At hospital'

class AmbulanceCapability(Enum):
    B = 'Basic'
    A = 'Advanced'
    R = 'Rescue'

class Ambulance(UpdatedByModel):

    # TODO: Should we consider denormalizing Ambulance to avoid duplication with AmbulanceUpdate?

    # ambulance properties
    identifier = models.CharField(max_length=50, unique=True)

    # TODO: Should we add an active flag?

    AMBULANCE_CAPABILITY_CHOICES = \
        [(m.name, m.value) for m in AmbulanceCapability]
    capability = models.CharField(max_length=1,
                                  choices=AMBULANCE_CAPABILITY_CHOICES)

    # status
    AMBULANCE_STATUS_CHOICES = \
        [(m.name, m.value) for m in AmbulanceStatus]
    status = models.CharField(max_length=2,
                              choices=AMBULANCE_STATUS_CHOICES,
                              default=AmbulanceStatus.UK.name)

    # location
    orientation = models.FloatField(default=0.0)
    location = models.PointField(srid=4326, default=defaults['location'])

    # timestamp
    timestamp = models.DateTimeField(default=timezone.now)

    # location client
    location_client = models.ForeignKey('login.Client',
                                        on_delete=models.CASCADE,
                                        blank=True, null=True,
                                        related_name='location_client_set')

    # default value for _loaded_values
    _loaded_values = None

    @classmethod
    def from_db(cls, db, field_names, values):

        # call super
        instance = super(Ambulance, cls).from_db(db, field_names, values)

        # store the original field values on the instance
        instance._loaded_values = dict(zip(field_names, values))

        # return instance
        return instance

    def save(self, *args, **kwargs):

        # creation?
        created = self.pk is None

        # loaded_values?
        loaded_values = self._loaded_values is not None

        # has location changed?
        has_moved = False
        if (not loaded_values) or \
                calculate_distance(self._loaded_values['location'], self.location) > stationary_radius:
            has_moved = True

        # calculate orientation only if location has changed and orientation has not changed
        if has_moved and loaded_values and self._loaded_values['orientation'] == self.orientation:
            # TODO: should we allow for a small radius before updating direction?
            self.orientation = calculate_orientation(self._loaded_values['location'], self.location)
            logger.debug('< {} - {} = {}'.format(self._loaded_values['location'],
                                                 self.location,
                                                 self.orientation))

        logger.debug('loaded_values: {}'.format(loaded_values))
        logger.debug('_loaded_values: {}'.format(self._loaded_values))
        logger.debug('self.location_client: {}'.format(self.location_client))

        # location_client changed?
        if self.location_client is None:
            location_client_id = None
        else:
            location_client_id = self.location_client.id
        location_client_changed = False
        if loaded_values and location_client_id != self._loaded_values['location_client_id']:
            location_client_changed = True

        logger.debug('location_client_changed: {}'.format(location_client_changed))
        # TODO: Check if client is logged with ambulance if setting location_client

        # if comment, status or location changed
        model_changed = False
        if has_moved or \
                self._loaded_values['status'] != self.status or \
                self._loaded_values['comment'] != self.comment:

            # save to Ambulance
            super().save(*args, **kwargs)

            logger.debug('SAVED')

            # save to AmbulanceUpdate
            data = {k: getattr(self, k)
                    for k in ('status', 'orientation',
                              'location', 'timestamp',
                              'comment', 'updated_by', 'updated_on')}
            data['ambulance'] = self
            obj = AmbulanceUpdate(**data)
            obj.save()

            logger.debug('UPDATE SAVED')

            # model changed
            model_changed = True

        # if identifier or capability changed
        # NOTE: self._loaded_values is NEVER None because has_moved is True
        elif (location_client_changed or
              self._loaded_values['identifier'] != self.identifier or
              self._loaded_values['capability'] != self.capability):

            # save only to Ambulance
            super().save(*args, **kwargs)

            logger.debug('SAVED')

            # model changed
            model_changed = True

        # Did the model change?
        if model_changed:

            # publish to mqtt
            from mqtt.publish import SingletonPublishClient
            SingletonPublishClient().publish_ambulance(self)

            logger.debug('PUBLISHED ON MQTT')

        # just created?
        if created:
            # invalidate permissions cache
            from login.permissions import cache_clear
            cache_clear()

    def delete(self, *args, **kwargs):

        # remove from mqtt
        from mqtt.publish import SingletonPublishClient
        SingletonPublishClient().remove_ambulance(self)

        # invalidate permissions cache
        from login.permissions import cache_clear
        cache_clear()

        # delete from Ambulance
        super().delete(*args, **kwargs)

    def get_absolute_url(self):
        return reverse('ambulance:detail', kwargs={'pk': self.id})

    def __str__(self):
        return ('Ambulance {}(id={}) ({}) [{}]:\n' +
                '    Status: {}\n' +
                '  Location: {} @ {}\n' +
                ' LocClient: {}\n' +
                '   Updated: {} by {}').format(self.identifier,
                                               self.id,
                                               AmbulanceCapability[self.capability].value,
                                               self.comment,
                                               AmbulanceStatus[self.status].value,
                                               self.location,
                                               self.timestamp,
                                               self.location_client,
                                               self.updated_by,
                                               self.updated_on)

# Call related models

class CallPriority(Enum):
    A = 'Resucitation'
    B = 'Emergent'
    C = 'Urgent'
    D = 'Less urgent'
    E = 'Not urgent'
    O = 'Omega'

class CallStatus(Enum):
    P = 'Pending'
    S = 'Started'
    E = 'Ended'

class CallPublishMixin:

    def save(self, *args, **kwargs):

        # publish?
        publish = kwargs.pop('publish', True)

        # save to Call
        super().save(*args, **kwargs)

        if publish:
            self.publish()

class Call(CallPublishMixin,
           AddressModel,
           UpdatedByModel):

    # status
    CALL_STATUS_CHOICES = \
        [(m.name, m.value) for m in CallStatus]
    status = models.CharField(max_length=1,
                              choices=CALL_STATUS_CHOICES,
                              default=CallStatus.P.name)

    # details
    details = models.CharField(max_length=500, default="")

    # call priority
    CALL_PRIORITY_CHOICES = \
        [(m.name, m.value) for m in CallPriority]
    priority = models.CharField(max_length=1,
                                choices=CALL_PRIORITY_CHOICES,
                                default=CallPriority.E.name)

    # timestamps
    pending_at = models.DateTimeField(null=True, blank=True)
    started_at = models.DateTimeField(null=True, blank=True)
    ended_at = models.DateTimeField(null=True, blank=True)

    # created at
    created_at = models.DateTimeField(auto_now_add=True)

    def save(self, *args, **kwargs):

        # publish?
        publish = kwargs.pop('publish', True)

        if self.status == CallStatus.E.name:
            # remove topic from mqtt server
            from mqtt.publish import SingletonPublishClient
            SingletonPublishClient().remove_call(self)

            # prevent publication
            publish = False

        # call super
        super().save(*args, **kwargs, publish=publish)

    def publish(self):

        # publish to mqtt
        from mqtt.publish import SingletonPublishClient
        SingletonPublishClient().publish_call(self)

    def __str__(self):
        return "{} ({})".format(self.location, self.priority)

class AmbulanceCallStatus(Enum):
    R = 'Requested'
    O = 'Ongoing'
    I = 'Interrupted'
    C = 'Completed'

class AmbulanceCall(CallPublishMixin,
                    models.Model):

    # status
    AMBULANCE_CALL_STATUS_CHOICES = \
        [(m.name, m.value) for m in AmbulanceCallStatus]
    status = models.CharField(max_length=1,
                              choices=AMBULANCE_CALL_STATUS_CHOICES,
                              default=AmbulanceCallStatus.R.name)

    # call
    call = models.ForeignKey(Call,
                             on_delete=models.CASCADE)

    # ambulance
    ambulance = models.ForeignKey(Ambulance,
                                  on_delete=models.CASCADE)

    # created at
    created_at = models.DateTimeField(auto_now_add=True)

    def save(self, *args, **kwargs):

        # publish?
        publish = kwargs.pop('publish', True)

        # changed to ongoing?
        if self.status == AmbulanceCallStatus.O.name:

            # retrieve call
            call = self.call

            if call.status != CallStatus.S.name:

                # change call status to started
                call.status = CallStatus.S.name
                call.save()

        # changed to complete?
        elif self.status == AmbulanceCallStatus.C.name:

            # retrieve call
            call = self.call

            # retrieve all ongoing ambulances
            ongoing_ambulancecalls = call.ambulancecall_set.exclude(status=AmbulanceCallStatus.C.name)

            set_size = len(ongoing_ambulancecalls)
            if (set_size == 0 or
                    (set_size == 1 and ongoing_ambulancecalls[0].ambulance is not self)):

                logger.debug('This is the last ambulance; will end call.')

                # change call status to finished
                call.status = CallStatus.E.name
                call.save()

                # prevent publication
                publish = False

            else:

                logger.debug('There are still {} ambulances in this call.'.format(set_size))

        # call super
        super().save(*args, **kwargs, publish=publish)

    def publish(self):

        # publish to mqtt
        from mqtt.publish import SingletonPublishClient
        SingletonPublishClient().publish_call_status(self)

    class Meta:
        unique_together = ('call', 'ambulance')

class AmbulanceUpdate(models.Model):

    # ambulance
    ambulance = models.ForeignKey(Ambulance,
                                  on_delete=models.CASCADE)

    # ambulance call
    # TODO: Is it possible to enforce that ambulance_call.ambulance == ambulance?
    ambulance_call = models.ForeignKey(AmbulanceCall,
                                       on_delete=models.SET_NULL,
                                       null=True)

    # ambulance status
    AMBULANCE_STATUS_CHOICES = \
        [(m.name, m.value) for m in AmbulanceStatus]
    status = models.CharField(max_length=2,
                              choices=AMBULANCE_STATUS_CHOICES,
                              default=AmbulanceStatus.UK.name)

    # location
    orientation = models.FloatField(default=0.0)
    location = models.PointField(srid=4326, default=defaults['location'])

    # timestamp, indexed
    timestamp = models.DateTimeField(db_index=True, default=timezone.now)

    # comment
    comment = models.CharField(max_length=254, null=True, blank=True)

    # updated by
    updated_by = models.ForeignKey(User,
                                   on_delete=models.CASCADE)
    updated_on = models.DateTimeField(default=timezone.now)

    class Meta:
        indexes = [
            models.Index(
                fields=['ambulance', 'timestamp'],
                name='ambulance_timestamp_idx',
            ),
        ]

# Patient might be expanded in the future

class Patient(CallPublishMixin,
              models.Model):
    """
    A model that provides patient fields.
    """

    call = models.ForeignKey(Call,
                             on_delete=models.CASCADE)

    name = models.CharField(max_length=254, default="")
    age = models.IntegerField(null=True)

    def publish(self):

        # publish to mqtt
        from mqtt.publish import SingletonPublishClient
        SingletonPublishClient().publish_call(self.call)

# Location related models

# noinspection PyPep8
class LocationType(Enum):
    b = 'Base'
    a = 'AED'
    o = 'Other'

class Location(AddressModel, UpdatedByModel):
    # location name
    name = models.CharField(max_length=254, unique=True)

    # location type
    LOCATION_TYPE_CHOICES = \
        [(m.name, m.value) for m in LocationType]
    type = models.CharField(max_length=1,
                            choices=LOCATION_TYPE_CHOICES,
                            default=LocationType.o.name)

    # location
    location = models.PointField(srid=4326, null=True)

    def get_absolute_url(self):
        return reverse('ambulance:location_detail', kwargs={'pk': self.id})

    def __str__(self):
        return "{} @{} ({})".format(self.name, self.location, self.comment)

# THOSE NEED REVIEWING

class Region(models.Model):
    name = models.CharField(max_length=254, unique=True)
    center = models.PointField(srid=4326, null=True)

    def __str__(self):
        return self.name
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
@package ion.agents.platform.oms.simulator.test.test_oms_simulator
@file    ion/agents/platform/oms/simulator/test/test_oms_simulator.py
@author  Carlos Rueda
@brief   Test cases for the simulator
"""

__author__ = 'Carlos Rueda'
__license__ = 'Apache 2.0'

from pyon.public import log
from ion.agents.platform.oms.simulator.logger import Logger
Logger.set_logger(log)

from pyon.util.unit_test import IonUnitTestCase

from ion.agents.platform.oms.simulator.oms_simulator import OmsSimulator
from ion.agents.platform.oms.test.oms_test_mixin import OmsTestMixin

from nose.plugins.attrib import attr

@attr('UNIT', group='sa')
class Test(IonUnitTestCase, OmsTestMixin):
    """
    Test cases for the simulator, which is instantiated directly (ie.,
    no connection to external simulator is involved).
    """

    @classmethod
    def setUpClass(cls):
        OmsTestMixin.setUpClass()
        OmsTestMixin.start_http_server()
        cls.oms = OmsSimulator()

    @classmethod
    def tearDownClass(cls):
        cls.oms._deactivate_simulator()
        event_notifications = OmsTestMixin.stop_http_server()
        log.info("event_notifications = %s" % str(event_notifications))
<commit_msg>allow to also test against our simulator running externally<commit_after>#!/usr/bin/env python

"""
@package ion.agents.platform.oms.simulator.test.test_oms_simulator
@file    ion/agents/platform/oms/simulator/test/test_oms_simulator.py
@author  Carlos Rueda
@brief   Test cases for the simulator, by default in embedded form,
         but the OMS environment variable can be used to indicate other.
"""

__author__ = 'Carlos Rueda'
__license__ = 'Apache 2.0'

from pyon.public import log
from ion.agents.platform.oms.simulator.logger import Logger
Logger.set_logger(log)

from pyon.util.unit_test import IonUnitTestCase

from ion.agents.platform.oms.oms_client_factory import OmsClientFactory
from ion.agents.platform.oms.simulator.oms_simulator import OmsSimulator
from ion.agents.platform.oms.test.oms_test_mixin import OmsTestMixin

from nose.plugins.attrib import attr

@attr('UNIT', group='sa')
class Test(IonUnitTestCase, OmsTestMixin):
    """
    Test cases for the simulator, which is instantiated directly (ie.,
    no connection to external simulator is involved).
    """

    @classmethod
    def setUpClass(cls):
        OmsTestMixin.setUpClass()
        OmsTestMixin.start_http_server()
        cls.oms = OmsClientFactory.create_instance()

    @classmethod
    def tearDownClass(cls):
        if isinstance(cls.oms, OmsSimulator):
            cls.oms._deactivate_simulator()
        event_notifications = OmsTestMixin.stop_http_server()
        log.info("event_notifications = %s" % str(event_notifications))
<|endoftext|>


<commit_before>from .models import Note
from rest_framework import serializers, viewsets, routers

class NoteSerializer(serializers.HyperlinkedModelSerializer):
    class Meta:
        model = Note
        fields = ('id', 'content', 'time', 'update_time')

class NoteViewSet(viewsets.ModelViewSet):
    queryset = Note.objects.all()
    serializer_class = NoteSerializer

router = routers.DefaultRouter()
router.register(r'notes', NoteViewSet)<commit_msg> update NoteViewSet to readonly<commit_after># -*- coding: utf-8 -*-
from .models import Note
from rest_framework import serializers, viewsets, routers

class NoteSerializer(serializers.HyperlinkedModelSerializer):
    class Meta:
        model = Note
        fields = ('id', 'content', 'time', 'update_time')

class NoteViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = Note.objects.all()
    serializer_class = NoteSerializer

router = routers.DefaultRouter()
router.register(r'notes', NoteViewSet)<|endoftext|>


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

# import glob
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import numpy as np
# import os
# import pandas as pd
# import pylab as pl
# import re
# import sys
# import trappy

# Configure logging
import logging

class Filters(object):

    def __init__(self, trace, tasks=None):
        self.trace = trace
        self.tasks = tasks

        self.big_tasks = {}

        self.big_frequent_tasks_pids = None
        self.wkp_frequent_tasks_pids = None

        self.big_cap = self.trace.platform['nrg_model']['big']['cpu']['cap_max']
        self.little_cap = self.trace.platform['nrg_model']['little']['cpu']['cap_max']
        self.big_cpus = self.trace.platform['clusters']['big']
        self.little_cpus = self.trace.platform['clusters']['little']

        # Minimum and Maximum x_time to use for all plots
        self.x_min = 0
        self.x_max = self.trace.time_range

        # Reset x axis time range to full scale
        self.setXTimeRange()

    def setXTimeRange(self, t_min=None, t_max=None):
        if t_min is None:
            self.x_min = 0
        else:
            self.x_min = t_min
        if t_max is None:
            self.x_max = self.trace.time_range
        else:
            self.x_max = t_max
        logging.info('Set plots time range to (%.6f, %.6f)[s]',
                self.x_min, self.x_max)

    def topBigTasks(self, max_tasks=10, min_samples=100, min_utilization=None):
        """
        Tasks which had a 'utilization' bigger than the specified threshold
        """

        if min_utilization is None:
            min_utilization = self.little_cap

        df = self.trace.df('sched_load_avg_task')
        big_tasks_events = df[df.util_avg > min_utilization]
        big_tasks = big_tasks_events.pid.unique()

        big_tasks_count = big_tasks.size
        print 'Total {} tasks with at least {} "utilization" samples > {}'\
                .format(big_tasks_count, min_samples, min_utilization)

        big_tasks_stats = big_tasks_events.groupby('pid')\
                            .describe(include=['object']);
        big_tasks_pids = big_tasks_stats.unstack()['comm']\
                            .sort_values(by=['count'], ascending=False)
        big_tasks_pids = big_tasks_pids[big_tasks_pids['count'] > min_samples]

        big_topmost = big_tasks_pids.head(max_tasks)
        print 'Top {} "big" tasks:'.format(max_tasks)
        print big_topmost

        self.big_frequent_tasks_pids = list(big_topmost.index)

        # Keep track of big tasks tload events
        self.big_tasks['sched_load_avg_task'] = big_tasks_events
        return self.big_frequent_tasks_pids

    def _taskIsBig(self, utilization):
        if utilization > self.little_cap:
            return self.big_cap
        return self.little_cap

    def plotBigTasks(self, max_tasks=10, min_samples=100, min_utilization=None):

        # Get the list of big and frequent tasks
        if self.big_frequent_tasks_pids is None:
            pids = self.topBigTasks(max_tasks, min_samples, min_utilization)

        big_frequent_tasks_count = len(self.big_frequent_tasks_pids)
        if big_frequent_tasks_count == 0:
            print "No big/frequent tasks to plot"
            return

        # Get the list of events for all big frequent tasks
        df = self.trace.df('sched_load_avg_task')
        big_frequent_tasks_events = df[df.pid.isin(self.big_frequent_tasks_pids)]

        # Add a column to represent big status
        big_frequent_tasks_events.loc[:,'isbig'] = \
                big_frequent_tasks_events['util_avg'].map(self._taskIsBig)

        # Define axes for side-by-side plottings
        fig, axes = plt.subplots(big_frequent_tasks_count, 1,
                                 figsize=(14, big_frequent_tasks_count*5));
        plt.subplots_adjust(wspace=0.1, hspace=0.2);

        plot_idx = 0
        for i, group in big_frequent_tasks_events.groupby('pid'):

            # Build task names (there could be multiple, during the task lifetime)
            big_frequent_task_i = big_frequent_tasks_events[big_frequent_tasks_events['pid'] == i]
            task_names = big_frequent_task_i.comm.unique()
            task_name = 'PID: ' + str(i)
            for s in task_names:
                task_name += ' | ' + s

            # Plot title
            if (big_frequent_tasks_count == 1):
                ax_ratio = axes
            else:
                 ax_ratio = axes[plot_idx]
            ax_ratio.set_title(task_name);

            # Left axis: utilization
            ax_ratio = group.plot(y=['util_avg', 'isbig'],
                            style=['r.', '-b'],
                            drawstyle='steps-post',
                            linewidth=1,
                            ax=ax_ratio)
            ax_ratio.set_xlim(self.x_min, self.x_max);
            ax_ratio.set_ylim(0, 1100)
            ax_ratio.set_ylabel('util_avg')

            plot_idx+=1

        print 'Tasks which have been a "utilization" of {0:d} for at least {1:d} samples'\
            .format(self.little_cap, min_samples)

    def topWakeupTasks(self, max_tasks=10, min_wakeups=100):
        """
        Tasks which wakeups more frequent than a specified threshold
        """

        df = self.trace.df('sched_wakeup')

        wkp_tasks_stats = df.groupby('pid').describe(include=['object'])
        wkp_tasks_pids = wkp_tasks_stats.unstack()['comm']\
                        .sort_values(by=['count'], ascending=False)
        wkp_tasks_pids = wkp_tasks_pids[wkp_tasks_pids['count'] > min_wakeups]

        wkp_topmost = wkp_tasks_pids.head(max_tasks)
        print 'Top {} "big" tasks:'.format(max_tasks)
        print wkp_topmost

        self.wkp_frequent_tasks_pids = list(wkp_topmost.index)
        return self.wkp_frequent_tasks_pids

    def plotWakeupTasks(self, max_tasks=10, min_wakeups=0, per_cluster=False):

        # Get the list of big and frequent tasks
        if self.wkp_frequent_tasks_pids is None:
            pids = self.topWakeupTasks(max_tasks, min_wakeups)

        wkp_frequent_tasks_count = len(self.wkp_frequent_tasks_pids)
        if wkp_frequent_tasks_count == 0:
            print "No big/frequent wakeups tasks to plot"
            return

        # Define axes for side-by-side plottings
        fig, axes = plt.subplots(2, 1, figsize=(14, 5));
        plt.subplots_adjust(wspace=0.2, hspace=0.3);

        if per_cluster:

            # Get per cluster wakeup events
            df = self.trace.df('sched_wakeup_new')
            big_frequent = (
                    (df.target_cpu.isin(self.big_cpus))
                    )
            ntbc = df[big_frequent]
            little_frequent = (
                    (df.target_cpu.isin(self.little_cpus))
                    )
            ntlc = df[little_frequent];

            ax = axes[0]
            ax.set_title('Tasks Forks on big CPUs');
            ax.set_xlim(self.x_min, self.x_max);
            ntbc.pid.plot(style=['g.'], ax=ax);

            ax = axes[1]
            ax.set_title('Tasks Forks on LITTLE CPUs');
            ax.set_xlim(self.x_min, self.x_max);
            ntlc.pid.plot(style=['g.'], ax=ax);

        else:

            ax = axes[0]
            ax.set_title('Tasks WakeUps Events');
            df = self.trace.df('sched_wakeup')
            df.pid.astype(int).plot(style=['b.'], ax=ax);
            ax.set_xlim(self.x_min, self.x_max);
            ax.xaxis.set_visible(False);

            ax = axes[1]
            ax.set_title('Tasks Forks Events');
            df = self.trace.df('sched_wakeup_new')
            df.pid.astype(int).plot(style=['r.'], ax=ax);
            ax.set_xlim(self.x_min, self.x_max);
            ax.xaxis.set_visible(False);

    def plotTasksVsFrequency(self, big_cluster=True):

        if big_cluster:
            cluster_correct = 'big'
            cluster_wrong = 'LITTLE'
            cpus = self.big_cpus
        else:
            cluster_correct = 'LITTLE'
            cluster_wrong = 'big'
            cpus = self.little_cpus

        fig, axes = plt.subplots(2, 1, figsize=(14, 5));
        plt.subplots_adjust(wspace=0.2, hspace=0.3);

        df_wkp  = self.big_tasks['sched_load_avg_task']
        # Add column of expected cluster, depending on utilization value and
        # capacity of the selected cluster
        bu_bc = ( \
                (df_wkp['util_avg'] > 500) & \
                (df_wkp['cpu'].isin([2,3]))
            )
        su_lc = ( \
                (df_wkp['util_avg'] <= 500) & \
                (df_wkp['cpu'].isin([0,1]))
            )
        df_wkp.loc[:,'ccap_mutil'] = np.select(
            [(bu_bc | su_lc)], [True], False)

        df_freq = self.trace.df('cpu_frequency')
        rd_freq = df_freq[df_freq['cpu'].isin(cpus)]

        ax = axes[0]
        ax.set_title('Big Tasks Utilization vs Allocation');
        for ucolor, umatch in zip('gr', [True, False]):
            cdata  = df_wkp[df_wkp['ccap_mutil'] == umatch]
            if (len(cdata) > 0):
                cdata['util_avg'].plot(ax=ax,
                        style=[ucolor+'.'], legend=False);
        ax.set_xlim(self.x_min, self.x_max);
        ax.xaxis.set_visible(False);

        ax = axes[1]
        ax.set_title('Frequencies on "{}" cluster'.format(cluster_correct))
        df_freq['frequency'].plot(style=['-b'], ax=ax, drawstyle='steps-post');
        ax.set_xlim(self.x_min, self.x_max);

    def rtTasks(self, max_prio = 100):
        df = self.trace.df('sched_switch')
        df = df[df.next_prio <= max_prio]
        df = df[['next_pid', 'next_comm']]
        df = df.drop_duplicates()
        rt_tasks = {}
        for pid,task in df.values:
                rt_tasks[pid] = task
        return rt_tasks
<commit_msg>libs/utils/filters: return tasks dictionary for all filters<commit_after># SPDX-License-Identifier: Apache-2.0
#
# Copyright (C) 2015, ARM Limited and contributors.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# import glob
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import numpy as np
# import os
# import pandas as pd
# import pylab as pl
# import re
# import sys
# import trappy

# Configure logging
import logging

class Filters(object):

    def __init__(self, trace, tasks=None):
        self.trace = trace
        self.tasks = tasks

        self.big_tasks = {}

        self.big_frequent_tasks_pids = None
        self.big_frequent_tasks_tasks = None
        self.wkp_frequent_tasks_pids = None
        self.wkp_frequent_tasks_tasks = None

        self.big_cap = self.trace.platform['nrg_model']['big']['cpu']['cap_max']
        self.little_cap = self.trace.platform['nrg_model']['little']['cpu']['cap_max']
        self.big_cpus = self.trace.platform['clusters']['big']
        self.little_cpus = self.trace.platform['clusters']['little']

        # Minimum and Maximum x_time to use for all plots
        self.x_min = 0
        self.x_max = self.trace.time_range

        # Reset x axis time range to full scale
        self.setXTimeRange()

    def setXTimeRange(self, t_min=None, t_max=None):
        if t_min is None:
            self.x_min = 0
        else:
            self.x_min = t_min
        if t_max is None:
            self.x_max = self.trace.time_range
        else:
            self.x_max = t_max
        logging.info('Set plots time range to (%.6f, %.6f)[s]',
                self.x_min, self.x_max)

    def topBigTasks(self, max_tasks=10, min_samples=100, min_utilization=None):
        """
        Tasks which had a 'utilization' bigger than the specified threshold
        """

        if min_utilization is None:
            min_utilization = self.little_cap

        df = self.trace.df('sched_load_avg_task')
        big_tasks_events = df[df.util_avg > min_utilization]
        big_tasks = big_tasks_events.pid.unique()

        big_tasks_count = big_tasks.size
        print 'Total {} tasks with at least {} "utilization" samples > {}'\
                .format(big_tasks_count, min_samples, min_utilization)

        big_tasks_stats = big_tasks_events.groupby('pid')\
                            .describe(include=['object']);
        big_tasks_pids = big_tasks_stats.unstack()['comm']\
                            .sort_values(by=['count'], ascending=False)
        big_tasks_pids = big_tasks_pids[big_tasks_pids['count'] > min_samples]

        big_topmost = big_tasks_pids.head(max_tasks)
        print 'Top {} "big" tasks:'.format(max_tasks)
        print big_topmost

        self.big_frequent_tasks_tasks = dict(zip(big_topmost.top.values, big_topmost.index))
        self.big_frequent_tasks_pids = list(big_topmost.index)

        # Keep track of big tasks tload events
        self.big_tasks['sched_load_avg_task'] = big_tasks_events

        return self.big_frequent_tasks_tasks

    def _taskIsBig(self, utilization):
        if utilization > self.little_cap:
            return self.big_cap
        return self.little_cap

    def plotBigTasks(self, max_tasks=10, min_samples=100, min_utilization=None):

        # Get the list of big and frequent tasks
        if self.big_frequent_tasks_pids is None:
            self.topBigTasks(max_tasks, min_samples, min_utilization)

        big_frequent_tasks_count = len(self.big_frequent_tasks_pids)
        if big_frequent_tasks_count == 0:
            print "No big/frequent tasks to plot"
            return

        # Get the list of events for all big frequent tasks
        df = self.trace.df('sched_load_avg_task')
        big_frequent_tasks_events = df[df.pid.isin(self.big_frequent_tasks_pids)]

        # Add a column to represent big status
        big_frequent_tasks_events.loc[:,'isbig'] = \
                big_frequent_tasks_events['util_avg'].map(self._taskIsBig)

        # Define axes for side-by-side plottings
        fig, axes = plt.subplots(big_frequent_tasks_count, 1,
                                 figsize=(14, big_frequent_tasks_count*5));
        plt.subplots_adjust(wspace=0.1, hspace=0.2);

        plot_idx = 0
        for i, group in big_frequent_tasks_events.groupby('pid'):

            # Build task names (there could be multiple, during the task lifetime)
            big_frequent_task_i = big_frequent_tasks_events[big_frequent_tasks_events['pid'] == i]
            task_names = big_frequent_task_i.comm.unique()
            task_name = 'PID: ' + str(i)
            for s in task_names:
                task_name += ' | ' + s

            # Plot title
            if (big_frequent_tasks_count == 1):
                ax_ratio = axes
            else:
                 ax_ratio = axes[plot_idx]
            ax_ratio.set_title(task_name);

            # Left axis: utilization
            ax_ratio = group.plot(y=['util_avg', 'isbig'],
                            style=['r.', '-b'],
                            drawstyle='steps-post',
                            linewidth=1,
                            ax=ax_ratio)
            ax_ratio.set_xlim(self.x_min, self.x_max);
            ax_ratio.set_ylim(0, 1100)
            ax_ratio.set_ylabel('util_avg')

            plot_idx+=1

        print 'Tasks which have been a "utilization" of {0:d} for at least {1:d} samples'\
            .format(self.little_cap, min_samples)

    def topWakeupTasks(self, max_tasks=10, min_wakeups=100):
        """
        Tasks which wakeups more frequent than a specified threshold
        """

        df = self.trace.df('sched_wakeup')

        wkp_tasks_stats = df.groupby('pid').describe(include=['object'])
        wkp_tasks_pids = wkp_tasks_stats.unstack()['comm']\
                        .sort_values(by=['count'], ascending=False)
        wkp_tasks_pids = wkp_tasks_pids[wkp_tasks_pids['count'] > min_wakeups]

        wkp_topmost = wkp_tasks_pids.head(max_tasks)
        print 'Top {} "big" tasks:'.format(max_tasks)
        print wkp_topmost

        self.wkp_frequent_tasks_tasks = dict(zip(wkp_topmost.top.values, wkp_topmost.index))
        self.wkp_frequent_tasks_pids = list(wkp_topmost.index)

        return self.wkp_frequent_tasks_tasks

    def plotWakeupTasks(self, max_tasks=10, min_wakeups=0, per_cluster=False):

        # Get the list of big and frequent tasks
        if self.wkp_frequent_tasks_pids is None:
            self.topWakeupTasks(max_tasks, min_wakeups)

        wkp_frequent_tasks_count = len(self.wkp_frequent_tasks_pids)
        if wkp_frequent_tasks_count == 0:
            print "No big/frequent wakeups tasks to plot"
            return

        # Define axes for side-by-side plottings
        fig, axes = plt.subplots(2, 1, figsize=(14, 5));
        plt.subplots_adjust(wspace=0.2, hspace=0.3);

        if per_cluster:

            # Get per cluster wakeup events
            df = self.trace.df('sched_wakeup_new')
            big_frequent = (
                    (df.target_cpu.isin(self.big_cpus))
                    )
            ntbc = df[big_frequent]
            little_frequent = (
                    (df.target_cpu.isin(self.little_cpus))
                    )
            ntlc = df[little_frequent];

            ax = axes[0]
            ax.set_title('Tasks Forks on big CPUs');
            ax.set_xlim(self.x_min, self.x_max);
            ntbc.pid.plot(style=['g.'], ax=ax);

            ax = axes[1]
            ax.set_title('Tasks Forks on LITTLE CPUs');
            ax.set_xlim(self.x_min, self.x_max);
            ntlc.pid.plot(style=['g.'], ax=ax);

        else:

            ax = axes[0]
            ax.set_title('Tasks WakeUps Events');
            df = self.trace.df('sched_wakeup')
            df.pid.astype(int).plot(style=['b.'], ax=ax);
            ax.set_xlim(self.x_min, self.x_max);
            ax.xaxis.set_visible(False);

            ax = axes[1]
            ax.set_title('Tasks Forks Events');
            df = self.trace.df('sched_wakeup_new')
            df.pid.astype(int).plot(style=['r.'], ax=ax);
            ax.set_xlim(self.x_min, self.x_max);
            ax.xaxis.set_visible(False);

    def plotTasksVsFrequency(self, big_cluster=True):

        if big_cluster:
            cluster_correct = 'big'
            cluster_wrong = 'LITTLE'
            cpus = self.big_cpus
        else:
            cluster_correct = 'LITTLE'
            cluster_wrong = 'big'
            cpus = self.little_cpus

        fig, axes = plt.subplots(2, 1, figsize=(14, 5));
        plt.subplots_adjust(wspace=0.2, hspace=0.3);

        df_wkp  = self.big_tasks['sched_load_avg_task']
        # Add column of expected cluster, depending on utilization value and
        # capacity of the selected cluster
        bu_bc = ( \
                (df_wkp['util_avg'] > 500) & \
                (df_wkp['cpu'].isin([2,3]))
            )
        su_lc = ( \
                (df_wkp['util_avg'] <= 500) & \
                (df_wkp['cpu'].isin([0,1]))
            )
        df_wkp.loc[:,'ccap_mutil'] = np.select(
            [(bu_bc | su_lc)], [True], False)

        df_freq = self.trace.df('cpu_frequency')
        rd_freq = df_freq[df_freq['cpu'].isin(cpus)]

        ax = axes[0]
        ax.set_title('Big Tasks Utilization vs Allocation');
        for ucolor, umatch in zip('gr', [True, False]):
            cdata  = df_wkp[df_wkp['ccap_mutil'] == umatch]
            if (len(cdata) > 0):
                cdata['util_avg'].plot(ax=ax,
                        style=[ucolor+'.'], legend=False);
        ax.set_xlim(self.x_min, self.x_max);
        ax.xaxis.set_visible(False);

        ax = axes[1]
        ax.set_title('Frequencies on "{}" cluster'.format(cluster_correct))
        df_freq['frequency'].plot(style=['-b'], ax=ax, drawstyle='steps-post');
        ax.set_xlim(self.x_min, self.x_max);

    def rtTasks(self, max_prio = 100):
        df = self.trace.df('sched_switch')
        df = df[df.next_prio <= max_prio]
        df = df[['next_pid', 'next_comm']]
        df = df.drop_duplicates()
        rt_tasks = {}
        for pid,task in df.values:
                rt_tasks[pid] = task
        return rt_tasks
<|endoftext|>