#pragma once

#include "rev/Utilities.h"
#include "rev/geometry/Tools.h"
#include <array>
#include <glm/glm.hpp>
#include <gsl/gsl_assert>
#include <iostream>
#include <limits>
#include <set>
#include <unordered_set>
#include <variant>
#include <vector>

namespace rev {

template <typename SurfaceData>
struct Triangle {
    std::array<glm::vec3, 3> vertices;
    SurfaceData data;

    AxisAlignedBoundingBox getBoundingBox() const
    {
        AxisAlignedBoundingBox box;
        for (const auto& vertex : vertices) {
            box.expandToVertex(vertex);
        }
        return box;
    }

    glm::vec3 getNormal() const
    {
        auto e1 = vertices[1] - vertices[0];
        auto e2 = vertices[2] - vertices[0];
        return glm::normalize(glm::cross(e1, e2));
    }

    bool isLeftOfPlane(const AxisAlignedPlane& plane) const
    {
        for (const auto& vertex : vertices) {
            if (vertex[plane.dimensionIndex] < plane.boundary) {
                return true;
            }
        }
        return false;
    }

    bool isRightOfPlane(const AxisAlignedPlane& plane) const
    {
        for (const auto& vertex : vertices) {
            if (vertex[plane.dimensionIndex] > plane.boundary) {
                return true;
            }
        }
        return false;
    }

    std::optional<AxisAlignedBoundingBox> clippedBoundingBox(const AxisAlignedBoundingBox& box)
    {
        // Alternate between these two buffers for input and output
        std::array<glm::vec3, 9> buf1{};
        std::array<glm::vec3, 9> buf2{};

        // Fill the input buffer
        for (int i = 0; i < 3; i++) {
            buf1[i] = vertices[i];
        }

        // Clip the triangle by all sides of our bounding box
        auto inputIter = buf1.data();
        auto inputEnd = inputIter + 3;
        auto outputIter = buf2.data();

        // First the min planes in each dimension
        for (uint8_t dim = 0; dim < 3; dim++) {
            gsl::span<glm::vec3> inputRange(inputIter, inputEnd);
            OutputIteratorPolygonBuilder builder(outputIter);
            AxisAlignedPlane boundary = { dim, box.minimum[dim] };
            boundary.splitConvexPolygon(inputRange, NullPolygonBuilder(), builder);

            if (std::distance(outputIter, builder.getIterator()) < 3) {
                return {};
            }

            std::swap(inputIter, outputIter);
            inputEnd = builder.getIterator();
        }

        // Now the max planes in each dimension
        for (uint8_t dim = 0; dim < 3; dim++) {
            gsl::span<glm::vec3> inputRange(inputIter, inputEnd);
            OutputIteratorPolygonBuilder builder(outputIter);
            AxisAlignedPlane boundary = { dim, box.maximum[dim] };
            boundary.splitConvexPolygon(inputRange, builder, NullPolygonBuilder());

            if (std::distance(outputIter, builder.getIterator()) < 3) {
                return {};
            }

            std::swap(inputIter, outputIter);
            inputEnd = builder.getIterator();
        }

        gsl::span<glm::vec3> inputRange(inputIter, inputEnd);
        auto clippedBox = smallestBoxContainingVertices(inputRange);
        clippedBox.reduceToBox(box);
        return clippedBox;
    }

    struct Hit {
        glm::vec2 uv;
        float t;
    };

    std::optional<Hit> castRay(const Ray& ray) const
    {
        glm::vec3 edge1 = vertices[1] - vertices[0];
        glm::vec3 edge2 = vertices[2] - vertices[0];

        glm::vec3 p = glm::cross(ray.direction, edge2);
        float determinant = glm::dot(edge1, p);
        if (abs(determinant) < std::numeric_limits<float>::epsilon()) {
            // The ray is parallel.
            return std::nullopt;
        }

        glm::vec3 fromV0 = ray.origin - vertices[0];
        float u = glm::dot(fromV0, p) / determinant;
        if (u < 0.0f || u > 1.0f) {
            return std::nullopt;
        }

        glm::vec3 q = glm::cross(fromV0, edge1);
        float v = glm::dot(ray.direction, q) / determinant;
        if (v < 0.0f || v > 1.0f) {
            return std::nullopt;
        }

        float t = glm::dot(edge2, q) / determinant;
        if (t < 0.0f) {
            return std::nullopt;
        }

        return Hit{ { u, v }, t };
    }

    glm::vec3 baryCentricToCartesian(glm::vec2 uv)
    {
        float u = uv[0];
        float v = uv[1];
        float w = 1.0f - (u + v);
        return w * vertices[0] + u * vertices[1] + v * vertices[2];
    }

    glm::vec2 cartesianToBarycentric(glm::vec3 position)
    {
        glm::vec3 v0 = vertices[1] - vertices[0];
        glm::vec3 v1 = vertices[2] - vertices[0];
        glm::vec3 v2 = position - vertices[0];

        float d00 = glm::dot(v0, v0);
        float d01 = glm::dot(v0, v1);
        float d11 = glm::dot(v1, v1);
        float d20 = glm::dot(v2, v0);
        float d21 = glm::dot(v2, v1);
        float denom = d00 * d11 - d01 * d01;

        float v = (d11 * d20 - d01 * d21) / denom;
        float w = (d00 * d21 - d01 * d20) / denom;
        float u = 1.0f - v - w;

        return { u, v };
    }

    std::optional<float> intersectsSphere(const Sphere& sphere)
    {
        glm::vec3 normal = getNormal();
        float t = glm::dot(normal, (vertices[0] - sphere.center));
        if (sphere.radius < abs(t)) {
            // Closest point to the triangle is outside the sphere.
            return std::nullopt;
        }

        glm::vec3 closestPoint = sphere.center - (normal * t);
        glm::vec2 cpBary = cartesianToBarycentric(closestPoint);
        size_t outsideCount = 0;
        glm::bvec3 isOutside{};
        if (cpBary[0] < 0.0f) {
            outsideCount++;
            isOutside[0] = true;
        }

        if (cpBary[1] < 0.0f) {
            outsideCount++;
            isOutside[1] = true;
        }

        if ((1.0f - (cpBary[0] + cpBary[1])) < 0.0f) {
            outsideCount++;
            isOutside[2] = true;
        }

        if (outsideCount == 0) {
            // Closest point is inside the triangle and inside the sphere.
            return t;
        }

        // Calculate the square of the radius of the cross section of the sphere passing through
        // the triangle's plane.
        float r2 = (sphere.radius * sphere.radius) - (t * t);

        if (outsideCount == 1) {
            // We need to find the distance to the edge (squared) here.
            glm::vec3 edge;
            glm::vec3 toVertex;
            if (isOutside[0]) {
                edge = vertices[2] - vertices[1];
                toVertex = closestPoint - vertices[1];
            } else if (isOutside[1]) {
                edge = vertices[2] - vertices[0];
                toVertex = closestPoint - vertices[0];
            } else {
                edge = vertices[1] - vertices[0];
                toVertex = closestPoint - vertices[0];
            }

            float toVertexLength = glm::length(toVertex);
            float legLength = glm::dot(edge, toVertex) / toVertexLength;

            return ((toVertexLength * toVertexLength) - (legLength * legLength)) < r2;
        }

        // Otherwise, one of the vertices is actually the closest point on the triangle.
        glm::vec3 toVertex;
        if (!isOutside[0]) {
            toVertex = closestPoint - vertices[0];
        } else if (!isOutside[1]) {
            toVertex = closestPoint - vertices[1];
        } else {
            toVertex = closestPoint - vertices[2];
        }

        return glm::dot(toVertex, toVertex) < r2;
    }
};

template <typename SurfaceData>
struct LeafNode {
    std::unordered_set<Triangle<SurfaceData>*> triangles;
};

template <typename SurfaceData>
struct BranchNode;

template <typename SurfaceData>
using MapNode = std::variant<LeafNode<SurfaceData>, BranchNode<SurfaceData>>;

template <typename SurfaceData>
struct BranchNode {
    AxisAlignedPlane split;

    std::unique_ptr<MapNode<SurfaceData>> left;
    std::unique_ptr<MapNode<SurfaceData>> right;
};

template <typename SurfaceData>
class KDTree {
public:
    KDTree(std::vector<std::unique_ptr<Triangle<SurfaceData>>> triangles,
        std::unique_ptr<MapNode<SurfaceData>> root, const AxisAlignedBoundingBox boundingBox)
        : _triangles(std::move(triangles))
        , _root(std::move(root))
        , _boundingBox(boundingBox)
    {
    }

    struct Hit {
        Triangle<SurfaceData>* triangle;
        glm::vec2 uv;
        float t;
    };

    std::optional<Hit> castRay(
        const Ray& ray, float maxDistance = std::numeric_limits<float>::infinity()) const
    {
        if (_boundingBox.containsPoint(ray.origin)) {
            AxisAlignedBoundingBox::Hit entryPoint{ ray.origin, 0.0f, 0 };
            return getHit(*_root, ray, _boundingBox, entryPoint, maxDistance);
        } else {
            auto hit = _boundingBox.castExternalRay(ray);
            if (!hit || (hit->t > maxDistance)) {
                return std::nullopt;
            }
            return getHit(*_root, ray, _boundingBox, *hit, maxDistance);
        }
    }

    template <typename Visitor>
    void visitTrianglesIntersectingSphere(const Sphere& sphere, Visitor&& visitor) const
    {
        std::visit(
            [this, &sphere, &visitor](const auto& node) {
                visitTrianglesIntersectingSphereInNode(
                    sphere, std::forward<Visitor>(visitor), _boundingBox, node);
            },
            *_root);
    }

    void dump() const
    {
        std::cout << "[KDTree]{" << std::endl;
        printBoundingBox(_boundingBox, 2);
        std::visit([this](const auto& root) { printNode(_boundingBox, root, 2); }, *_root);
        std::cout << "}" << std::endl;
    }

private:
    std::optional<Hit> getHit(const MapNode<SurfaceData>& node, const Ray& ray,
        const AxisAlignedBoundingBox& box, AxisAlignedBoundingBox::Hit& entryPoint,
        float maxDistance) const
    {
        return std::visit(
            [this, &ray, &box, &entryPoint, maxDistance](
                const auto& node) { return getHit(node, ray, box, entryPoint, maxDistance); },
            node);
    }

    std::optional<Hit> getHit(const BranchNode<SurfaceData>& node, const Ray& ray,
        const AxisAlignedBoundingBox& box, AxisAlignedBoundingBox::Hit& entryPoint,
        float maxDistance) const
    {
        auto [leftBox, rightBox] = box.split(node.split);
        float position = entryPoint.intersectionPoint[node.split.dimensionIndex];
        bool goLeft;
        if (position < node.split.boundary) {
            goLeft = true;
        } else if (position > node.split.boundary) {
            goLeft = false;
        } else if (ray.direction[node.split.dimensionIndex] < 0.0f) {
            goLeft = false;
        } else if (ray.direction[node.split.dimensionIndex] > 0.0f) {
            goLeft = true;
        } else {
            return std::nullopt;
        }

        MapNode<SurfaceData>* childNode = goLeft ? node.left.get() : node.right.get();
        auto hit = getHit(*childNode, ray, (goLeft ? leftBox : rightBox), entryPoint, maxDistance);
        if (hit) {
            return hit;
        }
        if (entryPoint.t > maxDistance) {
            return std::nullopt;
        }
        if (entryPoint.planeDimension != node.split.dimensionIndex) {
            return std::nullopt;
        }

        bool shouldTryOtherChild;
        if (goLeft) {
            shouldTryOtherChild = (ray.direction[node.split.dimensionIndex] > 0.0f);
        } else {
            shouldTryOtherChild = (ray.direction[node.split.dimensionIndex] < 0.0f);
        }

        if (shouldTryOtherChild) {
            childNode = goLeft ? node.right.get() : node.left.get();
            return getHit(*childNode, ray, (goLeft ? rightBox : leftBox), entryPoint, maxDistance);
        }
        return std::nullopt;
    }

    std::optional<Hit> getHit(const LeafNode<SurfaceData>& node, const Ray& ray,
        const AxisAlignedBoundingBox& box, AxisAlignedBoundingBox::Hit& entryPoint,
        float maxDistance) const
    {
        float t = std::numeric_limits<float>::infinity();
        std::optional<Hit> bestHit;
        for (const auto& triangle : node.triangles) {
            auto hit = triangle->castRay(ray);
            if (hit) {
                if (!bestHit || (hit->t < bestHit->t)) {
                    bestHit = Hit{ triangle, hit->uv, hit->t };
                }
            }
        }

        if (bestHit && !(bestHit->t > maxDistance)) {
            return bestHit;
        }

        auto boxHit = box.castInternalRay(ray);
        entryPoint = boxHit;

        return std::nullopt;
    }

    template <typename Visitor>
    void visitTrianglesIntersectingSphereInNode(const Sphere& sphere, Visitor&& visitor,
        const AxisAlignedBoundingBox& box, const LeafNode<SurfaceData>& node) const
    {
        if (!box.intersectsSphere(sphere)) {
            return;
        }

        for (const auto& triangle : node.triangles) {
            auto hit = triangle->intersectsSphere(sphere);
            if (hit) {
                visitor(*triangle, *hit);
            }
        }
    }

    template <typename Visitor>
    void visitTrianglesIntersectingSphereInNode(const Sphere& sphere, Visitor&& visitor,
        const AxisAlignedBoundingBox& box, const BranchNode<SurfaceData>& node) const
    {
        if (!box.intersectsSphere(sphere)) {
            return;
        }

        auto [leftBox, rightBox] = box.split(node.split);
        std::visit(
            [this, leftBox = leftBox, &sphere, &visitor](const auto& node) {
                visitTrianglesIntersectingSphereInNode(
                    sphere, std::forward<Visitor>(visitor), leftBox, node);
            },
            *node.left);
        std::visit(
            [this, rightBox = rightBox, &sphere, &visitor](const auto& node) {
                visitTrianglesIntersectingSphereInNode(
                    sphere, std::forward<Visitor>(visitor), rightBox, node);
            },
            *node.right);
    }

    void printIndent(size_t indent = 0) const
    {
        for (size_t i = 0; i < indent; i++) {
            std::cout << " ";
        }
    }

    void printVertex(const glm::vec3& vertex) const
    {
        std::cout << "(" << vertex.x << ", " << vertex.y << ", " << vertex.z << ")";
    }

    void printBoundingBox(const AxisAlignedBoundingBox& box, size_t indent = 0) const
    {
        printIndent(indent);
        std::cout << "[Box]{ Min: ";
        printVertex(box.minimum);
        std::cout << " Max: ";
        printVertex(box.maximum);
        std::cout << " }" << std::endl;
    }

    void printTriangle(const Triangle<SurfaceData>& triangle, size_t indent = 0) const
    {
        printIndent(indent);
        std::cout << "[Triangle]{" << std::endl;
        for (const auto& vertex : triangle.vertices) {
            printIndent(indent + 2);
            printVertex(vertex);
            std::cout << std::endl;
        }
        printIndent(indent);
        std::cout << "}" << std::endl;
    }

    void printNode(const AxisAlignedBoundingBox& box, const LeafNode<SurfaceData>& node,
        size_t indent = 0) const
    {
        printIndent(indent);
        std::cout << "[Leaf]{" << std::endl;
        printBoundingBox(box, indent + 2);
        for (const auto& triangle : node.triangles) {
            printTriangle(*triangle, indent + 2);
        }
        printIndent(indent);
        std::cout << "}" << std::endl;
    }

    void printNode(const AxisAlignedBoundingBox& box, const BranchNode<SurfaceData>& node,
        size_t indent = 0) const
    {
        printIndent(indent);
        std::cout << "[Branch]{" << std::endl;
        printBoundingBox(box, indent + 2);
        auto [leftBox, rightBox] = box.split(node.split);
        std::visit([this, indent, leftBox = leftBox](
                       const auto& childNode) { printNode(leftBox, childNode, indent + 2); },
            *node.left);
        std::visit([this, indent, rightBox = rightBox](
                       const auto& childNode) { printNode(rightBox, childNode, indent + 2); },
            *node.right);
        printIndent(indent);
        std::cout << "}" << std::endl;
    }

    std::vector<std::unique_ptr<Triangle<SurfaceData>>> _triangles;
    std::unique_ptr<MapNode<SurfaceData>> _root;
    AxisAlignedBoundingBox _boundingBox;
};

template <typename SurfaceData>
class KDTreeBuilder {
public:
    void addTriangle(std::array<glm::vec3, 3> vertices, SurfaceData data)
    {
        Expects(!glm::any(glm::isnan(vertices[0])));
        Expects(!glm::any(glm::isnan(vertices[1])));
        Expects(!glm::any(glm::isnan(vertices[2])));

        auto triangle
            = std::make_unique<Triangle<SurfaceData>>(Triangle<SurfaceData>{ vertices, data });
        auto boundingBox = smallestBoxContainingVertices(vertices);
        _boundingBox.expandToBox(boundingBox);

        for (const auto& event : buildTriangleEvents(boundingBox, triangle.get())) {
            _events.insert(event);
        }
        _triangles.push_back(std::move(triangle));
    }

    KDTree<SurfaceData> build()
    {
        Expects(!_triangles.empty());
        Expects(!_events.empty());

        std::unordered_set<Triangle<SurfaceData>*> triangleSet;
        for (const auto& triangle : _triangles) {
            triangleSet.insert(triangle.get());
        }
        auto rootNode = createNode(_boundingBox, std::move(_events), triangleSet);
        return KDTree<SurfaceData>{ std::move(_triangles), std::move(rootNode), _boundingBox };
    }

private:
    static constexpr float kTraversalCost = 15.0f;
    static constexpr float kIntersectionCost = 20.0f;
    static constexpr float kEmptySplitDiscount = 0.8f;

    enum class Side {
        Left,
        Right,
    };

    struct Event {
        enum class Type {
            Ending = 0,
            Planar = 1,
            Starting = 2,
        };

        bool operator<(const Event& other) const
        {
            auto tuplify = [](const Event& event) {
                return std::tie(event.separationPlane.boundary,
                    event.separationPlane.dimensionIndex, event.type, event.triangle);
            };
            return tuplify(*this) < tuplify(other);
        }

        Triangle<SurfaceData>* triangle;
        AxisAlignedPlane separationPlane;
        Type type;
    };

    std::unique_ptr<MapNode<SurfaceData>> createNode(const AxisAlignedBoundingBox& box,
        std::set<Event> events, std::unordered_set<Triangle<SurfaceData>*> triangles)
    {
        size_t triangleCount = triangles.size();
        if (triangleCount > 0) {
            auto [plane, side, cost] = findBestSplit(box, events, triangleCount);
            float terminateCost = static_cast<float>(triangleCount) * kIntersectionCost;
            if (terminateCost > cost) {
                std::unordered_set<Triangle<SurfaceData>*> leftTriangles;
                std::unordered_set<Triangle<SurfaceData>*> rightTriangles;

                for (const auto& event : events) {
                    const auto& eventPlane = event.separationPlane;
                    if (eventPlane.dimensionIndex != plane.dimensionIndex) {
                        continue;
                    }

                    if (eventPlane.boundary < plane.boundary) {
                        if (event.type != Event::Type::Starting) {
                            size_t erased = triangles.erase(event.triangle);
                            Expects(erased == 1);
                            leftTriangles.insert(event.triangle);

                            Expects(event.triangle->isLeftOfPlane(plane));
                        }
                    } else if (eventPlane.boundary > plane.boundary) {
                        if (event.type != Event::Type::Ending) {
                            size_t erased = triangles.erase(event.triangle);
                            Expects(erased == 1);
                            rightTriangles.insert(event.triangle);

                            Expects(event.triangle->isRightOfPlane(plane));
                        }
                    } else {
                        // On the boundary
                        size_t erased = triangles.erase(event.triangle);
                        Expects(erased == 1);
                        switch (event.type) {
                        case Event::Type::Starting:
                            rightTriangles.insert(event.triangle);
                            Expects(event.triangle->isRightOfPlane(plane));
                            break;
                        case Event::Type::Ending:
                            leftTriangles.insert(event.triangle);
                            Expects(event.triangle->isLeftOfPlane(plane));
                            break;
                        case Event::Type::Planar:
                            if (side == Side::Left) {
                                leftTriangles.insert(event.triangle);
                            } else {
                                rightTriangles.insert(event.triangle);
                            }
                            break;
                        }
                    }
                }

                std::set<Event> leftEvents;
                std::set<Event> rightEvents;
                for (const auto& event : events) {
                    if (leftTriangles.count(event.triangle)) {
                        leftEvents.insert(event);
                    } else if (rightTriangles.count(event.triangle)) {
                        rightEvents.insert(event);
                    }
                }

                auto [leftBox, rightBox] = box.split(plane);
                // Overlapping triangles remain
                for (const auto& triangle : triangles) {
                    Expects(triangle->isLeftOfPlane(plane));
                    Expects(triangle->isRightOfPlane(plane));

                    auto leftClipBox = triangle->clippedBoundingBox(leftBox);
                    if (leftClipBox) {
                        for (const auto& event : buildTriangleEvents(*leftClipBox, triangle))
                        {
                            leftEvents.insert(event);
                        }
                        leftTriangles.insert(triangle);
                    }

                    auto rightClipBox = triangle->clippedBoundingBox(rightBox);
                    if (rightClipBox) {
                        for (const auto& event : buildTriangleEvents(*rightClipBox, triangle))
                        {
                            rightEvents.insert(event);
                        }
                        rightTriangles.insert(triangle);
                    }
                }
                triangles.clear();
                events.clear();

                return std::make_unique<MapNode<SurfaceData>>(BranchNode<SurfaceData>{
                    plane,
                    createNode(leftBox, std::move(leftEvents), std::move(leftTriangles)),
                    createNode(rightBox, std::move(rightEvents), std::move(rightTriangles)),
                });
            }
        }
        return std::make_unique<MapNode<SurfaceData>>(
            LeafNode<SurfaceData>{ std::move(triangles) });
    }

    std::set<Event> buildTriangleEvents(
        const AxisAlignedBoundingBox& boundingBox, Triangle<SurfaceData>* triangle)
    {
        AxisAlignedBoundingBox actualBox = triangle->getBoundingBox();
        std::set<Event> events;
        for (uint8_t k = 0; k < 3; k++) {
            float minimum = boundingBox.minimum[k];
            float maximum = boundingBox.maximum[k];
            Expects(!(minimum < actualBox.minimum[k]));
            Expects(!(maximum > actualBox.maximum[k]));
            if (minimum < maximum) {
                events.insert(
                    Event{ triangle, AxisAlignedPlane{ k, maximum }, Event::Type::Ending });
                events.insert(
                    Event{ triangle, AxisAlignedPlane{ k, minimum }, Event::Type::Starting });
            } else {
                events.insert(
                    Event{ triangle, AxisAlignedPlane{ k, minimum }, Event::Type::Planar });
            }
        }
        return events;
    }

    std::tuple<AxisAlignedPlane, Side, float> findBestSplit(
        const AxisAlignedBoundingBox& box, const std::set<Event>& events, size_t triangleCount)
    {
        float bestCost = std::numeric_limits<float>::infinity();
        AxisAlignedPlane bestPlane;
        Side bestSide;

        std::array<size_t, 3> leftTriangleCounts{ 0 };
        std::array<size_t, 3> rightTriangleCounts{ triangleCount, triangleCount, triangleCount };
        auto iter = events.begin();
        auto eventsEnd = events.end();
        while (iter != eventsEnd) {
            const AxisAlignedPlane& candidatePlane = iter->separationPlane;
            uint8_t dimension = candidatePlane.dimensionIndex;
            size_t endingTriangleCount = 0;
            size_t planarTriangleCount = 0;
            size_t startingTriangleCount = 0;
            auto validate = [&eventsEnd, &candidatePlane, dimension](const auto& iter) {
                return (iter != eventsEnd)
                    && !(iter->separationPlane.boundary > candidatePlane.boundary)
                    && (iter->separationPlane.dimensionIndex == dimension);
            };
            while (validate(iter) && (iter->type == Event::Type::Ending)) {
                endingTriangleCount++;
                iter++;
            }
            while (validate(iter) && (iter->type == Event::Type::Planar)) {
                planarTriangleCount++;
                iter++;
            }
            while (validate(iter) && (iter->type == Event::Type::Starting)) {
                startingTriangleCount++;
                iter++;
            }
            rightTriangleCounts[dimension] -= planarTriangleCount;
            rightTriangleCounts[dimension] -= endingTriangleCount;
            auto [side, cost] = surfaceAreaCostHeuristic(box, candidatePlane,
                leftTriangleCounts[dimension], planarTriangleCount, rightTriangleCounts[dimension]);
            if (cost < bestCost) {
                bestCost = cost;
                bestPlane = candidatePlane;
                bestSide = side;
            }
            leftTriangleCounts[dimension] += planarTriangleCount;
            leftTriangleCounts[dimension] += startingTriangleCount;
        }
        return { bestPlane, bestSide, bestCost };
    }

    std::pair<Side, float> surfaceAreaCostHeuristic(const AxisAlignedBoundingBox& box,
        const AxisAlignedPlane& splitPlane, size_t leftTriangleCount, size_t planarTriangleCount,
        size_t rightTriangleCount)
    {
        auto [leftBox, rightBox] = box.split(splitPlane);
        float boxSurfaceArea = box.getSurfaceArea();

        float leftArea = leftBox.getSurfaceArea() / boxSurfaceArea;
        float rightArea = rightBox.getSurfaceArea() / boxSurfaceArea;
        bool removesVolume = (leftBox.getVolume() > 0.0f) && (rightBox.getVolume() > 0.0f);

        auto cost = [leftArea, rightArea, removesVolume](
                        size_t leftTriangleCount, size_t rightTriangleCount) {
            float cost = static_cast<float>(leftTriangleCount) * leftArea
                + static_cast<float>(rightTriangleCount) * rightArea;
            if (removesVolume && (!leftTriangleCount || !rightTriangleCount)) {
                // We slightly prefer splits that remove empty space.
                cost *= kEmptySplitDiscount;
            }
            return cost;
        };

        float leftCost = cost(leftTriangleCount + planarTriangleCount, rightTriangleCount);
        float rightCost = cost(leftTriangleCount, planarTriangleCount + rightTriangleCount);
        if (leftCost < rightCost) {
            return { Side::Left, (kIntersectionCost * leftCost) + kTraversalCost };
        } else {
            return { Side::Right, (kIntersectionCost * rightCost) + kTraversalCost };
        }
    }

    AxisAlignedBoundingBox _boundingBox;
    std::vector<std::unique_ptr<Triangle<SurfaceData>>> _triangles;
    std::set<Event> _events;
};
}


/* rng/taus113.c
 * Copyright (C) 2002 Atakan Gurkan
 * Based on the file taus.c which has the notice
 * Copyright (C) 1996, 1997, 1998, 1999, 2000, 2007 James Theiler, Brian Gough
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 *
 * This program 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.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/* This is a maximally equidistributed combined, collision free
   Tausworthe generator, with a period ~2^{113}. The sequence is,

   x_n = (z1_n ^ z2_n ^ z3_n ^ z4_n)

   b = (((z1_n <<  6) ^ z1_n) >> 13)
   z1_{n+1} = (((z1_n & 4294967294) << 18) ^ b)
   b = (((z2_n <<  2) ^ z2_n) >> 27)
   z2_{n+1} = (((z2_n & 4294967288) <<  2) ^ b)
   b = (((z3_n << 13) ^ z3_n) >> 21)
   z3_{n+1} = (((z3_n & 4294967280) <<  7) ^ b)
   b = (((z4_n <<  3)  ^ z4_n) >> 12)
   z4_{n+1} = (((z4_n & 4294967168) << 13) ^ b)

   computed modulo 2^32. In the formulas above '^' means exclusive-or
   (C-notation), not exponentiation.
   The algorithm is for 32-bit integers, hence a bitmask is used to clear
   all but least significant 32 bits, after left shifts, to make the code
   work on architectures where integers are 64-bit.

   The generator is initialized with
   zi = (69069 * z{i+1}) MOD 2^32 where z0 is the seed provided
   During initialization a check is done to make sure that the initial seeds
   have a required number of their most significant bits set.
   After this, the state is passed through the RNG 10 times to ensure the
   state satisfies a recurrence relation.

   References:
   P. L'Ecuyer, "Tables of Maximally-Equidistributed Combined LFSR Generators",
   Mathematics of Computation, 68, 225 (1999), 261--269.
     http://www.iro.umontreal.ca/~lecuyer/myftp/papers/tausme2.ps
   P. L'Ecuyer, "Maximally Equidistributed Combined Tausworthe Generators",
   Mathematics of Computation, 65, 213 (1996), 203--213.
     http://www.iro.umontreal.ca/~lecuyer/myftp/papers/tausme.ps
   the online version of the latter contains corrections to the print version.
*/

#include <config.h>
#include <stdlib.h>
#include <gsl/gsl_rng.h>

#define LCG(n) ((69069UL * n) & 0xffffffffUL)
#define MASK 0xffffffffUL

static inline unsigned long int taus113_get (void *vstate);
static double taus113_get_double (void *vstate);
static void taus113_set (void *state, unsigned long int s);

typedef struct
{
  unsigned long int z1, z2, z3, z4;
}
taus113_state_t;

static inline unsigned long
taus113_get (void *vstate)
{
  taus113_state_t *state = (taus113_state_t *) vstate;
  unsigned long b1, b2, b3, b4;

  b1 = ((((state->z1 << 6UL) & MASK) ^ state->z1) >> 13UL);
  state->z1 = ((((state->z1 & 4294967294UL) << 18UL) & MASK) ^ b1);

  b2 = ((((state->z2 << 2UL) & MASK) ^ state->z2) >> 27UL);
  state->z2 = ((((state->z2 & 4294967288UL) << 2UL) & MASK) ^ b2);

  b3 = ((((state->z3 << 13UL) & MASK) ^ state->z3) >> 21UL);
  state->z3 = ((((state->z3 & 4294967280UL) << 7UL) & MASK) ^ b3);

  b4 = ((((state->z4 << 3UL) & MASK) ^ state->z4) >> 12UL);
  state->z4 = ((((state->z4 & 4294967168UL) << 13UL) & MASK) ^ b4);

  return (state->z1 ^ state->z2 ^ state->z3 ^ state->z4);

}

static double
taus113_get_double (void *vstate)
{
  return taus113_get (vstate) / 4294967296.0;
}

static void
taus113_set (void *vstate, unsigned long int s)
{
  taus113_state_t *state = (taus113_state_t *) vstate;

  if (!s)
    s = 1UL;                    /* default seed is 1 */

  state->z1 = LCG (s);
  if (state->z1 < 2UL)
    state->z1 += 2UL;
  state->z2 = LCG (state->z1);
  if (state->z2 < 8UL)
    state->z2 += 8UL;
  state->z3 = LCG (state->z2);
  if (state->z3 < 16UL)
    state->z3 += 16UL;
  state->z4 = LCG (state->z3);
  if (state->z4 < 128UL)
    state->z4 += 128UL;

  /* Calling RNG ten times to satify recurrence condition */
  taus113_get (state);
  taus113_get (state);
  taus113_get (state);
  taus113_get (state);
  taus113_get (state);
  taus113_get (state);
  taus113_get (state);
  taus113_get (state);
  taus113_get (state);
  taus113_get (state);

  return;
}

static const gsl_rng_type taus113_type = {
  "taus113",                    /* name */
  0xffffffffUL,                 /* RAND_MAX */
  0,                            /* RAND_MIN */
  sizeof (taus113_state_t),
  &taus113_set,
  &taus113_get,
  &taus113_get_double
};

const gsl_rng_type *gsl_rng_taus113 = &taus113_type;

/*  Rules for analytic calculations using GNU Emacs Calc:
    (used to find the values for the test program)

  [ LCG(n) := n * 69069 mod (2^32) ]

  [ b1(x) := rsh(xor(lsh(x, 6), x), 13),
  q1(x) := xor(lsh(and(x, 4294967294), 18), b1(x)),
  b2(x) := rsh(xor(lsh(x, 2), x), 27),
  q2(x) := xor(lsh(and(x, 4294967288), 2), b2(x)),
  b3(x) := rsh(xor(lsh(x, 13), x), 21),
  q3(x) := xor(lsh(and(x, 4294967280), 7), b3(x)),
  b4(x) := rsh(xor(lsh(x, 3), x), 12),
  q4(x) := xor(lsh(and(x, 4294967168), 13), b4(x))
  ]

  [ S([z1,z2,z3,z4]) := [q1(z1), q2(z2), q3(z3), q4(z4)] ]
*/


/*
 * Copyright 2021 The DAPHNE Consortium
 *
 * 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.
 */

#ifndef SRC_RUNTIME_LOCAL_KERNELS_SYRK_H
#define SRC_RUNTIME_LOCAL_KERNELS_SYRK_H

#include <runtime/local/context/DaphneContext.h>
#include <runtime/local/datastructures/CSRMatrix.h>
#include <runtime/local/datastructures/DataObjectFactory.h>
#include <runtime/local/datastructures/DenseMatrix.h>

#include <cblas.h>

// ****************************************************************************
// Struct for partial template specialization
// ****************************************************************************

template<class DTRes, class DTArg>
struct Syrk {
    static void apply(DTRes *& res, const DTArg * arg, DCTX(ctx)) = delete;
};

// ****************************************************************************
// Convenience function
// ****************************************************************************

template<class DTRes, class DTArg>
void syrk(DTRes *& res, const DTArg * arg, DCTX(ctx)) {
    Syrk<DTRes, DTArg>::apply(res, arg, ctx);
}

// ****************************************************************************
// (Partial) template specializations for different data/value types
// ****************************************************************************

// ----------------------------------------------------------------------------
// DenseMatrix <- DenseMatrix
// ----------------------------------------------------------------------------

template<>
struct Syrk<DenseMatrix<double>, DenseMatrix<double>> {
    static void apply(DenseMatrix<double> *& res, const DenseMatrix<double> * arg, DCTX(ctx)) {
        const size_t numRows = arg->getNumRows();
        const size_t numCols = arg->getNumCols();

        if(res == nullptr)
            res = DataObjectFactory::create<DenseMatrix<double>>(numCols, numCols, false);

        cblas_dsyrk(CblasRowMajor,
            CblasUpper,
            CblasTrans,
            numCols,
            numRows,
            1.0,
            arg->getValues(),
            arg->getRowSkip(),
            0.0,
            res->getValues(),
            res->getRowSkip());
        for (auto r = 0u; r < numCols; ++r) {
            for (auto c = r + 1; c < numCols; ++c) {
                res->set(c, r, res->get(r, c));
            }
        }
    }
};

template<>
struct Syrk<DenseMatrix<float>, DenseMatrix<float>> {
    static void apply(DenseMatrix<float> *& res, const DenseMatrix<float> * arg, DCTX(ctx)) {
        const size_t numRows = arg->getNumRows();
        const size_t numCols = arg->getNumCols();

        if(res == nullptr)
            res = DataObjectFactory::create<DenseMatrix<float>>(numCols, numCols, false);

        cblas_ssyrk(CblasRowMajor,
            CblasUpper,
            CblasTrans,
            numCols,
            numRows,
            1.0,
            arg->getValues(),
            arg->getRowSkip(),
            0.0,
            res->getValues(),
            res->getRowSkip());
        for (auto r = 0u; r < numCols; ++r) {
            for (auto c = r + 1; c < numCols; ++c) {
                res->set(c, r, res->get(r, c));
            }
        }
    }
};

// ----------------------------------------------------------------------------
// CSRMatrix <- CSRMatrix
// ----------------------------------------------------------------------------

template<typename VT>
struct Syrk<CSRMatrix<VT>, CSRMatrix<VT>> {
    static void apply(CSRMatrix<VT> *& res, const CSRMatrix<VT> * arg, DCTX(ctx)) {
        const size_t numRows = arg->getNumRows();
        const size_t numCols = arg->getNumCols();

        if(res == nullptr)
            res = DataObjectFactory::create<CSRMatrix<VT>>(numCols, numRows, arg->getNumNonZeros(), false);
        assert(false && "TODO: Syrk for Sparse");
    }
};

#endif //SRC_RUNTIME_LOCAL_KERNELS_SYRK_H


/* rng/r250.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 James Theiler, Brian Gough
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or (at
 * your option) any later version.
 *
 * This program 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.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#include <config.h>
#include <stdlib.h>
#include <gsl/gsl_rng.h>

/* This is a shift-register random number generator. The sequence is

   x_n = x_{n-103} ^ x_{n-250}        ("^" means XOR)

   defined on 32-bit words.

   The first 250 elements x_1 .. x_250 are first initialized as x_n =
   s_n, where s_n = (69069*s_{n-1}) mod 2^32 and s_0=s is the
   user-supplied seed. To ensure that the sequence does not lie on a
   subspace we force 32 of the entries to be linearly independent.  We
   take the 32 elements x[3], x[10], x[17], x[24], ..., 213 and apply
   the following operations,

   x[3]   &= 11111111111111111111111111111111
   x[3]   |= 10000000000000000000000000000000
   x[10]  &= 01111111111111111111111111111111
   x[10]  |= 01000000000000000000000000000000
   x[17]  &= 00111111111111111111111111111111
   x[17]  |= 00100000000000000000000000000000
   ....      ...
   x[206] &= 00000000000000000000000000000111
   x[206] |= 00000000000000000000000000000100
   x[213] &= 00000000000000000000000000000011
   x[213] |= 00000000000000000000000000000010
   x[220] &= 00000000000000000000000000000001
   x[220] |= 00000000000000000000000000000001

   i.e. if we consider the bits of the 32 elements as forming a 32x32
   array then we are setting the diagonal bits of the array to one and
   masking the lower triangle below the diagonal to zero.

   With this initialization procedure the theoretical value of
   x_{10001} is 1100653588 for s = 1 (Actually I got this by running
   the original code). The subscript 10001 means (1) seed the
   generator with s = 1 and then do 10000 actual iterations.

   The period of this generator is about 2^250.

   The algorithm works for any number of bits. It is implemented here
   for 32 bits.

   From: S. Kirkpatrick and E. Stoll, "A very fast shift-register
   sequence random number generator", Journal of Computational Physics,
   40, 517-526 (1981). */

static inline unsigned long int r250_get (void *vstate);
static double r250_get_double (void *vstate);
static void r250_set (void *state, unsigned long int s);

typedef struct
  {
    int i;
    unsigned long x[250];
  }
r250_state_t;

static inline unsigned long int
r250_get (void *vstate)
{
  r250_state_t *state = (r250_state_t *) vstate;
  unsigned long int k;
  int j;

  int i = state->i;

  if (i >= 147)
    {
      j = i - 147;
    }
  else
    {
      j = i + 103;
    }

  k = state->x[i] ^ state->x[j];
  state->x[i] = k;

  if (i >= 249)
    {
      state->i = 0;
    }
  else
    {
      state->i = i + 1;
    }

  return k;
}

static double
r250_get_double (void *vstate)
{
  return r250_get (vstate) /  4294967296.0 ;
}

static void
r250_set (void *vstate, unsigned long int s)
{
  r250_state_t *state = (r250_state_t *) vstate;

  int i;

  if (s == 0)
    s = 1;	/* default seed is 1 */

  state->i = 0;

#define LCG(n) ((69069 * n) & 0xffffffffUL)

  for (i = 0; i < 250; i++)	/* Fill the buffer  */
    {
      s = LCG (s);
      state->x[i] = s;
    }

  {
    /* Masks for turning on the diagonal bit and turning off the
       leftmost bits */

    unsigned long int msb = 0x80000000UL;
    unsigned long int mask = 0xffffffffUL;

    for (i = 0; i < 32; i++)
      {
	int k = 7 * i + 3;	/* Select a word to operate on        */
	state->x[k] &= mask;	/* Turn off bits left of the diagonal */
	state->x[k] |= msb;	/* Turn on the diagonal bit           */
	mask >>= 1;
	msb >>= 1;
      }
  }

  return;
}

static const gsl_rng_type r250_type =
{"r250",			/* name */
 0xffffffffUL,			/* RAND_MAX */
 0,			        /* RAND_MIN */
 sizeof (r250_state_t),
 &r250_set,
 &r250_get,
 &r250_get_double};

const gsl_rng_type *gsl_rng_r250 = &r250_type;


/*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*
**                                                                                **
** This file forms part of the Underworld geophysics modelling application.         **
**                                                                                **
** For full license and copyright information, please refer to the LICENSE.md file  **
** located at the project root, or contact the authors.                             **
**                                                                                **
**~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*/
/*

Given a block system
A x = b,
or
( K G )(u) = (f)
( D C )(p)   (h)

We define a symmetrically scaled system,
L A R R^-1 x = L b,		A'x' = b'
where L = diag(L1, L2), R = diag(R1,R2) are block diagonal
and L_i and R_i are both diagonal matrirces.

The scaling produces
( L1   )( K G )( R1   )( R1   )^-1(u) = ( L1   )(f)
(   L2 )( D C )(   R2 )(   R2 )   (p)   (   L2 )(h)
and we solve
( L1   )( K G )( R1   )(v) = ( L1   )(f)
(   L2 )( D C )(   R2 )(q)   (   L2 )(h)
or
( L1KR1 L1GR2 )(v) = (L1f)
( L2DR1 L2CR2 )(q)   (L2h)

The solution u,p is recovered via
(u) = ( R1   )(v)
(p)   (   R2 )(q)

*/

#include <petsc.h>
#include <petscmat.h>
#include <petscvec.h>

#include <StGermain/StGermain.h>
#include <StgDomain/StgDomain.h>

#include "common-driver-utils.h"
#include "stokes_block_scaling.h"

/* private prototypes */
PetscErrorCode BSSCR_MatStokesBlock_ApplyScaling( MatStokesBlockScaling BA, Mat A, Vec b, Vec x, Mat S, PetscTruth sym );

/* A x = b -> A'x' = b' */
/*
Note this routine actually modifies the matrix and rhs b.
*/

// updated
PetscErrorCode BSSCR_MatStokesBlock_ApplyScaling( MatStokesBlockScaling BA, Mat A, Vec b, Vec x, Mat S, PetscTruth sym )
{
	Mat K,G,D,C;
	Vec L1,L2, R1,R2;
	Vec f,h, u,p;

	/* Get the scalings out the block mat data */
	VecNestGetSubVec( BA->Lz, 0, &L1 );
	VecNestGetSubVec( BA->Lz, 1, &L2 );
	VecNestGetSubVec( BA->Rz, 0, &R1 );
	VecNestGetSubVec( BA->Rz, 1, &R2 );

	/* get the subblock solution and rhs */
	if( x != PETSC_NULL ) {
		VecNestGetSubVec( x, 0, &u );
		VecNestGetSubVec( x, 1, &p );

		VecPointwiseDivide( u, u,R1); /* x <- x * 1/R1 */
		VecPointwiseDivide( p, p,R2);

	}
	if( b != PETSC_NULL ) {
		VecNestGetSubVec( b, 0, &f );
		VecNestGetSubVec( b, 1, &h );

		VecPointwiseMult( f, f,L1); /* f <- f * L1 */
		VecPointwiseMult( h, h,L2);

	}

	/* Scale matrices */
	MatNestGetSubMat( A, 0,0, &K );
	MatNestGetSubMat( A, 0,1, &G );
	MatNestGetSubMat( A, 1,0, &D );
	MatNestGetSubMat( A, 1,1, &C );

	if( K != PETSC_NULL ) {		MatDiagonalScale( K, L1,R1 );		}
	if( G != PETSC_NULL ) {		MatDiagonalScale( G, L1,R2 );		}
	if( D != PETSC_NULL && !sym ) {	MatDiagonalScale( D, L2,R1 );		}
	if( C != PETSC_NULL ) {		MatDiagonalScale( C, L2,R2 );		}
	if( S != PETSC_NULL ) {		MatDiagonalScale( S, L2,R2 );		}

	PetscFunctionReturn(0);
}
PetscErrorCode BSSCR_MatStokesBlock_ApplyScaling2( MatStokesBlockScaling BA, Mat A, Vec b, Vec x, Mat S, Mat M, PetscTruth sym )
{
	Mat K,G,D,C;
	Vec L1,L2, R1,R2;
	Vec f,h, u,p;

	/* Get the scalings out the block mat data */
	VecNestGetSubVec( BA->Lz, 0, &L1 );
	VecNestGetSubVec( BA->Lz, 1, &L2 );
	VecNestGetSubVec( BA->Rz, 0, &R1 );
	VecNestGetSubVec( BA->Rz, 1, &R2 );

	/* get the subblock solution and rhs */
	if( x != PETSC_NULL ) {
		VecNestGetSubVec( x, 0, &u );
		VecNestGetSubVec( x, 1, &p );

		VecPointwiseDivide( u, u,R1); /* x <- x * 1/R1 */
		VecPointwiseDivide( p, p,R2);

	}
	if( b != PETSC_NULL ) {
		VecNestGetSubVec( b, 0, &f );
		VecNestGetSubVec( b, 1, &h );

		VecPointwiseMult( f, f,L1); /* f <- f * L1 */
		VecPointwiseMult( h, h,L2);

	}

	/* Scale matrices */
	MatNestGetSubMat( A, 0,0, &K );
	MatNestGetSubMat( A, 0,1, &G );
	MatNestGetSubMat( A, 1,0, &D );
	MatNestGetSubMat( A, 1,1, &C );

	if( K != PETSC_NULL ) {		MatDiagonalScale( K, L1,R1 );		}
	if( G != PETSC_NULL ) {		MatDiagonalScale( G, L1,R2 );		}
	if( D != PETSC_NULL && !sym ) {	MatDiagonalScale( D, L2,R1 );		}
	if( C != PETSC_NULL ) {		MatDiagonalScale( C, L2,R2 );		}
	if( S != PETSC_NULL ) {		MatDiagonalScale( S, L2,R2 );		}
	if( M != PETSC_NULL ) {		MatDiagonalScale( M, L2,R2 );		}

	PetscFunctionReturn(0);
}

// updated
PetscErrorCode BSSCR_MatStokesBlockScalingCreate( MatStokesBlockScaling *_BA )
{
	MatStokesBlockScaling BA;

	PetscMalloc( sizeof(struct _p_MatStokesBlockScaling), &BA );

	BA->Lz = PETSC_NULL;
	BA->Rz = PETSC_NULL;

	BA->scaling_exists              = PETSC_FALSE;
	BA->scalings_have_been_inverted = PETSC_FALSE;
	BA->system_has_been_scaled      = PETSC_FALSE;

	*_BA = BA;

	PetscFunctionReturn(0);
}

// updated
PetscErrorCode BSSCR_MatStokesBlockScalingDestroy( MatStokesBlockScaling BA )
{

	if( BA->scaling_exists == PETSC_FALSE ) PetscFunctionReturn(0);

	if( BA->Lz != PETSC_NULL ) {
		Stg_VecDestroy(&BA->Lz );
		BA->Lz = PETSC_NULL;
	}

	if( BA->Rz != PETSC_NULL ) {
		Stg_VecDestroy(&BA->Rz );
		BA->Rz = PETSC_NULL;
	}

	PetscFree( BA );

	PetscFunctionReturn(0);
}

/*
A is 2x2 block matrix
b and x are 2x1 block vectors
*/
// updated
PetscErrorCode BSSCR_MatBlock_ConstructScaling( MatStokesBlockScaling BA, Mat A, Vec b, Vec x )
{
    if( BA->scaling_exists == PETSC_FALSE ) {

	VecDuplicate( x, &BA->Lz );
	VecDuplicate( x, &BA->Rz );

	BA->scaling_exists = PETSC_TRUE;
	BSSCR_MatStokesBlockDefaultBuildScaling( BA,A );
	BA->scalings_have_been_inverted = PETSC_FALSE;
    }

    PetscFunctionReturn(0);
}

// updated
PetscErrorCode BSSCR_mat_block_invert_scalings( MatStokesBlockScaling BA )
{
	Vec L1,L2, R1,R2;

	VecNestGetSubVec( BA->Lz, 0, &L1 );
	VecNestGetSubVec( BA->Lz, 1, &L2 );
	VecNestGetSubVec( BA->Rz, 0, &R1 );
	VecNestGetSubVec( BA->Rz, 1, &R2 );

	VecReciprocal(L1);
	VecReciprocal(L2);
	VecReciprocal(R1);
	VecReciprocal(R2);

	/* toggle inversion flag */
	if( BA->scalings_have_been_inverted == PETSC_TRUE ) {   BA->scalings_have_been_inverted = PETSC_FALSE;  }
	if( BA->scalings_have_been_inverted == PETSC_FALSE ) {  BA->scalings_have_been_inverted = PETSC_TRUE;   }

	PetscFunctionReturn(0);
}

/* updated */
PetscErrorCode BSSCR_MatStokesBlockScaleSystem( MatStokesBlockScaling BA, Mat A, Vec b, Vec x, Mat S, PetscTruth sym )
{
	if( BA->scaling_exists == PETSC_FALSE ) {
		BSSCR_MatBlock_ConstructScaling( BA,A,b,x );
	}

	if( BA->scalings_have_been_inverted == PETSC_TRUE ) {
		BSSCR_mat_block_invert_scalings(BA); /* to undo inversion */
	}

	BSSCR_MatStokesBlock_ApplyScaling(BA,A,b,x,S,sym);
	BA->system_has_been_scaled = PETSC_TRUE;

	/*
	PetscPrintf( PETSC_COMM_WORLD, "Post Scaling \n");
	MatBlock_ReportOperatorScales(A);
	*/

	PetscFunctionReturn(0);
}
PetscErrorCode BSSCR_MatStokesBlockScaleSystem2( MatStokesBlockScaling BA, Mat A, Vec b, Vec x, Mat S, Mat M, PetscTruth sym )
{
	if( BA->scaling_exists == PETSC_FALSE ) {
		BSSCR_MatBlock_ConstructScaling( BA,A,b,x );
	}

	if( BA->scalings_have_been_inverted == PETSC_TRUE ) {
		BSSCR_mat_block_invert_scalings(BA); /* to undo inversion */
	}

	BSSCR_MatStokesBlock_ApplyScaling2(BA,A,b,x,S,M,sym);
	BA->system_has_been_scaled = PETSC_TRUE;

	/*
	PetscPrintf( PETSC_COMM_WORLD, "Post Scaling \n");
	MatBlock_ReportOperatorScales(A);
	*/

	PetscFunctionReturn(0);
}

// updated
PetscErrorCode BSSCR_MatStokesBlockUnScaleSystem( MatStokesBlockScaling BA, Mat A, Vec b, Vec x, Mat S, PetscTruth sym )
{

	if( BA->system_has_been_scaled == PETSC_FALSE ) {
		printf("Warning: MatBlock has not been scaled !! \n");
		PetscFunctionReturn(0);
	}

	if( BA->scalings_have_been_inverted == PETSC_FALSE ) {
		BSSCR_mat_block_invert_scalings(BA);
	}

	BSSCR_MatStokesBlock_ApplyScaling(BA,A,b,x,S,sym);
	BA->system_has_been_scaled = PETSC_FALSE;

	PetscFunctionReturn(0);
}

PetscErrorCode BSSCR_MatStokesBlockReportOperatorScales( Mat A, PetscTruth sym )
{
	Vec rA, rG;
	PetscInt loc;
	PetscReal min, max;
	Mat K,G,D,C;

	MatNestGetSubMat( A, 0,0, &K );
	MatNestGetSubMat( A, 0,1, &G );
	MatNestGetSubMat( A, 1,0, &D );
	MatNestGetSubMat( A, 1,1, &C );

	MatGetVecs( K, PETSC_NULL, &rA );
	VecDuplicate( rA, &rG );

	/* Report the row max and mins */
	if (K!=PETSC_NULL) {
                PetscReal KNorm;
		MatGetRowMax( K, rA, PETSC_NULL );
		VecMax( rA, &loc, &max );
		PetscPrintf( PETSC_COMM_WORLD, "Sup_max(K) = %g \n", max );

		MatGetRowMinAbs( K, rA, PETSC_NULL );
		VecMin( rA, &loc, &min );
		PetscPrintf( PETSC_COMM_WORLD, "Sup_min(K) = %g \n", min );

		MatNorm(K,NORM_1,&KNorm);
		PetscPrintf( PETSC_COMM_WORLD, "norm_1(K)  = %g \n\n",KNorm);
	}

	if( G != PETSC_NULL ) {
		MatGetRowMax( G, rG, PETSC_NULL );
		VecMax( rG, &loc, &max );
		PetscPrintf( PETSC_COMM_WORLD, "Sup_max(G) = %g \n", max );

		MatGetRowMinAbs( G, rG, PETSC_NULL );
		VecMin( rG, &loc, &min );
		PetscPrintf( PETSC_COMM_WORLD, "Sup_min(G) = %g \n", min );
	}

	if( D != PETSC_NULL && !sym ) {
                Vec rD;

                MatGetVecs( D, PETSC_NULL, &rD );
		MatGetRowMax( D, rD, PETSC_NULL );
		VecMax( rD, &loc, &max );
		PetscPrintf( PETSC_COMM_WORLD, "Sup_max(D) = %g \n", max );

		MatGetRowMinAbs( D, rD, PETSC_NULL );
		VecMin( rD, &loc, &min );
		PetscPrintf( PETSC_COMM_WORLD, "Sup_min(D) = %g \n", min );

                Stg_VecDestroy(&rD );
	}

	if( C != PETSC_NULL ) {
		Vec cG;

		MatGetVecs( G, &cG, PETSC_NULL );
		MatGetRowMax( C, cG, PETSC_NULL );
		VecMax( cG, &loc, &max );
		PetscPrintf( PETSC_COMM_WORLD, "Sup_max(C) = %g \n", max );

		MatGetRowMin( C, cG, PETSC_NULL );
		VecMin( cG, &loc, &min );
		PetscPrintf( PETSC_COMM_WORLD, "Sup_min(C) = %g \n\n", min );

		Stg_VecDestroy(&cG);
	}

	Stg_VecDestroy(&rA );
	Stg_VecDestroy(&rG );

	PetscFunctionReturn(0);
}

// updated
PetscErrorCode BSSCR_MatStokesBlockDefaultBuildScaling( MatStokesBlockScaling BA, Mat A )
{
	Mat K,G,D,C;
	Vec rG;
	PetscScalar rg2, rg, ra;
	PetscInt N;
	Vec rA, rC;
	Vec L1,L2, R1,R2;

	VecNestGetSubVec( BA->Lz, 0, &L1 );
	VecNestGetSubVec( BA->Lz, 1, &L2 );

	VecNestGetSubVec( BA->Rz, 0, &R1 );
	VecNestGetSubVec( BA->Rz, 1, &R2 );

	rA = L1;
	rC = L2;

	MatNestGetSubMat( A, 0,0, &K );
	MatNestGetSubMat( A, 0,1, &G );
	MatNestGetSubMat( A, 1,0, &D );
	MatNestGetSubMat( A, 1,1, &C );

	VecDuplicate( rA, &rG );

	/* Get magnitude of K */
	//px_MatGetAbsRowSum( K, rA );
	MatGetRowMax( K, rA, PETSC_NULL );

	VecSqrt( rA );
	VecReciprocal( rA );

	VecDot( rA,rA, &ra );
	VecGetSize( rA, &N );
	ra = PetscSqrtScalar( ra/N );

	/* Get magnitude of G */
	//px_MatGetAbsRowSum( G, rG );
	MatGetRowMax( G, rG, PETSC_NULL );

	VecDot( rG, rG, &rg2 );
	VecGetSize( rG, &N );
	rg = PetscSqrtScalar(rg2/N);
	//      printf("rg = %f \n", rg );

	VecSet( rC, 1.0/(rg*ra) );

	Stg_VecDestroy(&rG );

	VecCopy( L1, R1 );
	VecCopy( L2, R2 );

	PetscFunctionReturn(0);
}


/* specfunc/legendre.h
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Gerard Jungman
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or (at
 * your option) any later version.
 *
 * This program 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.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

/* Author:  G. Jungman */

/* Declare private but non-local support functions
 * used in various Legendre function evaluations.
 */

#include <gsl/gsl_sf_result.h>

/* Large negative mu asymptotic
 * P^{-mu}_{-1/2 + I tau}, mu -> Inf
 * |x| < 1
 */
int
gsl_sf_conicalP_xlt1_large_neg_mu_e(double mu, double tau, double x,
                                       gsl_sf_result * result, double * ln_multiplier);

/* Large tau uniform asymptotics
 * P^{-mu}_{-1/2 + I tau}, tau -> Inf
 * 1 < x
 */
int
gsl_sf_conicalP_xgt1_neg_mu_largetau_e(const double mu, const double tau,
                                          const double x, double acosh_x,
                                          gsl_sf_result * result, double * ln_multiplier);

/* Large tau uniform asymptotics
 * P^{-mu}_{-1/2 + I tau}, tau -> Inf
 * -1 < x < 1
 */
int
gsl_sf_conicalP_xlt1_neg_mu_largetau_e(const double mu, const double tau,
                                          const double x, const double acos_x,
                                          gsl_sf_result * result, double * ln_multiplier);

/* P^{mu}_{-1/2 + I tau}
 * x->Inf
 *
 *  * This is effective to precision EPS for
 *
 *    (mu^2 + tau^2)/((1 + tau^2)^(1/2) x^2) < EPS^{1/3}
 *
 * since it goes only to a fixed order, based on the
 * representation in terms of hypegeometric functions
 * of argument 1/x^2.
 * [Zhurina+Karmazina, (3.8)]
 */
int
gsl_sf_conicalP_large_x_e(const double mu, const double tau, const double x,
                             gsl_sf_result * result, double * ln_multiplier);


/**
 * \author Sylvain Marsat, University of Maryland - NASA GSFC
 *
 * \brief C code for example generation of waveforms with the EOBNRv2HM reduced order model,
 * processing through the Fourier-domain LISA response and SNR calculation.
 *
 */

#define _XOPEN_SOURCE 500

#ifdef __GNUC__
#define UNUSED __attribute__ ((unused))
#else
#define UNUSED
#endif

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <complex.h>
#include <time.h>
#include <unistd.h>
#include <getopt.h>
#include <stdbool.h>
#include <string.h>

#include <gsl/gsl_errno.h>
#include <gsl/gsl_bspline.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_min.h>
#include <gsl/gsl_spline.h>
#include <gsl/gsl_complex.h>

#include "constants.h"
#include "struct.h"
#include "EOBNRv2HMROMstruct.h"
#include "EOBNRv2HMROM.h"
#include "wip.h"
#include "LISAgeometry.h"
#include "LISAFDresponse.h"

/* Parameters for the generation of a ROM waveform (in the form of a list of modes) */
/* All parameters are to be given in SI units! */
typedef struct tagROMParams {
  int nbmode;                /* Number of modes to generate (starting with the 22) - defaults to 1 (22 mode only) */
  double tRef;               /* shift in time with respect to the 22-fit-removed waveform */
  double phiRef;             /* phase at fRef */
  double fRef;               /* reference frequency */
  double m1;                 /* mass of companion 1 */
  double m2;                 /* mass of companion 2 */
  double distance;           /* distance of source */
  char outname[256];         /* file to which output should be written */
} ROMParams;

/* Parameters for the generation of a ROM waveform (in the form of a list of modes) */
/* All parameters are in SI units in the internals */
/* Angle definitions are taken from the Krolak&al paper gr-qc/0401108 */
typedef struct tagLISAParams {
  int nbmode;                /* Number of modes to generate (starting with the 22) - defaults to 1 (22 mode only) */
  double tRef;               /* shift in time with respect to the 22-fit-removed waveform */
  double phiRef;             /* phase at fRef */
  double fRef;               /* reference frequency */
  double m1;                 /* mass of companion 1 */
  double m2;                 /* mass of companion 2 */
  double distance;           /* distance of source */
  double inclination;        /* inclination of L relative to line of sight */
  double lambda;             /* First angle for the position in the sky of the source */
  double beta;               /* Second angle for the position in the sky of the source */
  double psi;                /* Polarization angle */
  char outname[256];         /* file to which output should be written */
} LISAParams;

/* Parse command line and return a newly allocated ROMParams object
 * Masses are input in solar masses and distances in Mpc - converted in SI for the internals */
static ROMParams* parse_args_ROM(ssize_t argc, char **argv) {
    ssize_t i;
    ROMParams* params;
    params = (ROMParams*) malloc(sizeof(ROMParams));
    memset(params, 0, sizeof(ROMParams));

    /* Set default values to the arguments */
    params->nbmode = 1;
    params->tRef = 0.;
    params->phiRef = 0.;
    params->fRef = 0.;
    params->m1 = 1. * 1e6 * MSUN_SI;
    params->m2 = 1. * 1e6 * MSUN_SI;
    params->distance = 1. * 1e9 * PC_SI;

    /* consume command line */
    for (i = 1; i < argc; ++i) {
        if (strcmp(argv[i], "--nbmode") == 0) {
            params->nbmode = atof(argv[++i]);
        } else if (strcmp(argv[i], "--tRef") == 0) {
            params->tRef = atof(argv[++i]);
        } else if (strcmp(argv[i], "--phiRef") == 0) {
            params->phiRef = atof(argv[++i]);
        } else if (strcmp(argv[i], "--fRef") == 0) {
            params->fRef = atof(argv[++i]);
        } else if (strcmp(argv[i], "--m1") == 0) {
            params->m1 = atof(argv[++i]) * MSUN_SI;
        } else if (strcmp(argv[i], "--m2") == 0) {
            params->m2 = atof(argv[++i]) * MSUN_SI;
        } else if (strcmp(argv[i], "--distance") == 0) {
            params->distance = atof(argv[++i]) * 1e6 * PC_SI;
        } else if (strcmp(argv[i], "--outname") == 0) {
            strncpy(params->outname, argv[++i], 256);
        } else {
            printf("Error: invalid option: %s\n", argv[i]);
            goto fail;
        }
    }

    return params;

    fail:
    free(params);
    exit(1);
}

/* Parse command line and return a newly allocated LISAParams object
 * Masses are input in solar masses and distances in Mpc - converted in SI for the internals */
static LISAParams* parse_args_LISA(ssize_t argc, char **argv) {
    ssize_t i;
    LISAParams* params;
    params = (LISAParams*) malloc(sizeof(LISAParams));
    memset(params, 0, sizeof(LISAParams));

    /* Set default values to the arguments */
    params->nbmode = 1;
    params->tRef = 0.;
    params->phiRef = 0.;
    params->fRef = 0.;
    params->m1 = 1. * 1e6 * MSUN_SI;
    params->m2 = 1. * 1e6 * MSUN_SI;
    params->distance = 1. * 1e9 * PC_SI;
    params->inclination = 0.;
    params->lambda = 0.;
    params->beta = 0.;
    params->psi = 0.;
    sprintf(params->outname, "");

    /* consume command line */
    for (i = 1; i < argc; ++i) {
        if (strcmp(argv[i], "--nbmode") == 0) {
            params->nbmode = atof(argv[++i]);
        } else if (strcmp(argv[i], "--tRef") == 0) {
            params->tRef = atof(argv[++i]);
        } else if (strcmp(argv[i], "--phiRef") == 0) {
            params->phiRef = atof(argv[++i]);
        } else if (strcmp(argv[i], "--fRef") == 0) {
            params->fRef = atof(argv[++i]);
        } else if (strcmp(argv[i], "--m1") == 0) {
            params->m1 = atof(argv[++i]) * MSUN_SI;
        } else if (strcmp(argv[i], "--m2") == 0) {
            params->m2 = atof(argv[++i]) * MSUN_SI;
        } else if (strcmp(argv[i], "--distance") == 0) {
            params->distance = atof(argv[++i]) * 1e6 * PC_SI;
        } else if (strcmp(argv[i], "--inclination") == 0) {
            params->inclination = atof(argv[++i]);
        } else if (strcmp(argv[i], "--lambda") == 0) {
            params->lambda = atof(argv[++i]);
        } else if (strcmp(argv[i], "--beta") == 0) {
            params->beta = atof(argv[++i]);
        } else if (strcmp(argv[i], "--psi") == 0) {
            params->psi = atof(argv[++i]);
        } else if (strcmp(argv[i], "--outname") == 0) {
            strncpy(params->outname, argv[++i], 256);
        } else {
            printf("Error: invalid option: %s\n", argv[i]);
            goto fail;
        }
    }

    return params;

    fail:
    free(params);
    exit(1);
}

/* Function to output to a file an AmpPhaseFrequencySeries; each mode will be output to a separate file using this function */
static int Write_CAmpPhaseFrequencySeries(FILE* f, CAmpPhaseFrequencySeries* freqseries) {
    gsl_vector* freq = freqseries->freq;
    gsl_vector* amp_real = freqseries->amp_real;
    gsl_vector* amp_imag = freqseries->amp_imag;
    gsl_vector* phase = freqseries->phase;

    int len = (int) freq->size;
    /*Here, we could add a check on the length of the gsl_vectors*/

    fprintf(f, "# f amp_re amp_im phase\n");
    for (int i=0; i<len; i++) {
      fprintf(f, "%.16e %.16e %.16e %.16e\n", gsl_vector_get(freq, i), gsl_vector_get(amp_real, i), gsl_vector_get(amp_imag, i), gsl_vector_get(phase, i));
    }

    return 0;
}

/*
 * Swhitenoise
 * a simple noise function
 */
double Swhitenoise(double x){
  return 1;
};

/*
 * main
 */
int main (int argc , char **argv) {
    FILE* f;
    ListmodesCAmpPhaseFrequencySeries* listROM = NULL;
    ListmodesCAmpPhaseFrequencySeries* listA = NULL;
    ListmodesCAmpPhaseFrequencySeries* listE = NULL;
    ListmodesCAmpPhaseFrequencySeries* listT = NULL;
    //ROMParams* params;
    LISAParams* params;

    /* parse commandline */
    //params = parse_args_ROM(argc, argv);
    params = parse_args_LISA(argc, argv);

    /* Generate the waveform with the ROM */
    SimEOBNRv2HMROM(&listROM, params->nbmode, params->tRef, params->phiRef, params->fRef, params->m1, params->m2, params->distance);
    /* Process the waveform through the LISA response */
    LISASimFDResponseTDI(&listROM, &listA, &listE, &listT, params->inclination, params->lambda, params->beta, params->psi);

    ListmodesCAmpPhaseFrequencySeries* listelementA;
    ListmodesCAmpPhaseFrequencySeries* listelementE;
    ListmodesCAmpPhaseFrequencySeries* listelementT;
    int l,m;
    int status = 0;
    /* Loop over the modes */
    for(int i=0; i<params->nbmode; i++){
      l = listmode[i][0];
      m = listmode[i][1];
      listelementA = ListmodesCAmpPhaseFrequencySeries_GetMode(listA, l, m);
      listelementE = ListmodesCAmpPhaseFrequencySeries_GetMode(listE, l, m);
      listelementT = ListmodesCAmpPhaseFrequencySeries_GetMode(listT, l, m);
      /* Write files - suffix _A,E,T_lm.dat imposed to each file name, attached to the string given by outname */
      /* If outname is still the default empty string, we do not output and skip this stage */
      if(!(strcmp(params->outname, "") == 0)){
	char *filenameA = malloc(strlen(params->outname)+64);
	char *filenameE = malloc(strlen(params->outname)+64);
	char *filenameT = malloc(strlen(params->outname)+64);
	sprintf(filenameA, "%s%s%d%d%s", params->outname, "_A_", l, m, ".dat");
	sprintf(filenameE, "%s%s%d%d%s", params->outname, "_E_", l, m, ".dat");
	sprintf(filenameT, "%s%s%d%d%s", params->outname, "_T_", l, m, ".dat");
	f = fopen(filenameA, "w");
	status |= Write_CAmpPhaseFrequencySeries(f, listelementA->freqseries);
	fclose(f);
	f = fopen(filenameE, "w");
	status |= Write_CAmpPhaseFrequencySeries(f, listelementE->freqseries);
	fclose(f);
	f = fopen(filenameT, "w");
	status |= Write_CAmpPhaseFrequencySeries(f, listelementT->freqseries);
	fclose(f);
	if (status) goto fail;
      }

      /* Example SNR calculation, taking the A observable */
      double *f1=listelementA->freqseries->freq->data;
      int n1=listelementA->freqseries->freq->size;
      double *s1Ar=listelementA->freqseries->amp_real->data;
      double *s1Ai=listelementA->freqseries->amp_imag->data;
      double *s1p=listelementA->freqseries->phase->data;
      double *f2=listelementA->freqseries->freq->data;
      int n2=listelementA->freqseries->freq->size;
      double *s2Ar=listelementA->freqseries->amp_real->data;
      double *s2Ai=listelementA->freqseries->amp_imag->data;
      double *s2p=listelementA->freqseries->phase->data;

      printf("n1: %d\n", n1);
      printf("s1Ar[299]: %g\n", s1Ar[299]);

      double start=((double)clock())/CLOCKS_PER_SEC;
      //JGB: Set these somehwere else where it makes sense.  f_min or f_max <= 0 means use intersection of signal domains.
      double f_min=-1.0;
      double f_max=-1.0;
      double rho2= wip_phase (f1, n1, f2, n2, s1Ar, s1Ai, s1p, s2Ar, s2Ai, s2p, Swhitenoise, 1.0, f_min, f_max);
      double end=((double)clock())/CLOCKS_PER_SEC;
      printf( "SNR2 = %g,   SNR time = %g\n", rho2, end-start);
      printf( "SNR = %g,   SNR time = %g\n",sqrt(rho2), end-start);

    }

    /* clean up */
    free(params);
    ListmodesCAmpPhaseFrequencySeries_Destroy(listROM);
    ListmodesCAmpPhaseFrequencySeries_Destroy(listA);
    ListmodesCAmpPhaseFrequencySeries_Destroy(listE);
    ListmodesCAmpPhaseFrequencySeries_Destroy(listT);
    return 0;

    fail:
    free(params);
    ListmodesCAmpPhaseFrequencySeries_Destroy(listROM);
    ListmodesCAmpPhaseFrequencySeries_Destroy(listA);
    ListmodesCAmpPhaseFrequencySeries_Destroy(listE);
    ListmodesCAmpPhaseFrequencySeries_Destroy(listT);
    return 1;
}


#ifndef _fim_h_included_
#define _fim_h_included_

#include <assert.h>
#include <stdlib.h>
#include <algorithm>
#include <cctype>
#include <fstream>
#include <functional>
#include <iostream>
#include <map>
#include <numeric>
#include <set>
#include <sstream>
#include <vector>
#include <sstream>
#include <string>
#include <boost/config.hpp>
#include <boost/foreach.hpp>
#include <boost/random.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/matrix.hpp>

#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>

#include <boost/numeric/ublas/io.hpp>

#include <jsc/bioinfo/gene_anno.hpp>
#include <jsc/util/log.hpp>

#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

#include "splicing_graph.h"
#include "read.h"
#include "linalg.h"
#include "accessible_read_starts.h"

using namespace std;
using namespace boost;
using namespace boost::numeric;
using namespace jsc::bioinfo;
using namespace jsc::util;

class FIM {
	public:
		FIM () {
			ofim_call_count = 0;
		}

		double get_ofim_call_count() const {
			return ofim_call_count;
		};

		static double estimate_mle_variance(
				ublas::matrix<double> const & I) {
			return estimate_mle_variance_by_diag(I);
		}

		static double estimate_mle_variance_by_diag(
				ublas::matrix<double> const & I) {
			unsigned long K = I.size1() + 1;
			double sum = 0;
			for (unsigned long p = 0; p < K - 1; ++p) {
				sum += 1.0 / I(p,p);
			}
			return sum;
		}

		static double estimate_mle_variance_by_inv(
				ublas::matrix<double> const & I) {
			assert(I.size1() == I.size2());
			L_(debug) << I;
			ublas::matrix<double> inv =
				ublas::zero_matrix<double>(I.size1(), I.size1());
			linalg::invert_matrix(I, inv);

			L_(debug) << inv;

			unsigned long K = I.size1() + 1;
			double sum = 0;
			for (unsigned long p = 0; p < K - 1; ++p) {
				for (unsigned long q = 0; q < K - 1; ++q) {
					sum += inv(p,q);
				}
				sum += inv(p,p);
			}

			return sum;
		}

		static void print_fim_for_R(ostream & os, ublas::matrix<double> I) {
			unsigned long K = I.size1() + 1;
			for (unsigned long p = 0; p < K - 1; ++p) {
				for (unsigned long q = 0; q < K - 1; ++q) {
					os << I(p,q) << "\t";
				}
				os << endl;
			}
		}

		static void print_fim_diag_for_R(ostream & os, ublas::matrix<double> I) {
			unsigned long K = I.size1() + 1;
			os << "(";
			for (unsigned long p = 0; p < K - 1; ++p) {
				os << I(p,p) << ",";
			}
			os << ")";
			os << endl;
		}

		ublas::matrix<double> bruteforce_fim(shared_ptr<Read> readp,
				vector<double> const & iso_probs) {
			L_(debug) << "FIM::bruteforce_fim";
			reset_ofim_call_count();
			shared_ptr<AccessibleReadStarts> ars = readp->get_ARS();
			unsigned long K = ars->get_num_isoforms();
			ublas::matrix<double> I = ublas::zero_matrix<double>(K-1, K-1);

			for (unsigned long k = 0; k < K; ++k) {
				if (iso_probs[k] == 0) {
					continue;
				}
				ublas::matrix<int> sign;
				L_(debug) << "Possible Isoform #" << k
					<< ": Generating reads with starts from 0 to "
					<< ars->get_iso_ARS_total_length(k)
					<< " in ARS";
				for (unsigned long a = 0;
						a < ars->get_iso_ARS_total_length(k);
						++a) {
					// generate a read starting from a in ARS
					readp->generate_read(k, ars->ARStart2IsoStart(k, a));
					double G = readp->prob_generated_by_iso(k);
					double log_scaler = log(G) + log(iso_probs[k]);
					ublas::matrix<double> m = ofim(readp, iso_probs, sign);
					for (unsigned long p = 0; p < K - 1; ++p) {
						for (unsigned long q = 0; q < K - 1; ++q) {
							if (sign(p,q) > 0) {
								m(p,q) = exp(m(p,q) + log_scaler);
							} else if (sign(p,q) < 0) {
								m(p,q) = - exp(m(p,q) + log_scaler);
							} else {
								// m(p,q) must be zero, no need to perform scaling
							}
							assert(!isinf(m(p,q)) && !isnan(m(p,q)));
						}
					}
					I += m;
				}
			}
			return I;
		}

		ublas::matrix<double> fast_fim(shared_ptr<Read> readp,
				vector<double> const & iso_probs) {
			L_(debug) << "FIM::fast_fim";
			reset_ofim_call_count();
			shared_ptr<AccessibleReadStarts> ars = readp->get_ARS();
			unsigned long K = ars->get_num_isoforms();
			ublas::matrix<double> I = ublas::zero_matrix<double>(K-1, K-1);
			for (unsigned long k = 0; k < K; ++k) {
				if (iso_probs[k] == 0) {
					continue;
				}
				ublas::matrix<int> sign;
				L_(debug) << "Possible Isoform #" << k
					<< ": Generating reads with starts from 0 to "
					<< ars->get_iso_ARS_total_length(k)
					<< " in ARS";

				unsigned long a = 0;
				while (a < ars->get_iso_ARS_total_length(k)) {
					// generate a read starting from a
					readp->generate_read(k, ars->ARStart2IsoStart(k, a));

					unsigned long N_overlapping_exons = readp->exon_indices.size();
					unsigned long N_eq_samples = min(
							ars->exon_lengths[readp->exon_indices[0]]
							- readp->start_at_first_exon,
							ars->exon_lengths[readp->exon_indices[N_overlapping_exons - 1]]
							- readp->end_at_last_exon + 1);

					while ( (a + N_eq_samples - 1 >= ars->get_iso_ARS_total_length(k)) || (ars->ARStart2IsoStart(k, a + N_eq_samples - 1) != (ars->ARStart2IsoStart(k, a) + N_eq_samples - 1)) ){
						N_eq_samples--;
					}

					double G = readp->prob_generated_by_iso(k);
					ublas::matrix<double> m = ofim(readp, iso_probs, sign);
					bool zero_scaler = ((G == 0) || (iso_probs[k] == 0));
					double log_scaler = log(N_eq_samples)
						+ log(G)
						+ log(iso_probs[k]);
					for (unsigned long p = 0; p < K - 1; ++p) {
						for (unsigned long q = 0; q < K - 1; ++q) {
							if (zero_scaler) {
								m(p,q) = 0;
							} else if (sign(p,q) > 0) {
								m(p,q) = exp(m(p,q) + log_scaler);
							} else if (sign(p,q) < 0) {
								m(p,q) = - exp(m(p,q) + log_scaler);
							} else {
								// m(p,q) must be zero, no need to perform scaling
							}
							assert(!isinf(m(p,q)) && !isnan(m(p,q)));
						}
					}
					I += m;
					a += N_eq_samples;
				}
			}
			return I;
		}

		ublas::matrix<double> faster_fim(shared_ptr<Read> readp,
				vector<double> const & iso_probs) {
			L_(debug) << "FIM::faster_fim";
			reset_ofim_call_count();
			shared_ptr<AccessibleReadStarts> ars = readp->get_ARS();
			unsigned long K = ars->get_num_isoforms();
			ublas::matrix<double> I = ublas::zero_matrix<double>(K-1, K-1);
			vector<interval_list<long > > covered_sample_starts;
			for (unsigned long k = 0; k < K; ++k) {
				interval_list<long> il;
				covered_sample_starts.push_back(il);
			}
			for (unsigned long k = 0; k < K; ++k) {
				if (iso_probs[k] == 0) {
					continue;
				}

				unsigned long a = covered_sample_starts[k].find_min_uncovered(0);
				L_(debug) << "Possible Isoform #" << k
					<< ": Generating reads with starts from " << a << " to "
					<< ars->get_iso_ARS_total_length(k)
					<< " in ARS";
				while (a < ars->get_iso_ARS_total_length(k)) {
					// generate a read starting from a
					readp->generate_read(k, ars->ARStart2IsoStart(k, a));

					unsigned long N_overlapping_exons = readp->exon_indices.size();
					unsigned long N_eq_samples = min(
							ars->exon_lengths[readp->exon_indices[0]]
							- readp->start_at_first_exon,
							ars->exon_lengths[readp->exon_indices[N_overlapping_exons - 1]]
							- readp->end_at_last_exon + 1);
					while ( (a + N_eq_samples - 1 >= ars->get_iso_ARS_total_length(k)) || (ars->ARStart2IsoStart(k, a + N_eq_samples - 1) != (ars->ARStart2IsoStart(k, a) + N_eq_samples - 1)) ){
						N_eq_samples--;
					}

					ublas::matrix<int> sign;
					ublas::matrix<double> m = ofim(readp, iso_probs, sign);

					double G = readp->prob_generated_by_iso(k);
					double scaler = G * N_eq_samples * iso_probs[k];

					if (covered_sample_starts[k].compute_overlap(a, a + N_eq_samples)
							> 0) {
						L_(debug2) << "Cov for iso # " << k
							<< ": "	<< covered_sample_starts[k];
						L_(error) << "Overlap > 0!";
					}
					covered_sample_starts[k].add_interval(a, a + N_eq_samples);
					L_(debug2) << "Adding [" << a
						<< "," << (a + N_eq_samples)
						<< ") to possible isoform #" << k;
					for (unsigned long j = 0; j < K; ++j) {
						if (j == k) continue;
						long first_matching_idx =
							is_connected_exons_compatible_with_isoform(
									readp->exon_indices,
									ars->iso_exon_indices[j]);
						if (first_matching_idx >= 0) {
							double G_j = readp->prob_generated_by_iso(j);
							scaler += G_j * N_eq_samples * iso_probs[j];
							unsigned long a2 = readp->start_at_first_exon
								+ ars->iso_exon_total_lengths[j][first_matching_idx]
								- ars->exon_lengths[readp->exon_indices[0]];
							a2 = ars->IsoStart2ARStart(j, a2);

							if (covered_sample_starts[j].compute_overlap(a2, a2 + N_eq_samples)
									> 0) {
								L_(debug2) << "Cov for iso # " << j
									<< ": "	<< covered_sample_starts[j];
								L_(error) << "Overlap > 0!";
							}
							covered_sample_starts[j].add_interval(a2, a2 + N_eq_samples);
							L_(debug2) << "Adding [" << a2
								<< "," << (a2 + N_eq_samples)
								<< ") to possible isoform #" << j;
						}
					}

					double log_scaler = log(scaler);
					for (unsigned long p = 0; p < K - 1; ++p) {
						for (unsigned long q = 0; q < K - 1; ++q) {
							if (sign(p,q) > 0) {
								m(p,q) = exp(m(p,q) + log_scaler);
							} else if (sign(p,q) < 0) {
								m(p,q) = - exp(m(p,q) + log_scaler);
							} else {
								// m(p,q) must be zero, no need to perform scaling
							}
							assert(!isinf(m(p,q)) && !isnan(m(p,q)));
						}
					}
					I += m;
					a = covered_sample_starts[k].find_min_uncovered(0);
				}
			}
			return I;
		}
	protected:
		void reset_ofim_call_count() {
			ofim_call_count = 0;
		};

		// \mathfrak{I}_{s}^{(m)}(\Theta)
		ublas::matrix<double> ofim(shared_ptr<const Read> readp,
				vector<double> const & iso_probs,
				ublas::matrix<int> & sign) {
			++ofim_call_count;
			shared_ptr<AccessibleReadStarts> ars = readp->get_ARS();
			unsigned long K = ars->get_num_isoforms();
			ublas::matrix<double> m = ublas::zero_matrix<double>(K-1,K-1);
			sign = ublas::zero_matrix<int>(K-1, K-1);
			vector<double> v_delta_G(K, 0);
			for (unsigned long k = 0; k < K; ++k) {
				if (readp->is_compatible_with_iso(k)) {
					v_delta_G[k] = readp->prob_generated_by_iso(k);
				}
			}

			double sum = 0;
			for (unsigned long k = 0; k < K; ++k) {
				if (v_delta_G[k] > 0 && iso_probs[k] > 0) {
					sum += iso_probs[k] * v_delta_G[k];
				}
			}

			assert(sum > 0);

			double log_prod = log(sum) + log(sum);
			for (unsigned long p = 0; p < K - 1; ++p) {
				for (unsigned long q = 0; q < K - 1; ++q) {
					m(p,q) = (v_delta_G[p] - v_delta_G[K - 1])*
						(v_delta_G[q] - v_delta_G[K - 1]);
					if (m(p,q) != 0) {
						sign(p,q) = m(p,q) > 0 ? 1 : -1;
						if (m(p,q) < 0) {
							m(p,q) = - m(p,q);
						}
						m(p,q) = log(m(p,q)) - log_prod;
					}
				}
			}

			if (readp->exon_indices[0] == 0
					&& readp->start_at_first_exon == 2730) {
				L_(info) << *readp << " " << m(0,0) << " " << sign(0,0);
			}

			return m;
		}

	private:
		double ofim_call_count;
};

#endif


#include "std_includes.h"

#ifndef MATRIX_H
#define MATRIX_H

#ifdef CRANIUM_USE_MKL
#define CRANIUM_USE_BLAS
#include "mkl.h"
#else
#ifdef CRANIUM_USE_CBLAS
#define CRANIUM_USE_BLAS
#include <cblas.h>
#endif
#endif

// represents user-supplied training data
typedef struct DataSet_ {
    size_t rows;
    size_t cols;
    float** data;
} DataSet;

// represents a matrix of data in row-major order
typedef struct Matrix_ {
    size_t rows;
    size_t cols;
    float* data;
} Matrix;

#ifndef RAND_MOD
#define RAND_MOD 2147483647
#endif

// create dataset given user data
static DataSet* createDataSet(size_t rows, size_t cols, float** data);

// uses memory of the original data to split dataset into batches
static DataSet** createBatches(DataSet* allData, int numBatches);

// split a dataset into row matrices
static Matrix** splitRows(DataSet* dataset);

// destroy dataset
static void destroyDataSet(DataSet* dataset);

// convert dataset to matrix
static Matrix* dataSetToMatrix(DataSet* dataset);

// creates a matrix given data
static Matrix* createMatrix(size_t rows, size_t cols, float* data);

// creates a matrix zeroed out
static Matrix* createMatrixZeroes(size_t rows, size_t cols);

// get an element of a matrix
static float getMatrix(Matrix* mat, size_t row, size_t col);

// set an element of a matrix
static void setMatrix(Matrix* mat, size_t row, size_t col, float val);

// sets the values in $to equal to values in $from
static void copyValuesInto(Matrix* from, Matrix* to);

// prints the entries of a matrix
static void printMatrix(Matrix* input);

// sets each entry in matrix to 0
static void zeroMatrix(Matrix* orig);

// returns transpose of matrix
static Matrix* transpose(Matrix* orig);

// transposes matrix and places data into $origT
static void transposeInto(Matrix* orig, Matrix* origT);

// adds two matrices and returns result
static Matrix* add(Matrix* A, Matrix* b);

// adds $from to $to and places result in $to
static void addTo(Matrix* from, Matrix* to);

// adds $B, a row vector, to each row of $A
static Matrix* addToEachRow(Matrix* A, Matrix* B);

// multiplies every element of $orig by $C
static void scalarMultiply(Matrix* orig, float c);

// multiplies $A and $B (ordering: AB) and returns product matrix
static Matrix* multiply(Matrix* A, Matrix* B);

// multiplies $A and $B (ordering: AB) and places values into $into
static void multiplyInto(Matrix* A, Matrix* B, Matrix* into);

// element-wise multiplcation
static Matrix* hadamard(Matrix* A, Matrix* B);

// places values of hadamard product of $A and $B into $into
static void hadamardInto(Matrix* A, Matrix* B, Matrix* into);

// returns a shallow copy of input matrix
static Matrix* copy(Matrix* orig);

// returns 1 if matrices are equal, 0 otherwise
static int equals(Matrix* A, Matrix* B);

// frees a matrix and its data
static void destroyMatrix(Matrix* matrix);

/*
    Begin functions.
*/

static DataSet* createDataSet(size_t rows, size_t cols, float** data){
    DataSet* dataset = (DataSet*)malloc(sizeof(DataSet));
    dataset->rows = rows;
    dataset->cols = cols;
    dataset->data = data;
    return dataset;
}

DataSet** createBatches(DataSet* allData, int numBatches){
    DataSet** batches = (DataSet**)malloc(sizeof(DataSet*) * numBatches);
    int remainder = allData->rows % numBatches;
    int i;
    int curRow = 0;
    for (i = 0; i < numBatches; i++){
        size_t batchSize = allData->rows / numBatches;
        if (remainder-- > 0){
            batchSize++;
        }
        batches[i] = createDataSet(batchSize, allData->cols, allData->data + curRow);
        curRow += batchSize;
    }
    return batches;
}

static Matrix** splitRows(DataSet* dataset){
    Matrix** rows = (Matrix**)malloc(sizeof(Matrix*) * dataset->rows);
    int i;
    for (i = 0; i < dataset->rows; i++){
        rows[i] = createMatrix(1, dataset->cols, dataset->data[i]);
    }
    return rows;
}

void shuffleTogether(DataSet* A, DataSet* B, std::minstd_rand &gen){
    assert(A->rows == B->rows);
    int i;
    for (i = 0; i < A->rows - 1; i++){
        size_t j = i + gen() / (RAND_MOD / (A->rows - i) + 1);
        float* tmpA = A->data[j];
        A->data[j] = A->data[i];
        A->data[i] = tmpA;
        float* tmpB = B->data[j];
        B->data[j] = B->data[i];
        B->data[i] = tmpB;
    }
}

static void destroyDataSet(DataSet* dataset){
    int i;
    for (i = 0; i < dataset->rows; i++){
        free(dataset->data[i]);
    }
    free(dataset->data);
    free(dataset);
}

static Matrix* dataSetToMatrix(DataSet* dataset){
    Matrix* convert = (Matrix*)malloc(sizeof(Matrix));
    convert->rows = dataset->rows;
    convert->cols = dataset->cols;
    convert->data = (float*)malloc(sizeof(float) * dataset->rows * dataset->cols);
    int i, j;
    for (i = 0; i < dataset->rows; i++){
        for (j = 0; j < dataset->cols; j++){
            setMatrix(convert, i, j, dataset->data[i][j]);
        }
    }
    return convert;
}

Matrix* createMatrix(size_t rows, size_t cols, float* data){
    assert(rows > 0 && cols > 0);
    Matrix* matrix = (Matrix*)malloc(sizeof(Matrix));
    matrix->rows = rows;
    matrix->cols = cols;
    matrix->data = data;
    return matrix;
}

Matrix* createMatrixZeroes(size_t rows, size_t cols){
    assert(rows > 0 && cols > 0);
    Matrix* matrix = (Matrix*)malloc(sizeof(Matrix));
    matrix->rows = rows;
    matrix->cols = cols;
    float* data = (float*)calloc(rows * cols, sizeof(float));
    matrix->data = data;
    return matrix;
}

inline float getMatrix(Matrix* mat, size_t row, size_t col){
    return mat->data[row * mat->cols + col];
}

inline void setMatrix(Matrix* mat, size_t row, size_t col, float val){
    mat->data[row * mat->cols + col] = val;
}

void copyValuesInto(Matrix* from, Matrix* to){
    assert(from->rows == to->rows && from->cols == to->cols);
    memcpy(to->data, from->data, sizeof(float) * to->rows * to->cols);
}

void printMatrix(Matrix* input){
    int i, j;
    for (i = 0; i < input->rows; i++){
        printf("\n");
        for (j = 0; j < input->cols; j++){
            printf("%.2f ", getMatrix(input, i, j));
        }
    }
    printf("\n");
}

void zeroMatrix(Matrix* orig){
    memset(orig->data, 0, orig->rows * orig->cols * sizeof(float));
}

Matrix* transpose(Matrix* orig){
    float* data = (float*)malloc(sizeof(float) * orig->rows * orig->cols);
    Matrix* transpose = createMatrix(orig->cols, orig->rows, data);
    int i, j;
    for (i = 0; i < orig->rows; i++){
        for (j = 0; j < orig->cols; j++){
            setMatrix(transpose, i, j, getMatrix(orig, i, j));
        }
    }
    return transpose;
}

void transposeInto(Matrix* orig, Matrix* origT){
    assert(orig->rows == origT->cols && orig->cols == origT->rows);
    int i, j;
    for (i = 0; i < orig->rows; i++){
        for (j = 0; j < orig->cols; j++){
            setMatrix(origT, j, i, getMatrix(orig, i, j));
        }
    }
}

Matrix* add(Matrix* A, Matrix* B){
    assert(A->rows == B->rows && A->cols == B->cols);
    float* data = (float*)malloc(sizeof(float) * A->rows * B->rows);
    Matrix* result = createMatrix(A->rows, A->cols, data);
    int i, j;
    for (i = 0; i < A->rows; i++){
        for (j = 0; j < A->cols; j++){
            setMatrix(result, i, j, getMatrix(B, i, j) + getMatrix(A, i, j));
        }
    }
    return result;
}

void matrix_mul_add(Matrix * to, float a, Matrix * x, float b, bool square = false) {
	assert(to->rows == x->rows && to->cols == x->cols);
	int i, j;
	for (i = 0; i < to->rows; i++) {
		if (!square) {
			for (j = 0; j < to->cols; j++)
				setMatrix(to, i, j, getMatrix(to, i, j) * a + getMatrix(x, i, j) * b);
		}
		else for (j = 0; j < to->cols; j++) {
			float val = getMatrix(x, i, j);
			setMatrix(to, i, j, getMatrix(to, i, j) * a + val * val * b);
		}
	}
}

void matrix_replace(Matrix * to, Matrix * x, float b) {
	assert(to->rows == x->rows && to->cols == x->cols);
	int i, j;
	for (i = 0; i < to->rows; i++) {
		for (j = 0; j < to->cols; j++)
			setMatrix(to, i, j, getMatrix(x, i, j) * b);
	}
}

void adam_update(Matrix * to, Matrix * m, Matrix * v, float a) {
	assert(to->rows == m->rows && to->cols == m->cols);
	assert(to->rows == v->rows && to->cols == v->cols);
	int i, j;
	for (i = 0; i < to->rows; i++) {
		for (j = 0; j < to->cols; j++)
			setMatrix(to, i, j, (getMatrix(m, i, j) * a) / (sqrt(getMatrix(v, i, j)) + 0.000000001));
	}
}

void addTo(Matrix* from, Matrix* to){
    assert(from->rows == to->rows && from->cols == to->cols);
    int i, j;
    for (i = 0; i < from->rows; i++){
        for (j = 0; j < from->cols; j++){
            setMatrix(to, i, j, getMatrix(from, i, j) + getMatrix(to, i, j));
        }
    }
}

// add B to each row of A
Matrix* addToEachRow(Matrix* A, Matrix* B){
    assert(A->cols == B->cols && B->rows == 1);
    float* data = (float*)malloc(sizeof(float) * A->rows * A->cols);
    Matrix* result = createMatrix(A->rows, A->cols, data);
    int i, j;
    for (i = 0; i < A->rows; i++){
        for (j = 0; j < A->cols; j++){
            setMatrix(result, i, j, getMatrix(A, i, j) + getMatrix(B, 0, j));
        }
    }
    return result;
}

void scalarMultiply(Matrix* orig, float c){
    int i, j;
    for (i = 0; i < orig->rows; i++){
        for (j = 0; j < orig->cols; j++){
            setMatrix(orig, i, j, getMatrix(orig, i, j) * c);
        }
    }
}

Matrix* multiply(Matrix* A, Matrix* B){
    assert(A->cols == B->rows);
    float* data = (float*)malloc(sizeof(float) * A->rows * B->cols);
    Matrix* result = createMatrix(A->rows, B->cols, data);
#ifdef CRANIUM_USE_BLAS
    zeroMatrix(result);
    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, A->rows, B->cols
    , A->cols, 1, A->data, A->cols, B->data, B->cols, 1, result->data, result->cols);
    return result;
#endif
    int i, j;
    for (i = 0; i < A->rows; i++){
        for (j = 0; j < B->cols; j++){
            float sum = 0;
            int k;
            for (k = 0; k < B->rows; k++){
                sum += getMatrix(A, i, k) * getMatrix(B, k, j);
            }
            setMatrix(result, i, j, sum);
        }
    }
    return result;
}

void multiplyInto(Matrix* A, Matrix* B, Matrix* into){
#ifdef CRANIUM_USE_BLAS
    zeroMatrix(into);
    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, A->rows, B->cols
    , A->cols, 1, A->data, A->cols, B->data, B->cols, 1, into->data, into->cols);
    return;
#endif
    assert(A->cols == B->rows);
    assert(A->rows == into->rows && B->cols == into->cols);
    int i, j;
    for (i = 0; i < A->rows; i++){
        for (j = 0; j < B->cols; j++){
            float sum = 0;
            int k;
            for (k = 0; k < B->rows; k++){
                sum += getMatrix(A, i, k) * getMatrix(B, k, j);
            }
            setMatrix(into, i, j, sum);
        }
    }
}

Matrix* hadamard(Matrix* A, Matrix* B){
    assert(A->rows == B->rows && A->cols == B->cols);
    float* data = (float*)malloc(sizeof(float) * A->rows * A->cols);
    Matrix* result = createMatrix(A->rows, A->cols, data);
    int i, j;
    for (i = 0; i < A->rows; i++){
        for (j = 0; j < A->cols; j++){
            setMatrix(result, i, j, getMatrix(A, i, j) * getMatrix(B, i, j));
        }
    }
    return result;
}

void hadamardInto(Matrix* A, Matrix* B, Matrix* into){
    assert(A->rows == B->rows && A->cols == B->cols);
    assert(A->rows == into->rows && A->cols == into->cols);
    int i, j;
    for (i = 0; i < A->rows; i++){
        for (j = 0; j < A->cols; j++){
            setMatrix(into, i, j, getMatrix(A, i, j) * getMatrix(B, i, j));
        }
    }
}

Matrix* copy(Matrix* orig){
    float* data = (float*)malloc(sizeof(float) * orig->rows * orig->cols);
    memcpy(data, orig->data, sizeof(float) * orig->cols * orig->rows);
    return createMatrix(orig->rows, orig->cols, data);
}

int equals(Matrix* A, Matrix* B){
    if (A->rows != B->rows){
        return 0;
    }
    if (A->cols != B->cols){
        return 0;
    }
    int i, j;
    for (i = 0; i < A->rows; i++){
        for (j = 0; j < A->cols; j++){
            if (getMatrix(A, i, j) != getMatrix(B, i, j)){
                return 0;
            }
        }
    }
    return 1;
}

void destroyMatrix(Matrix* matrix){
    free(matrix->data);
    free(matrix);
}

#endif


/* multifit/covar.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Brian Gough
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or (at
 * your option) any later version.
 *
 * This program 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.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#include <config.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_multifit_nlin.h>
#include <gsl/gsl_blas.h>

int
gsl_multifit_gradient (const gsl_matrix * J, const gsl_vector * f,
                       gsl_vector * g)
{
  int status = gsl_blas_dgemv (CblasTrans, 1.0, J, f, 0.0, g);
  return status;
}