#include <boost/config.hpp>
#include <boost/program_options.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <iostream>
#include <random>
#include <sys/time.h>
#include <exception>

namespace po = boost::program_options;

static int SIZEX = 800;
static int SIZEY = 600;
static int ITER = 20;
static int RADIUS = 15;
static double TEMPERATURE = 2.0;
static double alpha = 0.05;
static int mouse_x, mouse_y;
static bool pressed = false;
static bool run_evolve = false;
static bool DEBUG = false;

inline double double_rand(const double & min, const double & max)
{
    static thread_local std::mt19937 generator;
    std::uniform_real_distribution<double> distribution(min,max);
    return distribution(generator);
}

void mouse_callback(int event, int x, int y, int flags, void *data)
{
    switch(event)
    {
        case CV_EVENT_MOUSEMOVE:
            mouse_x = x;
            mouse_y = y;

            // printf("(%d,%d)\n", x, y);
            break;
        case CV_EVENT_LBUTTONDOWN:
            pressed = true;
            break;
        case CV_EVENT_LBUTTONUP:
            pressed = false;
            break;
    }
}

void evolve(double **data, double **buffer)
{
    // ~ 62 bits of accuracy (as in of conservation only)
    int i,j;
    double dr = 1. / SIZEX;
    double dt = dr * dr;
    double *u = *data;
    double *next_state = *buffer;

    struct timeval start, end;
    gettimeofday(&start, NULL);

    for(int k = 0; k < ITER; k++)
    {
    #pragma omp parallel for private(i, j), shared(u, next_state, dt, dr, alpha)
        for(i=0; i < SIZEY; i++)
        {
            if(i == 0)
            {
                next_state[0]= u[0] + alpha * (
                    u[1] +
                    u[SIZEX] -
                    2*u[0]); // j == 0

                next_state[SIZEX - 1]= u[SIZEX - 1] + alpha * (
                    u[2 * SIZEX -1] +
                    u[SIZEX - 2] -
                    2*u[SIZEX - 1]); // j == SIZEX - 1

                for(j = 1; j < SIZEX-1; j++)
                    next_state[j]= u[j] + alpha * (
                        u[SIZEX + j] +
                        u[j + 1] +
                        u[j - 1] -
                        3*u[j]);
            }
            else if(i == SIZEY - 1)
            {
                next_state[i*SIZEX]= u[i*SIZEX] + alpha * (
                    u[(i-1)*SIZEX] +
                    u[i*SIZEX + 1] -
                    2*u[i*SIZEX]); // j == 0

                next_state[i*SIZEX + SIZEX - 1]= u[i*SIZEX + SIZEX - 1] + alpha * (
                    u[(i-1)*SIZEX + SIZEX - 1] +
                    u[i*SIZEX + SIZEX - 2] -
                    2*u[i*SIZEX + SIZEX - 1]); // j == SIZEX - 1

                for(j = 1; j < SIZEX-1; j++)
                    next_state[i*SIZEX + j]= u[i*SIZEX + j] + alpha * (
                        u[(i-1)*SIZEX + j] +
                        u[i*SIZEX + j + 1] +
                        u[i*SIZEX + j - 1] -
                        3*u[i*SIZEX + j]);
            }
            else
            {
                next_state[i * SIZEX]= u[i * SIZEX] + alpha * (
                    u[(i-1)*SIZEX] +
                    u[(i+1)*SIZEX] +
                    u[i*SIZEX + 1] -
                    3*u[i*SIZEX]); // j == 0

                next_state[i * SIZEX + SIZEX - 1]= u[i * SIZEX + SIZEX - 1] + alpha * (
                    u[(i-1)*SIZEX + SIZEX - 1] +
                    u[(i+1)*SIZEX + SIZEX - 1] +
                    u[i*SIZEX + SIZEX -2] -
                    3*u[i*SIZEX + SIZEX -1]); // j == SIZEX - 1

                for(j=1; j < SIZEX-1; j++)
                    next_state[i*SIZEX + j]= u[i * SIZEX + j] + alpha * (
                        u[(i-1)*SIZEX + j] +
                        u[(i+1)*SIZEX + j] +
                        u[i*SIZEX + j + 1] +
                        u[i*SIZEX + j - 1] -
                        4*u[i*SIZEX + j]);
            }

        }
        if(k != ITER)
        {
            auto tmp = next_state;
            next_state = u;
            u = tmp;
        }
    }
    *buffer = *data;
    *data = next_state;

    if(DEBUG)
    {
        double sum = 0.0;
        #pragma omp parallel for private(i) reduction(+:sum), shared(next_state)
        for(i=0; i < SIZEY*SIZEX; i++) sum += next_state[i];

        gettimeofday(&end, NULL);

        double delta = ((end.tv_sec  - start.tv_sec) * 1000000u +
            end.tv_usec - start.tv_usec) / 1.e3;
        printf("time (ms): %.5f \t total: %.3f \t average: %.4f\n", delta, sum, sum / SIZEX / SIZEY);
        // std::cout << "time (ms): " << delta << "\t\ttotal: " << sum << std::endl;
    }
}

void parse_arguments(int argc, char** argv)
{
    try
    {
        po::options_description desc("Allowed arguments");
        desc.add_options()
            ("help,h", "print this message.")
            ("height,y", po::value<int>()->default_value(600), "set height of the window.")
            ("width,x", po::value<int>()->default_value(800), "set width of the window")
            ("alpha,a", po::value<double>()->default_value(0.05), "alpha parameter of heat equation.")
            ("iter,i", po::value<int>()->default_value(20), "number of iterations per render.")
            ("size,s", po::value<int>()->default_value(15), "radius of the brush when drawing.")
            ("temp,t", po::value<double>()->default_value(1.3), "temperature of the brush.")
            ("debug,d", "debug flag to print extra states' info");

        po::variables_map vmap;
        po::store(po::parse_command_line(argc, argv, desc), vmap);
        po::notify(vmap);

        if(vmap.count("help"))
        {
            std::cout << desc;
            exit(0);
        }
        SIZEY = vmap["height"].as<int>();
        SIZEX = vmap["width"].as<int>();
        alpha = vmap["alpha"].as<double>();
        ITER = vmap["iter"].as<int>();
        RADIUS = vmap["size"].as<int>();
        TEMPERATURE = vmap["temp"].as<double>();
        if(vmap.count("debug")) DEBUG = true;
    }
    catch(std::exception& e)
    {
        std::cerr << "error: " << e.what() << std::endl;
        exit(1);
    }
    catch(...)
    {
        std::cerr << "error: Unknown error" << std::endl;
        exit(2);
    }
}

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

    double *data = new double[SIZEX*SIZEY];
    double *buffer = new double[SIZEX*SIZEY];
    cv::Mat image(SIZEY, SIZEX, CV_64F, data);
    cv::Mat scaled, color;

    std::string name = "Heat equation visualization";
    cv::namedWindow(name, cv::WINDOW_AUTOSIZE);
    cv::setMouseCallback(name, mouse_callback, nullptr);
    int val,x,y;

    auto vcap = new cv::VideoWriter("heat2d_video.avi",
            CV_FOURCC('M','J','P','G'),
            120,
            cv::Size(SIZEX, SIZEY),
            true);

    while( (val = cv::waitKey(1000/60)) != 27) // escape
    {
        if(val == 32) //space
        {
            run_evolve = !run_evolve;
        }
        if((x != mouse_x || y != mouse_y) && pressed)
        {
            x = mouse_x;
            y = mouse_y;
            cv::circle(image, cv::Point2d(x,y), RADIUS, cv::Scalar(TEMPERATURE), -1);
        }
        image.convertTo(scaled, CV_8UC1,255, 0);
        cv::applyColorMap(scaled, color, cv::COLORMAP_JET);
        if(run_evolve)
        {
            evolve(&data, &buffer);
        }
        image = cv::Mat(SIZEY, SIZEX, CV_64F, data);
        vcap->write(color);
        cv::imshow(name, color);
    }

    vcap->release();

    delete[] data;
    delete[] buffer;
    return 0;
}


#include "Polygon.h"

#include <algorithm>
#include <numeric>

#include <boost/range/adaptors.hpp>
#include <boost/range/algorithm.hpp>

#include "Vector.h"
#include "Vertex.h"

Polygon::classification Polygon::classify(const Plane& plane) const
{
	size_t count = points.size();

	size_t front = 0, back = 0, onplane = 0;

	for (const auto& p : points)
	{
		auto test = plane.evaluate(p);

		if (test <= 0) back++;
		if (test >= 0) front++;
		if (test == 0) onplane++;
	}

	if (onplane == count) return Polygon::classification::onPlane;
	if (front == count) return Polygon::classification::front;
	if (back == count) return Polygon::classification::back;
	return Polygon::classification::spanning;
}

Vertex Polygon::origin() const
{
	return std::accumulate(points.cbegin(), points.cend(), Vertex{ 0, 0, 0 }) / double(points.size());
}

Plane Polygon::plane() const
{
	if(points.size() < 3)
		return Plane();

	return { points[0], points[1], points[2] };
}

void Polygon::rotate(const Vertex & point, const Matrix3d & rotmat)
{
	for (auto &p : points)
		p = p.rotate(point, rotmat);
}

void Polygon::move(const Vertex & v)
{
	for (auto &p : points)
		p += v;
}

void Polygon::moveTo(const Vertex & p)
{
	Vertex dist = p - origin();

	for (auto &p : points)
		p += dist;
}

void Polygon::scale(const Vertex &scale)
{
	this->scale(origin(), scale);
}

void Polygon::scale(const Vertex& origin, const Vertex& scale)
{
	for (auto &point : points)
	{
		Vertex vec = Vector::diff(origin, point).vec();
		vec.x(vec.x() * scale.x());
		vec.y(vec.y() * scale.y());
		vec.z(vec.z() * scale.z());

		point = origin + vec;
	}
}

void Polygon::sliceThis(const Plane &plane)
{
	auto[back, front] = slice(plane);
	if (!back.points.empty() && !front.points.empty())
		points = back.points;
}

std::pair<Polygon, Polygon> Polygon::slice(const Plane& plane) const
{
	auto classification = classify(plane);

	std::pair<Polygon, Polygon> ret;

	if (classification != classification::spanning)
	{
		if (classification == classification::back) ret.first = *this;
		else if (classification == classification::front) ret.second = *this;
		return ret;
	}

	size_t prev = 0;

	for (size_t i = 0; i <= points.size(); i++)
	{
		size_t index = i % points.size();
		Vertex end = points[index];
		auto c = plane.evaluate(end);

		if (i > 0 && c != 0 && prev != 0 && c != prev)
		{
			Vertex start = points[i - 1];
			Vector line = Vector::diff(start, end);
			Vertex intersect = Plane::intersectPoint(plane, line);
			if (!Vertex::isVertex(intersect))
				throw new std::exception("Expected intersection");
			ret.first.points.push_back(intersect);
			ret.second.points.push_back(intersect);
		}

		if (i < points.size())
		{
			if (c <= 0) ret.first.points.push_back(end);
			if (c >= 0) ret.second.points.push_back(end);
		}

		prev = c;
	}

	return ret;
}

void Polygon::flip()
{
	std::reverse(points.begin(), points.end());
}

void Polygon::roundPoints(size_t precision)
{
	double exp = std::pow(10, precision);
	std::transform(points.begin(), points.end(), points.begin(), [exp](Vertex v) {
		v.x(std::round(v.x() * exp) / exp);
		v.y(std::round(v.y() * exp) / exp);
		v.z(std::round(v.z() * exp) / exp);
		return v;
	});
}

Vertex Polygon::intersectPoint(const Vector& line, int flags) const
{
	Plane plane = this->plane();

	auto point = Plane::intersectPoint(plane, line);

	Vertex defaultRet = (flags & lineBoundsFlag::RETURN_END_ON_FAIL) > 0 ? line.end() : Vertex();

	flags = flags & lineBoundsFlag::ALLOW_BOTH;

	if (!Vertex::isVertex(point))
		return defaultRet;

	bool test = testCollision(point);

	if (!test)
		return defaultRet;

	if (flags != lineBoundsFlag::ALLOW_BOTH)
	{
		double position = line.calculatePosition(point);

		if ((lineBoundsFlag::ALLOW_BACK & flags) == 0 && (position < 0 || doubleeq(position, 0)))
			return defaultRet;

		if ((lineBoundsFlag::ALLOW_FRONT & flags) == 0 && (position > 1 || doubleeq(position, 1)))
			return defaultRet;
	}

	return point;
}

bool Polygon::testCollision(const Vertex& point) const
{
	double sum = 0;

	for (size_t n = 0; n < points.size(); n++)
	{
		Vertex p1 = points[n] - point;
		Vertex p2 = points[(n + 1) % points.size()] - point;

		double nom = p1.length() * p2.length();

		if (doubleeq(nom, 0))
			return false;

		sum += acos(p1.dotProduct(p2) / nom);
	}

	return doubleeq(sum, M_PI * 2);
}

bool Polygon::testCollision(const Vector& line, int flags) const
{
	// Ensure method returns NaN vector if it fails so we can test it.
	Vertex intersect = intersectPoint(line, flags);

	if (!Vertex::isVertex(intersect))
		return false;

	return true;
}

bool Polygon::testCollision(const Polygon& polygon) const
{
	if (&polygon == this)
		return true;

	if (points.size() < 3 || polygon.points.size() < 3)
		return false;

	Vertex thisNorm = this->plane().normal();
	Vertex polyNorm = polygon.plane().normal();

	if (thisNorm == polyNorm || thisNorm == -polyNorm)
		return false;

	for (size_t n = 0; n < points.size(); n++)
	{
		Vector line = Vector::diff(points[n], points[(n + 1) % points.size()]);

		if (polygon.testCollision(line, ALLOW_NONE))
			return true;
	}

	for (size_t n = 0; n < polygon.points.size(); n++)
	{
		Vector line = Vector::diff(polygon.points[n], polygon.points[(n + 1) % polygon.points.size()]);

		if (this->testCollision(line, ALLOW_NONE))
			return true;
	}

	return false;
}

Polygon::Polygon(const Plane & p)
{
	auto dir = p.closestAxisToNormal();
	auto tempV = dir == Vertex::unitZ ? -Vertex::unitY : -Vertex::unitZ;

	auto normal = p.normal();

	auto up = tempV.crossProduct(normal).normalize();
	auto right = normal.crossProduct(up).normalize();

	points = {
		p.p1() + right + up,
		p.p1() - right + up,
		p.p1() - right - up,
		p.p1() + right + up
	};

	auto orig = origin();

	std::transform(points.begin(), points.end(), points.begin(), [&orig](Vertex v) {
		return (v - orig).normalize() * 10'000'000.0 + orig;
	});

}

Polygon::Polygon(const std::initializer_list<Vertex>& points)
	: points(points)
{
}


#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <math.h>
#include <stdlib.h>

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

#define DADEPT_FLOATING_POINT_TYPE float
#include <adept_source.h>
#include <adept_arrays.h>
using adept::adouble;
using adept::aMatrix;
using adept::aVector;

using adept::Vector;
using adept::adouble;
using adept::aReal;

//extern "C" {
// from https://github.com/AndrewCarterUK/mnist-neural-network-plain-c

#define MNIST_LABEL_MAGIC 0x00000801
#define MNIST_IMAGE_MAGIC 0x00000803
#define MNIST_IMAGE_WIDTH 28
#define MNIST_IMAGE_HEIGHT 28
#define MNIST_IMAGE_SIZE MNIST_IMAGE_WIDTH * MNIST_IMAGE_HEIGHT
#define MNIST_LABELS 10

typedef struct mnist_label_file_header_t_ {
    uint32_t magic_number;
    uint32_t number_of_labels;
} __attribute__((packed)) mnist_label_file_header_t;

typedef struct mnist_image_file_header_t_ {
    uint32_t magic_number;
    uint32_t number_of_images;
    uint32_t number_of_rows;
    uint32_t number_of_columns;
} __attribute__((packed)) mnist_image_file_header_t;

typedef struct mnist_image_t_ {
    uint8_t pixels[MNIST_IMAGE_SIZE];
} __attribute__((packed)) mnist_image_t;

typedef struct mnist_dataset_t_ {
    mnist_image_t * images;
    uint8_t * labels;
    uint32_t size;
} mnist_dataset_t;

typedef struct neural_network_t_ {
    float b[MNIST_LABELS];
    float W[MNIST_LABELS][MNIST_IMAGE_SIZE];
} neural_network_t;

typedef struct aneural_network_t_ {
    adept::FixedArray<float,true,MNIST_LABELS> b;
    adept::FixedArray<float,true,MNIST_LABELS,MNIST_IMAGE_SIZE> W;
} aneural_network_t;

typedef struct neural_network_gradient_t_ {
    float b_grad[MNIST_LABELS];
    float W_grad[MNIST_LABELS][MNIST_IMAGE_SIZE];
} neural_network_gradient_t;

/**
 * Convert from the big endian format in the dataset if we're on a little endian
 * machine.
 */
uint32_t map_uint32(uint32_t in)
{
#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
    return (
        ((in & 0xFF000000) >> 24) |
        ((in & 0x00FF0000) >>  8) |
        ((in & 0x0000FF00) <<  8) |
        ((in & 0x000000FF) << 24)
    );
#else
    return in;
#endif
}

/**
 * Read labels from file.
 *
 * File format: http://yann.lecun.com/exdb/mnist/
 */
uint8_t * get_labels(const char * path, uint32_t * number_of_labels)
{
    FILE * stream;
    mnist_label_file_header_t header;
    uint8_t * labels;

    stream = fopen(path, "rb");

    if (NULL == stream) {
        fprintf(stderr, "Could not open file: %s\n", path);
        return NULL;
    }

    if (1 != fread(&header, sizeof(mnist_label_file_header_t), 1, stream)) {
        fprintf(stderr, "Could not read label file header from: %s\n", path);
        fclose(stream);
        return NULL;
    }

    header.magic_number = map_uint32(header.magic_number);
    header.number_of_labels = map_uint32(header.number_of_labels);

    if (MNIST_LABEL_MAGIC != header.magic_number) {
        fprintf(stderr, "Invalid header read from label file: %s (%08X not %08X)\n", path, header.magic_number, MNIST_LABEL_MAGIC);
        fclose(stream);
        return NULL;
    }

    *number_of_labels = header.number_of_labels;

    labels = (uint8_t*)malloc(*number_of_labels * sizeof(uint8_t));

    if (labels == NULL) {
        fprintf(stderr, "Could not allocated memory for %d labels\n", *number_of_labels);
        fclose(stream);
        return NULL;
    }

    if (*number_of_labels != fread(labels, 1, *number_of_labels, stream)) {
        fprintf(stderr, "Could not read %d labels from: %s\n", *number_of_labels, path);
        free(labels);
        fclose(stream);
        return NULL;
    }

    fclose(stream);

    return labels;
}

/**
 * Read images from file.
 *
 * File format: http://yann.lecun.com/exdb/mnist/
 */
mnist_image_t * get_images(const char * path, uint32_t * number_of_images)
{
    FILE * stream;
    mnist_image_file_header_t header;
    mnist_image_t * images;

    stream = fopen(path, "rb");

    if (NULL == stream) {
        fprintf(stderr, "Could not open file: %s\n", path);
        return NULL;
    }

    if (1 != fread(&header, sizeof(mnist_image_file_header_t), 1, stream)) {
        fprintf(stderr, "Could not read image file header from: %s\n", path);
        fclose(stream);
        return NULL;
    }

    header.magic_number = map_uint32(header.magic_number);
    header.number_of_images = map_uint32(header.number_of_images);
    header.number_of_rows = map_uint32(header.number_of_rows);
    header.number_of_columns = map_uint32(header.number_of_columns);

    if (MNIST_IMAGE_MAGIC != header.magic_number) {
        fprintf(stderr, "Invalid header read from image file: %s (%08X not %08X)\n", path, header.magic_number, MNIST_IMAGE_MAGIC);
        fclose(stream);
        return NULL;
    }

    if (MNIST_IMAGE_WIDTH != header.number_of_rows) {
        fprintf(stderr, "Invalid number of image rows in image file %s (%d not %d)\n", path, header.number_of_rows, MNIST_IMAGE_WIDTH);
    }

    if (MNIST_IMAGE_HEIGHT != header.number_of_columns) {
        fprintf(stderr, "Invalid number of image columns in image file %s (%d not %d)\n", path, header.number_of_columns, MNIST_IMAGE_HEIGHT);
    }

    *number_of_images = header.number_of_images;
    images = (mnist_image_t*)malloc(*number_of_images * sizeof(mnist_image_t));

    if (images == NULL) {
        fprintf(stderr, "Could not allocated memory for %d images\n", *number_of_images);
        fclose(stream);
        return NULL;
    }

    if (*number_of_images != fread(images, sizeof(mnist_image_t), *number_of_images, stream)) {
        fprintf(stderr, "Could not read %d images from: %s\n", *number_of_images, path);
        free(images);
        fclose(stream);
        return NULL;
    }

    fclose(stream);

    return images;
}

/**
 * Free all the memory allocated in a dataset. This should not be used on a
 * batched dataset as the memory is allocated to the parent.
 */
void mnist_free_dataset(mnist_dataset_t * dataset)
{
    free(dataset->images);
    free(dataset->labels);
    free(dataset);
}

mnist_dataset_t * mnist_get_dataset(const char * image_path, const char * label_path)
{
    mnist_dataset_t * dataset;
    uint32_t number_of_images, number_of_labels;

    dataset = (mnist_dataset_t*)calloc(1, sizeof(mnist_dataset_t));

    if (NULL == dataset) {
        return NULL;
    }

    dataset->images = get_images(image_path, &number_of_images);

    if (NULL == dataset->images) {
        mnist_free_dataset(dataset);
        return NULL;
    }

    dataset->labels = get_labels(label_path, &number_of_labels);

    if (NULL == dataset->labels) {
        mnist_free_dataset(dataset);
        return NULL;
    }

    if (number_of_images != number_of_labels) {
        fprintf(stderr, "Number of images does not match number of labels (%d != %d)\n", number_of_images, number_of_labels);
        mnist_free_dataset(dataset);
        return NULL;
    }

    dataset->size = number_of_images;

    return dataset;
}

/**
 * Fills the batch dataset with a subset of the parent dataset.
 */
int mnist_batch(mnist_dataset_t * dataset, mnist_dataset_t * batch, int size, int number)
{
    int start_offset;

    start_offset = size * number;

    if (start_offset >= dataset->size) {
        return 0;
    }

    batch->images = &dataset->images[start_offset];
    batch->labels = &dataset->labels[start_offset];
    batch->size = size;

    if (start_offset + batch->size > dataset->size) {
        batch->size = dataset->size - start_offset;
    }

    return 1;
}

#define STEPS 1000
#define BATCH_SIZE 100

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

// Convert a pixel value from 0-255 to one from 0 to 1
#define PIXEL_SCALE(x) (((float) (x)) / 255.0f)

// Returns a random value between 0 and 1
#define RAND_FLOAT() (((float) rand()) / ((float) RAND_MAX))

/**
 * Initialise the weights and bias vectors with values between 0 and 1
 */
void neural_network_random_weights(neural_network_t * network)
{
    int i, j;

    for (i = 0; i < MNIST_LABELS; i++) {
        network->b[i] = RAND_FLOAT();

        for (j = 0; j < MNIST_IMAGE_SIZE; j++) {
            network->W[i][j] = RAND_FLOAT();
        }
    }
}

/**
 * Calculate the softmax vector from the activations. This uses a more
 * numerically stable algorithm that normalises the activations to prevent
 * large exponents.
 */
void neural_network_softmax(float * activations, int length)
{
    int i;
    float sum, max;

    for (i = 1, max = activations[0]; i < length; i++) {
        if (activations[i] > max) {
            max = activations[i];
        }
    }

    for (i = 0, sum = 0; i < length; i++) {
        activations[i] = exp(activations[i] - max);
        sum += activations[i];
    }

    for (i = 0; i < length; i++) {
        activations[i] /= sum;
    }
}

/**
 * Calculate the softmax vector from the activations. This uses a more
 * numerically stable algorithm that normalises the activations to prevent
 * large exponents.
 */
void aneural_network_softmax(adept::FixedArray<float,true,MNIST_LABELS> &activations, int length)
{
    int i;
    aReal sum, max;

    for (i = 1, max = activations[0]; i < length; i++) {
        if (activations[i] > max) {
            max = activations[i];
        }
    }

    for (i = 0, sum = 0; i < length; i++) {
        activations[i] = exp(activations[i] - max);
        sum += activations[i];
    }

    activations /= sum;
    for (i = 0; i < length; i++) {
        activations[i] /= sum;
    }
}
static double maxval(const float *activations, int length) {
    float max = activations[0];

    for (int i = 1; i < length; i++) {
        if (activations[i] > max) {
            max = activations[i];
        }
    }
    return max;
}

static double sumval(const float *activations, int length) {
    float sum = 0;

    for (int i = 0; i < length; i++) {
        sum += activations[i];
    }
    return sum;
}

static    void makeexps(float* exps, const float* activations, int length, double max) {
    for (int i = 0; i < length; i++) {
        exps[i] = exp(activations[i] - max);
    }

}
/**
 * Calculate the softmax vector from the activations. This uses a more
 * numerically stable algorithm that normalises the activations to prevent
 * large exponents.
 */
static void neural_network_softmax_v2(const float * activations, float* outp, int length)
{
    int i;
    float sum, max;

    for (i = 1, max = activations[0]; i < length; i++) {
        if (activations[i] > max) {
            max = activations[i];
        }
    }

    for (i = 0, sum = 0; i < length; i++) {
        sum += exp(activations[i] - max);
    }

    for (i = 0; i < length; i++) {
        outp[i] = exp(activations[i] - max) / sum;
    }
#if 0
    float max = maxval(activations, length);
    float exps[length];
    makeexps(exps,activations, length, max);
    float sum = sumval(exps,length);
    /*
    for (int i = 0; i < length; i++) {
        double tmp = exps[i];//exp(activations[i] - max);
        sum += tmp;
    }
    */

    for (int i = 0; i < length; i++) {
        outp[i] = exps[i] / sum;
    }
#endif
}
/**
 * Use the weights and bias vector to forward propogate through the neural
 * network and calculate the activations.
 */
void neural_network_hypothesis(const mnist_image_t * image, const neural_network_t * network, float activations[MNIST_LABELS])
{
    int i, j;

    for (i = 0; i < MNIST_LABELS; i++) {
        activations[i] = network->b[i];

        for (j = 0; j < MNIST_IMAGE_SIZE; j++) {
            activations[i] += network->W[i][j] * PIXEL_SCALE(image->pixels[j]);
        }
    }

    neural_network_softmax(activations, MNIST_LABELS);
}

/**
 * Use the weights and bias vector to forward propogate through the neural
 * network and calculate the activations.
 */
static float neural_network_hypothesis_v2(const mnist_image_t * image, const neural_network_t * network, uint8_t label)
{
    float activations[MNIST_LABELS] = {0};
    int i, j;

    for (i = 0; i < MNIST_LABELS; i++) {
        activations[i] = network->b[i];

        for (j = 0; j < MNIST_IMAGE_SIZE; j++) {
            activations[i] += network->W[i][j] * PIXEL_SCALE(image->pixels[j]);
        }
    }

    float activations2[MNIST_LABELS] = { 0 };
    neural_network_softmax_v2(activations, activations2, MNIST_LABELS);
    return -log(activations2[label]);
}

static aReal neural_network_hypothesis_adept(const mnist_image_t * image, const aneural_network_t * network, uint8_t label)
{
    adept::FixedArray<float,true,MNIST_LABELS> activations = network->b;
    int i, j;

    for (i = 0; i < MNIST_LABELS; i++) {
        for (j = 0; j < MNIST_IMAGE_SIZE; j++) {
            activations(i) += network->W(i,j) * PIXEL_SCALE(image->pixels[j]);
        }
    }

    aneural_network_softmax(activations, MNIST_LABELS);
    return -log(activations[label]);
}

static void calculateDerivatives_adept(mnist_image_t * image, bool run, adept::Stack& stack, const aneural_network_t * anetwork, neural_network_t* gradient, uint8_t label) {

    if (!run) {
        stack.new_recording();
    //    run = true;
    } else
        stack.continue_recording();
    auto resa = neural_network_hypothesis_adept(image, anetwork, label);
    resa.set_gradient(1.0);
    stack.reverse();
    stack.pause_recording();

    for (int i = 0; i < MNIST_LABELS; i++) {
        gradient->b[i] = anetwork->b(i).get_gradient();
        for (int j = 0; j < MNIST_IMAGE_SIZE; j++) {
            gradient->W[i][j] = anetwork->W(i,j).get_gradient();
        }
    }
}

extern "C" {
#include <adBuffer.h>
}

void neural_network_softmax_b(float *activations, float *activationsb, int
        length) {
    float sum, max;
    float sumb, maxb;
    int branch;
    max = activations[0];
    for (int i = 1; i < length; ++i)
        if (activations[i] > max) {
            max = activations[i];
            pushControl1b(1);
        } else
            pushControl1b(0);
    sum = 0;
    for (int i = 0; i < length; ++i) {
        pushReal4(activations[i]);
        activations[i] = (float)exp(activations[i] - max);
        sum = sum + activations[i];
    }
    for (int i = 0; i < length; ++i) {
        pushReal4(activations[i]);
        activations[i] = activations[i]/sum;
    }
    sumb = 0.0;
    for (int i = length-1; i > -1; --i) {
        popReal4(&(activations[i]));
        sumb = sumb - activations[i]*activationsb[i]/(sum*sum);
        activationsb[i] = activationsb[i]/sum;
    }
    {
      float tempb;
      maxb = 0.0;
      for (int i = length-1; i > -1; --i) {
          activationsb[i] = activationsb[i] + sumb;
          popReal4(&(activations[i]));
          tempb = exp(activations[i]-max)*activationsb[i];
          maxb = maxb - tempb;
          activationsb[i] = tempb;
      }
    }
    for (int i = length-1; i > 0; --i) {
        popControl1b(&branch);
        if (branch != 0) {
            activationsb[i] = activationsb[i] + maxb;
            maxb = 0.0;
        }
    }
    activationsb[0] = activationsb[0] + maxb;
}

/**
 * Calculate the softmax vector from the activations. This uses a more
 * numerically stable algorithm that normalises the activations to prevent
 * large exponents.
 */
// Convert a pixel value from 0-255 to one from 0 to 1
// Returns a random value between 0 and 1
void neural_network_softmax_c(float *activations, int length) {
    float sum, max;
    max = activations[0];
	int i;
    for (i = 1; i < length; ++i)
        if (activations[i] > max)
            max = activations[i];
    sum = 0;
    for (i = 0; i < length; ++i) {
        activations[i] = (float)exp(activations[i] - max);
        sum += activations[i];
    }
    for (i = 0; i < length; ++i)
        activations[i] /= sum;
}

/*
  Differentiation of neural_network_hypothesis_tapenadesource in reverse (adjoint) mode:
   gradient     of useful results: neural_network_hypothesis_tapenadesource
                *network.b[0:10-1] *network.W[0:10-1][0:28*28-1]
   with respect to varying inputs: *network.b[0:10-1] *network.W[0:10-1][0:28*28-1]
   RW status of diff variables: neural_network_hypothesis_tapenadesource:in-killed
                *network.b[0:10-1]:incr *network.W[0:10-1][0:28*28-1]:incr
   Plus diff mem management of: network:in *network.b:in *network.W:in
                *network.W[0:10-1]:in
*/
static void neural_network_hypothesis_tapenadesource_b(const mnist_image_t *
        image, const neural_network_t *network, neural_network_t *networkb,
        uint8_t label, float neural_network_hypothesis_tapenadesourceb) {
    float activations[10];
    float activationsb[10];
    int ii1;
    float neural_network_hypothesis_tapenadesource;
    for (int i = 0; i < 10; ++i) {
        activations[i] = network->b[i];
        for (int j = 0; j < 784; ++j)
            activations[i] = activations[i] + network->W[i][j]*((float)image->
                pixels[j]/255.0f);
    }
    pushReal4Array(activations, 10);
    neural_network_softmax_c(activations, 10);
    for (ii1 = 0; ii1 < 10; ++ii1)
        activationsb[ii1] = 0.0;
    activationsb[(int)label] = activationsb[(int)label] -
        neural_network_hypothesis_tapenadesourceb/activations[(int)label];
    popReal4Array(activations, 10);
    neural_network_softmax_b(activations, activationsb, 10);
    for (int i = 9; i > -1; --i) {
        for (int j = 783; j > -1; --j)
            networkb->W[i][j] = networkb->W[i][j] + (float)image->pixels[j]*
                activationsb[i]/255.0f;
        networkb->b[i] = networkb->b[i] + activationsb[i];
        activationsb[i] = 0.0;
    }
}

/**
 * Update the gradients for this step of gradient descent using the gradient
 * contributions from a single training example (image).
 *
 * This function returns the loss ontribution from this training example.
 */
float neural_network_gradient_update(mnist_image_t * image, const neural_network_t * network, neural_network_gradient_t * gradient, uint8_t label)
{
    float activations[MNIST_LABELS];
    float b_grad, W_grad;
    int i, j;

    // First forward propagate through the network to calculate activations
    neural_network_hypothesis(image, network, activations);

    for (i = 0; i < MNIST_LABELS; i++) {
        // This is the gradient for a softmax bias input
        b_grad = (i == label) ? activations[i] - 1 : activations[i];

        for (j = 0; j < MNIST_IMAGE_SIZE; j++) {
            // The gradient for the neuron weight is the bias multiplied by the input weight
            W_grad = b_grad * PIXEL_SCALE(image->pixels[j]);

            // Update the weight gradient
            gradient->W_grad[i][j] += W_grad;
        }

        // Update the bias gradient
        gradient->b_grad[i] += b_grad;
    }

    // Cross entropy loss
    return 0.0f - log(activations[label]);
}

extern int enzyme_const;
template<typename Return, typename... T>
Return __enzyme_autodiff(T...);

static void calculateDerivatives(mnist_image_t * image, const neural_network_t * network, neural_network_t* gradient, uint8_t label) {
    __enzyme_autodiff<void>(neural_network_hypothesis_v2, enzyme_const, image, network, gradient, enzyme_const, label);
}

/**
 * Run one step of gradient descent and update the neural network.
 */
float neural_network_training_step(mnist_dataset_t * dataset, neural_network_t * network, float learning_rate)
{
    neural_network_t gradient = {0};
    neural_network_t gradient2 = {0};

    /*
    adept::Stack stack;
    aneural_network_t anetwork;

    for (int i = 0; i < MNIST_LABELS; i++) {
        anetwork.b[i] = network->b[i];
        for (int j = 0; j < MNIST_IMAGE_SIZE; j++) {
            anetwork.W[i][j] = network->W[i][j];
        }
    }*/

    float total_loss;
    int i, j;

    // Calculate the gradient and the loss by looping through the training set
    for (i = 0, total_loss = 0; i < dataset->size; i++) {
		mnist_image_t* image = &dataset->images[i];
		uint8_t label = dataset->labels[i];

    	// First forward propagate through the network to calculate activations

        //calculateDerivatives(image, network, &gradient, label);
        //calculateDerivatives_adept(image, i != 0, stack, &anetwork, &gradient, label);
		//neural_network_hypothesis_tapenadesource_b(image, network, &gradient, label, 1.0);

        total_loss +=neural_network_gradient_update(image, network, (neural_network_gradient_t*)&gradient, label);

        //total_loss +=neural_network_gradient_update(image, network, (neural_network_gradient_t*)&gradient2, label);

	    //float activations[MNIST_LABELS];
        //neural_network_hypothesis(image, network, activations);
    	//total_loss -= log(activations[label]);

    }

    // Apply gradient descent to the network
    for (i = 0; i < MNIST_LABELS; i++) {
        //printf("b'[i] %f %f\n", gradient.b[i], gradient2.b[i]);
        network->b[i] -= learning_rate * gradient.b[i] / ((float) dataset->size);

        for (j = 0; j < MNIST_IMAGE_SIZE + 1; j++) {
            network->W[i][j] -= learning_rate * gradient.W[i][j] / ((float) dataset->size);
        }
    }

    return total_loss;
}

/**
 * Run one step of gradient descent and update the neural network.
 */
float neural_network_training_step_enzyme(mnist_dataset_t * dataset, neural_network_t * network, float learning_rate)
{
    neural_network_t gradient = {0};
    neural_network_t gradient2 = {0};

    float total_loss;
    int i, j;

    // Calculate the gradient and the loss by looping through the training set
    for (i = 0, total_loss = 0; i < dataset->size; i++) {
		mnist_image_t* image = &dataset->images[i];
		uint8_t label = dataset->labels[i];

    	// First forward propagate through the network to calculate activations

        calculateDerivatives(image, network, &gradient, label);
        //calculateDerivatives_adept(image, i != 0, stack, &anetwork, &gradient, label);
		//neural_network_hypothesis_tapenadesource_b(image, network, &gradient, label, 1.0);

        //total_loss +=neural_network_gradient_update(image, network, (neural_network_gradient_t*)&gradient2, label);

	    float activations[MNIST_LABELS];
        neural_network_hypothesis(image, network, activations);
    	total_loss -= log(activations[label]);

    }

    // Apply gradient descent to the network
    for (i = 0; i < MNIST_LABELS; i++) {
        //printf("b'[i] %f %f\n", gradient.b[i], gradient2.b[i]);
        network->b[i] -= learning_rate * gradient.b[i] / ((float) dataset->size);

        for (j = 0; j < MNIST_IMAGE_SIZE + 1; j++) {
            network->W[i][j] -= learning_rate * gradient.W[i][j] / ((float) dataset->size);
        }
    }

    return total_loss;
}

/**
 * Run one step of gradient descent and update the neural network.
 */
float neural_network_training_step_adept(mnist_dataset_t * dataset, neural_network_t * network, float learning_rate)
{
    neural_network_t gradient = {0};
    neural_network_t gradient2 = {0};

    adept::Stack stack;
    aneural_network_t anetwork;

    for (int i = 0; i < MNIST_LABELS; i++) {
        anetwork.b[i] = network->b[i];
        for (int j = 0; j < MNIST_IMAGE_SIZE; j++) {
            anetwork.W[i][j] = network->W[i][j];
        }
    }

    float total_loss;
    int i, j;

    // Calculate the gradient and the loss by looping through the training set
    for (i = 0, total_loss = 0; i < dataset->size; i++) {
		mnist_image_t* image = &dataset->images[i];
		uint8_t label = dataset->labels[i];

    	// First forward propagate through the network to calculate activations

        calculateDerivatives_adept(image, i != 0, stack, &anetwork, &gradient, label);

        //total_loss +=neural_network_gradient_update(image, network, (neural_network_gradient_t*)&gradient2, label);

	    float activations[MNIST_LABELS];
        neural_network_hypothesis(image, network, activations);
    	total_loss -= log(activations[label]);

    }

    // Apply gradient descent to the network
    for (i = 0; i < MNIST_LABELS; i++) {
        //printf("b'[i] %f %f\n", gradient.b[i], gradient2.b[i]);
        network->b[i] -= learning_rate * gradient.b[i] / ((float) dataset->size);

        for (j = 0; j < MNIST_IMAGE_SIZE + 1; j++) {
            network->W[i][j] -= learning_rate * gradient.W[i][j] / ((float) dataset->size);
        }
    }

    return total_loss;
}

/**
 * Run one step of gradient descent and update the neural network.
 */
float neural_network_training_step_tapenade(mnist_dataset_t * dataset, neural_network_t * network, float learning_rate)
{
    neural_network_t gradient = {0};
    neural_network_t gradient2 = {0};

    float total_loss;
    int i, j;

    // Calculate the gradient and the loss by looping through the training set
    for (i = 0, total_loss = 0; i < dataset->size; i++) {
		mnist_image_t* image = &dataset->images[i];
		uint8_t label = dataset->labels[i];

    	// First forward propagate through the network to calculate activations

		neural_network_hypothesis_tapenadesource_b(image, network, &gradient, label, 1.0);
        //total_loss +=neural_network_gradient_update(image, network, (neural_network_gradient_t*)&gradient2, label);

	    float activations[MNIST_LABELS];
        neural_network_hypothesis(image, network, activations);
    	total_loss -= log(activations[label]);
    }

    // Apply gradient descent to the network
    for (i = 0; i < MNIST_LABELS; i++) {
        //printf("b'[i] %f %f\n", gradient.b[i], gradient2.b[i]);
        network->b[i] -= learning_rate * gradient.b[i] / ((float) dataset->size);

        for (j = 0; j < MNIST_IMAGE_SIZE + 1; j++) {
            network->W[i][j] -= learning_rate * gradient.W[i][j] / ((float) dataset->size);
        }
    }

    return total_loss;
}

/**
 * Downloaded from: http://yann.lecun.com/exdb/mnist/
 */
const char * train_images_file = "data/train-images-idx3-ubyte";
const char * train_labels_file = "data/train-labels-idx1-ubyte";
const char * test_images_file = "data/t10k-images-idx3-ubyte";
const char * test_labels_file = "data/t10k-labels-idx1-ubyte";

/**
 * Calculate the accuracy of the predictions of a neural network on a dataset.
 */
float calculate_accuracy(mnist_dataset_t * dataset, neural_network_t * network) {
    float activations[MNIST_LABELS], max_activation;
    int i, j, correct, predict;

    // Loop through the dataset
    for (i = 0, correct = 0; i < dataset->size; i++) {
        // Calculate the activations for each image using the neural network
        neural_network_hypothesis(&dataset->images[i], network, activations);

        // Set predict to the index of the greatest activation
        for (j = 0, predict = 0, max_activation = activations[0]; j < MNIST_LABELS; j++) {
            if (max_activation < activations[j]) {
                max_activation = activations[j];
                predict = j;
            }
        }

        // Increment the correct count if we predicted the right label
        if (predict == dataset->labels[i]) {
            correct++;
        }
    }

    // Return the percentage we predicted correctly as the accuracy
    return ((float) correct) / ((float) dataset->size);
}

#include <sys/time.h>
#include <stdlib.h>
#include <math.h>
#include <inttypes.h>
#include <string.h>

float tdiff(struct timeval *start, struct timeval *end) {
  return (end->tv_sec-start->tv_sec) + 1e-6*(end->tv_usec-start->tv_usec);
}

void run(float (*fn)(mnist_dataset_t*, neural_network_t*, float)) {
    mnist_dataset_t * train_dataset, * test_dataset;
    mnist_dataset_t batch;
    neural_network_t network;
    float loss, accuracy;
    int i, batches;

    // Read the datasets from the files
    train_dataset = mnist_get_dataset(train_images_file, train_labels_file);
    test_dataset = mnist_get_dataset(test_images_file, test_labels_file);

    // Initialise weights and biases with random values
    neural_network_random_weights(&network);

    // Calculate how many batches (so we know when to wrap around)
    batches = train_dataset->size / BATCH_SIZE;

    struct timeval start, end;
    gettimeofday(&start, NULL);

    for (i = 0; i < STEPS; i++) {
        // Initialise a new batch
        mnist_batch(train_dataset, &batch, 100, i % batches);

        // Run one step of gradient descent and calculate the loss
        loss = fn(&batch, &network, 0.5);

        // Calculate the accuracy using the whole test dataset
        accuracy = calculate_accuracy(test_dataset, &network);

        printf("Step %04d\tAverage Loss: %.2f\tAccuracy: %.3f\n", i, loss / batch.size, accuracy);
    }

    gettimeofday(&end, NULL);
    printf("%0.6f\n", tdiff(&start, &end));

    // Cleanup
    mnist_free_dataset(train_dataset);
    mnist_free_dataset(test_dataset);

}

int main(int argc, char *argv[])
{
    printf("Regular\n");
    run(neural_network_training_step);
    printf("Enzyme\n");
    run(neural_network_training_step_enzyme);
    printf("Adept\n");
    run(neural_network_training_step_adept);
    printf("Tapenade\n");
    run(neural_network_training_step_tapenade);
    return 0;
}


/**
 * Computing transforms.
 * \todo should be smarter about dimensions...what if 2d data, what if 3d+time?
 */
#include <stdlib.h>
#include <stdint.h>
#include "src/aabb.h"
#include "nd.h"
#include <Eigen/Core>
#include <Eigen/LU>
using namespace Eigen;
#include <iostream>
using namespace std;

//#define DEBUG

#ifndef restrict
  #define restrict __restrict
#endif

// output transform should be T s.t. rs=T*rd
// so, if r=Ts*rs=Td*rd
//   T = inv(Ts)*Td

/**
 * \param[in,out] out   Must be preallocated with at least (ndim+1)*(ndim+1)
 *                      elements
 */
extern "C"
void compose(float *restrict out,
             aabb_t bbox,
             float sx, float sy, float sz, // scale in aabb units
             float *restrict transform, unsigned ndim)
{
  MatrixXf                         dst2world(ndim+1,ndim+1);
  Map<Matrix<float,Dynamic,Dynamic,RowMajor> > src2world(transform,ndim+1,ndim+1);
  Map<Matrix<float,Dynamic,Dynamic,RowMajor> > T(out,ndim+1,ndim+1);

  int64_t      *o;
  AABBGet(bbox,0,&o,0);
  dst2world.setIdentity();
  dst2world.block<3,1>(0,ndim)<<(float)o[0],(float)o[1],(float)o[2];
  dst2world.block<3,3>(0,0).diagonal()<<sx,sy,sz;

  T=(src2world.inverse()*dst2world).eval();
#ifdef DEBUG
  #define show(e) cout<<#e" is "<<endl<<e<<endl<<endl
  show(dst2world);
  show(src2world);
  show(src2world.inverse());
  show(T);
  #undef show
#endif
}

extern "C"
void box2box(float *restrict out,
             nd_t dst, aabb_t dstbox,
             nd_t src, aabb_t srcbox)
{ size_t d;
  int64_t *ori,*shape;

  AABBGet(srcbox,0,&ori,&shape);
  d=ndndim(src);
  MatrixXf src2world(d+1,d+1);
  src2world.setIdentity().block<3,3>(0,0).diagonal()
    << (float)shape[0]/(float)ndshape(src)[0],
       (float)shape[1]/(float)ndshape(src)[1],
       (float)shape[2]/(float)ndshape(src)[2];
  src2world.block<3,1>(0,d)
    << (float)ori[0],(float)ori[1],(float)ori[2];

  AABBGet(dstbox,0,&ori,&shape);
  d=ndndim(dst);
  MatrixXf dst2world(d+1,d+1);
  dst2world.setIdentity().block<3,3>(0,0).diagonal()
    << (float)shape[0]/(float)ndshape(dst)[0],
       (float)shape[1]/(float)ndshape(dst)[1],
       (float)shape[2]/(float)ndshape(dst)[2];
  dst2world.block<3,1>(0,d)
    << (float)ori[0],(float)ori[1],(float)ori[2];

  Map<Matrix<float,Dynamic,Dynamic,RowMajor> > T(out,ndndim(src)+1,ndndim(dst)+1);
  T=src2world.inverse()*dst2world;
#ifdef DEBUG
  #define show(e) cout<<#e" is "<<endl<<e<<endl<<endl
  show(dst2world);
  show(src2world);
  show(src2world.inverse());
  show(T);
  #undef show
#endif
  //  [s,r_s] x [r_d,d]; r_s == r_d
}


//==============================================================================
//         Copyright 2003 - 2012   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2012   LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef NT2_CORE_FUNCTIONS_GLOBALNORM_HPP_INCLUDED
#define NT2_CORE_FUNCTIONS_GLOBALNORM_HPP_INCLUDED

#include <nt2/include/functor.hpp>
#include <nt2/sdk/meta/as_real.hpp>
#include <boost/mpl/int.hpp>

namespace nt2 { namespace tag
  {
    /*!
      @brief globalnorm generic tag

      Represents the globalnorm function in generic contexts.

      @par Models:
      Hierarchy
    **/
    struct globalnorm_ : ext::abstract_<globalnorm_>
    {
      /// @brief Parent hierarchy
      typedef ext::abstract_<globalnorm_> parent;
      template<class... Args>
      static BOOST_FORCEINLINE BOOST_AUTO_DECLTYPE dispatch(Args&&... args)
      BOOST_AUTO_DECLTYPE_BODY( dispatching_globalnorm_( ext::adl_helper(), static_cast<Args&&>(args)... ) )
    };
  }
  namespace ext
  {
    template<class Site>
    BOOST_FORCEINLINE generic_dispatcher<tag::globalnorm_, Site> dispatching_globalnorm_(adl_helper, boost::dispatch::meta::unknown_<Site>, ...)
    {
      return generic_dispatcher<tag::globalnorm_, Site>();
    }
    template<class... Args>
    struct impl_globalnorm_;
  }

  /*!
    @brief Global norm

    Computes the norm of a whole table expression with static or dynamic choice
    of the norm computation formula.

    Call protocols to globalnorm are summarized in the following table. We advise to use static
    calls whenever possible as it prevents cascaded run-time if clauses and goes directly
    to the right call at execution.

    @code
    |--------------------|-------------------|------------------------------|-------------------|
    | mnorm(a0, p)                                                                              |
    |--------------------|-------------------|------------------------------|-------------------|
    |    static p        |  dynamic p        |     formula (pseudo-code)    |  equivalent to    |
    |--------------------|-------------------|------------------------------|-------------------|
    | nt2::one_          | 1                 |       sum(abs(x(_)))         | globalnorm1(x)    |
    | nt2::two_          | 2                 |   sqrt(sum(sqr(abs(x(_)))))  | globalnorm2(x)    |
    |    -               | p (positive)      |    sum(abs(x)^p)^(1/p)       | globalnormp(x,p)  |
    | nt2::inf_          | nt2::Inf<T>()     |       max(abs((x))           | globalnorminf(x)  |
    | nt2::fro_          | -1                |   sqrt(sum(sqr(abs(x(_)))))  | globalnormfro(x)  |
    |--------------------|-------------------|------------------------------|-------------------|
    | mnorm<p>(a0)                                                                              |
    |--------------------|-------------------|------------------------------|-------------------|
    |    static p        |                   |     matrix                   |                   |
    |--------------------|-------------------|------------------------------|-------------------|
    | nt2::tag::one_ or 1|        -          |       sum(abs(x(_)))         |  globalnorm1(x)   |
    | nt2::tag::two_ or 2|        -          |   sqrt(sum(sqr(abs(x(_)))))  |  globalnorm2(x)   |
    | p  (integer only)  |        -          |    sum(abs(x)^p)^(1/p)       |  globalnormp(x,p) |
    | nt2::tag::inf_     |        -          |       max(abs((x))           |  globalnorminf(x) |
    | nt2::tag::fro_     |        -          |   sqrt(sum(sqr(abs(x(_)))))  |  globalnormfro(x) |
    |--------------------|-------------------|------------------------------|-------------------|
    @endcode

    @par Semantic:

    For any expression @c a0 of type @c A0, the following call:

    @code
    as_real<A0::value_type>::type x = globalnorm(a0);
    @endcode

    is equivalent to:

    @code
    as_real<A0::value_type>::type x = globalnorm2(a0);
    @endcode

    For any expression @c a0 of type @c A0 and any floating point value @c p, the
    following call:

    @code
    as_real<A0::value_type>::type x = globalnorm(a0,p);
    @endcode

    is equivalent to:

    @code
    as_real<A0::value_type>::type x = globalnormp(a0,p);
    @endcode

    if @c p is finite and to :

    @code
    as_real<A0::value_type>::type x = globalmax(abs(a0));
    @endcode

    if @c is +Inf and to :

    @code
    as_real<A0::value_type>::type x = globalmin(abs(a0));
    @endcode

    if @c p is -Inf.

    @note If 0 < p < 1 or p = -inf, globalnorm does not share the properties that
    define a mathematical norm,  but only a quasi-norm if  0 < p < 1 and a notation
    facility for p = -inf.

    @par Static Interface

    globalnorm can also be invoked with a template parameter which is either a
    functor tag describing the constant value to use instead of @c p or an
    Integral Constant. For example,

    @code
    as_real<A0::value_type>::type x = globalnorm<tag::two_>(a0);
    @endcode

    is equivalent to:

    @code
    as_real<A0::value_type>::type x = globalnorm2(a0);
    @endcode

    Similarly,

    @code
    as_real<A0::value_type>::type x = globalnorm<5>(a0);
    @endcode

    is equivalent to:

    @code
    as_real<A0::value_type>::type x = globalnormp(a0,5);
    @endcode

    @note Whenever a constant functor tag or an Integral Constant is used, compile
    time optimization is performed (if available) so the correct variant of globalnorm is
    called. For example, calls similar to globalnorm<2>(a0) will invoke
    globalnorm2(a0) instead of globalnormp(a0, 2), and globalnorm<5>(a0) will simply
    invoke directly globalnormp(a0, 5) but without any runtime selection.

    @param a0 Expression to compute the norm of
    @param a1 Type of norm to compute
  **/
  BOOST_DISPATCH_FUNCTION_IMPLEMENTATION(nt2::tag::globalnorm_, globalnorm, 2)

  /// @overload
  BOOST_DISPATCH_FUNCTION_IMPLEMENTATION(nt2::tag::globalnorm_, globalnorm, 1)

  /// @overload
  template<typename Tag, typename A0>
  BOOST_FORCEINLINE typename meta::as_real<typename A0::value_type>::type
  globalnorm(const A0& a0)
  {
    return globalnorm(a0, nt2::meta::as_<Tag>());
  }

  /// @overload
  template<int Value, typename A0>
  BOOST_FORCEINLINE typename meta::as_real<typename A0::value_type>::type
  globalnorm(const A0& a0)
  {
    return globalnorm(a0, boost::mpl::int_<Value>() );
  }
}

#endif


// Copyright (C) 2014 The Regents of the University of California (Regents).
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.
//
//     * Neither the name of The Regents or University of California nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
// Please contact the author of this library if you have any questions.
// Author: Chris Sweeney (cmsweeney@cs.ucsb.edu)

#include "theia/sfm/estimators/estimate_triangulation.h"

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <limits>

#include "theia/sfm/camera/camera.h"
#include "theia/sfm/create_and_initialize_ransac_variant.h"
#include "theia/sfm/triangulation/triangulation.h"
#include "theia/sfm/types.h"
#include "theia/solvers/estimator.h"
#include "theia/solvers/ransac.h"
#include "theia/solvers/sample_consensus_estimator.h"
#include "theia/util/util.h"

namespace theia {

namespace {
// The pixel observation and projection matrix needed in order to triangulate a
// 3D point.
struct PointObservation {
  Matrix3x4d projection_matrix;
  Camera camera;
  Eigen::Vector2d normalized_feature;
  Eigen::Vector2d observed_pixel;
};

// Returns true if the point is in front of the camera and false if the point is
// behind the camera.
bool IsPointInFrontOfCamera(const Matrix3x4d& projection_matrix,
                            const Eigen::Vector4d& point) {
  return point.dot(projection_matrix.row(2)) > 0;
}

class TriangulationEstimator
    : public Estimator<PointObservation, Eigen::Vector4d> {
 public:
  TriangulationEstimator() {}

  double SampleSize() const { return 2; }

  // Triangulates the 3D point from 2 observations.
  bool EstimateModel(const std::vector<PointObservation>& observations,
                     std::vector<Eigen::Vector4d>* triangulated_points) const {
    triangulated_points->resize(1);
    if (!Triangulate(observations[0].projection_matrix,
                     observations[1].projection_matrix,
                     observations[0].normalized_feature,
                     observations[1].normalized_feature,
                     &triangulated_points->at(0))) {
      return false;
    }

    // Only return true if the point is in front of both cameras and the
    // triangulation was a success.
    return IsPointInFrontOfCamera(observations[0].projection_matrix,
                                  triangulated_points->at(0)) &&
           IsPointInFrontOfCamera(observations[1].projection_matrix,
                                  triangulated_points->at(0));
  }

  double Error(const PointObservation& observation,
               const Eigen::Vector4d& triangulated_point) const {
    Eigen::Vector2d reprojection;
    const double depth =
        observation.camera.ProjectPoint(triangulated_point, &reprojection);
    if (depth <= 0) {
      return std::numeric_limits<double>::max();
    }
    return (observation.observed_pixel - reprojection).squaredNorm();
  }
};

}  // namespace

bool EstimateTriangulation(const RansacParameters& ransac_params,
                           const std::vector<Camera>& cameras,
                           const std::vector<theia::Feature>& features,
                           Eigen::Vector4d* triangulated_point,
                           RansacSummary* summary) {
  CHECK_EQ(cameras.size(), features.size());
  CHECK_NOTNULL(triangulated_point);

  // If we only have a few data points, then we should exhaustively search the
  // solution space for the best combination.
  static const int kMaxNumDataPointsForExhaustiveSearch = 15;
  if (cameras.size() < 2) {
    return false;
  }

  // Create point correspondences.
  std::vector<PointObservation> point_observations(cameras.size());
  for (int i = 0; i < point_observations.size(); i++) {
    // Create the projection atirx.
    Matrix3x4d projection_matrix;
    projection_matrix.leftCols<3>() =
        cameras[i].GetOrientationAsRotationMatrix();
    projection_matrix.rightCols<1>() =
        -projection_matrix.leftCols<3>() * cameras[i].GetPosition();

    point_observations[i].projection_matrix = projection_matrix;
    point_observations[i].camera = cameras[i];
    point_observations[i].normalized_feature =
        cameras[i].PixelToNormalizedCoordinates(features[i].point_).hnormalized();
    point_observations[i].observed_pixel = features[i].point_;
  }

  // RANSAC triangulation.
  TriangulationEstimator triangulation_estimator;
  std::unique_ptr<SampleConsensusEstimator<TriangulationEstimator> > ransac;
  if (cameras.size() <= kMaxNumDataPointsForExhaustiveSearch) {
    // Set the minimum number of iterations to be the exact number of possible
    // combinations. This forces all combinations to be tested.
    const int num_combinations = cameras.size() * (cameras.size() - 1) / 2;
    RansacParameters exhaustive_params = ransac_params;
    exhaustive_params.min_iterations = num_combinations;
    exhaustive_params.max_iterations = num_combinations;
    ransac = CreateAndInitializeRansacVariant<TriangulationEstimator>(
        RansacType::EXHAUSTIVE, exhaustive_params, triangulation_estimator);
  } else {
    ransac = CreateAndInitializeRansacVariant<TriangulationEstimator>(
        RansacType::RANSAC, ransac_params, triangulation_estimator);
  }

  // Run the RANSAC scheme.
  CHECK(ransac->Initialize());
  return ransac->Estimate(point_observations, triangulated_point, summary);
}

}  // namespace theia


//***************************************************************************
//* Copyright (c) 2015 Saint Petersburg State University
//* Copyright (c) 2011-2014 Saint Petersburg Academic University
//* All Rights Reserved
//* See file LICENSE for details.
//***************************************************************************

#include "kmer_coverage_model.hpp"

#include "utils/logger/logger.hpp"
#include "utils/verify.hpp"
#include "math/xmath.h"
#include "math/smooth.hpp"

#include <boost/math/special_functions/zeta.hpp>
#include <boost/math/distributions/normal.hpp>
#include <boost/math/distributions/skew_normal.hpp>
#include <boost/math/distributions/geometric.hpp>
#include <boost/math/distributions/pareto.hpp>

#include <nlopt/nlopt.hpp>

#include <vector>

#include <cstring>
#include <cstdint>
#include <cstddef>
#include <cmath>

namespace coverage_model {

using std::isfinite;

static const size_t MaxCopy = 10;

static double dzeta(double x, double p) {
    return pow(x, -p - 1) / boost::math::zeta(p + 1);
}

static double perr(size_t i, double scale, double shape) {
    return pow((1 + shape * ((double) (i - 1)) / scale), -1.0 / shape) -
           pow((1 + shape * ((double) i) / scale), -1.0 / shape);
}

static double pgood(size_t i, double zp, double u, double sd, double shape,
                    double* mixprobs = NULL) {
    double res = 0;

    for (unsigned copy = 0; copy < MaxCopy; ++copy) {
        boost::math::skew_normal snormal((copy + 1) * u, sd * sqrt(copy + 1), shape);
        // res += (mixprobs ? mixprobs[copy] : dzeta(copy + 1, zp)) * (boost::math::cdf(snormal, i + 1) - boost::math::cdf(snormal, i));
        res += (mixprobs ? mixprobs[copy] : dzeta(copy + 1, zp)) * boost::math::pdf(snormal, i);
    }

    return res;
}

class CovModelLogLike {
    const std::vector<size_t>& cov;

public:
    CovModelLogLike(const std::vector<size_t>& cov)
            : cov(cov) {}

    int getN() const { return 7; };

private:

    double eval_(const double* x) const {
        double zp = x[0], p = x[1], shape = x[2], u = x[3], sd = x[4], scale = x[5], shape2 = x[6];

        if (zp <= 1 || shape <= 0 || sd <= 0 || p < 1e-9 || p > 1 - 1e-9 || u <= 0 || scale <= 0 ||
            !isfinite(zp) || !isfinite(shape) || !isfinite(sd) || !isfinite(p) || !isfinite(u) ||
            !isfinite(scale) || !isfinite(shape2))
            return +std::numeric_limits<double>::infinity();

        std::vector<double> kmer_probs(cov.size());

        // Error
        for (size_t i = 0; i < kmer_probs.size(); ++i)
            kmer_probs[i] += p * perr(i + 1, scale, shape);

        // Good
        for (size_t i = 0; i < kmer_probs.size(); ++i)
            kmer_probs[i] += (1 - p) * pgood(i + 1, zp, u, sd, shape2);

        double res = 0;
        for (size_t i = 0; i < kmer_probs.size(); ++i)
            res += (double) (cov[i]) * log(kmer_probs[i]);

        return -res;
    }
};

struct CovModelLogLikeEMData {
    const std::vector<size_t>& cov;
    const std::vector<double>& z;
};

static double CovModelLogLikeEM(unsigned, const double* x, double*, void* data) {
    double zp = x[0], shape = x[1], u = x[2], sd = x[3], scale = x[4], shape2 = x[5];

    // INFO("Entry: " << x[0] << " " << x[1] << " " << x[2] << " " << x[3] << " " << x[4]);

    if (zp <= 1 || shape <= 0 || sd <= 0 || u <= 0 || scale <= 0 ||
        !isfinite(zp) || !isfinite(shape) || !isfinite(sd) || !isfinite(u) ||
        !isfinite(scale) || !isfinite(shape2))
        return -std::numeric_limits<double>::infinity();

    const std::vector<size_t>& cov = static_cast<CovModelLogLikeEMData*>(data)->cov;
    const std::vector<double>& z = static_cast<CovModelLogLikeEMData*>(data)->z;

    std::vector<double> kmer_probs(cov.size(), 0);

    // Error
    for (size_t i = 0; i < kmer_probs.size(); ++i) {
        if (cov[i] == 0)
            continue;

        kmer_probs[i] += z[i] * log(perr(i + 1, scale, shape));
    }

    // Good
    // Pre-compute mixing probabilities
    std::vector<double> mixprobs(MaxCopy, 0);
    for (unsigned copy = 0; copy < MaxCopy; ++copy)
        mixprobs[copy] = dzeta(copy + 1, zp);

    // Compute the density
    for (size_t i = 0; i < kmer_probs.size(); ++i) {
        if (cov[i] == 0)
            continue;

        double val = log(pgood(i + 1, zp, u, sd, shape2, &mixprobs[0]));
        if (!isfinite(val))
            val = -1000.0;
        kmer_probs[i] += (1 - z[i]) * val;
    }

    double res = 0;
    for (size_t i = 0; i < kmer_probs.size(); ++i)
        res += (double) (cov[i]) * kmer_probs[i];

    // INFO("f: " << res);
    return res;
}

static std::vector<double> EStep(const std::vector<double>& x,
                                 double p, size_t N) {
    double zp = x[0], shape = x[1], u = x[2], sd = x[3], scale = x[4], shape2 = x[5];

    std::vector<double> res(N);
    for (size_t i = 0; i < N; ++i) {
        double pe = p * perr(i + 1, scale, shape);
        res[i] = pe / (pe + (1 - p) * pgood(i + 1, zp, u, sd, shape2));
        if (!isfinite(res[i]))
            res[i] = 1.0;
    }

    return res;
}

// Estimate the coverage mean by finding the max past the
// first valley.
size_t KMerCoverageModel::EstimateValley() const {
    // Smooth the histogram
    std::vector<size_t> scov;
    math::Smooth3RS3R(scov, cov_);

    size_t Valley = scov[0];

    // Start finding the valley
    size_t Idx = 1;
    while (scov[Idx] < Valley && Idx < scov.size()) {
        Valley = scov[Idx];
        Idx += 1;
    }
    Idx -= 1;

    INFO("Kmer coverage valley at: " << Idx);

    return Idx;
}

void KMerCoverageModel::Fit() {
    CHECK_FATAL_ERROR(cov_.size() > 10, "Invalid kmer coverage histogram, make sure that the coverage is indeed uniform");

    // Find the minimal coverage point using smoothed histogram.
    Valley_ = EstimateValley();

    // First estimate of coverage is the first maximum after the valley.
    MaxCov_ = Valley_ + 1;
    size_t MaxHist = cov_[MaxCov_];
    for (size_t i = Valley_ + 1; i < cov_.size(); ++i) {
        if (cov_[i] > MaxHist) {
            MaxHist = cov_[i];
            MaxCov_ = i;
        }
    }
    INFO("K-mer histogram maximum: " << MaxCov_);

    // Refine the estimate via median
    size_t AfterValley = 0, SecondValley = std::min(2 * MaxCov_ - Valley_, cov_.size());
    for (size_t i = Valley_ + 1; i < SecondValley; ++i)
        AfterValley += cov_[i];

    size_t ccov = 0;
    for (size_t i = Valley_ + 1; i < SecondValley; ++i) {
        if (ccov > AfterValley / 2) {
            MaxCov_ = std::max(i, MaxCov_);
            break;
        }
        ccov += cov_[i];
    }

    if (MaxCov_ - Valley_ < 3)
        WARN("Too many erroneous kmers, the estimates might be unreliable");

    std::vector<size_t> mvals(1 + MaxCov_ - Valley_);
    mvals[0] = cov_[MaxCov_];
    size_t tmadcov = mvals[0];
    for (size_t i = 1; i < std::min(MaxCov_ - Valley_, cov_.size() - MaxCov_); ++i) {
        mvals[i] = cov_[MaxCov_ + i] + cov_[MaxCov_ - i];
        tmadcov += mvals[i];
    }
    size_t madcov = 0;
    double CovSd = sqrt((double) (5 * MaxCov_));
    for (size_t i = 0; i < MaxCov_ - Valley_; ++i) {
        if (madcov > tmadcov / 2) {
            CovSd = (double) i;
            break;
        }
        madcov += mvals[i];
    }
    CovSd *= 1.4826;
    INFO("Estimated median coverage: " << MaxCov_ << ". Coverage mad: " << CovSd);

    // Estimate error probability as ratio of kmers before the valley.
    size_t BeforeValley = 0, Total = 0;
    double ErrorProb = 0;
    for (size_t i = 0; i < cov_.size(); ++i) {
        if (i <= Valley_)
            BeforeValley += cov_[i];
        Total += cov_[i];
    }
    ErrorProb = (double) BeforeValley / (double) Total;
    // Allow some erroneous / good kmers.
    ErrorProb = std::min(1 - 1e-3, ErrorProb);
    ErrorProb = std::max(1e-3, ErrorProb);

    TRACE("Total: " << Total << ". Before: " << BeforeValley);
    TRACE("p: " << ErrorProb);

    std::vector<double> x = {3.0, 3.0, (double) MaxCov_, CovSd, 1.0, 0.0},
        lb = {0.0, 0.0, 0.0, (double) (MaxCov_ - Valley_), 0.0, -6.0},
        ub = {2000.0, 2000.0, (double) (2 * MaxCov_), (double) SecondValley, 2000.0, 6.0};

    INFO("Fitting coverage model");
    // Ensure that there will be at least 2 iterations.
    double PrevErrProb = 2;
    const double ErrProbThr = 1e-8;
    auto GoodCov = cov_;
    GoodCov.resize(std::min(cov_.size(), 5 * MaxCopy * MaxCov_ / 4));
    converged_ = true;
    unsigned it = 1;
    while (fabs(PrevErrProb - ErrorProb) > ErrProbThr) {
        // Recalculate the vector of posterior error probabilities
        std::vector<double> z = EStep(x, ErrorProb, GoodCov.size());

        // Recalculate the probability of error
        PrevErrProb = ErrorProb;
        ErrorProb = 0;
        for (size_t i = 0; i < GoodCov.size(); ++i)
            ErrorProb += z[i] * (double) GoodCov[i];
        ErrorProb /= (double) Total;

        bool LastIter = fabs(PrevErrProb - ErrorProb) <= ErrProbThr;

        nlopt::opt opt(nlopt::LN_NELDERMEAD, 6);
        CovModelLogLikeEMData data = {GoodCov, z};
        opt.set_max_objective(CovModelLogLikeEM, &data);
        if (!LastIter)
            opt.set_maxeval(5 * 6 * it);
        opt.set_xtol_rel(1e-8);
        opt.set_ftol_rel(1e-8);

        double fMin;
        nlopt::result Results = nlopt::FAILURE;
        try {
            Results = opt.optimize(x, fMin);
        } catch (nlopt::roundoff_limited&) {
        }

        VERBOSE_POWER_T2(it, 1, "... iteration " << it);
        TRACE("Results: ");
        TRACE("Converged: " << Results << " " << "F: " << fMin);

        double zp = x[0], shape = x[1], u = x[2], sd = x[3], scale = x[4], shape2 = x[5];
        TRACE("zp: " << zp << " p: " << ErrorProb << " shape: " << shape << " u: " << u << " sd: " << sd <<
                     " scale: " << scale << " shape2: " << shape2);

        it += 1;
    }

    double delta = x[5] / sqrt(1 + x[5] * x[5]);
    mean_coverage_ = x[2] + x[3] * delta * sqrt(2 / M_PI);
    sd_coverage_ = x[3] * sqrt(1 - 2 * delta * delta / M_PI);
    INFO("Fitted mean coverage: " << mean_coverage_ << ". Fitted coverage std. dev: " << sd_coverage_);

    // Now let us check whether we have sane results
    for (size_t i = 0; i < x.size(); ++i)
        if (!isfinite(x[i])) {
            converged_ = false;
            break;
        }

    if (!isfinite(ErrorProb))
        converged_ = false;

    // See, if we can deduce proper threshold

    // First, check whether initial estimate of Valley was sane.
    ErrorThreshold_ = 0;
    if (converged_ && Valley_ > x[2] && x[2] > 2) {
        Valley_ = (size_t) math::round(x[2] / 2.0);
        WARN("Valley value was estimated improperly, reset to " << Valley_);
    }

    // If the model converged, then use it to estimate the thresholds.
    if (converged_) {
        std::vector<double> z = EStep(x, ErrorProb, GoodCov.size());

        INFO("Probability of erroneous kmer at valley: " << z[Valley_]);
        converged_ = false;
        for (size_t i = 0; i < z.size(); ++i)
            if (z[i] > strong_probability_threshold_) //0.999
                LowThreshold_ = std::min(i + 1, Valley_);
            else if (z[i] < probability_threshold_) {//0.05?
                ErrorThreshold_ = std::max(i + 1, Valley_);
                converged_ = true;
                break;
            }

#if 0
        for (size_t i = 0; i < z.size(); ++i) {
            double zp = x[0], shape = x[1], u = x[2], sd = x[3], scale = x[4], shape2 = x[5];
            double pe = ErrorProb * perr(i + 1, scale, shape);
            double pg = (1 - ErrorProb) * pgood(i + 1, zp, u, sd, shape2);

            fprintf(stderr, "%e %e %e %e\n", pe, pg, z[i], perr(i + 1, scale, shape));
        }
#endif
    }

    // See, if we have sane ErrorThreshold_ and go down to something convervative, if not.
    if (converged_) {
        INFO("Preliminary threshold calculated as: " << ErrorThreshold_);
        ErrorThreshold_ = (Valley_ < mean_coverage_ ?
                           std::min(Valley_ + (size_t) (mean_coverage_ - (double) Valley_) / 2, ErrorThreshold_) :
                           Valley_);
        INFO("Threshold adjusted to: " << ErrorThreshold_);
    } else {
        ErrorThreshold_ = Valley_;
        LowThreshold_ = 1;
        WARN("Failed to determine erroneous kmer threshold. Threshold set to: " << ErrorThreshold_);
    }

    // Now the bonus: estimate the genome size!
    GenomeSize_ = 0;
    for (size_t i = ErrorThreshold_ - 1; i < GoodCov.size(); ++i)
        GenomeSize_ += GoodCov[i];
    GenomeSize_ /= 2;

    INFO("Estimated genome size (ignoring repeats): " << GenomeSize_);
}

}


#ifndef TVMTL_MATRIX_UTILS_HPP
#define TVMTL_MATRIX_UTILS_HPP

#include <cmath>
#include <complex>
#include <iostream>

#include <Eigen/Dense>
#include <unsupported/Eigen/MatrixFunctions>
#include <unsupported/Eigen/KroneckerProduct>

namespace tvmtl {

template <typename MatrixType>
    void SolveTriangularSylvester(const MatrixType& A, const MatrixType& B, const MatrixType& C, MatrixType& result)
    {
/*	Eigen::eigen_assert(A.rows() == A.cols());
	Eigen::eigen_assert(A.isUpperTriangular());
	Eigen::eigen_assert(B.rows() == B.cols());
	Eigen::eigen_assert(B.isUpperTriangular());
	Eigen::eigen_assert(C.rows() == A.rows());
	Eigen::eigen_assert(C.cols() == B.rows());
  */
      typedef typename MatrixType::Index Index;
      typedef typename MatrixType::Scalar Scalar;

      Index m = A.rows();
      Index n = B.rows();
      MatrixType X(m, n);

      for (Index i = m - 1; i >= 0; --i) {
          for (Index j = 0; j < n; ++j) {
	        // Compute AX = \sum_{k=i+1}^m A_{ik} X_{kj}
		Scalar AX;
		if (i == m - 1) {
		    AX = 0;
		}
		else {
		    Eigen::Matrix<Scalar,1,1> AXmatrix = A.row(i).tail(m-1-i) * X.col(j).tail(m-1-i);
		    AX = AXmatrix(0,0);
		}

		// Compute XB = \sum_{k=1}^{j-1} X_{ik} B_{kj}
		Scalar XB;
		if (j == 0) {
		    XB = 0;
		}
		else {
		    Eigen::Matrix<Scalar,1,1> XBmatrix = X.row(i).head(j) * B.col(j).head(j);
		    XB = XBmatrix(0,0);
		}

		X(i,j) = (C(i,j) - AX - XB) / (A(i,i) + B(j,j));
	}
    }
result = X;
}

template <typename DerivedX, typename DerivedY, typename DerivedZ>
void MatrixRootFrechetDerivative(const Eigen::MatrixBase<DerivedX>& X, const Eigen::MatrixBase<DerivedY>& E, Eigen::MatrixBase<DerivedZ>& result){

    // Matrix Parameters
    typedef Eigen::internal::traits<DerivedX> Traits;
    typedef typename Traits::Scalar Scalar;
    static const int Rows = Traits::RowsAtCompileTime, Cols = Traits::ColsAtCompileTime;
    static const int Options = DerivedX::Options;
    static const int MaxRows = Traits::MaxRowsAtCompileTime, MaxCols = Traits::MaxColsAtCompileTime;

    // Switch to complex arithmetic
    typedef std::complex<Scalar> ComplexScalar;
    typedef Eigen::Matrix<ComplexScalar, Rows, Cols, Options, MaxRows, MaxCols> ComplexMatrix;

    ComplexMatrix CX = X.template cast<ComplexScalar>();
    ComplexMatrix CE = E.template cast<ComplexScalar>();
    ComplexMatrix CResult;

    // Complex Schur Decomposition
    const Eigen::ComplexSchur<ComplexMatrix> SchurOfX(CX);
    ComplexMatrix T = SchurOfX.matrixT();
    ComplexMatrix U = SchurOfX.matrixU();

    ComplexMatrix sqrtT;
    CE = U.adjoint()*CE*U;

    Eigen::MatrixSquareRootTriangular<ComplexMatrix>(T).compute(sqrtT);
    SolveTriangularSylvester(sqrtT, sqrtT, CE, CResult);

    CResult = U * CResult * U.adjoint();

    result = CResult.real();
}

template <typename DerivedX, typename DerivedY, typename DerivedZ>
void MatrixLogarithmFrechetDerivative(const Eigen::MatrixBase<DerivedX>& X, const Eigen::MatrixBase<DerivedY>& E, Eigen::MatrixBase<DerivedZ>& result){

    // Matrix Parameters
    typedef Eigen::internal::traits<DerivedX> Traits;
    typedef typename Traits::Scalar Scalar;
    static const int Rows = Traits::RowsAtCompileTime, Cols = Traits::ColsAtCompileTime;
    static const int Options = DerivedX::Options;
    static const int MaxRows = Traits::MaxRowsAtCompileTime, MaxCols = Traits::MaxColsAtCompileTime;

    // Switch to complex arithmetic
    typedef std::complex<Scalar> ComplexScalar;
    typedef Eigen::Matrix<ComplexScalar, Rows, Cols, Options, MaxRows, MaxCols> ComplexMatrix;

    ComplexMatrix CX = X.template cast<ComplexScalar>();
    ComplexMatrix CE = E.template cast<ComplexScalar>();
    ComplexMatrix CResult;

    // Order of the Pade approximant
    // If this is changed, we also need new weights and nodes
    const int m = 7;

    // Complex Schur Decomposition
    const Eigen::ComplexSchur<ComplexMatrix> SchurOfX(CX);
    ComplexMatrix T = SchurOfX.matrixT();
    ComplexMatrix U = SchurOfX.matrixU();

    //Compute the number of square roots
    int s = 0;
    const int smax = 20;
    const double theta7 = 2.88e-1;
    double rho = theta7 + 1.0;

    Eigen::Matrix<ComplexScalar, Rows, 1> D;
    D = T.diagonal();

    #ifdef TVMTL_MATRIX_UTILS_DEBUG_VERBOSE
	std::cout << "Diagonal of Schur Decomposition: \n" << D << std::endl;
    #endif

    while(rho > theta7 && s < smax){
	D=D.cwiseSqrt();
	#ifdef TVMTL_MATRIX_UTILS_DEBUG_VERBOSE
	    std::cout << "D: " << D << std::endl;
	#endif
	++s;
	rho = (D - Eigen::Matrix<ComplexScalar, Rows, 1>::Constant(1)).cwiseAbs().maxCoeff();
	 #ifdef TVMTL_MATRIX_UTILS_DEBUG_VERBOSE
	    std::cout << "rho: " << rho << std::endl;
	 #endif
    }

    #ifdef TVMTL_MATRIX_UTILS_DEBUG
	std::cout << "Number of square roots  for dlog estimation: " << s << std::endl;
    #endif

    ComplexMatrix sqrtT;
    CE = U.adjoint()*CE*U;

    for(int k=0; k<s; k++){
	Eigen::MatrixSquareRootTriangular<ComplexMatrix>(T).compute(sqrtT);
	T=sqrtT;
	SolveTriangularSylvester(T, T, CE, CE);
    }

    const double nodes[]   = { 0.0254460438286207377369051579760744L, 0.1292344072003027800680676133596058L,
            0.2970774243113014165466967939615193L, 0.5000000000000000000000000000000000L,
            0.7029225756886985834533032060384807L, 0.8707655927996972199319323866403942L,
            0.9745539561713792622630948420239256L };
    const double weights[] = { 0.0647424830844348466353057163395410L, 0.1398526957446383339507338857118898L,
              0.1909150252525594724751848877444876L, 0.2089795918367346938775510204081633L,
              0.1909150252525594724751848877444876L, 0.1398526957446383339507338857118898L,
              0.0647424830844348466353057163395410L };

    ComplexMatrix TminusI = T - ComplexMatrix::Identity(T.rows(), T.rows());
    CResult.setZero();

    for (int i = 0; i < m; ++i){
  	ComplexMatrix IplusBetaTm = ComplexMatrix::Identity(T.rows(), T.rows()) + nodes[i] * TminusI;
  	ComplexMatrix X = IplusBetaTm.template triangularView< Eigen::Upper >().solve(CE).transpose();
  	CResult += weights[i] * (IplusBetaTm.template triangularView< Eigen::Upper >().transpose().solve(X)).transpose();
    }

    CResult = std::pow(2.0,s) * U * CResult * U.adjoint();

    result = CResult.real();
}

template <typename DerivedX, typename DerivedY>
void KroneckerDLog(const Eigen::MatrixBase<DerivedX>& X, Eigen::MatrixBase<DerivedY>& Result){

    typedef Eigen::internal::traits<DerivedX> Traits;
    typedef typename Traits::Scalar Scalar;
    static const int Rows = Traits::RowsAtCompileTime;

    DerivedX E, PartialDiff;
    DerivedY R;

    for (int i = 0; i < Rows; i++) {
    	for (int j = 0; j < Rows; j++) {
	    E = DerivedX::Zero();
	    E(i,j) = 1.0;
	    MatrixLogarithmFrechetDerivative(X, E, PartialDiff);
	    PartialDiff.transposeInPlace();
	    R.row(i*Rows+j) = Eigen::Map<Eigen::VectorXd>(PartialDiff.data(), PartialDiff.size());
    	}
    }
    Result = R;
}

template <typename DerivedX, typename DerivedY>
void KroneckerDSqrt2(const Eigen::MatrixBase<DerivedX>& X, Eigen::MatrixBase<DerivedY>& Result){

    typedef Eigen::internal::traits<DerivedX> Traits;
    typedef typename Traits::Scalar Scalar;
    static const int Rows = Traits::RowsAtCompileTime;

    DerivedX E, PartialDiff;
    DerivedY R;

    for (int i = 0; i < Rows; i++) {
    	for (int j = 0; j < Rows; j++) {
	    E = DerivedX::Zero();
	    E(i,j) = 1.0;
	    MatrixRootFrechetDerivative(X, E, PartialDiff);
	    PartialDiff.transposeInPlace();
	    R.row(i*Rows+j) = Eigen::Map<Eigen::VectorXd>(PartialDiff.data(), PartialDiff.size());
    	}
    }
    Result = R;
}

template <typename DerivedX, typename DerivedY>
void KroneckerDSqrt(const Eigen::MatrixBase<DerivedX>& X, Eigen::MatrixBase<DerivedY>& Result){

    DerivedX Xsqrt = X.sqrt();
    DerivedY R = Eigen::kroneckerProduct(DerivedX::Identity(),Xsqrt) + Eigen::kroneckerProduct(Xsqrt.transpose(), DerivedX::Identity());
    Result = R.inverse();
}

} // end namespace tvmtl

#endif


/**
 * Copyright 2021-2022 Huawei Technologies Co., Ltd
 *
 * 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 "plugin/device/cpu/kernel/lstsq_cpu_kernel.h"
#include <Eigen/Dense>
#include "plugin/device/cpu/hal/device/cpu_device_address.h"
#include "kernel/common_utils.h"

namespace mindspore {
namespace kernel {
namespace {
constexpr size_t kLstsqInputsNum = 2;
constexpr size_t kLstsqOutputsNum = 1;
constexpr size_t kXDimNum = 2;
constexpr size_t kADimNum_1 = 1;
constexpr size_t kADimNum_2 = 2;
}  // namespace

void LstsqCpuKernelMod::InitKernel(const CNodePtr &kernel_node) {
  MS_EXCEPTION_IF_NULL(kernel_node);
  input_0_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, 0);
  input_1_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, 1);
  if (input_0_shape_.size() != kXDimNum) {
    MS_LOG(EXCEPTION) << "The input x tensor's rank must be 2 for 'Lstsq' Op, but x tensor's rank is "
                      << input_0_shape_.size();
  }
  if (input_1_shape_.size() != kADimNum_2 && input_1_shape_.size() != kADimNum_1) {
    MS_LOG(EXCEPTION) << "The input a tensor's rank must be 2 or 1 for 'Lstsq' Op, but a tensor's rank is "
                      << input_1_shape_.size();
  }
  if (input_0_shape_[0] != input_1_shape_[0]) {
    MS_LOG(EXCEPTION) << "The length of x_dim[0]: " << input_0_shape_[0]
                      << " is not equal to the length of a_dims[0]: " << input_1_shape_[0] << ".";
  }
  dtype_0_ = AnfAlgo::GetInputDeviceDataType(kernel_node, 0);
  dtype_1_ = AnfAlgo::GetInputDeviceDataType(kernel_node, 1);
  if (dtype_0_ != dtype_1_) {
    MS_LOG(EXCEPTION) << "For Lstsq input's dtypes are not the same.";
  }
}

bool LstsqCpuKernelMod::Launch(const std::vector<kernel::AddressPtr> &inputs, const std::vector<kernel::AddressPtr> &,
                               const std::vector<kernel::AddressPtr> &outputs) {
  CHECK_KERNEL_INPUTS_NUM(inputs.size(), kLstsqInputsNum, kernel_name_);
  CHECK_KERNEL_OUTPUTS_NUM(outputs.size(), kLstsqOutputsNum, kernel_name_);
  if (dtype_0_ == kNumberTypeFloat16) {
    LaunchKernel<float, float16>(inputs, outputs);
  } else if (dtype_0_ == kNumberTypeFloat32) {
    LaunchKernel<float, float>(inputs, outputs);
  } else if (dtype_0_ == kNumberTypeFloat64) {
    LaunchKernel<double, double>(inputs, outputs);
  } else {
    MS_LOG(EXCEPTION) << "Unsupported input data type.";
  }
  return true;
}

template <typename T1, typename T2>
void LstsqCpuKernelMod::LaunchKernel(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &outputs) {
  auto input_0_addr = reinterpret_cast<T2 *>(inputs[0]->addr);
  auto input_1_addr = reinterpret_cast<T2 *>(inputs[1]->addr);
  auto output_addr = reinterpret_cast<T2 *>(outputs[0]->addr);
  size_t m = input_0_shape_[0];
  size_t n = input_0_shape_[1];
  size_t k = 0;
  if (input_1_shape_.size() == kADimNum_1) {
    k = 1;
  } else {
    k = input_1_shape_[1];
  }

  typedef Eigen::Matrix<T1, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> MartixXd;
  MartixXd A(m, n);
  MartixXd B(m, k);
  for (size_t i = 0; i < m * n; i++) {
    A.data()[i] = static_cast<T1>(input_0_addr[i]);
  }
  for (size_t i = 0; i < m * k; i++) {
    B.data()[i] = static_cast<T1>(input_1_addr[i]);
  }
  MartixXd result;
  if (m >= n) {
    result = A.colPivHouseholderQr().solve(B);
  } else {
    MartixXd A_Transpose = A.transpose();
    MartixXd temp = A * A_Transpose;
    MartixXd tempI = temp.inverse();
    MartixXd x = A_Transpose * tempI;
    MartixXd output = x * B;
    result = output;
  }
  for (size_t i = 0; i < n; i++)
    for (size_t j = 0; j < k; j++) {
      *(output_addr + i * k + j) = static_cast<T2>(result(i, j));
    }
}

MS_KERNEL_FACTORY_REG(NativeCpuKernelMod, Lstsq, LstsqCpuKernelMod);
}  // namespace kernel
}  // namespace mindspore


#include <boost/config/warning_disable.hpp>
#include <boost/spirit/include/qi.hpp>
#include <boost/spirit/include/karma.hpp>
#include <boost/spirit/include/phoenix_core.hpp>
#include <boost/spirit/include/phoenix_operator.hpp>
#include <boost/fusion/include/std_pair.hpp>

#include <iostream>
#include <string>
#include <complex>

namespace client{
  template <typename Iterator>
    bool parse_complex(Iterator first, Iterator last, std::complex<double>& c){
      using boost::spirit::qi::double_;
      using boost::spirit::qi::_1;
      using boost::spirit::qi::phrase_parse;
      using boost::spirit::ascii::space;
      using boost::phoenix::ref;

      double rN = 0.0;
      double iN = 0.0;
      bool r = phrase_parse(first, last,
          (
            '(' >> double_[ref(rN) = _1]
                >> -(',' >> double_[ref(iN) = _1]) >> ')'
            | double_[ref(rN) = _1]
          ),
          space);

      if(!r || first != last)
        return false;
      c = std::complex<double>(rN, iN);
      return r;
    }

  template <typename OutputIterator>
    bool generate_complex(OutputIterator sink, std::complex<double> const& c){
      using boost::spirit::karma::eps;
      using boost::spirit::karma::double_;
      using boost::spirit::karma::_1;
      using boost::spirit::karma::generate;

      return generate(sink,
          (
            eps(c.imag() != 0) << '(' << double_[_1 = c.real()] << ", " << double_[_1 = c.imag()] << ')'
            | double_[_1 = c.real()]
          )
      );
    }
}

int main(){
  std::string str;
  while(getline(std::cin, str)){
    if(str.empty() || str[0] == 'q' || str[0] == 'Q')
      break;

    std::complex<double> c;
    if(client::parse_complex(str.begin(), str.end(), c)){
      std::string generated;
      std::back_insert_iterator<std::string> sink(generated);
      if(!client::generate_complex(sink, c)){
        std::cout << "Generating failed\n";
      }
      else{
        std::cout << "Generated: " << generated << std::endl;
      }
    }
    else{
      std::cout << "Parsing failed\n";
    }
  }

  return 0;
}