<commit_before>/*
 * Copyright (c) 2008 Digital Bazaar, Inc.  All rights reserved.
 */
#include "db/compress/deflate/Deflater.h"

using namespace db::compress::deflate;
using namespace db::io;
using namespace db::rt;

Deflater::Deflater()
{
   mShouldFinish = false;
   mFinished = false;
}

Deflater::~Deflater()
{
   Deflater::cleanupStream();
}

void Deflater::cleanupStream()
{
   // clean up previous stream
   if(mFinished)
   {
      if(mDeflating)
      {
         deflateEnd(&mZipStream);
      }
      else
      {
         inflateEnd(&mZipStream);
      }
   }

   // re-initialize stream (use defaults via Z_NULL)
   mZipStream.zalloc = Z_NULL;
   mZipStream.zfree = Z_NULL;
   mZipStream.opaque = Z_NULL;
   mZipStream.next_in = Z_NULL;
   mZipStream.avail_in = 0;
}

bool Deflater::createException(int ret)
{
   bool rval = false;

   Exception* e = NULL;

   switch(ret)
   {
      case Z_OK:
      case Z_STREAM_END:
      case Z_BUF_ERROR:
         // not an error, buffer just too small, space will be allocated
         break;
      case Z_MEM_ERROR:
         // not enough memory
         e = new Exception(
            "Not enough memory for inflation/deflation.",
            "db.compress.deflate.InsufficientMemory");
         break;
      case Z_VERSION_ERROR:
         // zlib library version incompatible
         e = new Exception(
            "Incompatible zlib library version.",
            "db.compress.deflate.IncompatibleVersion");
         break;
      case Z_STREAM_ERROR:
         // invalid stream parameters
         e = new Exception(
            "Invalid zip stream parameters. Null pointer?",
            "db.compress.deflate.InvalidZipStreamParams");
         break;
      default:
         // something else went wrong
         e = new Exception(
            "Could not inflate/deflate.",
            "db.compress.deflate.Error");
         break;
   }

   if(e != NULL)
   {
      rval = true;

      if(mZipStream.msg != NULL)
      {
         // use zlib stream error message as cause
         ExceptionRef cause = new Exception(mZipStream.msg);
         e->setCause(cause);
      }

      ExceptionRef ref = e;
      Exception::setLast(ref, false);
   }

   return e;
}

bool Deflater::startDeflating(int level, bool raw)
{
   // clean up previous stream
   cleanupStream();

   int ret;

   if(raw)
   {
      // windowBits is negative to indicate a raw stream
      // 8 is the default memLevel, use default strategy as well
      ret = deflateInit2(
         &mZipStream, level, Z_DEFLATED, -15, 8, Z_DEFAULT_STRATEGY);
   }
   else
   {
      // windowBits set to use zlib header/trailer
      // 8 is the default memLevel, use default strategy as well
      ret = deflateInit2(
         &mZipStream, level, Z_DEFLATED, 15, 8, Z_DEFAULT_STRATEGY);
   }

   mDeflating = true;
   mShouldFinish = false;
   mFinished = false;

   return !createException(ret);
}

bool Deflater::startInflating(bool raw)
{
   // clean up previous stream
   cleanupStream();

   int ret;

   if(raw)
   {
      // windowBits is negative to indicate a raw stream
      ret = inflateInit2(&mZipStream, -15);
   }
   else
   {
      // add 32 to windowBits to support gzip or zlib decoding
      ret = inflateInit2(&mZipStream, 15 + 32);
   }

   mDeflating = false;
   mShouldFinish = false;
   mFinished = false;

   return !createException(ret);
}

void Deflater::setInput(const char* b, int length, bool finish)
{
   // set zip stream input buffer
   mZipStream.next_in = (unsigned char*)b;
   mZipStream.avail_in = length;
   mShouldFinish = finish;
}

int Deflater::process(ByteBuffer* dst, bool resize)
{
   int rval = 0;

   if(!mFinished)
   {
      // when deflating, let zlib determine flushing to maximize compression
      // when inflating, flush output whenever possible
      int flush = (mShouldFinish ?
         Z_FINISH : (mDeflating ? Z_NO_FLUSH : Z_SYNC_FLUSH));

      // keep processing while no error, no output data, while there is
      // input data or processing should finish, and while there is room
      // to store the output data or if the destination buffer can be
      // resized
      while(rval == 0 &&
            (mZipStream.avail_in > 0 || mShouldFinish) &&
            (resize || !dst->isFull()))
      {
         if(resize && dst->isFull())
         {
            // allocate space for output
            dst->allocateSpace(1024, resize);
         }

         // set output buffer, store old free space
         int freeSpace = dst->freeSpace();
         dst->allocateSpace(freeSpace, false);
         mZipStream.next_out = (unsigned char*)dst->data();
         mZipStream.avail_out = freeSpace;

         // perform deflation/inflation
         int ret = (mDeflating) ?
            ::deflate(&mZipStream, flush) :
            ::inflate(&mZipStream, flush);

         // extend destination buffer by data written to it
         int length = freeSpace - mZipStream.avail_out;
         dst->extend(length);

         // handle potential exception
         if(createException(ret))
         {
            rval = -1;
         }
         else
         {
            // return number of output bytes
            rval = length;
         }

         if(ret == Z_STREAM_END)
         {
            // finished
            mFinished = true;
         }
      }
   }

   return rval;
}

MutationAlgorithm::Result Deflater::mutateData(
   ByteBuffer* src, ByteBuffer* dst, bool finish)
{
   MutationAlgorithm::Result rval = MutationAlgorithm::Stepped;

   if(!isFinished())
   {
      if(inputAvailable() == 0)
      {
         if(src->isEmpty() && !finish)
         {
            // more data required
            rval = MutationAlgorithm::NeedsData;
         }
         else
         {
            // set input
            setInput(src->data(), src->length(), finish);
         }
      }

      // keep processing while no output data and algorithm stepped
      while(dst->isEmpty() && rval == MutationAlgorithm::Stepped)
      {
         // try to process existing input
         int ret = process(dst, false);
         if(ret == 0)
         {
            // either request more data or algorithm is complete
            rval = (!isFinished()) ?
               MutationAlgorithm::NeedsData :
               MutationAlgorithm::CompleteTruncate;
         }
         else if(ret == -1)
         {
            // exception occurred
            rval = MutationAlgorithm::Error;
         }
         else if(isFinished())
         {
            // algorithm complete
            rval = MutationAlgorithm::CompleteTruncate;
         }

         // clear source buffer of data that has been consumed
         src->clear(src->length() - inputAvailable());
      }
   }
   else
   {
      // algorithm completed
      rval = MutationAlgorithm::CompleteTruncate;
   }

   return rval;
}

unsigned int Deflater::inputAvailable()
{
   return mZipStream.avail_in;
}

bool Deflater::isFinished()
{
   return mFinished;
}

unsigned int Deflater::getTotalInputBytes()
{
   return mZipStream.total_in;
}

unsigned int Deflater::getTotalOutputBytes()
{
   return mZipStream.total_out;
}
<commit_msg>Added #define for exception names.<commit_after>/*
 * Copyright (c) 2008-2009 Digital Bazaar, Inc.  All rights reserved.
 */
#include "db/compress/deflate/Deflater.h"

using namespace db::compress::deflate;
using namespace db::io;
using namespace db::rt;

#define EXCEPTION_DEFLATE "db.compress.deflate"

Deflater::Deflater()
{
   mShouldFinish = false;
   mFinished = false;
}

Deflater::~Deflater()
{
   Deflater::cleanupStream();
}

void Deflater::cleanupStream()
{
   // clean up previous stream
   if(mFinished)
   {
      if(mDeflating)
      {
         deflateEnd(&mZipStream);
      }
      else
      {
         inflateEnd(&mZipStream);
      }
   }

   // re-initialize stream (use defaults via Z_NULL)
   mZipStream.zalloc = Z_NULL;
   mZipStream.zfree = Z_NULL;
   mZipStream.opaque = Z_NULL;
   mZipStream.next_in = Z_NULL;
   mZipStream.avail_in = 0;
}

bool Deflater::createException(int ret)
{
   bool rval = false;

   Exception* e = NULL;

   switch(ret)
   {
      case Z_OK:
      case Z_STREAM_END:
      case Z_BUF_ERROR:
         // not an error, buffer just too small, space will be allocated
         break;
      case Z_MEM_ERROR:
         // not enough memory
         e = new Exception(
            "Not enough memory for inflation/deflation.",
            EXCEPTION_DEFLATE ".InsufficientMemory");
         break;
      case Z_VERSION_ERROR:
         // zlib library version incompatible
         e = new Exception(
            "Incompatible zlib library version.",
            EXCEPTION_DEFLATE ".IncompatibleVersion");
         break;
      case Z_STREAM_ERROR:
         // invalid stream parameters
         e = new Exception(
            "Invalid zip stream parameters. Null pointer?",
            EXCEPTION_DEFLATE ".InvalidZipStreamParams");
         break;
      default:
         // something else went wrong
         e = new Exception(
            "Could not inflate/deflate.",
            EXCEPTION_DEFLATE ".Error");
         break;
   }

   if(e != NULL)
   {
      rval = true;

      if(mZipStream.msg != NULL)
      {
         // use zlib stream error message as cause
         ExceptionRef cause = new Exception(mZipStream.msg);
         e->setCause(cause);
      }

      ExceptionRef ref = e;
      Exception::setLast(ref, false);
   }

   return e;
}

bool Deflater::startDeflating(int level, bool raw)
{
   // clean up previous stream
   cleanupStream();

   int ret;

   if(raw)
   {
      // windowBits is negative to indicate a raw stream
      // 8 is the default memLevel, use default strategy as well
      ret = deflateInit2(
         &mZipStream, level, Z_DEFLATED, -15, 8, Z_DEFAULT_STRATEGY);
   }
   else
   {
      // windowBits set to use zlib header/trailer
      // 8 is the default memLevel, use default strategy as well
      ret = deflateInit2(
         &mZipStream, level, Z_DEFLATED, 15, 8, Z_DEFAULT_STRATEGY);
   }

   mDeflating = true;
   mShouldFinish = false;
   mFinished = false;

   return !createException(ret);
}

bool Deflater::startInflating(bool raw)
{
   // clean up previous stream
   cleanupStream();

   int ret;

   if(raw)
   {
      // windowBits is negative to indicate a raw stream
      ret = inflateInit2(&mZipStream, -15);
   }
   else
   {
      // add 32 to windowBits to support gzip or zlib decoding
      ret = inflateInit2(&mZipStream, 15 + 32);
   }

   mDeflating = false;
   mShouldFinish = false;
   mFinished = false;

   return !createException(ret);
}

void Deflater::setInput(const char* b, int length, bool finish)
{
   // set zip stream input buffer
   mZipStream.next_in = (unsigned char*)b;
   mZipStream.avail_in = length;
   mShouldFinish = finish;
}

int Deflater::process(ByteBuffer* dst, bool resize)
{
   int rval = 0;

   if(!mFinished)
   {
      // when deflating, let zlib determine flushing to maximize compression
      // when inflating, flush output whenever possible
      int flush = (mShouldFinish ?
         Z_FINISH : (mDeflating ? Z_NO_FLUSH : Z_SYNC_FLUSH));

      // keep processing while no error, no output data, while there is
      // input data or processing should finish, and while there is room
      // to store the output data or if the destination buffer can be
      // resized
      while(rval == 0 &&
            (mZipStream.avail_in > 0 || mShouldFinish) &&
            (resize || !dst->isFull()))
      {
         if(resize && dst->isFull())
         {
            // allocate space for output
            dst->allocateSpace(1024, resize);
         }

         // set output buffer, store old free space
         int freeSpace = dst->freeSpace();
         dst->allocateSpace(freeSpace, false);
         mZipStream.next_out = (unsigned char*)dst->data();
         mZipStream.avail_out = freeSpace;

         // perform deflation/inflation
         int ret = (mDeflating) ?
            ::deflate(&mZipStream, flush) :
            ::inflate(&mZipStream, flush);

         // extend destination buffer by data written to it
         int length = freeSpace - mZipStream.avail_out;
         dst->extend(length);

         // handle potential exception
         if(createException(ret))
         {
            rval = -1;
         }
         else
         {
            // return number of output bytes
            rval = length;
         }

         if(ret == Z_STREAM_END)
         {
            // finished
            mFinished = true;
         }
      }
   }

   return rval;
}

MutationAlgorithm::Result Deflater::mutateData(
   ByteBuffer* src, ByteBuffer* dst, bool finish)
{
   MutationAlgorithm::Result rval = MutationAlgorithm::Stepped;

   if(!isFinished())
   {
      if(inputAvailable() == 0)
      {
         if(src->isEmpty() && !finish)
         {
            // more data required
            rval = MutationAlgorithm::NeedsData;
         }
         else
         {
            // set input
            setInput(src->data(), src->length(), finish);
         }
      }

      // keep processing while no output data and algorithm stepped
      while(dst->isEmpty() && rval == MutationAlgorithm::Stepped)
      {
         // try to process existing input
         int ret = process(dst, false);
         if(ret == 0)
         {
            // either request more data or algorithm is complete
            rval = (!isFinished()) ?
               MutationAlgorithm::NeedsData :
               MutationAlgorithm::CompleteTruncate;
         }
         else if(ret == -1)
         {
            // exception occurred
            rval = MutationAlgorithm::Error;
         }
         else if(isFinished())
         {
            // algorithm complete
            rval = MutationAlgorithm::CompleteTruncate;
         }

         // clear source buffer of data that has been consumed
         src->clear(src->length() - inputAvailable());
      }
   }
   else
   {
      // algorithm completed
      rval = MutationAlgorithm::CompleteTruncate;
   }

   return rval;
}

unsigned int Deflater::inputAvailable()
{
   return mZipStream.avail_in;
}

bool Deflater::isFinished()
{
   return mFinished;
}

unsigned int Deflater::getTotalInputBytes()
{
   return mZipStream.total_in;
}

unsigned int Deflater::getTotalOutputBytes()
{
   return mZipStream.total_out;
}
<|endoftext|>


<commit_before>/*
 * Licensed under the MIT License (MIT)
 *
 * Copyright (c) 2013 AudioScience Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

/**
 * notification_imp.cpp
 *
 * Notification implementation
 */

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <signal.h>
#include <unistd.h>

#include <sys/socket.h>
#include <sys/ioctl.h>
#include <linux/if_packet.h>
#include <linux/if_ether.h>
#include <linux/if_arp.h>
#include <linux/filter.h>
#include <linux/if.h>
#include <arpa/inet.h>

#include "enumeration.h"
#include "notification_imp.h"

namespace avdecc_lib
{
    notification_imp *notification_imp_ref = new notification_imp();

    notification_imp::notification_imp()
    {
        notification_thread_init(); // Start notification thread
    }

    notification_imp::~notification_imp()
    {
        post_notification_event();
    }

    int notification_imp::notification_thread_init()
    {
        int rc;

        sem_init(&notify_waiting, 0, 0);

        rc = pthread_create(&h_thread, NULL, &notification_imp::dispatch_thread, (void *)this);
        if (rc)
        {
            printf("ERROR; return code from pthread_create() is %d\n", rc);
            exit(-1);
        }

        return 0;
    }

    void * notification_imp::dispatch_thread(void * param)
    {
        return ((notification_imp *)param)->dispatch_callbacks();

    }

    void * notification_imp::dispatch_callbacks(void)
    {
        while (true)
        {
            sem_wait(&notify_waiting);

            if((write_index - read_index) > 0)
            {
                notification_callback(user_obj,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].notification_type,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].guid,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].cmd_type,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].desc_type,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].desc_index,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].cmd_status,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].notification_id
                                     );
                read_index++;
            }
            else
            {
                break;
            }
        }

        return 0;
    }

    void notification_imp::post_notification_event()
    {
        sem_post(&notify_waiting);
    }
}
<commit_msg>forgot guid reference, renamed to entity_id<commit_after>/*
 * Licensed under the MIT License (MIT)
 *
 * Copyright (c) 2013 AudioScience Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

/**
 * notification_imp.cpp
 *
 * Notification implementation
 */

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <signal.h>
#include <unistd.h>

#include <sys/socket.h>
#include <sys/ioctl.h>
#include <linux/if_packet.h>
#include <linux/if_ether.h>
#include <linux/if_arp.h>
#include <linux/filter.h>
#include <linux/if.h>
#include <arpa/inet.h>

#include "enumeration.h"
#include "notification_imp.h"

namespace avdecc_lib
{
    notification_imp *notification_imp_ref = new notification_imp();

    notification_imp::notification_imp()
    {
        notification_thread_init(); // Start notification thread
    }

    notification_imp::~notification_imp()
    {
        post_notification_event();
    }

    int notification_imp::notification_thread_init()
    {
        int rc;

        sem_init(&notify_waiting, 0, 0);

        rc = pthread_create(&h_thread, NULL, &notification_imp::dispatch_thread, (void *)this);
        if (rc)
        {
            printf("ERROR; return code from pthread_create() is %d\n", rc);
            exit(-1);
        }

        return 0;
    }

    void * notification_imp::dispatch_thread(void * param)
    {
        return ((notification_imp *)param)->dispatch_callbacks();

    }

    void * notification_imp::dispatch_callbacks(void)
    {
        while (true)
        {
            sem_wait(&notify_waiting);

            if((write_index - read_index) > 0)
            {
                notification_callback(user_obj,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].notification_type,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].entity_id,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].cmd_type,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].desc_type,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].desc_index,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].cmd_status,
                                      notification_buf[read_index % NOTIFICATION_BUF_COUNT].notification_id
                                     );
                read_index++;
            }
            else
            {
                break;
            }
        }

        return 0;
    }

    void notification_imp::post_notification_event()
    {
        sem_post(&notify_waiting);
    }
}
<|endoftext|>


<commit_before>#include <iostream>
#include <cmath>
using namespace std;

#define PRIME 2017

template <unsigned int p>
class primefield{
	unsigned int a;
public:
	primefield<p>(unsigned int n):a(n%p){};
	primefield<p>& operator+= (const primefield<p>& q);
	primefield<p> operator* (const primefield<p>& q) const;
	primefield<p>& operator*= (const primefield<p>& q);
	primefield<p> operator^ (unsigned int n) const; // power
	primefield<p> operator/ (const primefield<p>& q) const;
	primefield<p> operator+ (const primefield<p>& q) const;
	primefield<p> operator- (const primefield<p>& q) const;
	operator unsigned int() const;
	void print(ostream& os) const;
};

template <unsigned int p>
primefield<p>& primefield<p>::operator+= (const primefield<p>& q){
	a += q.a;
	a %= p;
	return (*this);
}

template <unsigned int p>
primefield<p> primefield<p>::operator+ (const primefield<p>& q) const{
	primefield<p> r(*this);
	r += q;
	return r;
}

template <unsigned int p>
primefield<p> primefield<p>::operator- (const primefield<p>& q) const{
	primefield<p> r(a - q.a);
	return r;
}

template <unsigned int p>
primefield<p> primefield<p>::operator* (const primefield<p>& q) const {
	primefield<p> r(*this);
	r *= q;
	return r;
}

template <unsigned int p>
primefield<p>& primefield<p>::operator*= (const primefield<p>& q) {
	a *= q.a;
	a %= p;
	return (*this);
}

template <unsigned int p>
primefield<p> primefield<p>::operator^ (unsigned int n) const {
	primefield<p> r(1);
	for (size_t i = 0; i < n; i++)
	{
		r *= (*this);
	}
	return r;
}

template <unsigned int p>
primefield<p> primefield<p>::operator/ (const primefield<p>& q) const{
	return (*this)*(q ^ (p - 2));
}

template <unsigned int p>
primefield<p>::operator unsigned int() const{
	return static_cast<unsigned int>(a);
}

template <unsigned int p>
void primefield<p>::print(ostream& os) const{
	os << a;
}

template <unsigned int p>
ostream& operator<<(ostream& os, const primefield<p> q){
	q.print(os);
	return os;
}

template <unsigned int p>
primefield<p> power(const primefield<p>& q, unsigned int n){
	return q^n;
}

template <unsigned int p>
primefield<p> factorial(const primefield<p>& q){
	if (q == 1) return primefield<p>(1);
	return q*factorial(q - primefield<p>(1));
}

int main() {
	//primefield<PRIME> partial_sum(0);
	//for (unsigned int i = 1; i <= 50; ++i) {
	//	partial_sum += primefield<PRIME>(1) / primefield<PRIME>(i);
	//}
	//cout << "sum_i=1^50 1/i = " << partial_sum << endl;
	//unsigned int aa, bb;
	//cout << "a, b = "; cin >> aa >> bb;
	//primefield<PRIME> a(aa), b(bb);
	//cout << "a + b = " << a + b << endl;
	//cout << "a - b = " << a - b << endl;
	//cout << "a * b = " << a * b << endl;
	//cout << "a / b = " << a / b << endl;
	//cout << "a^10 = " << power(a, 10) << endl;
	//cout << "a! = " << factorial(a) << endl;
	cout << power(primefield<79>(53),10) << endl;
	cout << primefield<79>(3)/primefield<79>(21)+primefield<79>(5)/primefield<79>(67) << endl;
	cout << primefield<79>(3)/primefield<79>(21)-primefield<79>(5)/primefield<79>(67) << endl;
	cout << factorial(primefield<79>(50)) << endl;
	primefield<79> partial_sum(0);
	for (unsigned int i = 1; i <= 50; ++i) {
		partial_sum += primefield<79>(1) / primefield<79>(i);
	}
	cout << partial_sum << endl;
	return 0;
}
<commit_msg>fix negative modulo<commit_after>#include <iostream>
#include <cmath>
using namespace std;

#define PRIME 2017

template <unsigned int p>
class primefield{
	unsigned int a;
public:
	primefield<p>(int n) {
		a = n < 0 ? (n+p)%p : n%p;
	};
	primefield<p>& operator+= (const primefield<p>& q);
	primefield<p> operator* (const primefield<p>& q) const;
	primefield<p>& operator*= (const primefield<p>& q);
	primefield<p> operator^ (unsigned int n) const; // power
	primefield<p> operator/ (const primefield<p>& q) const;
	primefield<p> operator+ (const primefield<p>& q) const;
	primefield<p> operator- (const primefield<p>& q) const;
	operator unsigned int() const;
	void print(ostream& os) const;
};

template <unsigned int p>
primefield<p>& primefield<p>::operator+= (const primefield<p>& q){
	a += q.a;
	a %= p;
	return (*this);
}

template <unsigned int p>
primefield<p> primefield<p>::operator+ (const primefield<p>& q) const{
	primefield<p> r(*this);
	r += q;
	return r;
}

template <unsigned int p>
primefield<p> primefield<p>::operator- (const primefield<p>& q) const{
	primefield<p> r(a - q.a);
	return r;
}

template <unsigned int p>
primefield<p> primefield<p>::operator* (const primefield<p>& q) const {
	primefield<p> r(*this);
	r *= q;
	return r;
}

template <unsigned int p>
primefield<p>& primefield<p>::operator*= (const primefield<p>& q) {
	a *= q.a;
	a %= p;
	return (*this);
}

template <unsigned int p>
primefield<p> primefield<p>::operator^ (unsigned int n) const {
	primefield<p> r(1);
	for (size_t i = 0; i < n; i++)
	{
		r *= (*this);
	}
	return r;
}

template <unsigned int p>
primefield<p> primefield<p>::operator/ (const primefield<p>& q) const{
	return (*this)*(q ^ (p - 2));
}

template <unsigned int p>
primefield<p>::operator unsigned int() const{
	return static_cast<unsigned int>(a);
}

template <unsigned int p>
void primefield<p>::print(ostream& os) const{
	os << a;
}

template <unsigned int p>
ostream& operator<<(ostream& os, const primefield<p> q){
	q.print(os);
	return os;
}

template <unsigned int p>
primefield<p> power(const primefield<p>& q, unsigned int n){
	return q^n;
}

template <unsigned int p>
primefield<p> factorial(const primefield<p>& q){
	if (q == 1) return primefield<p>(1);
	return q*factorial(q - primefield<p>(1));
}

int main() {
	//primefield<PRIME> partial_sum(0);
	//for (unsigned int i = 1; i <= 50; ++i) {
	//	partial_sum += primefield<PRIME>(1) / primefield<PRIME>(i);
	//}
	//cout << "sum_i=1^50 1/i = " << partial_sum << endl;
	//unsigned int aa, bb;
	//cout << "a, b = "; cin >> aa >> bb;
	//primefield<PRIME> a(aa), b(bb);
	//cout << "a + b = " << a + b << endl;
	//cout << "a - b = " << a - b << endl;
	//cout << "a * b = " << a * b << endl;
	//cout << "a / b = " << a / b << endl;
	//cout << "a^10 = " << power(a, 10) << endl;
	//cout << "a! = " << factorial(a) << endl;
	cout << power(primefield<79>(53), 10) << endl;
	cout << primefield<79>(3) / primefield<79>(21) + primefield<79>(5) / primefield<79>(67) << endl;
	cout << primefield<79>(3) / primefield<79>(21) - primefield<79>(5) / primefield<79>(67) << endl;
	cout << factorial(primefield<79>(50)) << endl;
	primefield<79> partial_sum(0);
	for (unsigned int i = 1; i <= 50; ++i) {
		partial_sum += primefield<79>(1) / primefield<79>(i);
	}
	cout << partial_sum << endl;
	return 0;
}
<|endoftext|>


<commit_before>#ifndef COFFEE_MILL_DCD_CONVERT
#define COFFEE_MILL_DCD_CONVERT
#include <mill/dcd/DCDReader.hpp>
#include <mill/pdb/PDBReader.hpp>
#include <mill/pdb/PDBWriter.hpp>
#include <toml/toml.hpp>

namespace mill
{

inline const char* dcd_convert_usage() noexcept
{
    return "usage: mill dcd convert <format> <dcdfile> [reference(optional)]\n"
           "    $ mill dcd convert pdb traj.dcd [model.pdb(optional)]\n"
           "      convert traj.dcd to pdb format.\n"
           "      without pdb file, it consider all the particles as CA.\n"
           "      the pdb file must have the same number of particles with snapshot in dcd.\n";
}

// argv := arrayof{ "convert", "pdb", "filename", [pdb] }
template<typename vectorT>
int mode_dcd_convert(int argument_c, const char **argument_v)
{
    if(argument_c < 2)
    {
        log::error("mill dcd convert: too few arguments");
        log::error(dcd_convert_usage());
        return 1;
    }

    const std::string format(argument_v[1]);
    if(format == "help")
    {
        log::info(dcd_convert_usage());
        return 0;
    }
    else if(format != "pdb")
    {
        log::error("mill dcd convert: unknown format: ", format);
        log::error(dcd_convert_usage());
        return 1;
    }

    if(argument_c < 3)
    {
        log::error("mill dcd convert: too few arguments");
        log::error(dcd_convert_usage());
        return 1;
    }
    const std::string fname(argument_v[2]);

    if(fname.substr(fname.size()-4, 4) == ".dcd")
    {
        std::string pdbname;
        if(argument_c >= 4)
        {
            pdbname = std::string(argument_v[3]);
        }

        const std::string outname = fname.substr(0, fname.size()-4) +
                                    "_converted.pdb";
        std::ofstream ofs(outname);
        if(not ofs.good())
        {
            log::error("mill dcd convert: file open error: ", outname);
            log::error(dcd_convert_usage());
            return 1;
        }

        DCDReader reader(fname);
        PDBWriter<vectorT> writer;

        const auto header = reader.read_header();

        const std::size_t num_particles = header.at("nparticle").as_integer();
        std::vector<PDBAtom<vectorT>> atoms(num_particles);
        if(not pdbname.empty()) // reference pdb exists.
        {
            PDBReader<vectorT> pdbreader;
            std::ifstream pdbfile(pdbname);
            if(not pdbfile.good())
            {
                log::error("mill dcd convert: file open error: ", pdbname);
                log::error(dcd_convert_usage());
                return 1;
            }
            atoms = pdbreader.read(pdbname);
            if(atoms.size() != static_cast<std::size_t>(num_particles))
            {
                log::error("mill dcd convert: file open error: "
                           "pdb file may have different structure: ",
                           "num particle in dcd = ", num_particles,
                           "num particle in pdb = ", atoms.size());
                log::error(dcd_convert_usage());
                return 1;
            }
        }
        else
        {
            std::size_t idx = 1;
            for(auto& atom : atoms)
            {
                atom = make_default_atom(static_cast<int>(idx), vectorT(0,0,0));
                ++idx;
            }
        }

        std::size_t model = 0;
        for(const auto& frame : reader)
        {
            ++model;
            ofs << "MODEL     " << std::setw(4) << model << '\n';
            std::size_t idx = 0;
            for(const auto& p : frame)
            {
                atoms.at(idx).position = p.position();
                ++idx;
            }
            writer.write(ofs, atoms);
            ofs << "ENDMDL\n";
        }
        ofs << std::flush;
        ofs.close();
        return 0;
    }
    else
    {
        log::error("error: mill dcd convert: unknown file extension: ", fname);
        log::error(dcd_convert_usage());
        return 1;
    }
}

} // mill
#endif// COFFEE_MILL_DCD_CONVERT
<commit_msg>:recycle: use extension_of(filename)<commit_after>#ifndef COFFEE_MILL_DCD_CONVERT
#define COFFEE_MILL_DCD_CONVERT
#include <mill/dcd/DCDReader.hpp>
#include <mill/pdb/PDBReader.hpp>
#include <mill/pdb/PDBWriter.hpp>
#include <mill/util/file_extension.hpp>
#include <toml/toml.hpp>

namespace mill
{

inline const char* dcd_convert_usage() noexcept
{
    return "usage: mill dcd convert <format> <dcdfile> [reference(optional)]\n"
           "    $ mill dcd convert pdb traj.dcd [model.pdb(optional)]\n"
           "      convert traj.dcd to pdb format.\n"
           "      without pdb file, it consider all the particles as CA.\n"
           "      the pdb file must have the same number of particles with snapshot in dcd.\n";
}

// argv := arrayof{ "convert", "pdb", "filename", [pdb] }
template<typename vectorT>
int mode_dcd_convert(int argument_c, const char **argument_v)
{
    if(argument_c < 2)
    {
        log::error("mill dcd convert: too few arguments");
        log::error(dcd_convert_usage());
        return 1;
    }

    const std::string format(argument_v[1]);
    if(format == "help")
    {
        log::info(dcd_convert_usage());
        return 0;
    }
    else if(format != "pdb")
    {
        log::error("mill dcd convert: unknown format: ", format);
        log::error(dcd_convert_usage());
        return 1;
    }

    if(argument_c < 3)
    {
        log::error("mill dcd convert: too few arguments");
        log::error(dcd_convert_usage());
        return 1;
    }
    const std::string fname(argument_v[2]);

    if(extension_of(fname) == ".dcd")
    {
        std::string pdbname;
        if(argument_c >= 4)
        {
            pdbname = std::string(argument_v[3]);
        }

        const std::string outname = fname.substr(0, fname.size()-4) +
                                    "_converted.pdb";
        std::ofstream ofs(outname);
        if(not ofs.good())
        {
            log::error("mill dcd convert: file open error: ", outname);
            log::error(dcd_convert_usage());
            return 1;
        }

        DCDReader reader(fname);
        PDBWriter<vectorT> writer;

        const auto header = reader.read_header();

        const std::size_t num_particles = header.at("nparticle").as_integer();
        std::vector<PDBAtom<vectorT>> atoms(num_particles);
        if(not pdbname.empty()) // reference pdb exists.
        {
            PDBReader<vectorT> pdbreader;
            std::ifstream pdbfile(pdbname);
            if(not pdbfile.good())
            {
                log::error("mill dcd convert: file open error: ", pdbname);
                log::error(dcd_convert_usage());
                return 1;
            }
            atoms = pdbreader.read(pdbname);
            if(atoms.size() != static_cast<std::size_t>(num_particles))
            {
                log::error("mill dcd convert: file open error: "
                           "pdb file may have different structure: ",
                           "num particle in dcd = ", num_particles,
                           "num particle in pdb = ", atoms.size());
                log::error(dcd_convert_usage());
                return 1;
            }
        }
        else
        {
            std::size_t idx = 1;
            for(auto& atom : atoms)
            {
                atom = make_default_atom(static_cast<int>(idx), vectorT(0,0,0));
                ++idx;
            }
        }

        std::size_t model = 0;
        for(const auto& frame : reader)
        {
            ++model;
            ofs << "MODEL     " << std::setw(4) << model << '\n';
            std::size_t idx = 0;
            for(const auto& p : frame)
            {
                atoms.at(idx).position = p.position();
                ++idx;
            }
            writer.write(ofs, atoms);
            ofs << "ENDMDL\n";
        }
        ofs << std::flush;
        ofs.close();
        return 0;
    }
    else
    {
        log::error("error: mill dcd convert: unknown file extension: ", fname);
        log::error(dcd_convert_usage());
        return 1;
    }
}

} // mill
#endif// COFFEE_MILL_DCD_CONVERT
<|endoftext|>


<commit_before>#include <iostream>
#include <map>

#include <TROOT.h>
#include <TFile.h>
#include <TTree.h>
#include <TString.h>

#include "kTracker/SRawEvent.h"
#include "kTracker/SRecEvent.h"

#include "DataStruct.h"

using namespace std;

int main(int argc, char* argv[])
{
    // input data structure
    SRecEvent* recEvent = new SRecEvent;

    TFile* dataFile = new TFile(argv[1], "READ");
    TTree* dataTree = (TTree*)dataFile->Get(argv[3]);

    dataTree->SetBranchAddress("recEvent", &recEvent);

    //check if raw event exists
    bool mcEvent = false;
    bool dataEvent = false;
    SRawEvent* rawEvent;
    SRawMCEvent* rawMCEvent;
    if(dataTree->FindBranch("rawEvent") != NULL)
    {
        if(TString(dataTree->FindBranch("rawEvent")->GetClassName()) == "SRawMCEvent")
        {
            rawMCEvent = new SRawMCEvent;
            dataTree->SetBranchAddress("rawEvent", &rawMCEvent);
            mcEvent = true;
        }
        else
        {
            rawEvent = new SRawEvent;
            dataTree->SetBranchAddress("rawEvent", &rawEvent);
            dataEvent = true;
        }
    }

    // output data structure
    Dimuon* p_dimuon = new Dimuon; Dimuon& dimuon = *p_dimuon;
    Spill* p_spill = new Spill; Spill& spill = *p_spill;
    Event* p_event = new Event; Event& event = *p_event;
    Track* p_posTrack = new Track; Track& posTrack = *p_posTrack;
    Track* p_negTrack = new Track; Track& negTrack = *p_negTrack;

    TFile* saveFile = new TFile(argv[2], "recreate");
    TTree* saveTree = new TTree("save", "save");

    saveTree->Branch("dimuon", &p_dimuon, 256000, 99);
    saveTree->Branch("event", &p_event, 256000, 99);
    saveTree->Branch("spill", &p_spill, 256000, 99);
    saveTree->Branch("posTrack", &p_posTrack, 256000, 99);
    saveTree->Branch("negTrack", &p_negTrack, 256000, 99);

    //Initialize spill information accordingly
    spill.mcflag = mcEvent;
    map<int, Spill> spillBank;
    if(!mcEvent)
    {
        TFile* spillFile = new TFile(argv[4]);
        TTree* spillTree = (TTree*)spillFile->Get("save");

        spillTree->SetBranchAddress("spill", &p_spill);

        for(int i = 0; i < spillTree->GetEntries(); ++i)
        {
            spillTree->GetEntry(i);
            spillBank.insert(map<int, Spill>::value_type(spill.spillID, spill));
        }
    }

    int nDimuons = 0;
    double x_dummy, y_dummy, z_dummy;
    for(int i = 0; i < dataTree->GetEntries(); ++i)
    {
        dataTree->GetEntry(i);
        if(i % 1000 == 0)
        {
            cout << "Converting " << i << "/" << dataTree->GetEntries() << endl;
        }

        event.runID = recEvent->getRunID();
        event.spillID = recEvent->getSpillID();
        event.eventID = recEvent->getEventID();
        event.intensity = dataEvent ? rawEvent->getIntensity() : 0;
        if(mcEvent)
        {
            event.weight = rawMCEvent->weight;
            event.MATRIX1 = rawMCEvent->isEmuTriggered() ? 1 : -1;

            spill.spillID = recEvent->getSpillID();
            spill.targetPos = recEvent->getTargetPos();
            spill.TARGPOS_CONTROL = recEvent->getTargetPos();
        }
        else if(dataEvent)
        {
            event.MATRIX1 = recEvent->isTriggeredBy(SRawEvent::MATRIX1) ? 1 : -1;
            spill = spillBank[recEvent->getSpillID()];
        }
        else
        {
            event.MATRIX1 = 1;
        }

        for(int j = 0; j < recEvent->getNDimuons(); ++j)
        {
            ++nDimuons;
            SRecDimuon recDimuon = recEvent->getDimuon(j);

            dimuon.dimuonID = nDimuons;
            dimuon.posTrackID = recDimuon.trackID_pos;
            dimuon.negTrackID = recDimuon.trackID_neg;
            dimuon.chisq_dimuon = recDimuon.chisq_kf;
            dimuon.trackSeparation = recDimuon.vtx_pos.Z() - recDimuon.vtx_neg.Z();
            dimuon.mass = recDimuon.mass;
            dimuon.xF = recDimuon.xF;
            dimuon.x1 = recDimuon.x1;
            dimuon.x2 = recDimuon.x2;
            dimuon.pT = recDimuon.pT;

            dimuon.px1 = recDimuon.p_pos.Px();
            dimuon.py1 = recDimuon.p_pos.Py();
            dimuon.pz1 = recDimuon.p_pos.Pz();
            dimuon.px2 = recDimuon.p_neg.Px();
            dimuon.py2 = recDimuon.p_neg.Py();
            dimuon.pz2 = recDimuon.p_neg.Pz();

            dimuon.dx = recDimuon.vtx.X();
            dimuon.dy = recDimuon.vtx.Y();
            dimuon.dz = recDimuon.vtx.Z();
            dimuon.dpx = dimuon.px1 + dimuon.px2;
            dimuon.dpy = dimuon.py1 + dimuon.py2;
            dimuon.dpz = dimuon.pz1 + dimuon.pz2;

            if(!dimuon.goodDimuon()) continue;

            SRecTrack recPosTrack = recEvent->getTrack(dimuon.posTrackID);
            posTrack.trackID = recDimuon.trackID_pos;
            posTrack.triggerID = recPosTrack.getTriggerRoad();
            posTrack.charge = 1;
            posTrack.nHits = recPosTrack.getNHits();
            posTrack.nHitsSt1 = recPosTrack.getNHitsInStation(1);
            posTrack.nHitsSt2 = recPosTrack.getNHitsInStation(2);
            posTrack.nHitsSt3 = recPosTrack.getNHitsInStation(3);
            posTrack.nHitsSt4H = recPosTrack.getNHitsInPTY();
            posTrack.nHitsSt4V = recPosTrack.getNHitsInPTX();
            posTrack.chisq = recPosTrack.getChisq();
            posTrack.chisq_dump = recPosTrack.getChisqDump();
            posTrack.chisq_target = recPosTrack.getChisqTarget();
            posTrack.chisq_upstream = recPosTrack.getChisqUpstream();
            posTrack.px0 = recDimuon.p_pos.Px();
            posTrack.py0 = recDimuon.p_pos.Py();
            posTrack.pz0 = recDimuon.p_pos.Pz();
            posTrack.x_vertex = recPosTrack.getVtxPar(0);
            posTrack.y_vertex = recPosTrack.getVtxPar(1);
            posTrack.z_vertex = recPosTrack.getVtxPar(2);
            posTrack.x_target = recPosTrack.getTargetPos().X();
            posTrack.y_target = recPosTrack.getTargetPos().Y();
            posTrack.z_target = recPosTrack.getTargetPos().Z();
            posTrack.x_dump = recPosTrack.getDumpPos().X();
            posTrack.y_dump = recPosTrack.getDumpPos().Y();
            posTrack.z_dump = recPosTrack.getDumpPos().Z();
            recPosTrack.getMomentumVertex(x_dummy, y_dummy, z_dummy);
            posTrack.px_vertex = x_dummy;
            posTrack.py_vertex = y_dummy;
            posTrack.pz_vertex = z_dummy;
            recPosTrack.getExpMomentumFast(650., x_dummy, y_dummy, z_dummy);
            posTrack.px_st1 = x_dummy;
            posTrack.py_st1 = y_dummy;
            posTrack.pz_st1 = z_dummy;
            recPosTrack.getExpMomentumFast(1900., x_dummy, y_dummy, z_dummy);
            posTrack.px_st3 = x_dummy;
            posTrack.py_st3 = y_dummy;
            posTrack.pz_st3 = z_dummy;
            recPosTrack.getExpPositionFast(650., x_dummy, y_dummy);
            posTrack.x_st1 = x_dummy;
            posTrack.y_st1 = y_dummy;
            posTrack.z_st1 = 650.;
            recPosTrack.getExpPositionFast(1900., x_dummy, y_dummy);
            posTrack.x_st3 = x_dummy;
            posTrack.y_st3 = y_dummy;
            posTrack.z_st1 = 1900.;

            SRecTrack recNegTrack = recEvent->getTrack(dimuon.negTrackID);
            negTrack.trackID = recDimuon.trackID_neg;
            negTrack.triggerID = recNegTrack.getTriggerRoad();
            negTrack.charge = -1;
            negTrack.nHits = recNegTrack.getNHits();
            negTrack.nHitsSt1 = recNegTrack.getNHitsInStation(1);
            negTrack.nHitsSt2 = recNegTrack.getNHitsInStation(2);
            negTrack.nHitsSt3 = recNegTrack.getNHitsInStation(3);
            negTrack.nHitsSt4H = recNegTrack.getNHitsInPTY();
            negTrack.nHitsSt4V = recNegTrack.getNHitsInPTX();
            negTrack.chisq = recNegTrack.getChisq();
            negTrack.chisq_dump = recNegTrack.getChisqDump();
            negTrack.chisq_target = recNegTrack.getChisqTarget();
            negTrack.chisq_upstream = recNegTrack.getChisqUpstream();
            negTrack.px0 = recDimuon.p_neg.Px();
            negTrack.py0 = recDimuon.p_neg.Py();
            negTrack.pz0 = recDimuon.p_neg.Pz();
            negTrack.x_vertex = recNegTrack.getVtxPar(0);
            negTrack.y_vertex = recNegTrack.getVtxPar(1);
            negTrack.z_vertex = recNegTrack.getVtxPar(2);
            negTrack.x_target = recNegTrack.getTargetPos().X();
            negTrack.y_target = recNegTrack.getTargetPos().Y();
            negTrack.z_target = recNegTrack.getTargetPos().Z();
            negTrack.x_dump = recNegTrack.getDumpPos().X();
            negTrack.y_dump = recNegTrack.getDumpPos().Y();
            negTrack.z_dump = recNegTrack.getDumpPos().Z();
            recNegTrack.getMomentumVertex(x_dummy, y_dummy, z_dummy);
            negTrack.px_vertex = x_dummy;
            negTrack.py_vertex = y_dummy;
            negTrack.pz_vertex = z_dummy;
            recNegTrack.getExpMomentumFast(650., x_dummy, y_dummy, z_dummy);
            negTrack.px_st1 = x_dummy;
            negTrack.py_st1 = y_dummy;
            negTrack.pz_st1 = z_dummy;
            recNegTrack.getExpMomentumFast(1900., x_dummy, y_dummy, z_dummy);
            negTrack.px_st3 = x_dummy;
            negTrack.py_st3 = y_dummy;
            negTrack.pz_st3 = z_dummy;
            recNegTrack.getExpPositionFast(650., x_dummy, y_dummy);
            negTrack.x_st1 = x_dummy;
            negTrack.y_st1 = y_dummy;
            negTrack.z_st1 = 650.;
            recNegTrack.getExpPositionFast(1900., x_dummy, y_dummy);
            negTrack.x_st3 = x_dummy;
            negTrack.y_st3 = y_dummy;
            negTrack.z_st3 = 1900.;

            saveTree->Fill();
        }

        recEvent->clear();
    }

    saveFile->cd();
    saveTree->Write();
    saveFile->Close();

    return EXIT_SUCCESS;
}
<commit_msg>kTrackerResReader will confirm the spillID exists before read the event, if not it will quit the event and alert<commit_after>#include <iostream>
#include <map>

#include <TROOT.h>
#include <TFile.h>
#include <TTree.h>
#include <TString.h>

#include "kTracker/SRawEvent.h"
#include "kTracker/SRecEvent.h"

#include "DataStruct.h"

using namespace std;

int main(int argc, char* argv[])
{
    // input data structure
    SRecEvent* recEvent = new SRecEvent;

    TFile* dataFile = new TFile(argv[1], "READ");
    TTree* dataTree = (TTree*)dataFile->Get(argv[3]);

    dataTree->SetBranchAddress("recEvent", &recEvent);

    //check if raw event exists
    bool mcEvent = false;
    bool dataEvent = false;
    SRawEvent* rawEvent;
    SRawMCEvent* rawMCEvent;
    if(dataTree->FindBranch("rawEvent") != NULL)
    {
        if(TString(dataTree->FindBranch("rawEvent")->GetClassName()) == "SRawMCEvent")
        {
            rawMCEvent = new SRawMCEvent;
            dataTree->SetBranchAddress("rawEvent", &rawMCEvent);
            mcEvent = true;
        }
        else
        {
            rawEvent = new SRawEvent;
            dataTree->SetBranchAddress("rawEvent", &rawEvent);
            dataEvent = true;
        }
    }

    // output data structure
    Dimuon* p_dimuon = new Dimuon; Dimuon& dimuon = *p_dimuon;
    Spill* p_spill = new Spill; Spill& spill = *p_spill;
    Event* p_event = new Event; Event& event = *p_event;
    Track* p_posTrack = new Track; Track& posTrack = *p_posTrack;
    Track* p_negTrack = new Track; Track& negTrack = *p_negTrack;

    TFile* saveFile = new TFile(argv[2], "recreate");
    TTree* saveTree = new TTree("save", "save");

    saveTree->Branch("dimuon", &p_dimuon, 256000, 99);
    saveTree->Branch("event", &p_event, 256000, 99);
    saveTree->Branch("spill", &p_spill, 256000, 99);
    saveTree->Branch("posTrack", &p_posTrack, 256000, 99);
    saveTree->Branch("negTrack", &p_negTrack, 256000, 99);

    //Initialize spill information accordingly
    spill.mcflag = mcEvent;
    map<int, Spill> spillBank;
    if(!mcEvent)
    {
        TFile* spillFile = new TFile(argv[4]);
        TTree* spillTree = (TTree*)spillFile->Get("save");

        spillTree->SetBranchAddress("spill", &p_spill);

        for(int i = 0; i < spillTree->GetEntries(); ++i)
        {
            spillTree->GetEntry(i);
            spillBank.insert(map<int, Spill>::value_type(spill.spillID, spill));
        }
    }

    int nDimuons = 0;
    double x_dummy, y_dummy, z_dummy;
    for(int i = 0; i < dataTree->GetEntries(); ++i)
    {
        dataTree->GetEntry(i);
        if(i % 1000 == 0)
        {
            cout << "Converting " << i << "/" << dataTree->GetEntries() << endl;
        }

        event.runID = recEvent->getRunID();
        event.spillID = recEvent->getSpillID();
        event.eventID = recEvent->getEventID();
        event.intensity = dataEvent ? rawEvent->getIntensity() : 0;
        if(mcEvent)
        {
            event.weight = rawMCEvent->weight;
            event.MATRIX1 = rawMCEvent->isEmuTriggered() ? 1 : -1;

            spill.spillID = recEvent->getSpillID();
            spill.targetPos = recEvent->getTargetPos();
            spill.TARGPOS_CONTROL = recEvent->getTargetPos();
        }
        else if(dataEvent)
        {
            event.MATRIX1 = recEvent->isTriggeredBy(SRawEvent::MATRIX1) ? 1 : -1;

            if(spillBank.find(recEvent->getSpillID()) == spillBank.end())
            {
                event.log("spill ID doesn't exist!");
                continue;
            }
            spill = spillBank[recEvent->getSpillID()];
        }
        else
        {
            event.MATRIX1 = 1;
        }

        for(int j = 0; j < recEvent->getNDimuons(); ++j)
        {
            ++nDimuons;
            SRecDimuon recDimuon = recEvent->getDimuon(j);

            dimuon.dimuonID = nDimuons;
            dimuon.posTrackID = recDimuon.trackID_pos;
            dimuon.negTrackID = recDimuon.trackID_neg;
            dimuon.chisq_dimuon = recDimuon.chisq_kf;
            dimuon.trackSeparation = recDimuon.vtx_pos.Z() - recDimuon.vtx_neg.Z();
            dimuon.mass = recDimuon.mass;
            dimuon.xF = recDimuon.xF;
            dimuon.x1 = recDimuon.x1;
            dimuon.x2 = recDimuon.x2;
            dimuon.pT = recDimuon.pT;

            dimuon.px1 = recDimuon.p_pos.Px();
            dimuon.py1 = recDimuon.p_pos.Py();
            dimuon.pz1 = recDimuon.p_pos.Pz();
            dimuon.px2 = recDimuon.p_neg.Px();
            dimuon.py2 = recDimuon.p_neg.Py();
            dimuon.pz2 = recDimuon.p_neg.Pz();

            dimuon.dx = recDimuon.vtx.X();
            dimuon.dy = recDimuon.vtx.Y();
            dimuon.dz = recDimuon.vtx.Z();
            dimuon.dpx = dimuon.px1 + dimuon.px2;
            dimuon.dpy = dimuon.py1 + dimuon.py2;
            dimuon.dpz = dimuon.pz1 + dimuon.pz2;

            if(!dimuon.goodDimuon()) continue;

            SRecTrack recPosTrack = recEvent->getTrack(dimuon.posTrackID);
            posTrack.trackID = recDimuon.trackID_pos;
            posTrack.triggerID = recPosTrack.getTriggerRoad();
            posTrack.charge = 1;
            posTrack.nHits = recPosTrack.getNHits();
            posTrack.nHitsSt1 = recPosTrack.getNHitsInStation(1);
            posTrack.nHitsSt2 = recPosTrack.getNHitsInStation(2);
            posTrack.nHitsSt3 = recPosTrack.getNHitsInStation(3);
            posTrack.nHitsSt4H = recPosTrack.getNHitsInPTY();
            posTrack.nHitsSt4V = recPosTrack.getNHitsInPTX();
            posTrack.chisq = recPosTrack.getChisq();
            posTrack.chisq_dump = recPosTrack.getChisqDump();
            posTrack.chisq_target = recPosTrack.getChisqTarget();
            posTrack.chisq_upstream = recPosTrack.getChisqUpstream();
            posTrack.px0 = recDimuon.p_pos.Px();
            posTrack.py0 = recDimuon.p_pos.Py();
            posTrack.pz0 = recDimuon.p_pos.Pz();
            posTrack.x_vertex = recPosTrack.getVtxPar(0);
            posTrack.y_vertex = recPosTrack.getVtxPar(1);
            posTrack.z_vertex = recPosTrack.getVtxPar(2);
            posTrack.x_target = recPosTrack.getTargetPos().X();
            posTrack.y_target = recPosTrack.getTargetPos().Y();
            posTrack.z_target = recPosTrack.getTargetPos().Z();
            posTrack.x_dump = recPosTrack.getDumpPos().X();
            posTrack.y_dump = recPosTrack.getDumpPos().Y();
            posTrack.z_dump = recPosTrack.getDumpPos().Z();
            recPosTrack.getMomentumVertex(x_dummy, y_dummy, z_dummy);
            posTrack.px_vertex = x_dummy;
            posTrack.py_vertex = y_dummy;
            posTrack.pz_vertex = z_dummy;
            recPosTrack.getExpMomentumFast(650., x_dummy, y_dummy, z_dummy);
            posTrack.px_st1 = x_dummy;
            posTrack.py_st1 = y_dummy;
            posTrack.pz_st1 = z_dummy;
            recPosTrack.getExpMomentumFast(1900., x_dummy, y_dummy, z_dummy);
            posTrack.px_st3 = x_dummy;
            posTrack.py_st3 = y_dummy;
            posTrack.pz_st3 = z_dummy;
            recPosTrack.getExpPositionFast(650., x_dummy, y_dummy);
            posTrack.x_st1 = x_dummy;
            posTrack.y_st1 = y_dummy;
            posTrack.z_st1 = 650.;
            recPosTrack.getExpPositionFast(1900., x_dummy, y_dummy);
            posTrack.x_st3 = x_dummy;
            posTrack.y_st3 = y_dummy;
            posTrack.z_st1 = 1900.;

            SRecTrack recNegTrack = recEvent->getTrack(dimuon.negTrackID);
            negTrack.trackID = recDimuon.trackID_neg;
            negTrack.triggerID = recNegTrack.getTriggerRoad();
            negTrack.charge = -1;
            negTrack.nHits = recNegTrack.getNHits();
            negTrack.nHitsSt1 = recNegTrack.getNHitsInStation(1);
            negTrack.nHitsSt2 = recNegTrack.getNHitsInStation(2);
            negTrack.nHitsSt3 = recNegTrack.getNHitsInStation(3);
            negTrack.nHitsSt4H = recNegTrack.getNHitsInPTY();
            negTrack.nHitsSt4V = recNegTrack.getNHitsInPTX();
            negTrack.chisq = recNegTrack.getChisq();
            negTrack.chisq_dump = recNegTrack.getChisqDump();
            negTrack.chisq_target = recNegTrack.getChisqTarget();
            negTrack.chisq_upstream = recNegTrack.getChisqUpstream();
            negTrack.px0 = recDimuon.p_neg.Px();
            negTrack.py0 = recDimuon.p_neg.Py();
            negTrack.pz0 = recDimuon.p_neg.Pz();
            negTrack.x_vertex = recNegTrack.getVtxPar(0);
            negTrack.y_vertex = recNegTrack.getVtxPar(1);
            negTrack.z_vertex = recNegTrack.getVtxPar(2);
            negTrack.x_target = recNegTrack.getTargetPos().X();
            negTrack.y_target = recNegTrack.getTargetPos().Y();
            negTrack.z_target = recNegTrack.getTargetPos().Z();
            negTrack.x_dump = recNegTrack.getDumpPos().X();
            negTrack.y_dump = recNegTrack.getDumpPos().Y();
            negTrack.z_dump = recNegTrack.getDumpPos().Z();
            recNegTrack.getMomentumVertex(x_dummy, y_dummy, z_dummy);
            negTrack.px_vertex = x_dummy;
            negTrack.py_vertex = y_dummy;
            negTrack.pz_vertex = z_dummy;
            recNegTrack.getExpMomentumFast(650., x_dummy, y_dummy, z_dummy);
            negTrack.px_st1 = x_dummy;
            negTrack.py_st1 = y_dummy;
            negTrack.pz_st1 = z_dummy;
            recNegTrack.getExpMomentumFast(1900., x_dummy, y_dummy, z_dummy);
            negTrack.px_st3 = x_dummy;
            negTrack.py_st3 = y_dummy;
            negTrack.pz_st3 = z_dummy;
            recNegTrack.getExpPositionFast(650., x_dummy, y_dummy);
            negTrack.x_st1 = x_dummy;
            negTrack.y_st1 = y_dummy;
            negTrack.z_st1 = 650.;
            recNegTrack.getExpPositionFast(1900., x_dummy, y_dummy);
            negTrack.x_st3 = x_dummy;
            negTrack.y_st3 = y_dummy;
            negTrack.z_st3 = 1900.;

            saveTree->Fill();
        }

        recEvent->clear();
    }

    saveFile->cd();
    saveTree->Write();
    saveFile->Close();

    return EXIT_SUCCESS;
}
<|endoftext|>


<commit_before>#include <chrono>
#include <glm/gtx/io.hpp>
#include <iostream>
#include <thread>

#include <mos/aud/renderer.hpp>
#include <mos/aud/sound_stream.hpp>
#include <mos/aud/sound.hpp>
#include <mos/aud/listener.hpp>
#include <mos/aud/scene.hpp>

namespace mos::aud {

Renderer::Renderer()
    : reverb_properties(EFX_REVERB_PRESET_LIVINGROOM) {
  static constexpr ALint sampling_frequency{44100};
  std::array context_attributes{ALC_FREQUENCY, sampling_frequency, ALC_HRTF_SOFT, ALC_TRUE, 0};
  device_ = alcOpenDevice(nullptr);
  context_ = alcCreateContext(device_, context_attributes.data());
  alcMakeContextCurrent(context_);

#ifdef MOS_EFX
  if (!alcIsExtensionPresent(alcGetContextsDevice(alcGetCurrentContext()),
                             "ALC_EXT_EFX")) {
    throw std::runtime_error("OpenAL EFX not supported.");
  }

  ALuint reverb_effect = AL_NONE;
  alGenEffects(1, &reverb_effect);
  alEffecti(reverb_effect, AL_EFFECT_TYPE, AL_EFFECT_REVERB);
  alEffectf(reverb_effect, AL_REVERB_DENSITY, reverb_properties.flDensity);
  alEffectf(reverb_effect, AL_REVERB_DIFFUSION, reverb_properties.flDiffusion);
  alEffectf(reverb_effect, AL_REVERB_GAIN, reverb_properties.flGain);
  alEffectf(reverb_effect, AL_REVERB_GAINHF, reverb_properties.flGainHF);
  alEffectf(reverb_effect, AL_REVERB_DECAY_TIME, reverb_properties.flDecayTime);
  alEffectf(reverb_effect, AL_REVERB_DECAY_HFRATIO,
            reverb_properties.flDecayHFRatio);
  alEffectf(reverb_effect, AL_REVERB_REFLECTIONS_GAIN,
            reverb_properties.flReflectionsGain);
  alEffectf(reverb_effect, AL_REVERB_REFLECTIONS_DELAY,
            reverb_properties.flReflectionsDelay);
  alEffectf(reverb_effect, AL_REVERB_LATE_REVERB_GAIN,
            reverb_properties.flLateReverbGain);
  alEffectf(reverb_effect, AL_REVERB_LATE_REVERB_DELAY,
            reverb_properties.flLateReverbDelay);
  alEffectf(reverb_effect, AL_REVERB_AIR_ABSORPTION_GAINHF,
            reverb_properties.flAirAbsorptionGainHF);
  alEffectf(reverb_effect, AL_REVERB_ROOM_ROLLOFF_FACTOR,
            reverb_properties.flRoomRolloffFactor);
  alEffecti(reverb_effect, AL_REVERB_DECAY_HFLIMIT,
            reverb_properties.iDecayHFLimit);

  if (reverb_effect == AL_NONE) {
    throw std::runtime_error("Could not create reverb effect.");
  }

  alGenAuxiliaryEffectSlots(1, &reverb_slot);
  if (reverb_slot == AL_NONE) {
    throw std::runtime_error("Could not create reverb effect slot.");
  }

  alAuxiliaryEffectSloti(reverb_slot, AL_EFFECTSLOT_EFFECT, reverb_effect);

  alGenFilters(1, &lowpass_filter1);
  if (lowpass_filter1 == AL_NONE) {
    std::runtime_error("Could not create lowpass filter.");
  }
  alFilteri(lowpass_filter1, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
  alFilterf(lowpass_filter1, AL_LOWPASS_GAIN, 0.6f);
  alFilterf(lowpass_filter1, AL_LOWPASS_GAINHF, 0.05f);

  alGenFilters(1, &lowpass_filter2);
  if (lowpass_filter2 == AL_NONE) {
    std::runtime_error("Could not create lowpass filter.");
  }
  alFilteri(lowpass_filter2, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
  alFilterf(lowpass_filter2, AL_LOWPASS_GAIN, 0.3f);
  alFilterf(lowpass_filter2, AL_LOWPASS_GAINHF, 0.01f);
#endif

  listener(Listener());
}

Renderer::~Renderer() {
  for (auto source : sources_) {
    alSourceStop(source.second);
    alDeleteSources(1, &source.second);
  }
  for (auto buffer : buffers_) {
    alDeleteBuffers(1, &buffer.second);
  }
  alcMakeContextCurrent(nullptr);
  alcDestroyContext(context_);
  alcCloseDevice(device_);
}

void Renderer::render_sound(const Sound &sound, const float dt) {
  if (sources_.find(sound.source.id()) == sources_.end()) {
    ALuint al_source{AL_NONE};
    alGenSources(1, &al_source);
    sources_.insert(SourcePair(sound.source.id(), al_source));

#ifdef MOS_EFX
    alSource3i(al_source, AL_AUXILIARY_SEND_FILTER, reverb_slot, 0,
               AL_FILTER_NULL);
    ALuint al_filter{AL_NONE};
    alGenFilters(1, &al_filter);
    filters_.insert(SourcePair(sound.source.id(), al_filter));
    alFilteri(al_filter, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
    alSourcei(al_source, AL_DIRECT_FILTER, al_filter);
#endif
  }

  ALuint al_source = sources_.at(sound.source.id());

  auto buffer = sound.buffer;
  if (buffer) {
    auto format =
        buffer->channels() == 1 ? AL_FORMAT_MONO16 : AL_FORMAT_STEREO16;
    if (buffers_.find(buffer->id()) == buffers_.end()) {
        ALuint al_buffer{0u};
      alGenBuffers(1, &al_buffer);
      {
        long data_size = std::distance(buffer->begin(), buffer->end());
        const ALvoid *data = buffer->data();
        alBufferData(al_buffer, format, data, data_size * sizeof(short),
                     buffer->sample_rate());
      }
      buffers_.insert(BufferPair(buffer->id(), al_buffer));
    }

    ALuint al_buffer = buffers_.at(sound.buffer->id());

    int v{0u};
    alGetSourcei(al_source, AL_BUFFER, &v);
    if (v == 0) {
      alSourcei(al_source, AL_BUFFER, al_buffer);
    }
  }

  alSourcei(al_source, AL_LOOPING, sound.source.loop);
  alSourcef(al_source, AL_PITCH, sound.source.pitch);
  //alSourcef(al_source, AL_GAIN, buffer_source.source.gain);
  alSource3f(al_source, AL_POSITION, sound.source.position.x,
             sound.source.position.y, sound.source.position.z);
  alSource3f(al_source, AL_VELOCITY, sound.source.velocity.x,
             sound.source.velocity.y, sound.source.velocity.z);

#ifdef MOS_EFX
  auto al_filter = filters_[sound.source.id()];
  ALfloat al_gain{0.0f};
  alGetFilterf(al_filter, AL_LOWPASS_GAIN, &al_gain);
  auto error = glm::clamp((1.0f - sound.source.obstructed), 0.0f, 1.0f) - al_gain;
  float gain = al_gain + error * dt;

  ALfloat al_gain_hf{0.0f};
  alGetFilterf(al_filter, AL_LOWPASS_GAINHF, &al_gain_hf);
  auto error_hf = glm::clamp((1.0f - sound.source.obstructed), 0.0f, 1.0f) - al_gain_hf;
  float gain_hf = al_gain_hf + error_hf * dt;

  alSourcef(al_source, AL_GAIN, gain * sound.source.gain);

  alFilteri(al_filter, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
  alFilterf(al_filter, AL_LOWPASS_GAIN, gain);      // 0.5f
  alFilterf(al_filter, AL_LOWPASS_GAINHF, gain_hf); // 0.01f
  alSourcei(al_source, AL_DIRECT_FILTER, al_filter);
#endif

  ALenum state{AL_STOPPED};
  alGetSourcei(al_source, AL_SOURCE_STATE, &state);

  if (sound.source.playing && (state != AL_PLAYING)) {
    alSourcePlay(al_source);
  }

  if (!sound.source.playing && (state == AL_PLAYING)) {
    alSourceStop(al_source);
  }

  if (state == AL_STOPPED) {
    alSourceRewind(al_source);
  }
}

void Renderer::render_sound_stream(const Sound_stream &sound_stream, const float dt) {
  if (sources_.find(sound_stream.source.id()) == sources_.end()) {
    ALuint al_source{0u};
    alGenSources(1, &al_source);
    sources_.insert(SourcePair(sound_stream.source.id(), al_source));

#ifdef MOS_EFX
    alSource3i(al_source, AL_AUXILIARY_SEND_FILTER, reverb_slot, 0,
               AL_FILTER_NULL);
    ALuint al_filter{0u};
    alGenFilters(1, &al_filter);
    filters_.insert(SourcePair(sound_stream.source.id(), al_filter));
    alFilteri(al_filter, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
    alSourcei(al_source, AL_DIRECT_FILTER, al_filter);
#endif
  };

  ALuint al_source = sources_.at(sound_stream.source.id());
  alSourcei(al_source, AL_LOOPING, sound_stream.source.loop);
  alSourcef(al_source, AL_PITCH, sound_stream.source.pitch);
  //alSourcef(al_source, AL_GAIN, stream_source.source.gain);
  alSource3f(al_source, AL_POSITION, sound_stream.source.position.x,
             sound_stream.source.position.y, sound_stream.source.position.z);
  alSource3f(al_source, AL_VELOCITY, sound_stream.source.velocity.x,
             sound_stream.source.velocity.y, sound_stream.source.velocity.z);

#ifdef MOS_EFX
  auto al_filter = filters_[sound_stream.source.id()];
  ALfloat al_gain{0.0f};
  alGetFilterf(al_filter, AL_LOWPASS_GAIN, &al_gain);
  auto error = glm::clamp((1.0f - sound_stream.source.obstructed), 0.0f, 1.0f) - al_gain;
  float gain = al_gain + error * dt;

  ALfloat al_gain_hf{0.0f};
  alGetFilterf(al_filter, AL_LOWPASS_GAINHF, &al_gain_hf);
  auto error_hf = glm::clamp((1.0f - sound_stream.source.obstructed), 0.0f, 1.0f) - al_gain_hf;
  float gain_hf = al_gain_hf + error_hf * dt;

  alSourcef(al_source, AL_GAIN, gain * sound_stream.source.gain);

  alFilteri(al_filter, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
  alFilterf(al_filter, AL_LOWPASS_GAIN , gain);      // 0.5f
  alFilterf(al_filter, AL_LOWPASS_GAINHF, gain_hf); // 0.01f
  alSourcei(al_source, AL_DIRECT_FILTER, al_filter);
#endif

  ALenum state{AL_STOPPED};
  alGetSourcei(al_source, AL_SOURCE_STATE, &state);

  ALint processed = 0;
  alGetSourcei(al_source, AL_BUFFERS_PROCESSED, &processed);
  ALuint buffer = 0;

  if(sound_stream.stream) {
    auto format = sound_stream.stream->channels() == 1 ? AL_FORMAT_MONO16 : AL_FORMAT_STEREO16;

    std::array<ALuint,4> buffers{0,0,0,0};
    if (sound_stream.source.playing && (state != AL_PLAYING)) {
      alGenBuffers(4, buffers.data());
      const auto samples = sound_stream.stream->read();
      const int size = sound_stream.stream->buffer_size;
      for (const auto & buffer : buffers) {
        alBufferData(
            buffer, format, samples.data(),
            size * sizeof(ALshort), sound_stream.stream->sample_rate());
        alSourceQueueBuffers(al_source, 1, &buffer);
      }
      alSourcePlay(al_source);
    }

    while (processed--) {
		  alSourceUnqueueBuffers(al_source, 1, &buffer);
                  const auto samples = sound_stream.stream->read();
                  const int size = sound_stream.stream->buffer_size;
			if (buffer != 0) {
	  			alBufferData(buffer, format, samples.data(),
					size * sizeof(ALshort),
					sound_stream.stream->sample_rate());
				alSourceQueueBuffers(al_source, 1, &buffer);
			}
    }
    if (sound_stream.source.loop && sound_stream.stream->done()) {
      sound_stream.stream->seek_start();
    }
    if (!sound_stream.source.playing && (state == AL_PLAYING)) {
      alSourceStop(al_source);
      ALint count{0};
      alGetSourcei(al_source, AL_BUFFERS_QUEUED, &count);
      alSourceUnqueueBuffers(al_source, count, &buffer);
      alDeleteBuffers(4, buffers.data());
    }
  }
}

auto Renderer::listener() const -> Listener {
  Listener listener;
  alGetListener3f(AL_POSITION, &listener.position.x, &listener.position.y,
                  &listener.position.z);

  std::array<float,6> orientation{};
  alGetListenerf(AL_ORIENTATION, orientation.data());
  listener.direction.x = orientation[0];
  listener.direction.y = orientation[1];
  listener.direction.z = orientation[2];
  listener.up.x = orientation[3];
  listener.up.y = orientation[4];
  listener.up.z = orientation[5];

  alGetListener3f(AL_VELOCITY, &listener.velocity.x, &listener.velocity.y,
                  &listener.velocity.z);

  return listener;
}

void Renderer::listener(const Listener &listener) {
  alListener3f(AL_POSITION, listener.position.x, listener.position.y,
               listener.position.z);

  alListener3f(AL_VELOCITY, listener.velocity.x, listener.velocity.y,
               listener.velocity.z);

  std::array<float, 6> orientation {listener.direction.x, listener.direction.y,
                          listener.direction.z, listener.up.x,
                          listener.up.y, listener.up.z};
  alListenerfv(AL_ORIENTATION, orientation.data());
  alListenerf(AL_GAIN, listener.gain);
}

auto Renderer::render(const Scene &scene, const float dt) -> void {
  listener(scene.listener);
  for (const auto &sound : scene.sounds) {
    render_sound(sound, dt);
  }
  for (const auto &sound_stream : scene.sound_streams) {
    render_sound_stream(sound_stream, dt);
  }
}

void Renderer::clear() {
  for (auto source : sources_) {
    alSourceStop(source.second);
    alDeleteSources(1, &source.second);
  }
  for (auto buffer : buffers_) {
    alDeleteBuffers(1, &buffer.second);
  }
  sources_.clear();
  buffers_.clear();
}
}
<commit_msg>Check for device<commit_after>#include <chrono>
#include <glm/gtx/io.hpp>
#include <iostream>
#include <thread>

#include <mos/aud/renderer.hpp>
#include <mos/aud/sound_stream.hpp>
#include <mos/aud/sound.hpp>
#include <mos/aud/listener.hpp>
#include <mos/aud/scene.hpp>

namespace mos::aud {

Renderer::Renderer()
    : reverb_properties(EFX_REVERB_PRESET_LIVINGROOM) {
  static constexpr ALint sampling_frequency{44100};
  std::array context_attributes{ALC_FREQUENCY, sampling_frequency, ALC_HRTF_SOFT, ALC_TRUE, 0};
  device_ = alcOpenDevice(nullptr);
  if (device_ == AL_NONE) {
    throw std::runtime_error("No valid audio device");
  }
  context_ = alcCreateContext(device_, context_attributes.data());
  alcMakeContextCurrent(context_);

#ifdef MOS_EFX
  if (!alcIsExtensionPresent(alcGetContextsDevice(alcGetCurrentContext()),
                             "ALC_EXT_EFX")) {
    throw std::runtime_error("OpenAL EFX not supported.");
  }

  ALuint reverb_effect = AL_NONE;
  alGenEffects(1, &reverb_effect);
  alEffecti(reverb_effect, AL_EFFECT_TYPE, AL_EFFECT_REVERB);
  alEffectf(reverb_effect, AL_REVERB_DENSITY, reverb_properties.flDensity);
  alEffectf(reverb_effect, AL_REVERB_DIFFUSION, reverb_properties.flDiffusion);
  alEffectf(reverb_effect, AL_REVERB_GAIN, reverb_properties.flGain);
  alEffectf(reverb_effect, AL_REVERB_GAINHF, reverb_properties.flGainHF);
  alEffectf(reverb_effect, AL_REVERB_DECAY_TIME, reverb_properties.flDecayTime);
  alEffectf(reverb_effect, AL_REVERB_DECAY_HFRATIO,
            reverb_properties.flDecayHFRatio);
  alEffectf(reverb_effect, AL_REVERB_REFLECTIONS_GAIN,
            reverb_properties.flReflectionsGain);
  alEffectf(reverb_effect, AL_REVERB_REFLECTIONS_DELAY,
            reverb_properties.flReflectionsDelay);
  alEffectf(reverb_effect, AL_REVERB_LATE_REVERB_GAIN,
            reverb_properties.flLateReverbGain);
  alEffectf(reverb_effect, AL_REVERB_LATE_REVERB_DELAY,
            reverb_properties.flLateReverbDelay);
  alEffectf(reverb_effect, AL_REVERB_AIR_ABSORPTION_GAINHF,
            reverb_properties.flAirAbsorptionGainHF);
  alEffectf(reverb_effect, AL_REVERB_ROOM_ROLLOFF_FACTOR,
            reverb_properties.flRoomRolloffFactor);
  alEffecti(reverb_effect, AL_REVERB_DECAY_HFLIMIT,
            reverb_properties.iDecayHFLimit);

  if (reverb_effect == AL_NONE) {
    throw std::runtime_error("Could not create reverb effect.");
  }

  alGenAuxiliaryEffectSlots(1, &reverb_slot);
  if (reverb_slot == AL_NONE) {
    throw std::runtime_error("Could not create reverb effect slot.");
  }

  alAuxiliaryEffectSloti(reverb_slot, AL_EFFECTSLOT_EFFECT, reverb_effect);

  alGenFilters(1, &lowpass_filter1);
  if (lowpass_filter1 == AL_NONE) {
    std::runtime_error("Could not create lowpass filter.");
  }
  alFilteri(lowpass_filter1, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
  alFilterf(lowpass_filter1, AL_LOWPASS_GAIN, 0.6f);
  alFilterf(lowpass_filter1, AL_LOWPASS_GAINHF, 0.05f);

  alGenFilters(1, &lowpass_filter2);
  if (lowpass_filter2 == AL_NONE) {
    std::runtime_error("Could not create lowpass filter.");
  }
  alFilteri(lowpass_filter2, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
  alFilterf(lowpass_filter2, AL_LOWPASS_GAIN, 0.3f);
  alFilterf(lowpass_filter2, AL_LOWPASS_GAINHF, 0.01f);
#endif

  listener(Listener());
}

Renderer::~Renderer() {
  for (auto source : sources_) {
    alSourceStop(source.second);
    alDeleteSources(1, &source.second);
  }
  for (auto buffer : buffers_) {
    alDeleteBuffers(1, &buffer.second);
  }
  alcMakeContextCurrent(nullptr);
  alcDestroyContext(context_);
  alcCloseDevice(device_);
}

void Renderer::render_sound(const Sound &sound, const float dt) {
  if (sources_.find(sound.source.id()) == sources_.end()) {
    ALuint al_source{AL_NONE};
    alGenSources(1, &al_source);
    sources_.insert(SourcePair(sound.source.id(), al_source));

#ifdef MOS_EFX
    alSource3i(al_source, AL_AUXILIARY_SEND_FILTER, reverb_slot, 0,
               AL_FILTER_NULL);
    ALuint al_filter{AL_NONE};
    alGenFilters(1, &al_filter);
    filters_.insert(SourcePair(sound.source.id(), al_filter));
    alFilteri(al_filter, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
    alSourcei(al_source, AL_DIRECT_FILTER, al_filter);
#endif
  }

  ALuint al_source = sources_.at(sound.source.id());

  auto buffer = sound.buffer;
  if (buffer) {
    auto format =
        buffer->channels() == 1 ? AL_FORMAT_MONO16 : AL_FORMAT_STEREO16;
    if (buffers_.find(buffer->id()) == buffers_.end()) {
        ALuint al_buffer{0u};
      alGenBuffers(1, &al_buffer);
      {
        long data_size = std::distance(buffer->begin(), buffer->end());
        const ALvoid *data = buffer->data();
        alBufferData(al_buffer, format, data, data_size * sizeof(short),
                     buffer->sample_rate());
      }
      buffers_.insert(BufferPair(buffer->id(), al_buffer));
    }

    ALuint al_buffer = buffers_.at(sound.buffer->id());

    int v{0u};
    alGetSourcei(al_source, AL_BUFFER, &v);
    if (v == 0) {
      alSourcei(al_source, AL_BUFFER, al_buffer);
    }
  }

  alSourcei(al_source, AL_LOOPING, sound.source.loop);
  alSourcef(al_source, AL_PITCH, sound.source.pitch);
  //alSourcef(al_source, AL_GAIN, buffer_source.source.gain);
  alSource3f(al_source, AL_POSITION, sound.source.position.x,
             sound.source.position.y, sound.source.position.z);
  alSource3f(al_source, AL_VELOCITY, sound.source.velocity.x,
             sound.source.velocity.y, sound.source.velocity.z);

#ifdef MOS_EFX
  auto al_filter = filters_[sound.source.id()];
  ALfloat al_gain{0.0f};
  alGetFilterf(al_filter, AL_LOWPASS_GAIN, &al_gain);
  auto error = glm::clamp((1.0f - sound.source.obstructed), 0.0f, 1.0f) - al_gain;
  float gain = al_gain + error * dt;

  ALfloat al_gain_hf{0.0f};
  alGetFilterf(al_filter, AL_LOWPASS_GAINHF, &al_gain_hf);
  auto error_hf = glm::clamp((1.0f - sound.source.obstructed), 0.0f, 1.0f) - al_gain_hf;
  float gain_hf = al_gain_hf + error_hf * dt;

  alSourcef(al_source, AL_GAIN, gain * sound.source.gain);

  alFilteri(al_filter, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
  alFilterf(al_filter, AL_LOWPASS_GAIN, gain);      // 0.5f
  alFilterf(al_filter, AL_LOWPASS_GAINHF, gain_hf); // 0.01f
  alSourcei(al_source, AL_DIRECT_FILTER, al_filter);
#endif

  ALenum state{AL_STOPPED};
  alGetSourcei(al_source, AL_SOURCE_STATE, &state);

  if (sound.source.playing && (state != AL_PLAYING)) {
    alSourcePlay(al_source);
  }

  if (!sound.source.playing && (state == AL_PLAYING)) {
    alSourceStop(al_source);
  }

  if (state == AL_STOPPED) {
    alSourceRewind(al_source);
  }
}

void Renderer::render_sound_stream(const Sound_stream &sound_stream, const float dt) {
  if (sources_.find(sound_stream.source.id()) == sources_.end()) {
    ALuint al_source{0u};
    alGenSources(1, &al_source);
    sources_.insert(SourcePair(sound_stream.source.id(), al_source));

#ifdef MOS_EFX
    alSource3i(al_source, AL_AUXILIARY_SEND_FILTER, reverb_slot, 0,
               AL_FILTER_NULL);
    ALuint al_filter{0u};
    alGenFilters(1, &al_filter);
    filters_.insert(SourcePair(sound_stream.source.id(), al_filter));
    alFilteri(al_filter, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
    alSourcei(al_source, AL_DIRECT_FILTER, al_filter);
#endif
  };

  ALuint al_source = sources_.at(sound_stream.source.id());
  alSourcei(al_source, AL_LOOPING, sound_stream.source.loop);
  alSourcef(al_source, AL_PITCH, sound_stream.source.pitch);
  //alSourcef(al_source, AL_GAIN, stream_source.source.gain);
  alSource3f(al_source, AL_POSITION, sound_stream.source.position.x,
             sound_stream.source.position.y, sound_stream.source.position.z);
  alSource3f(al_source, AL_VELOCITY, sound_stream.source.velocity.x,
             sound_stream.source.velocity.y, sound_stream.source.velocity.z);

#ifdef MOS_EFX
  auto al_filter = filters_[sound_stream.source.id()];
  ALfloat al_gain{0.0f};
  alGetFilterf(al_filter, AL_LOWPASS_GAIN, &al_gain);
  auto error = glm::clamp((1.0f - sound_stream.source.obstructed), 0.0f, 1.0f) - al_gain;
  float gain = al_gain + error * dt;

  ALfloat al_gain_hf{0.0f};
  alGetFilterf(al_filter, AL_LOWPASS_GAINHF, &al_gain_hf);
  auto error_hf = glm::clamp((1.0f - sound_stream.source.obstructed), 0.0f, 1.0f) - al_gain_hf;
  float gain_hf = al_gain_hf + error_hf * dt;

  alSourcef(al_source, AL_GAIN, gain * sound_stream.source.gain);

  alFilteri(al_filter, AL_FILTER_TYPE, AL_FILTER_LOWPASS);
  alFilterf(al_filter, AL_LOWPASS_GAIN , gain);      // 0.5f
  alFilterf(al_filter, AL_LOWPASS_GAINHF, gain_hf); // 0.01f
  alSourcei(al_source, AL_DIRECT_FILTER, al_filter);
#endif

  ALenum state{AL_STOPPED};
  alGetSourcei(al_source, AL_SOURCE_STATE, &state);

  ALint processed = 0;
  alGetSourcei(al_source, AL_BUFFERS_PROCESSED, &processed);
  ALuint buffer = 0;

  if(sound_stream.stream) {
    auto format = sound_stream.stream->channels() == 1 ? AL_FORMAT_MONO16 : AL_FORMAT_STEREO16;

    std::array<ALuint,4> buffers{0,0,0,0};
    if (sound_stream.source.playing && (state != AL_PLAYING)) {
      alGenBuffers(4, buffers.data());
      const auto samples = sound_stream.stream->read();
      const int size = sound_stream.stream->buffer_size;
      for (const auto & buffer : buffers) {
        alBufferData(
            buffer, format, samples.data(),
            size * sizeof(ALshort), sound_stream.stream->sample_rate());
        alSourceQueueBuffers(al_source, 1, &buffer);
      }
      alSourcePlay(al_source);
    }

    while (processed--) {
		  alSourceUnqueueBuffers(al_source, 1, &buffer);
                  const auto samples = sound_stream.stream->read();
                  const int size = sound_stream.stream->buffer_size;
			if (buffer != 0) {
	  			alBufferData(buffer, format, samples.data(),
					size * sizeof(ALshort),
					sound_stream.stream->sample_rate());
				alSourceQueueBuffers(al_source, 1, &buffer);
			}
    }
    if (sound_stream.source.loop && sound_stream.stream->done()) {
      sound_stream.stream->seek_start();
    }
    if (!sound_stream.source.playing && (state == AL_PLAYING)) {
      alSourceStop(al_source);
      ALint count{0};
      alGetSourcei(al_source, AL_BUFFERS_QUEUED, &count);
      alSourceUnqueueBuffers(al_source, count, &buffer);
      alDeleteBuffers(4, buffers.data());
    }
  }
}

auto Renderer::listener() const -> Listener {
  Listener listener;
  alGetListener3f(AL_POSITION, &listener.position.x, &listener.position.y,
                  &listener.position.z);

  std::array<float,6> orientation{};
  alGetListenerf(AL_ORIENTATION, orientation.data());
  listener.direction.x = orientation[0];
  listener.direction.y = orientation[1];
  listener.direction.z = orientation[2];
  listener.up.x = orientation[3];
  listener.up.y = orientation[4];
  listener.up.z = orientation[5];

  alGetListener3f(AL_VELOCITY, &listener.velocity.x, &listener.velocity.y,
                  &listener.velocity.z);

  return listener;
}

void Renderer::listener(const Listener &listener) {
  alListener3f(AL_POSITION, listener.position.x, listener.position.y,
               listener.position.z);

  alListener3f(AL_VELOCITY, listener.velocity.x, listener.velocity.y,
               listener.velocity.z);

  std::array<float, 6> orientation {listener.direction.x, listener.direction.y,
                          listener.direction.z, listener.up.x,
                          listener.up.y, listener.up.z};
  alListenerfv(AL_ORIENTATION, orientation.data());
  alListenerf(AL_GAIN, listener.gain);
}

auto Renderer::render(const Scene &scene, const float dt) -> void {
  listener(scene.listener);
  for (const auto &sound : scene.sounds) {
    render_sound(sound, dt);
  }
  for (const auto &sound_stream : scene.sound_streams) {
    render_sound_stream(sound_stream, dt);
  }
}

void Renderer::clear() {
  for (auto source : sources_) {
    alSourceStop(source.second);
    alDeleteSources(1, &source.second);
  }
  for (auto buffer : buffers_) {
    alDeleteBuffers(1, &buffer.second);
  }
  sources_.clear();
  buffers_.clear();
}
}
<|endoftext|>


<commit_before>#include <offboard_control/offboard_control.h>
#include <offboard_control/State.h>

#include <std_msgs/String.h>
#include <sstream>

static const int QUEUE_SIZE = 1;
static const std::string TOPIC_CONTROL_TAKEOFF = "offboard_control/takeoff";
static const std::string TOPIC_CONTROL_LAND = "offboard_control/land";
static const std::string TOPIC_CONTROL_RTL = "offboard_control/rtl";
static const std::string TOPIC_CONTROL_WAYPOINT = "offboard_control/waypoint";
static const std::string TOPIC_CONTROL_WAYPOINT_ARRIVED = "offboard_control/waypoint/arrived";
static const std::string TOPIC_CONTROL_VELOCITY = "offboard_control/velocity";
static const std::string TOPIC_CONTROL_VELOCITY_ALERT = "offboard_control/velocity/alert";
static const std::string TOPIC_CONTROL_RESUME_MISSION = "offboard_control/resume_mission";
static const std::string TOPIC_CONTROL_STOP = "offboard_control/stop";
static const std::string TOPIC_CONTROL_STATE = "offboard_control/state";
static const std::string TOPIC_CONTROL_GRIPPER = "offboard_control/gripper";
static const std::string TOPIC_LOGGER = "/ocs/log_message";
static const float TAKEOFF_HEIGHT = 10.0;
static const float CLOSE_ENOUGH = 2.0;
static const float VELOCITY_ALERT_HEIGHT = 3.0;
static const float PI = 3.14159265359;

OffboardControl::OffboardControl(
  ros::NodeHandle &rosNode,
  ros::Rate rosRate,
  std::string odometryTopic,
  MavrosAdapter::Autopilot autopilot
) : mMavrosAdapter(rosNode, rosRate, autopilot, TAKEOFF_HEIGHT, this, &OffboardControl::armingEvent, &OffboardControl::localPoseEvent) {
  this->mNodeHandle = &rosNode;
  this->mTakeoffService = this->mNodeHandle->advertiseService(TOPIC_CONTROL_TAKEOFF, &OffboardControl::takeoffService, this);
  this->mLandService = this->mNodeHandle->advertiseService(TOPIC_CONTROL_LAND, &OffboardControl::landService, this);
  this->mRtlService = this->mNodeHandle->advertiseService(TOPIC_CONTROL_RTL, &OffboardControl::rtlService, this);
  this->mWaypointSubscriber = this->mNodeHandle->subscribe<geometry_msgs::Pose>(TOPIC_CONTROL_WAYPOINT, QUEUE_SIZE, &OffboardControl::waypointCallback, this);
  this->mVelocitySubscriber = this->mNodeHandle->subscribe<geometry_msgs::Twist>(TOPIC_CONTROL_VELOCITY, QUEUE_SIZE, &OffboardControl::velocityCallback, this);
  this->mResumeMissionSubscriber = this->mNodeHandle->subscribe<std_msgs::Empty>(TOPIC_CONTROL_RESUME_MISSION, QUEUE_SIZE, &OffboardControl::resumeMissionCallback, this);
  this->mStopSubscriber = this->mNodeHandle->subscribe<std_msgs::Empty>(TOPIC_CONTROL_STOP, QUEUE_SIZE, &OffboardControl::stopCallback, this);
  this->mOdometrySubscriber = this->mNodeHandle->subscribe<nav_msgs::Odometry>(odometryTopic, QUEUE_SIZE, &OffboardControl::odometryCallback, this);
  this->mGripperSubscriber = this->mNodeHandle->subscribe<std_msgs::Bool>(TOPIC_CONTROL_GRIPPER, QUEUE_SIZE, &OffboardControl::gripperCallback, this);
  this->mWaypointArrivedPublisher = this->mNodeHandle->advertise<std_msgs::Empty>(TOPIC_CONTROL_WAYPOINT_ARRIVED, QUEUE_SIZE);
  this->mVelocityAlertPublisher = this->mNodeHandle->advertise<geometry_msgs::Pose>(TOPIC_CONTROL_VELOCITY_ALERT, QUEUE_SIZE);
  this->mLoggerPublisher = this->mNodeHandle->advertise<std_msgs::String>(TOPIC_LOGGER, QUEUE_SIZE);
  this->mStatePublisher = this->mNodeHandle->advertise<offboard_control::State>("offboard_control/state", QUEUE_SIZE);
  this->mEnRouteToWaypoint = false;
  this->mReceivedOdometry = false;
  this->initializeParameters();
}

void OffboardControl::initializeMavros() {
  this->mMavrosAdapter.initialize();
  this->logMessage("Initialized");
}

void OffboardControl::armingEvent() {
  this->mGlobalWaypoint = this->mCurrentOdometry.pose.pose;
  this->mGlobalWaypoint.position.z += TAKEOFF_HEIGHT;
  this->setWaypoint(this->mGlobalWaypoint);
  std::stringstream ss;
  ss << "Armed at (" << this->mCurrentOdometry.pose.pose.position.x << ", " << this->mCurrentOdometry.pose.pose.position.y << ", " << this->mCurrentOdometry.pose.pose.position.z << ")";
  this->logMessage(ss.str());
}

void OffboardControl::localPoseEvent(geometry_msgs::Pose pose) {
  this->publishState();
}

void OffboardControl::controlEffortEvent(geometry_msgs::Twist twist) {
  if (this->mMavrosAdapter.isInVelocityMode()) {
    this->mMavrosAdapter.executeMoveWithVelocity(twist);
  }
}

void OffboardControl::initializeParameters() {
  if (this->mNodeHandle->getParam("/arena/angle", this->mAngleOffset)) {
    std::stringstream ss;
    ss << "Angular offset set to \"" << this->mAngleOffset << "\".";
    this->logMessage(ss.str());
  }
  else {
    this->mAngleOffset = 0.0;
    this->logMessage("Angular offset parameter for arena not configured.", true);
  }
}

bool OffboardControl::takeoffService(std_srvs::Trigger::Request& request, std_srvs::Trigger::Response& response) {
  if (this->mReceivedOdometry) {
    std::stringstream ss;
    ss << "Takeing off to " << TAKEOFF_HEIGHT << " meters.";
    response.success = true;
    response.message = ss.str();
    this->mMavrosAdapter.executeTakeoffSequence(this->mAngleOffset);
  }
  else {
    response.success = false;
    response.message = "Takeoff request denied - no odometry has been received.";
  }
  return true;
}

bool OffboardControl::landService(std_srvs::Trigger::Request& request, std_srvs::Trigger::Response& response) {
  std::stringstream ss;
  ss << "Landing at (" << this->mCurrentOdometry.pose.pose.position.x << ", " << this->mCurrentOdometry.pose.pose.position.y << ")";
  response.success = true;
  response.message = ss.str();
  this->mMavrosAdapter.executeLandingSequence();
  return true;
}

bool OffboardControl::rtlService(std_srvs::Trigger::Request& request, std_srvs::Trigger::Response& response) {
  response.success = true;
  response.message = "Returning to launch.";
  this->mMavrosAdapter.executeRtlSequence();
  return true;
}

void OffboardControl::waypointCallback(const geometry_msgs::Pose::ConstPtr& msg) {
  this->setWaypoint(*msg);
}

void OffboardControl::velocityCallback(const geometry_msgs::Twist::ConstPtr& msg) {
  geometry_msgs::Twist twist;
  float sinTheta = sin(this->mAngleOffset);
  float cosTheta = cos(this->mAngleOffset);
  twist.linear.x = (msg->linear.x * cosTheta) - (msg->linear.y * sinTheta);
  twist.linear.y = (msg->linear.y * cosTheta) + (msg->linear.x * sinTheta);
  twist.linear.z = msg->linear.z;
  this->mMavrosAdapter.executeMoveWithVelocity(twist);
}

void OffboardControl::resumeMissionCallback(const std_msgs::Empty::ConstPtr& msg) {
  this->mEnRouteToWaypoint = false;
  // Set waypoint to current position when returning to mission.
  this->mGlobalWaypoint = this->mCurrentOdometry.pose.pose;
  this->setWaypoint(this->mGlobalWaypoint);
  this->mMavrosAdapter.executeMissionResume();
}

void OffboardControl::stopCallback(const std_msgs::Empty::ConstPtr& msg) {
  this->mGlobalWaypoint = this->mCurrentOdometry.pose.pose;
  this->setWaypoint(this->mGlobalWaypoint);
}

void OffboardControl::odometryCallback(const nav_msgs::Odometry::ConstPtr& msg) {
  this->mCurrentOdometry = *msg;
  this->mReceivedOdometry = true;
  if (this->mEnRouteToWaypoint &&
      fabs(msg->pose.pose.position.x - this->mGlobalWaypoint.position.x) < CLOSE_ENOUGH &&
      fabs(msg->pose.pose.position.y - this->mGlobalWaypoint.position.y) < CLOSE_ENOUGH &&
      fabs(msg->pose.pose.position.z - this->mGlobalWaypoint.position.z) < CLOSE_ENOUGH) {
    this->mWaypointArrivedPublisher.publish(std_msgs::Empty());
    this->mEnRouteToWaypoint = false;
  }
  else if (this->mMavrosAdapter.isInVelocityMode() && msg->pose.pose.position.z < VELOCITY_ALERT_HEIGHT) {
    this->mVelocityAlertPublisher.publish(msg->pose.pose);
  }
}

void OffboardControl::gripperCallback(const std_msgs::Bool::ConstPtr& msg) {
  this->mMavrosAdapter.enableGripper(msg->data);
  this->logMessage("Magnetic gripper: " + msg->data);
}

void OffboardControl::setWaypoint(const geometry_msgs::Pose pose) {
  this->mMavrosAdapter.executeMoveWithWaypoint(pose);
}

void OffboardControl::logMessage(std::string message, bool isError) const {
  std_msgs::String logMessage;
  if (isError) {
    logMessage.data = "[ERROR] ";
  }
  logMessage.data += "[" + this->getNamespace() + "] [OffboardControl] " + message;
  this->mLoggerPublisher.publish(logMessage);
}

std::string OffboardControl::getNamespace() const {
  std::string ns = this->mNodeHandle->getNamespace();
  ns.erase(0, 2);
  return ns;
}

void OffboardControl::publishState() const {
  offboard_control::State stateMsg;
  stateMsg.mode = this->mMavrosAdapter.getActionString();
  stateMsg.en_route_to_waypoint = this->mEnRouteToWaypoint;
  stateMsg.global_waypoint = this->mGlobalWaypoint;
  stateMsg.local_waypoint = this->mLocalWaypoint;
  this->mStatePublisher.publish(stateMsg);
}
<commit_msg>Refactored setup for ROS messaging.<commit_after>#include <offboard_control/offboard_control.h>
#include <offboard_control/State.h>

#include <std_msgs/String.h>
#include <sstream>

static const int QUEUE_SIZE = 1;
static const float TAKEOFF_HEIGHT = 10.0;
static const float CLOSE_ENOUGH = 2.0;
static const float VELOCITY_ALERT_HEIGHT = 3.0;
static const float PI = 3.14159265359;

OffboardControl::OffboardControl(
  ros::NodeHandle &rosNode,
  ros::Rate rosRate,
  std::string odometryTopic,
  MavrosAdapter::Autopilot autopilot
) : mMavrosAdapter(rosNode, rosRate, autopilot, TAKEOFF_HEIGHT, this, &OffboardControl::armingEvent, &OffboardControl::localPoseEvent) {
  this->mNodeHandle = &rosNode;
  this->mTakeoffService = this->mNodeHandle->advertiseService("offboard_control/takeoff", &OffboardControl::takeoffService, this);
  this->mLandService = this->mNodeHandle->advertiseService("offboard_control/land", &OffboardControl::landService, this);
  this->mRtlService = this->mNodeHandle->advertiseService("offboard_control/rtl", &OffboardControl::rtlService, this);
  this->mWaypointSubscriber = this->mNodeHandle->subscribe<geometry_msgs::Pose>("offboard_control/waypoint", QUEUE_SIZE, &OffboardControl::waypointCallback, this);
  this->mVelocitySubscriber = this->mNodeHandle->subscribe<geometry_msgs::Twist>("offboard_control/velocity", QUEUE_SIZE, &OffboardControl::velocityCallback, this);
  this->mResumeMissionSubscriber = this->mNodeHandle->subscribe<std_msgs::Empty>("offboard_control/resume_mission", QUEUE_SIZE, &OffboardControl::resumeMissionCallback, this);
  this->mStopSubscriber = this->mNodeHandle->subscribe<std_msgs::Empty>("offboard_control/stop", QUEUE_SIZE, &OffboardControl::stopCallback, this);
  this->mOdometrySubscriber = this->mNodeHandle->subscribe<nav_msgs::Odometry>(odometryTopic, QUEUE_SIZE, &OffboardControl::odometryCallback, this);
  this->mGripperSubscriber = this->mNodeHandle->subscribe<std_msgs::Bool>("offboard_control/gripper", QUEUE_SIZE, &OffboardControl::gripperCallback, this);
  this->mTakeoffCompletedPublisher = this->mNodeHandle->advertise<std_msgs::Empty>("offboard_control/takeoff/completed", QUEUE_SIZE);
  this->mWaypointArrivedPublisher = this->mNodeHandle->advertise<std_msgs::Empty>("offboard_control/waypoint/arrived", QUEUE_SIZE);
  this->mVelocityAlertPublisher = this->mNodeHandle->advertise<geometry_msgs::Pose>("offboard_control/velocity/alert", QUEUE_SIZE);
  this->mStatePublisher = this->mNodeHandle->advertise<offboard_control::State>("offboard_control/state", QUEUE_SIZE);
  this->mLoggerPublisher = this->mNodeHandle->advertise<std_msgs::String>("offboard_control/log", QUEUE_SIZE);
  this->mEnRouteToWaypoint = false;
  this->mReceivedOdometry = false;
  this->initializeParameters();
}

void OffboardControl::initializeMavros() {
  this->mMavrosAdapter.initialize();
  this->logMessage("Initialized");
}

void OffboardControl::armingEvent() {
  this->mGlobalWaypoint = this->mCurrentOdometry.pose.pose;
  this->mGlobalWaypoint.position.z += TAKEOFF_HEIGHT;
  this->setWaypoint(this->mGlobalWaypoint);
  std::stringstream ss;
  ss << "Armed at (" << this->mCurrentOdometry.pose.pose.position.x << ", " << this->mCurrentOdometry.pose.pose.position.y << ", " << this->mCurrentOdometry.pose.pose.position.z << ")";
  this->logMessage(ss.str());
}

void OffboardControl::localPoseEvent(geometry_msgs::Pose pose) {
  this->publishState();
}

void OffboardControl::controlEffortEvent(geometry_msgs::Twist twist) {
  if (this->mMavrosAdapter.isInVelocityMode()) {
    this->mMavrosAdapter.executeMoveWithVelocity(twist);
  }
}

void OffboardControl::initializeParameters() {
  if (this->mNodeHandle->getParam("/arena/angle", this->mAngleOffset)) {
    std::stringstream ss;
    ss << "Angular offset set to \"" << this->mAngleOffset << "\".";
    this->logMessage(ss.str());
  }
  else {
    this->mAngleOffset = 0.0;
    this->logMessage("Angular offset parameter for arena not configured.", true);
  }
}

bool OffboardControl::takeoffService(std_srvs::Trigger::Request& request, std_srvs::Trigger::Response& response) {
  if (this->mReceivedOdometry) {
    std::stringstream ss;
    ss << "Taking off to " << TAKEOFF_HEIGHT << " meters.";
    response.success = true;
    response.message = ss.str();
    this->mMavrosAdapter.executeTakeoffSequence(this->mAngleOffset);
  }
  else {
    response.success = false;
    response.message = "Takeoff request denied - no odometry has been received.";
  }
  return true;
}

bool OffboardControl::landService(std_srvs::Trigger::Request& request, std_srvs::Trigger::Response& response) {
  std::stringstream ss;
  ss << "Landing at (" << this->mCurrentOdometry.pose.pose.position.x << ", " << this->mCurrentOdometry.pose.pose.position.y << ")";
  response.success = true;
  response.message = ss.str();
  this->mMavrosAdapter.executeLandingSequence();
  return true;
}

bool OffboardControl::rtlService(std_srvs::Trigger::Request& request, std_srvs::Trigger::Response& response) {
  response.success = true;
  response.message = "Returning to launch.";
  this->mMavrosAdapter.executeRtlSequence();
  return true;
}

void OffboardControl::waypointCallback(const geometry_msgs::Pose::ConstPtr& msg) {
  this->setWaypoint(*msg);
}

void OffboardControl::velocityCallback(const geometry_msgs::Twist::ConstPtr& msg) {
  geometry_msgs::Twist twist;
  float sinTheta = sin(this->mAngleOffset);
  float cosTheta = cos(this->mAngleOffset);
  twist.linear.x = (msg->linear.x * cosTheta) - (msg->linear.y * sinTheta);
  twist.linear.y = (msg->linear.y * cosTheta) + (msg->linear.x * sinTheta);
  twist.linear.z = msg->linear.z;
  this->mMavrosAdapter.executeMoveWithVelocity(twist);
}

void OffboardControl::resumeMissionCallback(const std_msgs::Empty::ConstPtr& msg) {
  this->mEnRouteToWaypoint = false;
  // Set waypoint to current position when returning to mission.
  this->mGlobalWaypoint = this->mCurrentOdometry.pose.pose;
  this->setWaypoint(this->mGlobalWaypoint);
  this->mMavrosAdapter.executeMissionResume();
}

void OffboardControl::stopCallback(const std_msgs::Empty::ConstPtr& msg) {
  this->mGlobalWaypoint = this->mCurrentOdometry.pose.pose;
  this->setWaypoint(this->mGlobalWaypoint);
}

void OffboardControl::odometryCallback(const nav_msgs::Odometry::ConstPtr& msg) {
  this->mCurrentOdometry = *msg;
  this->mReceivedOdometry = true;
  if (this->mEnRouteToWaypoint &&
      fabs(msg->pose.pose.position.x - this->mGlobalWaypoint.position.x) < CLOSE_ENOUGH &&
      fabs(msg->pose.pose.position.y - this->mGlobalWaypoint.position.y) < CLOSE_ENOUGH &&
      fabs(msg->pose.pose.position.z - this->mGlobalWaypoint.position.z) < CLOSE_ENOUGH) {
    this->mWaypointArrivedPublisher.publish(std_msgs::Empty());
    this->mEnRouteToWaypoint = false;
  }
  else if (this->mMavrosAdapter.isInVelocityMode() && msg->pose.pose.position.z < VELOCITY_ALERT_HEIGHT) {
    this->mVelocityAlertPublisher.publish(msg->pose.pose);
  }
}

void OffboardControl::gripperCallback(const std_msgs::Bool::ConstPtr& msg) {
  this->mMavrosAdapter.enableGripper(msg->data);
  this->logMessage("Magnetic gripper: " + msg->data);
}

void OffboardControl::setWaypoint(const geometry_msgs::Pose pose) {
  this->mMavrosAdapter.executeMoveWithWaypoint(pose);
}

void OffboardControl::logMessage(std::string message, bool isError) const {
  std_msgs::String logMessage;
  if (isError) {
    logMessage.data = "[ERROR] ";
  }
  logMessage.data += "[" + this->getNamespace() + "] [OffboardControl] " + message;
  this->mLoggerPublisher.publish(logMessage);
}

std::string OffboardControl::getNamespace() const {
  std::string ns = this->mNodeHandle->getNamespace();
  ns.erase(0, 2);
  return ns;
}

void OffboardControl::publishState() const {
  offboard_control::State stateMsg;
  stateMsg.mode = this->mMavrosAdapter.getActionString();
  stateMsg.en_route_to_waypoint = this->mEnRouteToWaypoint;
  stateMsg.global_waypoint = this->mGlobalWaypoint;
  stateMsg.local_waypoint = this->mLocalWaypoint;
  this->mStatePublisher.publish(stateMsg);
}
<|endoftext|>