<commit_before>import time

from django.conf.urls.defaults import *
from django.http import HttpResponse
from django.utils.translation import ugettext as _

from ella import newman
from ella.core.models.main import Category

from ella.newman import models as m
from ella.newman.filterspecs import filter_spec
from ella.newman.permission import is_category_fk, is_site_fk, applicable_categories
from ella.newman.utils import user_category_filter, get_user_config
from ella.newman.config import CATEGORY_FILTER

class DevMessageAdmin(newman.NewmanModelAdmin):
    list_display = ('title', 'author', 'version', 'ts',)
    search_fields = ('title', 'summary', 'details',)
    list_filter = ('author', 'ts',)
    prepopulated_fields = {'slug': ('title',)}

    def save_model(self, request, obj, form, change):
        if not obj.id:
            obj.author = request.user
        obj.save()

class HelpItemAdmin(newman.NewmanModelAdmin):
    list_display = ('__unicode__',)
    list_filter = ('ct', 'lang',)
    rich_text_fields = {'': ('long',)}
    list_select_related = False

class CategoryUserRoleAdmin(newman.NewmanModelAdmin):
    list_filter = ('user', 'group',)
    list_display = ('user', 'group',)
    search_fields = ('user__username', 'category__title', 'category__slug')
    suggest_fields = {'category': ('__unicode__', 'title', 'slug',)}

    def get_urls(self):
        urls = patterns('',
            url(r'^refresh/$',
                self.refresh_view,
                name='categoryuserrole-refresh'),
        )
        urls += super(CategoryUserRoleAdmin, self).get_urls()
        return urls

    def refresh_view(self, request, extra_context=None):
        from ella.newman.management.commands.syncroles import denormalize
        # TODO: don't wait for denormalize()
        denormalize()
        return HttpResponse(_('All roles is now refreshed.'))

class CategoryUserRoleInline(newman.NewmanTabularInline):
    model = m.CategoryUserRole
    max_num = 3
    suggest_fields = {'category': ('__unicode__', 'title', 'tree_path') }

newman.site.register(m.DevMessage, DevMessageAdmin)
newman.site.register(m.AdminHelpItem, HelpItemAdmin)
newman.site.register(m.CategoryUserRole, CategoryUserRoleAdmin)

# Category filter -- restricted categories accordingly to CategoryUserRoles and categories filtered via AdminSettings.
# custom registered DateField filter. Filter is inserted to the beginning of filter chain.
category_lookup = lambda fspec: '%s__%s__exact' % (fspec.f.name, fspec.f.rel.get_related_field().name)

@filter_spec(lambda field: is_category_fk(field), category_lookup)
def category_field_filter(fspec):
    qs = Category.objects.filter(pk__in=applicable_categories(fspec.user))
    for cat in user_category_filter(qs, fspec.user):
        lookup_var = '%s__%s__exact' % (fspec.f.name, fspec.f.rel.to._meta.pk.name)
        link = ( cat, {lookup_var: cat.pk})
        fspec.links.append(link)
    return True

site_lookup = lambda fspec: '%s__%s__exact' % (fspec.f.name, fspec.f.rel.get_related_field().name)
@filter_spec(lambda field: is_site_fk(field), site_lookup)
def site_field_filter(fspec):
    category_ids = get_user_config(fspec.user, CATEGORY_FILTER)
    if not category_ids:
        if not fspec.user.is_superuser:
            category_ids = m.DenormalizedCategoryUserRole.objects.root_categories_by_user(fspec.user)
        else:
            category_ids = Category.objects.filter(tree_parent=None)
    qs = Category.objects.filter(pk__in=category_ids)
    sites = map(lambda c: c.site, qs)
    for site in sites:
        #category__site__id__exact=1
        lookup_var = '%s__%s__exact' % (fspec.f.name, fspec.f.rel.get_related_field().name)
        link = ( site, {lookup_var: site.pk})
        fspec.links.append(link)
    return True

# TODO: register some non-ella apps, fix it

class SiteAdmin(newman.NewmanModelAdmin):
    list_display = ('domain', 'name',)

from django.contrib.sites.models import Site
newman.site.register(Site, SiteAdmin)

from django.contrib.auth.admin import UserAdmin, GroupAdmin
class UserAdmin(UserAdmin, newman.NewmanModelAdmin):
    inlines = [CategoryUserRoleInline]
    suggest_fields = {'groups': ('name',)}

class GroupAdmin(GroupAdmin, newman.NewmanModelAdmin):
    pass

from django.contrib.auth.models import User, Group
newman.site.register(User, UserAdmin)
newman.site.register(Group, GroupAdmin)
<commit_msg>Added 'slug' into CategoryUserRoleInline suggest_field.<commit_after>import time

from django.conf.urls.defaults import *
from django.http import HttpResponse
from django.utils.translation import ugettext as _

from ella import newman
from ella.core.models.main import Category

from ella.newman import models as m
from ella.newman.filterspecs import filter_spec
from ella.newman.permission import is_category_fk, is_site_fk, applicable_categories
from ella.newman.utils import user_category_filter, get_user_config
from ella.newman.config import CATEGORY_FILTER

class DevMessageAdmin(newman.NewmanModelAdmin):
    list_display = ('title', 'author', 'version', 'ts',)
    search_fields = ('title', 'summary', 'details',)
    list_filter = ('author', 'ts',)
    prepopulated_fields = {'slug': ('title',)}

    def save_model(self, request, obj, form, change):
        if not obj.id:
            obj.author = request.user
        obj.save()

class HelpItemAdmin(newman.NewmanModelAdmin):
    list_display = ('__unicode__',)
    list_filter = ('ct', 'lang',)
    rich_text_fields = {'': ('long',)}
    list_select_related = False

class CategoryUserRoleAdmin(newman.NewmanModelAdmin):
    list_filter = ('user', 'group',)
    list_display = ('user', 'group',)
    search_fields = ('user__username', 'category__title', 'category__slug')
    suggest_fields = {'category': ('__unicode__', 'title', 'slug',)}

    def get_urls(self):
        urls = patterns('',
            url(r'^refresh/$',
                self.refresh_view,
                name='categoryuserrole-refresh'),
        )
        urls += super(CategoryUserRoleAdmin, self).get_urls()
        return urls

    def refresh_view(self, request, extra_context=None):
        from ella.newman.management.commands.syncroles import denormalize
        # TODO: don't wait for denormalize()
        denormalize()
        return HttpResponse(_('All roles is now refreshed.'))

class CategoryUserRoleInline(newman.NewmanTabularInline):
    model = m.CategoryUserRole
    max_num = 3
    suggest_fields = {'category': ('__unicode__', 'title', 'tree_path', 'slug', ) }

newman.site.register(m.DevMessage, DevMessageAdmin)
newman.site.register(m.AdminHelpItem, HelpItemAdmin)
newman.site.register(m.CategoryUserRole, CategoryUserRoleAdmin)

# Category filter -- restricted categories accordingly to CategoryUserRoles and categories filtered via AdminSettings.
# custom registered DateField filter. Filter is inserted to the beginning of filter chain.
category_lookup = lambda fspec: '%s__%s__exact' % (fspec.f.name, fspec.f.rel.get_related_field().name)

@filter_spec(lambda field: is_category_fk(field), category_lookup)
def category_field_filter(fspec):
    qs = Category.objects.filter(pk__in=applicable_categories(fspec.user))
    for cat in user_category_filter(qs, fspec.user):
        lookup_var = '%s__%s__exact' % (fspec.f.name, fspec.f.rel.to._meta.pk.name)
        link = ( cat, {lookup_var: cat.pk})
        fspec.links.append(link)
    return True

site_lookup = lambda fspec: '%s__%s__exact' % (fspec.f.name, fspec.f.rel.get_related_field().name)
@filter_spec(lambda field: is_site_fk(field), site_lookup)
def site_field_filter(fspec):
    category_ids = get_user_config(fspec.user, CATEGORY_FILTER)
    if not category_ids:
        if not fspec.user.is_superuser:
            category_ids = m.DenormalizedCategoryUserRole.objects.root_categories_by_user(fspec.user)
        else:
            category_ids = Category.objects.filter(tree_parent=None)
    qs = Category.objects.filter(pk__in=category_ids)
    sites = map(lambda c: c.site, qs)
    for site in sites:
        #category__site__id__exact=1
        lookup_var = '%s__%s__exact' % (fspec.f.name, fspec.f.rel.get_related_field().name)
        link = ( site, {lookup_var: site.pk})
        fspec.links.append(link)
    return True

# TODO: register some non-ella apps, fix it

class SiteAdmin(newman.NewmanModelAdmin):
    list_display = ('domain', 'name',)

from django.contrib.sites.models import Site
newman.site.register(Site, SiteAdmin)

from django.contrib.auth.admin import UserAdmin, GroupAdmin
class UserAdmin(UserAdmin, newman.NewmanModelAdmin):
    inlines = [CategoryUserRoleInline]
    suggest_fields = {'groups': ('name',)}

class GroupAdmin(GroupAdmin, newman.NewmanModelAdmin):
    pass

from django.contrib.auth.models import User, Group
newman.site.register(User, UserAdmin)
newman.site.register(Group, GroupAdmin)
<|endoftext|>


<commit_before>from urllib import urlencode
from logging import getLogger

from django.contrib import admin
from django.core.urlresolvers import reverse
from django.http import HttpResponseRedirect
from django.shortcuts import redirect

from ..settings import SOCIAL_INSTAGRAM_CLIENT_ID, SOCIAL_INSTAGRAM_REDIRECT_URI
from ..models import TwitterSetting, FacebookSetting, InstagramSetting, RSSSetting, \
                     APPROVED, PENDING, FAVORITED, REJECTED, LEGAL
from ..views import begin_auth

log = getLogger(__name__)

def approve_message(modeladmin, request, queryset):
    for item in queryset:
        item.status = APPROVED
        item.save()
approve_message.short_description = "Mark As Approved"

def rejected_message(modeladmin, request, queryset):
    for item in queryset:
        item.status = REJECTED
        item.save()
rejected_message.short_description = "Mark As Rejected"

def favorite_message(modeladmin, request, queryset):
    for item in queryset:
        item.status = FAVORITED
        item.save()
favorite_message.short_description = "Mark As Favorite"

def pending_message(modeladmin, request, queryset):
    for item in queryset:
        item.status = PENDING
        item.save()
pending_message.short_description = "Mark As Pending"

def legal_message(modeladmin, request, queryset):
    for item in queryset:
        item.status = LEGAL
        item.save()
legal_message.short_description = "Mark As Legal"

def admin_url(model, url, object_id=None):
    """
    Returns the URL for the given model and admin url name.
    """
    opts = model._meta
    url = "admin:%s_%s_%s" % (opts.app_label, opts.object_name.lower(), url)
    args = ()
    if object_id is not None:
        args = (object_id,)
    return reverse(url, args=args)

class SingletonAdmin(admin.ModelAdmin):
    """
    Admin class for models that should only contain a single instance
    in the database. Redirect all views to the change view when the
    instance exists, and to the add view when it doesn't.
    """

    def handle_save(self, request, response):
        """
        Handles redirect back to the dashboard when save is clicked
        (eg not save and continue editing), by checking for a redirect
        response, which only occurs if the form is valid.
        """
        form_valid = isinstance(response, HttpResponseRedirect)
        if request.POST.get("_save") and form_valid:
            return redirect("admin:index")
        return response

    def add_view(self, *args, **kwargs):
        """
        Redirect to the change view if the singleton instance exists.
        """
        try:
            singleton = self.model.objects.get()
        except (self.model.DoesNotExist, self.model.MultipleObjectsReturned):
            kwargs.setdefault("extra_context", {})
            kwargs["extra_context"]["singleton"] = True
            response = super(SingletonAdmin, self).add_view(*args, **kwargs)
            return self.handle_save(args[0], response)
        return redirect(admin_url(self.model, "change", singleton.id))

    def changelist_view(self, *args, **kwargs):
        """
        Redirect to the add view if no records exist or the change
        view if the singleton instance exists.
        """
        try:
            singleton = self.model.objects.get()
        except self.model.MultipleObjectsReturned:
            return super(SingletonAdmin, self).changelist_view(*args, **kwargs)
        except self.model.DoesNotExist:
            return redirect(admin_url(self.model, "add"))
        return redirect(admin_url(self.model, "change", singleton.id))

    def change_view(self, *args, **kwargs):
        """
        If only the singleton instance exists, pass ``True`` for
        ``singleton`` into the template which will use CSS to hide
        the "save and add another" button.
        """
        kwargs.setdefault("extra_context", {})
        kwargs["extra_context"]["singleton"] = self.model.objects.count() == 1
        response = super(SingletonAdmin, self).change_view(*args, **kwargs)
        return self.handle_save(args[0], response)

    class Media:
        js = [
              settings.STATIC_URL+'social-singleton.js',
        ]

class HideableAdmin(admin.ModelAdmin):
    def get_model_perms(self, request):
        if self.pref_model.objects.count():
            return super(HideableAdmin, self).get_model_perms(request)
        else:
            return {
                'add': False,
                'change': False,
                'delete': False,
            }

class MessageAdmin(admin.ModelAdmin):
    actions = [approve_message, rejected_message, favorite_message, pending_message, legal_message]
    list_display = ('id', 'admin_media_preview', 'message', 'status', 'media_type', 'network', 'date', 'tags', )
    list_filter = ('network', 'media_type', 'status', 'date', '_tags')
    list_display_links = ('id', 'admin_media_preview' )
    readonly_fields = ('admin_media_preview', 'reply_to',)
    exclude = ('tags', )
    ordering = ('-status', '-date')

    class Media:
        js = ("jquery.multiselect.js", "tag_multiselect.js", "json2.js", )
        css = { 'all': ("jquery.multiselect.css", ) }

class FacebookAccountAdmin(HideableAdmin):
    pref_model = FacebookSetting

class FacebookMessageAdmin(MessageAdmin, HideableAdmin):
    pref_model = FacebookSetting
    list_display = [c for c in MessageAdmin.list_display if c != 'tags']
    list_filter = [c for c in MessageAdmin.list_filter if c != '_tags']

class FacebookSearchAdmin(HideableAdmin):
    pref_model = FacebookSetting

class TwitterAccountAdmin(HideableAdmin):
    list_display = ('screen_name',)
    pref_model = TwitterSetting

    def add_view(self, request, form_url='', extra_context=None):
        log.debug("request: %s", request)
        log.debug("form_url: %s", form_url)
        log.debug("extra_context: %s", extra_context)
        return begin_auth(request)

class TwitterPublicAccountAdmin(HideableAdmin):
    pref_model = TwitterSetting

class TwitterMessageAdmin(MessageAdmin, HideableAdmin):
    pref_model = TwitterSetting
    list_display = ('id', 'admin_media_preview', 'message', 'status', 'media_type', 'date', )
    list_filter = ('twitter_search__search_term', 'twitter_account__screen_name', 'status', 'media_type', 'date')

class TwitterSearchAdmin(HideableAdmin):
    pref_model = TwitterSetting

class InstagramAccountAdmin(HideableAdmin):
    pref_model = InstagramSetting
    list_display = ('id', 'admin_image', 'username', 'name')
    list_filter = ('username', 'name')
    search_fields = ('username', 'name')

    def add_view(self, *args, **kwargs):
        if SOCIAL_INSTAGRAM_CLIENT_ID and SOCIAL_INSTAGRAM_REDIRECT_URI:
            params = urlencode({
                'client_id': SOCIAL_INSTAGRAM_CLIENT_ID,
                'redirect_uri': SOCIAL_INSTAGRAM_REDIRECT_URI,
                'response_type': 'code',
            })
            url = '?'.join(['https://api.instagram.com/oauth/authorize/', params])
            log.info('Login via Instagram API')
            log.info(url)
            return HttpResponseRedirect(url)
        else:
            return HttpResponseRedirect(reverse('admin:social_instagramsetting_add'))

class InstagramPublicAccountAdmin(HideableAdmin):
    pref_model = InstagramSetting

class InstagramSearchAdmin(HideableAdmin):
    pref_model = InstagramSetting

class InstagramMessageAdmin(MessageAdmin, HideableAdmin):
    pref_model = InstagramSetting
    list_display = ('id', 'admin_image_low','message', 'status', 'media_type', 'date', )
    list_display_links = ('id', 'admin_image_low' )
    list_filter = ('status', 'media_type', 'date',)

class RSSAccountAdmin(HideableAdmin):
    pref_model = RSSSetting

class RSSMessageAdmin(HideableAdmin):
    pref_model = RSSSetting
    list_display = ('id', 'message', 'date')
    list_filter = ('rss_account__feed_name', 'status', 'media_type', 'date',)
<commit_msg>add settings import<commit_after>from urllib import urlencode
from logging import getLogger

from django.conf import settings

from django.contrib import admin
from django.core.urlresolvers import reverse
from django.http import HttpResponseRedirect
from django.shortcuts import redirect

from ..settings import SOCIAL_INSTAGRAM_CLIENT_ID, SOCIAL_INSTAGRAM_REDIRECT_URI
from ..models import TwitterSetting, FacebookSetting, InstagramSetting, RSSSetting, \
                     APPROVED, PENDING, FAVORITED, REJECTED, LEGAL
from ..views import begin_auth

log = getLogger(__name__)

def approve_message(modeladmin, request, queryset):
    for item in queryset:
        item.status = APPROVED
        item.save()
approve_message.short_description = "Mark As Approved"

def rejected_message(modeladmin, request, queryset):
    for item in queryset:
        item.status = REJECTED
        item.save()
rejected_message.short_description = "Mark As Rejected"

def favorite_message(modeladmin, request, queryset):
    for item in queryset:
        item.status = FAVORITED
        item.save()
favorite_message.short_description = "Mark As Favorite"

def pending_message(modeladmin, request, queryset):
    for item in queryset:
        item.status = PENDING
        item.save()
pending_message.short_description = "Mark As Pending"

def legal_message(modeladmin, request, queryset):
    for item in queryset:
        item.status = LEGAL
        item.save()
legal_message.short_description = "Mark As Legal"

def admin_url(model, url, object_id=None):
    """
    Returns the URL for the given model and admin url name.
    """
    opts = model._meta
    url = "admin:%s_%s_%s" % (opts.app_label, opts.object_name.lower(), url)
    args = ()
    if object_id is not None:
        args = (object_id,)
    return reverse(url, args=args)

class SingletonAdmin(admin.ModelAdmin):
    """
    Admin class for models that should only contain a single instance
    in the database. Redirect all views to the change view when the
    instance exists, and to the add view when it doesn't.
    """

    def handle_save(self, request, response):
        """
        Handles redirect back to the dashboard when save is clicked
        (eg not save and continue editing), by checking for a redirect
        response, which only occurs if the form is valid.
        """
        form_valid = isinstance(response, HttpResponseRedirect)
        if request.POST.get("_save") and form_valid:
            return redirect("admin:index")
        return response

    def add_view(self, *args, **kwargs):
        """
        Redirect to the change view if the singleton instance exists.
        """
        try:
            singleton = self.model.objects.get()
        except (self.model.DoesNotExist, self.model.MultipleObjectsReturned):
            kwargs.setdefault("extra_context", {})
            kwargs["extra_context"]["singleton"] = True
            response = super(SingletonAdmin, self).add_view(*args, **kwargs)
            return self.handle_save(args[0], response)
        return redirect(admin_url(self.model, "change", singleton.id))

    def changelist_view(self, *args, **kwargs):
        """
        Redirect to the add view if no records exist or the change
        view if the singleton instance exists.
        """
        try:
            singleton = self.model.objects.get()
        except self.model.MultipleObjectsReturned:
            return super(SingletonAdmin, self).changelist_view(*args, **kwargs)
        except self.model.DoesNotExist:
            return redirect(admin_url(self.model, "add"))
        return redirect(admin_url(self.model, "change", singleton.id))

    def change_view(self, *args, **kwargs):
        """
        If only the singleton instance exists, pass ``True`` for
        ``singleton`` into the template which will use CSS to hide
        the "save and add another" button.
        """
        kwargs.setdefault("extra_context", {})
        kwargs["extra_context"]["singleton"] = self.model.objects.count() == 1
        response = super(SingletonAdmin, self).change_view(*args, **kwargs)
        return self.handle_save(args[0], response)

    class Media:
        js = [
              settings.STATIC_URL+'social-singleton.js',
        ]

class HideableAdmin(admin.ModelAdmin):
    def get_model_perms(self, request):
        if self.pref_model.objects.count():
            return super(HideableAdmin, self).get_model_perms(request)
        else:
            return {
                'add': False,
                'change': False,
                'delete': False,
            }

class MessageAdmin(admin.ModelAdmin):
    actions = [approve_message, rejected_message, favorite_message, pending_message, legal_message]
    list_display = ('id', 'admin_media_preview', 'message', 'status', 'media_type', 'network', 'date', 'tags', )
    list_filter = ('network', 'media_type', 'status', 'date', '_tags')
    list_display_links = ('id', 'admin_media_preview' )
    readonly_fields = ('admin_media_preview', 'reply_to',)
    exclude = ('tags', )
    ordering = ('-status', '-date')

    class Media:
        js = ("jquery.multiselect.js", "tag_multiselect.js", "json2.js", )
        css = { 'all': ("jquery.multiselect.css", ) }

class FacebookAccountAdmin(HideableAdmin):
    pref_model = FacebookSetting

class FacebookMessageAdmin(MessageAdmin, HideableAdmin):
    pref_model = FacebookSetting
    list_display = [c for c in MessageAdmin.list_display if c != 'tags']
    list_filter = [c for c in MessageAdmin.list_filter if c != '_tags']

class FacebookSearchAdmin(HideableAdmin):
    pref_model = FacebookSetting

class TwitterAccountAdmin(HideableAdmin):
    list_display = ('screen_name',)
    pref_model = TwitterSetting

    def add_view(self, request, form_url='', extra_context=None):
        log.debug("request: %s", request)
        log.debug("form_url: %s", form_url)
        log.debug("extra_context: %s", extra_context)
        return begin_auth(request)

class TwitterPublicAccountAdmin(HideableAdmin):
    pref_model = TwitterSetting

class TwitterMessageAdmin(MessageAdmin, HideableAdmin):
    pref_model = TwitterSetting
    list_display = ('id', 'admin_media_preview', 'message', 'status', 'media_type', 'date', )
    list_filter = ('twitter_search__search_term', 'twitter_account__screen_name', 'status', 'media_type', 'date')

class TwitterSearchAdmin(HideableAdmin):
    pref_model = TwitterSetting

class InstagramAccountAdmin(HideableAdmin):
    pref_model = InstagramSetting
    list_display = ('id', 'admin_image', 'username', 'name')
    list_filter = ('username', 'name')
    search_fields = ('username', 'name')

    def add_view(self, *args, **kwargs):
        if SOCIAL_INSTAGRAM_CLIENT_ID and SOCIAL_INSTAGRAM_REDIRECT_URI:
            params = urlencode({
                'client_id': SOCIAL_INSTAGRAM_CLIENT_ID,
                'redirect_uri': SOCIAL_INSTAGRAM_REDIRECT_URI,
                'response_type': 'code',
            })
            url = '?'.join(['https://api.instagram.com/oauth/authorize/', params])
            log.info('Login via Instagram API')
            log.info(url)
            return HttpResponseRedirect(url)
        else:
            return HttpResponseRedirect(reverse('admin:social_instagramsetting_add'))

class InstagramPublicAccountAdmin(HideableAdmin):
    pref_model = InstagramSetting

class InstagramSearchAdmin(HideableAdmin):
    pref_model = InstagramSetting

class InstagramMessageAdmin(MessageAdmin, HideableAdmin):
    pref_model = InstagramSetting
    list_display = ('id', 'admin_image_low','message', 'status', 'media_type', 'date', )
    list_display_links = ('id', 'admin_image_low' )
    list_filter = ('status', 'media_type', 'date',)

class RSSAccountAdmin(HideableAdmin):
    pref_model = RSSSetting

class RSSMessageAdmin(HideableAdmin):
    pref_model = RSSSetting
    list_display = ('id', 'message', 'date')
    list_filter = ('rss_account__feed_name', 'status', 'media_type', 'date',)
<|endoftext|>


<commit_before>from . import constants as c
from .events import Event

class Session:
    """
    Represents a client session
    """
    def __init__(self, shark, client, info={}):
        """
        Initialize a session with
        - `shark`: a SocketShark instance,
        - `client`: a Websocket-backend-specific client object which implements
                    an async send() method that takes a JSON dict, and
        - `info`: a dict with any other information that should be logged (e.g.
                  the client's remote address).
        """
        self.auth_info = {}
        self.shark = shark
        self.config = shark.config
        self.client = client
        self.log = self.shark.log.bind(session=id(self))
        self.log.debug('new session', **info)
        self.subscriptions = {}  # dict of Subscription objects by name
        self.active = True
        shark.sessions.add(self)

    async def on_client_event(self, data):
        """
        Called by the WebSocket backend when a new client messages comes in.o
        Expects a JSON dict.
        """
        assert self.active
        self.log.debug('client event', data=data)
        event = Event.from_data(self, data)
        try:
            await event.full_process()
        except:
            self.shark.log.exception('unhandled event processing exception',
                                     exc_info=True)
            await event.send_error(c.ERR_UNHANDLED_EXCEPTION)
            await self.close()
            raise

    async def on_service_event(self, data):
        """
        Called by the ServiceReceiver with a JSON dict on messages published by
        a service.
        """
        # Don't attempt to process messages to closed connections.
        if not self.active:
            return

        self.log.debug('service event', data=data)

        # Filter by comparing filter_fields to auth_info
        subscription_name = data['subscription']
        subscription = self.subscriptions.get(subscription_name)
        if not subscription:
            return

        filter_fields = subscription.service_config.get('filter_fields', [])
        for field in filter_fields:
            if field in data:
                if self.auth_info.get(field) != data[field]:
                    # Message filtered.
                    return

        msg = {
            'event': 'message',
            'subscription': subscription.name,
            'data': data['data'],
        }
        msg.update(subscription.extra_data)
        await self.send(msg)

    async def send(self, data):
        """
        Sends a JSON message to the client.
        """
        await self.client.send(data)

    async def close(self):
        if self.active:
            self.log.info('closing connection')
            self.active = False
            await self.client.close()
        else:
            self.log.info('connection already closing')

    async def on_close(self):
        """
        Called by the WebSocket backend to indicate the connection was closed.
        """
        self.active = False
        self.log.info('connection closed')
        await self.unsubscribe_all()
        self.shark.sessions.remove(self)

    async def unsubscribe_all(self):
        """
        Force-unsubscribe all subscriptions of the session.
        """
        while self.subscriptions:
            name, subscription = self.subscriptions.popitem()
            await subscription.force_unsubscribe()
<commit_msg>Don't re-raise WebSocket exception<commit_after>from . import constants as c
from .events import Event

class Session:
    """
    Represents a client session
    """
    def __init__(self, shark, client, info={}):
        """
        Initialize a session with
        - `shark`: a SocketShark instance,
        - `client`: a Websocket-backend-specific client object which implements
                    an async send() method that takes a JSON dict, and
        - `info`: a dict with any other information that should be logged (e.g.
                  the client's remote address).
        """
        self.auth_info = {}
        self.shark = shark
        self.config = shark.config
        self.client = client
        self.log = self.shark.log.bind(session=id(self))
        self.log.debug('new session', **info)
        self.subscriptions = {}  # dict of Subscription objects by name
        self.active = True
        shark.sessions.add(self)

    async def on_client_event(self, data):
        """
        Called by the WebSocket backend when a new client messages comes in.o
        Expects a JSON dict.
        """
        assert self.active
        self.log.debug('client event', data=data)
        event = Event.from_data(self, data)
        try:
            await event.full_process()
        except:
            self.shark.log.exception('unhandled event processing exception',
                                     exc_info=True)
            await event.send_error(c.ERR_UNHANDLED_EXCEPTION)
            await self.close()

    async def on_service_event(self, data):
        """
        Called by the ServiceReceiver with a JSON dict on messages published by
        a service.
        """
        # Don't attempt to process messages to closed connections.
        if not self.active:
            return

        self.log.debug('service event', data=data)

        # Filter by comparing filter_fields to auth_info
        subscription_name = data['subscription']
        subscription = self.subscriptions.get(subscription_name)
        if not subscription:
            return

        filter_fields = subscription.service_config.get('filter_fields', [])
        for field in filter_fields:
            if field in data:
                if self.auth_info.get(field) != data[field]:
                    # Message filtered.
                    return

        msg = {
            'event': 'message',
            'subscription': subscription.name,
            'data': data['data'],
        }
        msg.update(subscription.extra_data)
        await self.send(msg)

    async def send(self, data):
        """
        Sends a JSON message to the client.
        """
        await self.client.send(data)

    async def close(self):
        if self.active:
            self.log.info('closing connection')
            self.active = False
            await self.client.close()
        else:
            self.log.info('connection already closing')

    async def on_close(self):
        """
        Called by the WebSocket backend to indicate the connection was closed.
        """
        self.active = False
        self.log.info('connection closed')
        await self.unsubscribe_all()
        self.shark.sessions.remove(self)

    async def unsubscribe_all(self):
        """
        Force-unsubscribe all subscriptions of the session.
        """
        while self.subscriptions:
            name, subscription = self.subscriptions.popitem()
            await subscription.force_unsubscribe()
<|endoftext|>


<commit_before>import RPi.GPIO as GPIO
from time import sleep
from bmbb_fish import BmBB

my_fish = BmBB()

while True:
    my_fish.tail_up()
    sleep(.5)
    my_fish.tail_back()
    sleep(.5)
    my_fish.speak("hello, Dave")
    sleep(.5)
    my_fish.head_up()
    sleep(.5)
    my_fish.head_back()
    sleep(.5)
    my_fish.mouth_open()
    sleep(.5)
    my_fish.mouth_close()
    sleep(.5)

<commit_msg>too out while true, added kill fish<commit_after>import RPi.GPIO as GPIO
from time import sleep
from bmbb_fish import BmBB

my_fish = BmBB()

my_fish.tail_up()
sleep(.5)
my_fish.tail_back()
sleep(.5)
my_fish.speak("hello, Dave")
sleep(.5)
my_fish.head_up()
sleep(.5)
my_fish.head_back()
sleep(.5)
my_fish.mouth_open()
sleep(.5)
my_fish.mouth_close()
sleep(.5)
my_fish.kill_fish()
<|endoftext|>


<commit_before>#!/usr/bin/python

## TODO:
# Save rDAT and trial info
# Implement Error Handling
# Add support for correction trials
# Add support for probe trials

import os, sys, random, csv, logging, logging.handlers, json
import numpy as np
import datetime as dt
from pyoperant import hwio, utils

class GoodNite(Exception):
    """ exception for when the lights should be off """
    pass

def get_stimulus(trial_class,options):
    """ take trial class and experimental options and return a tuple containing the wav filename & additional info to play

    take in the trial class and the options dict
    returns a stimulus dictionary

    """
    # determine stimulus for next trial, based on class
    stim = dict()
    stim['class'] = trial_class

    stim['string'] = []
    input_file_list = []
    trial_isi = []
    motif_id = 0
    cum_matrix = np.cumsum(options['models'][trial_class], axis=1)
    for item in range(options['strlen_max']):
        motif_id = (cum_matrix[motif_id] < random.random()).sum()
        motif_name = options['stim_map'][str(motif_id)]
        isi = random.gauss(options['isi_mean'], options['isi_stdev'])
        if isi < 0.0: isi = 0.0
        input_file_list.append((options['stims'][motif_name], isi))
        stim['string'].append(motif_name)
    input_file_list[-1] = (input_file_list[-1][0], 0.0)
    stim['input_file_list'] = input_file_list

    (stim['filename'], stim['epochs']) = utils.concat_wav(stim['input_file_list'], os.path.join(options['stim_path'], 'temp_concat.wav'))

    stim['dur'] = stim['epochs'][-1][-1]

    return stim

def get_options(cmd_line):
    """ get all of the configuration options for the experiment """

    # set path variables
    CODE_PATH = os.path.dirname(os.path.realpath(__file__))
    USER_PATH = os.path.expanduser('~')
    BIRD_PATH = os.path.join(USER_PATH,'opdat', 'B' + cmd_line['subj'])
    STIM_PATH = os.path.join(BIRD_PATH,'stims')

    with open(os.path.join(BIRD_PATH, cmd_line['config_file'])) as config_f:
        options = json.load(config_f)
    options['box_id'] = cmd_line['box']
    options['subject_id'] = cmd_line['subj']
    options['code_path'] = CODE_PATH
    options['user_path'] = USER_PATH
    options['bird_path'] = BIRD_PATH
    options['stim_path'] = STIM_PATH
    options['script_fname'] = os.path.basename(__file__)

    # run through stim files
    for name, filename in options['stims'].items():
        filename_full = os.path.join(options['stim_path'], filename)
        options['stims'][name] = filename_full

    # run through the transitions for each stimulus and generate transition matrixes
    n = len(options['stims'])
    options['models'] = {}
    for stim_class, pair_set in options['diagnostic_transitions'].items():
        mat = np.zeros((n+1,n+1))
        mat[0][1:] = 1.0/n
        for pair in pair_set:
            mat[pair[0]][1:] = (1-options['alpha'])/(n-1)
            mat[pair[0]][pair[1]] = options['alpha']
        options['models'][stim_class] = mat

    # define log files, rDAT files
    filetime_fmt = '%Y%m%d%H%M%S'
    timestamp = dt.datetime.now().strftime(filetime_fmt)

    options['log_file'] = os.path.join(options['bird_path'], options['subject_id'] + '.log')
    options['data_csv'] = os.path.join(options['bird_path'], options['subject_id'] + '_' + options['script_fname'] + '_trials_' + timestamp + '.csv')
    options['config_snapshot'] = os.path.join(options['bird_path'], options['subject_id'] + '_' + options['script_fname'] + '_config_' + timestamp + '.json')
    options['summaryDAT'] = os.path.join(options['bird_path'], options['subject_id'] + '.summaryDAT')
    options['fields_to_save'] = ['number','type','trial_start','class','stim_start','stim_string','response','feed','timeout','cum_correct','cum_correct_thresh']

    # oreo
    with open(options['data_csv'], 'wb') as data_fh:
        trialWriter = csv.writer(data_fh)
        trialWriter.writerow(options['fields_to_save'])
    with open(options['config_snapshot'], 'wb') as config_snap:
        json.dump(options, config_snap, sort_keys=True, indent=4)

    return options

def run_trial(trial, options):
    pass

def save_trial(trial, options):
    # write trial results to CSV
    with open(options['data_csv'],'ab') as data_fh:
        trialWriter = csv.DictWriter(data_fh,fieldnames=options['fields_to_save'],extrasaction='ignore')
        trialWriter.writerow(trial)

def wait_for_peck(box, options):
    """ runs a loop, querying for pecks. returns peck time or "GoodNite" exception """
    no_peck = True
    box.write(box.dio['LED_center'],True)
    while no_peck:
        no_peck = not box.read(box.dio['IR_center'])

        if not utils.check_time(options['light_schedule']):
            raise GoodNite()
    return dt.datetime.now()

if __name__ == "__main__":

    options = get_options(utils.parse_commandline())

    if  options['debug']:
        log_level = logging.DEBUG
    else:
        log_level = logging.INFO

    logging.basicConfig(filename=options['log_file'],
                        level=log_level,
                        format='%(asctime)s:%(levelname)s:%(message)s')
    log = logging.getLogger()
    #email_handler = logging.handlers.SMTPHandler(mailhost='localhost',
    #                                             fromaddr='bird@vogel.ucsd.edu',
    #                                             toaddrs=[options['experimenter']['email'],],
    #                                             subject='error notice',
    #                                             )
    #email_handler.setlevel(logging.ERROR)
    #log.addHandler(email_handler)

    # initialize box
    box = hwio.OperantBox(options['box_id'])
    log.debug('box %i initialized' % options['box_id'])
    box.reset()

    trial = dict()
    trial['correct'] = False
    trial['cum_correct'] = 0
    trial['cum_correct_thresh'] = 1
    trial['number'] = 0

    do_correction = False

    summary = utils.init_summary()

    # start experiment
    do_experiment = True
    while do_experiment:

        try:
            # first, check if we should be running trials, otherwise lightsout
            if not utils.check_time(options['light_schedule']):
                raise GoodNite()

            # make sure lights are on at the beginning of each trial, prep for trial
            box.lights_on()

            trial['feed'] = False
            trial['timeout'] = False
            trial['number'] += 1

            if do_correction and options['correction_trials']:
                trial['type'] = 'correction'
                log.debug("correction trial: class is %s" % trial['class'])
            else:
                trial['type'] = 'normal'
                trial['class'] = random.choice(options['models'].keys())
                trial_stim = get_stimulus(trial['class'],options)
                log.debug("trial class is %s" % trial['class'])

            # wait for bird to peck
            log.debug('waiting for peck...')
            trial['trial_start'] = wait_for_peck(box, options)

            # record trial initiation
            summary['trials'] += 1
            summary['last_trial_time'] = trial['trial_start'].ctime()
            log.info("trial started at %s" % trial['trial_start'].ctime())

            # play temp stimulus
            trial['stim_start'] = dt.datetime.now()
            wave_proc = box.play_wav(trial_stim['filename'])
            box.write(box.dio['LED_center'],False)

            # wait for response
            wait_min = trial_stim['epochs'][options['strlen_min']-1][-1]
            hwio.wait(wait_min)

            # check for responses
            wait_max = trial_stim['epochs'][-1][-1] + options['response_win']
            check_peck = True
            box.write(box.dio['LED_left'],True)
            box.write(box.dio['LED_right'],True)
            while check_peck:
                elapsed_time = dt.datetime.now() - trial['stim_start']
                if elapsed_time > dt.timedelta(seconds=wait_max):
                    trial['response'] = 'none'
                    check_peck = False
                elif box.read(box.dio['IR_left']):
                    trial['response_time'] = dt.datetime.now()
                    wave_proc.terminate()
                    trial['response'] = 'L'
                    check_peck = False
                    summary['responses'] += 1
                elif box.read(box.dio['IR_right']):
                    trial['response_time'] = dt.datetime.now()
                    wave_proc.terminate()
                    trial['response'] = 'R'
                    check_peck = False
                    summary['responses'] += 1
            # TODO: note response in event file
            box.write(box.dio['LED_left'],False)
            box.write(box.dio['LED_right'],False)
            trial['response_timedelta'] = elapsed_time

            # calculate the number of motifs the bird heard
            num_mots = 0
            for epoch in trial_stim['epochs']:
                if elapsed_time > dt.timedelta(seconds=epoch[0]):
                    num_mots += 1
            # determine the string of motifs the bird heard
            trial['stim_string'] = ';'.join(trial_stim['string'][0:num_mots])
            trial['stim_epochs'] = trial_stim['epochs'][0:num_mots]

            # decide how to respond to the subject for normal trials
            if trial['type'] == 'normal':
                if trial['response'] == trial['class']:
                    # correct response
                    do_correction = False
                    trial['cum_correct'] += 1

                    if options['secondary_reinf']:
                        # give secondary reinforcer
                        trial['flash_epoch'] = box.flash(dur=0.5)

                    if (trial['cum_correct'] >= trial['cum_correct_thresh']):
                        # if cum currect reaches feed threshold, then feed
                        summary['feeds'] += 1
                        trial['feed'] = True

                        try:
                            trial['feed_epoch'] = box.feed(options['feed_dur'][trial['class']])

                        # but catch the feed errors
                        except hwio.ResponseDuringFeedError as err:
                            trial['feed'] = 'Error'
                            summary['responses_during_feed'] += 1
                            log.error("response during feed on box %s" % str(err))
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()

                        except hwio.HopperAlreadyUpError as err:
                            trial['feed'] = 'Error'
                            log.warning("hopper already up on box %s" % str(err))
                            summary['hopper_already_up'] += 1
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()

                        except hwio.HopperDidntRaiseError as err:
                            trial['feed'] = 'Error'
                            log.error("hopper didn't come up on box %s" % str(err))
                            summary['hopper_failures'] += 1
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()

                        except hwio.HopperDidntDropError as err:
                            trial['feed'] = 'Error'
                            log.warning("hopper didn't go down on box %s" % str(err))
                            summary['hopper_wont_go_down'] += 1
                            box.reset()
                            box.lights_on()

                        finally:
                            trial['cum_correct_thresh'] = random.randint(1, 2*options['variable_ratio']-1)

                elif trial['response'] == 'none':
                    # ignore non-responses
                    pass

                else:
                    trial['timeout_epoch'] = box.timeout(options['timeout_dur'][trial['class']])
                    trial['timeout'] = True
                    trial['cum_correct'] = 0
                    do_correction = True

            # correction trial logic.
            elif trial['type'] == 'correction':
                if trial['response'] == trial['class']:
                    # correct response
                    do_correction = False

                    if options['secondary_reinf']:
                        # give secondary reinforcer
                        trial['flash_epoch'] = box.flash(dur=0.5)

                    if (trial['cum_correct'] >= trial['cum_correct_thresh']):
                        # if cum currect reaches feed threshold, then feed
                        summary['feeds'] += 1
                        trial['feed'] = True

                        try:
                            trial['feed_epoch'] = box.feed(options['feed_dur'][trial['class']])
                        # but catch the feed errors
                        except hwio.ResponseDuringFeedError as err:
                            trial['feed'] = 'Error'
                            summary['responses_during_feed'] += 1
                            log.error("response during feed on box %s" % str(err))
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()
                        except hwio.HopperAlreadyUpError as err:
                            trial['feed'] = 'Error'
                            log.warning("hopper already up on box %s" % str(err))
                            summary['hopper_already_up'] += 1
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()
                        except hwio.HopperDidntRaiseError as err:
                            trial['feed'] = 'Error'
                            log.error("hopper didn't come up on box %s" % str(err))
                            summary['hopper_failures'] += 1
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()
                        except hwio.HopperDidntDropError as err:
                            trial['feed'] = 'Error'
                            log.warning("hopper didn't go down on box %s" % str(err))
                            summary['hopper_wont_go_down'] += 1
                            box.reset()
                            box.lights_on()

                elif trial['response'] == 'none':
                    # ignore non-responses
                    pass

                else:
                    trial['timeout_epoch'] = box.timeout(options['timeout_dur'][trial['class']])
                    trial['timeout'] = True
                    trial['cum_correct'] = 0
                    do_correction = True

            trial['trial_end'] = dt.datetime.now()
            save_trial(trial, options)

            if trial['feed']:
                trial['cum_correct'] = 0

        except GoodNite:
            """ reset experimental parameters for the next day """
            poll_int = 60.0
            box.lights_off()
            summary = utils.init_summary()
            log.debug('waiting %f seconds before checking light schedule...' % (poll_int))
            hwio.wait(poll_int)

        except hwio.ComediError as err:
            log.critical(str(err) + ", terminating operant control script")
            do_experiment = False

        except hwio.OperantError as err:
            log.error(str(err))

        finally:
            utils.write_summary(summary,options['summaryDAT'])

<commit_msg>HACK: moved transition calculation to allow save of options to json<commit_after>#!/usr/bin/python

## TODO:
# Save rDAT and trial info
# Implement Error Handling
# Add support for correction trials
# Add support for probe trials

import os, sys, random, csv, logging, logging.handlers, json
import numpy as np
import datetime as dt
from pyoperant import hwio, utils

class GoodNite(Exception):
    """ exception for when the lights should be off """
    pass

def get_stimulus(trial_class,options):
    """ take trial class and experimental options and return a tuple containing the wav filename & additional info to play

    take in the trial class and the options dict
    returns a stimulus dictionary

    """
    # determine stimulus for next trial, based on class
    stim = dict()
    stim['class'] = trial_class

    stim['string'] = []
    input_file_list = []
    trial_isi = []
    motif_id = 0
    cum_matrix = np.cumsum(options['models'][trial_class], axis=1)
    for item in range(options['strlen_max']):
        motif_id = (cum_matrix[motif_id] < random.random()).sum()
        motif_name = options['stim_map'][str(motif_id)]
        isi = random.gauss(options['isi_mean'], options['isi_stdev'])
        if isi < 0.0: isi = 0.0
        input_file_list.append((options['stims'][motif_name], isi))
        stim['string'].append(motif_name)
    input_file_list[-1] = (input_file_list[-1][0], 0.0)
    stim['input_file_list'] = input_file_list

    (stim['filename'], stim['epochs']) = utils.concat_wav(stim['input_file_list'], os.path.join(options['stim_path'], 'temp_concat.wav'))

    stim['dur'] = stim['epochs'][-1][-1]

    return stim

def get_options(cmd_line):
    """ get all of the configuration options for the experiment """

    # set path variables
    CODE_PATH = os.path.dirname(os.path.realpath(__file__))
    USER_PATH = os.path.expanduser('~')
    BIRD_PATH = os.path.join(USER_PATH,'opdat', 'B' + cmd_line['subj'])
    STIM_PATH = os.path.join(BIRD_PATH,'stims')

    with open(os.path.join(BIRD_PATH, cmd_line['config_file'])) as config_f:
        options = json.load(config_f)
    options['box_id'] = cmd_line['box']
    options['subject_id'] = cmd_line['subj']
    options['code_path'] = CODE_PATH
    options['user_path'] = USER_PATH
    options['bird_path'] = BIRD_PATH
    options['stim_path'] = STIM_PATH
    options['script_fname'] = os.path.basename(__file__)

    # run through stim files
    for name, filename in options['stims'].items():
        filename_full = os.path.join(options['stim_path'], filename)
        options['stims'][name] = filename_full

    # define log files, rDAT files
    filetime_fmt = '%Y%m%d%H%M%S'
    timestamp = dt.datetime.now().strftime(filetime_fmt)

    options['log_file'] = os.path.join(options['bird_path'], options['subject_id'] + '.log')
    options['data_csv'] = os.path.join(options['bird_path'], options['subject_id'] + '_' + options['script_fname'] + '_trials_' + timestamp + '.csv')
    options['config_snapshot'] = os.path.join(options['bird_path'], options['subject_id'] + '_' + options['script_fname'] + '_config_' + timestamp + '.json')
    options['summaryDAT'] = os.path.join(options['bird_path'], options['subject_id'] + '.summaryDAT')
    options['fields_to_save'] = ['number','type','trial_start','class','stim_start','stim_string','response','feed','timeout','cum_correct','cum_correct_thresh']

    # oreo
    with open(options['data_csv'], 'wb') as data_fh:
        trialWriter = csv.writer(data_fh)
        trialWriter.writerow(options['fields_to_save'])
    with open(options['config_snapshot'], 'wb') as config_snap:
        json.dump(options, config_snap, sort_keys=True, indent=4)

    # run through the transitions for each stimulus and generate transition matrixes
    n = len(options['stims'])
    options['models'] = {}
    for stim_class, pair_set in options['diagnostic_transitions'].items():
        mat = np.zeros((n+1,n+1))
        mat[0][1:] = 1.0/n
        for pair in pair_set:
            mat[pair[0]][1:] = (1-options['alpha'])/(n-1)
            mat[pair[0]][pair[1]] = options['alpha']
        options['models'][stim_class] = mat

    return options

def run_trial(trial, options):
    pass

def save_trial(trial, options):
    # write trial results to CSV
    with open(options['data_csv'],'ab') as data_fh:
        trialWriter = csv.DictWriter(data_fh,fieldnames=options['fields_to_save'],extrasaction='ignore')
        trialWriter.writerow(trial)

def wait_for_peck(box, options):
    """ runs a loop, querying for pecks. returns peck time or "GoodNite" exception """
    no_peck = True
    box.write(box.dio['LED_center'],True)
    while no_peck:
        no_peck = not box.read(box.dio['IR_center'])

        if not utils.check_time(options['light_schedule']):
            raise GoodNite()
    return dt.datetime.now()

if __name__ == "__main__":

    options = get_options(utils.parse_commandline())

    if  options['debug']:
        log_level = logging.DEBUG
    else:
        log_level = logging.INFO

    logging.basicConfig(filename=options['log_file'],
                        level=log_level,
                        format='%(asctime)s:%(levelname)s:%(message)s')
    log = logging.getLogger()
    #email_handler = logging.handlers.SMTPHandler(mailhost='localhost',
    #                                             fromaddr='bird@vogel.ucsd.edu',
    #                                             toaddrs=[options['experimenter']['email'],],
    #                                             subject='error notice',
    #                                             )
    #email_handler.setlevel(logging.ERROR)
    #log.addHandler(email_handler)

    # initialize box
    box = hwio.OperantBox(options['box_id'])
    log.debug('box %i initialized' % options['box_id'])
    box.reset()

    trial = dict()
    trial['correct'] = False
    trial['cum_correct'] = 0
    trial['cum_correct_thresh'] = 1
    trial['number'] = 0

    do_correction = False

    summary = utils.init_summary()

    # start experiment
    do_experiment = True
    while do_experiment:

        try:
            # first, check if we should be running trials, otherwise lightsout
            if not utils.check_time(options['light_schedule']):
                raise GoodNite()

            # make sure lights are on at the beginning of each trial, prep for trial
            box.lights_on()

            trial['feed'] = False
            trial['timeout'] = False
            trial['number'] += 1

            if do_correction and options['correction_trials']:
                trial['type'] = 'correction'
                log.debug("correction trial: class is %s" % trial['class'])
            else:
                trial['type'] = 'normal'
                trial['class'] = random.choice(options['models'].keys())
                trial_stim = get_stimulus(trial['class'],options)
                log.debug("trial class is %s" % trial['class'])

            # wait for bird to peck
            log.debug('waiting for peck...')
            trial['trial_start'] = wait_for_peck(box, options)

            # record trial initiation
            summary['trials'] += 1
            summary['last_trial_time'] = trial['trial_start'].ctime()
            log.info("trial started at %s" % trial['trial_start'].ctime())

            # play temp stimulus
            trial['stim_start'] = dt.datetime.now()
            wave_proc = box.play_wav(trial_stim['filename'])
            box.write(box.dio['LED_center'],False)

            # wait for response
            wait_min = trial_stim['epochs'][options['strlen_min']-1][-1]
            hwio.wait(wait_min)

            # check for responses
            wait_max = trial_stim['epochs'][-1][-1] + options['response_win']
            check_peck = True
            box.write(box.dio['LED_left'],True)
            box.write(box.dio['LED_right'],True)
            while check_peck:
                elapsed_time = dt.datetime.now() - trial['stim_start']
                if elapsed_time > dt.timedelta(seconds=wait_max):
                    trial['response'] = 'none'
                    check_peck = False
                elif box.read(box.dio['IR_left']):
                    trial['response_time'] = dt.datetime.now()
                    wave_proc.terminate()
                    trial['response'] = 'L'
                    check_peck = False
                    summary['responses'] += 1
                elif box.read(box.dio['IR_right']):
                    trial['response_time'] = dt.datetime.now()
                    wave_proc.terminate()
                    trial['response'] = 'R'
                    check_peck = False
                    summary['responses'] += 1
            # TODO: note response in event file
            box.write(box.dio['LED_left'],False)
            box.write(box.dio['LED_right'],False)
            trial['response_timedelta'] = elapsed_time

            # calculate the number of motifs the bird heard
            num_mots = 0
            for epoch in trial_stim['epochs']:
                if elapsed_time > dt.timedelta(seconds=epoch[0]):
                    num_mots += 1
            # determine the string of motifs the bird heard
            trial['stim_string'] = ';'.join(trial_stim['string'][0:num_mots])
            trial['stim_epochs'] = trial_stim['epochs'][0:num_mots]

            # decide how to respond to the subject for normal trials
            if trial['type'] == 'normal':
                if trial['response'] == trial['class']:
                    # correct response
                    do_correction = False
                    trial['cum_correct'] += 1

                    if options['secondary_reinf']:
                        # give secondary reinforcer
                        trial['flash_epoch'] = box.flash(dur=0.5)

                    if (trial['cum_correct'] >= trial['cum_correct_thresh']):
                        # if cum currect reaches feed threshold, then feed
                        summary['feeds'] += 1
                        trial['feed'] = True

                        try:
                            trial['feed_epoch'] = box.feed(options['feed_dur'][trial['class']])

                        # but catch the feed errors
                        except hwio.ResponseDuringFeedError as err:
                            trial['feed'] = 'Error'
                            summary['responses_during_feed'] += 1
                            log.error("response during feed on box %s" % str(err))
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()

                        except hwio.HopperAlreadyUpError as err:
                            trial['feed'] = 'Error'
                            log.warning("hopper already up on box %s" % str(err))
                            summary['hopper_already_up'] += 1
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()

                        except hwio.HopperDidntRaiseError as err:
                            trial['feed'] = 'Error'
                            log.error("hopper didn't come up on box %s" % str(err))
                            summary['hopper_failures'] += 1
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()

                        except hwio.HopperDidntDropError as err:
                            trial['feed'] = 'Error'
                            log.warning("hopper didn't go down on box %s" % str(err))
                            summary['hopper_wont_go_down'] += 1
                            box.reset()
                            box.lights_on()

                        finally:
                            trial['cum_correct_thresh'] = random.randint(1, 2*options['variable_ratio']-1)

                elif trial['response'] == 'none':
                    # ignore non-responses
                    pass

                else:
                    trial['timeout_epoch'] = box.timeout(options['timeout_dur'][trial['class']])
                    trial['timeout'] = True
                    trial['cum_correct'] = 0
                    do_correction = True

            # correction trial logic.
            elif trial['type'] == 'correction':
                if trial['response'] == trial['class']:
                    # correct response
                    do_correction = False

                    if options['secondary_reinf']:
                        # give secondary reinforcer
                        trial['flash_epoch'] = box.flash(dur=0.5)

                    if (trial['cum_correct'] >= trial['cum_correct_thresh']):
                        # if cum currect reaches feed threshold, then feed
                        summary['feeds'] += 1
                        trial['feed'] = True

                        try:
                            trial['feed_epoch'] = box.feed(options['feed_dur'][trial['class']])
                        # but catch the feed errors
                        except hwio.ResponseDuringFeedError as err:
                            trial['feed'] = 'Error'
                            summary['responses_during_feed'] += 1
                            log.error("response during feed on box %s" % str(err))
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()
                        except hwio.HopperAlreadyUpError as err:
                            trial['feed'] = 'Error'
                            log.warning("hopper already up on box %s" % str(err))
                            summary['hopper_already_up'] += 1
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()
                        except hwio.HopperDidntRaiseError as err:
                            trial['feed'] = 'Error'
                            log.error("hopper didn't come up on box %s" % str(err))
                            summary['hopper_failures'] += 1
                            hwio.wait(options['feed_dur'][trial['class']])
                            box.reset()
                            box.lights_on()
                        except hwio.HopperDidntDropError as err:
                            trial['feed'] = 'Error'
                            log.warning("hopper didn't go down on box %s" % str(err))
                            summary['hopper_wont_go_down'] += 1
                            box.reset()
                            box.lights_on()

                elif trial['response'] == 'none':
                    # ignore non-responses
                    pass

                else:
                    trial['timeout_epoch'] = box.timeout(options['timeout_dur'][trial['class']])
                    trial['timeout'] = True
                    trial['cum_correct'] = 0
                    do_correction = True

            trial['trial_end'] = dt.datetime.now()
            save_trial(trial, options)

            if trial['feed']:
                trial['cum_correct'] = 0

        except GoodNite:
            """ reset experimental parameters for the next day """
            poll_int = 60.0
            box.lights_off()
            summary = utils.init_summary()
            log.debug('waiting %f seconds before checking light schedule...' % (poll_int))
            hwio.wait(poll_int)

        except hwio.ComediError as err:
            log.critical(str(err) + ", terminating operant control script")
            do_experiment = False

        except hwio.OperantError as err:
            log.error(str(err))

        finally:
            utils.write_summary(summary,options['summaryDAT'])

<|endoftext|>


<commit_before>import numpy as np

from .. import util
from .. import transformations

from ..geometry import align_vectors

class Camera(object):

    def __init__(
            self,
            name=None,
            resolution=None,
            focal=None,
            fov=None,
            scene=None,
            transform=None):
        """
        Create a new Camera object that stores camera intrinsic
        and extrinsic parameters.

        TODO: skew is not supported
        TODO: cx and cy that are not half of width and height

        Parameters
        ------------
        name : str or None
          Name for camera to be used as node name
        resolution : (2,) int
          Pixel size in (height, width)
        focal : (2,) float
          Focal length in pixels. Either pass this OR FOV
          but not both.  focal = (K[0][0], K[1][1])
        fov : (2,) float
          Field of view (fovx, fovy) in degrees
        transform : (4, 4) float
          Camera transform, extrinsic homogoenous
          transformation matrix.
        """

        if name is None:
            # if name is not passed, make it something unique
            self.name = 'camera_{}'.format(util.unique_id(6).upper())
        else:
            # otherwise assign it
            self.name = name

        if fov is None and focal is None:
            raise ValueError('either focal length or FOV required!')

        # store whether or not we computed the focal length
        self._focal_computed = False

        # set the passed (2,) float focal length
        self.focal = focal

        # set the passed (2,) float FOV in degrees
        self.fov = fov

        if resolution is None:
            # if unset make resolution 15 pixels per degree
            resolution = (self.fov * 15.0).astype(np.int64)
        self.resolution = resolution

        # add a back- reference to scene object
        self.scene = scene
        # set transform
        self.transform = transform

    @property
    def resolution(self):
        """
        Get the camera resolution in pixels.

        Returns
        ------------
        resolution (2,) float
          Camera resolution in pixels
        """
        return self._resolution

    @resolution.setter
    def resolution(self, values):
        """
        Set the camera resolution in pixels.

        Parameters
        ------------
        resolution (2,) float
          Camera resolution in pixels
        """
        values = np.asanyarray(values, dtype=np.int64)
        if values.shape != (2,):
            raise ValueError('resolution must be (2,) float')
        values.flags.writeable = False
        self._resolution = values
        # unset computed value that depends on the other plus resolution
        if self._focal_computed:
            self._focal = None
        else:
            # fov must be computed
            self._fov = None

    @property
    def scene(self):
        """
        Get a reference to the scene that this camera is in.

        Returns
        -------------
        scene : None, or trimesh.Scene
          Scene where this camera is attached
        """
        return self._scene

    @scene.setter
    def scene(self, value):
        """
        Set the reference to the scene that this camera is in.

        Parameters
        -------------
        scene : None, or trimesh.Scene
          Scene where this camera is attached
        """

        # save the scene reference
        self._scene = value

        # check if we have local not None transform
        # an if we can apply it to the scene graph
        # also check here that scene is a real scene
        if (hasattr(self, '_transform') and
            self._transform is not None and
                hasattr(value, 'graph')):
            # set scene transform to locally saved transform
            self._scene.graph[self.name] = self._transform
            # set local transform to None
            self._transform = None

    @property
    def transform(self):
        """
        Get the (4, 4) homogenous transformation from the
        world frame to this camera object.

        Returns
        ------------
        transform : (4, 4) float
          Transform from world to camera
        """
        # no scene set
        if self._scene is None:
            # no transform saved locally
            if not hasattr(self, '_transform') or self._transform is None:
                return np.eye(4)
            # transform saved locally
            return self._transform

        # get the transform from the scene
        return self._scene.graph[self.name][0]

    @transform.setter
    def transform(self, values):
        """

        Set the (4, 4) homogenous transformation from the
        world frame to this camera object.

        Parameters
        ------------
        transform : (4, 4) float
          Transform from world to camera
        """
        if values is None:
            return
        matrix = np.asanyarray(values, dtype=np.float64)
        if matrix.shape != (4, 4):
            raise ValueError('transform must be (4, 4) float!')
        matrix.flags.writeable = False
        if self._scene is None:
            # if no scene, save transform locally
            self._transform = matrix
        else:
            # assign passed values to transform
            self._scene.graph[self.name] = matrix

    @property
    def focal(self):
        """
        Get the focal length in pixels for the camera.

        Returns
        ------------
        focal : (2,) float
          Focal length in pixels
        """
        if self._focal is None:
            # calculate focal length from FOV
            focal = (
                self._resolution / (2.0 * np.tan(np.radians(self._fov / 2.0))))
            focal.flags.writeable = False
            self._focal = focal

        return self._focal

    @focal.setter
    def focal(self, values):
        """
        Set the focal length in pixels for the camera.

        Returns
        ------------
        focal : (2,) float
          Focal length in pixels.
        """
        if values is None:
            self._focal = None
        else:
            # flag this as not computed (hence fov must be)
            # this is necessary so changes to resolution can reset the
            # computed quantity without changing the explicitly set quantity
            self._focal_computed = False
            values = np.asanyarray(values, dtype=np.float64)
            if values.shape != (2,):
                raise ValueError('focal length must be (2,) float')
            values.flags.writeable = False
            # assign passed values to focal length
            self._focal = values
            # focal overrides FOV
            self._fov = None

    @property
    def K(self):
        """
        Get the intrinsic matrix for the Camera object.

        Returns
        -----------
        K : (3, 3) float
          Intrinsic matrix for camera
        """
        K = np.eye(3, dtype=np.float64)
        K[0, 0] = self.focal[0]
        K[1, 1] = self.focal[1]
        K[:2, 2] = self.resolution / 2.0
        return K

    @K.setter
    def K(self, values):
        if values is None:
            return
        values = np.asanyarray(values, dtype=np.float64)
        if values.shape != (3, 3):
            raise ValueError('matrix must be (3,3)!')

        if not np.allclose(values.flatten()[[1, 3, 6, 7, 8]],
                           [0, 0, 0, 0, 1]):
            raise ValueError(
                'matrix should only have focal length and resolution!')

        # set focal length from matrix
        self.focal = [values[0, 0], values[1, 1]]
        # set resolution from matrix
        self.resolution = values[:2, 2] * 2

    @property
    def fov(self):
        """
        Get the field of view in degrees.

        Returns
        -------------
        fov : (2,) float
          XY field of view in degrees
        """
        if self._fov is None:
            fov = 2.0 * np.degrees(
                np.arctan((self._resolution / 2.0) / self._focal))
            fov.flags.writeable = False
            self._fov = fov
        return self._fov

    @fov.setter
    def fov(self, values):
        """
        Set the field of view in degrees.

        Parameters
        -------------
        values : (2,) float
          Size of FOV to set in degrees
        """
        if values is None:
            self._fov = None
        else:
            # flag this as computed (hence fov must not be)
            # this is necessary so changes to resolution can reset the
            # computed quantity without changing the explicitly set quantity
            self._focal_computed = True
            values = np.asanyarray(values, dtype=np.float64)
            if values.shape != (2,):
                raise ValueError('focal length must be (2,) int')
            values.flags.writeable = False
            # assign passed values to FOV
            self._fov = values
            # fov overrides focal
            self._focal = None

    def to_rays(self):
        """
        Convert a trimesh.scene.Camera object to ray origins
        and direction vectors. Will return one ray per pixel,
        as set in camera.resolution.

        Returns
        --------------
        origins : (n, 3) float
          Ray origins in space
        vectors : (n, 3) float
          Ray direction unit vectors
        angles : (n, 2) float
          Ray spherical coordinate angles in radians
        """
        return camera_to_rays(self)

def look_at(points, fov, rotation=None, distance=None, center=None):
    """
    Generate transform for a camera to keep a list
    of points in the camera's field of view.

    Parameters
    -------------
    points : (n, 3) float
      Points in space
    fov : (2,) float
      Field of view, in DEGREES
    rotation : None, or (4, 4) float
      Rotation matrix for initial rotation
    center : None, or (3,) float
      Center of field of view.

    Returns
    --------------
    transform : (4, 4) float
      Transformation matrix with points in view
    """

    if rotation is None:
        rotation = np.eye(4)
    else:
        rotation = np.asanyarray(rotation, dtype=np.float64)
    points = np.asanyarray(points, dtype=np.float64)

    # Transform points to camera frame (just use the rotation part)
    rinv = rotation[:3, :3].T
    points_c = rinv.dot(points.T).T

    if center is None:
        # Find the center of the points' AABB in camera frame
        center_c = points_c.min(axis=0) + 0.5 * points_c.ptp(axis=0)
    else:
        # Transform center to camera frame
        center_c = rinv.dot(center)

    # Re-center the points around the camera-frame origin
    points_c -= center_c

    # Find the minimum distance for the camera from the origin
    # so that all points fit in the view frustrum
    tfov = np.tan(np.radians(fov) / 2.0)

    if distance is None:
        distance = np.max(np.abs(points_c[:, :2]) /
                          tfov + points_c[:, 2][:, np.newaxis])

    # set the pose translation
    center_w = rotation[:3, :3].dot(center_c)
    cam_pose = rotation.copy()
    cam_pose[:3, 3] = center_w + distance * cam_pose[:3, 2]

    return cam_pose

def camera_to_rays(camera):
    """
    Convert a trimesh.scene.Camera object to ray origins
    and direction vectors. Will return one ray per pixel,
    as set in camera.resolution.

    Parameters
    --------------
    camera : trimesh.scene.Camera
      Camera with transform defined

    Returns
    --------------
    origins : (n, 3) float
      Ray origins in space
    vectors : (n, 3) float
      Ray direction unit vectors
    angles : (n, 2) float
      Ray spherical coordinate angles in radians
    """
    # radians of half the field of view
    half = np.radians(camera.fov / 2.0)
    # scale it down by two pixels to keep image under resolution
    half *= (camera.resolution - 2) / camera.resolution

    # create an evenly spaced list of angles
    angles = util.grid_linspace(bounds=[-half, half],
                                count=camera.resolution)

    # turn the angles into unit vectors
    vectors = util.unitize(np.column_stack((
        np.sin(angles),
        np.ones(len(angles)))))

    # flip the camera transform to change sign of Z
    transform = np.dot(
        camera.transform,
        align_vectors([1, 0, 0], [-1, 0, 0]))

    # apply the rotation to the direction vectors
    vectors = transformations.transform_points(
        vectors,
        transform,
        translate=False)

    # camera origin is single point, extract from transform
    origin = transformations.translation_from_matrix(transform)
    # tile it into corresponding list of ray vectorsy
    origins = np.ones_like(vectors) * origin

    return origins, vectors, angles
<commit_msg>Add copy method to scene.Camera<commit_after>import copy

import numpy as np

from .. import util
from .. import transformations

from ..geometry import align_vectors

class Camera(object):

    def __init__(
            self,
            name=None,
            resolution=None,
            focal=None,
            fov=None,
            scene=None,
            transform=None):
        """
        Create a new Camera object that stores camera intrinsic
        and extrinsic parameters.

        TODO: skew is not supported
        TODO: cx and cy that are not half of width and height

        Parameters
        ------------
        name : str or None
          Name for camera to be used as node name
        resolution : (2,) int
          Pixel size in (height, width)
        focal : (2,) float
          Focal length in pixels. Either pass this OR FOV
          but not both.  focal = (K[0][0], K[1][1])
        fov : (2,) float
          Field of view (fovx, fovy) in degrees
        transform : (4, 4) float
          Camera transform, extrinsic homogoenous
          transformation matrix.
        """

        if name is None:
            # if name is not passed, make it something unique
            self.name = 'camera_{}'.format(util.unique_id(6).upper())
        else:
            # otherwise assign it
            self.name = name

        if fov is None and focal is None:
            raise ValueError('either focal length or FOV required!')

        # store whether or not we computed the focal length
        self._focal_computed = False

        # set the passed (2,) float focal length
        self.focal = focal

        # set the passed (2,) float FOV in degrees
        self.fov = fov

        if resolution is None:
            # if unset make resolution 15 pixels per degree
            resolution = (self.fov * 15.0).astype(np.int64)
        self.resolution = resolution

        # add a back- reference to scene object
        self.scene = scene
        # set transform
        self.transform = transform

    def copy(self):
        """
        Get a copy of the current camera.
        """
        return copy.deepcopy(self)

    @property
    def resolution(self):
        """
        Get the camera resolution in pixels.

        Returns
        ------------
        resolution (2,) float
          Camera resolution in pixels
        """
        return self._resolution

    @resolution.setter
    def resolution(self, values):
        """
        Set the camera resolution in pixels.

        Parameters
        ------------
        resolution (2,) float
          Camera resolution in pixels
        """
        values = np.asanyarray(values, dtype=np.int64)
        if values.shape != (2,):
            raise ValueError('resolution must be (2,) float')
        values.flags.writeable = False
        self._resolution = values
        # unset computed value that depends on the other plus resolution
        if self._focal_computed:
            self._focal = None
        else:
            # fov must be computed
            self._fov = None

    @property
    def scene(self):
        """
        Get a reference to the scene that this camera is in.

        Returns
        -------------
        scene : None, or trimesh.Scene
          Scene where this camera is attached
        """
        return self._scene

    @scene.setter
    def scene(self, value):
        """
        Set the reference to the scene that this camera is in.

        Parameters
        -------------
        scene : None, or trimesh.Scene
          Scene where this camera is attached
        """

        # save the scene reference
        self._scene = value

        # check if we have local not None transform
        # an if we can apply it to the scene graph
        # also check here that scene is a real scene
        if (hasattr(self, '_transform') and
            self._transform is not None and
                hasattr(value, 'graph')):
            # set scene transform to locally saved transform
            self._scene.graph[self.name] = self._transform
            # set local transform to None
            self._transform = None

    @property
    def transform(self):
        """
        Get the (4, 4) homogenous transformation from the
        world frame to this camera object.

        Returns
        ------------
        transform : (4, 4) float
          Transform from world to camera
        """
        # no scene set
        if self._scene is None:
            # no transform saved locally
            if not hasattr(self, '_transform') or self._transform is None:
                return np.eye(4)
            # transform saved locally
            return self._transform

        # get the transform from the scene
        return self._scene.graph[self.name][0]

    @transform.setter
    def transform(self, values):
        """

        Set the (4, 4) homogenous transformation from the
        world frame to this camera object.

        Parameters
        ------------
        transform : (4, 4) float
          Transform from world to camera
        """
        if values is None:
            return
        matrix = np.asanyarray(values, dtype=np.float64)
        if matrix.shape != (4, 4):
            raise ValueError('transform must be (4, 4) float!')
        matrix.flags.writeable = False
        if self._scene is None:
            # if no scene, save transform locally
            self._transform = matrix
        else:
            # assign passed values to transform
            self._scene.graph[self.name] = matrix

    @property
    def focal(self):
        """
        Get the focal length in pixels for the camera.

        Returns
        ------------
        focal : (2,) float
          Focal length in pixels
        """
        if self._focal is None:
            # calculate focal length from FOV
            focal = (
                self._resolution / (2.0 * np.tan(np.radians(self._fov / 2.0))))
            focal.flags.writeable = False
            self._focal = focal

        return self._focal

    @focal.setter
    def focal(self, values):
        """
        Set the focal length in pixels for the camera.

        Returns
        ------------
        focal : (2,) float
          Focal length in pixels.
        """
        if values is None:
            self._focal = None
        else:
            # flag this as not computed (hence fov must be)
            # this is necessary so changes to resolution can reset the
            # computed quantity without changing the explicitly set quantity
            self._focal_computed = False
            values = np.asanyarray(values, dtype=np.float64)
            if values.shape != (2,):
                raise ValueError('focal length must be (2,) float')
            values.flags.writeable = False
            # assign passed values to focal length
            self._focal = values
            # focal overrides FOV
            self._fov = None

    @property
    def K(self):
        """
        Get the intrinsic matrix for the Camera object.

        Returns
        -----------
        K : (3, 3) float
          Intrinsic matrix for camera
        """
        K = np.eye(3, dtype=np.float64)
        K[0, 0] = self.focal[0]
        K[1, 1] = self.focal[1]
        K[:2, 2] = self.resolution / 2.0
        return K

    @K.setter
    def K(self, values):
        if values is None:
            return
        values = np.asanyarray(values, dtype=np.float64)
        if values.shape != (3, 3):
            raise ValueError('matrix must be (3,3)!')

        if not np.allclose(values.flatten()[[1, 3, 6, 7, 8]],
                           [0, 0, 0, 0, 1]):
            raise ValueError(
                'matrix should only have focal length and resolution!')

        # set focal length from matrix
        self.focal = [values[0, 0], values[1, 1]]
        # set resolution from matrix
        self.resolution = values[:2, 2] * 2

    @property
    def fov(self):
        """
        Get the field of view in degrees.

        Returns
        -------------
        fov : (2,) float
          XY field of view in degrees
        """
        if self._fov is None:
            fov = 2.0 * np.degrees(
                np.arctan((self._resolution / 2.0) / self._focal))
            fov.flags.writeable = False
            self._fov = fov
        return self._fov

    @fov.setter
    def fov(self, values):
        """
        Set the field of view in degrees.

        Parameters
        -------------
        values : (2,) float
          Size of FOV to set in degrees
        """
        if values is None:
            self._fov = None
        else:
            # flag this as computed (hence fov must not be)
            # this is necessary so changes to resolution can reset the
            # computed quantity without changing the explicitly set quantity
            self._focal_computed = True
            values = np.asanyarray(values, dtype=np.float64)
            if values.shape != (2,):
                raise ValueError('focal length must be (2,) int')
            values.flags.writeable = False
            # assign passed values to FOV
            self._fov = values
            # fov overrides focal
            self._focal = None

    def to_rays(self):
        """
        Convert a trimesh.scene.Camera object to ray origins
        and direction vectors. Will return one ray per pixel,
        as set in camera.resolution.

        Returns
        --------------
        origins : (n, 3) float
          Ray origins in space
        vectors : (n, 3) float
          Ray direction unit vectors
        angles : (n, 2) float
          Ray spherical coordinate angles in radians
        """
        return camera_to_rays(self)

def look_at(points, fov, rotation=None, distance=None, center=None):
    """
    Generate transform for a camera to keep a list
    of points in the camera's field of view.

    Parameters
    -------------
    points : (n, 3) float
      Points in space
    fov : (2,) float
      Field of view, in DEGREES
    rotation : None, or (4, 4) float
      Rotation matrix for initial rotation
    center : None, or (3,) float
      Center of field of view.

    Returns
    --------------
    transform : (4, 4) float
      Transformation matrix with points in view
    """

    if rotation is None:
        rotation = np.eye(4)
    else:
        rotation = np.asanyarray(rotation, dtype=np.float64)
    points = np.asanyarray(points, dtype=np.float64)

    # Transform points to camera frame (just use the rotation part)
    rinv = rotation[:3, :3].T
    points_c = rinv.dot(points.T).T

    if center is None:
        # Find the center of the points' AABB in camera frame
        center_c = points_c.min(axis=0) + 0.5 * points_c.ptp(axis=0)
    else:
        # Transform center to camera frame
        center_c = rinv.dot(center)

    # Re-center the points around the camera-frame origin
    points_c -= center_c

    # Find the minimum distance for the camera from the origin
    # so that all points fit in the view frustrum
    tfov = np.tan(np.radians(fov) / 2.0)

    if distance is None:
        distance = np.max(np.abs(points_c[:, :2]) /
                          tfov + points_c[:, 2][:, np.newaxis])

    # set the pose translation
    center_w = rotation[:3, :3].dot(center_c)
    cam_pose = rotation.copy()
    cam_pose[:3, 3] = center_w + distance * cam_pose[:3, 2]

    return cam_pose

def camera_to_rays(camera):
    """
    Convert a trimesh.scene.Camera object to ray origins
    and direction vectors. Will return one ray per pixel,
    as set in camera.resolution.

    Parameters
    --------------
    camera : trimesh.scene.Camera
      Camera with transform defined

    Returns
    --------------
    origins : (n, 3) float
      Ray origins in space
    vectors : (n, 3) float
      Ray direction unit vectors
    angles : (n, 2) float
      Ray spherical coordinate angles in radians
    """
    # radians of half the field of view
    half = np.radians(camera.fov / 2.0)
    # scale it down by two pixels to keep image under resolution
    half *= (camera.resolution - 2) / camera.resolution

    # create an evenly spaced list of angles
    angles = util.grid_linspace(bounds=[-half, half],
                                count=camera.resolution)

    # turn the angles into unit vectors
    vectors = util.unitize(np.column_stack((
        np.sin(angles),
        np.ones(len(angles)))))

    # flip the camera transform to change sign of Z
    transform = np.dot(
        camera.transform,
        align_vectors([1, 0, 0], [-1, 0, 0]))

    # apply the rotation to the direction vectors
    vectors = transformations.transform_points(
        vectors,
        transform,
        translate=False)

    # camera origin is single point, extract from transform
    origin = transformations.translation_from_matrix(transform)
    # tile it into corresponding list of ray vectorsy
    origins = np.ones_like(vectors) * origin

    return origins, vectors, angles
<|endoftext|>