<commit_before>#include "PNGReadWrite.hpp"

#include <memory>

#include <cstdio>
#include <png.h>

#include <boost/log/trivial.hpp>
#include <boost/nowide/cstdio.hpp>

namespace Slic3r { namespace png {

struct PNGDescr {
    png_struct *png = nullptr; png_info *info = nullptr;

    PNGDescr() = default;
    PNGDescr(const PNGDescr&) = delete;
    PNGDescr(PNGDescr&&) = delete;
    PNGDescr& operator=(const PNGDescr&) = delete;
    PNGDescr& operator=(PNGDescr&&) = delete;

    ~PNGDescr()
    {
        if (png && info) png_destroy_info_struct(png, &info);
        if (png) png_destroy_read_struct( &png, nullptr, nullptr);
    }
};

bool is_png(const ReadBuf &rb)
{
    static const constexpr int PNG_SIG_BYTES = 8;

#if PNG_LIBPNG_VER_MINOR <= 2
    // Earlier libpng versions had png_sig_cmp(png_bytep, ...) which is not
    // a const pointer. It is not possible to cast away the const qualifier from
    // the input buffer so... yes... life is challenging...
    png_byte buf[PNG_SIG_BYTES];
    auto inbuf = static_cast<const std::uint8_t *>(rb.buf);
    std::copy(inbuf, inbuf + PNG_SIG_BYTES, buf);
#else
    auto buf = static_cast<png_const_bytep>(rb.buf);
#endif

    return rb.sz >= PNG_SIG_BYTES && !png_sig_cmp(buf, 0, PNG_SIG_BYTES);
}

// Buffer read callback for libpng. It provides an allocated output buffer and
// the amount of data it desires to read from the input.
static void png_read_callback(png_struct *png_ptr,
                              png_bytep   outBytes,
                              png_size_t  byteCountToRead)
{
    // Retrieve our input buffer through the png_ptr
    auto reader = static_cast<IStream *>(png_get_io_ptr(png_ptr));

    if (!reader || !reader->is_ok()) return;

    reader->read(static_cast<std::uint8_t *>(outBytes), byteCountToRead);
}

bool decode_png(IStream &in_buf, ImageGreyscale &out_img)
{
    static const constexpr int PNG_SIG_BYTES = 8;

    std::vector<png_byte> sig(PNG_SIG_BYTES, 0);
    in_buf.read(sig.data(), PNG_SIG_BYTES);
    if (!png_check_sig(sig.data(), PNG_SIG_BYTES))
        return false;

    PNGDescr dsc;
    dsc.png = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr,
                                     nullptr);

    if(!dsc.png) return false;

    dsc.info = png_create_info_struct(dsc.png);
    if(!dsc.info) return false;

    png_set_read_fn(dsc.png, static_cast<void *>(&in_buf), png_read_callback);

    // Tell that we have already read the first bytes to check the signature
    png_set_sig_bytes(dsc.png, PNG_SIG_BYTES);

    png_read_info(dsc.png, dsc.info);

    out_img.cols = png_get_image_width(dsc.png, dsc.info);
    out_img.rows = png_get_image_height(dsc.png, dsc.info);
    size_t color_type = png_get_color_type(dsc.png, dsc.info);
    size_t bit_depth  = png_get_bit_depth(dsc.png, dsc.info);

    if (color_type != PNG_COLOR_TYPE_GRAY || bit_depth != 8)
        return false;

    out_img.buf.resize(out_img.rows * out_img.cols);

    auto readbuf = static_cast<png_bytep>(out_img.buf.data());
    for (size_t r = 0; r < out_img.rows; ++r)
        png_read_row(dsc.png, readbuf + r * out_img.cols, nullptr);

    return true;
}

// Down to earth function to store a packed RGB image to file. Mostly useful for debugging purposes.
// Based on https://www.lemoda.net/c/write-png/
bool write_rgb_to_file(const char *file_name_utf8, size_t width, size_t height, const uint8_t *data_rgb)
{
    bool result = false;

    FILE *fp = boost::nowide::fopen(file_name_utf8, "wb");
    if (! fp) {
        BOOST_LOG_TRIVIAL(error) << "write_png_file: File could not be opened for writing: " << file_name_utf8;
        goto fopen_failed;
    }

    png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    if (! png_ptr) {
        BOOST_LOG_TRIVIAL(error) << "write_png_file: png_create_write_struct() failed";
        goto png_create_write_struct_failed;
    }

    png_infop info_ptr = png_create_info_struct(png_ptr);
    if (! info_ptr) {
        BOOST_LOG_TRIVIAL(error) << "write_png_file: png_create_info_struct() failed";
        goto png_create_info_struct_failed;
    }

    // Set up error handling.
    if (setjmp(png_jmpbuf(png_ptr))) {
        BOOST_LOG_TRIVIAL(error) << "write_png_file: setjmp() failed";
        goto png_failure;
    }

    // Set image attributes.
    png_set_IHDR(png_ptr,
        info_ptr,
        png_uint_32(width),
        png_uint_32(height),
        8, // depth
        PNG_COLOR_TYPE_RGB,
        PNG_INTERLACE_NONE,
        PNG_COMPRESSION_TYPE_DEFAULT,
        PNG_FILTER_TYPE_DEFAULT);

    // Initialize rows of PNG.
    auto row_pointers = reinterpret_cast<png_byte**>(::png_malloc(png_ptr, height * sizeof(png_byte*)));
    for (size_t y = 0; y < height; ++ y) {
        auto row = reinterpret_cast<png_byte*>(::png_malloc(png_ptr, sizeof(uint8_t) * width * 3));
        row_pointers[y] = row;
        memcpy(row, data_rgb + width * y * 3, sizeof(uint8_t) * width * 3);
    }

    // Write the image data to "fp".
    png_init_io(png_ptr, fp);
    png_set_rows(png_ptr, info_ptr, row_pointers);
    png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, nullptr);

    for (size_t y = 0; y < height; ++ y)
        png_free(png_ptr, row_pointers[y]);
    png_free(png_ptr, row_pointers);
    result = true;

png_failure:
png_create_info_struct_failed:
    ::png_destroy_write_struct(&png_ptr, &info_ptr);
png_create_write_struct_failed:
    ::fclose(fp);
fopen_failed:
    return result;
}

bool write_rgb_to_file(const std::string &file_name_utf8, size_t width, size_t height, const uint8_t *data_rgb)
{
    return write_rgb_to_file(file_name_utf8.c_str(), width, height, data_rgb);
}

bool write_rgb_to_file(const std::string &file_name_utf8, size_t width, size_t height, const std::vector<uint8_t> &data_rgb)
{
    assert(width * height * 3 == data_rgb.size());
    return write_rgb_to_file(file_name_utf8.c_str(), width, height, data_rgb.data());
}

// Scaled variants are mostly useful for debugging purposes, for example to export images of low resolution distance fileds.
// Scaling is done by multiplying rows and columns without any smoothing to emphasise the original pixels.
bool write_rgb_to_file_scaled(const char *file_name_utf8, size_t width, size_t height, const uint8_t *data_rgb, size_t scale)
{
    if (scale <= 1)
        return write_rgb_to_file(file_name_utf8, width, height, data_rgb);
    else {
        std::vector<uint8_t> scaled(width * height * 3 * scale * scale);
        uint8_t *dst = scaled.data();
        for (size_t r = 0; r < height; ++ r) {
            for (size_t repr = 0; repr < scale; ++ repr) {
                const uint8_t *row = data_rgb + width * 3 * r;
                for (size_t c = 0; c < width; ++ c) {
                    for (size_t repc = 0; repc < scale; ++ repc) {
                        *dst ++ = row[0];
                        *dst ++ = row[1];
                        *dst ++ = row[2];
                    }
                    row += 3;
                }
            }
        }
        return write_rgb_to_file(file_name_utf8, width * scale, height * scale, scaled.data());
    }
}

bool write_rgb_to_file_scaled(const std::string &file_name_utf8, size_t width, size_t height, const uint8_t *data_rgb, size_t scale)
{
    return write_rgb_to_file_scaled(file_name_utf8.c_str(), width, height, data_rgb, scale);
}

bool write_rgb_to_file_scaled(const std::string &file_name_utf8, size_t width, size_t height, const std::vector<uint8_t> &data_rgb, size_t scale)
{
    assert(width * height * 3 == data_rgb.size());
    return write_rgb_to_file_scaled(file_name_utf8.c_str(), width, height, data_rgb.data(), scale);
}

}} // namespace Slic3r::png
<commit_msg>Fixing compilation issues on conforming compilers<commit_after>#include "PNGReadWrite.hpp"

#include <memory>

#include <cstdio>
#include <png.h>

#include <boost/log/trivial.hpp>
#include <boost/nowide/cstdio.hpp>

namespace Slic3r { namespace png {

struct PNGDescr {
    png_struct *png = nullptr; png_info *info = nullptr;

    PNGDescr() = default;
    PNGDescr(const PNGDescr&) = delete;
    PNGDescr(PNGDescr&&) = delete;
    PNGDescr& operator=(const PNGDescr&) = delete;
    PNGDescr& operator=(PNGDescr&&) = delete;

    ~PNGDescr()
    {
        if (png && info) png_destroy_info_struct(png, &info);
        if (png) png_destroy_read_struct( &png, nullptr, nullptr);
    }
};

bool is_png(const ReadBuf &rb)
{
    static const constexpr int PNG_SIG_BYTES = 8;

#if PNG_LIBPNG_VER_MINOR <= 2
    // Earlier libpng versions had png_sig_cmp(png_bytep, ...) which is not
    // a const pointer. It is not possible to cast away the const qualifier from
    // the input buffer so... yes... life is challenging...
    png_byte buf[PNG_SIG_BYTES];
    auto inbuf = static_cast<const std::uint8_t *>(rb.buf);
    std::copy(inbuf, inbuf + PNG_SIG_BYTES, buf);
#else
    auto buf = static_cast<png_const_bytep>(rb.buf);
#endif

    return rb.sz >= PNG_SIG_BYTES && !png_sig_cmp(buf, 0, PNG_SIG_BYTES);
}

// Buffer read callback for libpng. It provides an allocated output buffer and
// the amount of data it desires to read from the input.
static void png_read_callback(png_struct *png_ptr,
                              png_bytep   outBytes,
                              png_size_t  byteCountToRead)
{
    // Retrieve our input buffer through the png_ptr
    auto reader = static_cast<IStream *>(png_get_io_ptr(png_ptr));

    if (!reader || !reader->is_ok()) return;

    reader->read(static_cast<std::uint8_t *>(outBytes), byteCountToRead);
}

bool decode_png(IStream &in_buf, ImageGreyscale &out_img)
{
    static const constexpr int PNG_SIG_BYTES = 8;

    std::vector<png_byte> sig(PNG_SIG_BYTES, 0);
    in_buf.read(sig.data(), PNG_SIG_BYTES);
    if (!png_check_sig(sig.data(), PNG_SIG_BYTES))
        return false;

    PNGDescr dsc;
    dsc.png = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr,
                                     nullptr);

    if(!dsc.png) return false;

    dsc.info = png_create_info_struct(dsc.png);
    if(!dsc.info) return false;

    png_set_read_fn(dsc.png, static_cast<void *>(&in_buf), png_read_callback);

    // Tell that we have already read the first bytes to check the signature
    png_set_sig_bytes(dsc.png, PNG_SIG_BYTES);

    png_read_info(dsc.png, dsc.info);

    out_img.cols = png_get_image_width(dsc.png, dsc.info);
    out_img.rows = png_get_image_height(dsc.png, dsc.info);
    size_t color_type = png_get_color_type(dsc.png, dsc.info);
    size_t bit_depth  = png_get_bit_depth(dsc.png, dsc.info);

    if (color_type != PNG_COLOR_TYPE_GRAY || bit_depth != 8)
        return false;

    out_img.buf.resize(out_img.rows * out_img.cols);

    auto readbuf = static_cast<png_bytep>(out_img.buf.data());
    for (size_t r = 0; r < out_img.rows; ++r)
        png_read_row(dsc.png, readbuf + r * out_img.cols, nullptr);

    return true;
}

// Down to earth function to store a packed RGB image to file. Mostly useful for debugging purposes.
// Based on https://www.lemoda.net/c/write-png/
bool write_rgb_to_file(const char *file_name_utf8, size_t width, size_t height, const uint8_t *data_rgb)
{
    bool         result       = false;

    // Forward declaration due to the gotos.
    png_structp  png_ptr      = nullptr;
    png_infop    info_ptr     = nullptr;
    png_byte   **row_pointers = nullptr;

    FILE        *fp = boost::nowide::fopen(file_name_utf8, "wb");
    if (! fp) {
        BOOST_LOG_TRIVIAL(error) << "write_png_file: File could not be opened for writing: " << file_name_utf8;
        goto fopen_failed;
    }

    png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    if (! png_ptr) {
        BOOST_LOG_TRIVIAL(error) << "write_png_file: png_create_write_struct() failed";
        goto png_create_write_struct_failed;
    }

    info_ptr = png_create_info_struct(png_ptr);
    if (! info_ptr) {
        BOOST_LOG_TRIVIAL(error) << "write_png_file: png_create_info_struct() failed";
        goto png_create_info_struct_failed;
    }

    // Set up error handling.
    if (setjmp(png_jmpbuf(png_ptr))) {
        BOOST_LOG_TRIVIAL(error) << "write_png_file: setjmp() failed";
        goto png_failure;
    }

    // Set image attributes.
    png_set_IHDR(png_ptr,
        info_ptr,
        png_uint_32(width),
        png_uint_32(height),
        8, // depth
        PNG_COLOR_TYPE_RGB,
        PNG_INTERLACE_NONE,
        PNG_COMPRESSION_TYPE_DEFAULT,
        PNG_FILTER_TYPE_DEFAULT);

    // Initialize rows of PNG.
    row_pointers = reinterpret_cast<png_byte**>(::png_malloc(png_ptr, height * sizeof(png_byte*)));
    for (size_t y = 0; y < height; ++ y) {
        auto row = reinterpret_cast<png_byte*>(::png_malloc(png_ptr, sizeof(uint8_t) * width * 3));
        row_pointers[y] = row;
        memcpy(row, data_rgb + width * y * 3, sizeof(uint8_t) * width * 3);
    }

    // Write the image data to "fp".
    png_init_io(png_ptr, fp);
    png_set_rows(png_ptr, info_ptr, row_pointers);
    png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, nullptr);

    for (size_t y = 0; y < height; ++ y)
        png_free(png_ptr, row_pointers[y]);
    png_free(png_ptr, row_pointers);

    result = true;

png_failure:
png_create_info_struct_failed:
    ::png_destroy_write_struct(&png_ptr, &info_ptr);
png_create_write_struct_failed:
    ::fclose(fp);
fopen_failed:
    return result;
}

bool write_rgb_to_file(const std::string &file_name_utf8, size_t width, size_t height, const uint8_t *data_rgb)
{
    return write_rgb_to_file(file_name_utf8.c_str(), width, height, data_rgb);
}

bool write_rgb_to_file(const std::string &file_name_utf8, size_t width, size_t height, const std::vector<uint8_t> &data_rgb)
{
    assert(width * height * 3 == data_rgb.size());
    return write_rgb_to_file(file_name_utf8.c_str(), width, height, data_rgb.data());
}

// Scaled variants are mostly useful for debugging purposes, for example to export images of low resolution distance fileds.
// Scaling is done by multiplying rows and columns without any smoothing to emphasise the original pixels.
bool write_rgb_to_file_scaled(const char *file_name_utf8, size_t width, size_t height, const uint8_t *data_rgb, size_t scale)
{
    if (scale <= 1)
        return write_rgb_to_file(file_name_utf8, width, height, data_rgb);
    else {
        std::vector<uint8_t> scaled(width * height * 3 * scale * scale);
        uint8_t *dst = scaled.data();
        for (size_t r = 0; r < height; ++ r) {
            for (size_t repr = 0; repr < scale; ++ repr) {
                const uint8_t *row = data_rgb + width * 3 * r;
                for (size_t c = 0; c < width; ++ c) {
                    for (size_t repc = 0; repc < scale; ++ repc) {
                        *dst ++ = row[0];
                        *dst ++ = row[1];
                        *dst ++ = row[2];
                    }
                    row += 3;
                }
            }
        }
        return write_rgb_to_file(file_name_utf8, width * scale, height * scale, scaled.data());
    }
}

bool write_rgb_to_file_scaled(const std::string &file_name_utf8, size_t width, size_t height, const uint8_t *data_rgb, size_t scale)
{
    return write_rgb_to_file_scaled(file_name_utf8.c_str(), width, height, data_rgb, scale);
}

bool write_rgb_to_file_scaled(const std::string &file_name_utf8, size_t width, size_t height, const std::vector<uint8_t> &data_rgb, size_t scale)
{
    assert(width * height * 3 == data_rgb.size());
    return write_rgb_to_file_scaled(file_name_utf8.c_str(), width, height, data_rgb.data(), scale);
}

}} // namespace Slic3r::png
<|endoftext|>


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

#include "cloginclient.h"
#include "cloginserver.h"
#include "cloginisc.h"
#include "croseisc.h"
#include "packetfactory.h"
#include "isccommon.h"
#include "isc_server_register.h"
#include "cli_login_req.h"
#include "connection.h"
#include "connectionpool.h"
#include "croseserver.h"
#include "cli_channel_list_req.h"
#include "cli_srv_select_req.h"
#include "srv_srv_select_reply.h"
#include "epackettype.h"
#include "srv_channel_list_reply.h"
#include "config.h"

using namespace RoseCommon;
using namespace RoseCommon::Packet;

CLoginClient::CLoginClient()
    : CRoseClient(),
      access_rights_(0),
      login_state_(eSTATE::DEFAULT),
      userid_(0),
      session_id_(0),
      server_(nullptr) {}

CLoginClient::CLoginClient(CLoginServer* server, std::unique_ptr<Core::INetwork> _sock)
    : CRoseClient(std::move(_sock)),
      access_rights_(0),
      login_state_(eSTATE::DEFAULT),
      userid_(0),
      session_id_(0),
      server_(server) {}

void CLoginClient::sendLoginReply(SrvLoginReply::Result Result) {
  logger_->debug("sendLoginReply({})", Result);
  auto packet = SrvLoginReply::create(Result, 0, 0);

  if (Result == SrvLoginReply::OK) {
    login_state_ = eSTATE::LOGGEDIN;
    packet.set_right(access_rights_);

    // loop the server list here
    std::lock_guard<std::mutex> lock(server_->GetISCListMutex());
    for (auto& server : server_->GetISCList()) {
      if (server->get_type() == Isc::ServerType::CHAR) {
        CLoginISC* svr = static_cast<CLoginISC*>(server.get());
        if (!svr) {
            continue;
        }

        // This if check is needed since the client actually looks for this.
        auto info = SrvLoginReply::ServerInfo();
        info.set_name(svr->getName());
        info.set_id(svr->get_id() + 1);
        info.set_test(svr->isTestServer() ? '@' : ' ');
        packet.add_serversInfo(info);
      }
    }
  }

  send(packet);
}

bool CLoginClient::userLogin(CliLoginReq&& P) {
  if (login_state_ != eSTATE::DEFAULT) {
    logger_->warn("Client {} is attempting to login when already logged in.",
                  get_id());
    return true;
  }
  uint32_t serverCount = 0;

  server_->GetISCListMutex().lock();
  for (auto& server : server_->GetISCList()) {
    if (server && server->get_type() == Isc::ServerType::CHAR) serverCount++;
  }
  server_->GetISCListMutex().unlock();

  logger_->debug("Found {} type 1 (CHAR) servers", serverCount);

  if (serverCount < 1) {
    // Servers are under inspection
    sendLoginReply(SrvLoginReply::FAILED);
    return true;
  }

  username_ = Core::escapeData(P.get_username());
  std::string clientpass = Core::escapeData(P.get_password());

  logger_->debug("Client sent '{}' as the password", clientpass);

  auto conn = Core::connectionPool.getConnection(Core::osirose);
  Core::AccountTable table{};
  Core::SessionTable session{};
  try {
    const auto res = conn(sqlpp::select(table.id, table.password, table.access, table.active, table.online, table.loginCount)
              .from(table).where(table.accountType == "user" and table.username == username_
                  and table.password == sqlpp::verbatim<sqlpp::varchar>(fmt::format("SHA2(CONCAT('{}', salt), 256)", clientpass))));

        if (!res.empty()) {
            const auto &row = res.front();
            const auto ses = conn(sqlpp::select(session.id).from(session).where(session.userid == row.id));
            if (!row.access.is_null())
                access_rights_ = row.access;

            if (access_rights_ < 1) {
                // Banned
                sendLoginReply(SrvLoginReply::NO_RIGHT_TO_CONNECT);
                return false;
            }

            if (!row.online && ses.empty()) {
                // Okay to login!!
                userid_ = row.id;
                session_id_ = std::time(nullptr);
                conn(sqlpp::update(table).set(table.online = 1,
                                              table.loginCount = row.loginCount + 1,
                                              table.lastip = get_address(),
                                              table.lasttime = std::chrono::system_clock::now())
                     .where(table.id == userid_));
                this->set_name(username_);
                sendLoginReply(SrvLoginReply::OK);
            } else {
                // Online already
                sendLoginReply(SrvLoginReply::ALREADY_LOGGEDIN);
            }
        } else {
            if (!conn(sqlpp::select(table.id).from(table).where(table.username == username_)).empty())
                sendLoginReply(SrvLoginReply::INVALID_PASSWORD);
            else {
                // The user doesn't exist or server is down.
                auto& config = Core::Config::getInstance();
                if (config.loginServer().createAccountOnFail) {
                  logger_->debug("Creating account");
                    std::string query = fmt::format("CALL create_account('{}', '{}');", username_, clientpass);
                    conn->execute(query);
                }
                sendLoginReply(SrvLoginReply::UNKNOWN_ACCOUNT);
            }
        }
  } catch (const sqlpp::exception &e) {
    logger_->error("Error while accessing the database: {}", e.what());
        sendLoginReply(SrvLoginReply::FAILED);
  }
  return true;
}

void CLoginClient::onDisconnected() {
  Core::SessionTable session{};
  auto conn = Core::connectionPool.getConnection(Core::osirose);
  auto res = conn(sqlpp::select(session.id).from(session)
                  .where(session.id == session_id_));
  if (res.empty()) {
    Core::AccountTable account{};
    conn(sqlpp::update(account).set(account.online = 0)
         .where(account.id == userid_));
  }
}

bool CLoginClient::channelList(CliChannelListReq&& P) {
  if (login_state_ != eSTATE::LOGGEDIN) {
    logger_->warn(
        "Client {} is attempting to get channel list before logging in.",
        get_id());
    return true;
  }
  uint32_t ServerID = P.get_serverId()-1;

  auto packet = SrvChannelListReply::create(ServerID+1);
  std::lock_guard<std::mutex> lock(server_->GetISCListMutex());
  for (auto& obj : server_->GetISCList()) {
    CLoginISC* server = static_cast<CLoginISC*>(obj.get());
    if (!server) {
        continue;
    }
    if (server->get_type() == Isc::ServerType::CHAR &&
        server->get_id() == ServerID) {
      for (tChannelInfo& info : server->getChannelList()) {
        auto channel = SrvChannelListReply::ChannelInfo();
        channel.set_id(info.ChannelID + 1);
        channel.set_name(info.channelName);
        channel.set_lowAge(0);
        channel.set_highAge(0);
        channel.set_capacity(0);
        packet.add_channels(channel);
      }
    }
  }

  this->send(packet);
  logger_->trace("Client {}: Channel List end.", get_id());
  return true;
}

bool CLoginClient::serverSelect(CliSrvSelectReq&& P) {
  if (login_state_ != eSTATE::LOGGEDIN) {
    logger_->warn(
        "Client {} is attempting to select a server before logging in.",
        get_id());
    return true;
  }
  uint32_t serverID = P.get_serverId()-1;
  uint8_t channelID = P.get_channelId()-1;
  login_state_ = eSTATE::TRANSFERING;

  // 0 = Good to go
  // 1 = Failed
  // 2 = Full
  // 3 = Invalid channel
  // 4 = Channel not active
  std::lock_guard<std::mutex> lock(server_->GetISCListMutex());
  for (auto& obj : server_->GetISCList()) {
    CLoginISC* server = static_cast<CLoginISC*>(obj.get());
    if (server->get_type() == Isc::ServerType::CHAR &&
        server->get_id() == serverID) {
      Core::SessionTable session{};
      auto conn = Core::connectionPool.getConnection(Core::osirose);
      conn(sqlpp::insert_into(session).set(
                    session.id = session_id_,
                    session.userid = userid_,
                    session.channelid = channelID,
                    session.time = std::chrono::system_clock::now()));

      auto packet = SrvSrvSelectReply::create(
          SrvSrvSelectReply::OK, session_id_, 0,
          server->get_address(), server->get_port());
      this->send(packet);
      break;
    }
  }
  logger_->trace("Client {}: Server Select end.", get_id());
  return true;
}

bool CLoginClient::handlePacket(uint8_t* _buffer) {
  logger_->trace("CLoginClient::handlePacket start");
  switch (CRosePacket::type(_buffer)) {
    case ePacketType::PAKCS_CHANNEL_LIST_REQ:
      return channelList(
          CliChannelListReq::create(_buffer));
    case ePacketType::PAKCS_SRV_SELECT_REQ:
      return serverSelect(
          CliSrvSelectReq::create(_buffer));
    case ePacketType::PAKCS_LOGIN_REQ:
      return userLogin(CliLoginReq::create(_buffer));

    default:
      return CRoseClient::handlePacket(_buffer);
  }
  logger_->trace("CLoginClient::handlePacket end");
  return true;
}
<commit_msg>Update cloginclient.cpp<commit_after>// Copyright 2016 Chirstopher Torres (Raven), L3nn0x
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http ://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "cloginclient.h"
#include "cloginserver.h"
#include "cloginisc.h"
#include "croseisc.h"
#include "packetfactory.h"
#include "isccommon.h"
#include "isc_server_register.h"
#include "cli_login_req.h"
#include "connection.h"
#include "connectionpool.h"
#include "croseserver.h"
#include "cli_channel_list_req.h"
#include "cli_srv_select_req.h"
#include "srv_srv_select_reply.h"
#include "epackettype.h"
#include "srv_channel_list_reply.h"
#include "config.h"

using namespace RoseCommon;
using namespace RoseCommon::Packet;

CLoginClient::CLoginClient()
    : CRoseClient(),
      access_rights_(0),
      login_state_(eSTATE::DEFAULT),
      userid_(0),
      session_id_(0),
      server_(nullptr) {}

CLoginClient::CLoginClient(CLoginServer* server, std::unique_ptr<Core::INetwork> _sock)
    : CRoseClient(std::move(_sock)),
      access_rights_(0),
      login_state_(eSTATE::DEFAULT),
      userid_(0),
      session_id_(0),
      server_(server) {}

void CLoginClient::sendLoginReply(SrvLoginReply::Result Result) {
  logger_->debug("sendLoginReply({})", Result);
  auto packet = SrvLoginReply::create(Result, 0, 0);

  if (Result == SrvLoginReply::OK) {
    login_state_ = eSTATE::LOGGEDIN;
    packet.set_right(access_rights_);

    // loop the server list here
    std::lock_guard<std::mutex> lock(server_->GetISCListMutex());
    for (auto& server : server_->GetISCList()) {
      if (server->get_type() == Isc::ServerType::CHAR) {
        CLoginISC* svr = static_cast<CLoginISC*>(server.get());
        if (!svr) {
            continue;
        }

        // This if check is needed since the client actually looks for this.
        auto info = SrvLoginReply::ServerInfo();
        info.set_name(svr->getName());
        info.set_id(svr->get_id() + 1);
        info.set_test(svr->isTestServer() ? '@' : ' ');
        packet.add_serversInfo(info);
      }
    }
  }

  send(packet);
}

bool CLoginClient::userLogin(CliLoginReq&& P) {
  if (login_state_ != eSTATE::DEFAULT) {
    logger_->warn("Client {} is attempting to login when already logged in.",
                  get_id());
    return true;
  }
  uint32_t serverCount = 0;

  server_->GetISCListMutex().lock();
  for (auto& server : server_->GetISCList()) {
    if (server && server->get_type() == Isc::ServerType::CHAR) serverCount++;
  }
  server_->GetISCListMutex().unlock();

  logger_->debug("Found {} type 1 (CHAR) servers", serverCount);

  if (serverCount < 1) {
    // Servers are under inspection
    sendLoginReply(SrvLoginReply::FAILED);
    return true;
  }

  username_ = Core::escapeData(P.get_username());
  std::string clientpass = Core::escapeData(P.get_password());

  logger_->debug("Client sent '{}' as the password", clientpass);

  auto conn = Core::connectionPool.getConnection<Core::Osirose>();
  Core::AccountTable table{};
  Core::SessionTable session{};
  try {
    const auto res = conn(sqlpp::select(table.id, table.password, table.access, table.active, table.online, table.loginCount)
              .from(table).where(table.accountType == "user" and table.username == username_
                  and table.password == sqlpp::verbatim<sqlpp::varchar>(fmt::format("SHA2(CONCAT('{}', salt), 256)", clientpass))));

        if (!res.empty()) {
            const auto &row = res.front();
            const auto ses = conn(sqlpp::select(session.id).from(session).where(session.userid == row.id));
            if (!row.access.is_null())
                access_rights_ = row.access;

            if (access_rights_ < 1) {
                // Banned
                sendLoginReply(SrvLoginReply::NO_RIGHT_TO_CONNECT);
                return false;
            }

            if (!row.online && ses.empty()) {
                // Okay to login!!
                userid_ = row.id;
                session_id_ = std::time(nullptr);
                conn(sqlpp::update(table).set(table.online = 1,
                                              table.loginCount = row.loginCount + 1,
                                              table.lastip = get_address(),
                                              table.lasttime = std::chrono::system_clock::now())
                     .where(table.id == userid_));
                this->set_name(username_);
                sendLoginReply(SrvLoginReply::OK);
            } else {
                // Online already
                sendLoginReply(SrvLoginReply::ALREADY_LOGGEDIN);
            }
        } else {
            if (!conn(sqlpp::select(table.id).from(table).where(table.username == username_)).empty())
                sendLoginReply(SrvLoginReply::INVALID_PASSWORD);
            else {
                // The user doesn't exist or server is down.
                auto& config = Core::Config::getInstance();
                if (config.loginServer().createAccountOnFail) {
                  logger_->debug("Creating account");
                    std::string query = fmt::format("CALL create_account('{}', '{}');", username_, clientpass);
                    conn->execute(query);
                }
                sendLoginReply(SrvLoginReply::UNKNOWN_ACCOUNT);
            }
        }
  } catch (const sqlpp::exception &e) {
    logger_->error("Error while accessing the database: {}", e.what());
        sendLoginReply(SrvLoginReply::FAILED);
  }
  return true;
}

void CLoginClient::onDisconnected() {
  Core::SessionTable session{};
  auto conn = Core::connectionPool.getConnection<Core::Osirose>();
  auto res = conn(sqlpp::select(session.id).from(session)
                  .where(session.id == session_id_));
  if (res.empty()) {
    Core::AccountTable account{};
    conn(sqlpp::update(account).set(account.online = 0)
         .where(account.id == userid_));
  }
}

bool CLoginClient::channelList(CliChannelListReq&& P) {
  if (login_state_ != eSTATE::LOGGEDIN) {
    logger_->warn(
        "Client {} is attempting to get channel list before logging in.",
        get_id());
    return true;
  }
  uint32_t ServerID = P.get_serverId()-1;

  auto packet = SrvChannelListReply::create(ServerID+1);
  std::lock_guard<std::mutex> lock(server_->GetISCListMutex());
  for (auto& obj : server_->GetISCList()) {
    CLoginISC* server = static_cast<CLoginISC*>(obj.get());
    if (!server) {
        continue;
    }
    if (server->get_type() == Isc::ServerType::CHAR &&
        server->get_id() == ServerID) {
      for (tChannelInfo& info : server->getChannelList()) {
        auto channel = SrvChannelListReply::ChannelInfo();
        channel.set_id(info.ChannelID + 1);
        channel.set_name(info.channelName);
        channel.set_lowAge(0);
        channel.set_highAge(0);
        channel.set_capacity(0);
        packet.add_channels(channel);
      }
    }
  }

  this->send(packet);
  logger_->trace("Client {}: Channel List end.", get_id());
  return true;
}

bool CLoginClient::serverSelect(CliSrvSelectReq&& P) {
  if (login_state_ != eSTATE::LOGGEDIN) {
    logger_->warn(
        "Client {} is attempting to select a server before logging in.",
        get_id());
    return true;
  }
  uint32_t serverID = P.get_serverId()-1;
  uint8_t channelID = P.get_channelId()-1;
  login_state_ = eSTATE::TRANSFERING;

  // 0 = Good to go
  // 1 = Failed
  // 2 = Full
  // 3 = Invalid channel
  // 4 = Channel not active
  std::lock_guard<std::mutex> lock(server_->GetISCListMutex());
  for (auto& obj : server_->GetISCList()) {
    CLoginISC* server = static_cast<CLoginISC*>(obj.get());
    if (server->get_type() == Isc::ServerType::CHAR &&
        server->get_id() == serverID) {
      Core::SessionTable session{};
      auto conn = Core::connectionPool.getConnection<Core::Osirose>();
      conn(sqlpp::insert_into(session).set(
                    session.id = session_id_,
                    session.userid = userid_,
                    session.channelid = channelID,
                    session.time = std::chrono::system_clock::now()));

      auto packet = SrvSrvSelectReply::create(
          SrvSrvSelectReply::OK, session_id_, 0,
          server->get_address(), server->get_port());
      this->send(packet);
      break;
    }
  }
  logger_->trace("Client {}: Server Select end.", get_id());
  return true;
}

bool CLoginClient::handlePacket(uint8_t* _buffer) {
  logger_->trace("CLoginClient::handlePacket start");
  switch (CRosePacket::type(_buffer)) {
    case ePacketType::PAKCS_CHANNEL_LIST_REQ:
      return channelList(
          CliChannelListReq::create(_buffer));
    case ePacketType::PAKCS_SRV_SELECT_REQ:
      return serverSelect(
          CliSrvSelectReq::create(_buffer));
    case ePacketType::PAKCS_LOGIN_REQ:
      return userLogin(CliLoginReq::create(_buffer));

    default:
      return CRoseClient::handlePacket(_buffer);
  }
  logger_->trace("CLoginClient::handlePacket end");
  return true;
}
<|endoftext|>


<commit_before>#include "Arduino.h"
#include "Display.h"

#define DIGIT_ON  HIGH
#define DIGIT_OFF  LOW

#define SEGMENT_ON  LOW
#define SEGMENT_OFF HIGH

Display::Display() {
  pinMode(PIN_DISPLAY_SEGMENT_A, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_B, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_C, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_D, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_E, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_F, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_G, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_COLON, OUTPUT);

  pinMode(PIN_DISPLAY_DIGIT_1, OUTPUT);
  pinMode(PIN_DISPLAY_DIGIT_2, OUTPUT);
  pinMode(PIN_DISPLAY_DIGIT_3, OUTPUT);
  pinMode(PIN_DISPLAY_DIGIT_4, OUTPUT);

  digitalWrite(PIN_DISPLAY_SEGMENT_COLON, SEGMENT_ON);
}

void Display::display_text(char* text) {
  // Serial.println(text);

  _pause_if_required();

  _display_digit(PIN_DISPLAY_DIGIT_1, text[0]);
  _display_digit(PIN_DISPLAY_DIGIT_2, text[1]);
  _display_digit(PIN_DISPLAY_DIGIT_3, text[2]);
  _display_digit(PIN_DISPLAY_DIGIT_4, text[3]);
}

void Display::display_time(int toDisplay) {

  int minutes = toDisplay / 60;
  int seconds = toDisplay % 60;
  char numbers_as_letters[] = "0123456789";

  _pause_if_required();
  _display_digit(PIN_DISPLAY_DIGIT_1, numbers_as_letters[minutes / 10 % 10]);
  _display_digit(PIN_DISPLAY_DIGIT_2, numbers_as_letters[minutes % 10     ]);
  _display_digit(PIN_DISPLAY_DIGIT_3, numbers_as_letters[seconds / 10     ]);
  _display_digit(PIN_DISPLAY_DIGIT_4, numbers_as_letters[seconds % 10     ]);
}

void Display::_pause_if_required() {
  // Pause if the display has been painted too recently;
  while( (millis() - _display_last_painted) < DISPLAY_LOOP_TIME) ;
  _display_last_painted = millis();
}

void Display::_display_digit(int digit, char toDisplay) {
  digitalWrite(digit, DIGIT_ON);
  _turn_segments_on(toDisplay);
  delayMicroseconds(DISPLAY_BRIGHTNESS);
  digitalWrite(digit, DIGIT_OFF);
  _turn_all_segments_off();
}

void Display::_turn_all_segments_off () {
  digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_OFF);
}

// Given a number, turns on those segments
void Display::_turn_segments_on(char charToDisplay) {

  // TODO - change all these to use a byte per char and the bits in the byte to
  // represent the segments in the LED display.

  switch (charToDisplay) {

  case ' ':
  case 0: // nul
    _turn_all_segments_off();
    break;

  case '0':
  case 'O':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    break;

  case '1':
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    break;

  case '2':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '3':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '4':
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '5':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '6':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '7':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    break;

  case '8':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '9':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case 'B':
  case 'b':
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case 'E':
  case 'e':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case 'F':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case 'Y':
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  default:
    Serial.print("ERROR - can't light up segment for ascii decimal '");
    Serial.print((int)charToDisplay);
    Serial.println("'");
    break;

  }
}
<commit_msg>Add letters (b, r, i, v, o, l) for later<commit_after>#include "Arduino.h"
#include "Display.h"

#define DIGIT_ON  HIGH
#define DIGIT_OFF  LOW

#define SEGMENT_ON  LOW
#define SEGMENT_OFF HIGH

Display::Display() {
  pinMode(PIN_DISPLAY_SEGMENT_A, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_B, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_C, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_D, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_E, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_F, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_G, OUTPUT);
  pinMode(PIN_DISPLAY_SEGMENT_COLON, OUTPUT);

  pinMode(PIN_DISPLAY_DIGIT_1, OUTPUT);
  pinMode(PIN_DISPLAY_DIGIT_2, OUTPUT);
  pinMode(PIN_DISPLAY_DIGIT_3, OUTPUT);
  pinMode(PIN_DISPLAY_DIGIT_4, OUTPUT);

  digitalWrite(PIN_DISPLAY_SEGMENT_COLON, SEGMENT_ON);
}

void Display::display_text(char* text) {
  // Serial.println(text);

  _pause_if_required();

  _display_digit(PIN_DISPLAY_DIGIT_1, text[0]);
  _display_digit(PIN_DISPLAY_DIGIT_2, text[1]);
  _display_digit(PIN_DISPLAY_DIGIT_3, text[2]);
  _display_digit(PIN_DISPLAY_DIGIT_4, text[3]);
}

void Display::display_time(int toDisplay) {

  int minutes = toDisplay / 60;
  int seconds = toDisplay % 60;
  char numbers_as_letters[] = "0123456789";

  _pause_if_required();
  _display_digit(PIN_DISPLAY_DIGIT_1, numbers_as_letters[minutes / 10 % 10]);
  _display_digit(PIN_DISPLAY_DIGIT_2, numbers_as_letters[minutes % 10     ]);
  _display_digit(PIN_DISPLAY_DIGIT_3, numbers_as_letters[seconds / 10     ]);
  _display_digit(PIN_DISPLAY_DIGIT_4, numbers_as_letters[seconds % 10     ]);
}

void Display::_pause_if_required() {
  // Pause if the display has been painted too recently;
  while( (millis() - _display_last_painted) < DISPLAY_LOOP_TIME) ;
  _display_last_painted = millis();
}

void Display::_display_digit(int digit, char toDisplay) {
  digitalWrite(digit, DIGIT_ON);
  _turn_segments_on(toDisplay);
  delayMicroseconds(DISPLAY_BRIGHTNESS);
  digitalWrite(digit, DIGIT_OFF);
  _turn_all_segments_off();
}

void Display::_turn_all_segments_off () {
  digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_OFF);
  digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_OFF);
}

// Given a number, turns on those segments
void Display::_turn_segments_on(char charToDisplay) {

  // TODO - change all these to use a byte per char and the bits in the byte to
  // represent the segments in the LED display.

  switch (charToDisplay) {

  case ' ':
  case 0: // nul
    _turn_all_segments_off();
    break;

  case '0':
  case 'O':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    break;

  case '1':
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    break;

  case '2':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '3':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '4':
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '5':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '6':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '7':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    break;

  case '8':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case '9':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case 'B':
  case 'b':
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case 'i':
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    break;

  case 'E':
  case 'e':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case 'F':
    digitalWrite(PIN_DISPLAY_SEGMENT_A, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case 'l':
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    break;

  case 'o':
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case 'R':
  case 'r':
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case 'Y':
    digitalWrite(PIN_DISPLAY_SEGMENT_B, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_F, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_G, SEGMENT_ON);
    break;

  case 'v':
    digitalWrite(PIN_DISPLAY_SEGMENT_C, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_D, SEGMENT_ON);
    digitalWrite(PIN_DISPLAY_SEGMENT_E, SEGMENT_ON);
    break;

  default:
    Serial.print("ERROR - can't light up segment for ascii decimal '");
    Serial.print((int)charToDisplay);
    Serial.println("'");
    break;

  }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2008-2018 the MRtrix3 contributors.
 *
 * 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/
 *
 * MRtrix3 is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty
 * of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 *
 * For more details, see http://www.mrtrix.org/
 */

#include "command.h"
#include "image.h"

#include <Eigen/Dense>
#include <Eigen/Eigenvalues>

#define DEFAULT_SIZE 5

using namespace MR;
using namespace App;

const char* const dtypes[] = { "float32", "float64", NULL };

void usage ()
{
  SYNOPSIS = "Denoise DWI data and estimate the noise level based on the optimal threshold for PCA";

  DESCRIPTION
    + "DWI data denoising and noise map estimation by exploiting data redundancy in the PCA domain "
    "using the prior knowledge that the eigenspectrum of random covariance matrices is described by "
    "the universal Marchenko-Pastur (MP) distribution. Fitting the MP distribution to the spectrum "
    "of patch-wise signal matrices hence provides an estimator of the noise level 'sigma', as was "
    "first shown in Veraart et al. (2016) and later improved in Cordero-Grande et al. (2019). This "
    "noise level estimate then determines the optimal cut-off for PCA denoising."

    + "Important note: image denoising must be performed as the first step of the image processing pipeline. "
    "The routine will fail if interpolation or smoothing has been applied to the data prior to denoising."

    + "Note that this function does not correct for non-Gaussian noise biases.";

  AUTHOR = "Daan Christiaens (daan.christiaens@kcl.ac.uk) & "
           "Jelle Veraart (jelle.veraart@nyumc.org) & "
           "J-Donald Tournier (jdtournier@gmail.com)";

  REFERENCES
    + "Veraart, J.; Novikov, D.S.; Christiaens, D.; Ades-aron, B.; Sijbers, J. & Fieremans, E. " // Internal
    "Denoising of diffusion MRI using random matrix theory. "
    "NeuroImage, 2016, 142, 394-406, doi: 10.1016/j.neuroimage.2016.08.016"

    + "Veraart, J.; Fieremans, E. & Novikov, D.S. " // Internal
    "Diffusion MRI noise mapping using random matrix theory. "
    "Magn. Res. Med., 2016, 76(5), 1582-1593, doi: 10.1002/mrm.26059"

    + "Cordero-Grande, L.; Christiaens, D.; Hutter, J.; Price, A.N.; Hajnal, J.V. " // Internal
    "Complex diffusion-weighted image estimation via matrix recovery under general noise models. "
    "NeuroImage, 2019, 200, 391-404, doi: 10.1016/j.neuroimage.2019.06.039";

  ARGUMENTS
  + Argument ("dwi", "the input diffusion-weighted image.").type_image_in ()

  + Argument ("out", "the output denoised DWI image.").type_image_out ();

  OPTIONS
    + Option ("mask", "only perform computation within the specified binary brain mask image.")
    +   Argument ("image").type_image_in()

    + Option ("extent", "set the window size of the denoising filter. (default = " + str(DEFAULT_SIZE) + "," + str(DEFAULT_SIZE) + "," + str(DEFAULT_SIZE) + ")")
    +   Argument ("window").type_sequence_int ()

    + Option ("noise", "the output noise map.")
    +   Argument ("level").type_image_out()

    + Option ("datatype", "datatype for SVD (float32 or float64).")
    +   Argument ("spec").type_choice(dtypes)

    + Option ("exp1", "use the original noise estimator (exp1) as defined in Veraart et al. (2016), "
                      "rather than the corrected estimator (exp2) in Cordero-Grande et al. (2019).");

  COPYRIGHT = "Copyright (c) 2016 New York University, University of Antwerp, and the MRtrix3 contributors \n \n"
      "Permission is hereby granted, free of charge, to any non-commercial entity ('Recipient') obtaining a copy of this software and "
      "associated documentation files (the 'Software'), to the Software solely for non-commercial research, including the rights to "
      "use, copy and modify the Software, subject to the following conditions: \n \n"
      "\t 1. The above copyright notice and this permission notice shall be included by Recipient in all copies or substantial portions of "
      "the Software. \n \n"
      "\t 2. THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES"
      "OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE"
      "LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR"
      "IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. \n \n"
      "\t 3. In no event shall NYU be liable for direct, indirect, special, incidental or consequential damages in connection with the Software. "
      "Recipient will defend, indemnify and hold NYU harmless from any claims or liability resulting from the use of the Software by recipient. \n \n"
      "\t 4. Neither anything contained herein nor the delivery of the Software to recipient shall be deemed to grant the Recipient any right or "
      "licenses under any patents or patent application owned by NYU. \n \n"
      "\t 5. The Software may only be used for non-commercial research and may not be used for clinical care. \n \n"
      "\t 6. Any publication by Recipient of research involving the Software shall cite the references listed below.";

}

using value_type = float;

template <class ImageType, typename F = float>
class DenoisingFunctor {
  MEMALIGN(DenoisingFunctor)

public:

  typedef Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic> MatrixX;
  typedef Eigen::Matrix<F, Eigen::Dynamic, 1> VectorX;

  DenoisingFunctor (ImageType& dwi, vector<int> extent, Image<bool>& mask, ImageType& noise, bool exp1)
    : extent {{extent[0]/2, extent[1]/2, extent[2]/2}},
      m (dwi.size(3)), n (extent[0]*extent[1]*extent[2]),
      r (std::min(m,n)), q (std::max(m,n)), exp1(exp1),
      X (m,n), pos {{0, 0, 0}},
      mask (mask), noise (noise)
  { }

  void operator () (ImageType& dwi, ImageType& out)
  {
    // Process voxels in mask only
    if (mask.valid()) {
      assign_pos_of (dwi).to (mask);
      if (!mask.value())
        return;
    }

    // Load data in local window
    load_data (dwi);

    // Compute Eigendecomposition:
    MatrixX XtX (r,r);
    if (m <= n)
      XtX.template triangularView<Eigen::Lower>() = X * X.adjoint();
    else
      XtX.template triangularView<Eigen::Lower>() = X.adjoint() * X;
    Eigen::SelfAdjointEigenSolver<MatrixX> eig (XtX);
    // eigenvalues provide squared singular values, sorted in increasing order:
    VectorX s = eig.eigenvalues();

    // Marchenko-Pastur optimal threshold
    const double lam_r = std::max(double(s[0]), 0.0) / q;
    double clam = 0.0;
    sigma2 = 0.0;
    ssize_t cutoff_p = 0;
    for (ssize_t p = 0; p < r; ++p)     // p+1 is the number of noise components
    {                                   // (as opposed to the paper where p is defined as the number of signal components)
      double lam = std::max(double(s[p]), 0.0) / q;
      clam += lam;
      double gam = double(p+1) / (exp1 ? q : q-(r-p-1));
      double sigsq1 = clam / double(p+1);
      double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
      // sigsq2 > sigsq1 if signal else noise
      if (sigsq2 < sigsq1) {
        sigma2 = sigsq1;
        cutoff_p = p+1;
      }
    }

    if (cutoff_p > 0) {
      // recombine data using only eigenvectors above threshold:
      s.head (cutoff_p).setZero();
      s.tail (r-cutoff_p).setOnes();
      if (m <= n)
        X.col (n/2) = eig.eigenvectors() * ( s.asDiagonal() * ( eig.eigenvectors().adjoint() * X.col(n/2) ));
      else
        X.col (n/2) = X * ( eig.eigenvectors() * ( s.asDiagonal() * eig.eigenvectors().adjoint().col(n/2) ));
    }

    // Store output
    assign_pos_of(dwi).to(out);
    if (m <= n) {
      for (auto l = Loop (3) (out); l; ++l)
        out.value() = value_type (X(out.index(3), n/2));
    }
    else {
      for (auto l = Loop (3) (out); l; ++l)
        out.value() = value_type (X(out.index(3), n/2));
    }

    // store noise map if requested:
    if (noise.valid()) {
      assign_pos_of(dwi).to(noise);
      noise.value() = value_type (std::sqrt(sigma2));
    }
  }

  void load_data (ImageType& dwi)
  {
    pos[0] = dwi.index(0); pos[1] = dwi.index(1); pos[2] = dwi.index(2);
    // fill patch
    X.setZero();
    size_t k = 0;
    for (int z = -extent[2]; z <= extent[2]; z++) {
      dwi.index(2) = wrapindex(z, 2, dwi.size(2));
      for (int y = -extent[1]; y <= extent[1]; y++) {
        dwi.index(1) = wrapindex(y, 1, dwi.size(1));
        for (int x = -extent[0]; x <= extent[0]; x++, k++) {
          dwi.index(0) = wrapindex(x, 0, dwi.size(0));
          X.col(k) = dwi.row(3);
        }
      }
    }
    // reset image position
    dwi.index(0) = pos[0];
    dwi.index(1) = pos[1];
    dwi.index(2) = pos[2];
  }

private:
  const std::array<ssize_t, 3> extent;
  const ssize_t m, n, r, q;
  const bool exp1;
  MatrixX X;
  VectorX Xm;
  std::array<ssize_t, 3> pos;
  double sigma2;
  Image<bool> mask;
  ImageType noise;

  inline size_t wrapindex(int r, int axis, int max) const {
    int rr = pos[axis] + r;
    if (rr < 0)    rr = extent[axis] - r;
    if (rr >= max) rr = (max-1) - extent[axis] - r;
    return rr;
  }

};

void run ()
{
  auto dwi_in = Image<value_type>::open (argument[0]).with_direct_io(3);

  if (dwi_in.ndim() != 4 || dwi_in.size(3) <= 1)
    throw Exception ("input image must be 4-dimensional");

  Image<bool> mask;
  auto opt = get_options ("mask");
  if (opt.size()) {
    mask = Image<bool>::open (opt[0][0]);
    check_dimensions (mask, dwi_in, 0, 3);
  }

  auto header = Header (dwi_in);
  header.datatype() = DataType::Float32;
  auto dwi_out = Image<value_type>::create (argument[1], header);

  opt = get_options("extent");
  vector<int> extent = { DEFAULT_SIZE, DEFAULT_SIZE, DEFAULT_SIZE };
  if (opt.size()) {
    extent = parse_ints(opt[0][0]);
    if (extent.size() == 1)
      extent = {extent[0], extent[0], extent[0]};
    if (extent.size() != 3)
      throw Exception ("-extent must be either a scalar or a list of length 3");
    for (auto &e : extent)
      if (!(e & 1))
        throw Exception ("-extent must be a (list of) odd numbers");
  }

  bool exp1 = get_options("exp1").size();

  Image<value_type> noise;
  opt = get_options("noise");
  if (opt.size()) {
    header.ndim() = 3;
    noise = Image<value_type>::create (opt[0][0], header);
  }

  opt = get_options("datatype");
  if (!opt.size() || (int(opt[0][0]) == 0)) {
    DEBUG("Computing SVD with single precision.");
    DenoisingFunctor< Image<value_type> , float > func (dwi_in, extent, mask, noise, exp1);
    ThreadedLoop ("running MP-PCA denoising", dwi_in, 0, 3)
      .run (func, dwi_in, dwi_out);
  }
  else if (int(opt[0][0]) == 1) {
    DEBUG("Computing SVD with double precision.");
    DenoisingFunctor< Image<value_type> , double > func (dwi_in, extent, mask, noise, exp1);
    ThreadedLoop ("running MP-PCA denoising", dwi_in, 0, 3)
      .run (func, dwi_in, dwi_out);
  }
  else {
    assert(0);
  }

}

<commit_msg>dwidenoise: change interface to select choice instead of flag.<commit_after>/*
 * Copyright (c) 2008-2018 the MRtrix3 contributors.
 *
 * 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/
 *
 * MRtrix3 is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty
 * of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 *
 * For more details, see http://www.mrtrix.org/
 */

#include "command.h"
#include "image.h"

#include <Eigen/Dense>
#include <Eigen/Eigenvalues>

#define DEFAULT_SIZE 5

using namespace MR;
using namespace App;

const char* const dtypes[] = { "float32", "float64", NULL };

const char* const estimators[] = { "Exp1", "Exp2", NULL };

void usage ()
{
  SYNOPSIS = "Denoise DWI data and estimate the noise level based on the optimal threshold for PCA";

  DESCRIPTION
    + "DWI data denoising and noise map estimation by exploiting data redundancy in the PCA domain "
    "using the prior knowledge that the eigenspectrum of random covariance matrices is described by "
    "the universal Marchenko-Pastur (MP) distribution. Fitting the MP distribution to the spectrum "
    "of patch-wise signal matrices hence provides an estimator of the noise level 'sigma', as was "
    "first shown in Veraart et al. (2016) and later improved in Cordero-Grande et al. (2019). This "
    "noise level estimate then determines the optimal cut-off for PCA denoising."

    + "Important note: image denoising must be performed as the first step of the image processing pipeline. "
    "The routine will fail if interpolation or smoothing has been applied to the data prior to denoising."

    + "Note that this function does not correct for non-Gaussian noise biases.";

  AUTHOR = "Daan Christiaens (daan.christiaens@kcl.ac.uk) & "
           "Jelle Veraart (jelle.veraart@nyumc.org) & "
           "J-Donald Tournier (jdtournier@gmail.com)";

  REFERENCES
    + "Veraart, J.; Novikov, D.S.; Christiaens, D.; Ades-aron, B.; Sijbers, J. & Fieremans, E. " // Internal
    "Denoising of diffusion MRI using random matrix theory. "
    "NeuroImage, 2016, 142, 394-406, doi: 10.1016/j.neuroimage.2016.08.016"

    + "Veraart, J.; Fieremans, E. & Novikov, D.S. " // Internal
    "Diffusion MRI noise mapping using random matrix theory. "
    "Magn. Res. Med., 2016, 76(5), 1582-1593, doi: 10.1002/mrm.26059"

    + "Cordero-Grande, L.; Christiaens, D.; Hutter, J.; Price, A.N.; Hajnal, J.V. " // Internal
    "Complex diffusion-weighted image estimation via matrix recovery under general noise models. "
    "NeuroImage, 2019, 200, 391-404, doi: 10.1016/j.neuroimage.2019.06.039";

  ARGUMENTS
  + Argument ("dwi", "the input diffusion-weighted image.").type_image_in ()

  + Argument ("out", "the output denoised DWI image.").type_image_out ();

  OPTIONS
    + Option ("mask", "only perform computation within the specified binary brain mask image.")
    +   Argument ("image").type_image_in()

    + Option ("extent", "set the window size of the denoising filter. (default = " + str(DEFAULT_SIZE) + "," + str(DEFAULT_SIZE) + "," + str(DEFAULT_SIZE) + ")")
    +   Argument ("window").type_sequence_int ()

    + Option ("noise", "the output noise map.")
    +   Argument ("level").type_image_out()

    + Option ("datatype", "datatype for SVD (single or double precision).")
    +   Argument ("float32/float64").type_choice(dtypes)

    + Option ("estimator", "select noise level estimator (default = Exp2), either \n"
                           "Exp1: the original estimator used in Veraart et al. (2016), or \n"
                           "Exp2: the improved estimator introduced in Cordero-Grande et al. (2019).")
    +   Argument ("Exp1/Exp2").type_choice(estimators);

  COPYRIGHT = "Copyright (c) 2016 New York University, University of Antwerp, and the MRtrix3 contributors \n \n"
      "Permission is hereby granted, free of charge, to any non-commercial entity ('Recipient') obtaining a copy of this software and "
      "associated documentation files (the 'Software'), to the Software solely for non-commercial research, including the rights to "
      "use, copy and modify the Software, subject to the following conditions: \n \n"
      "\t 1. The above copyright notice and this permission notice shall be included by Recipient in all copies or substantial portions of "
      "the Software. \n \n"
      "\t 2. THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES"
      "OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE"
      "LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR"
      "IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. \n \n"
      "\t 3. In no event shall NYU be liable for direct, indirect, special, incidental or consequential damages in connection with the Software. "
      "Recipient will defend, indemnify and hold NYU harmless from any claims or liability resulting from the use of the Software by recipient. \n \n"
      "\t 4. Neither anything contained herein nor the delivery of the Software to recipient shall be deemed to grant the Recipient any right or "
      "licenses under any patents or patent application owned by NYU. \n \n"
      "\t 5. The Software may only be used for non-commercial research and may not be used for clinical care. \n \n"
      "\t 6. Any publication by Recipient of research involving the Software shall cite the references listed below.";

}

using value_type = float;

template <class ImageType, typename F = float>
class DenoisingFunctor {
  MEMALIGN(DenoisingFunctor)

public:

  typedef Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic> MatrixX;
  typedef Eigen::Matrix<F, Eigen::Dynamic, 1> VectorX;

  DenoisingFunctor (ImageType& dwi, vector<int> extent, Image<bool>& mask, ImageType& noise, bool exp1)
    : extent {{extent[0]/2, extent[1]/2, extent[2]/2}},
      m (dwi.size(3)), n (extent[0]*extent[1]*extent[2]),
      r (std::min(m,n)), q (std::max(m,n)), exp1(exp1),
      X (m,n), pos {{0, 0, 0}},
      mask (mask), noise (noise)
  { }

  void operator () (ImageType& dwi, ImageType& out)
  {
    // Process voxels in mask only
    if (mask.valid()) {
      assign_pos_of (dwi).to (mask);
      if (!mask.value())
        return;
    }

    // Load data in local window
    load_data (dwi);

    // Compute Eigendecomposition:
    MatrixX XtX (r,r);
    if (m <= n)
      XtX.template triangularView<Eigen::Lower>() = X * X.adjoint();
    else
      XtX.template triangularView<Eigen::Lower>() = X.adjoint() * X;
    Eigen::SelfAdjointEigenSolver<MatrixX> eig (XtX);
    // eigenvalues provide squared singular values, sorted in increasing order:
    VectorX s = eig.eigenvalues();

    // Marchenko-Pastur optimal threshold
    const double lam_r = std::max(double(s[0]), 0.0) / q;
    double clam = 0.0;
    sigma2 = 0.0;
    ssize_t cutoff_p = 0;
    for (ssize_t p = 0; p < r; ++p)     // p+1 is the number of noise components
    {                                   // (as opposed to the paper where p is defined as the number of signal components)
      double lam = std::max(double(s[p]), 0.0) / q;
      clam += lam;
      double gam = double(p+1) / (exp1 ? q : q-(r-p-1));
      double sigsq1 = clam / double(p+1);
      double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
      // sigsq2 > sigsq1 if signal else noise
      if (sigsq2 < sigsq1) {
        sigma2 = sigsq1;
        cutoff_p = p+1;
      }
    }

    if (cutoff_p > 0) {
      // recombine data using only eigenvectors above threshold:
      s.head (cutoff_p).setZero();
      s.tail (r-cutoff_p).setOnes();
      if (m <= n)
        X.col (n/2) = eig.eigenvectors() * ( s.asDiagonal() * ( eig.eigenvectors().adjoint() * X.col(n/2) ));
      else
        X.col (n/2) = X * ( eig.eigenvectors() * ( s.asDiagonal() * eig.eigenvectors().adjoint().col(n/2) ));
    }

    // Store output
    assign_pos_of(dwi).to(out);
    if (m <= n) {
      for (auto l = Loop (3) (out); l; ++l)
        out.value() = value_type (X(out.index(3), n/2));
    }
    else {
      for (auto l = Loop (3) (out); l; ++l)
        out.value() = value_type (X(out.index(3), n/2));
    }

    // store noise map if requested:
    if (noise.valid()) {
      assign_pos_of(dwi).to(noise);
      noise.value() = value_type (std::sqrt(sigma2));
    }
  }

  void load_data (ImageType& dwi)
  {
    pos[0] = dwi.index(0); pos[1] = dwi.index(1); pos[2] = dwi.index(2);
    // fill patch
    X.setZero();
    size_t k = 0;
    for (int z = -extent[2]; z <= extent[2]; z++) {
      dwi.index(2) = wrapindex(z, 2, dwi.size(2));
      for (int y = -extent[1]; y <= extent[1]; y++) {
        dwi.index(1) = wrapindex(y, 1, dwi.size(1));
        for (int x = -extent[0]; x <= extent[0]; x++, k++) {
          dwi.index(0) = wrapindex(x, 0, dwi.size(0));
          X.col(k) = dwi.row(3);
        }
      }
    }
    // reset image position
    dwi.index(0) = pos[0];
    dwi.index(1) = pos[1];
    dwi.index(2) = pos[2];
  }

private:
  const std::array<ssize_t, 3> extent;
  const ssize_t m, n, r, q;
  const bool exp1;
  MatrixX X;
  VectorX Xm;
  std::array<ssize_t, 3> pos;
  double sigma2;
  Image<bool> mask;
  ImageType noise;

  inline size_t wrapindex(int r, int axis, int max) const {
    int rr = pos[axis] + r;
    if (rr < 0)    rr = extent[axis] - r;
    if (rr >= max) rr = (max-1) - extent[axis] - r;
    return rr;
  }

};

void run ()
{
  auto dwi_in = Image<value_type>::open (argument[0]).with_direct_io(3);

  if (dwi_in.ndim() != 4 || dwi_in.size(3) <= 1)
    throw Exception ("input image must be 4-dimensional");

  Image<bool> mask;
  auto opt = get_options ("mask");
  if (opt.size()) {
    mask = Image<bool>::open (opt[0][0]);
    check_dimensions (mask, dwi_in, 0, 3);
  }

  auto header = Header (dwi_in);
  header.datatype() = DataType::Float32;
  auto dwi_out = Image<value_type>::create (argument[1], header);

  opt = get_options("extent");
  vector<int> extent = { DEFAULT_SIZE, DEFAULT_SIZE, DEFAULT_SIZE };
  if (opt.size()) {
    extent = parse_ints(opt[0][0]);
    if (extent.size() == 1)
      extent = {extent[0], extent[0], extent[0]};
    if (extent.size() != 3)
      throw Exception ("-extent must be either a scalar or a list of length 3");
    for (auto &e : extent)
      if (!(e & 1))
        throw Exception ("-extent must be a (list of) odd numbers");
  }

  bool exp1 = false;
  opt = get_options("estimator");
  if (opt.size()) {
    exp1 = (int(opt[0][0]) == 0);
  }

  Image<value_type> noise;
  opt = get_options("noise");
  if (opt.size()) {
    header.ndim() = 3;
    noise = Image<value_type>::create (opt[0][0], header);
  }

  opt = get_options("datatype");
  if (!opt.size() || (int(opt[0][0]) == 0)) {
    DEBUG("Computing SVD with single precision.");
    DenoisingFunctor< Image<value_type> , float > func (dwi_in, extent, mask, noise, exp1);
    ThreadedLoop ("running MP-PCA denoising", dwi_in, 0, 3)
      .run (func, dwi_in, dwi_out);
  }
  else if (int(opt[0][0]) == 1) {
    DEBUG("Computing SVD with double precision.");
    DenoisingFunctor< Image<value_type> , double > func (dwi_in, extent, mask, noise, exp1);
    ThreadedLoop ("running MP-PCA denoising", dwi_in, 0, 3)
      .run (func, dwi_in, dwi_out);
  }
  else {
    assert(0);
  }

}

<|endoftext|>


<commit_before>//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
//  This source code is licensed under the BSD-style license found in the
//  LICENSE file in the root directory of this source tree. An additional grant
//  of patent rights can be found in the PATENTS file in the same directory.
//
#ifndef ROCKSDB_LITE
#include "rocksdb/ldb_tool.h"
#include "rocksdb/utilities/ldb_cmd.h"
#include "tools/ldb_cmd_impl.h"

namespace rocksdb {

LDBOptions::LDBOptions() {}

void LDBCommandRunner::PrintHelp(const char* exec_name) {
  std::string ret;

  ret.append("ldb - LevelDB Tool");
  ret.append("\n\n");
  ret.append("commands MUST specify --" + LDBCommand::ARG_DB +
             "=<full_path_to_db_directory> when necessary\n");
  ret.append("\n");
  ret.append(
      "The following optional parameters control if keys/values are "
      "input/output as hex or as plain strings:\n");
  ret.append("  --" + LDBCommand::ARG_KEY_HEX +
             " : Keys are input/output as hex\n");
  ret.append("  --" + LDBCommand::ARG_VALUE_HEX +
             " : Values are input/output as hex\n");
  ret.append("  --" + LDBCommand::ARG_HEX +
             " : Both keys and values are input/output as hex\n");
  ret.append(
      "  --" + LDBCommand::ARG_CF_NAME +
      " : name of the column family to operate on. default: default column "
      "family\n");
  ret.append("\n");

  ret.append(
      "The following optional parameters control the database "
      "internals:\n");
  ret.append("  --" + LDBCommand::ARG_TTL +
             " with 'put','get','scan','dump','query','batchput'"
             " : DB supports ttl and value is internally timestamp-suffixed\n");
  ret.append("  --" + LDBCommand::ARG_BLOOM_BITS + "=<int,e.g.:14>\n");
  ret.append("  --" + LDBCommand::ARG_FIX_PREFIX_LEN + "=<int,e.g.:14>\n");
  ret.append("  --" + LDBCommand::ARG_COMPRESSION_TYPE +
             "=<no|snappy|zlib|bzip2|lz4|lz4hc|xpress|zstd>\n");
  ret.append("  --" + LDBCommand::ARG_COMPRESSION_MAX_DICT_BYTES +
             "=<int,e.g.:14>\n");
  ret.append("  --" + LDBCommand::ARG_BLOCK_SIZE + "=<block_size_in_bytes>\n");
  ret.append("  --" + LDBCommand::ARG_AUTO_COMPACTION + "=<true|false>\n");
  ret.append("  --" + LDBCommand::ARG_DB_WRITE_BUFFER_SIZE +
             "=<int,e.g.:16777216>\n");
  ret.append("  --" + LDBCommand::ARG_WRITE_BUFFER_SIZE +
             "=<int,e.g.:4194304>\n");
  ret.append("  --" + LDBCommand::ARG_FILE_SIZE + "=<int,e.g.:2097152>\n");

  ret.append("\n\n");
  ret.append("Data Access Commands:\n");
  PutCommand::Help(ret);
  GetCommand::Help(ret);
  BatchPutCommand::Help(ret);
  ScanCommand::Help(ret);
  DeleteCommand::Help(ret);
  DBQuerierCommand::Help(ret);
  ApproxSizeCommand::Help(ret);
  CheckConsistencyCommand::Help(ret);

  ret.append("\n\n");
  ret.append("Admin Commands:\n");
  WALDumperCommand::Help(ret);
  CompactorCommand::Help(ret);
  ReduceDBLevelsCommand::Help(ret);
  ChangeCompactionStyleCommand::Help(ret);
  DBDumperCommand::Help(ret);
  DBLoaderCommand::Help(ret);
  ManifestDumpCommand::Help(ret);
  ListColumnFamiliesCommand::Help(ret);
  DBFileDumperCommand::Help(ret);
  InternalDumpCommand::Help(ret);
  RepairCommand::Help(ret);
  BackupCommand::Help(ret);
  RestoreCommand::Help(ret);

  fprintf(stderr, "%s\n", ret.c_str());
}

void LDBCommandRunner::RunCommand(
    int argc, char** argv, Options options, const LDBOptions& ldb_options,
    const std::vector<ColumnFamilyDescriptor>* column_families) {
  if (argc <= 2) {
    PrintHelp(argv[0]);
    exit(1);
  }

  LDBCommand* cmdObj = LDBCommand::InitFromCmdLineArgs(
      argc, argv, options, ldb_options, column_families);
  if (cmdObj == nullptr) {
    fprintf(stderr, "Unknown command\n");
    PrintHelp(argv[0]);
    exit(1);
  }

  if (!cmdObj->ValidateCmdLineOptions()) {
    exit(1);
  }

  cmdObj->Run();
  LDBCommandExecuteResult ret = cmdObj->GetExecuteState();
  fprintf(stderr, "%s\n", ret.ToString().c_str());
  delete cmdObj;

  exit(ret.IsFailed());
}

void LDBTool::Run(int argc, char** argv, Options options,
                  const LDBOptions& ldb_options,
                  const std::vector<ColumnFamilyDescriptor>* column_families) {
  LDBCommandRunner::RunCommand(argc, argv, options, ldb_options,
                               column_families);
}
} // namespace rocksdb

#endif  // ROCKSDB_LITE
<commit_msg>16384 as e.g .value for compression_max_dict_bytes<commit_after>//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
//  This source code is licensed under the BSD-style license found in the
//  LICENSE file in the root directory of this source tree. An additional grant
//  of patent rights can be found in the PATENTS file in the same directory.
//
#ifndef ROCKSDB_LITE
#include "rocksdb/ldb_tool.h"
#include "rocksdb/utilities/ldb_cmd.h"
#include "tools/ldb_cmd_impl.h"

namespace rocksdb {

LDBOptions::LDBOptions() {}

void LDBCommandRunner::PrintHelp(const char* exec_name) {
  std::string ret;

  ret.append("ldb - LevelDB Tool");
  ret.append("\n\n");
  ret.append("commands MUST specify --" + LDBCommand::ARG_DB +
             "=<full_path_to_db_directory> when necessary\n");
  ret.append("\n");
  ret.append(
      "The following optional parameters control if keys/values are "
      "input/output as hex or as plain strings:\n");
  ret.append("  --" + LDBCommand::ARG_KEY_HEX +
             " : Keys are input/output as hex\n");
  ret.append("  --" + LDBCommand::ARG_VALUE_HEX +
             " : Values are input/output as hex\n");
  ret.append("  --" + LDBCommand::ARG_HEX +
             " : Both keys and values are input/output as hex\n");
  ret.append(
      "  --" + LDBCommand::ARG_CF_NAME +
      " : name of the column family to operate on. default: default column "
      "family\n");
  ret.append("\n");

  ret.append(
      "The following optional parameters control the database "
      "internals:\n");
  ret.append("  --" + LDBCommand::ARG_TTL +
             " with 'put','get','scan','dump','query','batchput'"
             " : DB supports ttl and value is internally timestamp-suffixed\n");
  ret.append("  --" + LDBCommand::ARG_BLOOM_BITS + "=<int,e.g.:14>\n");
  ret.append("  --" + LDBCommand::ARG_FIX_PREFIX_LEN + "=<int,e.g.:14>\n");
  ret.append("  --" + LDBCommand::ARG_COMPRESSION_TYPE +
             "=<no|snappy|zlib|bzip2|lz4|lz4hc|xpress|zstd>\n");
  ret.append("  --" + LDBCommand::ARG_COMPRESSION_MAX_DICT_BYTES +
             "=<int,e.g.:16384>\n");
  ret.append("  --" + LDBCommand::ARG_BLOCK_SIZE + "=<block_size_in_bytes>\n");
  ret.append("  --" + LDBCommand::ARG_AUTO_COMPACTION + "=<true|false>\n");
  ret.append("  --" + LDBCommand::ARG_DB_WRITE_BUFFER_SIZE +
             "=<int,e.g.:16777216>\n");
  ret.append("  --" + LDBCommand::ARG_WRITE_BUFFER_SIZE +
             "=<int,e.g.:4194304>\n");
  ret.append("  --" + LDBCommand::ARG_FILE_SIZE + "=<int,e.g.:2097152>\n");

  ret.append("\n\n");
  ret.append("Data Access Commands:\n");
  PutCommand::Help(ret);
  GetCommand::Help(ret);
  BatchPutCommand::Help(ret);
  ScanCommand::Help(ret);
  DeleteCommand::Help(ret);
  DBQuerierCommand::Help(ret);
  ApproxSizeCommand::Help(ret);
  CheckConsistencyCommand::Help(ret);

  ret.append("\n\n");
  ret.append("Admin Commands:\n");
  WALDumperCommand::Help(ret);
  CompactorCommand::Help(ret);
  ReduceDBLevelsCommand::Help(ret);
  ChangeCompactionStyleCommand::Help(ret);
  DBDumperCommand::Help(ret);
  DBLoaderCommand::Help(ret);
  ManifestDumpCommand::Help(ret);
  ListColumnFamiliesCommand::Help(ret);
  DBFileDumperCommand::Help(ret);
  InternalDumpCommand::Help(ret);
  RepairCommand::Help(ret);
  BackupCommand::Help(ret);
  RestoreCommand::Help(ret);

  fprintf(stderr, "%s\n", ret.c_str());
}

void LDBCommandRunner::RunCommand(
    int argc, char** argv, Options options, const LDBOptions& ldb_options,
    const std::vector<ColumnFamilyDescriptor>* column_families) {
  if (argc <= 2) {
    PrintHelp(argv[0]);
    exit(1);
  }

  LDBCommand* cmdObj = LDBCommand::InitFromCmdLineArgs(
      argc, argv, options, ldb_options, column_families);
  if (cmdObj == nullptr) {
    fprintf(stderr, "Unknown command\n");
    PrintHelp(argv[0]);
    exit(1);
  }

  if (!cmdObj->ValidateCmdLineOptions()) {
    exit(1);
  }

  cmdObj->Run();
  LDBCommandExecuteResult ret = cmdObj->GetExecuteState();
  fprintf(stderr, "%s\n", ret.ToString().c_str());
  delete cmdObj;

  exit(ret.IsFailed());
}

void LDBTool::Run(int argc, char** argv, Options options,
                  const LDBOptions& ldb_options,
                  const std::vector<ColumnFamilyDescriptor>* column_families) {
  LDBCommandRunner::RunCommand(argc, argv, options, ldb_options,
                               column_families);
}
} // namespace rocksdb

#endif  // ROCKSDB_LITE
<|endoftext|>


<commit_before>// Copyright(c) 2019 Jounayd Id Salah
// Distributed under the MIT License (See accompanying file LICENSE.md file or copy at http://opensource.org/licenses/MIT).
#include "math/pch.h"
#include "coQEF4.h"
#include "../../matrix/coSymMat4_f.h"
#include "../../matrix/coMat4_f.h"
#include "../../vector/coVec4_f.h"

// Implementations shared by:
// - https://github.com/nickgildea/qef (public domain)
// - https://github.com/Lin20/isosurface/tree/master/Isosurface/Isosurface/QEFSolver (MIT License)
// Docs:
// - https://web.stanford.edu/class/cme335/lecture7.pdf

const coFloat coQEF4::tolerance = 1e-6f;

void coSVD_Rotate_XY(coFloat& x, coFloat& y, coFloat c, coFloat s)
{
	const coFloat u = x;
	const coFloat v = y;
	x = c * u - s * v;
	y = s * u + c * v;
}

void coSVD_RotateQ_XY(coFloat& x, coFloat& y, coFloat& a, coFloat c, coFloat s)
{
	const coFloat cc = c * c;
	const coFloat ss = s * s;
	const coFloat mx = 2.0f * c * s * a;
	const coFloat u = x;
	const coFloat v = y;
	x = cc * u - mx + ss * v;
	y = ss * u + mx + cc * v;
}

void coSVD_Rotate(coMat4& vtav, coMat4& v, coInt a, coInt  b)
{
	coASSERT(false);
	if (vtav[a][b] == 0.0f)
		return;

	coFloat c, s;
	coCalcSymmetricGivensCoefficients(vtav[a][a], vtav[a][b], vtav[b][b], c, s);
	coSVD_RotateQ_XY(vtav[a][a], vtav[b][b], vtav[a][b], c, s);
	coSVD_Rotate_XY(vtav[0][3 - b], vtav[1 - a][2], c, s);
	vtav[a][b] = 0.0f;

	coSVD_Rotate_XY(v[0][a], v[0][b], c, s);
	coSVD_Rotate_XY(v[1][a], v[1][b], c, s);
	coSVD_Rotate_XY(v[2][a], v[2][b], c, s);
}

void coRot01_post(coMat4& this_, coFloat c, coFloat s)
{
// 	const coFloat m00 = this_[0][0], m01 = this_[0][1], m10 = this_[1][0], m11 = this_[1][1], m20 = this_[2][0],
// 		m21 = this_[2][1];
// 	this_ = coMat3(
// 		coFloatx3(c * m00 - s * m01, s * m00 + c * m01, this_[0][2]),
// 		coFloatx3(c * m10 - s * m11, s * m10 + c * m11, this_[1][2]),
// 		coFloatx3(c * m20 - s * m21, s * m20 + c * m21, this_[2][2]));
}

void coRot02_post(coMat4& this_, coFloat c, coFloat s)
{
// 	const coFloat m00 = this_[0][0], m02 = this_[0][2], m10 = this_[1][0], m12 = this_[1][2],
// 		m20 = this_[2][0], m22 = this_[2][2];
// 	this_ = coMat3(
// 		coFloatx3(c * m00 - s * m02, this_[0][1], s * m00 + c * m02),
// 		coFloatx3(c * m10 - s * m12, this_[1][1], s * m10 + c * m12),
// 		coFloatx3(c * m20 - s * m22, this_[2][1], s * m20 + c * m22));
}

void coRot03_post(coMat4& this_, coFloat c, coFloat s)
{
}

void coRot12_post(coMat4& this_, coFloat c, coFloat s)
{
// 	const coFloat m01 = this_[0][1], m02 = this_[0][2], m11 = this_[1][1], m12 = this_[1][2],
// 		m21 = this_[2][1], m22 = this_[2][2];
// 	this_ = coMat3(
// 		coFloatx3(this_[0][0], c * m01 - s * m02, s * m01 + c * m02),
// 		coFloatx3(this_[1][0], c * m11 - s * m12, s * m11 + c * m12),
// 		coFloatx3(this_[2][0], c * m21 - s * m22, s * m21 + c * m22));
}

void coRot13_post(coMat4& this_, coFloat c, coFloat s)
{
}

void coRot23_post(coMat4& this_, coFloat c, coFloat s)
{
}

void coSVD_Rotate01(coSymMat4& vtav, coMat4& v)
{
	if (vtav.m01 == 0.0f)
		return;

 	coFloat c = 0.0f;
 	coFloat s = 0.0f;
 	coRot01(vtav, c, s);
 	c = 0.0f;
 	s = 0.0f;
 	coRot01_post(v, c, s);
}

void coSVD_Rotate02(coSymMat4& vtav, coMat4& v)
{
	if (vtav.m02 == 0.0f)
		return;

	coFloat c = 0.0f;
	coFloat s = 0.0f;
	coRot02(vtav, c, s);
	c = 0.0f;
	s = 0.0f;
	coRot02_post(v, c, s);
}

void coSVD_Rotate03(coSymMat4& vtav, coMat4& v)
{
	if (vtav.m03 == 0.0f)
		return;

	coFloat c = 0.0f;
	coFloat s = 0.0f;
	coRot03(vtav, c, s);
	c = 0.0f;
	s = 0.0f;
	coRot03_post(v, c, s);
}

void coSVD_Rotate12(coSymMat4& vtav, coMat4& v)
{
	if (vtav.m12 == 0)
		return;

	coFloat c = 0.0f;
	coFloat s = 0.0f;
	coRot12(vtav, c, s);
	c = 0.0f;
	s = 0.0f;
	coRot12_post(v, c, s);
}

void coSVD_Rotate13(coSymMat4& vtav, coMat4& v)
{
	if (vtav.m13 == 0)
		return;

	coFloat c = 0.0f;
	coFloat s = 0.0f;
	coRot13(vtav, c, s);
	c = 0.0f;
	s = 0.0f;
	coRot13_post(v, c, s);
}

void coSVD_Rotate23(coSymMat4& vtav, coMat4& v)
{
	if (vtav.m23 == 0)
		return;

	coFloat c = 0.0f;
	coFloat s = 0.0f;
	coRot23(vtav, c, s);
	c = 0.0f;
	s = 0.0f;
	coRot23_post(v, c, s);
}

void coSVD_SolveSym(coVec4& sigma, coMat4& v, const coMat4& a)
{
 	coSymMat4 vtav(a);

 	const coFloat delta = coQEF4::tolerance * coFrobeniusNorm(vtav);
 	for (coInt i = 0; i < 10 && coOff(vtav) > delta; ++i)
 	{
 		coSVD_Rotate01(vtav, v);
		coSVD_Rotate02(vtav, v);
		coSVD_Rotate03(vtav, v);
		coSVD_Rotate12(vtav, v);
		coSVD_Rotate13(vtav, v);
		coSVD_Rotate23(vtav, v);
 	}
 	sigma = coVec4(vtav.m00, vtav.m11, vtav.m22, vtav.m33);
}

void coSVD_SolveSym2(coVec4& sigma, coMat4& V, const coMat4& A)
{
	coMat4 U = A;
	coMat4 v2;

	coFloat converge = 777.0f;
	for (coUint it = 0; it < 40; ++it)
	//while (converge > 0.0001f)
	{
		converge = 0.0f;
		for (coUint i = 1; i < 4; ++i)
		{
			for (coUint j = 0; j < i; ++j)
			{
				const coFloat alpha = coSquareLength(U[i]).x;
				const coFloat beta = coSquareLength(U[j]).x;
				const coFloat gamma = coDot(U[i], U[j]).x;
				converge = coMax(converge, coAbs(gamma) / coSquareRoot(alpha*beta));

				if (coAbs(gamma) < 1e-12f)
				{
					continue;
				}
				const coFloat zeta = (beta - alpha) / (2.0f * gamma);
				const coFloat t = coSign(zeta) / (coAbs(zeta) + coSquareRoot(1.0f + (zeta*zeta)));
				const coFloat c = 1.0f / coSquareRoot(1.0f + (t*t));
				const coFloat s = c * t;

				const coVec4 ui = U[i];
				U[i] = c*ui - s*U[j];
				U[j] = s*ui + c*U[j];

				const coVec4 vi = v2[i];
				v2[i] = c*vi - s*v2[j];
				v2[j] = s*vi + c*v2[j];
			}
		}
	}

	for (coUint i = 0; i < 4; ++i)
	{
		sigma[i] = coLength(U[i]).x;
	}

	V = v2;
}

void coSVD_ComputePseudoInverse(coMat4& out, const coVec4& sigma, const coMat4& v)
{
	coVec4 invSigma;
	invSigma.x = (coAbs(sigma.x) < coQEF4::tolerance || coAbs(1.0f / sigma.x) < coQEF4::tolerance) ? 0.0f : (1.0f / sigma.x);
	invSigma.y = (coAbs(sigma.y) < coQEF4::tolerance || coAbs(1.0f / sigma.y) < coQEF4::tolerance) ? 0.0f : (1.0f / sigma.y);
	invSigma.z = (coAbs(sigma.z) < coQEF4::tolerance || coAbs(1.0f / sigma.z) < coQEF4::tolerance) ? 0.0f : (1.0f / sigma.z);
	invSigma.w = (coAbs(sigma.w) < coQEF4::tolerance || coAbs(1.0f / sigma.w) < coQEF4::tolerance) ? 0.0f : (1.0f / sigma.w);
	out = v * coMat4(invSigma) * coTranspose(v);
}

void coSVD_SolveSymmetric(coVec4& out, const coMat4& A, const coVec4& b)
{
	coMat4 V;
	coVec4 sigma(nullptr);
	coSVD_SolveSym2(sigma, V, A);
	coMat4 invV(nullptr);
	coSVD_ComputePseudoInverse(invV, sigma, V);
	out = invV * b;
}

coVec4 coQEF4::Solve()
{
	coVec4 x;
	coSVD_SolveSymmetric(x, ATA, ATb);
	return x;

// 	coASSERT(nb);
// 	massPoint = pointAccum / coFloat(nb);
// 	const coVec4 ATb2 = ATb - ATA * massPoint;
// 	coVec4 x;
// 	coSVD_SolveSymmetric(x, ATA, ATb2);
// 	return x + massPoint;
}

coFloatx4 coQEF4::ComputeError(const coVec4& x) const
{
	// E(x) = (Ax-b)T(Ax-b)
	//      = x.ATAx - 2x.ATb + btb
	const coVec4 atax = ATA * x;
	return coDot(x, atax) - 2.0f * coDot(x, ATb) + btb;
}

void coQEF4::Add(const coVec4& p, const coVec4& n)
{
	const coVec4 nx = coBroadcastX(n);
	const coVec4 ny = coBroadcastY(n);
	const coVec4 nz = coBroadcastZ(n);
	const coVec4 nw = coBroadcastW(n);
	ATA.c0 += nx * n;
	ATA.c1 += ny * n;
	ATA.c2 += nz * n;
	ATA.c3 += nw * n;
	const coFloatx4 d = coDot(p, n);
	ATb += n * d;
	btb += (p * p).x;
	pointAccum += p;
	++nb;
}

void coQEF4::Add(const coQEF4& qef)
{
	ATA += qef.ATA;
	ATb += qef.ATb;
	btb += qef.btb;
	pointAccum += qef.pointAccum;
	nb += qef.nb;
}
<commit_msg>More qef work.<commit_after>// Copyright(c) 2019 Jounayd Id Salah
// Distributed under the MIT License (See accompanying file LICENSE.md file or copy at http://opensource.org/licenses/MIT).
#include "math/pch.h"
#include "coQEF4.h"
#include "../../matrix/coSymMat4_f.h"
#include "../../matrix/coMat4_f.h"
#include "../../vector/coVec4_f.h"

// Implementations shared by:
// - https://github.com/nickgildea/qef (public domain)
// - https://github.com/Lin20/isosurface/tree/master/Isosurface/Isosurface/QEFSolver (MIT License)
// Docs:
// - https://web.stanford.edu/class/cme335/lecture7.pdf

const coFloat coQEF4::tolerance = 1e-6f;

void coSVD_SolveSym2(coVec4& sigma, coMat4& V, const coMat4& A)
{
	coMat4 U = A;
	coMat4 v2;

	coFloat converge = 777.0f;
	for (coUint it = 0; it < 40; ++it)
	//while (converge > 0.0001f)
	{
		converge = 0.0f;
		for (coUint i = 1; i < 4; ++i)
		{
			for (coUint j = 0; j < i; ++j)
			{
				const coFloat alpha = coSquareLength(U[i]).x;
				const coFloat beta = coSquareLength(U[j]).x;
				const coFloat gamma = coDot(U[i], U[j]).x;
				converge = coMax(converge, coAbs(gamma) / coSquareRoot(alpha*beta));

				if (coAbs(gamma) < 1e-12f)
				{
					continue;
				}
				const coFloat zeta = (beta - alpha) / (2.0f * gamma);
				const coFloat t = coSign(zeta) / (coAbs(zeta) + coSquareRoot(1.0f + (zeta*zeta)));
				const coFloat c = 1.0f / coSquareRoot(1.0f + (t*t));
				const coFloat s = c * t;

				const coVec4 ui = U[i];
				U[i] = c*ui - s*U[j];
				U[j] = s*ui + c*U[j];

				const coVec4 vi = v2[i];
				v2[i] = c*vi - s*v2[j];
				v2[j] = s*vi + c*v2[j];
			}
		}
	}

	for (coUint i = 0; i < 4; ++i)
	{
		sigma[i] = coLength(U[i]).x;
	}

	V = v2;
}

void coSVD_ComputePseudoInverse(coMat4& out, const coVec4& sigma, const coMat4& v)
{
	coVec4 invSigma;
	invSigma.x = (coAbs(sigma.x) < coQEF4::tolerance || coAbs(1.0f / sigma.x) < coQEF4::tolerance) ? 0.0f : (1.0f / sigma.x);
	invSigma.y = (coAbs(sigma.y) < coQEF4::tolerance || coAbs(1.0f / sigma.y) < coQEF4::tolerance) ? 0.0f : (1.0f / sigma.y);
	invSigma.z = (coAbs(sigma.z) < coQEF4::tolerance || coAbs(1.0f / sigma.z) < coQEF4::tolerance) ? 0.0f : (1.0f / sigma.z);
	invSigma.w = (coAbs(sigma.w) < coQEF4::tolerance || coAbs(1.0f / sigma.w) < coQEF4::tolerance) ? 0.0f : (1.0f / sigma.w);
	out = v * coMat4(invSigma) * coTranspose(v);
}

void coSVD_SolveSymmetric(coVec4& out, const coMat4& A, const coVec4& b)
{
	coMat4 V;
	coVec4 sigma(nullptr);
	coSVD_SolveSym2(sigma, V, A);
	coMat4 invV(nullptr);
	coSVD_ComputePseudoInverse(invV, sigma, V);
	out = invV * b;
}

coVec4 coQEF4::Solve()
{
	coASSERT(nb);
	massPoint = pointAccum / coFloat(nb);
	const coVec4 ATb2 = ATb - ATA * massPoint;
	coVec4 x;
	coSVD_SolveSymmetric(x, ATA, ATb2);
	return x + massPoint;
}

coFloatx4 coQEF4::ComputeError(const coVec4& x) const
{
	// E(x) = (Ax-b)T(Ax-b)
	//      = x.ATAx - 2x.ATb + btb
	const coVec4 atax = ATA * x;
	return coDot(x, atax) - 2.0f * coDot(x, ATb) + btb;
}

void coQEF4::Add(const coVec4& p, const coVec4& n)
{
	const coVec4 nx = coBroadcastX(n);
	const coVec4 ny = coBroadcastY(n);
	const coVec4 nz = coBroadcastZ(n);
	const coVec4 nw = coBroadcastW(n);
	ATA.c0 += nx * n;
	ATA.c1 += ny * n;
	ATA.c2 += nz * n;
	ATA.c3 += nw * n;
	const coFloatx4 d = coDot(p, n);
	ATb += n * d;
	btb += (p * p).x;
	pointAccum += p;
	++nb;
}

void coQEF4::Add(const coQEF4& qef)
{
	ATA += qef.ATA;
	ATb += qef.ATb;
	btb += qef.btb;
	pointAccum += qef.pointAccum;
	nb += qef.nb;
}
<|endoftext|>