<commit_before>/* bzflag
 * Copyright (c) 1993 - 2004 Tim Riker
 *
 * This package is free software;  you can redistribute it and/or
 * modify it under the terms of the license found in the file
 * named LICENSE that should have accompanied this file.
 *
 * THIS PACKAGE IS PROVIDED ``AS IS'' AND WITHOUT ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED
 * WARRANTIES OF MERCHANTIBILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 */

// Provide BZFS with a list server connection

/* class header */
#include "ListServerConnection.h"

/* system implementation headers */
#include <string.h>
#include <string>
#include <math.h>
#include <errno.h>

/* common implementation headers */
#include "bzfio.h"
#include "version.h"
#include "TextUtils.h"
#include "Protocol.h"
#include "GameKeeper.h"

extern Address serverAddress;
extern PingPacket getTeamCounts();

ListServerLink::ListServerLink(std::string listServerURL, std::string publicizedAddress, std::string publicizedTitle)
{
  // parse url
  std::string protocol, hostname, pathname;
  int port = 80;
  bool useDefault = false;

  // use default if it can't be parsed
  if (!BzfNetwork::parseURL(listServerURL, protocol, hostname, port, pathname))
    useDefault = true;

  // use default if wrong protocol or invalid port
  if ((protocol != "http") || (port < 1) || (port > 65535))
    useDefault = true;

  // use default if bad address
  Address address = Address::getHostAddress(hostname);
  if (address.isAny())
    useDefault = true;

  // parse default list server URL if we need to; assume default works
  if (useDefault) {
    BzfNetwork::parseURL(DefaultListServerURL, protocol, hostname, port, pathname);
    DEBUG1("Provided list server URL (%s) is invalid.  Using default of %s.\n", listServerURL.c_str(), DefaultListServerURL);
  }

  // add to list
  this->address	   = address;
  this->port	   = port;
  this->pathname   = pathname;
  this->hostname   = hostname;
  this->linkSocket = NotConnected;

  this->publicizeAddress     = publicizedAddress;
  this->publicizeDescription = publicizedTitle;
  this->publicizeServer	     = true;  //if this c'tor is called, it's safe to publicize

  // schedule initial ADD message
  queueMessage(ListServerLink::ADD);
}

ListServerLink::ListServerLink()
{
  // does not create a usable link, so checks should be placed
  // in  all public member functions to ensure that nothing tries
  // to happen if publicizeServer is false
  this->linkSocket = NotConnected;
  this->publicizeServer = false;
}

ListServerLink::~ListServerLink()
{
  // now tell the list server that we're going away.  this can
  // take some time but we don't want to wait too long.  we do
  // our own multiplexing loop and wait for a maximum of 3 seconds
  // total.

  // if we aren't supposed to be publicizing, skip the whole thing
  // and don't waste 3 seconds.
  if (!publicizeServer)
    return;

  queueMessage(ListServerLink::REMOVE);
  TimeKeeper start = TimeKeeper::getCurrent();
  do {
    // compute timeout
    float waitTime = 3.0f - (TimeKeeper::getCurrent() - start);
    if (waitTime <= 0.0f)
      break;
    if (!isConnected()) //queueMessage should have connected us
      break;
    // check for list server socket connection
    int fdMax = -1;
    fd_set write_set;
    fd_set read_set;
    FD_ZERO(&write_set);
    FD_ZERO(&read_set);
    if (phase == ListServerLink::CONNECTING)
      _FD_SET(linkSocket, &write_set);
    else
      _FD_SET(linkSocket, &read_set);
    fdMax = linkSocket;

    // wait for socket to connect or timeout
    struct timeval timeout;
    timeout.tv_sec = long(floorf(waitTime));
    timeout.tv_usec = long(1.0e+6f * (waitTime - floorf(waitTime)));
    int nfound = select(fdMax + 1, (fd_set*)&read_set, (fd_set*)&write_set,
			0, &timeout);
    if (nfound == 0)
      // Time has gone, close and go
      break;
    // check for connection to list server
    if (FD_ISSET(linkSocket, &write_set))
      sendQueuedMessages();
    else if (FD_ISSET(linkSocket, &read_set))
      read();
  } while (true);

  // stop list server communication
  closeLink();
}

void ListServerLink::closeLink()
{
  if (isConnected()) {
    close(linkSocket);
    DEBUG4("Closing List Server\n");
    linkSocket = NotConnected;
  }
}

void ListServerLink::read()
{
  if (isConnected()) {
    char    buf[256];
    recv(linkSocket, buf, sizeof(buf), 0);
    closeLink();
    if (nextMessageType != ListServerLink::NONE)
      // There was a pending request arrived after we write:
      // we should redo all the stuff
      openLink();
  }
}

void ListServerLink::openLink()
{
  // start opening connection if not already doing so
  if (!isConnected()) {
    linkSocket = socket(AF_INET, SOCK_STREAM, 0);
    DEBUG4("Opening List Server\n");
    if (!isConnected()) {
      return;
    }

    // set to non-blocking for connect
    if (BzfNetwork::setNonBlocking(linkSocket) < 0) {
      closeLink();
      return;
    }

    // Make our connection come from our serverAddress in case we have
    // multiple/masked IPs so the list server can verify us.
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr = serverAddress;

    // assign the address to the socket
    if (bind(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
      closeLink();
      return;
    }

    // connect.  this should fail with EINPROGRESS but check for
    // success just in case.
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port   = htons(port);
    addr.sin_addr   = address;
    if (connect(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
#if defined(_WIN32)
#undef EINPROGRESS
#define EINPROGRESS EWOULDBLOCK
#endif
      if (getErrno() != EINPROGRESS) {
	nerror("connecting to list server");
	// try to lookup dns name again in case it moved
	this->address = Address::getHostAddress(this->hostname);
	closeLink();
      } else {
	phase = CONNECTING;
      }
    } else {
      // shouldn't arrive here. Just in case, clean
      closeLink();
    }
  }
}

void ListServerLink::queueMessage(MessageType type)
{
  // ignore if the server is not public
  if (!publicizeServer) return;

  // Open network connection only if closed
  if (!isConnected()) openLink();

  // record next message to send.
  nextMessageType = type;
}

void ListServerLink::sendQueuedMessages()
{
  if (!isConnected())
    return;

  if (nextMessageType == ListServerLink::ADD) {
    DEBUG3("Queuing ADD message to list server\n");
    addMe(getTeamCounts(), publicizeAddress, TextUtils::url_encode(publicizeDescription));
    lastAddTime = TimeKeeper::getCurrent();
  } else if (nextMessageType == ListServerLink::REMOVE) {
    DEBUG3("Queuing REMOVE message to list server\n");
    removeMe(publicizeAddress);
  }
}

void ListServerLink::addMe(PingPacket pingInfo,
			   std::string publicizedAddress,
			   std::string publicizedTitle)
{
  std::string msg;

  // encode ping reply as ascii hex digits plus NULL
  char gameInfo[PingPacketHexPackedSize + 1];
  pingInfo.packHex(gameInfo);

  // TODO loop through and send any player tokens that need approval
  // &checktokens=CallSign0%3D01234567%0D%0ACallSign1%3D90abcdef%0D%0A
  // TODO loop through groups we are interested and request those
  // *groups=GROUP0%0D%0AGROUP1%0D%0A
  // TODO we probably should convert to a POST instead. List server now allows either
  // send ADD message (must send blank line)
  msg = TextUtils::format("GET %s?action=ADD&nameport=%s&version=%s&gameinfo=%s&build=%s",
    pathname.c_str(), publicizedAddress.c_str(),
    getServerVersion(), gameInfo,
    getAppVersion());
// FIXME should be in a header someplace NOT an extern
extern uint16_t curMaxPlayers;
  msg += "&checktokens=";
  for (int i = 0; i < curMaxPlayers; i++) {
    GameKeeper::Player *playerData = GameKeeper::Player::getPlayerByIndex(i);
    if (playerData && strlen(playerData->player.getCallSign()) && strlen(playerData->player.getToken())) {
      msg += TextUtils::format(playerData->player.getCallSign());
      msg += "=";
      msg += TextUtils::format(playerData->player.getToken());
      msg += "%%0D%%0A";
    }
  }
  msg += TextUtils::format("&title=%s HTTP/1.1\r\nUser-Agent: bzfs %s\r\nHost: %s\r\nCache-Control: no-cache\r\n\r\n",
    publicizedTitle.c_str(),
    getAppVersion(),
    hostname.c_str());
  // TODO need to listen for user info replies and setup callsign for isAllowedToEnter()
  sendMessage(msg);
}

void ListServerLink::removeMe(std::string publicizedAddress)
{
  std::string msg;
  // send REMOVE (must send blank line)
  msg = TextUtils::format("GET %s?action=REMOVE&nameport=%s HTTP/1.1\r\n"
    "User-Agent: bzfs %s\r\nHost: %s\r\nCache-Control: no-cache\r\n\r\n",
    pathname.c_str(),
    publicizedAddress.c_str(),
    getAppVersion(),
    hostname.c_str());
  sendMessage(msg);
}

void ListServerLink::sendMessage(std::string message)
{
  const int bufsize = 4096;
  char msg[bufsize];
  strncpy(msg, message.c_str(), bufsize);
  msg[bufsize - 1] = 0;
  if (strlen(msg) > 0) {
    DEBUG3("%s\n", msg);
    if (send(linkSocket, msg, strlen(msg), 0) == -1) {
      perror("List server send failed");
      DEBUG3("Unable to send to the list server!\n");
      closeLink();
    } else {
      nextMessageType = ListServerLink::NONE;
      phase	   = ListServerLink::WRITTEN;
    }
  } else {
    closeLink();
  }
}

// Local Variables: ***
// mode:C++ ***
// tab-width: 8 ***
// c-basic-offset: 2 ***
// indent-tabs-mode: t ***
// End: ***
// ex: shiftwidth=2 tabstop=8
<commit_msg>send Connection: close, DEBUG4 shows list server replies<commit_after>/* bzflag
 * Copyright (c) 1993 - 2004 Tim Riker
 *
 * This package is free software;  you can redistribute it and/or
 * modify it under the terms of the license found in the file
 * named LICENSE that should have accompanied this file.
 *
 * THIS PACKAGE IS PROVIDED ``AS IS'' AND WITHOUT ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED
 * WARRANTIES OF MERCHANTIBILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 */

// Provide BZFS with a list server connection

/* class header */
#include "ListServerConnection.h"

/* system implementation headers */
#include <string.h>
#include <string>
#include <math.h>
#include <errno.h>

/* common implementation headers */
#include "bzfio.h"
#include "version.h"
#include "TextUtils.h"
#include "Protocol.h"
#include "GameKeeper.h"

extern Address serverAddress;
extern PingPacket getTeamCounts();

ListServerLink::ListServerLink(std::string listServerURL, std::string publicizedAddress, std::string publicizedTitle)
{
  // parse url
  std::string protocol, hostname, pathname;
  int port = 80;
  bool useDefault = false;

  // use default if it can't be parsed
  if (!BzfNetwork::parseURL(listServerURL, protocol, hostname, port, pathname))
    useDefault = true;

  // use default if wrong protocol or invalid port
  if ((protocol != "http") || (port < 1) || (port > 65535))
    useDefault = true;

  // use default if bad address
  Address address = Address::getHostAddress(hostname);
  if (address.isAny())
    useDefault = true;

  // parse default list server URL if we need to; assume default works
  if (useDefault) {
    BzfNetwork::parseURL(DefaultListServerURL, protocol, hostname, port, pathname);
    DEBUG1("Provided list server URL (%s) is invalid.  Using default of %s.\n", listServerURL.c_str(), DefaultListServerURL);
  }

  // add to list
  this->address	   = address;
  this->port	   = port;
  this->pathname   = pathname;
  this->hostname   = hostname;
  this->linkSocket = NotConnected;

  this->publicizeAddress     = publicizedAddress;
  this->publicizeDescription = publicizedTitle;
  this->publicizeServer	     = true;  //if this c'tor is called, it's safe to publicize

  // schedule initial ADD message
  queueMessage(ListServerLink::ADD);
}

ListServerLink::ListServerLink()
{
  // does not create a usable link, so checks should be placed
  // in  all public member functions to ensure that nothing tries
  // to happen if publicizeServer is false
  this->linkSocket = NotConnected;
  this->publicizeServer = false;
}

ListServerLink::~ListServerLink()
{
  // now tell the list server that we're going away.  this can
  // take some time but we don't want to wait too long.  we do
  // our own multiplexing loop and wait for a maximum of 3 seconds
  // total.

  // if we aren't supposed to be publicizing, skip the whole thing
  // and don't waste 3 seconds.
  if (!publicizeServer)
    return;

  queueMessage(ListServerLink::REMOVE);
  TimeKeeper start = TimeKeeper::getCurrent();
  do {
    // compute timeout
    float waitTime = 3.0f - (TimeKeeper::getCurrent() - start);
    if (waitTime <= 0.0f)
      break;
    if (!isConnected()) //queueMessage should have connected us
      break;
    // check for list server socket connection
    int fdMax = -1;
    fd_set write_set;
    fd_set read_set;
    FD_ZERO(&write_set);
    FD_ZERO(&read_set);
    if (phase == ListServerLink::CONNECTING)
      _FD_SET(linkSocket, &write_set);
    else
      _FD_SET(linkSocket, &read_set);
    fdMax = linkSocket;

    // wait for socket to connect or timeout
    struct timeval timeout;
    timeout.tv_sec = long(floorf(waitTime));
    timeout.tv_usec = long(1.0e+6f * (waitTime - floorf(waitTime)));
    int nfound = select(fdMax + 1, (fd_set*)&read_set, (fd_set*)&write_set,
			0, &timeout);
    if (nfound == 0)
      // Time has gone, close and go
      break;
    // check for connection to list server
    if (FD_ISSET(linkSocket, &write_set))
      sendQueuedMessages();
    else if (FD_ISSET(linkSocket, &read_set))
      read();
  } while (true);

  // stop list server communication
  closeLink();
}

void ListServerLink::closeLink()
{
  if (isConnected()) {
    close(linkSocket);
    DEBUG4("Closing List Server\n");
    linkSocket = NotConnected;
  }
}

void ListServerLink::read()
{
  if (isConnected()) {
    char    buf[2048];
    int bytes = recv(linkSocket, buf, sizeof(buf)-1, 0);
    buf[bytes]=0;
    DEBUG4("ListServerStart\n%sListServerEnd\n", buf);
    closeLink();
    if (nextMessageType != ListServerLink::NONE) {
      // There was a pending request arrived after we write:
      // we should redo all the stuff
      openLink();
    }
  }
}

void ListServerLink::openLink()
{
  // start opening connection if not already doing so
  if (!isConnected()) {
    linkSocket = socket(AF_INET, SOCK_STREAM, 0);
    DEBUG4("Opening List Server\n");
    if (!isConnected()) {
      return;
    }

    // set to non-blocking for connect
    if (BzfNetwork::setNonBlocking(linkSocket) < 0) {
      closeLink();
      return;
    }

    // Make our connection come from our serverAddress in case we have
    // multiple/masked IPs so the list server can verify us.
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr = serverAddress;

    // assign the address to the socket
    if (bind(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
      closeLink();
      return;
    }

    // connect.  this should fail with EINPROGRESS but check for
    // success just in case.
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port   = htons(port);
    addr.sin_addr   = address;
    if (connect(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
#if defined(_WIN32)
#undef EINPROGRESS
#define EINPROGRESS EWOULDBLOCK
#endif
      if (getErrno() != EINPROGRESS) {
	nerror("connecting to list server");
	// try to lookup dns name again in case it moved
	this->address = Address::getHostAddress(this->hostname);
	closeLink();
      } else {
	phase = CONNECTING;
      }
    } else {
      // shouldn't arrive here. Just in case, clean
      DEBUG3("list server connect and close?");
      closeLink();
    }
  }
}

void ListServerLink::queueMessage(MessageType type)
{
  // ignore if the server is not public
  if (!publicizeServer) return;

  // Open network connection only if closed
  if (!isConnected()) openLink();

  // record next message to send.
  nextMessageType = type;
}

void ListServerLink::sendQueuedMessages()
{
  if (!isConnected())
    return;

  if (nextMessageType == ListServerLink::ADD) {
    DEBUG3("Queuing ADD message to list server\n");
    addMe(getTeamCounts(), publicizeAddress, TextUtils::url_encode(publicizeDescription));
    lastAddTime = TimeKeeper::getCurrent();
  } else if (nextMessageType == ListServerLink::REMOVE) {
    DEBUG3("Queuing REMOVE message to list server\n");
    removeMe(publicizeAddress);
  }
}

void ListServerLink::addMe(PingPacket pingInfo,
			   std::string publicizedAddress,
			   std::string publicizedTitle)
{
  std::string msg;

  // encode ping reply as ascii hex digits plus NULL
  char gameInfo[PingPacketHexPackedSize + 1];
  pingInfo.packHex(gameInfo);

  // TODO loop through and send any player tokens that need approval
  // &checktokens=CallSign0%3D01234567%0D%0ACallSign1%3D90abcdef%0D%0A
  // TODO loop through groups we are interested and request those
  // *groups=GROUP0%0D%0AGROUP1%0D%0A
  // TODO we probably should convert to a POST instead. List server now allows either
  // send ADD message (must send blank line)
  msg = TextUtils::format("GET %s?action=ADD&nameport=%s&version=%s&gameinfo=%s&build=%s",
    pathname.c_str(), publicizedAddress.c_str(),
    getServerVersion(), gameInfo,
    getAppVersion());
// FIXME should be in a header someplace NOT an extern
extern uint16_t curMaxPlayers;
  msg += "&checktokens=";
  for (int i = 0; i < curMaxPlayers; i++) {
    GameKeeper::Player *playerData = GameKeeper::Player::getPlayerByIndex(i);
    if (playerData && strlen(playerData->player.getCallSign()) && strlen(playerData->player.getToken())) {
      msg += TextUtils::format(playerData->player.getCallSign());
      msg += "=";
      msg += TextUtils::format(playerData->player.getToken());
      msg += "%%0D%%0A";
    }
  }
  msg += TextUtils::format("&title=%s HTTP/1.1\r\n"
      "User-Agent: bzfs %s\r\n"
      "Host: %s\r\n"
      "Cache-Control: no-cache\r\n"
      "Connection: close\r\n"
      "\r\n",
    publicizedTitle.c_str(),
    getAppVersion(),
    hostname.c_str());
  sendMessage(msg);
}

void ListServerLink::removeMe(std::string publicizedAddress)
{
  std::string msg;
  // send REMOVE (must send blank line)
  msg = TextUtils::format("GET %s?action=REMOVE&nameport=%s HTTP/1.1\r\n"
    "User-Agent: bzfs %s\r\n"
    "Host: %s\r\n"
    "Cache-Control: no-cache\r\n"
    "Connection: close\r\n"
    "\r\n",
    pathname.c_str(),
    publicizedAddress.c_str(),
    getAppVersion(),
    hostname.c_str());
  sendMessage(msg);
}

void ListServerLink::sendMessage(std::string message)
{
  const int bufsize = 4096;
  char msg[bufsize];
  strncpy(msg, message.c_str(), bufsize);
  msg[bufsize - 1] = 0;
  if (strlen(msg) > 0) {
    DEBUG3("%s\n", msg);
    if (send(linkSocket, msg, strlen(msg), 0) == -1) {
      perror("List server send failed");
      DEBUG3("Unable to send to the list server!\n");
      closeLink();
    } else {
      nextMessageType = ListServerLink::NONE;
      phase	   = ListServerLink::WRITTEN;
    }
  } else {
    closeLink();
  }
}

// Local Variables: ***
// mode:C++ ***
// tab-width: 8 ***
// c-basic-offset: 2 ***
// indent-tabs-mode: t ***
// End: ***
// ex: shiftwidth=2 tabstop=8
<|endoftext|>


<commit_before>/* bzflag
 * Copyright (c) 1993 - 2005 Tim Riker
 *
 * This package is free software;  you can redistribute it and/or
 * modify it under the terms of the license found in the file
 * named LICENSE that should have accompanied this file.
 *
 * THIS PACKAGE IS PROVIDED ``AS IS'' AND WITHOUT ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED
 * WARRANTIES OF MERCHANTIBILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 */

// Provide BZFS with a list server connection

/* class header */
#include "ListServerConnection.h"

/* system implementation headers */
#include <string.h>
#include <string>
#include <math.h>
#include <errno.h>

/* common implementation headers */
#include "bzfio.h"
#include "version.h"
#include "TextUtils.h"
#include "Protocol.h"
#include "GameKeeper.h"

/* local implementation headers */
#include "CmdLineOptions.h"

// FIXME remove externs!
extern Address serverAddress;
extern PingPacket getTeamCounts();
extern uint16_t curMaxPlayers;
extern int getTarget(const char *victimname);
extern void sendMessage(int playerIndex, PlayerId targetPlayer, const char *message);
extern void sendPlayerInfo(void);
extern void sendIPUpdate(int targetPlayer, int playerIndex);
extern CmdLineOptions *clOptions;

ListServerLink::ListServerLink(std::string listServerURL, std::string publicizedAddress, std::string publicizedTitle)
{
  // parse url
  std::string protocol, hostname, pathname;
  int port = 80;
  bool useDefault = false;

  // use default if it can't be parsed
  if (!BzfNetwork::parseURL(listServerURL, protocol, hostname, port, pathname))
    useDefault = true;

  // use default if wrong protocol or invalid port
  if ((protocol != "http") || (port < 1) || (port > 65535))
    useDefault = true;

  // use default if bad address
  Address address = Address::getHostAddress(hostname);
  if (address.isAny())
    useDefault = true;

  // parse default list server URL if we need to; assume default works
  if (useDefault) {
    BzfNetwork::parseURL(DefaultListServerURL, protocol, hostname, port, pathname);
    DEBUG1("Provided list server URL (%s) is invalid.  Using default of %s.\n", listServerURL.c_str(), DefaultListServerURL);
  }

  // add to list
  this->address	   = address;
  this->port	   = port;
  this->pathname   = pathname;
  this->hostname   = hostname;
  this->linkSocket = NotConnected;

  this->publicizeAddress     = publicizedAddress;
  this->publicizeDescription = publicizedTitle;
  this->publicizeServer	     = true;  //if this c'tor is called, it's safe to publicize

  // schedule initial ADD message
  queueMessage(ListServerLink::ADD);
}

ListServerLink::ListServerLink()
{
  // does not create a usable link, so checks should be placed
  // in  all public member functions to ensure that nothing tries
  // to happen if publicizeServer is false
  this->linkSocket = NotConnected;
  this->publicizeServer = false;
}

ListServerLink::~ListServerLink()
{
  // now tell the list server that we're going away.  this can
  // take some time but we don't want to wait too long.  we do
  // our own multiplexing loop and wait for a maximum of 3 seconds
  // total.

  // if we aren't supposed to be publicizing, skip the whole thing
  // and don't waste 3 seconds.
  if (!publicizeServer)
    return;

  queueMessage(ListServerLink::REMOVE);
  TimeKeeper start = TimeKeeper::getCurrent();
  do {
    // compute timeout
    float waitTime = 3.0f - (TimeKeeper::getCurrent() - start);
    if (waitTime <= 0.0f)
      break;
    if (!isConnected()) //queueMessage should have connected us
      break;
    // check for list server socket connection
    int fdMax = -1;
    fd_set write_set;
    fd_set read_set;
    FD_ZERO(&write_set);
    FD_ZERO(&read_set);
    if (phase == ListServerLink::CONNECTING)
      _FD_SET(linkSocket, &write_set);
    else
      _FD_SET(linkSocket, &read_set);
    fdMax = linkSocket;

    // wait for socket to connect or timeout
    struct timeval timeout;
    timeout.tv_sec = long(floorf(waitTime));
    timeout.tv_usec = long(1.0e+6f * (waitTime - floorf(waitTime)));
    int nfound = select(fdMax + 1, (fd_set*)&read_set, (fd_set*)&write_set,
			0, &timeout);
    if (nfound == 0)
      // Time has gone, close and go
      break;
    // check for connection to list server
    if (FD_ISSET(linkSocket, &write_set))
      sendQueuedMessages();
    else if (FD_ISSET(linkSocket, &read_set))
      read();
  } while (true);

  // stop list server communication
  closeLink();
}

void ListServerLink::closeLink()
{
  if (isConnected()) {
    close(linkSocket);
    DEBUG4("Closing List Server\n");
    linkSocket = NotConnected;
  }
}

void ListServerLink::read()
{
  if (isConnected()) {
    char buf[2048];
    int bytes = recv(linkSocket, buf, sizeof(buf)-1, 0);
    // TODO don't close unless we've got it all
    closeLink();
    buf[bytes]=0;
    char* base = buf;
    static char *tokGoodIdentifier = "TOKGOOD: ";
    static char *tokBadIdentifier = "TOKBAD: ";
    // walks entire reply including HTTP headers
    while (*base) {
      // find next newline
      char* scan = base;
      while (*scan && *scan != '\r' && *scan != '\n') scan++;
      // if no newline then no more complete replies
      if (*scan != '\r' && *scan != '\n') break;
      while (*scan && (*scan == '\r' || *scan == '\n')) *scan++ = '\0';
      DEBUG4("Got line: \"%s\"\n", base);
      // TODO don't do this if we don't want central logins
      if (strncmp(base, tokGoodIdentifier, strlen(tokGoodIdentifier)) == 0) {
	char *callsign, *group;
	callsign = base + strlen(tokGoodIdentifier);
	DEBUG3("Got: %s\n", base);
	group = callsign;
	while (*group && (*group != ':')) group++;
	while (*group && (*group == ':')) *group++ = 0;
	int playerIndex = getTarget(callsign);
	if (playerIndex < curMaxPlayers) {
	  GameKeeper::Player *playerData = GameKeeper::Player::getPlayerByIndex(playerIndex);
	  // don't accept the global auth if there's a local account of the same name
	  // and the local account is not marked as being the same as the global account
	  if (!playerData->accessInfo.hasRealPassword()
	      || playerData->accessInfo.getUserInfo(callsign).hasGroup("LOCAL.GLOBAL")) {
	    if (!playerData->accessInfo.isRegistered())
	      playerData->accessInfo.storeInfo(NULL);
	    playerData->accessInfo.setPermissionRights();
	    while (*group) {
	      char *nextgroup = group;
	      while (*nextgroup && (*nextgroup != ':')) nextgroup++;
	      while (*nextgroup && (*nextgroup == ':')) *nextgroup++ = 0;
	      playerData->accessInfo.addGroup(group);
	      //DEBUG3("Got: [%d] \"%s\" \"%s\"\n", playerIndex, callsign, group);
	      group = nextgroup;
	    }
	    // automatically give global users permission to use local accounts
	    // since they either already have it, or there's no existing local account.
	    playerData->accessInfo.addGroup("LOCAL.GLOBAL");
	    sendMessage(ServerPlayer, playerIndex, "Global login approved!");
	    sendIPUpdate(playerIndex, -1);
	    sendPlayerInfo();
	  } else {
	    sendMessage(ServerPlayer, playerIndex, "Global login rejected. "
			"This callsign is registered locally on this server.  If it is "
			"yours, please ask an admin to grant it membership in the "
			"LOCAL.GLOBAL group.  If it is not yours, please ask an admin "
			"to deregister it so that you may use your global callsign.");
	  }
	  playerData->player.clearToken();
	}
      } else if (strncmp(base, tokBadIdentifier, strlen(tokBadIdentifier)) == 0) {
	char *callsign;
	callsign = base + strlen(tokBadIdentifier);
	int playerIndex = getTarget(callsign);
	DEBUG3("Got: [%d] %s\n", playerIndex, base);
	if (playerIndex < curMaxPlayers) {
	  GameKeeper::Player *playerData = GameKeeper::Player::getPlayerByIndex(playerIndex);
	  sendMessage(ServerPlayer, playerIndex, "Global login rejected, bad token.");
	  playerData->player.clearToken();
	}
      }

      // next reply
      base = scan;
    }
    if (nextMessageType != ListServerLink::NONE) {
      // There was a pending request arrived after we write:
      // we should redo all the stuff
      openLink();
    }
  }
}

void ListServerLink::openLink()
{
  // start opening connection if not already doing so
  if (!isConnected()) {
    linkSocket = socket(AF_INET, SOCK_STREAM, 0);
    DEBUG4("Opening List Server\n");
    if (!isConnected()) {
      return;
    }

    // set to non-blocking for connect
    if (BzfNetwork::setNonBlocking(linkSocket) < 0) {
      closeLink();
      return;
    }

    // Make our connection come from our serverAddress in case we have
    // multiple/masked IPs so the list server can verify us.
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr = serverAddress;

    // assign the address to the socket
    if (bind(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
      closeLink();
      return;
    }

    // connect.  this should fail with EINPROGRESS but check for
    // success just in case.
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port   = htons(port);
    addr.sin_addr   = address;
    if (connect(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
#if defined(_WIN32)
#undef EINPROGRESS
#define EINPROGRESS EWOULDBLOCK
#endif
      if (getErrno() != EINPROGRESS) {
	nerror("connecting to list server");
	// try to lookup dns name again in case it moved
	this->address = Address::getHostAddress(this->hostname);
	closeLink();
      } else {
	phase = CONNECTING;
      }
    } else {
      // shouldn't arrive here. Just in case, clean
      DEBUG3("list server connect and close?\n");
      closeLink();
    }
  }
}

void ListServerLink::queueMessage(MessageType type)
{
  // ignore if the server is not public
  if (!publicizeServer) return;

  // Open network connection only if closed
  if (!isConnected()) openLink();

  // record next message to send.
  nextMessageType = type;
}

void ListServerLink::sendQueuedMessages()
{
  if (!isConnected())
    return;

  if (nextMessageType == ListServerLink::ADD) {
    DEBUG3("Queuing ADD message to list server\n");
    addMe(getTeamCounts(), publicizeAddress, TextUtils::url_encode(publicizeDescription));
    lastAddTime = TimeKeeper::getCurrent();
  } else if (nextMessageType == ListServerLink::REMOVE) {
    DEBUG3("Queuing REMOVE message to list server\n");
    removeMe(publicizeAddress);
  }
}

void ListServerLink::addMe(PingPacket pingInfo,
			   std::string publicizedAddress,
			   std::string publicizedTitle)
{
  std::string msg;

  // encode ping reply as ascii hex digits plus NULL
  char gameInfo[PingPacketHexPackedSize + 1];
  pingInfo.packHex(gameInfo);

  // TODO we probably should convert to a POST. List server now allows either
  // send ADD message (must send blank line)
  msg = TextUtils::format("GET %s?action=ADD&nameport=%s&version=%s&gameinfo=%s&build=%s",
    pathname.c_str(), publicizedAddress.c_str(),
    getServerVersion(), gameInfo,
    getAppVersion());
  msg += "&checktokens=";
  // callsign1@ip1=token1%0D%0Acallsign2@ip2=token2%0D%0A
  for (int i = 0; i < curMaxPlayers; i++) {
    GameKeeper::Player *playerData = GameKeeper::Player::getPlayerByIndex(i);
    if (!playerData)
      continue;
    NetHandler *handler = playerData->netHandler;
    std::string encodedCallsign = TextUtils::url_encode(playerData->player.getCallSign());
    if (strlen(playerData->player.getCallSign()) && strlen(playerData->player.getToken())) {
      msg += encodedCallsign;
      Address addr = handler->getIPAddress();
      if (!addr.isPrivate()) {
        msg += "@";
        msg += handler->getTargetIP();
      }
      msg += "=";
      msg += playerData->player.getToken();
      msg += "%0D%0A";
    }
  }
  // *groups=GROUP0%0D%0AGROUP1%0D%0A
  msg += "&groups=";
  PlayerAccessMap::iterator itr = groupAccess.begin();
  for ( ; itr != groupAccess.end(); itr++) {
    if (itr->first.substr(0, 6) != "LOCAL.") {
      msg += itr->first.c_str();
      msg += "%0D%0A";
    }
  }

  msg += TextUtils::format("&title=%s HTTP/1.1\r\n"
      "User-Agent: bzfs %s\r\n"
      "Host: %s\r\n"
      "Cache-Control: no-cache\r\n"
      "Connection: close\r\n"
      "\r\n",
    publicizedTitle.c_str(),
    getAppVersion(),
    hostname.c_str());
  sendLSMessage(msg);
}

void ListServerLink::removeMe(std::string publicizedAddress)
{
  std::string msg;
  // send REMOVE (must send blank line)
  msg = TextUtils::format("GET %s?action=REMOVE&nameport=%s HTTP/1.1\r\n"
    "User-Agent: bzfs %s\r\n"
    "Host: %s\r\n"
    "Cache-Control: no-cache\r\n"
    "Connection: close\r\n"
    "\r\n",
    pathname.c_str(),
    publicizedAddress.c_str(),
    getAppVersion(),
    hostname.c_str());
  sendLSMessage(msg);
}

void ListServerLink::sendLSMessage(std::string message)
{
  const int bufsize = 4096;
  char msg[bufsize];
  strncpy(msg, message.c_str(), bufsize);
  msg[bufsize - 1] = 0;
  if (strlen(msg) > 0) {
    DEBUG3("%s\n", msg);
    if (send(linkSocket, msg, strlen(msg), 0) == -1) {
      perror("List server send failed");
      DEBUG3("Unable to send to the list server!\n");
      closeLink();
    } else {
      nextMessageType = ListServerLink::NONE;
      phase	   = ListServerLink::WRITTEN;
    }
  } else {
    closeLink();
  }
}

// Local Variables: ***
// mode:C++ ***
// tab-width: 8 ***
// c-basic-offset: 2 ***
// indent-tabs-mode: t ***
// End: ***
// ex: shiftwidth=2 tabstop=8
<commit_msg>break long "this-account-exists-locally-so-ask-an-admin-to-do-something-about-it" message into three parts so that it all gets there<commit_after>/* bzflag
 * Copyright (c) 1993 - 2005 Tim Riker
 *
 * This package is free software;  you can redistribute it and/or
 * modify it under the terms of the license found in the file
 * named LICENSE that should have accompanied this file.
 *
 * THIS PACKAGE IS PROVIDED ``AS IS'' AND WITHOUT ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED
 * WARRANTIES OF MERCHANTIBILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 */

// Provide BZFS with a list server connection

/* class header */
#include "ListServerConnection.h"

/* system implementation headers */
#include <string.h>
#include <string>
#include <math.h>
#include <errno.h>

/* common implementation headers */
#include "bzfio.h"
#include "version.h"
#include "TextUtils.h"
#include "Protocol.h"
#include "GameKeeper.h"

/* local implementation headers */
#include "CmdLineOptions.h"

// FIXME remove externs!
extern Address serverAddress;
extern PingPacket getTeamCounts();
extern uint16_t curMaxPlayers;
extern int getTarget(const char *victimname);
extern void sendMessage(int playerIndex, PlayerId targetPlayer, const char *message);
extern void sendPlayerInfo(void);
extern void sendIPUpdate(int targetPlayer, int playerIndex);
extern CmdLineOptions *clOptions;

ListServerLink::ListServerLink(std::string listServerURL, std::string publicizedAddress, std::string publicizedTitle)
{
  // parse url
  std::string protocol, hostname, pathname;
  int port = 80;
  bool useDefault = false;

  // use default if it can't be parsed
  if (!BzfNetwork::parseURL(listServerURL, protocol, hostname, port, pathname))
    useDefault = true;

  // use default if wrong protocol or invalid port
  if ((protocol != "http") || (port < 1) || (port > 65535))
    useDefault = true;

  // use default if bad address
  Address address = Address::getHostAddress(hostname);
  if (address.isAny())
    useDefault = true;

  // parse default list server URL if we need to; assume default works
  if (useDefault) {
    BzfNetwork::parseURL(DefaultListServerURL, protocol, hostname, port, pathname);
    DEBUG1("Provided list server URL (%s) is invalid.  Using default of %s.\n", listServerURL.c_str(), DefaultListServerURL);
  }

  // add to list
  this->address	   = address;
  this->port	   = port;
  this->pathname   = pathname;
  this->hostname   = hostname;
  this->linkSocket = NotConnected;

  this->publicizeAddress     = publicizedAddress;
  this->publicizeDescription = publicizedTitle;
  this->publicizeServer	     = true;  //if this c'tor is called, it's safe to publicize

  // schedule initial ADD message
  queueMessage(ListServerLink::ADD);
}

ListServerLink::ListServerLink()
{
  // does not create a usable link, so checks should be placed
  // in  all public member functions to ensure that nothing tries
  // to happen if publicizeServer is false
  this->linkSocket = NotConnected;
  this->publicizeServer = false;
}

ListServerLink::~ListServerLink()
{
  // now tell the list server that we're going away.  this can
  // take some time but we don't want to wait too long.  we do
  // our own multiplexing loop and wait for a maximum of 3 seconds
  // total.

  // if we aren't supposed to be publicizing, skip the whole thing
  // and don't waste 3 seconds.
  if (!publicizeServer)
    return;

  queueMessage(ListServerLink::REMOVE);
  TimeKeeper start = TimeKeeper::getCurrent();
  do {
    // compute timeout
    float waitTime = 3.0f - (TimeKeeper::getCurrent() - start);
    if (waitTime <= 0.0f)
      break;
    if (!isConnected()) //queueMessage should have connected us
      break;
    // check for list server socket connection
    int fdMax = -1;
    fd_set write_set;
    fd_set read_set;
    FD_ZERO(&write_set);
    FD_ZERO(&read_set);
    if (phase == ListServerLink::CONNECTING)
      _FD_SET(linkSocket, &write_set);
    else
      _FD_SET(linkSocket, &read_set);
    fdMax = linkSocket;

    // wait for socket to connect or timeout
    struct timeval timeout;
    timeout.tv_sec = long(floorf(waitTime));
    timeout.tv_usec = long(1.0e+6f * (waitTime - floorf(waitTime)));
    int nfound = select(fdMax + 1, (fd_set*)&read_set, (fd_set*)&write_set,
			0, &timeout);
    if (nfound == 0)
      // Time has gone, close and go
      break;
    // check for connection to list server
    if (FD_ISSET(linkSocket, &write_set))
      sendQueuedMessages();
    else if (FD_ISSET(linkSocket, &read_set))
      read();
  } while (true);

  // stop list server communication
  closeLink();
}

void ListServerLink::closeLink()
{
  if (isConnected()) {
    close(linkSocket);
    DEBUG4("Closing List Server\n");
    linkSocket = NotConnected;
  }
}

void ListServerLink::read()
{
  if (isConnected()) {
    char buf[2048];
    int bytes = recv(linkSocket, buf, sizeof(buf)-1, 0);
    // TODO don't close unless we've got it all
    closeLink();
    buf[bytes]=0;
    char* base = buf;
    static char *tokGoodIdentifier = "TOKGOOD: ";
    static char *tokBadIdentifier = "TOKBAD: ";
    // walks entire reply including HTTP headers
    while (*base) {
      // find next newline
      char* scan = base;
      while (*scan && *scan != '\r' && *scan != '\n') scan++;
      // if no newline then no more complete replies
      if (*scan != '\r' && *scan != '\n') break;
      while (*scan && (*scan == '\r' || *scan == '\n')) *scan++ = '\0';
      DEBUG4("Got line: \"%s\"\n", base);
      // TODO don't do this if we don't want central logins
      if (strncmp(base, tokGoodIdentifier, strlen(tokGoodIdentifier)) == 0) {
	char *callsign, *group;
	callsign = base + strlen(tokGoodIdentifier);
	DEBUG3("Got: %s\n", base);
	group = callsign;
	while (*group && (*group != ':')) group++;
	while (*group && (*group == ':')) *group++ = 0;
	int playerIndex = getTarget(callsign);
	if (playerIndex < curMaxPlayers) {
	  GameKeeper::Player *playerData = GameKeeper::Player::getPlayerByIndex(playerIndex);
	  // don't accept the global auth if there's a local account of the same name
	  // and the local account is not marked as being the same as the global account
	  if (!playerData->accessInfo.hasRealPassword()
	      || playerData->accessInfo.getUserInfo(callsign).hasGroup("LOCAL.GLOBAL")) {
	    if (!playerData->accessInfo.isRegistered())
	      playerData->accessInfo.storeInfo(NULL);
	    playerData->accessInfo.setPermissionRights();
	    while (*group) {
	      char *nextgroup = group;
	      while (*nextgroup && (*nextgroup != ':')) nextgroup++;
	      while (*nextgroup && (*nextgroup == ':')) *nextgroup++ = 0;
	      playerData->accessInfo.addGroup(group);
	      //DEBUG3("Got: [%d] \"%s\" \"%s\"\n", playerIndex, callsign, group);
	      group = nextgroup;
	    }
	    // automatically give global users permission to use local accounts
	    // since they either already have it, or there's no existing local account.
	    playerData->accessInfo.addGroup("LOCAL.GLOBAL");
	    sendMessage(ServerPlayer, playerIndex, "Global login approved!");
	    sendIPUpdate(playerIndex, -1);
	    sendPlayerInfo();
	  } else {
	    sendMessage(ServerPlayer, playerIndex, "Global login rejected. "
			"This callsign is registered locally on this server.");
	    sendMessage(ServerPlayer, playerIndex, "If the local account is yours, "
			"please ask an admin to grant it membership in the LOCAL.GLOBAL group.");
	    sendMessage(ServerPlayer, playerIndex, "If it is not yours, please ask an admin "
			"to deregister it so that you may use your global callsign.");
	  }
	  playerData->player.clearToken();
	}
      } else if (strncmp(base, tokBadIdentifier, strlen(tokBadIdentifier)) == 0) {
	char *callsign;
	callsign = base + strlen(tokBadIdentifier);
	int playerIndex = getTarget(callsign);
	DEBUG3("Got: [%d] %s\n", playerIndex, base);
	if (playerIndex < curMaxPlayers) {
	  GameKeeper::Player *playerData = GameKeeper::Player::getPlayerByIndex(playerIndex);
	  sendMessage(ServerPlayer, playerIndex, "Global login rejected, bad token.");
	  playerData->player.clearToken();
	}
      }

      // next reply
      base = scan;
    }
    if (nextMessageType != ListServerLink::NONE) {
      // There was a pending request arrived after we write:
      // we should redo all the stuff
      openLink();
    }
  }
}

void ListServerLink::openLink()
{
  // start opening connection if not already doing so
  if (!isConnected()) {
    linkSocket = socket(AF_INET, SOCK_STREAM, 0);
    DEBUG4("Opening List Server\n");
    if (!isConnected()) {
      return;
    }

    // set to non-blocking for connect
    if (BzfNetwork::setNonBlocking(linkSocket) < 0) {
      closeLink();
      return;
    }

    // Make our connection come from our serverAddress in case we have
    // multiple/masked IPs so the list server can verify us.
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr = serverAddress;

    // assign the address to the socket
    if (bind(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
      closeLink();
      return;
    }

    // connect.  this should fail with EINPROGRESS but check for
    // success just in case.
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port   = htons(port);
    addr.sin_addr   = address;
    if (connect(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
#if defined(_WIN32)
#undef EINPROGRESS
#define EINPROGRESS EWOULDBLOCK
#endif
      if (getErrno() != EINPROGRESS) {
	nerror("connecting to list server");
	// try to lookup dns name again in case it moved
	this->address = Address::getHostAddress(this->hostname);
	closeLink();
      } else {
	phase = CONNECTING;
      }
    } else {
      // shouldn't arrive here. Just in case, clean
      DEBUG3("list server connect and close?\n");
      closeLink();
    }
  }
}

void ListServerLink::queueMessage(MessageType type)
{
  // ignore if the server is not public
  if (!publicizeServer) return;

  // Open network connection only if closed
  if (!isConnected()) openLink();

  // record next message to send.
  nextMessageType = type;
}

void ListServerLink::sendQueuedMessages()
{
  if (!isConnected())
    return;

  if (nextMessageType == ListServerLink::ADD) {
    DEBUG3("Queuing ADD message to list server\n");
    addMe(getTeamCounts(), publicizeAddress, TextUtils::url_encode(publicizeDescription));
    lastAddTime = TimeKeeper::getCurrent();
  } else if (nextMessageType == ListServerLink::REMOVE) {
    DEBUG3("Queuing REMOVE message to list server\n");
    removeMe(publicizeAddress);
  }
}

void ListServerLink::addMe(PingPacket pingInfo,
			   std::string publicizedAddress,
			   std::string publicizedTitle)
{
  std::string msg;

  // encode ping reply as ascii hex digits plus NULL
  char gameInfo[PingPacketHexPackedSize + 1];
  pingInfo.packHex(gameInfo);

  // TODO we probably should convert to a POST. List server now allows either
  // send ADD message (must send blank line)
  msg = TextUtils::format("GET %s?action=ADD&nameport=%s&version=%s&gameinfo=%s&build=%s",
    pathname.c_str(), publicizedAddress.c_str(),
    getServerVersion(), gameInfo,
    getAppVersion());
  msg += "&checktokens=";
  // callsign1@ip1=token1%0D%0Acallsign2@ip2=token2%0D%0A
  for (int i = 0; i < curMaxPlayers; i++) {
    GameKeeper::Player *playerData = GameKeeper::Player::getPlayerByIndex(i);
    if (!playerData)
      continue;
    NetHandler *handler = playerData->netHandler;
    std::string encodedCallsign = TextUtils::url_encode(playerData->player.getCallSign());
    if (strlen(playerData->player.getCallSign()) && strlen(playerData->player.getToken())) {
      msg += encodedCallsign;
      Address addr = handler->getIPAddress();
      if (!addr.isPrivate()) {
        msg += "@";
        msg += handler->getTargetIP();
      }
      msg += "=";
      msg += playerData->player.getToken();
      msg += "%0D%0A";
    }
  }
  // *groups=GROUP0%0D%0AGROUP1%0D%0A
  msg += "&groups=";
  PlayerAccessMap::iterator itr = groupAccess.begin();
  for ( ; itr != groupAccess.end(); itr++) {
    if (itr->first.substr(0, 6) != "LOCAL.") {
      msg += itr->first.c_str();
      msg += "%0D%0A";
    }
  }

  msg += TextUtils::format("&title=%s HTTP/1.1\r\n"
      "User-Agent: bzfs %s\r\n"
      "Host: %s\r\n"
      "Cache-Control: no-cache\r\n"
      "Connection: close\r\n"
      "\r\n",
    publicizedTitle.c_str(),
    getAppVersion(),
    hostname.c_str());
  sendLSMessage(msg);
}

void ListServerLink::removeMe(std::string publicizedAddress)
{
  std::string msg;
  // send REMOVE (must send blank line)
  msg = TextUtils::format("GET %s?action=REMOVE&nameport=%s HTTP/1.1\r\n"
    "User-Agent: bzfs %s\r\n"
    "Host: %s\r\n"
    "Cache-Control: no-cache\r\n"
    "Connection: close\r\n"
    "\r\n",
    pathname.c_str(),
    publicizedAddress.c_str(),
    getAppVersion(),
    hostname.c_str());
  sendLSMessage(msg);
}

void ListServerLink::sendLSMessage(std::string message)
{
  const int bufsize = 4096;
  char msg[bufsize];
  strncpy(msg, message.c_str(), bufsize);
  msg[bufsize - 1] = 0;
  if (strlen(msg) > 0) {
    DEBUG3("%s\n", msg);
    if (send(linkSocket, msg, strlen(msg), 0) == -1) {
      perror("List server send failed");
      DEBUG3("Unable to send to the list server!\n");
      closeLink();
    } else {
      nextMessageType = ListServerLink::NONE;
      phase	   = ListServerLink::WRITTEN;
    }
  } else {
    closeLink();
  }
}

// Local Variables: ***
// mode:C++ ***
// tab-width: 8 ***
// c-basic-offset: 2 ***
// indent-tabs-mode: t ***
// End: ***
// ex: shiftwidth=2 tabstop=8
<|endoftext|>


<commit_before>//================================
// ==========   Dialing Lemmatizer (www.aot.ru)
//================================

#pragma warning(disable:4786)
#pragma warning(disable:4530)

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <locale.h>

#include "../common/utilit.h"
#include "../AgramtabLib/RusGramTab.h"
#include "../AgramtabLib/EngGramTab.h"
#include "../AgramtabLib/GerGramTab.h"
#include "../LemmatizerLib/Lemmatizers.h"
#include "../LemmatizerLib/Paradigm.h"

bool				bPrintIds = true;;
bool				bPrintForms = false;
MorphLanguageEnum	Language;
CLemmatizer*		pLemmatizer;
CAgramtab*			pAgramtab;

string GetGrammems (const char* tab_str)
{
	QWORD G;
	pAgramtab->GetGrammems(tab_str, G);
	string s = pAgramtab->GrammemsToStr (G);
	if (!s.empty() && (s[s.length() -1] == ','))
		s.erase(s.length() - 1);
	return s;
}

std::string  GetMorphInfo (std::string  Form)
{

	std::string  Result;

	bool bCapital   = is_upper_alpha((BYTE)Form[0], Language);

	std::vector<CFormInfo> Paradigms;

	pLemmatizer->CreateParadigmCollection(false, Form, bCapital, true, Paradigms);

	for (int i=0; i < Paradigms.size(); i++)
	{
        if (i > 0)
            Result += "\t";
		const CFormInfo& F = Paradigms[i];
		Result += F.m_bFound ? "+ " : "- ";

		Result += F.GetWordForm(0) + " ";

		{
			string GramCodes =	F.GetSrcAncode();
			BYTE  PartOfSpeech = pAgramtab->GetPartOfSpeech (GramCodes.c_str());
			Result	+= pAgramtab->GetPartOfSpeechStr(PartOfSpeech) + std::string(" ") ;

			string  CommonAncode = F.GetCommonAncode();
			Result += Format ("%s ", (CommonAncode.empty()) ? "" : GetGrammems(CommonAncode.c_str()).c_str());

            for (long i=0; i < GramCodes.length(); i+=2) {
                if (i > 0)
                    Result += ";";
				Result += Format ("%s", GetGrammems(GramCodes.c_str()+i).c_str());
            }

	   }

		if (bPrintIds)
			Result += Format(" %i",F.GetParadigmId());

		BYTE Accent = F.GetSrcAccentedVowel();
		if (Accent != 0xff)
			Result += Format(" %s'%s", Form.substr(0, Accent+1).c_str(), Form.substr(Accent+1).c_str());

        if (bPrintForms) {
		    Result += " ";
		    for (int k=0; k < F.GetCount (); k++)
		    {
                if (k > 0)
                    Result += ",";
			    Result += Paradigms[i].GetWordForm(k);
		    };
        }
   	};

	return Result;
};

void ProcessFile(std::string Filename)
{
	printf ("read %s\n", Filename.c_str());

	StringVector Forms;
	{
		FILE* fp = fopen (Filename.c_str(), "r");
		if (!fp)
		{
			printf ("cannot open file %s\n", Filename.c_str());
			return;
		};
		char buffer[256];
		while (fgets(buffer,256, fp))
		{
			std::string q =  buffer;
			Trim(q);
			if (q.empty()) continue;
			RmlMakeUpper(q,pLemmatizer->GetLanguage());
			Forms.push_back(q);
		};
		fclose (fp);
	};
	size_t Count = Forms.size();

	StringVector Results;
	Results.resize(Count, string (255,' ') );
	printf ("process %s\n", Filename.c_str());
	clock_t start = clock();
	const size_t OutBufferSize = 2000;
	char OutBuffer[OutBufferSize];
	for (size_t i=0; i < Count; i++)
	{
		if (!pLemmatizer->GetAllAncodesAndLemmasQuick(Forms[i], true, OutBuffer, OutBufferSize, true))
		{
			printf ("Too many variants for %s\n", Forms[i].c_str());
		};
		//pLemmatizer->GetAllAncodesQuick((const BYTE*)Forms[i].c_str(), true, OutBuffer);
		Results[i] = (char*)OutBuffer;
	};
	long ticks = clock() - start;
	printf ("Count of words = %i\n", Count);
	double seconds = (double)ticks/(double)CLOCKS_PER_SEC;
	printf ("Time = %g seconds; %i ticks\n", seconds,ticks);
	if (seconds > 0)
		printf ("Speed = %g words per seconds\n", (((double)Count)/seconds));
	else
		printf ("too few words to measure the speed\n");

	std::string OutFileName = MakeFName(Filename,"lem");
	printf ("writing to %s\n",OutFileName.c_str());
	{
		FILE* out_fp = fopen (OutFileName.c_str(), "w");
		if (!out_fp)
		{
			printf ("cannot write file %s\n", OutFileName.c_str());
			return;
		};
		for (size_t i=0; i < Count; i++)
			fprintf (out_fp,"%s -> %s\n", Forms[i].c_str(), Results[i].c_str());
		fclose (out_fp);
	};

};

template <class  T, class Y>
bool InitMorphologySystem()
{
	std::string langua_str = GetStringByLanguage(Language);
	pLemmatizer = new T;
	string strError;
	if (!pLemmatizer->LoadDictionariesRegistry(strError))
	{
   		fprintf (stderr, "Cannot load %s  morphological dictionary\n", langua_str.c_str());
		fprintf (stderr, "Return error: %s\n", strError.c_str());
		return false;
	};
	pAgramtab = new Y;
	if (!pAgramtab->LoadFromRegistry())
	{
   		fprintf (stderr, "Cannot load %s  gramtab\n", langua_str.c_str() );
		return false;
	};
	return true;
};

void PrintUsage()
{
		printf ("Dialing Command Line Morphology (www.aot.ru)\n");
		printf ("Usage: TestLem (RUSSIAN|ENGLISH|GERMAN) [file]\n");
		printf ("if [file] is specified then TestLem builds a lem-file by [file] and exits\n");
		printf ("otherwise it reads from standard input\n");
		printf ("Example: TestLem Russian\n");
		printf ("Example: TestLem English test.txt\n");
		exit(-1);
};

int main(int argc, char **argv)
{

    vector<string> ARGV;
    bool bEchoInput = false;
    for (size_t i=1; i <  argc; i++)
    {
		string s = argv[i];
        if (s == "-h" || s == "-help" ||  s == "/h" || s == "/help") {
            PrintUsage();
        }
        else if (s == "-echo")
        {
            bEchoInput = true;
        }
		else if (s == "-noids")
		{
			bPrintIds  = false;
		}
		else if (s == "-forms")
		{
			bPrintForms  = true;
		}
        else {
            ARGV.push_back(s);
        }
    }
	if (ARGV.size() < 1)
		PrintUsage();

    if (!GetLanguageByString(ARGV[0], Language))
	{
		PrintUsage();
	};

	std::string FileName;
	if (ARGV.size() > 1)
		FileName = ARGV[1];

	// ===============  LOADING DICTS ================
	fprintf (stderr,"Loading..\n");
	bool bResult = false;
	switch (Language)
	{
		case morphRussian:
					bResult = InitMorphologySystem<CLemmatizerRussian, CRusGramTab>();
					break;
		case morphEnglish :
					bResult = InitMorphologySystem<CLemmatizerEnglish, CEngGramTab>();
					break;
		case morphGerman:
					bResult = InitMorphologySystem<CLemmatizerGerman, CGerGramTab>();
					break;
		default : PrintUsage();

	};
	if (!bResult)
		return 1;

	if (!FileName.empty())
	{
		ProcessFile(FileName);
		delete pLemmatizer;
		delete pAgramtab;
		return 0;
	};

	// ===============  WORKING ===============

	fprintf (stderr,"Input a word..\n");
    char Form[100];
	while (fgets(Form, 1000, stdin))
	{
		std::string s = Form;
		Trim(s);
		if (s.empty()) break;
		string Result = GetMorphInfo(s);

        if (bEchoInput)  {
            printf ("%s\t", s.c_str());
        }
		printf ("%s\n", Result.c_str());

	};
	delete pLemmatizer;
	delete pAgramtab;
	return 0;
}

<commit_msg>sort paradigm under -sort<commit_after>//================================
// ==========   Dialing Lemmatizer (www.aot.ru)
//================================

#pragma warning(disable:4786)
#pragma warning(disable:4530)

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <locale.h>

#include "../common/utilit.h"
#include "../AgramtabLib/RusGramTab.h"
#include "../AgramtabLib/EngGramTab.h"
#include "../AgramtabLib/GerGramTab.h"
#include "../LemmatizerLib/Lemmatizers.h"
#include "../LemmatizerLib/Paradigm.h"

bool				bPrintIds = true;;
bool				bPrintForms = false;
bool                bSortParadigms = false;
MorphLanguageEnum	Language;
CLemmatizer*		pLemmatizer;
CAgramtab*			pAgramtab;

string GetGrammems (const char* tab_str)
{
	QWORD G;
	pAgramtab->GetGrammems(tab_str, G);
	string s = pAgramtab->GrammemsToStr (G);
	if (!s.empty() && (s[s.length() -1] == ','))
		s.erase(s.length() - 1);
	return s;
}

std::string  GetMorphInfo (std::string  Form)
{

	bool bCapital   = is_upper_alpha((BYTE)Form[0], Language);

	std::vector<CFormInfo> Paradigms;

	pLemmatizer->CreateParadigmCollection(false, Form, bCapital, true, Paradigms);

	std::vector< std::string>  Results;
	for (int i=0; i < Paradigms.size(); i++)
	{
        std::string Result;
		const CFormInfo& F = Paradigms[i];
		Result += F.m_bFound ? "+ " : "- ";

		Result += F.GetWordForm(0) + " ";

		{
			string GramCodes =	F.GetSrcAncode();
			BYTE  PartOfSpeech = pAgramtab->GetPartOfSpeech (GramCodes.c_str());
			Result	+= pAgramtab->GetPartOfSpeechStr(PartOfSpeech) + std::string(" ") ;

			string  CommonAncode = F.GetCommonAncode();
			Result += Format ("%s ", (CommonAncode.empty()) ? "" : GetGrammems(CommonAncode.c_str()).c_str());

            for (long i=0; i < GramCodes.length(); i+=2) {
                if (i > 0)
                    Result += ";";
				Result += Format ("%s", GetGrammems(GramCodes.c_str()+i).c_str());
            }

	   }

		if (bPrintIds)
			Result += Format(" %i",F.GetParadigmId());

		BYTE Accent = F.GetSrcAccentedVowel();
		if (Accent != 0xff)
			Result += Format(" %s'%s", Form.substr(0, Accent+1).c_str(), Form.substr(Accent+1).c_str());

        if (bPrintForms) {
		    Result += " ";
		    for (int k=0; k < F.GetCount (); k++)
		    {
                if (k > 0)
                    Result += ",";
			    Result += Paradigms[i].GetWordForm(k);
		    };
        };
        Results.push_back(Result);
   	};

    if (bSortParadigms) {
        std::sort( Results.begin(), Results.end() );
    };
    std::string Result;
	for (int i=0; i < Results.size(); i++) {
        if (i > 0)
            Result += "\t";
        Result  +=  Results[i] + "\n";
    }
	return Result;
};

void ProcessFile(std::string Filename)
{
	printf ("read %s\n", Filename.c_str());

	StringVector Forms;
	{
		FILE* fp = fopen (Filename.c_str(), "r");
		if (!fp)
		{
			printf ("cannot open file %s\n", Filename.c_str());
			return;
		};
		char buffer[256];
		while (fgets(buffer,256, fp))
		{
			std::string q =  buffer;
			Trim(q);
			if (q.empty()) continue;
			RmlMakeUpper(q,pLemmatizer->GetLanguage());
			Forms.push_back(q);
		};
		fclose (fp);
	};
	size_t Count = Forms.size();

	StringVector Results;
	Results.resize(Count, string (255,' ') );
	printf ("process %s\n", Filename.c_str());
	clock_t start = clock();
	const size_t OutBufferSize = 2000;
	char OutBuffer[OutBufferSize];
	for (size_t i=0; i < Count; i++)
	{
		if (!pLemmatizer->GetAllAncodesAndLemmasQuick(Forms[i], true, OutBuffer, OutBufferSize, true))
		{
			printf ("Too many variants for %s\n", Forms[i].c_str());
		};
		//pLemmatizer->GetAllAncodesQuick((const BYTE*)Forms[i].c_str(), true, OutBuffer);
		Results[i] = (char*)OutBuffer;
	};
	long ticks = clock() - start;
	printf ("Count of words = %i\n", Count);
	double seconds = (double)ticks/(double)CLOCKS_PER_SEC;
	printf ("Time = %g seconds; %i ticks\n", seconds,ticks);
	if (seconds > 0)
		printf ("Speed = %g words per seconds\n", (((double)Count)/seconds));
	else
		printf ("too few words to measure the speed\n");

	std::string OutFileName = MakeFName(Filename,"lem");
	printf ("writing to %s\n",OutFileName.c_str());
	{
		FILE* out_fp = fopen (OutFileName.c_str(), "w");
		if (!out_fp)
		{
			printf ("cannot write file %s\n", OutFileName.c_str());
			return;
		};
		for (size_t i=0; i < Count; i++)
			fprintf (out_fp,"%s -> %s\n", Forms[i].c_str(), Results[i].c_str());
		fclose (out_fp);
	};

};

template <class  T, class Y>
bool InitMorphologySystem()
{
	std::string langua_str = GetStringByLanguage(Language);
	pLemmatizer = new T;
	string strError;
	if (!pLemmatizer->LoadDictionariesRegistry(strError))
	{
   		fprintf (stderr, "Cannot load %s  morphological dictionary\n", langua_str.c_str());
		fprintf (stderr, "Return error: %s\n", strError.c_str());
		return false;
	};
	pAgramtab = new Y;
	if (!pAgramtab->LoadFromRegistry())
	{
   		fprintf (stderr, "Cannot load %s  gramtab\n", langua_str.c_str() );
		return false;
	};
	return true;
};

void PrintUsage()
{
		printf ("Dialing Command Line Morphology (www.aot.ru)\n");
		printf ("Usage: TestLem (RUSSIAN|ENGLISH|GERMAN) [file]\n");
		printf ("if [file] is specified then TestLem builds a lem-file by [file] and exits\n");
		printf ("otherwise it reads from standard input\n");
		printf ("Example: TestLem Russian\n");
		printf ("Example: TestLem English test.txt\n");
		exit(-1);
};

int main(int argc, char **argv)
{

    vector<string> ARGV;
    bool bEchoInput = false;
    for (size_t i=1; i <  argc; i++)
    {
		string s = argv[i];
        if (s == "-h" || s == "-help" ||  s == "/h" || s == "/help") {
            PrintUsage();
        }
        else if (s == "-echo")
        {
            bEchoInput = true;
        }
		else if (s == "-noids")
		{
			bPrintIds  = false;
		}
		else if (s == "-sort")
		{
			bSortParadigms  = true;
		}
		else if (s == "-forms")
		{
			bPrintForms  = true;
		}
        else {
            ARGV.push_back(s);
        }
    }
	if (ARGV.size() < 1)
		PrintUsage();

    if (!GetLanguageByString(ARGV[0], Language))
	{
		PrintUsage();
	};

	std::string FileName;
	if (ARGV.size() > 1)
		FileName = ARGV[1];

	// ===============  LOADING DICTS ================
	fprintf (stderr,"Loading..\n");
	bool bResult = false;
	switch (Language)
	{
		case morphRussian:
					bResult = InitMorphologySystem<CLemmatizerRussian, CRusGramTab>();
					break;
		case morphEnglish :
					bResult = InitMorphologySystem<CLemmatizerEnglish, CEngGramTab>();
					break;
		case morphGerman:
					bResult = InitMorphologySystem<CLemmatizerGerman, CGerGramTab>();
					break;
		default : PrintUsage();

	};
	if (!bResult)
		return 1;

	if (!FileName.empty())
	{
		ProcessFile(FileName);
		delete pLemmatizer;
		delete pAgramtab;
		return 0;
	};

	// ===============  WORKING ===============

	fprintf (stderr,"Input a word..\n");
    char Form[100];
	while (fgets(Form, 1000, stdin))
	{
		std::string s = Form;
		Trim(s);
		if (s.empty()) break;
		string Result = GetMorphInfo(s);

        if (bEchoInput)  {
            printf ("%s\t", s.c_str());
        }
		printf ("%s\n", Result.c_str());

	};
	delete pLemmatizer;
	delete pAgramtab;
	return 0;
}

<|endoftext|>


<commit_before>/**
 * @file confusion_matrix.cpp
 * @author Sean Massung
 */

#include <algorithm>
#include <iomanip>
#include <vector>

#include "classify/confusion_matrix.h"
#include "util/mapping.h"
#include "util/printing.h"

namespace meta
{
namespace classify
{

confusion_matrix::confusion_matrix()
    : predictions_{32, string_pair_hash}, total_{0}
{
    /* nothing */
}

void confusion_matrix::add(const class_label& predicted,
                           const class_label& actual, size_t times)
{
    std::pair<class_label, class_label> prediction{predicted, actual};
    predictions_[prediction] += times;
    counts_[actual] += times;
    classes_.insert(predicted);
    classes_.insert(actual);
    total_ += times;
}

void confusion_matrix::print_result_pairs(std::ostream& out) const
{
    for (auto& p : predictions_)
        for (size_t i = 0; i < p.second; ++i)
            out << p.first.first << " " << p.first.second << "\n";
}

void confusion_matrix::print(std::ostream& out) const
{
    auto max_label = std::max_element(classes_.begin(), classes_.end(),
        [](const class_label& a, const class_label& b)
        {
            return static_cast<std::string>(a).size() <
                static_cast<std::string>(b).size();
        });

    size_t w = std::max<size_t>(7, static_cast<std::string>(*max_label).size())
               + 2;
    out << std::left << std::endl << std::setw(w) << "";
    out << "  ";
    for (auto& aClass : classes_)
        out << std::setw(w) << aClass;
    out << std::endl;
    out << std::string(w, ' ') << std::string(classes_.size() * w, '-')
        << std::endl;

    for (auto& aClass : classes_)
    {
        out << std::setw(w - 1) << std::right << aClass << std::left << " | ";
        for (auto& pClass : classes_)
        {
            auto predIt = predictions_.find(std::make_pair(pClass, aClass));
            if (predIt != predictions_.end())
            {
                size_t numPred = predIt->second;
                double percent = static_cast<double>(numPred)
                                 / counts_.at(aClass);
                out.precision(3);
                std::stringstream ss;
                ss.precision(3);
                ss << percent;
                if (aClass == pClass)
                {
                    auto str = printing::make_bold(ss.str());
                    auto diff = str.size() - ss.str().size();
                    out << std::setw(w + diff) << str;
                }
                else
                {
                    out << std::setw(w) << ss.str();
                }
            }
            else
                out << std::setw(w) << "- ";
        }
        out << std::endl;
    }
    out << std::endl;
}

void confusion_matrix::print_class_stats(std::ostream& out,
                                         const class_label& label, double& prec,
                                         double& rec, double& f1,
                                         size_t width) const
{
    for (auto& cls : classes_)
    {
        prec += map::safe_at(predictions_, std::make_pair(cls, label));
        rec += map::safe_at(predictions_, std::make_pair(label, cls));
    }

    double correct = map::safe_at(predictions_, std::make_pair(label, label));

    if (rec != 0.0)
        rec = correct / rec;

    if (prec != 0.0)
        prec = correct / prec;

    if (prec + rec != 0.0)
        f1 = (2.0 * prec * rec) / (prec + rec);

    auto w1 = std::setw(width);
    auto w2 = std::setw(12);
    out << std::left << w1 << label << std::left << w2 << f1 << std::left << w2
        << prec << std::left << w2 << rec << std::endl;
}

void confusion_matrix::print_stats(std::ostream& out) const
{
    double t_prec = 0.0;
    double t_rec = 0.0;
    double t_f1 = 0.0;
    double t_corr = 0.0;

    auto max_label
        = std::max_element(classes_.begin(), classes_.end(),
            [](const class_label& a, const class_label& b)
            {
                return static_cast<std::string>(a).size()
                    < static_cast<std::string>(b).size();
            });

    size_t width = std::max<size_t>(
            12, static_cast<std::string>(*max_label).size() + 2);

    auto w1 = std::setw(width + printing::make_bold("").size());
    auto w2 = std::setw(12 + printing::make_bold("").size());
    out.precision(3);
    out << std::string(width + 12 * 3, '-') << std::endl << std::left << w1
        << printing::make_bold("Class") << std::left << w2
        << printing::make_bold("F1 Score") << std::left << w2
        << printing::make_bold("Precision") << std::left << w2
        << printing::make_bold("Recall") << std::endl
        << std::string(width + 12 * 3, '-') << std::endl;

    for (auto& cls : classes_)
    {
        double prec = 0.0, rec = 0.0, f1 = 0.0;
        auto it = predictions_.find(std::make_pair(cls, cls));
        if (it != predictions_.end())
            t_corr += it->second;
        print_class_stats(out, cls, prec, rec, f1, width);
        t_prec += prec;
        t_rec += rec;
        t_f1 += f1;
    }

    auto limit = [](double val)
    {
        std::stringstream ss;
        ss.precision(3);
        ss << val;
        return ss.str();
    };

    out << std::string(width + 12 * 3, '-') << std::endl << w1
        << printing::make_bold("Total") << w2
        << printing::make_bold(limit(t_f1 / classes_.size())) << w2
        << printing::make_bold(limit(t_prec / classes_.size())) << w2
        << printing::make_bold(limit(t_rec / classes_.size())) << std::endl
        << std::string(width + 12 * 3, '-') << std::endl << total_
        << " predictions attempted, overall accuracy: " << t_corr / total_
        << std::endl;
}

double confusion_matrix::accuracy() const
{
    double correct = 0.0;
    for (auto& cls : classes_)
        correct += map::safe_at(predictions_, std::make_pair(cls, cls));
    return correct / total_;
}

size_t confusion_matrix::string_pair_hash(const std::pair
                                          <std::string, std::string>& str_pair)
{
    return std::hash<std::string>()(str_pair.first + str_pair.second);
}

const confusion_matrix::prediction_counts& confusion_matrix::predictions() const
{
    return predictions_;
}

confusion_matrix confusion_matrix::operator+(const confusion_matrix
                                             & other) const
{
    confusion_matrix sum{*this};

    for (auto& pred : other.predictions())
        sum.add(pred.first.first, pred.first.second, pred.second);

    return sum;
}

confusion_matrix& confusion_matrix::operator+=(const confusion_matrix& other)
{
    *this = *this + other;
    return *this;
}

bool confusion_matrix::mcnemar_significant(const confusion_matrix& a,
                                           const confusion_matrix& b)
{
    std::set<class_label> classes = a.classes_;
    classes.insert(b.classes_.begin(), b.classes_.end());

    double a_adv = 0;
    double b_adv = 0;

    for (auto& cls : classes)
    {
        auto a_count = map::safe_at(a.predictions_, std::make_pair(cls, cls));
        auto b_count = map::safe_at(b.predictions_, std::make_pair(cls, cls));
        if (a_count > b_count)
            a_adv += (a_count - b_count);
        else if (b_count > a_count)
            b_adv += (b_count - a_count);
    }

    // does not approximate well if less than 25 samples
    if (a_adv + b_adv < 25)
        return false;

    double numerator = std::abs(a_adv - b_adv) - 0.5;
    double chi_square = (numerator * numerator) / (a_adv + b_adv);

    // check if significant with chi square, 1 degree of freedom, alpha == 0.05
    return chi_square > 3.84;
}
}
}
<commit_msg>fix bug where precision in stream is not restored to original value<commit_after>/**
 * @file confusion_matrix.cpp
 * @author Sean Massung
 */

#include <algorithm>
#include <iomanip>
#include <vector>

#include "classify/confusion_matrix.h"
#include "util/mapping.h"
#include "util/printing.h"

namespace meta
{
namespace classify
{

confusion_matrix::confusion_matrix()
    : predictions_{32, string_pair_hash}, total_{0}
{
    /* nothing */
}

void confusion_matrix::add(const class_label& predicted,
                           const class_label& actual, size_t times)
{
    std::pair<class_label, class_label> prediction{predicted, actual};
    predictions_[prediction] += times;
    counts_[actual] += times;
    classes_.insert(predicted);
    classes_.insert(actual);
    total_ += times;
}

void confusion_matrix::print_result_pairs(std::ostream& out) const
{
    for (auto& p : predictions_)
        for (size_t i = 0; i < p.second; ++i)
            out << p.first.first << " " << p.first.second << "\n";
}

void confusion_matrix::print(std::ostream& out) const
{
    auto saved_precision = out.precision();

    auto max_label = std::max_element(classes_.begin(), classes_.end(),
        [](const class_label& a, const class_label& b)
        {
            return static_cast<std::string>(a).size() <
                static_cast<std::string>(b).size();
        });

    size_t w = std::max<size_t>(7, static_cast<std::string>(*max_label).size())
               + 2;
    out << std::left << std::endl << std::setw(w) << "";
    out << "  ";
    for (auto& aClass : classes_)
        out << std::setw(w) << aClass;
    out << std::endl;
    out << std::string(w, ' ') << std::string(classes_.size() * w, '-')
        << std::endl;

    for (auto& aClass : classes_)
    {
        out << std::setw(w - 1) << std::right << aClass << std::left << " | ";
        for (auto& pClass : classes_)
        {
            auto predIt = predictions_.find(std::make_pair(pClass, aClass));
            if (predIt != predictions_.end())
            {
                size_t numPred = predIt->second;
                double percent = static_cast<double>(numPred)
                                 / counts_.at(aClass);
                out.precision(3);
                std::stringstream ss;
                ss.precision(3);
                ss << percent;
                if (aClass == pClass)
                {
                    auto str = printing::make_bold(ss.str());
                    auto diff = str.size() - ss.str().size();
                    out << std::setw(w + diff) << str;
                }
                else
                {
                    out << std::setw(w) << ss.str();
                }
            }
            else
                out << std::setw(w) << "- ";
        }
        out << std::endl;
    }
    out << std::endl;

    out.precision(saved_precision);
}

void confusion_matrix::print_class_stats(std::ostream& out,
                                         const class_label& label, double& prec,
                                         double& rec, double& f1,
                                         size_t width) const
{
    for (auto& cls : classes_)
    {
        prec += map::safe_at(predictions_, std::make_pair(cls, label));
        rec += map::safe_at(predictions_, std::make_pair(label, cls));
    }

    double correct = map::safe_at(predictions_, std::make_pair(label, label));

    if (rec != 0.0)
        rec = correct / rec;

    if (prec != 0.0)
        prec = correct / prec;

    if (prec + rec != 0.0)
        f1 = (2.0 * prec * rec) / (prec + rec);

    auto w1 = std::setw(width);
    auto w2 = std::setw(12);
    out << std::left << w1 << label << std::left << w2 << f1 << std::left << w2
        << prec << std::left << w2 << rec << std::endl;
}

void confusion_matrix::print_stats(std::ostream& out) const
{
    double t_prec = 0.0;
    double t_rec = 0.0;
    double t_f1 = 0.0;
    double t_corr = 0.0;
    auto saved_precision = out.precision();

    auto max_label
        = std::max_element(classes_.begin(), classes_.end(),
            [](const class_label& a, const class_label& b)
            {
                return static_cast<std::string>(a).size()
                    < static_cast<std::string>(b).size();
            });

    size_t width = std::max<size_t>(
            12, static_cast<std::string>(*max_label).size() + 2);

    auto w1 = std::setw(width + printing::make_bold("").size());
    auto w2 = std::setw(12 + printing::make_bold("").size());
    out.precision(3);
    out << std::string(width + 12 * 3, '-') << std::endl << std::left << w1
        << printing::make_bold("Class") << std::left << w2
        << printing::make_bold("F1 Score") << std::left << w2
        << printing::make_bold("Precision") << std::left << w2
        << printing::make_bold("Recall") << std::endl
        << std::string(width + 12 * 3, '-') << std::endl;

    for (auto& cls : classes_)
    {
        double prec = 0.0, rec = 0.0, f1 = 0.0;
        auto it = predictions_.find(std::make_pair(cls, cls));
        if (it != predictions_.end())
            t_corr += it->second;
        print_class_stats(out, cls, prec, rec, f1, width);
        t_prec += prec;
        t_rec += rec;
        t_f1 += f1;
    }

    auto limit = [](double val)
    {
        std::stringstream ss;
        ss.precision(3);
        ss << val;
        return ss.str();
    };

    out << std::string(width + 12 * 3, '-') << std::endl << w1
        << printing::make_bold("Total") << w2
        << printing::make_bold(limit(t_f1 / classes_.size())) << w2
        << printing::make_bold(limit(t_prec / classes_.size())) << w2
        << printing::make_bold(limit(t_rec / classes_.size())) << std::endl
        << std::string(width + 12 * 3, '-') << std::endl << total_
        << " predictions attempted, overall accuracy: " << t_corr / total_
        << std::endl;

    out.precision(saved_precision);
}

double confusion_matrix::accuracy() const
{
    double correct = 0.0;
    for (auto& cls : classes_)
        correct += map::safe_at(predictions_, std::make_pair(cls, cls));
    return correct / total_;
}

size_t confusion_matrix::string_pair_hash(const std::pair
                                          <std::string, std::string>& str_pair)
{
    return std::hash<std::string>()(str_pair.first + str_pair.second);
}

const confusion_matrix::prediction_counts& confusion_matrix::predictions() const
{
    return predictions_;
}

confusion_matrix confusion_matrix::operator+(const confusion_matrix
                                             & other) const
{
    confusion_matrix sum{*this};

    for (auto& pred : other.predictions())
        sum.add(pred.first.first, pred.first.second, pred.second);

    return sum;
}

confusion_matrix& confusion_matrix::operator+=(const confusion_matrix& other)
{
    *this = *this + other;
    return *this;
}

bool confusion_matrix::mcnemar_significant(const confusion_matrix& a,
                                           const confusion_matrix& b)
{
    std::set<class_label> classes = a.classes_;
    classes.insert(b.classes_.begin(), b.classes_.end());

    double a_adv = 0;
    double b_adv = 0;

    for (auto& cls : classes)
    {
        auto a_count = map::safe_at(a.predictions_, std::make_pair(cls, cls));
        auto b_count = map::safe_at(b.predictions_, std::make_pair(cls, cls));
        if (a_count > b_count)
            a_adv += (a_count - b_count);
        else if (b_count > a_count)
            b_adv += (b_count - a_count);
    }

    // does not approximate well if less than 25 samples
    if (a_adv + b_adv < 25)
        return false;

    double numerator = std::abs(a_adv - b_adv) - 0.5;
    double chi_square = (numerator * numerator) / (a_adv + b_adv);

    // check if significant with chi square, 1 degree of freedom, alpha == 0.05
    return chi_square > 3.84;
}
}
}
<|endoftext|>


<commit_before>/*
 *  qicomplex.cpp
 *
 *  Copyright (c) 2015 Tobias Wood.
 *
 *  This Source Code Form is subject to the terms of the Mozilla Public
 *  License, v. 2.0. If a copy of the MPL was not distributed with this
 *  file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 */

#include <iostream>

#include "itkImage.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkImageSliceConstIteratorWithIndex.h"
#include "itkImageSliceIteratorWithIndex.h"
#include "itkMultiplyImageFilter.h"
#include "itkComplexToPhaseImageFilter.h"
#include "itkComplexToModulusImageFilter.h"
#include "itkComplexToRealImageFilter.h"
#include "itkComplexToImaginaryImageFilter.h"
#include "itkComposeImageFilter.h"
#include "itkMagnitudeAndPhaseToComplexImageFilter.h"
#include "itkRegionOfInterestImageFilter.h"

#include "Util.h"
#include "ImageIO.h"
#include "Args.h"

namespace itk {
template<typename TImage>
class NegateFilter : public InPlaceImageFilter<TImage> {
public:
    typedef NegateFilter                Self;
    typedef InPlaceImageFilter<TImage>  Superclass;
    typedef SmartPointer<Self>          Pointer;
    typedef typename TImage::RegionType        TRegion;
    typedef typename TImage::InternalPixelType TPixel;

    itkNewMacro(Self);
    itkTypeMacro(NegateFilter, InPlaceImageFilter);

private:
    NegateFilter(const Self &); //purposely not implemented
    void operator=(const Self &);  //purposely not implemented
    bool m_all = false, m_alternate = false;
protected:
    NegateFilter() {}
    ~NegateFilter() {}

public:
    void SetNegate(const bool all, const bool alternate) {
        m_all = all;
        m_alternate = alternate;
    }
    void ThreadedGenerateData(const TRegion &region, ThreadIdType threadId) ITK_OVERRIDE {
        typedef typename TImage::PixelType PixelType;
        ImageSliceConstIteratorWithIndex<TImage> inIt(this->GetInput(), region);
        ImageSliceIteratorWithIndex<TImage>      outIt(this->GetOutput(), region);

        inIt.SetFirstDirection(0);
        inIt.SetSecondDirection(1);
        outIt.SetFirstDirection(0);
        outIt.SetSecondDirection(1);
        inIt.GoToBegin();
        outIt.GoToBegin();
        PixelType mult(m_all ? -1 : 1, 0);
        while(!inIt.IsAtEnd()) {
            while (!inIt.IsAtEndOfSlice()) {
                while (!inIt.IsAtEndOfLine()) {
                    outIt.Set(mult * inIt.Get());
                    ++inIt;
                    ++outIt;
                }
                inIt.NextLine();
                outIt.NextLine();
            }
            if (m_alternate) {
                mult = mult * PixelType(-1, 0);
            }
            inIt.NextSlice();
            outIt.NextSlice();
        }
        // std::std::cout << "End " << __PRETTY_FUNCTION__ << std::std::endl;
    }
};

} // End namespace itk

/* Arguments defined here so they are available in the templated run function */
args::ArgumentParser parser(
    "Input is specified with lower case letters. A valid combination of inputs "
    " must be specified, e.g. real & imaginary or magnitude & phase.\n"
    "Output is specified with upper case letters. Any combination can be given\n"
    "http://github.com/spinicist/QUIT");

args::ValueFlag<std::string> in_mag(parser, "IN_MAG", "Input magnitude file", {'m', "mag"});
args::ValueFlag<std::string> in_pha(parser, "IN_PHA", "Input phase file", {'p', "pha"});
args::ValueFlag<std::string> in_real(parser, "IN_REAL", "Input real file", {'r', "real"});
args::ValueFlag<std::string> in_imag(parser, "IN_IMAG", "Input imaginary file", {'i', "imag"});
args::ValueFlag<std::string> in_complex(parser, "IN_CPLX", "Input complex file", {'x', "complex"});
args::ValueFlag<std::string> in_realimag(parser, "IN_REALIMAG", "Input real & imaginary file", {'l', "realimag"});

args::ValueFlag<std::string> out_mag(parser, "OUT_MAG", "Output magnitude file", {'M', "MAG"});
args::ValueFlag<std::string> out_pha(parser, "OUT_PHA", "Output phase file", {'P', "PHA"});
args::ValueFlag<std::string> out_real(parser, "OUT_REAL", "Output real file", {'R', "REAL"});
args::ValueFlag<std::string> out_imag(parser, "OUT_IMAG", "Output imaginary file", {'I', "IMAG"});
args::ValueFlag<std::string> out_complex(parser, "OUT_CPLX", "Output complex file", {'X', "COMPLEX"});

args::HelpFlag help(parser, "HELP", "Show this help menu", {'h', "help"});
args::Flag     verbose(parser, "VERBOSE", "Print more information", {'v', "verbose"});
args::Flag     use_double(parser, "DOUBLE", "Process & output at double precision", {'d', "double"});
args::Flag     fixge(parser, "FIX_GE", "Negate alternate slices (fixes lack of FFT shift)", {"fixge"});
args::Flag     negate(parser, "NEGATE", "Negate entire volume", {"negate"});

template<typename TPixel>
void Run() {
    typedef itk::Image<TPixel, 4>          TImage;
    typedef itk::Image<std::complex<TPixel>, 4> TXImage;
    typedef itk::ImageFileWriter<TImage>   TWriter;

    typename TImage::Pointer img1 = ITK_NULLPTR, img2 = ITK_NULLPTR;
    typename TXImage::Pointer imgX = ITK_NULLPTR;

    if (in_real) {
        if (verbose) std::cout << "Reading real file: " << in_real.Get() << std::endl;
        img1 = QI::ReadImage<TImage>(in_real.Get());
        if (in_imag) {
            if (verbose) std::cout << "Reading imaginary file: " << in_imag.Get() << std::endl;
            img2 = QI::ReadImage<TImage>(in_imag.Get());
        } else {
            QI_FAIL("Must set real and imaginary inputs together");
        }
        auto compose = itk::ComposeImageFilter<TImage, TXImage>::New();
        compose->SetInput(0, img1);
        compose->SetInput(1, img2);
        compose->Update();
        imgX = compose->GetOutput();
        imgX->DisconnectPipeline();
    } else if (in_mag) {
        if (verbose) std::cout << "Reading magnitude file: " << in_mag.Get() << std::endl;
        img1 = QI::ReadImage<TImage>(in_mag.Get());
        if (in_pha) {
            if (verbose) std::cout << "Reading phase file: " << in_pha.Get() << std::endl;
            img2 = QI::ReadImage<TImage>(in_pha.Get());
        } else {
            QI_FAIL("Must set magnitude and phase inputs together");
        }
        auto compose = itk::MagnitudeAndPhaseToComplexImageFilter<TImage, TImage, TXImage>::New();
        compose->SetInput(0, img1);
        compose->SetInput(1, img2);
        compose->Update();
        imgX = compose->GetOutput();
        imgX->DisconnectPipeline();
    } else if (in_complex) {
        if (verbose) std::cout << "Reading complex file: " << in_complex.Get() << std::endl;
        imgX = QI::ReadImage<TXImage>(in_complex.Get());
    } else if (in_realimag) {
        if (verbose) std::cout << "Reading real/imaginary file: " << in_realimag.Get() << std::endl;
        auto img_both = QI::ReadImage<TImage>(in_realimag.Get());
        auto real_region = img_both->GetLargestPossibleRegion();
        auto imag_region = img_both->GetLargestPossibleRegion();
        real_region.GetModifiableSize()[3] = real_region.GetSize()[3] / 2;
        imag_region.GetModifiableSize()[3] = real_region.GetSize()[3];
        imag_region.GetModifiableIndex()[3] = real_region.GetSize()[3];
        auto extract_real = itk::RegionOfInterestImageFilter<TImage, TImage>::New();
        extract_real->SetRegionOfInterest(real_region);
        extract_real->SetInput(img_both);
        extract_real->Update();
        auto extract_imag = itk::RegionOfInterestImageFilter<TImage, TImage>::New();
        extract_imag->SetRegionOfInterest(imag_region);
        extract_imag->SetInput(img_both);
        extract_imag->Update();
        /* Hack round ITK changing the origin */
        extract_imag->GetOutput()->SetOrigin(extract_real->GetOutput()->GetOrigin());
        auto compose = itk::ComposeImageFilter<TImage, TXImage>::New();
        compose->SetInput(0, extract_real->GetOutput());
        compose->SetInput(1, extract_imag->GetOutput());
        compose->Update();
        imgX = compose->GetOutput();
        imgX->DisconnectPipeline();
    } else {
        QI_FAIL("No input files specified, use --help to see usage");
    }

    if (fixge || negate) {
        if (verbose && negate) std::cout << "Negating values" << std::endl;
        if (verbose && fixge)  std::cout << "Fixing GE lack of FFT-shift bug" << std::endl;
        auto fix = itk::NegateFilter<TXImage>::New();
        fix->SetInput(imgX);
        fix->SetNegate(negate, fixge);
        fix->Update();
        imgX = fix->GetOutput();
        imgX->DisconnectPipeline();
    }

    if (verbose) std::cout << "Writing output files" << std::endl;
    typename TWriter::Pointer write = TWriter::New();

    if (out_mag) {
        auto o = itk::ComplexToModulusImageFilter<TXImage, TImage>::New();
        o->SetInput(imgX);
        o->Update();
        QI::WriteImage(o->GetOutput(), out_mag.Get());
        if (verbose) std::cout << "Wrote magnitude image " << out_mag.Get() << std::endl;
    }
    if (out_pha) {
        auto o = itk::ComplexToPhaseImageFilter<TXImage, TImage>::New();
        o->SetInput(imgX);
        o->Update();
        QI::WriteImage(o->GetOutput(), out_pha.Get());
        if (verbose) std::cout << "Wrote phase image " << out_pha.Get() << std::endl;
    }
    if (out_real) {
        auto o = itk::ComplexToRealImageFilter<TXImage, TImage>::New();
        o->SetInput(imgX);
        o->Update();
        QI::WriteImage(o->GetOutput(), out_real.Get());
        if (verbose) std::cout << "Wrote real image " << out_real.Get() << std::endl;
    }
    if (out_imag) {
        auto o = itk::ComplexToImaginaryImageFilter<TXImage, TImage>::New();
        o->SetInput(imgX);
        o->Update();
        QI::WriteImage(o->GetOutput(), out_imag.Get());
        if (verbose) std::cout << "Wrote imaginary image " << out_imag.Get() << std::endl;
    }
    if (out_complex) {
        if (verbose) std::cout << "Writing complex image " << out_complex.Get() << std::endl;
        QI::WriteImage(imgX, out_complex.Get());
    }
}

int main(int argc, char **argv) {
    QI::ParseArgs(parser, argc, argv);
    if (use_double) {
        if (verbose) std::cout << "Using double precision" << std::endl;
        Run<double>();
    } else {
        if (verbose) std::cout << "Using float precision" << std::endl;
        Run<float>();
    }
    return EXIT_SUCCESS;
}
<commit_msg>STYLE: Changed order of arguments<commit_after>/*
 *  qicomplex.cpp
 *
 *  Copyright (c) 2015 Tobias Wood.
 *
 *  This Source Code Form is subject to the terms of the Mozilla Public
 *  License, v. 2.0. If a copy of the MPL was not distributed with this
 *  file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 */

#include <iostream>

#include "itkImage.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkImageSliceConstIteratorWithIndex.h"
#include "itkImageSliceIteratorWithIndex.h"
#include "itkMultiplyImageFilter.h"
#include "itkComplexToPhaseImageFilter.h"
#include "itkComplexToModulusImageFilter.h"
#include "itkComplexToRealImageFilter.h"
#include "itkComplexToImaginaryImageFilter.h"
#include "itkComposeImageFilter.h"
#include "itkMagnitudeAndPhaseToComplexImageFilter.h"
#include "itkRegionOfInterestImageFilter.h"

#include "Util.h"
#include "ImageIO.h"
#include "Args.h"

namespace itk {
template<typename TImage>
class NegateFilter : public InPlaceImageFilter<TImage> {
public:
    typedef NegateFilter                Self;
    typedef InPlaceImageFilter<TImage>  Superclass;
    typedef SmartPointer<Self>          Pointer;
    typedef typename TImage::RegionType        TRegion;
    typedef typename TImage::InternalPixelType TPixel;

    itkNewMacro(Self);
    itkTypeMacro(NegateFilter, InPlaceImageFilter);

private:
    NegateFilter(const Self &); //purposely not implemented
    void operator=(const Self &);  //purposely not implemented
    bool m_all = false, m_alternate = false;
protected:
    NegateFilter() {}
    ~NegateFilter() {}

public:
    void SetNegate(const bool all, const bool alternate) {
        m_all = all;
        m_alternate = alternate;
    }
    void ThreadedGenerateData(const TRegion &region, ThreadIdType threadId) ITK_OVERRIDE {
        typedef typename TImage::PixelType PixelType;
        ImageSliceConstIteratorWithIndex<TImage> inIt(this->GetInput(), region);
        ImageSliceIteratorWithIndex<TImage>      outIt(this->GetOutput(), region);

        inIt.SetFirstDirection(0);
        inIt.SetSecondDirection(1);
        outIt.SetFirstDirection(0);
        outIt.SetSecondDirection(1);
        inIt.GoToBegin();
        outIt.GoToBegin();
        PixelType mult(m_all ? -1 : 1, 0);
        while(!inIt.IsAtEnd()) {
            while (!inIt.IsAtEndOfSlice()) {
                while (!inIt.IsAtEndOfLine()) {
                    outIt.Set(mult * inIt.Get());
                    ++inIt;
                    ++outIt;
                }
                inIt.NextLine();
                outIt.NextLine();
            }
            if (m_alternate) {
                mult = mult * PixelType(-1, 0);
            }
            inIt.NextSlice();
            outIt.NextSlice();
        }
        // std::std::cout << "End " << __PRETTY_FUNCTION__ << std::std::endl;
    }
};

} // End namespace itk

/* Arguments defined here so they are available in the templated run function */
args::ArgumentParser parser(
    "Input is specified with lower case letters. A valid combination of inputs "
    " must be specified, e.g. real & imaginary or magnitude & phase.\n"
    "Output is specified with upper case letters. Any combination can be given\n"
    "http://github.com/spinicist/QUIT");

args::HelpFlag help(parser, "HELP", "Show this help menu", {'h', "help"});
args::Flag     verbose(parser, "VERBOSE", "Print more information", {'v', "verbose"});
args::ValueFlag<int> threads(parser, "THREADS", "Use N threads (default=4, 0=hardware limit)", {'T', "threads"}, 4);
args::Flag     use_double(parser, "DOUBLE", "Process & output at double precision", {'d', "double"});
args::Flag     fixge(parser, "FIX_GE", "Negate alternate slices (fixes lack of FFT shift)", {"fixge"});
args::Flag     negate(parser, "NEGATE", "Negate entire volume", {"negate"});

args::ValueFlag<std::string> in_mag(parser, "IN_MAG", "Input magnitude file", {'m', "mag"});
args::ValueFlag<std::string> in_pha(parser, "IN_PHA", "Input phase file", {'p', "pha"});
args::ValueFlag<std::string> in_real(parser, "IN_REAL", "Input real file", {'r', "real"});
args::ValueFlag<std::string> in_imag(parser, "IN_IMAG", "Input imaginary file", {'i', "imag"});
args::ValueFlag<std::string> in_complex(parser, "IN_CPLX", "Input complex file", {'x', "complex"});
args::ValueFlag<std::string> in_realimag(parser, "IN_REALIMAG", "Input real & imaginary file", {'l', "realimag"});

args::ValueFlag<std::string> out_mag(parser, "OUT_MAG", "Output magnitude file", {'M', "MAG"});
args::ValueFlag<std::string> out_pha(parser, "OUT_PHA", "Output phase file", {'P', "PHA"});
args::ValueFlag<std::string> out_real(parser, "OUT_REAL", "Output real file", {'R', "REAL"});
args::ValueFlag<std::string> out_imag(parser, "OUT_IMAG", "Output imaginary file", {'I', "IMAG"});
args::ValueFlag<std::string> out_complex(parser, "OUT_CPLX", "Output complex file", {'X', "COMPLEX"});

template<typename TPixel>
void Run() {
    typedef itk::Image<TPixel, 4>          TImage;
    typedef itk::Image<std::complex<TPixel>, 4> TXImage;
    typedef itk::ImageFileWriter<TImage>   TWriter;

    typename TImage::Pointer img1 = ITK_NULLPTR, img2 = ITK_NULLPTR;
    typename TXImage::Pointer imgX = ITK_NULLPTR;

    if (in_real) {
        if (verbose) std::cout << "Reading real file: " << in_real.Get() << std::endl;
        img1 = QI::ReadImage<TImage>(in_real.Get());
        if (in_imag) {
            if (verbose) std::cout << "Reading imaginary file: " << in_imag.Get() << std::endl;
            img2 = QI::ReadImage<TImage>(in_imag.Get());
        } else {
            QI_FAIL("Must set real and imaginary inputs together");
        }
        auto compose = itk::ComposeImageFilter<TImage, TXImage>::New();
        compose->SetInput(0, img1);
        compose->SetInput(1, img2);
        compose->Update();
        imgX = compose->GetOutput();
        imgX->DisconnectPipeline();
    } else if (in_mag) {
        if (verbose) std::cout << "Reading magnitude file: " << in_mag.Get() << std::endl;
        img1 = QI::ReadImage<TImage>(in_mag.Get());
        if (in_pha) {
            if (verbose) std::cout << "Reading phase file: " << in_pha.Get() << std::endl;
            img2 = QI::ReadImage<TImage>(in_pha.Get());
        } else {
            QI_FAIL("Must set magnitude and phase inputs together");
        }
        auto compose = itk::MagnitudeAndPhaseToComplexImageFilter<TImage, TImage, TXImage>::New();
        compose->SetInput(0, img1);
        compose->SetInput(1, img2);
        compose->Update();
        imgX = compose->GetOutput();
        imgX->DisconnectPipeline();
    } else if (in_complex) {
        if (verbose) std::cout << "Reading complex file: " << in_complex.Get() << std::endl;
        imgX = QI::ReadImage<TXImage>(in_complex.Get());
    } else if (in_realimag) {
        if (verbose) std::cout << "Reading real/imaginary file: " << in_realimag.Get() << std::endl;
        auto img_both = QI::ReadImage<TImage>(in_realimag.Get());
        auto real_region = img_both->GetLargestPossibleRegion();
        auto imag_region = img_both->GetLargestPossibleRegion();
        real_region.GetModifiableSize()[3] = real_region.GetSize()[3] / 2;
        imag_region.GetModifiableSize()[3] = real_region.GetSize()[3];
        imag_region.GetModifiableIndex()[3] = real_region.GetSize()[3];
        auto extract_real = itk::RegionOfInterestImageFilter<TImage, TImage>::New();
        extract_real->SetRegionOfInterest(real_region);
        extract_real->SetInput(img_both);
        extract_real->Update();
        auto extract_imag = itk::RegionOfInterestImageFilter<TImage, TImage>::New();
        extract_imag->SetRegionOfInterest(imag_region);
        extract_imag->SetInput(img_both);
        extract_imag->Update();
        /* Hack round ITK changing the origin */
        extract_imag->GetOutput()->SetOrigin(extract_real->GetOutput()->GetOrigin());
        auto compose = itk::ComposeImageFilter<TImage, TXImage>::New();
        compose->SetInput(0, extract_real->GetOutput());
        compose->SetInput(1, extract_imag->GetOutput());
        compose->Update();
        imgX = compose->GetOutput();
        imgX->DisconnectPipeline();
    } else {
        QI_FAIL("No input files specified, use --help to see usage");
    }

    if (fixge || negate) {
        if (verbose && negate) std::cout << "Negating values" << std::endl;
        if (verbose && fixge)  std::cout << "Fixing GE lack of FFT-shift bug" << std::endl;
        auto fix = itk::NegateFilter<TXImage>::New();
        fix->SetInput(imgX);
        fix->SetNegate(negate, fixge);
        fix->Update();
        imgX = fix->GetOutput();
        imgX->DisconnectPipeline();
    }

    if (verbose) std::cout << "Writing output files" << std::endl;
    typename TWriter::Pointer write = TWriter::New();

    if (out_mag) {
        auto o = itk::ComplexToModulusImageFilter<TXImage, TImage>::New();
        o->SetInput(imgX);
        o->Update();
        QI::WriteImage(o->GetOutput(), out_mag.Get());
        if (verbose) std::cout << "Wrote magnitude image " << out_mag.Get() << std::endl;
    }
    if (out_pha) {
        auto o = itk::ComplexToPhaseImageFilter<TXImage, TImage>::New();
        o->SetInput(imgX);
        o->Update();
        QI::WriteImage(o->GetOutput(), out_pha.Get());
        if (verbose) std::cout << "Wrote phase image " << out_pha.Get() << std::endl;
    }
    if (out_real) {
        auto o = itk::ComplexToRealImageFilter<TXImage, TImage>::New();
        o->SetInput(imgX);
        o->Update();
        QI::WriteImage(o->GetOutput(), out_real.Get());
        if (verbose) std::cout << "Wrote real image " << out_real.Get() << std::endl;
    }
    if (out_imag) {
        auto o = itk::ComplexToImaginaryImageFilter<TXImage, TImage>::New();
        o->SetInput(imgX);
        o->Update();
        QI::WriteImage(o->GetOutput(), out_imag.Get());
        if (verbose) std::cout << "Wrote imaginary image " << out_imag.Get() << std::endl;
    }
    if (out_complex) {
        if (verbose) std::cout << "Writing complex image " << out_complex.Get() << std::endl;
        QI::WriteImage(imgX, out_complex.Get());
    }
}

int main(int argc, char **argv) {
    QI::ParseArgs(parser, argc, argv);
    itk::MultiThreader::SetGlobalMaximumNumberOfThreads(threads.Get());
    if (use_double) {
        if (verbose) std::cout << "Using double precision" << std::endl;
        Run<double>();
    } else {
        if (verbose) std::cout << "Using float precision" << std::endl;
        Run<float>();
    }
    return EXIT_SUCCESS;
}
<|endoftext|>