<commit_before>import os
import re
import models_ida
import models_parser
import util_parser

from config import CONFIG
from sqlalchemy import select
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.sql.expression import func
from pida_types.types import IDA_TYPES

class IdaInfoParser(object):
    def __init__(self, db_file):
        self.db_file = db_file
        self.Session = sessionmaker()
        self.pattern_member = re.compile('(.*) .*\w+')

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.session.commit()

    def start(self):
        self.__create_connection()
        self.__drop_tables()
        self.__create_tables()
        self.__create_session()
        self.__parsing()
        self.__fetch_depend()
        self.__linking()

    def __parsing(self):
        self.__parsing_functions()
        self.__parsing_local_types()

    def __fetch_depend(self):
        self.__fetch_depend_functions()
        self.__fetch_depend_local_types()
        # todo : add fetch depend for typedef

    def __linking(self):
        self.__linking_functions()
        self.__linking_local_types()
        # todo : add link for typedef

    def __drop_tables(self):
        models_parser.Base.metadata.drop_all(self.engine_db)

    def __create_tables(self):
        models_ida.Base.metadata.create_all(self.engine_db)
        models_parser.Base.metadata.create_all(self.engine_db)

    def __create_session(self):
        self.Session.configure(bind=self.engine_db, autocommit=False)
        self.session = self.Session()

    def __create_connection(self):
        base_dir = os.path.dirname(self.db_file)
        if not os.path.exists(base_dir):
            os.makedirs(base_dir)

        if not os.path.exists(self.db_file):
            open(self.db_file, 'a').close()

        self.engine_db = create_engine('sqlite:///' + self.db_file, echo=CONFIG['sql_verbose'])

    def __parsing_functions(self):
        query = self.session.query(models_ida.IdaRawFunctions)
        count = query.count()
        if CONFIG['verbose']:
            print 'count functions: {count}'.format(count=count)

        index_item = 0
        functions = []
        for item in query.all():
            index_item += 1
            function = models_parser.Function(
                id_ida=item.get_id(),
                raw_name=item.get_name(),
                ida_type=item.get_type(),
                ida_fields=item.get_args_name(),
            )
            try:
                function.parsing()
                functions.append(function)
            except:
                print '[Error] function with id = {id}. catch exception'.format(id=item.get_id())

            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(functions)
        self.session.commit()

    def __parsing_local_types(self):
        query = self.session.query(models_ida.IdaRawLocalType)
        count = query.count()
        if CONFIG['verbose']:
            print 'count local_type: {count}'.format(count=count)

        index_item = 0
        local_types = []
        for item in query.all():
            index_item += 1
            local_type = models_parser.LocalType(
                id_ida=item.get_id(),
                raw_multi=item.get_type(),
            )
            local_type.parsing()
            local_types.append(local_type)

            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(local_types)
        self.session.commit()

    def __fetch_local_type(self, all_types):
        local_types = []
        for t in all_types:
            a_type = t
            while True:
                value = a_type['value']
                if a_type['idt'] == IDA_TYPES['local_type']:
                    local_types.append(a_type)
                    break

                elif a_type['idt'] == IDA_TYPES['function']:
                    local_types.extend(self.__fetch_local_type(value['args_type']))
                    value = value['ret_type']

                elif a_type['idt'] == IDA_TYPES['array']:
                    value = value['type']

                if type(value) is not type({}):
                    break

                a_type = value

        return local_types

    def __fetch_depend_functions(self):
        query = self.session.query(models_parser.Function)
        count = query.count()
        if CONFIG['verbose']:
            print 'count functions: {count}'.format(count=count)

        function_deps = []
        index_item = 0
        for item in query.all():
            index_item += 1
            all_types = item.get_args_type()
            all_types.append(item.get_return_type())

            dep_types = self.__fetch_local_type(all_types)

            for dep_type in dep_types:
                depend = models_parser.DependFunction(
                    id_function=item.get_id(),
                    id_local_type=dep_type['value']
                )
                function_deps.append(depend)

            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(function_deps)
        self.session.commit()

    def __fetch_members(self, one_line):
        local_types = set()

        bracket = [one_line.find('{'), one_line.rfind('}')]
        if -1 in bracket:
            return local_types

        raw_parents = one_line[:bracket[0] - 1]
        inheritance = raw_parents.find(':')
        if inheritance != -1:
            parents = raw_parents[inheritance + 1:]
            tmpl_pairs = list(util_parser.get_pairs_sym(parents, '<', '>'))
            if 0 == len(tmpl_pairs):
                local_types.update(parents.split(','))
            else:
                sets = set()
                for tmpl_pair in tmpl_pairs:
                    count_join = 0
                    for tmpl_pair2 in tmpl_pairs:
                        if not tmpl_pair2[0] <= tmpl_pair[0] <= tmpl_pair2[1]:
                            continue

                        if not tmpl_pair2[0] <= tmpl_pair[1] <= tmpl_pair2[1]:
                            continue

                        count_join += 1

                    sets.add((count_join, tmpl_pair))

                parents += ','
                for v in sets:
                    (count, pair) = v
                    if count != 1:
                        continue

                    if parents[pair[1] + 1] == ',':
                        parents = parents[:pair[1] + 1] + '!' + parents[pair[1] + 1:]

                for parent in parents.split('!,'):
                    local_types.add(parent.strip())

        raw_members = one_line[bracket[0] + 1:bracket[1]]
        for member in raw_members.split(';'):
            type_member = self.pattern_member.search(member)
            if type_member is None:
                continue

            type_member = type_member.group(1).strip()
            local_types.add(type_member)

        ret_local_types = set()
        for lt in local_types:
            ret_local_types.add(lt.strip())

        return ret_local_types

    def __fetch_depend_local_types(self):
        local_types_table = (
            select([models_parser.LocalType.id_ida]).where(models_parser.LocalType.e_type.in_(['struct', 'union']))
        )

        query = self.session.query(models_ida.IdaRawLocalType) \
            .filter(models_ida.IdaRawLocalType.id_ida.in_(local_types_table)) \
            .order_by(models_ida.IdaRawLocalType.id)

        count = query.count()
        if CONFIG['verbose']:
            print 'count local_type: {count}'.format(count=count)

        index_item = 0
        local_type_deps = []
        for item in query.all():
            index_item += 1
            members_type = self.__fetch_members(item.get_one_line())
            if len(members_type):
                id_members_q = (
                    select([models_ida.IdaRawLocalType.id_ida]).where(models_ida.IdaRawLocalType.name.in_(members_type))
                )

                dep_types = self.session.query(id_members_q).all()

                for dep_type in dep_types:
                    depend = models_parser.DependLocalType(
                        id_local_type=item.get_id(),
                        id_depend=dep_type.id_ida
                    )
                    local_type_deps.append(depend)

            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(local_type_deps)
        self.session.commit()

    def __linking_functions(self):
        query = self.session.query(models_parser.Function)
        count = query.count()
        if CONFIG['verbose']:
            print 'count local_type: {count}'.format(count=count)

        index_item = 0
        function_link = []
        for item in query.all():
            index_item += 1
            owner_name = item.get_owner_name()

            depend = models_parser.LinkFunctions(
                id_function=item.get_id(),
                owner_name=owner_name
            )
            function_link.append(depend)

            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(function_link)
        self.session.commit()

    def __link_local_types(self, parts):
        s_names = list()
        search_name = ''
        for p in parts:
            if len(search_name):
                search_name += '::'

            search_name += p
            s_names.append(search_name)

        q = self.session.query(models_ida.IdaRawLocalType)\
            .filter(models_ida.IdaRawLocalType.name.in_(s_names))\
            .order_by(func.length(models_ida.IdaRawLocalType.name))

        for lt in q.all():
            yield (lt.get_id(), lt.get_name())

    def __linking_local_types(self):
        query = self.session.query(models_ida.IdaRawLocalType)

        count = query.count()
        if CONFIG['verbose']:
            print 'count local_type: {count}'.format(count=count)

        index_item = 0
        link_local_types = []
        link_namespaces = []
        for item in query.all():
            index_item += 1

            parts = list(util_parser.split_name(item.get_name()))
            id_local_types = list(self.__link_local_types(parts))

            if len(id_local_types) >= 2:
                link_lt = models_parser.LinkLocalType(
                    id_parent=id_local_types[-2][0],
                    id_child=item.get_id()
                )
                link_local_types.append(link_lt)

            link_namespace = ''
            _diff = len(parts) - len(id_local_types)
            if _diff == 0:
                link_namespace = None
            else:
                for p in parts[:_diff]:
                    if len(link_namespace):
                        link_namespace += '::'

                    link_namespace += p

            link = models_parser.LinkNamespace(
                id_local_type=item.get_id(),
                namespace=link_namespace
            )

            link_namespaces.append(link)
            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(link_local_types)
        self.session.bulk_save_objects(link_namespaces)
        self.session.commit()
<commit_msg>Fixed infinity loop in dependencies<commit_after>import os
import re
import models_ida
import models_parser
import util_parser

from config import CONFIG
from sqlalchemy import select
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.sql.expression import func
from pida_types.types import IDA_TYPES

class IdaInfoParser(object):
    def __init__(self, db_file):
        self.db_file = db_file
        self.Session = sessionmaker()
        self.pattern_member = re.compile('(.*) .*\w+')

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.session.commit()

    def start(self):
        self.__create_connection()
        self.__drop_tables()
        self.__create_tables()
        self.__create_session()
        self.__parsing()
        self.__fetch_depend()
        self.__linking()

    def __parsing(self):
        self.__parsing_functions()
        self.__parsing_local_types()

    def __fetch_depend(self):
        self.__fetch_depend_functions()
        self.__fetch_depend_local_types()
        # todo : add fetch depend for typedef

    def __linking(self):
        self.__linking_functions()
        self.__linking_local_types()
        # todo : add link for typedef

    def __drop_tables(self):
        models_parser.Base.metadata.drop_all(self.engine_db)

    def __create_tables(self):
        models_ida.Base.metadata.create_all(self.engine_db)
        models_parser.Base.metadata.create_all(self.engine_db)

    def __create_session(self):
        self.Session.configure(bind=self.engine_db, autocommit=False)
        self.session = self.Session()

    def __create_connection(self):
        base_dir = os.path.dirname(self.db_file)
        if not os.path.exists(base_dir):
            os.makedirs(base_dir)

        if not os.path.exists(self.db_file):
            open(self.db_file, 'a').close()

        self.engine_db = create_engine('sqlite:///' + self.db_file, echo=CONFIG['sql_verbose'])

    def __parsing_functions(self):
        query = self.session.query(models_ida.IdaRawFunctions)
        count = query.count()
        if CONFIG['verbose']:
            print 'count functions: {count}'.format(count=count)

        index_item = 0
        functions = []
        for item in query.all():
            index_item += 1
            function = models_parser.Function(
                id_ida=item.get_id(),
                raw_name=item.get_name(),
                ida_type=item.get_type(),
                ida_fields=item.get_args_name(),
            )
            try:
                function.parsing()
                functions.append(function)
            except:
                print '[Error] function with id = {id}. catch exception'.format(id=item.get_id())

            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(functions)
        self.session.commit()

    def __parsing_local_types(self):
        query = self.session.query(models_ida.IdaRawLocalType)
        count = query.count()
        if CONFIG['verbose']:
            print 'count local_type: {count}'.format(count=count)

        index_item = 0
        local_types = []
        for item in query.all():
            index_item += 1
            local_type = models_parser.LocalType(
                id_ida=item.get_id(),
                raw_multi=item.get_type(),
            )
            local_type.parsing()
            local_types.append(local_type)

            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(local_types)
        self.session.commit()

    def __fetch_local_type(self, all_types):
        local_types = []
        for t in all_types:
            a_type = t
            while True:
                value = a_type['value']
                if a_type['idt'] == IDA_TYPES['local_type']:
                    local_types.append(a_type)
                    break

                elif a_type['idt'] == IDA_TYPES['function']:
                    local_types.extend(self.__fetch_local_type(value['args_type']))
                    value = value['ret_type']

                elif a_type['idt'] == IDA_TYPES['array']:
                    value = value['type']

                if type(value) is not type({}):
                    break

                a_type = value

        return local_types

    def __fetch_depend_functions(self):
        query = self.session.query(models_parser.Function)
        count = query.count()
        if CONFIG['verbose']:
            print 'count functions: {count}'.format(count=count)

        function_deps = []
        index_item = 0
        for item in query.all():
            index_item += 1
            all_types = item.get_args_type()
            all_types.append(item.get_return_type())

            dep_types = self.__fetch_local_type(all_types)

            for dep_type in dep_types:
                depend = models_parser.DependFunction(
                    id_function=item.get_id(),
                    id_local_type=dep_type['value']
                )
                function_deps.append(depend)

            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(function_deps)
        self.session.commit()

    def __fetch_members(self, one_line):
        local_types = set()

        bracket = [one_line.find('{'), one_line.rfind('}')]
        if -1 in bracket:
            return local_types

        raw_parents = one_line[:bracket[0] - 1]
        inheritance = raw_parents.find(':')
        if inheritance != -1:
            parents = raw_parents[inheritance + 1:]
            tmpl_pairs = list(util_parser.get_pairs_sym(parents, '<', '>'))
            if 0 == len(tmpl_pairs):
                local_types.update(parents.split(','))
            else:
                sets = set()
                for tmpl_pair in tmpl_pairs:
                    count_join = 0
                    for tmpl_pair2 in tmpl_pairs:
                        if not tmpl_pair2[0] <= tmpl_pair[0] <= tmpl_pair2[1]:
                            continue

                        if not tmpl_pair2[0] <= tmpl_pair[1] <= tmpl_pair2[1]:
                            continue

                        count_join += 1

                    sets.add((count_join, tmpl_pair))

                parents += ','
                for v in sets:
                    (count, pair) = v
                    if count != 1:
                        continue

                    if parents[pair[1] + 1] == ',':
                        parents = parents[:pair[1] + 1] + '!' + parents[pair[1] + 1:]

                for parent in parents.split('!,'):
                    local_types.add(parent.strip())

        raw_members = one_line[bracket[0] + 1:bracket[1]]
        for member in raw_members.split(';'):
            type_member = self.pattern_member.search(member)
            if type_member is None:
                continue

            type_member = type_member.group(1).strip()
            local_types.add(type_member)

        ret_local_types = set()
        for lt in local_types:
            ret_local_types.add(lt.strip())

        return ret_local_types

    def __fetch_depend_local_types(self):
        local_types_table = (
            select([models_parser.LocalType.id_ida]).where(models_parser.LocalType.e_type.in_(['struct', 'union']))
        )

        query = self.session.query(models_ida.IdaRawLocalType) \
            .filter(models_ida.IdaRawLocalType.id_ida.in_(local_types_table)) \
            .order_by(models_ida.IdaRawLocalType.id)

        count = query.count()
        if CONFIG['verbose']:
            print 'count local_type: {count}'.format(count=count)

        index_item = 0
        local_type_deps = []
        for item in query.all():
            index_item += 1
            members_type = self.__fetch_members(item.get_one_line())
            if len(members_type):
                id_members_q = (
                    select([models_ida.IdaRawLocalType.id_ida]).where(models_ida.IdaRawLocalType.name.in_(members_type))
                )

                dep_types = self.session.query(id_members_q).all()

                for dep_type in dep_types:
                    if item.get_id() != dep_type.id_ida:
                        depend = models_parser.DependLocalType(
                            id_local_type=item.get_id(),
                            id_depend=dep_type.id_ida
                        )
                        local_type_deps.append(depend)

            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(local_type_deps)
        self.session.commit()

    def __linking_functions(self):
        query = self.session.query(models_parser.Function)
        count = query.count()
        if CONFIG['verbose']:
            print 'count local_type: {count}'.format(count=count)

        index_item = 0
        function_link = []
        for item in query.all():
            index_item += 1
            owner_name = item.get_owner_name()

            depend = models_parser.LinkFunctions(
                id_function=item.get_id(),
                owner_name=owner_name
            )
            function_link.append(depend)

            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(function_link)
        self.session.commit()

    def __link_local_types(self, parts):
        s_names = list()
        search_name = ''
        for p in parts:
            if len(search_name):
                search_name += '::'

            search_name += p
            s_names.append(search_name)

        q = self.session.query(models_ida.IdaRawLocalType)\
            .filter(models_ida.IdaRawLocalType.name.in_(s_names))\
            .order_by(func.length(models_ida.IdaRawLocalType.name))

        for lt in q.all():
            yield (lt.get_id(), lt.get_name())

    def __linking_local_types(self):
        query = self.session.query(models_ida.IdaRawLocalType)

        count = query.count()
        if CONFIG['verbose']:
            print 'count local_type: {count}'.format(count=count)

        index_item = 0
        link_local_types = []
        link_namespaces = []
        for item in query.all():
            index_item += 1

            parts = list(util_parser.split_name(item.get_name()))
            id_local_types = list(self.__link_local_types(parts))

            if len(id_local_types) >= 2:
                link_lt = models_parser.LinkLocalType(
                    id_parent=id_local_types[-2][0],
                    id_child=item.get_id()
                )
                link_local_types.append(link_lt)

            link_namespace = ''
            _diff = len(parts) - len(id_local_types)
            if _diff == 0:
                link_namespace = None
            else:
                for p in parts[:_diff]:
                    if len(link_namespace):
                        link_namespace += '::'

                    link_namespace += p

            link = models_parser.LinkNamespace(
                id_local_type=item.get_id(),
                namespace=link_namespace
            )

            link_namespaces.append(link)
            if CONFIG['verbose']:
                if (index_item % CONFIG['page_size'] == 0) or (count - index_item == 0):
                    print 'items({current}/{count_item})'.format(current=index_item, count_item=count)

        self.session.bulk_save_objects(link_local_types)
        self.session.bulk_save_objects(link_namespaces)
        self.session.commit()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

"""
Supporting functions for Bishop
"""

__author__ = "Julian Jara-Ettinger"
__license__ = "MIT"

import ConfigParser
import os
import pickle
import pkg_resources
from PosteriorContainer import *
from Observer import *
from Map import *
from Agent import *

def SaveSamples(Container, Name):
    """
    Save object as a pickle file.

    Args:
        Container (PosteriorContainer): PosteriorContainer object
        Name (string): Filename. Function adds ".p" extension if it's not provided
    """
    try:
        if Name[-2:] != ".p":
            Name = Name + ".p"
        pickle.dump(Container, open(Name, "wb"))
    except Exception as error:
        print error

def LoadSamples(FileName):
    """
    Load samples from a pickle file

    Args:
        FileName (str): filename

    returns:
        Samples
    """
    try:
        Samples = pickle.load(open(FileName, "rb"))
        return Samples
    except Exception as error:
        print error

def AnalyzeSamples(FileName):
    """
    Print sample summary from a pickle file

    Args:
        FileName (str): filename
    """
    try:
        Samples = pickle.load(open(FileName, "rb"))
        Samples.LongSummary()
    except Exception as error:
        print error

def LoadObserverFromPC(PostCont):
    """
    Load observer object from a PosteriorContainer object.

    Args:
        Postcont (PosteriorContainer): Posterior container object

    Returns:
        Observer object
    """
    if PostCont.MapFile is None:
        print "No map associated with samples. Cannot load observer."
        return None
    else:
        try:
            Observer = LoadObserver(PostCont.MapFile, False, True)
            Observer.SetStartingPoint(PostCont.StartingPoint, False)
            Observer.PrintMap()
            return Observer
        except Exception as error:
            print error

def GetMapList(CurrDir):
    """
    Get list of map files in a directory.

    Args:
        CurrDir (str): Search directory.

    Returns:
        files (list): List of files.å
    """
    files = []
    for File in os.listdir(CurrDir):
        # Check if it's a file
        if File.endswith(".ini"):
            files.append(File)
    return files

def ShowAvailableMaps(Match=""):
    """
    Print list of maps in Bishop library.

    Args:
        Match (String): Only print maps that contain string.
    """
    # Create an empty dictionary.
    results = {}
    BaseDirectory = os.path.dirname(__file__) + "/Maps/"
    Files = GetMapList(BaseDirectory)
    if Files != []:
        results['Bishop main maps'] = Files
    for Item in os.listdir(BaseDirectory):
        TempDir = os.path.join(BaseDirectory, Item)
        if os.path.isdir(TempDir):
            # Add to dictionary
            Files = GetMapList(TempDir)
            if Files != []:
                results[Item] = GetMapList(TempDir)
    # print maps
    for key in results:
        sys.stdout.write(key + ":\n")
        for i in results[key]:
            if Match in i:
                sys.stdout.write("\t" + i + "\n")
    sys.stdout.write("\n")

def LocateFile(CurrDir, filename):
    """
    Search for file recursively.

    Args:
        CurrDir (str): Search directory
        filename (str): filename

    Returns:
        dir (str): Location of the file (if found; None otherwise)
    """
    for Currfile in os.listdir(CurrDir):
        test = os.path.join(CurrDir, Currfile)
        if os.path.isdir(test):
            res = LocateFile(test, filename)
            if res is not None:
                return res
        else:
            # it's a file
            if Currfile == filename:
                return CurrDir

def LoadObserver(MapConfig, Revise=False, Silent=False):
    """
    Load a map. If map isn't found in Bishop's library the
    function searches for the map in your working directory.

    Args:
        MapConfig (str): Name of map to load
        Revise (bool): When true, user manually confirms or overrides parameters.
        Silent (bool): If false then function doesn't print map.

    Returns:
        Observer object
    """
    try:
        Local = False
        if Revise:
            sys.stdout.write(
                "\nPress enter to accept the argument or type in the new value to replace it.\n\n")
        Config = ConfigParser.ConfigParser()
        FilePath = os.path.dirname(__file__) + "/Maps/"
        FilePath = LocateFile(FilePath, MapConfig + ".ini")
        if FilePath is not None:
            FilePath = FilePath + "/" + MapConfig + ".ini"
        #########################
        ## Load .ini map first ##
        #########################
        else:
            print "Map not in library. Checking local directory..."
            FilePath = MapConfig + ".ini"
            Local = True
            if not os.path.isfile(FilePath):
                print "ERROR: Map not found."
                return None
        Config.read(FilePath)
    except Exception as error:
        print error

    # Agent parameter section
    #########################
    if not Config.has_section("AgentParameters"):
        print "ERROR: AgentParameters block missing."
        return None
    if Config.has_option("AgentParameters", "Prior"):
        CostPrior = Config.get("AgentParameters", "Prior")
        RewardPrior = CostPrior
        if Revise:
            temp = raw_input("Prior (" + str(Prior) + "):")
            if temp != '':
                CostPrior = str(temp)
                RewardPrior = CostPrior
    else:
        if Config.has_option("AgentParameters", "CostPrior"):
            CostPrior = Config.get("AgentParameters", "CostPrior")
            if Revise:
                temp = raw_input("CostPrior (" + str(CostPrior) + "):")
                if temp != '':
                    CostPrior = str(temp)
        else:
            print "WARNING: No cost prior specified in AgentParameters. Use Agent.Priors() to see list of priors"
            return None
        if Config.has_option("AgentParameters", "RewardPrior"):
            RewardPrior = Config.get("AgentParameters", "RewardPrior")
            if Revise:
                temp = raw_input("RewardPrior (" + str(RewardPrior) + "):")
                if temp != '':
                    RewardPrior = str(temp)
        else:
            print "WARNING: No reward prior specified in AgentParameters. Use Agent.Priors() to see list of priors"
            return None
    if Config.has_option("AgentParameters", "Capacity"):
        Capacity = Config.getint("AgentParameters", "Capacity")
    else:
        Capacity = -1
    if Revise:
        temp = raw_input(
            "Agent capacity (" + str(Capacity) + "; -1 means unlimited):")
        if temp != '':
            Capacity = int(temp)
    if Config.has_option("AgentParameters", "Restrict"):
        Restrict = Config.getboolean("AgentParameters", "Restrict")
    else:
        print "Setting restrict to false (i.e., all terrains are equal)"
        Restrict = False
    if Config.has_option("AgentParameters", "SoftmaxChoice"):
        SoftmaxChoice = Config.getboolean("AgentParameters", "SoftmaxChoice")
    else:
        print "Softmaxing choices"
        SoftmaxChoice = True
    if Revise:
        temp = raw_input("Softmax choices (" + str(SoftmaxChoice) + "):")
        if temp != '':
            if temp == 'True':
                SoftmaxChoice = True
            elif temp == 'False':
                SoftmaxChoice = False
            else:
                sys.stdout.write("Not a valid choice. Ignoring.\n")
    if Config.has_option("AgentParameters", "SoftmaxAction"):
        SoftmaxAction = Config.getboolean("AgentParameters", "SoftmaxAction")
    else:
        print "Softmaxing actions"
        SoftmaxAction = True
    if Revise:
        temp = raw_input("Softmax actions (" + str(SoftmaxAction) + "):")
        if temp != '':
            if temp == 'True':
                SoftmaxAction = True
            elif temp == 'False':
                SoftmaxAction = False
            else:
                sys.stdout.write("Not a valid choice. Ignoring.\n")
    if Config.has_option("AgentParameters", "choiceTau"):
        choiceTau = Config.getfloat("AgentParameters", "choiceTau")
    else:
        if SoftmaxChoice:
            print "Setting choice softmax to 0.01"
            choiceTau = 0.01
        else:
            # Doesn't matter; won't be used.
            choiceTau = 0
    if Revise:
        temp = raw_input("Choice tau (" + str(choiceTau) + "):")
        if temp != '':
            choiceTau = float(temp)
    if Config.has_option("AgentParameters", "actionTau"):
        actionTau = Config.getfloat("AgentParameters", "actionTau")
    else:
        if SoftmaxAction:
            print "Setting action softmax to 0.01"
            actionTau = 0.01
        else:
            # Doesn't matter; won't be used.
            actionTau = 0
    if Revise:
        temp = raw_input("Action tau (" + str(actionTau) + "):")
        if temp != '':
            actionTau = float(temp)
    if Config.has_option("AgentParameters", "CostParameters"):
        CostParameters = Config.get("AgentParameters", "CostParameters")
        if Revise:
            temp = raw_input("Cost parameters (" + str(CostParameters) + "):")
            if temp != '':
                CostParameters = temp
        CostParameters = CostParameters.split()
        CostParameters = [float(i) for i in CostParameters]
    else:
        print "ERROR: Missing cost parameters for prior sampling in AgentParameters block."
        return None
    if Config.has_option("AgentParameters", "RewardParameters"):
        RewardParameters = Config.get("AgentParameters", "RewardParameters")
        if Revise:
            temp = raw_input(
                "Reward parameters (" + str(RewardParameters) + "):")
            if temp != '':
                RewardParameters = temp
        RewardParameters = [float(i) for i in RewardParameters.split()]
    else:
        print "ERROR: Missing cost parameters for prior sampling in AgentParameters block."
        return None
    if Config.has_option("AgentParameters", "PNull"):
        CNull = Config.getfloat("AgentParameters", "PNull")
        RNull = CNull
    else:
        if Config.has_option("AgentParameters", "CNull"):
            CNull = Config.getfloat("AgentParameters", "CNull")
        else:
            print "WARNING: No probability of terrains having null cost. Setting to 0."
            CNull = 0
        if Config.has_option("AgentParameters", "RNull"):
            RNull = Config.getfloat("AgentParameters", "RNull")
        else:
            print "WARNING: No probability of terrains having null cost. Setting to 0."
            RNull = 0
    if Revise:
        temp = raw_input("Null cost paramter (" + str(CNull) + "):")
        if temp != '':
            CNull = float(temp)
        temp = raw_input("Null reward paramter (" + str(RNull) + "):")
        if temp != '':
            RNull = float(temp)
    # Map parameter section
    #######################
    if not Config.has_section("MapParameters"):
        print "ERROR: MapParameters block missing."
        return None
    if Config.has_option("MapParameters", "DiagonalTravel"):
        DiagonalTravel = Config.getboolean(
            "MapParameters", "DiagonalTravel")
    else:
        print "Allowing diagonal travel"
        DiagonalTravel = True
    if Revise:
        temp = raw_input("Diagonal travel (" + str(DiagonalTravel) + "):")
        if temp != '':
            if temp == "True":
                DiagonalTravel = True
            elif temp == "False":
                DiagonalTravel = False
            else:
                sys.stdout.write("Not a valid choice. Ignoring.\n")
    if Config.has_option("MapParameters", "StartingPoint"):
        StartingPoint = Config.getint(
            "MapParameters", "StartingPoint")
    else:
        print "ERROR: Missing starting point in MapParameters block."
        return None
    if Revise:
        temp = raw_input("Starting point (" + str(StartingPoint) + "):")
        if temp != '':
            StartingPoint = int(temp)
    if Config.has_option("MapParameters", "ExitState"):
        ExitState = Config.getint(
            "MapParameters", "ExitState")
    else:
        print "ERROR: Missing exit state in MapParameters block."
        return None
    if Revise:
        temp = raw_input("Exit state (" + str(ExitState) + "):")
        if temp != '':
            ExitState = int(temp)
    try:
        if Config.has_option("MapParameters", "MapName"):
            MapName = Config.get(
                "MapParameters", "MapName")
            if not Local:
                TerrainPath = os.path.dirname(__file__) + "/Maps/"
                TerrainPath = os.path.join(
                    LocateFile(TerrainPath, MapName), MapName)
            else:
                TerrainPath = MapName
            f = open(TerrainPath, "r")
            MapLoad = True
            StateTypes = []
            StateNames = []
            mapheight = 0
            for line in iter(f):
                if MapLoad:
                    states = [int(i) for i in list(line.rstrip())]
                    if states == []:
                        MapLoad = False
                    else:
                        mapheight += 1
                        StateTypes.extend(states)
                else:
                    statename = line.rstrip()
                    if statename != "":
                        StateNames.append(statename)
            f.close()
            mapwidth = len(StateTypes) / mapheight
        else:
            print "ERROR: Missing map name"
            return None
    except:
        print "ERROR: Cannot load map layout."
        # raise
        return None
    # Load object information
    #########################
    if not Config.has_section("Objects"):
        print "ERROR: Objects block missing."
        return None
    else:
        if Config.has_option("Objects", "ObjectLocations"):
            ObjectLocations = Config.get("Objects", "ObjectLocations")
            ObjectLocations = [int(i) for i in ObjectLocations.split()]
            HasObjects = True
        else:
            print "WARNING: No objects in map (Agent will always go straight home)."
            HasObjects = False
        if HasObjects:
            if Config.has_option("Objects", "ObjectTypes"):
                ObjectTypes = Config.get("Objects", "ObjectTypes")
                ObjectTypes = [int(i) for i in ObjectTypes.split()]
                if len(ObjectTypes) != len(ObjectLocations):
                    print "Error: ObjectLocations and ObjectTypes should have the same length"
                    return None
            else:
                print "WARNING: No information about object types. Setting all to same kind."
                ObjectTypes = [0] * len(ObjectLocations)
            if Config.has_option("Objects", "ObjectNames"):
                ObjectNames = Config.get("Objects", "ObjectNames")
                ObjectNames = [str(i) for i in ObjectNames.split()]
            else:
                ObjectNames = None
        else:
            ObjectTypes = []
            ObjectNames = None
    # Create objects!
    try:
        MyMap = Map()
        MyMap.BuildGridWorld(mapwidth, mapheight, DiagonalTravel)
        MyMap.InsertObjects(ObjectLocations, ObjectTypes, ObjectNames)
        MyMap.StateTypes = StateTypes
        MyMap.StateNames = StateNames
        MyMap.AddStartingPoint(StartingPoint)
        MyMap.AddExitState(ExitState)
        if not Silent:
            sys.stdout.write("\n")
            MyMap.PrintMap()
        MyAgent = Agent(MyMap, CostPrior, RewardPrior, CostParameters, RewardParameters, Capacity,
                        SoftmaxChoice, SoftmaxAction, choiceTau, actionTau, CNull, RNull, Restrict)
        return Observer(MyAgent, MyMap)
    except Exception as error:
        print error

def AboutBishop():
    """
    About.
    """
    sys.stdout.write(
        "    ___      ___      ___      ___      ___      ___   \n")
    sys.stdout.write(
        "   /\\  \\    /\\  \\    /\\  \\    /\\__\\    /\\  \\    /\\  \\  \n")
    sys.stdout.write(
        "  /::\\  \\  _\\:\\  \\  /::\\  \\  /:/__/_  /::\\  \\  /::\\  \\ \n")
    sys.stdout.write(
        " /::\\:\\__\\/\\/::\\__\\/\\:\\:\\__\\/::\\/\\__\\/:/\\:\\__\\/::\\:\\__\\\n")
    sys.stdout.write(
        " \\:\\::/  /\\::/\\/__/\\:\\:\\/__/\\/\\::/  /\\:\\/:/  /\\/\\::/  /\n")
    sys.stdout.write(
        "  \\::/  /  \\:\\__\\   \\::/  /   /:/  /  \\::/  /    \\/__/ \n")
    sys.stdout.write(
        "   \\/__/    \\/__/    \\/__/    \\/__/    \\/__/           \n\n")
    sys.stdout.write(
        "Bishop. V " + str(pkg_resources.get_distribution("Bishop").version) + "\n")
    sys.stdout.write("http://github.com/julianje/Bishop\n")
    sys.stdout.write("Julian Jara-Ettinger. jjara@mit.edu\n\n")
    sys.stdout.write(
        "Results using Bishop 1.0.0: Jara-Ettinger, J., Schulz, L. E., & Tenenbaum J. B. (2015). The naive utility calculus: Joint inferences about the costs and rewards of actions. In Proceedings of the 37th Annual Conference of the Cognitive Science Society.\n\n")
<commit_msg>Minor bug fix<commit_after># -*- coding: utf-8 -*-

"""
Supporting functions for Bishop
"""

__author__ = "Julian Jara-Ettinger"
__license__ = "MIT"

import ConfigParser
import os
import pickle
import pkg_resources
from PosteriorContainer import *
from Observer import *
from Map import *
from Agent import *

def SaveSamples(Container, Name):
    """
    Save object as a pickle file.

    Args:
        Container (PosteriorContainer): PosteriorContainer object
        Name (string): Filename. Function adds ".p" extension if it's not provided
    """
    try:
        if Name[-2:] != ".p":
            Name = Name + ".p"
        pickle.dump(Container, open(Name, "wb"))
    except Exception as error:
        print error

def LoadSamples(FileName):
    """
    Load samples from a pickle file

    Args:
        FileName (str): filename

    returns:
        Samples
    """
    try:
        Samples = pickle.load(open(FileName, "rb"))
        return Samples
    except Exception as error:
        print error

def AnalyzeSamples(FileName):
    """
    Print sample summary from a pickle file

    Args:
        FileName (str): filename
    """
    try:
        Samples = pickle.load(open(FileName, "rb"))
        Samples.LongSummary()
    except Exception as error:
        print error

def LoadObserverFromPC(PostCont):
    """
    Load observer object from a PosteriorContainer object.

    Args:
        Postcont (PosteriorContainer): Posterior container object

    Returns:
        Observer object
    """
    if PostCont.MapFile is None:
        print "No map associated with samples. Cannot load observer."
        return None
    else:
        try:
            Observer = LoadObserver(PostCont.MapFile, False, True)
            Observer.SetStartingPoint(PostCont.StartingPoint, False)
            Observer.PrintMap()
            return Observer
        except Exception as error:
            print error

def GetMapList(CurrDir):
    """
    Get list of map files in a directory.

    Args:
        CurrDir (str): Search directory.

    Returns:
        files (list): List of files.å
    """
    files = []
    for File in os.listdir(CurrDir):
        # Check if it's a file
        if File.endswith(".ini"):
            files.append(File)
    return files

def ShowAvailableMaps(Match=""):
    """
    Print list of maps in Bishop library.

    Args:
        Match (String): Only print maps that contain string.
    """
    # Create an empty dictionary.
    results = {}
    BaseDirectory = os.path.dirname(__file__) + "/Maps/"
    Files = GetMapList(BaseDirectory)
    if Files != []:
        results['Bishop main maps'] = Files
    for Item in os.listdir(BaseDirectory):
        TempDir = os.path.join(BaseDirectory, Item)
        if os.path.isdir(TempDir):
            # Add to dictionary
            Files = GetMapList(TempDir)
            if Files != []:
                results[Item] = GetMapList(TempDir)
    # print maps
    for key in results:
        sys.stdout.write(key + ":\n")
        for i in results[key]:
            if Match in i:
                sys.stdout.write("\t" + i + "\n")
    sys.stdout.write("\n")

def LocateFile(CurrDir, filename):
    """
    Search for file recursively.

    Args:
        CurrDir (str): Search directory
        filename (str): filename

    Returns:
        dir (str): Location of the file (if found; None otherwise)
    """
    for Currfile in os.listdir(CurrDir):
        test = os.path.join(CurrDir, Currfile)
        if os.path.isdir(test):
            res = LocateFile(test, filename)
            if res is not None:
                return res
        else:
            # it's a file
            if Currfile == filename:
                return CurrDir

def LoadObserver(MapConfig, Revise=False, Silent=False):
    """
    Load a map. If map isn't found in Bishop's library the
    function searches for the map in your working directory.

    Args:
        MapConfig (str): Name of map to load
        Revise (bool): When true, user manually confirms or overrides parameters.
        Silent (bool): If false then function doesn't print map.

    Returns:
        Observer object
    """
    try:
        Local = False
        if Revise:
            sys.stdout.write(
                "\nPress enter to accept the argument or type in the new value to replace it.\n\n")
        Config = ConfigParser.ConfigParser()
        FilePath = os.path.dirname(__file__) + "/Maps/"
        FilePath = LocateFile(FilePath, MapConfig + ".ini")
        if FilePath is not None:
            FilePath = FilePath + "/" + MapConfig + ".ini"
        #########################
        ## Load .ini map first ##
        #########################
        else:
            print "Map not in library. Checking local directory..."
            FilePath = MapConfig + ".ini"
            Local = True
            if not os.path.isfile(FilePath):
                print "ERROR: Map not found."
                return None
        Config.read(FilePath)
    except Exception as error:
        print error

    # Agent parameter section
    #########################
    if not Config.has_section("AgentParameters"):
        print "ERROR: AgentParameters block missing."
        return None
    if Config.has_option("AgentParameters", "Prior"):
        CostPrior = Config.get("AgentParameters", "Prior")
        RewardPrior = CostPrior
        if Revise:
            temp = raw_input("Prior (" + str(CostPrior) + "):")
            if temp != '':
                CostPrior = str(temp)
                RewardPrior = CostPrior
    else:
        if Config.has_option("AgentParameters", "CostPrior"):
            CostPrior = Config.get("AgentParameters", "CostPrior")
            if Revise:
                temp = raw_input("CostPrior (" + str(CostPrior) + "):")
                if temp != '':
                    CostPrior = str(temp)
        else:
            print "WARNING: No cost prior specified in AgentParameters. Use Agent.Priors() to see list of priors"
            return None
        if Config.has_option("AgentParameters", "RewardPrior"):
            RewardPrior = Config.get("AgentParameters", "RewardPrior")
            if Revise:
                temp = raw_input("RewardPrior (" + str(RewardPrior) + "):")
                if temp != '':
                    RewardPrior = str(temp)
        else:
            print "WARNING: No reward prior specified in AgentParameters. Use Agent.Priors() to see list of priors"
            return None
    if Config.has_option("AgentParameters", "Capacity"):
        Capacity = Config.getint("AgentParameters", "Capacity")
    else:
        Capacity = -1
    if Revise:
        temp = raw_input(
            "Agent capacity (" + str(Capacity) + "; -1 means unlimited):")
        if temp != '':
            Capacity = int(temp)
    if Config.has_option("AgentParameters", "Restrict"):
        Restrict = Config.getboolean("AgentParameters", "Restrict")
    else:
        print "Setting restrict to false (i.e., all terrains are equal)"
        Restrict = False
    if Config.has_option("AgentParameters", "SoftmaxChoice"):
        SoftmaxChoice = Config.getboolean("AgentParameters", "SoftmaxChoice")
    else:
        print "Softmaxing choices"
        SoftmaxChoice = True
    if Revise:
        temp = raw_input("Softmax choices (" + str(SoftmaxChoice) + "):")
        if temp != '':
            if temp == 'True':
                SoftmaxChoice = True
            elif temp == 'False':
                SoftmaxChoice = False
            else:
                sys.stdout.write("Not a valid choice. Ignoring.\n")
    if Config.has_option("AgentParameters", "SoftmaxAction"):
        SoftmaxAction = Config.getboolean("AgentParameters", "SoftmaxAction")
    else:
        print "Softmaxing actions"
        SoftmaxAction = True
    if Revise:
        temp = raw_input("Softmax actions (" + str(SoftmaxAction) + "):")
        if temp != '':
            if temp == 'True':
                SoftmaxAction = True
            elif temp == 'False':
                SoftmaxAction = False
            else:
                sys.stdout.write("Not a valid choice. Ignoring.\n")
    if Config.has_option("AgentParameters", "choiceTau"):
        choiceTau = Config.getfloat("AgentParameters", "choiceTau")
    else:
        if SoftmaxChoice:
            print "Setting choice softmax to 0.01"
            choiceTau = 0.01
        else:
            # Doesn't matter; won't be used.
            choiceTau = 0
    if Revise:
        temp = raw_input("Choice tau (" + str(choiceTau) + "):")
        if temp != '':
            choiceTau = float(temp)
    if Config.has_option("AgentParameters", "actionTau"):
        actionTau = Config.getfloat("AgentParameters", "actionTau")
    else:
        if SoftmaxAction:
            print "Setting action softmax to 0.01"
            actionTau = 0.01
        else:
            # Doesn't matter; won't be used.
            actionTau = 0
    if Revise:
        temp = raw_input("Action tau (" + str(actionTau) + "):")
        if temp != '':
            actionTau = float(temp)
    if Config.has_option("AgentParameters", "CostParameters"):
        CostParameters = Config.get("AgentParameters", "CostParameters")
        if Revise:
            temp = raw_input("Cost parameters (" + str(CostParameters) + "):")
            if temp != '':
                CostParameters = temp
        CostParameters = CostParameters.split()
        CostParameters = [float(i) for i in CostParameters]
    else:
        print "ERROR: Missing cost parameters for prior sampling in AgentParameters block."
        return None
    if Config.has_option("AgentParameters", "RewardParameters"):
        RewardParameters = Config.get("AgentParameters", "RewardParameters")
        if Revise:
            temp = raw_input(
                "Reward parameters (" + str(RewardParameters) + "):")
            if temp != '':
                RewardParameters = temp
        RewardParameters = [float(i) for i in RewardParameters.split()]
    else:
        print "ERROR: Missing cost parameters for prior sampling in AgentParameters block."
        return None
    if Config.has_option("AgentParameters", "PNull"):
        CNull = Config.getfloat("AgentParameters", "PNull")
        RNull = CNull
    else:
        if Config.has_option("AgentParameters", "CNull"):
            CNull = Config.getfloat("AgentParameters", "CNull")
        else:
            print "WARNING: No probability of terrains having null cost. Setting to 0."
            CNull = 0
        if Config.has_option("AgentParameters", "RNull"):
            RNull = Config.getfloat("AgentParameters", "RNull")
        else:
            print "WARNING: No probability of terrains having null cost. Setting to 0."
            RNull = 0
    if Revise:
        temp = raw_input("Null cost paramter (" + str(CNull) + "):")
        if temp != '':
            CNull = float(temp)
        temp = raw_input("Null reward paramter (" + str(RNull) + "):")
        if temp != '':
            RNull = float(temp)
    # Map parameter section
    #######################
    if not Config.has_section("MapParameters"):
        print "ERROR: MapParameters block missing."
        return None
    if Config.has_option("MapParameters", "DiagonalTravel"):
        DiagonalTravel = Config.getboolean(
            "MapParameters", "DiagonalTravel")
    else:
        print "Allowing diagonal travel"
        DiagonalTravel = True
    if Revise:
        temp = raw_input("Diagonal travel (" + str(DiagonalTravel) + "):")
        if temp != '':
            if temp == "True":
                DiagonalTravel = True
            elif temp == "False":
                DiagonalTravel = False
            else:
                sys.stdout.write("Not a valid choice. Ignoring.\n")
    if Config.has_option("MapParameters", "StartingPoint"):
        StartingPoint = Config.getint(
            "MapParameters", "StartingPoint")
    else:
        print "ERROR: Missing starting point in MapParameters block."
        return None
    if Revise:
        temp = raw_input("Starting point (" + str(StartingPoint) + "):")
        if temp != '':
            StartingPoint = int(temp)
    if Config.has_option("MapParameters", "ExitState"):
        ExitState = Config.getint(
            "MapParameters", "ExitState")
    else:
        print "ERROR: Missing exit state in MapParameters block."
        return None
    if Revise:
        temp = raw_input("Exit state (" + str(ExitState) + "):")
        if temp != '':
            ExitState = int(temp)
    try:
        if Config.has_option("MapParameters", "MapName"):
            MapName = Config.get(
                "MapParameters", "MapName")
            if not Local:
                TerrainPath = os.path.dirname(__file__) + "/Maps/"
                TerrainPath = os.path.join(
                    LocateFile(TerrainPath, MapName), MapName)
            else:
                TerrainPath = MapName
            f = open(TerrainPath, "r")
            MapLoad = True
            StateTypes = []
            StateNames = []
            mapheight = 0
            for line in iter(f):
                if MapLoad:
                    states = [int(i) for i in list(line.rstrip())]
                    if states == []:
                        MapLoad = False
                    else:
                        mapheight += 1
                        StateTypes.extend(states)
                else:
                    statename = line.rstrip()
                    if statename != "":
                        StateNames.append(statename)
            f.close()
            mapwidth = len(StateTypes) / mapheight
        else:
            print "ERROR: Missing map name"
            return None
    except:
        print "ERROR: Cannot load map layout."
        # raise
        return None
    # Load object information
    #########################
    if not Config.has_section("Objects"):
        print "ERROR: Objects block missing."
        return None
    else:
        if Config.has_option("Objects", "ObjectLocations"):
            ObjectLocations = Config.get("Objects", "ObjectLocations")
            ObjectLocations = [int(i) for i in ObjectLocations.split()]
            HasObjects = True
        else:
            print "WARNING: No objects in map (Agent will always go straight home)."
            HasObjects = False
        if HasObjects:
            if Config.has_option("Objects", "ObjectTypes"):
                ObjectTypes = Config.get("Objects", "ObjectTypes")
                ObjectTypes = [int(i) for i in ObjectTypes.split()]
                if len(ObjectTypes) != len(ObjectLocations):
                    print "Error: ObjectLocations and ObjectTypes should have the same length"
                    return None
            else:
                print "WARNING: No information about object types. Setting all to same kind."
                ObjectTypes = [0] * len(ObjectLocations)
            if Config.has_option("Objects", "ObjectNames"):
                ObjectNames = Config.get("Objects", "ObjectNames")
                ObjectNames = [str(i) for i in ObjectNames.split()]
            else:
                ObjectNames = None
        else:
            ObjectTypes = []
            ObjectNames = None
    # Create objects!
    try:
        MyMap = Map()
        MyMap.BuildGridWorld(mapwidth, mapheight, DiagonalTravel)
        MyMap.InsertObjects(ObjectLocations, ObjectTypes, ObjectNames)
        MyMap.StateTypes = StateTypes
        MyMap.StateNames = StateNames
        MyMap.AddStartingPoint(StartingPoint)
        MyMap.AddExitState(ExitState)
        if not Silent:
            sys.stdout.write("\n")
            MyMap.PrintMap()
        MyAgent = Agent(MyMap, CostPrior, RewardPrior, CostParameters, RewardParameters, Capacity,
                        SoftmaxChoice, SoftmaxAction, choiceTau, actionTau, CNull, RNull, Restrict)
        return Observer(MyAgent, MyMap)
    except Exception as error:
        print error

def AboutBishop():
    """
    About.
    """
    sys.stdout.write(
        "    ___      ___      ___      ___      ___      ___   \n")
    sys.stdout.write(
        "   /\\  \\    /\\  \\    /\\  \\    /\\__\\    /\\  \\    /\\  \\  \n")
    sys.stdout.write(
        "  /::\\  \\  _\\:\\  \\  /::\\  \\  /:/__/_  /::\\  \\  /::\\  \\ \n")
    sys.stdout.write(
        " /::\\:\\__\\/\\/::\\__\\/\\:\\:\\__\\/::\\/\\__\\/:/\\:\\__\\/::\\:\\__\\\n")
    sys.stdout.write(
        " \\:\\::/  /\\::/\\/__/\\:\\:\\/__/\\/\\::/  /\\:\\/:/  /\\/\\::/  /\n")
    sys.stdout.write(
        "  \\::/  /  \\:\\__\\   \\::/  /   /:/  /  \\::/  /    \\/__/ \n")
    sys.stdout.write(
        "   \\/__/    \\/__/    \\/__/    \\/__/    \\/__/           \n\n")
    sys.stdout.write(
        "Bishop. V " + str(pkg_resources.get_distribution("Bishop").version) + "\n")
    sys.stdout.write("http://github.com/julianje/Bishop\n")
    sys.stdout.write("Julian Jara-Ettinger. jjara@mit.edu\n\n")
    sys.stdout.write(
        "Results using Bishop 1.0.0: Jara-Ettinger, J., Schulz, L. E., & Tenenbaum J. B. (2015). The naive utility calculus: Joint inferences about the costs and rewards of actions. In Proceedings of the 37th Annual Conference of the Cognitive Science Society.\n\n")
<|endoftext|>


<commit_before>
from pykml import parser
from pykml.parser import Schema
from pykml.factory import KML_ElementMaker as KML
from os import path
import json
import geojson
from lxml import etree
import sys, os, random
from datetime import datetime
from random import randrange

sampleTrip = '{"date": "20140413", "lastUpdate": "20140414T064442Z", "segments": [], "summary": null}'
sampleSegment = '{"place": {"type": "unknown", "id": 54095848, "location": {"lat": 37.3910149202, "lon": -122.0865010796}}, "endTime": "20140413T234434-0700", "type": "place", "startTime": "20140412T190446-0700", "lastUpdate": "20140414T064442Z"}'

def kml_multiple_to_geojson(infile_path, outdir_path, geojson_properties={}):
        """
        Converts a KML file with multiple Documents into geojson to store in an output file
        """
        data = __read_file(infile_path)
        coord_dict = __get_all_coords(data)
        if not os.path.exists(outdir_path):
                os.makedirs(outdir_path)
        for section_id, coords in coord_dict.items():
                filename = "%s.json" % section_id
                path = os.path.join(outdir_path, filename)
                outfile = file(path,'w')
                dump = __to_geojson(coords)
                outfile.write(dump)
                outfile.close()

def kml_to_geojson(infile_path, outfile_path, geojson_properties={}):
	"""
	Converts a KML file to geojson to store in an output file.
	"""
        data = __read_file(infile_path)
      	coords = __get_coords(data)
	outfile = open(outfile_path, 'w')
	outfile.write(__to_geojson(coords))
	outfile.close()

def __to_geojson(coordinates, geojson_properties={}):
        coordinatePoints = geojson.MultiPoint(coordinates)
	dump = geojson.dumps(geojson.Feature(geometry=coordinatePoints, properties=geojson_properties)) # can add properties and id to feature (perhaps trip/section id?)
        return dump

def kml_to_json(infile_path, outfile_path):
	"""
	Converts a KML file to geojson to store in an output file.
	"""
        data = __read_file(infile_path)
	coordinatePoints = __get_coords(data)
	segmentsList = []
	for coord in coordinatePoints:
		j = json.loads(sampleSegment)
		j["place"]["location"]["lat"] = coord[0] # 0th index might be longitude
		j["place"]["location"]["lon"] = coord[1]
		segmentsList.append(j)
	j = json.loads(sampleTrip)
	j["segments"] = segmentsList
	with open(outfile_path, 'w') as outfile:
		json.dump(j, outfile)

def __read_file(infile_path):
        data = ""
	with open(infile_path, "r") as temp:
		data = temp.read()
	assert(data != "")
        return data

def __get_coords(data):
	doc = parser.fromstring(data)
	dataPoints = doc.Document.Placemark.LineString.coordinates.text.split(" ")
	return [[float(i) for i in x.split(",")[:2]] for x in dataPoints]

def __get_all_coords(data):
        doc = parser.fromstring(data)
        coords = {}
        for folder in doc.Document.getchildren():
               	dataPoints = folder.Placemark.LineString.coordinates.text.split(" ")
                section_id = folder.Placemark.name.text
                coords[section_id] = [[float(i) for i in x.split(",")[:2]] for x in dataPoints]
        return coords

def validate_kml(filepath, schema = "https://developers.google.com/kml/schema/kml21.xsd"):
        """
        Validates kml located at filepath.
        TODO: Even imported kml files do not validate against default xml
        """
        schema_gomaps = Schema(schema)
	data = ""
	with open(filepath, "r") as temp:
		data = temp.read()
	assert(data != "")
	doc = parser.fromstring(data)
        return schema_ogc.assertValid(doc)

def section_to_kml(section, color, outfile_path="", write=True):
        """
        Converts a section into a kml file
        """
        line_style_id = "line-%s-5" % color
        red = "FF1212"
        green = "00B80C"
        start_icon_style_id = "icon-%s" % green
        end_icon_style_id = "icon-%s" % red
        make_coord = lambda p: (",".join(map(lambda x: str(x),
                         p["track_location"]["coordinates"]) + ["0.0"]))
        make_coord_point = lambda p: (",".join(map(lambda x: str(x),
                         p["coordinates"]) + ["0.0"]))
	style_id = "style-%s" % section['section_start_time']
        pm = KML.Placemark(
                KML.styleUrl("#%s" % line_style_id),
                KML.name(section['_id']),
                KML.LineString(
                        KML.tessellate(1),
                        KML.coordinates(" ".join(
                                map(lambda track_point: make_coord(track_point)
                                    ,section['track_points'][1:-1])))
                )
        )
        start_track_point = section['track_points'][0]
        end_track_point = section['track_points'][-1]
        start_time = mongodate_to_datetime(start_track_point['time'])
        end_time = mongodate_to_datetime(end_track_point['time'])
        start_point = KML.Placemark(
                KML.styleUrl("#%s" % start_icon_style_id),
                KML.name("Start: %s" % start_time),
                KML.description("Starting point"),
                KML.Point(KML.coordinates(make_coord(start_track_point)))
        )
        end_point = KML.Placemark(
                KML.styleUrl("#%s" % end_icon_style_id),
                KML.name("End: %s" % end_time),
                KML.description("Ending point"),
                KML.Point(KML.coordinates(make_coord(end_track_point)))
        )
        line_style = KML.Style(
                KML.LineStyle(
                        KML.color("ff%s" % color),
                        KML.width("5")
                )
        )
        line_style.set("id", line_style_id)
        start_icon_style = KML.Style(
                KML.IconStyle(
                        KML.color("ff%s" % green),
                        KML.scale("1.1")
                        # KML.Icon(
                        #         KML.href("http://www.gstatic.com/mapspro/images/stock/503-wht-blank_maps.png")
                        # )
                )
        )
        start_icon_style.set("id", start_icon_style_id)
        end_icon_style = KML.Style(
                KML.IconStyle(
                        KML.color("ff%s" % red),
                        KML.scale("1.1")
                        # KML.Icon(
                        #         KML.href("http://www.gstatic.com/mapspro/images/stock/503-wht-blank_maps.png")
                        # )
                )
        )
        end_icon_style.set("id", end_icon_style_id)
        fld = KML.Folder(
                KML.name(section['_id']),
                KML.description("From %s \nto %s" % (start_time, end_time)),
                pm,
                start_point,
                end_point
        )
        if write:
                kml = KML.kml(KML.Document(fld, section["user_id"]))
                path = os.path.join(outfile_path, str(section['user_id']) +'.kml')
                outfile = file(path,'w')
                outfile.write(etree.tostring(kml, pretty_print=True))
        else:
                return fld, line_style, start_icon_style, end_icon_style

def sections_to_kml(filename, sections, outfile_path=""):
        r = lambda: random.randint(0,255)
        l_tuples = map(lambda section: section_to_kml(section, '%02X%02X%02X' % (r(),r(),r()),  write=False) if len(section['track_points']) > 1 else None ,sections)
        flat = []

        for f,s,i_s,i_e in l_tuples:
                flat.append(f)
                flat.append(s)
                flat.append(i_s)
                flat.append(i_e)
        kml = KML.kml(KML.Document(sections[0]["user_id"],*flat))
        path = os.path.join(outfile_path, filename +'.kml')
        outfile = file(path,'w')
        outfile.write(etree.tostring(kml, pretty_print=True))

def chunks(l,n):
        """
        Generates evenly sized chunks of a list
        """
        for i in xrange(0, len(l), n):
                yield l[i:i+n]

def mongodate_to_datetime(date):
        return datetime.strptime(date[:-5], "%Y%m%dT%H%M%S")

if __name__ == "__main__":
        sys.path.append("%s/../" % os.getcwd())
        from get_database import get_section_db
        from uuid import UUID
        Sections = get_section_db()
        user_id = "__ID_GOES_HERE__"
        sections = Sections.find({'$and': [{'user_id':user_id}]})
        sections_to_kml(user_id,sections)
<commit_msg>Revert "removing url reference that may be interfereing with importing trips"<commit_after>
from pykml import parser
from pykml.parser import Schema
from pykml.factory import KML_ElementMaker as KML
from os import path
import json
import geojson
from lxml import etree
import sys, os, random
from datetime import datetime
from random import randrange

sampleTrip = '{"date": "20140413", "lastUpdate": "20140414T064442Z", "segments": [], "summary": null}'
sampleSegment = '{"place": {"type": "unknown", "id": 54095848, "location": {"lat": 37.3910149202, "lon": -122.0865010796}}, "endTime": "20140413T234434-0700", "type": "place", "startTime": "20140412T190446-0700", "lastUpdate": "20140414T064442Z"}'

def kml_multiple_to_geojson(infile_path, outdir_path, geojson_properties={}):
        """
        Converts a KML file with multiple Documents into geojson to store in an output file
        """
        data = __read_file(infile_path)
        coord_dict = __get_all_coords(data)
        if not os.path.exists(outdir_path):
                os.makedirs(outdir_path)
        for section_id, coords in coord_dict.items():
                filename = "%s.json" % section_id
                path = os.path.join(outdir_path, filename)
                outfile = file(path,'w')
                dump = __to_geojson(coords)
                outfile.write(dump)
                outfile.close()

def kml_to_geojson(infile_path, outfile_path, geojson_properties={}):
	"""
	Converts a KML file to geojson to store in an output file.
	"""
        data = __read_file(infile_path)
      	coords = __get_coords(data)
	outfile = open(outfile_path, 'w')
	outfile.write(__to_geojson(coords))
	outfile.close()

def __to_geojson(coordinates, geojson_properties={}):
        coordinatePoints = geojson.MultiPoint(coordinates)
	dump = geojson.dumps(geojson.Feature(geometry=coordinatePoints, properties=geojson_properties)) # can add properties and id to feature (perhaps trip/section id?)
        return dump

def kml_to_json(infile_path, outfile_path):
	"""
	Converts a KML file to geojson to store in an output file.
	"""
        data = __read_file(infile_path)
	coordinatePoints = __get_coords(data)
	segmentsList = []
	for coord in coordinatePoints:
		j = json.loads(sampleSegment)
		j["place"]["location"]["lat"] = coord[0] # 0th index might be longitude
		j["place"]["location"]["lon"] = coord[1]
		segmentsList.append(j)
	j = json.loads(sampleTrip)
	j["segments"] = segmentsList
	with open(outfile_path, 'w') as outfile:
		json.dump(j, outfile)

def __read_file(infile_path):
        data = ""
	with open(infile_path, "r") as temp:
		data = temp.read()
	assert(data != "")
        return data

def __get_coords(data):
	doc = parser.fromstring(data)
	dataPoints = doc.Document.Placemark.LineString.coordinates.text.split(" ")
	return [[float(i) for i in x.split(",")[:2]] for x in dataPoints]

def __get_all_coords(data):
        doc = parser.fromstring(data)
        coords = {}
        for folder in doc.Document.getchildren():
               	dataPoints = folder.Placemark.LineString.coordinates.text.split(" ")
                section_id = folder.Placemark.name.text
                coords[section_id] = [[float(i) for i in x.split(",")[:2]] for x in dataPoints]
        return coords

def validate_kml(filepath, schema = "https://developers.google.com/kml/schema/kml21.xsd"):
        """
        Validates kml located at filepath.
        TODO: Even imported kml files do not validate against default xml
        """
        schema_gomaps = Schema(schema)
	data = ""
	with open(filepath, "r") as temp:
		data = temp.read()
	assert(data != "")
	doc = parser.fromstring(data)
        return schema_ogc.assertValid(doc)

def section_to_kml(section, color, outfile_path="", write=True):
        """
        Converts a section into a kml file
        """
        line_style_id = "line-%s-5" % color
        red = "FF1212"
        green = "00B80C"
        start_icon_style_id = "icon-%s" % green
        end_icon_style_id = "icon-%s" % red
        make_coord = lambda p: (",".join(map(lambda x: str(x),
                         p["track_location"]["coordinates"]) + ["0.0"]))
        make_coord_point = lambda p: (",".join(map(lambda x: str(x),
                         p["coordinates"]) + ["0.0"]))
	style_id = "style-%s" % section['section_start_time']
        pm = KML.Placemark(
                KML.styleUrl("#%s" % line_style_id),
                KML.name(section['_id']),
                KML.LineString(
                        KML.tessellate(1),
                        KML.coordinates(" ".join(
                                map(lambda track_point: make_coord(track_point)
                                    ,section['track_points'][1:-1])))
                )
        )
        start_track_point = section['track_points'][0]
        end_track_point = section['track_points'][-1]
        start_time = mongodate_to_datetime(start_track_point['time'])
        end_time = mongodate_to_datetime(end_track_point['time'])
        start_point = KML.Placemark(
                KML.styleUrl("#%s" % start_icon_style_id),
                KML.name("Start: %s" % start_time),
                KML.description("Starting point"),
                KML.Point(KML.coordinates(make_coord(start_track_point)))
        )
        end_point = KML.Placemark(
                KML.styleUrl("#%s" % end_icon_style_id),
                KML.name("End: %s" % end_time),
                KML.description("Ending point"),
                KML.Point(KML.coordinates(make_coord(end_track_point)))
        )
        style = KML.Style(
                KML.LineStyle(
                        KML.color("ff%s" % color),
                        KML.width("5")
                )
        )
        line_style.set("id", line_style_id)
        start_icon_style = KML.Style(
                KML.IconStyle(
                        KML.color("ff%s" % green),
                        KML.scale("1.1"),
                        KML.Icon(
                                KML.href("http://www.gstatic.com/mapspro/images/stock/503-wht-blank_maps.png")
                        )
                )
        )
        start_icon_style.set("id", start_icon_style_id)
        end_icon_style = KML.Style(
                KML.IconStyle(
                        KML.color("ff%s" % red),
                        KML.scale("1.1"),
                        KML.Icon(
                                KML.href("http://www.gstatic.com/mapspro/images/stock/503-wht-blank_maps.png")
                        )
                )
        )
        end_icon_style.set("id", end_icon_style_id)
        fld = KML.Folder(
                KML.name(section['_id']),
                KML.description("From %s \nto %s" % (start_time, end_time)),
                pm,
                start_point,
                end_point
        )
        if write:
                kml = KML.kml(KML.Document(fld, section["user_id"]))
                path = os.path.join(outfile_path, str(section['user_id']) +'.kml')
                outfile = file(path,'w')
                outfile.write(etree.tostring(kml, pretty_print=True))
        else:
                return fld, line_style, start_icon_style, end_icon_style

def sections_to_kml(filename, sections, outfile_path=""):
        r = lambda: random.randint(0,255)
        l_tuples = map(lambda section: section_to_kml(section, '%02X%02X%02X' % (r(),r(),r()),  write=False) if len(section['track_points']) > 1 else None ,sections)
        flat = []

        for f,s,i_s,i_e in l_tuples:
                flat.append(f)
                flat.append(s)
                flat.append(i_s)
                flat.append(i_e)
        kml = KML.kml(KML.Document(sections[0]["user_id"],*flat))
        path = os.path.join(outfile_path, filename +'.kml')
        outfile = file(path,'w')
        outfile.write(etree.tostring(kml, pretty_print=True))

def chunks(l,n):
        """
        Generates evenly sized chunks of a list
        """
        for i in xrange(0, len(l), n):
                yield l[i:i+n]

def mongodate_to_datetime(date):
        return datetime.strptime(date[:-5], "%Y%m%dT%H%M%S")

if __name__ == "__main__":
        sys.path.append("%s/../" % os.getcwd())
        from get_database import get_section_db
        from uuid import UUID
        Sections = get_section_db()
        user_id = "__ID_GOES_HERE__"
        sections = Sections.find({'$and': [{'user_id':user_id}]})
        sections_to_kml(user_id,sections)
<|endoftext|>