#include "cpu/graph_proc.h"

#include <set>
#include <vector>
#include <numeric> //std::iota
#include <algorithm>
#include <random>
#include <iostream>

#include <Eigen/Dense>

using std::vector;

namespace graph_proc {

    py::array_t<bool> erode_mesh(const py::array_t<float>& vertexPositions, const py::array_t<int>& faceIndices, int nIterations, int minNeighbors) {
        int nVertices = vertexPositions.shape(0);
        int nFaces = faceIndices.shape(0);

        // Init output
        py::array_t<bool> nonErodedVertices = py::array_t<bool>({ nVertices, 1 });
        std::vector<bool> nonErodedVerticesVec(nVertices, false);

        // Init list of eroded face indices with original list
        std::vector<Eigen::Vector3i> erodedFaceIndicesVec;
        erodedFaceIndicesVec.reserve(nFaces);
        for (int FaceIdx = 0; FaceIdx < nFaces; ++FaceIdx) {
            Eigen::Vector3i face(*faceIndices.data(FaceIdx, 0), *faceIndices.data(FaceIdx, 1), *faceIndices.data(FaceIdx, 2));
            erodedFaceIndicesVec.push_back(face);
        }

        // Erode mesh for a total of nIterations
        for (int i = 0; i < nIterations; i++) {
            nFaces = erodedFaceIndicesVec.size();

            // We compute the number of neighboring vertices for each vertex.
            vector<int> numNeighbors(nVertices, 0);
            for (int i = 0; i < nFaces; i++) {
                const auto& face = erodedFaceIndicesVec[i];
                numNeighbors[face[0]] += 1;
                numNeighbors[face[1]] += 1;
                numNeighbors[face[2]] += 1;
            }

            std::vector<Eigen::Vector3i> tmp;
            tmp.reserve(nFaces);

            for (int i = 0; i < nFaces; i++) {
                const auto& face = erodedFaceIndicesVec[i];
                if (numNeighbors[face[0]] >= minNeighbors && numNeighbors[face[1]] >= minNeighbors && numNeighbors[face[2]] >= minNeighbors) {
                    tmp.push_back(face);
                }
            }

            // We kill the faces with border vertices.
            erodedFaceIndicesVec.clear();
            erodedFaceIndicesVec = std::move(tmp);
        }

        // Mark non isolated vertices as not eroded.
        nFaces = erodedFaceIndicesVec.size();

        for (int i = 0; i < nFaces; i++) {
            const auto& face = erodedFaceIndicesVec[i];
            nonErodedVerticesVec[face[0]] = true;
            nonErodedVerticesVec[face[1]] = true;
            nonErodedVerticesVec[face[2]] = true;
        }

        // Store into python array
        for (int i = 0; i < nVertices; i++) {
            *nonErodedVertices.mutable_data(i, 0) = nonErodedVerticesVec[i];
        }

        return nonErodedVertices;
    }

    int sample_nodes(
        const py::array_t<float>& vertexPositions, const py::array_t<bool>& nonErodedVertices,
        py::array_t<float>& nodePositions, py::array_t<int>& nodeIndices,
        float nodeCoverage,
        const bool useOnlyNonErodedIndices=true,
        const bool randomShuffle=true
    ) {
        // assert(vertexPositions.ndim() == 2);

        float nodeCoverage2 = nodeCoverage * nodeCoverage;
        int nVertices = vertexPositions.shape(0);
        // assert(vertexPositions.shape(1) == 3);
        // assert(nodePositions.shape(0) == nVertices);
        // assert(nodePositions.shape(1) == 3);
        // assert(nodeIndices.shape(0) == nVertices);
        // assert(nodeIndices.shape(1) == 1);

        nodePositions.resize({ nVertices, 3 }, false);
        nodeIndices.resize({ nVertices, 1 }, false);

        // create list of shuffled indices
        std::vector<int> shuffledVertices(nVertices);
        std::iota(std::begin(shuffledVertices), std::end(shuffledVertices), 0);

        if (randomShuffle) {
            std::default_random_engine re{std::random_device{}()};
            std::shuffle(std::begin(shuffledVertices), std::end(shuffledVertices), re);
        }

        std::vector<Eigen::Vector3f> nodePositionsVec;
        for (int vertexIdx : shuffledVertices) {
        // for (int vertexIdx = 0; vertexIdx < nVertices; ++vertexIdx) {
            Eigen::Vector3f point(*vertexPositions.data(vertexIdx, 0), *vertexPositions.data(vertexIdx, 1), *vertexPositions.data(vertexIdx, 2));

            if (useOnlyNonErodedIndices && !(*nonErodedVertices.data(vertexIdx))) {
                continue;
            }

            bool bIsNode = true;
            for (int nodeIdx = 0; nodeIdx < nodePositionsVec.size(); ++nodeIdx) {
                if ((point - nodePositionsVec[nodeIdx]).squaredNorm() <= nodeCoverage2) {
                    bIsNode = false;
                    break;
                }
            }

            if (bIsNode) {
                nodePositionsVec.push_back(point);
                int newNodeIdx = nodePositionsVec.size() - 1;
                *nodePositions.mutable_data(newNodeIdx, 0) = point.x();
                *nodePositions.mutable_data(newNodeIdx, 1) = point.y();
                *nodePositions.mutable_data(newNodeIdx, 2) = point.z();
                *nodeIndices.mutable_data(newNodeIdx, 0) = vertexIdx;
            }
        }

        return nodePositionsVec.size();
    }

    /**
     * Custom comparison operator for geodesic priority queue.
     */
    struct CustomCompare {
        bool operator()(const std::pair<int, float>& left, const std::pair<int, float>& right) {
            return left.second > right.second;
        }
    };

    inline float compute_anchor_weight(const Eigen::Vector3f& pointPosition, const Eigen::Vector3f& nodePosition, float nodeCoverage) {
        return std::exp(-(nodePosition - pointPosition).squaredNorm() / (2.f * nodeCoverage * nodeCoverage));
    }

    inline float compute_anchor_weight(float dist, float nodeCoverage) {
        return std::exp(- (dist * dist) / (2.f * nodeCoverage * nodeCoverage));
    }

    void compute_edges_geodesic(
		const py::array_t<float>& vertexPositions,
		const py::array_t<bool>& validVertices,
		const py::array_t<int>& faceIndices,
		const py::array_t<int>& nodeIndices,
		const int nMaxNeighbors, const float nodeCoverage,
        py::array_t<int>& graphEdges,
        py::array_t<float>& graphEdgesWeights,
        py::array_t<float>& graphEdgesDistances,
        py::array_t<float>& nodeToVertexDistances,
        const bool allow_only_valid_vertices,
        const bool enforce_total_num_neighbors
	) {
		int nVertices = vertexPositions.shape(0);
		int nFaces = faceIndices.shape(0);
        int nNodes = nodeIndices.shape(0);

        float maxInfluence = 2.f * nodeCoverage;

        // Preprocess vertex neighbors.
		vector<std::set<int>> vertexNeighbors(nVertices);
        for (int faceIdx = 0; faceIdx < nFaces; faceIdx++) {
            for (int j = 0; j < 3; j++) {
                int v_idx = *faceIndices.data(faceIdx, j);

                for (int k = 0; k < 3; k++) {
                    int n_idx = *faceIndices.data(faceIdx, k);

                    if (v_idx == n_idx) continue;
                    vertexNeighbors[v_idx].insert(n_idx);
                }
            }
        }

		// Compute inverse vertex -> node relationship.
		vector<int> mapVertexToNode(nVertices, -1);

		for (int nodeId = 0; nodeId < nNodes; nodeId++) {
			int vertexIdx = *nodeIndices.data(nodeId);
			if (vertexIdx >= 0) {
				mapVertexToNode[vertexIdx] = nodeId;
			}
		}

        // #pragma omp parallel for
		for (int nodeId = 0; nodeId < nNodes; nodeId++) {
			// vertex queue
            std::priority_queue<
				std::pair<int, float>,
				vector<std::pair<int, float>>,
				CustomCompare
			> nextVerticesWithIds;

			std::set<int> visitedVertices;

			// Add node vertex as the first vertex to be visited.
			int nodeVertexIdx = *nodeIndices.data(nodeId);
			if (nodeVertexIdx < 0) continue;
			nextVerticesWithIds.push(std::make_pair(nodeVertexIdx, 0.f));

			// Traverse all neighbors in the monotonically increasing order.
            vector<int>   neighborNodeIds;
            vector<float> neighborNodeWeights;
            vector<float> neighborNodeDistances;
			while (!nextVerticesWithIds.empty()) {
				auto nextVertex = nextVerticesWithIds.top();
				nextVerticesWithIds.pop();

				int nextVertexIdx = nextVertex.first;
				float nextVertexDist = nextVertex.second;

				// We skip the vertex, if it was already visited before.
				if (visitedVertices.find(nextVertexIdx) != visitedVertices.end()) continue;

                if (allow_only_valid_vertices && !*validVertices.data(nextVertexIdx)) {
                    std::cout << "compute_edges_geodesic:: ufff... we shouldn't be checking out this vertex" << std::endl;
                    exit(0);
                }

				// We check if the vertex is a node.
				int nextNodeId = mapVertexToNode[nextVertexIdx];
				if (nextNodeId >= 0 && nextNodeId != nodeId) {
                    neighborNodeIds.push_back(nextNodeId);
                    neighborNodeWeights.push_back(compute_anchor_weight(nextVertexDist, nodeCoverage));
                    neighborNodeDistances.push_back(nextVertexDist);
                    if (neighborNodeIds.size() >= nMaxNeighbors) break;
				}

				// Note down the node-vertex distance.
				*nodeToVertexDistances.mutable_data(nodeId, nextVertexIdx) = nextVertexDist;

				// We visit the vertex, and check all his neighbors.
				// We add only valid vertices under a certain distance
				visitedVertices.insert(nextVertexIdx);
				Eigen::Vector3f nextVertexPos(*vertexPositions.data(nextVertexIdx, 0), *vertexPositions.data(nextVertexIdx, 1), *vertexPositions.data(nextVertexIdx, 2));

				const auto& nextNeighbors = vertexNeighbors[nextVertexIdx];
				for (int neighborIdx : nextNeighbors) {

                    bool is_valid_vertex = *validVertices.data(neighborIdx);
                    if (allow_only_valid_vertices && !is_valid_vertex) {
                        continue;
                    }

					Eigen::Vector3f neighborVertexPos(*vertexPositions.data(neighborIdx, 0), *vertexPositions.data(neighborIdx, 1), *vertexPositions.data(neighborIdx, 2));
					float dist = nextVertexDist + (nextVertexPos - neighborVertexPos).norm();

                    if (enforce_total_num_neighbors) {
						nextVerticesWithIds.push(std::make_pair(neighborIdx, dist));
                    }
					else {
                        // std::cout << dist << " " << maxInfluence << std::endl;
                        if (dist <= maxInfluence) {
						    nextVerticesWithIds.push(std::make_pair(neighborIdx, dist));
					    }
                    }
				}
			}

            // Store the nearest neighbors.
            int nNeighbors = neighborNodeIds.size();

            float weightSum = 0.f;
            for (int i = 0; i < nNeighbors; i++) {
                *graphEdges.mutable_data(nodeId, i) = neighborNodeIds[i];
                weightSum += neighborNodeWeights[i];
            }

            // Normalize weights
            if (weightSum > 0) {
                for (int i = 0; i < nNeighbors; i++) {
                    *graphEdgesWeights.mutable_data(nodeId, i) = neighborNodeWeights[i] / weightSum;
                }
            }
            else if (nNeighbors > 0) {
                for (int i = 0; i < nNeighbors; i++) {
                    *graphEdgesWeights.mutable_data(nodeId, i) = neighborNodeWeights[i] / nNeighbors;
                }
            }

            // Store edge distance.
            for (int i = 0; i < nNeighbors; i++) {
                *graphEdgesDistances.mutable_data(nodeId, i) = neighborNodeDistances[i];
            }
		}
    }

    py::array_t<int> compute_edges_euclidean(const py::array_t<float>& nodePositions, int nMaxNeighbors) {
        int nNodes = nodePositions.shape(0);

        py::array_t<int> graphEdges = py::array_t<int>({ nNodes, nMaxNeighbors });

        // Find nearest Euclidean neighbors for each node.
        for (int nodeId = 0; nodeId < nNodes; nodeId++) {
            Eigen::Vector3f nodePos(*nodePositions.data(nodeId, 0), *nodePositions.data(nodeId, 1), *nodePositions.data(nodeId, 2));

            // Keep only the k nearest Euclidean neighbors.
            std::list<std::pair<int, float>> nearestNodesWithSquaredDistances;

            for (int neighborId = 0; neighborId < nNodes; neighborId++) {
                if (neighborId == nodeId) continue;

                Eigen::Vector3f neighborPos(*nodePositions.data(neighborId, 0), *nodePositions.data(neighborId, 1), *nodePositions.data(neighborId, 2));

                float distance2 = (nodePos - neighborPos).squaredNorm();
                bool bInserted = false;
                for (auto it = nearestNodesWithSquaredDistances.begin(); it != nearestNodesWithSquaredDistances.end(); ++it) {
                    // We insert the element at the first position where its distance is smaller than the other
                    // element's distance, which enables us to always keep a sorted list of at most k nearest
                    // neighbors.
                    if (distance2 <= it->second) {
                        it = nearestNodesWithSquaredDistances.insert(it, std::make_pair(neighborId, distance2));
                        bInserted = true;
                        break;
                    }
                }

                if (!bInserted && nearestNodesWithSquaredDistances.size() < nMaxNeighbors) {
                    nearestNodesWithSquaredDistances.emplace_back(std::make_pair(neighborId, distance2));
                }

                // We keep only the list of k nearest elements.
                if (bInserted && nearestNodesWithSquaredDistances.size() > nMaxNeighbors) {
                    nearestNodesWithSquaredDistances.pop_back();
                }
            }

            // Store nearest neighbor ids.
            int idx = 0;
            for (auto it = nearestNodesWithSquaredDistances.begin(); it != nearestNodesWithSquaredDistances.end(); ++it) {
                int neighborId = it->first;
                *graphEdges.mutable_data(nodeId, idx) = neighborId;
                idx++;
            }

            for (idx = nearestNodesWithSquaredDistances.size(); idx < nMaxNeighbors; idx++) {
                *graphEdges.mutable_data(nodeId, idx) = -1;
            }
        }

        return graphEdges;
    }

    inline int traverse_neighbors(const std::vector<std::set<int>>& node_neighbors, std::vector<int>& cluster_ids, int cluster_id, int node_id) {
        if (cluster_ids[node_id] != -1) return 0;

        std::set<int> active_node_indices;

        // Initialize with current node.
        int cluster_size = 0;
        active_node_indices.insert(node_id);

        // Process until we have no active nodes anymore.
        while (!active_node_indices.empty()) {
            int active_node_id = *active_node_indices.begin();
            active_node_indices.erase(active_node_indices.begin());

            if (cluster_ids[active_node_id] == -1) {
                cluster_ids[active_node_id] = cluster_id;
                ++cluster_size;
            }

            // Look if we need to process any of the neighbors
            for (const auto& n_idx : node_neighbors[active_node_id]) {
                if (cluster_ids[n_idx] == -1) {	// If it doesn't have a cluster yet
                    active_node_indices.insert(n_idx);
                }
            }
        }

        return cluster_size;
    }

    void node_and_edge_clean_up(const py::array_t<int>& graph_edges, py::array_t<bool>& valid_nodes_mask) {
        int num_nodes = graph_edges.shape(0);
        int max_num_neighbors = graph_edges.shape(1);

        std::list<int> removed_nodes;

        while (true) {
            int num_newly_removed_nodes = 0;

            for (int node_id = 0; node_id < num_nodes; ++node_id) {

                if (*valid_nodes_mask.data(node_id, 0) == false) {
                    // if node has been already removed, continue
                    continue;
                }

                int num_neighbors = 0;
                for (int i = 0; i < max_num_neighbors; ++i) {

                    int neighbor_id = *graph_edges.data(node_id, i);

                    // if neighboring node is -1, break, since by design 'graph_edges' has
                    // the shape [2, 3, 6, -1, -1, -1, -1, -1]
                    if (neighbor_id == -1) {
                        break;
                    }

                    // if neighboring node has been marked as invalid, continue
                    if (std::find(removed_nodes.begin(), removed_nodes.end(), neighbor_id) != removed_nodes.end()) {
                        continue;
                    }

                    ++num_neighbors;
                }

                if (num_neighbors <= 1) {
                    // remove node
                    *valid_nodes_mask.mutable_data(node_id, 0) = false;
                    removed_nodes.emplace_back(node_id);
                    // std::cout << "\tremoving node_id " << node_id << std::endl;
                    ++num_newly_removed_nodes;
                }
            }

            // std::cout << "num_newly_removed_nodes: " << num_newly_removed_nodes << std::endl;

            if (num_newly_removed_nodes == 0) {
                break;
            }
        }
    }

    std::vector<int> compute_clusters(
        const py::array_t<int> graph_edges,
        py::array_t<int> graph_clusters
    ) {
        int num_nodes = graph_edges.shape(0);
        int max_num_neighbors = graph_edges.shape(1);

        // convert graph_edges to a vector of sets
        std::vector<std::set<int>> node_neighbors(num_nodes);

        for (int node_id = 0; node_id < num_nodes; ++node_id) {
            for (int neighbor_idx = 0; neighbor_idx < max_num_neighbors; ++neighbor_idx) {

                int neighbor_id = *graph_edges.data(node_id, neighbor_idx);

                if (neighbor_id == -1) {
                    break;
                }

                node_neighbors[node_id].insert(neighbor_id);
                node_neighbors[neighbor_id].insert(node_id);
            }
        }

        std::vector<int> cluster_ids(num_nodes, -1);
        std::vector<int> clusters_size;

        int cluster_id = 0;
        for (int node_id = 0; node_id < num_nodes; ++node_id) {
            int cluster_size = traverse_neighbors(node_neighbors, cluster_ids, cluster_id, node_id);
            if (cluster_size > 0) {
                cluster_id++;
                clusters_size.push_back(cluster_size);
            }
        }

        for (int node_id = 0; node_id < num_nodes; ++node_id) {
            *graph_clusters.mutable_data(node_id, 0) = cluster_ids[node_id];
        }

        return clusters_size;
    }

    inline void compute_nearest_geodesic_nodes(
        const py::array_t<float>&  node_to_vertex_distance,
        const py::array_t<int>& valid_nodes_mask,
        const int vertex_id,
        std::vector<int>& nearest_geodesic_node_ids,
        std::vector<float>& dist_to_nearest_geodesic_nodes
    ) {
        int num_nodes = node_to_vertex_distance.shape(0);

        std::map<int, float> node_map;

        for (int n = 0; n < num_nodes; ++n) {

            // discard node if it was marked as invalid (due to not having enough neighbors)
            if (*valid_nodes_mask.data(n, 0) == false) {
                continue;
            }

            float dist = *node_to_vertex_distance.data(n, vertex_id);

            if (dist >= 0) {
                node_map.emplace(n, dist);
            }
        }

        // Sort the map by distance
        // Declaring the type of Predicate that accepts 2 pairs and return a bool
        typedef std::function<bool(std::pair<int, float>, std::pair<int, float>)> Comparator;

        // Defining a lambda function to compare two pairs. It will compare two pairs using second field
        Comparator comp_functor =
            [](std::pair<int, float> node1 ,std::pair<int, float> node2) {
                return node1.second < node2.second;
            };

        // Declaring a set that will store the pairs using above comparision logic
        std::set<std::pair<int, float>, Comparator> node_set(
            node_map.begin(), node_map.end(), comp_functor
        );

        for (auto n : node_set) {
            nearest_geodesic_node_ids.push_back(n.first);
            dist_to_nearest_geodesic_nodes.push_back(n.second);

            if (nearest_geodesic_node_ids.size() == GRAPH_K) {
                break;
            }
        }
    }

    void compute_pixel_anchors_geodesic(
        const py::array_t<float> &node_to_vertex_distance,
        const py::array_t<int> &valid_nodes_mask,
        const py::array_t<float> &vertices,
        const py::array_t<int> &vertex_pixels,
        py::array_t<int>& pixel_anchors,
        py::array_t<float>& pixel_weights,
        const int width, const int height,
        const float node_coverage
    ) {
        // Allocate graph node ids and corresponding skinning weights.
        // Initialize with invalid anchors.
        pixel_anchors.resize({ height, width, GRAPH_K }, false);
        pixel_weights.resize({ height, width, GRAPH_K }, false);

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                for (int k = 0; k < GRAPH_K; k++) {
                    *pixel_anchors.mutable_data(y, x, k) = -1;
                    *pixel_weights.mutable_data(y, x, k) = 0.f;
                }
            }
        }

        int num_vertices = vertices.shape(0);

        for (int vertex_id = 0; vertex_id < num_vertices; vertex_id++) {
            // Get corresponding pixel location
            int u = *vertex_pixels.data(vertex_id, 0);
            int v = *vertex_pixels.data(vertex_id, 1);

            // Initialize some variables
            std::vector<int> nearest_geodesic_node_ids;
            std::vector<float> dist_to_nearest_geodesic_nodes;
            std::vector<float> skinning_weights;

            nearest_geodesic_node_ids.reserve(GRAPH_K);
            dist_to_nearest_geodesic_nodes.reserve(GRAPH_K);
            skinning_weights.reserve(GRAPH_K);

            // Find closest geodesic nodes
            compute_nearest_geodesic_nodes(
                node_to_vertex_distance, valid_nodes_mask, vertex_id,
                nearest_geodesic_node_ids, dist_to_nearest_geodesic_nodes
            );

            int num_anchors = nearest_geodesic_node_ids.size();

            // Compute skinning weights.
            float weight_sum{ 0.f };
            for (int i = 0; i < num_anchors; ++i) {
                float geodesic_dist_to_node = dist_to_nearest_geodesic_nodes[i];

                float weight = compute_anchor_weight(geodesic_dist_to_node, node_coverage);
                weight_sum += weight;

                skinning_weights.push_back(weight);
            }

            // Normalize the skinning weights.
            if (weight_sum > 0) {
                for (int i = 0; i < num_anchors; i++)
                    skinning_weights[i] /= weight_sum;
            }
            else if (num_anchors > 0) {
                for (int i = 0; i < num_anchors; i++)
                    skinning_weights[i] = 1.f / num_anchors;
            }

            // Store the results.
            for (int i = 0; i < num_anchors; i++) {
                *pixel_anchors.mutable_data(v, u, i) = nearest_geodesic_node_ids[i];
                *pixel_weights.mutable_data(v, u, i) = skinning_weights[i];
            }
        }
    }

    void compute_pixel_anchors_euclidean(
        const py::array_t<float>& graphNodes,
        const py::array_t<float>& pointImage,
        float nodeCoverage,
        py::array_t<int>& pixelAnchors,
        py::array_t<float>& pixelWeights
    ) {
        int nNodes = graphNodes.shape(0);
        int width = pointImage.shape(2);
        int height = pointImage.shape(1);
        // int nChannels = pointImage.shape(0);

        // Allocate graph node ids and corresponding skinning weights.
        // Initialize with invalid anchors.
        pixelAnchors.resize({ height, width, GRAPH_K }, false);
        pixelWeights.resize({ height, width, GRAPH_K }, false);

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                for (int k = 0; k < GRAPH_K; k++) {
                    *pixelAnchors.mutable_data(y, x, k) = -1;
                    *pixelWeights.mutable_data(y, x, k) = 0.f;
                }
            }
        }

        // Compute anchors for every pixel.
        #pragma omp parallel for
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                // Query 3d pixel position.
                Eigen::Vector3f pixelPos(*pointImage.data(0, y, x), *pointImage.data(1, y, x), *pointImage.data(2, y, x));
                if (pixelPos.z() <= 0) continue;

                // Keep only the k nearest Euclidean neighbors.
                std::list<std::pair<int, float>> nearestNodesWithSquaredDistances;

                for (int nodeId = 0; nodeId < nNodes; nodeId++) {
                    Eigen::Vector3f nodePos(*graphNodes.data(nodeId, 0), *graphNodes.data(nodeId, 1), *graphNodes.data(nodeId, 2));

                    float distance2 = (pixelPos - nodePos).squaredNorm();
                    bool bInserted = false;
                    for (auto it = nearestNodesWithSquaredDistances.begin(); it != nearestNodesWithSquaredDistances.end(); ++it) {
                        // We insert the element at the first position where its distance is smaller than the other
                        // element's distance, which enables us to always keep a sorted list of at most k nearest
                        // neighbors.
                        if (distance2 <= it->second) {
                            it = nearestNodesWithSquaredDistances.insert(it, std::make_pair(nodeId, distance2));
                            bInserted = true;
                            break;
                        }
                    }

                    if (!bInserted && nearestNodesWithSquaredDistances.size() < GRAPH_K) {
                        nearestNodesWithSquaredDistances.emplace_back(std::make_pair(nodeId, distance2));
                    }

                    // We keep only the list of k nearest elements.
                    if (bInserted && nearestNodesWithSquaredDistances.size() > GRAPH_K) {
                        nearestNodesWithSquaredDistances.pop_back();
                    }
                }

                // Compute skinning weights.
                std::vector<int> nearestEuclideanNodeIds;
                nearestEuclideanNodeIds.reserve(nearestNodesWithSquaredDistances.size());

                std::vector<float> skinningWeights;
                skinningWeights.reserve(nearestNodesWithSquaredDistances.size());

                float weightSum{ 0.f };
                for (auto it = nearestNodesWithSquaredDistances.begin(); it != nearestNodesWithSquaredDistances.end(); ++it) {
                    int nodeId = it->first;

                    Eigen::Vector3f nodePos(*graphNodes.data(nodeId, 0), *graphNodes.data(nodeId, 1), *graphNodes.data(nodeId, 2));
                    float weight = compute_anchor_weight(pixelPos, nodePos, nodeCoverage);
                    weightSum += weight;

                    nearestEuclideanNodeIds.push_back(nodeId);
                    skinningWeights.push_back(weight);
                }

                // Normalize the skinning weights.
                int nAnchors = nearestEuclideanNodeIds.size();

                if (weightSum > 0) {
                    for (int i = 0; i < nAnchors; i++)	skinningWeights[i] /= weightSum;
                }
                else if (nAnchors > 0) {
                    for (int i = 0; i < nAnchors; i++)	skinningWeights[i] = 1.f / nAnchors;
                }

                // Store the results.
                for (int i = 0; i < nAnchors; i++) {
                    *pixelAnchors.mutable_data(y, x, i) = nearestEuclideanNodeIds[i];
                    *pixelWeights.mutable_data(y, x, i) = skinningWeights[i];
                }
            }
        }
    }

    void construct_regular_graph(
        const py::array_t<float>& pointImage,
        int xNodes, int yNodes,
        float edgeThreshold,
        float maxPointToNodeDistance,
        float maxDepth,
        py::array_t<float>& graphNodes,
        py::array_t<int>& graphEdges,
        py::array_t<int>& pixelAnchors,
        py::array_t<float>& pixelWeights
    ) {
        int width = pointImage.shape(2);
        int height = pointImage.shape(1);
        int nChannels = pointImage.shape(0);

        float xStep = float(width - 1) / (xNodes - 1);
        float yStep = float(height - 1) / (yNodes - 1);

        // Sample graph nodes.
        // We need to maintain the mapping from all -> valid nodes ids.
        int nNodes = xNodes * yNodes;
        std::vector<int> sampledNodeMapping(nNodes, -1);

        std::vector<Eigen::Vector3f> nodePositions;
        nodePositions.reserve(nNodes);

        int nodeId = 0;
        for (int y = 0; y < yNodes; y++) {
            for (int x = 0; x < xNodes; x++) {
                int nodeIdx = y * xNodes + x;

                // We use nearest neighbor interpolation for node position
                // computation.
                int xPixel = std::round(x * xStep);
                int yPixel = std::round(y * yStep);

                Eigen::Vector3f pixelPos(*pointImage.data(0, yPixel, xPixel), *pointImage.data(1, yPixel, xPixel), *pointImage.data(2, yPixel, xPixel));
                if (pixelPos.z() <= 0 || pixelPos.z() > maxDepth) continue;

                nodePositions.push_back(pixelPos);
                sampledNodeMapping[nodeIdx] = nodeId;
                nodeId++;
            }
        }
        int nSampledNodes = nodeId;

        // Compute graph edges using pixel-wise connectivity. Each node
        // is connected with at most 8 neighboring pixels.
        int numNeighbors = 8;
        float edgeThreshold2 = edgeThreshold * edgeThreshold;

        std::vector<int> sampledNodeEdges(nSampledNodes * numNeighbors, -1);
        std::vector<bool> connectedNodes(nSampledNodes, false);

        int nConnectedNodes = 0;
        for (int y = 0; y < yNodes; y++) {
            for (int x = 0; x < xNodes; x++) {
                int nodeIdx = y * xNodes + x;
                int nodeId = sampledNodeMapping[nodeIdx];

                if (nodeId >= 0) {
                    Eigen::Vector3f nodePosition = nodePositions[nodeId];

                    int neighborCount = 0;
                    for (int yDelta = -1; yDelta <= 1; yDelta++) {
                        for (int xDelta = -1; xDelta <= 1; xDelta++) {
                            int xNeighbor = x + xDelta;
                            int yNeighbor = y + yDelta;
                            if (xNeighbor < 0 || xNeighbor >= xNodes || yNeighbor < 0 || yNeighbor >= yNodes)
                                continue;

                            int neighborIdx = yNeighbor * xNodes + xNeighbor;

                            if (neighborIdx == nodeIdx || neighborIdx < 0)
                                continue;

                            int neighborId = sampledNodeMapping[neighborIdx];
                            if (neighborId >= 0) {
                                Eigen::Vector3f neighborPosition = nodePositions[neighborId];

                                if ((neighborPosition - nodePosition).squaredNorm() <= edgeThreshold2) {
                                    sampledNodeEdges[nodeId * numNeighbors + neighborCount] = neighborId;
                                    neighborCount++;
                                }
                            }
                        }
                    }

                    for (int i = neighborCount; i < numNeighbors; i++) {
                        sampledNodeEdges[nodeId * numNeighbors + i] = -1;
                    }

                    if (neighborCount > 0) {
                        connectedNodes[nodeId] = true;
                        nConnectedNodes += 1;
                    }
                }
            }
        }

        // Filter out nodes with no edges.
        // After changing node ids the edge ids need to be changed as well.
        std::vector<int> validNodeMapping(nSampledNodes, -1);

        graphNodes.resize({ nConnectedNodes, 3 }, false);
        graphEdges.resize({ nConnectedNodes, numNeighbors }, false);

        int validNodeId = 0;
        for (int y = 0; y < yNodes; y++) {
            for (int x = 0; x < xNodes; x++) {
                int nodeIdx = y * xNodes + x;
                int nodeId = sampledNodeMapping[nodeIdx];

                if (nodeId >= 0 && connectedNodes[nodeId]) {
                    validNodeMapping[nodeId] = validNodeId;

                    Eigen::Vector3f nodePosition = nodePositions[nodeId];
                    *graphNodes.mutable_data(validNodeId, 0) = nodePosition.x();
                    *graphNodes.mutable_data(validNodeId, 1) = nodePosition.y();
                    *graphNodes.mutable_data(validNodeId, 2) = nodePosition.z();

                    validNodeId++;
                }
            }
        }

        for (int y = 0; y < yNodes; y++) {
            for (int x = 0; x < xNodes; x++) {
                int nodeIdx = y * xNodes + x;
                int nodeId = sampledNodeMapping[nodeIdx];

                if (nodeId >= 0 && connectedNodes[nodeId]) {
                    int validNodeId = validNodeMapping[nodeId];

                    if (validNodeId >= 0) {
                        for (int i = 0; i < numNeighbors; i++) {
                            int sampledNeighborId = sampledNodeEdges[nodeId * numNeighbors + i];
                            if (sampledNeighborId >= 0) {
                                *graphEdges.mutable_data(validNodeId, i) = validNodeMapping[sampledNeighborId];
                            }
                            else {
                                *graphEdges.mutable_data(validNodeId, i) = -1;
                            }
                        }
                    }
                }
            }
        }

        // Compute pixel anchors and weights.
        pixelAnchors.resize({ height, width, 4 }, false);
        pixelWeights.resize({ height, width, 4 }, false);

        float maxPointToNodeDistance2 = maxPointToNodeDistance * maxPointToNodeDistance;

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                // Initialize with invalid values.
                for (int k = 0; k < 4; k++) {
                    *pixelAnchors.mutable_data(y, x, k) = -1;
                    *pixelWeights.mutable_data(y, x, k) = 0.f;
                }

                // Compute 4 nearest nodes.
                float xNode = float(x) / xStep;
                float yNode = float(y) / yStep;

                int x0 = std::floor(xNode), x1 = x0 + 1;
                int y0 = std::floor(yNode), y1 = y0 + 1;

                // Check that all neighboring nodes are valid.
                if (x0 < 0 || x1 >= xNodes || y0 < 0 || y1 >= yNodes)
                    continue;

                int sampledNode00 = sampledNodeMapping[y0 * xNodes + x0];
                int sampledNode01 = sampledNodeMapping[y1 * xNodes + x0];
                int sampledNode10 = sampledNodeMapping[y0 * xNodes + x1];
                int sampledNode11 = sampledNodeMapping[y1 * xNodes + x1];

                if (sampledNode00 < 0 || sampledNode01 < 0 || sampledNode10 < 0 || sampledNode11 < 0)
                    continue;

                int validNode00 = validNodeMapping[sampledNode00];
                int validNode01 = validNodeMapping[sampledNode01];
                int validNode10 = validNodeMapping[sampledNode10];
                int validNode11 = validNodeMapping[sampledNode11];

                if (validNode00 < 0 || validNode01 < 0 || validNode10 < 0 || validNode11 < 0)
                    continue;

                // Check that all nodes are close enough to the point.
                Eigen::Vector3f pixelPos(*pointImage.data(0, y, x), *pointImage.data(1, y, x), *pointImage.data(2, y, x));
                if (pixelPos.z() <= 0 || pixelPos.z() > maxDepth) continue;

                if ((pixelPos - nodePositions[sampledNode00]).squaredNorm() > maxPointToNodeDistance2 ||
                    (pixelPos - nodePositions[sampledNode01]).squaredNorm() > maxPointToNodeDistance2 ||
                    (pixelPos - nodePositions[sampledNode10]).squaredNorm() > maxPointToNodeDistance2 ||
                    (pixelPos - nodePositions[sampledNode11]).squaredNorm() > maxPointToNodeDistance2
                ) {
                    continue;
                }

                // Compute bilinear weights.
                float dx = xNode - x0;
                float dy = yNode - y0;

                float w00 = (1 - dx) * (1 - dy);
                float w01 = (1 - dx) * dy;
                float w10 = dx * (1 - dy);
                float w11 = dx * dy;

                *pixelAnchors.mutable_data(y, x, 0) = validNode00;
                *pixelWeights.mutable_data(y, x, 0) = w00;
                *pixelAnchors.mutable_data(y, x, 1) = validNode01;
                *pixelWeights.mutable_data(y, x, 1) = w01;
                *pixelAnchors.mutable_data(y, x, 2) = validNode10;
                *pixelWeights.mutable_data(y, x, 2) = w10;
                *pixelAnchors.mutable_data(y, x, 3) = validNode11;
                *pixelWeights.mutable_data(y, x, 3) = w11;
            }
        }
    }

    void update_pixel_anchors(
        const std::map<int, int>& node_id_mapping,
        py::array_t<int>& pixel_anchors
    ) {
        int height      = pixel_anchors.shape(0);
        int width       = pixel_anchors.shape(1);
        int num_anchors = pixel_anchors.shape(2);

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {

                for (int a = 0; a < num_anchors; a++) {

                    int current_anchor_id = *pixel_anchors.data(y, x, a);

                    if (current_anchor_id != -1) {
                        int mapped_anchor_id = node_id_mapping.at(current_anchor_id);

                        // update anchor only if it would actually change something
                        if (mapped_anchor_id != current_anchor_id) {
                            *pixel_anchors.mutable_data(y, x, a) = mapped_anchor_id;
                        }
                    }
                }
            }
        }
    }

} // namespace graph_proc


/*
This file is a part of Raman-Scattering-Code-Conversion.
<https://github.com/Kirbologist/Raman-Scattering-Code-Conversion>

Written by Siwan Li for the UQ School of Maths and Physics.
Copyright (C) 2021-2022 Siwan Li

This source code form is subject to the terms of the MIT License.
If a copy of the MIT License was not distributed with this file,
you can obtain one at <https://opensource.org/licenses/MIT>.

This code is used to calculate Raman scattering off of spheroids.
None of the code was included in the original SMARTIES package.
Many functions are used to provide I/O support to the main function.
*/

#ifndef RAMAN_ELASTIC_SCATTERING_HPP
#define RAMAN_ELASTIC_SCATTERING_HPP

#include "core.hpp"
#include "smarties.hpp"
#include <boost/lexical_cast.hpp>
#include <chrono>
#include <filesystem>
#include <fstream>
#include <exception>

using namespace std;
using namespace Eigen;
using namespace Smarties;
using boost::lexical_cast;
using boost::bad_lexical_cast;

/* Possible floating-point types used in calculations */
enum CalcType {SINGLE, DOUBLE, QUAD, CUSTOM, NONE};

/*
Controllable parameters that are used for calculations.
_rm suffix denotes parameters used for excitation T-matrix,
absense of _rm suffix denotes parameters used for Raman T-matrix.
*/
template <class Real>
struct RamanParams {
  // Describes the minimum and maximum of the particle diameter
  Real dia_min = 1000;
  Real dia_max = 2000;

  int N_rad = 100; // Number of particle sizes to calculate for
  int N_theta_p = 19; // Number of particle orientations to calculate for
  Real phi_p = 0; // Angle of rotation of particle about the z semiaxis

  // Number of spherical coordinates used throughout the volume of the particle
  int N_r = 100;
  int N_theta = 320;
  int N_phi = 320;

  // Values of h used, where h = a/c. N_h is the number of h values
  int N_h = 1;
  Real h_min = static_cast<Real>(1.0)/3;
  Real h_max = static_cast<Real>(1.0)/3;

  // Incident light parameters. Used for stParams struct
  Real epsilon1 = 1;
  Real epsilon2 = pow(static_cast<Real>(1.35), 2);
  Real epsilon2_rm = pow(static_cast<Real>(1.344), 2);
  Real lambda = 355;
  Real lambda_rm = 403.7;

  // Used for stParams struct
  int Nb_theta = 1000;
  int Nb_theta_pst = 1;

  // Used for stOptions struct
  int delta = 0;

  // Vectors of all possible radii/theta/h values
  ArrayXr<Real> rad_var;
  ArrayXr<Real> theta_p_var;
  ArrayXr<Real> h_var;
};

// All non-template functions defined in raman_elastic_scattering.cpp

/*
Get the 2 types to be used for floating-point calculations from a text file.
Inputs:
  in_file_name: path of the text file which contains the calculation type
Output:
  String array of size 2 that specifies 2 calculation types.
*/
std::array<CalcType, 2> GetCalcType(string in_file_name);

/*
Get the value given for some parameter from a text file.
Inputs:
  in_file_name: path of the text file which contains the option
  option: The exact string of the specific option/parameter to look for
Output:
  The user-defined value given to the option, as a string
*/
string GetOption(string in_file_name, string option);

/*
Determine if it's okay to write calculation info to a file from an option in a text file.
I.e. it checks the "Print output to file:" parameter.
Inputs:
  in_file_name: path of the text file which specifies if it's okay to write to output or not
Output:
  True if the file says "yes", i.e. it is okay, false otherwise
*/
bool CanWriteOutput(string in_file_name);

/*
Determine number of CPUs to use from a text file. I.e. it checks the "No. of CPUs" parameter.
Inputs:
  in_file_name: path of the text file which specifies the number of CPUs.
Output:
  The number of CPUs that are allowed
*/
int GetNumCPUs(string in_file_name);

/*
Determine number of threads to split the particle radii for from a text file.
I.e. it checks the "No. of CPUs to partition particle radii for:" parameter.
Note that any one of these threads can still fork into more threads if there are any available.
Inputs:
  in_file_name: path of the text file which specifies the number of CPUs.
Output:
  The number of CPUs to partition particle radii for
*/
int GetNumParticleCPUs(string in_file_name);

/*
Prints a string to both standard output (the terminal) and appends it to a text ilfe (if allowed).
Inputs:
  out_string: the string to print/write
  out_file_name: path of the text file to write to
  write_output: true enables writing to output, false disables it.
*/
void MultiPrint(string out_string, string out_file_name, bool write_output = false);

/*
Converts a string that's a fraction of two floats into a single float value
Throws a bad_lexical_cast error if string can't be read or converted.
Input:
  `frac` - must either be a floating-point literal, or two floating-point literals deliminated by a '/' character.
Output:
  the calculated value as a floating-point variable.
*/
template <class Real>
Real Frac2Float(string frac) {
  Real output;
  size_t offset = frac.find("/");
  if (offset == string::npos) {
    output = static_cast<Real>(lexical_cast<Real>(frac));
    return output;
  }
  output = static_cast<Real>(lexical_cast<Real>(frac.substr(0, offset)));
  output /= static_cast<Real>(lexical_cast<Real>(frac.substr(offset + 1)));
  return output;
}

/*
Creates RamanParams struct by reading values from text file.
The file can be opened or unopened, but must not be modified during this function call.
Inputs:
  `in_file_name` - path to the file
Output:
  A unique pointer to a new RamanParams struct with member values based on the parameters given in `in_file_name`.
Dependencies:
  GetOption, Frac2Float
*/
template <class Real>
unique_ptr<RamanParams<Real>> LoadParams(string in_file_name) {
  ifstream in_file;
  in_file.open(in_file_name, ios::in);
  if (!in_file.is_open())
    throw runtime_error("Error: cannot open " + in_file_name);

  auto output = make_unique<RamanParams<Real>>();

  string line;
  // Lists all lines/options to check for. Options with floating-point values are listed first,
  // then options with integral values are checked.
  vector<string> options {
    "Minimum diameter:", "Maximum diameter:", "Particle phi:", "Minimum h:", "Maximum h:",
    "epsilon1:", "epsilon2:", "Raman epsilon2:", "lambda:", "Raman lambda:",
    "No. of particle radii:", "No. of particle thetas:", "No. of h ratios:",
    "No. of r-coordinates:", "No. of theta-coordinates:", "No. of phi-coordinates:",
    "Nb_theta:", "Nb_theta_pst:", "Delta:"
  };
  std::array<Real, 10> float_params = {
    output->dia_min, output->dia_max, output->phi_p, output->h_min, output->h_max,
    output->epsilon1, output->epsilon2, output->epsilon2_rm, output->lambda, output->lambda_rm
  };
  std::array<int, 9> int_params = {
    output->N_rad, output->N_theta_p, output->N_h, output->N_r, output->N_theta, output->N_phi,
    output->Nb_theta, output->Nb_theta_pst, output->delta
  };
  for (size_t i = 0; i < options.size(); i++) {
    string option = options[i];
    string value = GetOption(in_file_name, option);

    try {
      if (i < float_params.size())
        float_params[i] = Frac2Float<Real>(value);
      else if (i < float_params.size() + int_params.size())
        int_params[i - float_params.size()] = lexical_cast<int>(value);
    } catch(bad_lexical_cast&) {
      cerr << "Cannot read value of option \"" << option << "\". Using default value.";
    }

    output->dia_min = float_params[0];
    output->dia_max = float_params[1];
    output->phi_p = float_params[2];
    output->h_min = float_params[3];
    output->h_max = float_params[4];
    output->epsilon1 = float_params[5];
    output->epsilon2 = float_params[6];
    output->epsilon2_rm = float_params[7];
    output->lambda = float_params[8];
    output->lambda_rm = float_params[9];
    output->N_rad = int_params[0];
    output->N_theta_p = int_params[1];
    output->N_h = int_params[2];
    output->N_r = int_params[3];
    output->N_theta = int_params[4];
    output->N_phi = int_params[5];
    output->Nb_theta = int_params[6];
    output->Nb_theta_pst = int_params[7];
    output->delta = int_params[8];
  }
  ArrayXr<Real> dia_var = ArrayXr<Real>::LinSpaced(output->N_rad, output->dia_min, output->dia_max);
  output->rad_var = dia_var/2;
  output->theta_p_var = ArrayXr<Real>::LinSpaced(output->N_theta_p, 0, mp_pi<Real>()/2);
  output->h_var = ArrayXr<Real>::LinSpaced(output->N_h, output->h_min, output->h_max);
  return output;
}

/*
Takes a RamanParams struct and uses it to generate an stParams struct for use with SMARTIES functions.
stParams member values are based on RamanParams member values, and on the radius and h value
at rad_ind and h_ind of rad_var and h_var respectively.
Inputs:
  `raman_params` - unique pointer to a RamanParams struct
  `rad_ind` - the index of the entry to use in `raman_params`->rad_var
  `h_ind` - the index of the entry to use in `raman_params`->h_var
  `type` - the type of values the output struct should have. If it's "rm",
           then the values denote the parameters of the Raman T-matrix. Otherwise they denote the parameters
           of the excitation T-matrix.
Output:
  returns a unique pointer to a new stParams struct
Dependencies:
  mp_pi
*/
template <class Real>
unique_ptr<stParams<Real>> Raman2SmartiesParams(
    const unique_ptr<RamanParams<Real>>& raman_params, int rad_ind, int h_ind, string type = string()) {
  auto params = make_unique<stParams<Real>>();
  params->epsilon1 = raman_params->epsilon1;
  params->lambda = ArrayXr<Real>(1);
  params->epsilon2 = ArrayXr<Real>(1);
  params->k1 = ArrayXr<Real>(1);
  params->s = ArrayXr<Real>(1);
  if (type == "rm") {
    params->lambda(0) = raman_params->lambda_rm;
    params->epsilon2(0) = raman_params->epsilon2_rm;
    params->k1(0) = 2*mp_pi<Real>() / raman_params->lambda_rm * sqrt(raman_params->epsilon1);
    params->s(0) = sqrt(raman_params->epsilon2_rm) / sqrt(raman_params->epsilon1);
  } else {
    params->lambda(0) = raman_params->lambda;
    params->epsilon2(0) = raman_params->epsilon2;
    params->k1(0) = 2*mp_pi<Real>() / raman_params->lambda * sqrt(raman_params->epsilon1);
    params->s(0) = sqrt(raman_params->epsilon2) / sqrt(raman_params->epsilon1);
  }
  params->Nb_theta = raman_params->Nb_theta;
  params->Nb_theta_pst = raman_params->Nb_theta_pst;

  Real h = raman_params->h_var(h_ind);
  Real radius = raman_params->rad_var(rad_ind);
  if (h > 1) {
    params->a = radius;
    params->c = radius / h;
  } else {
    params->a = radius * h;
    params->c = radius;
  }
  // The greater N is, the better the T-matrix converges.
  params->N = 6 + 2*static_cast<int>(ceil(max(params->a, params->c)/40));
  return params;
}

/*
Appends a stamp with time and parameter details of Raman scattering calculations at the end of a text file.
Function opens the file itself. It does not check for any errors.
`Real1` and `Real2` are the two selectable calculation types used for Raman scattering calculations.
DEV NOTE: Can make this into a non-template function using c++20 'concepts' keyword
Inputs:
  - out_file_name: text file to write stamp to
  - raman_params: unique pointer to a RamanParams struct containing the parameters to be written
Dependencies:
  - GetTypeName
*/
template <class Real1, class Real2>
void CreateTimeStamp(string out_file_name, const unique_ptr<RamanParams<Real1>>& raman_params) {
  string main_calc_type = GetTypeName<Real1>();
  string second_calc_type = GetTypeName<Real2>();

  ofstream out_file(out_file_name, ios::out | ios::app);
  auto sys_time = chrono::system_clock::now();
  time_t sys_time_date = chrono::system_clock::to_time_t(sys_time);
  out_file << "Session began at system time: " << ctime(&sys_time_date);
  out_file << "Running RamanElasticScattering with type " << main_calc_type;
  out_file << " and " << second_calc_type << endl;
  out_file.flush();
  out_file.close();
}

/*
Converts an existing stTR struct containing entries of one type into a new stTR struct
with the same entries converted into another type. This could theoretically be
implemented as a template cast operator or template copy constructor.
`From` is the entry type of the source struct and `To` is the entry type of the new struct.
Inputs:
  - st_TR_list: the source stTR struct, with entries of type `From`
Outputs:
  - Returns a new stTR struct, with entries of type `To`.
Dependencies:
  mp_im_unit
*/
template <class From, class To>
vector<unique_ptr<stTR<To>>> ConvertStTRList(const vector<unique_ptr<stTR<From>>>& st_TR_list) {
  vector<unique_ptr<stTR<To>>> output(st_TR_list.size());
  complex<To> I = mp_im_unit<To>();
  for (size_t i = 0; i < st_TR_list.size(); i++) {
    assert(st_TR_list[i]);
    assert(st_TR_list[i]->mat_list.size() == 2);
    assert(st_TR_list[i]->mat_list[0] == "st_4M_T");
    assert(st_TR_list[i]->mat_list[1] == "st_4M_R");
    output[i] = make_unique<stTR<To>>();
    output[i]->mat_list = st_TR_list[i]->mat_list;

    // Direct casting from a complex type to a complex type containing a Boost MPFR type isn't supported.
    output[i]->st_4M_T_eo().M11 = st_TR_list[i]->st_4M_T_eo().M11.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_T_eo().M11.imag().template cast<To>();
    output[i]->st_4M_T_eo().M12 = st_TR_list[i]->st_4M_T_eo().M12.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_T_eo().M12.imag().template cast<To>();
    output[i]->st_4M_T_eo().M21 = st_TR_list[i]->st_4M_T_eo().M21.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_T_eo().M21.imag().template cast<To>();
    output[i]->st_4M_T_eo().M22 = st_TR_list[i]->st_4M_T_eo().M22.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_T_eo().M22.imag().template cast<To>();
    output[i]->st_4M_T_eo().m = st_TR_list[i]->st_4M_T_eo().m;
    output[i]->st_4M_T_eo().ind1 = st_TR_list[i]->st_4M_T_eo().ind1;
    output[i]->st_4M_T_eo().ind2 = st_TR_list[i]->st_4M_T_eo().ind2;

    output[i]->st_4M_T_oe().M11 = st_TR_list[i]->st_4M_T_oe().M11.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_T_oe().M11.imag().template cast<To>();
    output[i]->st_4M_T_oe().M12 = st_TR_list[i]->st_4M_T_oe().M12.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_T_oe().M12.imag().template cast<To>();
    output[i]->st_4M_T_oe().M21 = st_TR_list[i]->st_4M_T_oe().M21.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_T_oe().M21.imag().template cast<To>();
    output[i]->st_4M_T_oe().M22 = st_TR_list[i]->st_4M_T_oe().M22.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_T_oe().M22.imag().template cast<To>();
    output[i]->st_4M_T_oe().m = st_TR_list[i]->st_4M_T_oe().m;
    output[i]->st_4M_T_oe().ind1 = st_TR_list[i]->st_4M_T_oe().ind1;
    output[i]->st_4M_T_oe().ind2 = st_TR_list[i]->st_4M_T_oe().ind2;

    output[i]->st_4M_R_eo().M11 = st_TR_list[i]->st_4M_R_eo().M11.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_R_eo().M11.imag().template cast<To>();
    output[i]->st_4M_R_eo().M12 = st_TR_list[i]->st_4M_R_eo().M12.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_R_eo().M12.imag().template cast<To>();
    output[i]->st_4M_R_eo().M21 = st_TR_list[i]->st_4M_R_eo().M21.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_R_eo().M21.imag().template cast<To>();
    output[i]->st_4M_R_eo().M22 = st_TR_list[i]->st_4M_R_eo().M22.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_R_eo().M22.imag().template cast<To>();
    output[i]->st_4M_R_eo().m = st_TR_list[i]->st_4M_R_eo().m;
    output[i]->st_4M_R_eo().ind1 = st_TR_list[i]->st_4M_R_eo().ind1;
    output[i]->st_4M_R_eo().ind2 = st_TR_list[i]->st_4M_R_eo().ind2;

    output[i]->st_4M_R_oe().M11 = st_TR_list[i]->st_4M_R_oe().M11.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_R_oe().M11.imag().template cast<To>();
    output[i]->st_4M_R_oe().M12 = st_TR_list[i]->st_4M_R_oe().M12.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_R_oe().M12.imag().template cast<To>();
    output[i]->st_4M_R_oe().M21 = st_TR_list[i]->st_4M_R_oe().M21.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_R_oe().M21.imag().template cast<To>();
    output[i]->st_4M_R_oe().M22 = st_TR_list[i]->st_4M_R_oe().M22.real().template cast<To>() +
        I * st_TR_list[i]->st_4M_R_oe().M22.imag().template cast<To>();
    output[i]->st_4M_R_oe().m = st_TR_list[i]->st_4M_R_oe().m;
    output[i]->st_4M_R_oe().ind1 = st_TR_list[i]->st_4M_R_oe().ind1;
    output[i]->st_4M_R_oe().ind2 = st_TR_list[i]->st_4M_R_oe().ind2;
  }
  return output;
}

/*
Performs calculations of the Raman scattering of light off of spheroids with multiprocessing.
If calculations are done for multiple spheroid radii, the calculations for each radii may be
allocated to the threads using dynamic scheduling.
The parameters are based on some input file, the T-matrices are calculated using floating-point type `Real1`,
and the field calculations are calculated using floating-point type `Real2`.
It then prints a summary of results to the standard output (terminal) and in a set of files in a given directory.
Each thread prints its own calculations to its own output file, so each output file is named
based on the parameters on the number of the thread.
Inputs:
  - in_file_name: path to the file containing parameters for use in the Raman scattering calculations.
                  Parameter details are given in the README file.
  - out_dir: directory to write output files to. If empty, output files are written to.
Dependencies:
  mp_pi, ArrayMap, TensorCast, TensorConj, LoadParams, CreateTimeStamp, MultiPrint,
  Raman2SmartiesParams, ConvertStTRList, slvForT, pstScatteringMatrixOA, vshMakeIncidentParams,
  rvhGetFieldCoefficients, pstMakeStructForField, vshEgenThetaAllPhi, vshEthetaForPhi
*/
template <class Real1, class Real2>
void RamanElasticScattering(string in_file_name, string out_dir = "", int cpus = 1) {
  // Initialise all parameters and variables
  Real2 PI = mp_pi<Real2>();

  unique_ptr<RamanParams<Real1>> raman_params1 = LoadParams<Real1>(in_file_name);
  unique_ptr<RamanParams<Real2>> raman_params2 = LoadParams<Real2>(in_file_name);
  bool write_output = (out_dir != "");

  Tensor3r<Real2> sigma_yz(raman_params2->N_rad, raman_params2->N_h, raman_params2->N_theta_p);
  Tensor3r<Real2> sigma_zy(raman_params2->N_rad, raman_params2->N_h, raman_params2->N_theta_p);
  Tensor3r<Real2> sigma_zz(raman_params2->N_rad, raman_params2->N_h, raman_params2->N_theta_p);
  Tensor3r<Real2> sigma_yy(raman_params2->N_rad, raman_params2->N_h, raman_params2->N_theta_p);
  ArrayXr<Real1> C_sca = ArrayXr<Real1>::Zero(raman_params1->N_rad);
  ArrayXr<Real1> C_ext = ArrayXr<Real1>::Zero(raman_params1->N_rad);
  ArrayXr<Real1> C_abs = ArrayXr<Real1>::Zero(raman_params1->N_rad);
  vector<unique_ptr<stSM<Real2>>> stSM_list(raman_params2->N_rad);

  auto options = make_unique<stOptions>();
  options->get_R = true; // Needed for near fields and will be overridden in any case
  options->delta = raman_params1->delta; // Use delta=-1 to estimate it automatically.
  options->NB = 0; // NB will be estimated automatically
  options->get_symmetric_T = false;
  Real2 phi_p = raman_params2->phi_p;

  int N_r = raman_params2->N_r;
  int N_theta = raman_params2->N_theta;
  int N_phi = raman_params2->N_phi;

  // Initialises a set regularly spaced theta coordinates that are used for internal field calculations
  RowArrayXr<Real2> theta_surf = RowArrayXr<Real2>::LinSpaced(N_theta, 0, PI);
  // Initialises a non-uniform set of N_r radii inside a unit ball
  // such that spherical shells between adjacent radii are equal volume
  ArrayXr<Real2> r_surf_u = ArrayXr<Real2>::LinSpaced(
      N_r, 1/static_cast<Real2>(N_r), 1).pow(static_cast<Real2>(1.0)/3);
  ArrayXXr<Real2> r_mat_u = r_surf_u.replicate(1, N_theta);
  ArrayXXr<Real2> theta_mat = theta_surf.replicate(N_r, 1);

  // Initialise output files
  if (write_output)
    std::filesystem::create_directory(out_dir);

  for (int h_ind = 0; h_ind < raman_params1->N_h; h_ind++) {
    // Split the iterations of the next for-loop among threads for parallel processing.
    #pragma omp parallel for num_threads(cpus) schedule(dynamic)
    for (int k = 0; k < raman_params1->N_rad; k++) {
      string out_file_name = out_dir + "/a2c_" + to_string(raman_params1->h_var(0)) +
          "_to_" + to_string(raman_params1->h_var(last)) + "_dia_" + to_string(raman_params1->dia_min) +
          "_to_" + to_string(raman_params1->dia_max) + "_range_" + to_string(omp_get_thread_num()) + ".txt";
      if (write_output) {
        ofstream out_file(out_file_name, ios::out | ios::app);
        out_file << endl;
        if (!out_file.good()) {
          cerr << "Warning: cannot create or open output file. Some output won't be written." << endl;
          out_file.close();
          write_output = false;
        }
      }
      if (write_output)
        CreateTimeStamp<Real1, Real2>(out_file_name, raman_params1);
      stringstream out_stream;

      out_stream.precision(4);
      out_stream << "aspect ratio " << raman_params1->h_var(h_ind) << endl;
      MultiPrint(out_stream.str(), out_file_name, write_output);
      out_stream.str(string());

      unique_ptr<stParams<Real1>> params1 = Raman2SmartiesParams(raman_params1, k, h_ind);
      unique_ptr<stParams<Real1>> params_rm1 = Raman2SmartiesParams(raman_params1, k, h_ind, "rm");
      unique_ptr<stParams<Real2>> params2 = Raman2SmartiesParams(raman_params2, k, h_ind);
      unique_ptr<stParams<Real2>> params_rm2 = Raman2SmartiesParams(raman_params2, k, h_ind, "rm");
      Real2 a = params2->a;
      Real2 c = params2->c;

      int N = params2->N;
      // Initialises array of radii coordinates that are used for internal field calculations
      // by conforming r_mat_u to the spheroid shape. The ratios of radii along the same angle theta is kept invariant.
      ArrayXXr<Real2> r_mat = a*c*r_mat_u/sqrt(pow(c, 2)*sin(theta_mat).pow(2) + pow(a, 2)*cos(theta_mat).pow(2));
      ArrayXXr<Real2> d_theta = theta_mat(all, seq(1, last)) - theta_mat(all, seq(0, last - 1));
      ArrayXXr<Real2> dr = r_mat(seq(1, last), all) - r_mat(seq(0, last - 1), all);
      ArrayXXr<Real2> dt_dr = d_theta(seq(1, last), all) * dr(all, seq(1, last)); // Jacobian of the integral
      RowArrayXr<Real2> r_row = r_mat.reshaped().transpose();
      RowArrayXr<Real2> theta_row = theta_mat.reshaped().transpose();

      chrono::steady_clock::time_point begin, end;
      chrono::duration<double> elapsed_seconds;

      // Calculating T-matrix for excitation
      begin = chrono::steady_clock::now();
      unique_ptr<stTmatrix<Real1>> T_mat = slvForT(params1, options);
      end = chrono::steady_clock::now();
      elapsed_seconds = end - begin;
      out_stream << elapsed_seconds.count() << endl;
      out_stream << "T-matrix calculatred for excitation" << endl;
      MultiPrint(out_stream.str(), out_file_name, write_output);
      out_stream.str(string());

      // Calculating T-matrix for Raman
      begin = chrono::steady_clock::now();
      unique_ptr<stTmatrix<Real1>> T_mat_rm = slvForT(params_rm1, options);
      end = chrono::steady_clock::now();
      elapsed_seconds = end - begin;
      out_stream << elapsed_seconds.count() << endl;
      out_stream << "T-matrix calculatred for Raman" << endl;
      MultiPrint(out_stream.str(), out_file_name, write_output);
      out_stream.str(string());

      // Converting variables to secondary calculation type
      vector<unique_ptr<stTR<Real2>>> st_TR_list = ConvertStTRList<Real1, Real2>(T_mat->st_TR_list);
      vector<unique_ptr<stTR<Real2>>> st_TR_list_rm = ConvertStTRList<Real1, Real2>(T_mat_rm->st_TR_list);

      C_sca(k) = T_mat->st_C_oa->C_sca(0); // st_C_oa->sca should only contain 1 element
      C_ext(k) = T_mat->st_C_oa->C_ext(0); // st_C_oa->ext should only contain 1 element
      C_abs(k) = T_mat->st_C_oa->C_abs(0); // st_C_oa->abs should only contain 1 element
      out_stream.precision(10);
      out_stream << "C_sca " << C_sca(k) << endl;
      out_stream << "C_ext " << C_ext(k) << endl;
      out_stream << "C_abs " << C_abs(k) << endl;
      MultiPrint(out_stream.str(), out_file_name, write_output);
      out_stream.str(string());

      // Calculating scattering matrix for 2 scattering angles 0 and pi
      stSM_list[k] = pstScatteringMatrixOA(st_TR_list, params2->lambda(0), static_cast<Real2>(C_sca(k)));

      for (int t = 0; t < raman_params2->N_theta_p; t++) {
        begin = chrono::steady_clock::now();
        Real2 theta_p = raman_params2->theta_p_var(t);
        std::array<long int, 3> new_dims = {N_r, N_theta, 3};
        ArrayXr<Real2> phi_var = ArrayXr<Real2>::LinSpaced(N_phi + 1, 0, 2*PI);

        // Calculate field at excitation wavelength
        Real2 alpha_p = 0; // Defines the orientation of the electric field, in the plane orthogonal to wavevector k.
        params2->inc_par = vshMakeIncidentParams(sIncType::GENERAL, N, theta_p, phi_p, alpha_p);
        // Internal field calculations for the incident field defined by stIncPar
        unique_ptr<stAbcdnm<Real2>> st_abcdnm = rvhGetFieldCoefficients(N, st_TR_list, params2->inc_par);
        unique_ptr<stRes<Real2>> st_res_E = pstMakeStructForField(st_abcdnm, params2);
        ArrayXXc<Real2> c_nm = Map<RowArrayXc<Real2>>(st_res_E->c_nm.transpose().data(), st_res_E->c_nm.size());
        ArrayXXc<Real2> d_nm = Map<RowArrayXc<Real2>>(st_res_E->d_nm.transpose().data(), st_res_E->d_nm.size());
        unique_ptr<stEAllPhi<Real2>> st_E_surf = vshEgenThetaAllPhi(st_res_E->lambda, st_res_E->epsilon2,
            c_nm, d_nm, r_row, theta_row, sBessel::J);
        // The field is calculated for a set of phi, from 0 and 2*pi
        Tensor4c<Real2> E_field_z(N_r, N_theta, N_phi + 1, 3);
        E_field_z.setZero();
        for (int m = 0; m <= N_phi; m++) {
          Real2 phi = phi_var(m);
          unique_ptr<stEforPhi<Real2>> st_E_for_phi = vshEthetaForPhi(st_E_surf, phi);
          long int dims[2] = {st_E_for_phi->E_r.cols(), 3*st_E_for_phi->E_r.rows()};
          ArrayXXc<Real2> E_field_phi(dims[0], dims[1]);
          E_field_phi << st_E_for_phi->E_r.matrix().adjoint(),
              st_E_for_phi->E_t.matrix().adjoint(), st_E_for_phi->E_f.matrix().adjoint();
          E_field_z.chip(m, 2) = TensorCast(E_field_phi).reshape(new_dims);
        }

        alpha_p = PI/2; // Internal field calculations for different polarisation alpha_p
        params2->inc_par = vshMakeIncidentParams(sIncType::GENERAL, N, theta_p, phi_p, alpha_p);
        st_abcdnm = rvhGetFieldCoefficients(N, st_TR_list, params2->inc_par);
        st_res_E = pstMakeStructForField(st_abcdnm, params2);
        c_nm = Map<RowArrayXc<Real2>>(st_res_E->c_nm.transpose().data(), st_res_E->c_nm.size());
        d_nm = Map<RowArrayXc<Real2>>(st_res_E->d_nm.transpose().data(), st_res_E->d_nm.size());
        st_E_surf = vshEgenThetaAllPhi(st_res_E->lambda, st_res_E->epsilon2,
            c_nm, d_nm, r_row, theta_row, sBessel::J);
        // The field is calculated for a set of phi, from 0 and 2*pi
        Tensor4c<Real2> E_field_y(N_r, N_theta, N_phi + 1, 3);
        E_field_y.setZero();
        for (int m = 0; m <= N_phi; m++) {
          Real2 phi = phi_var(m);
          unique_ptr<stEforPhi<Real2>> st_E_for_phi = vshEthetaForPhi(st_E_surf, phi);
          long int dims[2] = {st_E_for_phi->E_r.cols(), 3*st_E_for_phi->E_r.rows()};
          ArrayXXc<Real2> E_field_phi(dims[0], dims[1]);
          E_field_phi << st_E_for_phi->E_r.matrix().adjoint(),
              st_E_for_phi->E_t.matrix().adjoint(), st_E_for_phi->E_f.matrix().adjoint();
          E_field_y.chip(m, 2) = TensorCast(E_field_phi).reshape(new_dims);
        }

        // Calculate field at Raman wavelength
        alpha_p = 0;
        params_rm2->inc_par = vshMakeIncidentParams(sIncType::GENERAL, N, theta_p, phi_p, alpha_p);
        st_abcdnm = rvhGetFieldCoefficients(N, st_TR_list_rm, params_rm2->inc_par);
        st_res_E = pstMakeStructForField(st_abcdnm, params_rm2);
        c_nm = Map<RowArrayXc<Real2>>(st_res_E->c_nm.transpose().data(), st_res_E->c_nm.size());
        d_nm = Map<RowArrayXc<Real2>>(st_res_E->d_nm.transpose().data(), st_res_E->d_nm.size());
        st_E_surf = vshEgenThetaAllPhi(st_res_E->lambda, st_res_E->epsilon2,
            c_nm, d_nm, r_row, theta_row, sBessel::J);
        // The field is calculated for a set of phi, from 0 and 2*pi
        Tensor4c<Real2> E_field_rm_z(N_r, N_theta, N_phi + 1, 3);
        E_field_rm_z.setZero();
        for (int m = 0; m <= N_phi; m++) {
          Real2 phi = phi_var(m);
          unique_ptr<stEforPhi<Real2>> st_E_for_phi = vshEthetaForPhi(st_E_surf, phi);
          long int dims[2] = {st_E_for_phi->E_r.cols(), 3*st_E_for_phi->E_r.rows()};
          ArrayXXc<Real2> E_field_phi(dims[0], dims[1]);
          E_field_phi << st_E_for_phi->E_r.matrix().adjoint(),
              st_E_for_phi->E_t.matrix().adjoint(), st_E_for_phi->E_f.matrix().adjoint();
          E_field_rm_z.chip(m, 2) = TensorCast(E_field_phi).reshape(new_dims);
        }

        alpha_p = PI/2;
        params_rm2->inc_par = vshMakeIncidentParams(sIncType::GENERAL, N, theta_p, phi_p, alpha_p);
        st_abcdnm = rvhGetFieldCoefficients(N, st_TR_list_rm, params_rm2->inc_par);
        st_res_E = pstMakeStructForField(st_abcdnm, params_rm2);
        c_nm = Map<RowArrayXc<Real2>>(st_res_E->c_nm.transpose().data(), st_res_E->c_nm.size());
        d_nm = Map<RowArrayXc<Real2>>(st_res_E->d_nm.transpose().data(), st_res_E->d_nm.size());
        st_E_surf = vshEgenThetaAllPhi(st_res_E->lambda, st_res_E->epsilon2,
            c_nm, d_nm, r_row, theta_row, sBessel::J);
        // The field is calculated for a set of phi, from 0 and 2*pi
        Tensor4c<Real2> E_field_rm_y(N_r, N_theta, N_phi + 1, 3);
        E_field_rm_y.setZero();
        for (int m = 0; m <= N_phi; m++) {
          Real2 phi = phi_var(m);
          unique_ptr<stEforPhi<Real2>> st_E_for_phi = vshEthetaForPhi(st_E_surf, phi);
          long int dims[2] = {st_E_for_phi->E_r.cols(), 3*st_E_for_phi->E_r.rows()};
          ArrayXXc<Real2> E_field_phi(dims[0], dims[1]);
          E_field_phi << st_E_for_phi->E_r.matrix().adjoint(),
              st_E_for_phi->E_t.matrix().adjoint(), st_E_for_phi->E_f.matrix().adjoint();
          E_field_rm_y.chip(m, 2) = TensorCast(E_field_phi).reshape(new_dims);
        }

        // Calculate scattering cross-sections
        std::array<int, 1> dim3 = {3}, dim2 = {2};
        Tensor<Real2, 2> M_tens = 2*PI/N_phi*(E_field_rm_z * TensorConj(E_field_z))
            .sum(dim3).abs().pow(static_cast<Real2>(2.0)).sum(dim2);
        Map<ArrayXXr<Real2>> M_mat = ArrayMap(M_tens);
        ArrayXXr<Real2> F = M_mat * r_mat.pow(2) * sin(theta_mat);
        ArrayXXr<Real2> intergrand = (F(seq(1, last), seq(1, last)) + F(seq(0, last - 1), seq(1, last)) +
            F(seq(1, last), seq(0, last - 1)) + F(seq(0, last - 1), seq(0, last - 1))) / 4 * dt_dr;
        sigma_zz(k, h_ind, t) = (intergrand).sum()/(PI*a*a*c*4/3);

        M_tens = 2*PI/N_phi*(E_field_rm_y * TensorConj(E_field_z))
            .sum(dim3).abs().pow(static_cast<Real2>(2.0)).sum(dim2);
        M_mat = ArrayMap(M_tens);
        F = M_mat * r_mat.pow(2) * sin(theta_mat);
        intergrand = (F(seq(1, last), seq(1, last)) + F(seq(0, last - 1), seq(1, last)) +
            F(seq(1, last), seq(0, last - 1)) + F(seq(0, last - 1), seq(0, last - 1))) / 4 * dt_dr;
        sigma_yz(k, h_ind, t) = (intergrand).sum()/(PI*a*a*c*4/3);

        M_tens = 2*PI/N_phi*(E_field_rm_y * TensorConj(E_field_y))
            .sum(dim3).abs().pow(static_cast<Real2>(2.0)).sum(dim2);
        M_mat = ArrayMap(M_tens);
        F = M_mat * r_mat.pow(2) * sin(theta_mat);
        intergrand = (F(seq(1, last), seq(1, last)) + F(seq(0, last - 1), seq(1, last)) +
            F(seq(1, last), seq(0, last - 1)) + F(seq(0, last - 1), seq(0, last - 1))) / 4 * dt_dr;
        sigma_yy(k, h_ind, t) = (intergrand).sum()/(PI*a*a*c*4/3);

        M_tens = 2*PI/N_phi*(E_field_rm_z * TensorConj(E_field_y))
            .sum(dim3).abs().pow(static_cast<Real2>(2.0)).sum(dim2);
        M_mat = ArrayMap(M_tens);
        F = M_mat * r_mat.pow(2) * sin(theta_mat);
        intergrand = (F(seq(1, last), seq(1, last)) + F(seq(0, last - 1), seq(1, last)) +
            F(seq(1, last), seq(0, last - 1)) + F(seq(0, last - 1), seq(0, last - 1))) / 4 * dt_dr;
        sigma_zy(k, h_ind, t) = (intergrand).sum()/(PI*a*a*c*4/3);

        out_stream.precision(5);
        out_stream << "max diameter = " << max(a, c)*2e-3;
        out_stream.precision(10);
        out_stream << " µm, theta = " << theta_p * 180/PI << " degrees" << endl;
        out_stream.precision(11);
        out_stream << "--- relative raman zz = " << sigma_zz(k, h_ind, t) << endl;
        out_stream << "--- relative raman yz = " << sigma_yz(k, h_ind, t) << endl;
        out_stream << "--- relative raman zy = " << sigma_zy(k, h_ind, t) << endl;
        out_stream << "--- relative raman yy = " << sigma_yy(k, h_ind, t) << endl;

        end = chrono::steady_clock::now();
        elapsed_seconds = end - begin;
        out_stream << elapsed_seconds.count() << endl;

        MultiPrint(out_stream.str(), out_file_name, write_output);
        out_stream.str(string());
      }
    }
  }
}

#endif


#include <vector>
#include <boost/math/tools/polynomial.hpp>
#include <fstream>
#include <immintrin.h>
#include <map>
#include "cpp-btree/btree_map.h"
//---------------------------------------------------------------------------
typedef std::pair<double,double> Coord;
typedef std::pair<double, unsigned> segCoord;
typedef btree::btree_map<double, double> btree_map;
typedef btree::btree_map<double, unsigned> btree_map_segments;
//---------------------------------------------------------------------------
class SplineEvaluation;
//---------------------------------------------------------------------------
/// A linear spline for evaluation
class SplineEvaluation
{
    private:
    static const unsigned countOfBuckets = 2 + 1;
    static const unsigned limitSplineSize = 100000;
    int32_t offset;
    int32_t length;

    void transformSpline();
    double selfInterpolate(unsigned segment, double& x) const;
    unsigned searchLeft(int segment, double& x) const;
    unsigned searchRight(int segment, double& x) const;
    unsigned binarySearch(unsigned lower, unsigned upper, double& x) const;
    void changeIntoArray(long double** moveHere, boost::math::tools::polynomial<long double> poly);
    void convertValues();
    void splitSegmentSpline();
    public:
    static const unsigned polyGrade = 2;
    /// Constructor
    SplineEvaluation();
    /// Deconstructor
    ~SplineEvaluation();
    /// Constructor with the function
    SplineEvaluation(const std::vector<Coord>& function, unsigned desiredSize, const unsigned useSplineMode, const unsigned searchMode, const unsigned testMode);
    /// Constructor with given data
    SplineEvaluation(unsigned size, const double* x, const double* y);

    void save(std::ofstream& file_out);
    void load(std::ifstream& file_in);

    unsigned size() const;
    unsigned findExactSegment(double segment, double& x) const;
    /// Return the approximate value of function(x)
    double evaluate(double& x) const;
    double hornerEvaluate(double& x) const;

    double splittedHornerEvaluate(double& x) const;
    double hornerEvaluateByBuckets(unsigned bucketIndex, double& x) const;

    double chebyshevEvaluate(double& x) const;
    double splittedChebyshevEvaluate(double& x) const;

    double mapEvaluate(double& x) const;
    std::map<double, double> coordinates;

    double btreeEvaluate(double& x) const;
    btree_map* btreeMap;

    double btreeSegmentsEvaluate(double& x) const;
    btree_map_segments* btreeSegments;

    std::vector<Coord> spline;
    /// Spline which has to be fitted
    std::vector<segCoord> segmentSpline;
    /// Fitted polynomial to the spline
    long double* poly;
    double* values;
    unsigned splineSize;

    std::vector<segCoord>* splineBuckets;
    long double **polyBuckets;
    double splittedPositions[countOfBuckets];
    __m128d saveLimits;

    /// Fit a polynomial to spline
    boost::math::tools::polynomial<long double> fitSpline(const std::vector<segCoord>& spline);
};
//---------------------------------------------------------------------------


/*
 * convnn.cpp
 *
 */
#include <string>

#include <boost/assert.hpp>

#include "core/utils.h"
#include "core/functions.h"
#include "layers/contlayer.h"
#include "layers/convlayer.h"
#include "layers/fclayer.h"
#include "layers/polllayer.h"
#include "layers/smaxlayer.h"

#include "cnn.h"

using namespace std;
using namespace yann;

////////////////////////////////////////////////////////////////////////////////////////////////
//
// ConvolutionalNetwork implementation
//
yann::ConvolutionalNetwork::ConvPollParams::ConvPollParams() :
    _output_frames_num(0),
    _conv_filter_size(0),
    _polling_mode(PollingLayer::PollMode_Avg),
    _polling_filter_size(0)
{
}

void yann::ConvolutionalNetwork::append(unique_ptr<SequentialLayer> & container,
                   const MatrixSize & input_rows,
                   const MatrixSize & input_cols,
                   const MatrixSize & input_frames_num,
                   const ConvPollParams & params)
{
  YANN_CHECK(container);
  YANN_CHECK_GT(input_rows, 0);
  YANN_CHECK_GT(input_cols, 0);
  YANN_CHECK_GT(input_frames_num, 0);
  YANN_CHECK(input_frames_num == 1 || !params._mappings.empty());
  YANN_CHECK_GT(params._output_frames_num, 0 || !params._mappings.empty());
  YANN_CHECK_GT(params._conv_filter_size, 0);
  YANN_CHECK_GT(params._polling_filter_size, 0);

  // ConvolutionalLayer
  if(input_frames_num == 1) {
    auto conv_bcast_layer = ConvolutionalLayer::create_conv_bcast_layer(
        params._output_frames_num,   // N output frames
        input_rows,
        input_cols,
        params._conv_filter_size,
        params._conv_activation_funtion
    );
    YANN_CHECK(conv_bcast_layer);
    container->append_layer(std::move(conv_bcast_layer));
  } else {
    YANN_CHECK(!params._mappings.empty());
    YANN_CHECK(params._output_frames_num == 0 || params._output_frames_num == params._mappings.size());

    auto conv_mapping_layer = make_unique<MappingLayer>(input_frames_num);
    YANN_CHECK(conv_mapping_layer);
    for(auto & layer_mappings : params._mappings) {
      auto conv_layer = make_unique<ConvolutionalLayer>(
          input_rows,
          input_cols,
          params._conv_filter_size);
      YANN_CHECK(conv_layer);

      if(params._conv_activation_funtion) {
        conv_layer->set_activation_function(params._conv_activation_funtion);
      }
      conv_mapping_layer->append_layer(std::move(conv_layer), layer_mappings);
    }
    container->append_layer(std::move(conv_mapping_layer));
  }

  // PollingLayer
  MatrixSize poll_input_rows = ConvolutionalLayer::get_conv_output_rows(input_rows, params._conv_filter_size);
  MatrixSize poll_input_cols = ConvolutionalLayer::get_conv_output_cols(input_cols, params._conv_filter_size);
  auto poll_layer = PollingLayer::create_poll_parallel_layer(
     params._output_frames_num,
     poll_input_rows,
     poll_input_cols,
     params._polling_filter_size,
     params._polling_mode,
     params._polling_activation_funtion);
  YANN_CHECK(poll_layer);
  container->append_layer(std::move(poll_layer));
}

unique_ptr<SequentialLayer> yann::ConvolutionalNetwork::create(
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    const ConvPollParams & params,
    const std::unique_ptr<ActivationFunction> & fc_activation_funtion,
    const MatrixSize & output_size)
{
  YANN_CHECK_GT(output_size, 0);

  auto container = make_unique<SequentialLayer>();
  YANN_CHECK(container);

  // append one layer: we have 1 input frame
  append(container, input_rows, input_cols, 1, params);

  // FullyConnectedLayer
  auto fc_layer = make_unique<FullyConnectedLayer>(
      container->get_output_size(),
      output_size);
  YANN_CHECK(fc_layer);
  if(fc_activation_funtion) {
    fc_layer->set_activation_function(fc_activation_funtion);
  }
  container->append_layer(std::move(fc_layer));

  // done
  return container;
}

std::unique_ptr<SequentialLayer> yann::ConvolutionalNetwork::create(
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    const ConvPollParams & params1,
    const ConvPollParams & params2,
    const std::unique_ptr<ActivationFunction> & fc_activation_funtion,
    const MatrixSize & output_size)
{
  YANN_CHECK_GT(output_size, 0);

  auto container = make_unique<SequentialLayer>();
  YANN_CHECK(container);

  // append one layer: we have 1 input frame
  append(container, input_rows, input_cols, 1, params1);

  // append second layer: the output frames from first layer are our inputs
  MatrixSize poll1_input_rows = ConvolutionalLayer::get_conv_output_rows(input_rows, params1._conv_filter_size);
  MatrixSize poll1_input_cols = ConvolutionalLayer::get_conv_output_cols(input_cols, params1._conv_filter_size);
  MatrixSize conv2_input_rows = PollingLayer::get_output_rows(poll1_input_rows, params1._polling_filter_size);
  MatrixSize conv2_input_cols = PollingLayer::get_output_cols(poll1_input_cols, params1._polling_filter_size);
  append(container, conv2_input_rows, conv2_input_cols, params1._output_frames_num, params2);

  // FullyConnectedLayer
  auto fc_layer = make_unique<FullyConnectedLayer>(
      container->get_output_size(),
      output_size);
  YANN_CHECK(fc_layer);
  if(fc_activation_funtion) {
    fc_layer->set_activation_function(fc_activation_funtion);
  }
  container->append_layer(std::move(fc_layer));

  // done
  return container;
}

std::unique_ptr<SequentialLayer> yann::ConvolutionalNetwork::create_lenet1(
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    PollingLayer::Mode polling_mode,
    const MatrixSize & fc_size,
    const MatrixSize & output_size,
    const std::unique_ptr<ActivationFunction> & conv_activation_funtion,
    const std::unique_ptr<ActivationFunction> & poll_activation_funtion,
    const std::unique_ptr<ActivationFunction> & fc_activation_funtion)
{
  // see http://yann.lecun.com/exdb/publis/pdf/lecun-90c.pdf
  ConvPollParams params1;
  params1._output_frames_num = 4;
  params1._conv_filter_size  = 5;
  if(conv_activation_funtion) {
    params1._conv_activation_funtion = conv_activation_funtion->copy();
  }
  params1._polling_mode = polling_mode;
  params1._polling_filter_size = 2;
  if(poll_activation_funtion) {
    params1._polling_activation_funtion = poll_activation_funtion->copy();
  }

  ConvPollParams params2;
  params2._output_frames_num = 12;
  params2._conv_filter_size  = 5;
  if(conv_activation_funtion) {
    params2._conv_activation_funtion = conv_activation_funtion->copy();
  }
  params2._polling_mode = polling_mode;
  params2._polling_filter_size = 2;
  if(poll_activation_funtion) {
    params2._polling_activation_funtion = poll_activation_funtion->copy();
  }
  params2._mappings.push_back({ 0 });
  params2._mappings.push_back({ 0, 1 });
  params2._mappings.push_back({ 0, 1 });
  params2._mappings.push_back({ 1 });
  params2._mappings.push_back({ 0, 1 });
  params2._mappings.push_back({ 0, 1 });
  params2._mappings.push_back({ 2 });
  params2._mappings.push_back({ 2, 3 });
  params2._mappings.push_back({ 2, 3 });
  params2._mappings.push_back({ 3 });
  params2._mappings.push_back({ 2, 3 });
  params2._mappings.push_back({ 2, 3 });

  auto container = create(
      input_rows, input_cols, params1, params2,
      fc_activation_funtion, fc_size > 0 ? fc_size : output_size);
  YANN_CHECK(container);

  // add one more FC layer if needed
  if(fc_size > 0) {
    auto fc_layer = make_unique<FullyConnectedLayer>(fc_size, output_size);
    YANN_CHECK(fc_layer);
    if(fc_activation_funtion) {
      fc_layer->set_activation_function(fc_activation_funtion);
    }
    container->append_layer(std::move(fc_layer));
  }

  // done
  return container;
}

std::unique_ptr<SequentialLayer> yann::ConvolutionalNetwork::create_boosted_lenet1(
    const MatrixSize & paths_num,
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    PollingLayer::Mode polling_mode,
    const MatrixSize & fc_size,
    const MatrixSize & output_size,
    const std::unique_ptr<ActivationFunction> & conv_activation_funtion,
    const std::unique_ptr<ActivationFunction> & poll_activation_funtion,
    const std::unique_ptr<ActivationFunction> & fc_activation_funtion)
{
  // create N paths
  auto bcast_layer = make_unique<BroadcastLayer>();
  YANN_CHECK(bcast_layer);
  for(auto ii = paths_num; ii > 0; --ii) {
    auto path = create_lenet1(
        input_rows, input_cols, polling_mode,
        0, fc_size, // don't create fc layer, we will add one for the merge
        conv_activation_funtion, poll_activation_funtion,
        fc_activation_funtion);
    YANN_CHECK(path);
    bcast_layer->append_layer(std::move(path));
  }

  // merge them together
  auto fc_layer = make_unique<FullyConnectedLayer>(paths_num * fc_size, output_size);
  YANN_CHECK(fc_layer);
  fc_layer->set_activation_function(fc_activation_funtion);

  // and finally create container
  auto container = make_unique<SequentialLayer>();
  YANN_CHECK(container);
  container->append_layer(std::move(bcast_layer));
  container->append_layer(std::move(fc_layer));

  // done
  return container;
}

std::unique_ptr<SequentialLayer> yann::ConvolutionalNetwork::create_lenet5(
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    PollingLayer::Mode polling_mode,
    const MatrixSize & fc1_size,
    const MatrixSize & fc2_size,
    const MatrixSize & output_size,
    const std::unique_ptr<ActivationFunction> & conv_activation_funtion,
    const std::unique_ptr<ActivationFunction> & poll_activation_funtion,
    const std::unique_ptr<ActivationFunction> & fc_activation_funtion)
{
  // see http://vision.stanford.edu/cs598_spring07/papers/Lecun98.pdf
  ConvPollParams params1;
  params1._output_frames_num = 6;
  params1._conv_filter_size  = 5;
  if(conv_activation_funtion) {
    params1._conv_activation_funtion = conv_activation_funtion->copy();
  }
  params1._polling_mode = polling_mode;
  params1._polling_filter_size = 2;
  if(poll_activation_funtion) {
    params1._polling_activation_funtion = poll_activation_funtion->copy();
  }

  ConvPollParams params2;
  params2._output_frames_num = 16;
  params2._conv_filter_size  = 5;
  if(conv_activation_funtion) {
    params2._conv_activation_funtion = conv_activation_funtion->copy();
  }
  params2._polling_mode = polling_mode;
  params2._polling_filter_size = 2;
  if(poll_activation_funtion) {
    params2._polling_activation_funtion = poll_activation_funtion->copy();
  }
  params2._mappings.push_back({ 0, 1, 2 }); // 0
  params2._mappings.push_back({ 1, 2, 3 });
  params2._mappings.push_back({ 2, 3, 4 });
  params2._mappings.push_back({ 3, 4, 5 });
  params2._mappings.push_back({ 0, 4, 5 }); // 4
  params2._mappings.push_back({ 0, 1, 5 });
  params2._mappings.push_back({ 0, 1, 2, 3 });
  params2._mappings.push_back({ 1, 2, 3, 4 });
  params2._mappings.push_back({ 2, 3, 4, 5 }); // 8
  params2._mappings.push_back({ 0, 3, 4, 5 });
  params2._mappings.push_back({ 0, 1, 4, 5 });
  params2._mappings.push_back({ 0, 1, 2, 5 });
  params2._mappings.push_back({ 0, 1, 3, 4 }); // 12
  params2._mappings.push_back({ 1, 2, 4, 5 });
  params2._mappings.push_back({ 0, 2, 3, 5 });
  params2._mappings.push_back({ 0, 1, 2, 3, 4, 5 });

  auto container = create(input_rows, input_cols, params1, params2, fc_activation_funtion, fc1_size);
  YANN_CHECK(container);

  // add one more FC layer
  auto fc_layer1 = make_unique<FullyConnectedLayer>(fc1_size, fc2_size);
  YANN_CHECK(fc_layer1);
  if(fc_activation_funtion) {
    fc_layer1->set_activation_function(fc_activation_funtion);
  }
  container->append_layer(std::move(fc_layer1));

  // and one more FC layer
  auto fc_layer2 = make_unique<FullyConnectedLayer>(fc2_size, output_size);
  YANN_CHECK(fc_layer2);
  if(fc_activation_funtion) {
    fc_layer2->set_activation_function(fc_activation_funtion);
  }
  container->append_layer(std::move(fc_layer2));

  // done
  return container;
}


#include "stdafx.h"
#include "NRSolver.h"
#include "PrimitiveNetwork.h"
#include "Exceptions.h"
#include <Eigen/Sparse>
#include <Eigen/SparseLU>

using Eigen::Triplet;
using Eigen::SparseLU;
using Eigen::SparseQR;
using Eigen::SparseMatrix;

using namespace std;
using namespace PowerSolutions::ObjectModel;

namespace PowerSolutions
{
	namespace PowerFlow
	{

		NRSolver::NRSolver()
		{ }

		NRSolver::~NRSolver()
		{ }

		void NRSolver::BeforeIterations()
		{
			assert(Block1EquationCount() == PQNodeCount + PVNodeCount);
			assert(EquationCount() == PQNodeCount * 2 + PVNodeCount);
			//PQ节点：[dP dQ] = J * [V theta]
			//PV节点：[dP] = J * [theta]
			//       PQ/PV    |   PQ
			//Δy = [dP ... dP | dQ ... dQ] = ConstraintPower - CurrentPower
			//Δx = [theta ... theta | V ... V]
			// 分界线/起始索引：
			// H[0,0], N[NodeCount - 1, 0]
			//初始化矩阵维数
			//注意向量需要手动清零
			ConstraintPowerInjection.resize(EquationCount());
			CurrentAnswer.resize(EquationCount());
			CorrectionAnswer.resize(EquationCount());
			PowerInjectionDeviation.resize(EquationCount());
			//估计雅可比矩阵非零元的数量
			//按照最不理想的情况来考虑雅可比矩阵的空间
			JocobianReservedValuesCount = 0;
			for (int n = 0; n < Block1EquationCount(); n++)
			{
				auto& node = PNetwork->Nodes(n);
				JocobianReservedValuesCount += min(node.Degree() * 2, EquationCount());
				if (node.Type() == NodeType::PQNode)
				{
					JocobianReservedValuesCount += min(node.Degree() * 2, EquationCount());
				}
			}
			Jocobian.resize(EquationCount(), EquationCount());

			//生成目标注入功率向量 y，以及迭代初值向量。
			for (auto& node : PNetwork->Nodes())
			{
				if (node->Type() != NodeType::SlackNode)
				{
					int subIndex = Block1EquationCount() + node->SubIndex();
					ConstraintPowerInjection(node->Index()) = node->ActivePowerInjection();
					CurrentAnswer(node->Index()) = arg(node->Bus()->InitialVoltage());
					if (node->Type() == NodeType::PQNode)
					{
						ConstraintPowerInjection(subIndex) = node->ReactivePowerInjection();
						CurrentAnswer(subIndex) = abs(node->Bus()->InitialVoltage());
					}
				}
			}
			_PS_TRACE("目标函数值 Y ==========");
			_PS_TRACE(ConstraintPowerInjection);
		}

		pair<const PrimitiveNetwork::NodeInfo*, double> NRSolver::EvalDeviation()
		{
			EvalPowerInjection();
			//注意，最小系数不一定是绝对值最小的系数
			int row;
			double dev;
			dev = PowerInjectionDeviation.cwiseAbs().maxCoeff(&row);
			if (row < Block1EquationCount())
				return make_pair(&PNetwork->Nodes(row), dev);
			else
				return make_pair(&PNetwork->PQNodes(row - Block1EquationCount()), dev);
		}

		bool NRSolver::OnIteration()
		{
			GenerateJacobian();
			if (!GenerateNextAnswer()) return false;
			return true;
		}

		void NRSolver::AfterIterations()
		{
			for (int i = 0; i < NodeCount; i++)
				UpdateNodeStatus(i);
		}

		inline double NRSolver::NodeVoltage(int NodeIndex)
		{
			auto& node = PNetwork->Nodes(NodeIndex);
			if (node.Type() == NodeType::PQNode)
				return CurrentAnswer(Block1EquationCount() + node.SubIndex());
			else
				return node.Voltage();
		}

		inline double NRSolver::NodeAngle(int NodeIndex)
		{
			assert(NodeIndex < NodeCount);
			//注意此处区分平衡节点。
			return NodeIndex < NodeCount - 1 ? CurrentAnswer(NodeIndex) : 0;
		}

		void NRSolver::EvalPowerInjection()
		{
			//TODO 计入导纳矩阵可能的不对称性
			//计算各节点的实际注入功率，以及功率偏差 PowerInjectionDeviation
			PowerInjectionDeviation = ConstraintPowerInjection;
			for (auto& node : PSolution->NodeStatus()) node.ClearPowerInjection();
			//遍历所有节点，包括平衡节点
			//此处使用 for 而非 for-each 是为了与数学表达保持一致
			auto &Admittance = PNetwork->Admittance;
			for (int m = 0; m < NodeCount; m++)
			{
				auto &statusM = PSolution->NodeStatus(m);
				auto Um = NodeVoltage(m);
				auto thetaM = NodeAngle(m);
				int subM = Block1EquationCount() + statusM.SubIndex();
				//计算导纳矩阵中非对角元素对应的功率。
				for (int n = m + 1; n < NodeCount; n++)
				{
					// 导纳矩阵是上三角矩阵，row < col
					// TODO 取消对导纳矩阵对称的假定（移相变压器）。
					complexd Y = Admittance.coeff(m, n);
					complexd Y1 = Admittance.coeff(n, m);
					//if (abs(Y) < 1e-10) continue;
					auto UmUn = Um * NodeVoltage(n);
					auto thetaMn = thetaM - NodeAngle(n);
					auto sinMn = sin(thetaMn);
					auto cosMn = cos(thetaMn);
					auto &statusN = PSolution->NodeStatus(n);
					//累计上一次迭代结果对应的注入功率
					statusM.AddPowerInjections(UmUn * (Y.real() * cosMn + Y.imag() * sinMn),
						UmUn * (Y.real() * sinMn - Y.imag() * cosMn));
					statusN.AddPowerInjections(UmUn * (Y1.real() * cosMn - Y1.imag() * sinMn),
						UmUn * (-Y1.real() * sinMn - Y1.imag() * cosMn));
				}
				//计算导纳矩阵中对角元素对应的功率。
				auto UmSqr = Um * Um;
				auto Y = Admittance.coeff(m, m);
				statusM.AddPowerInjections(UmSqr * Y.real(), -UmSqr * Y.imag());
				//生成功率偏差向量 Δy'
				if (statusM.Type() != NodeType::SlackNode)
				{
					PowerInjectionDeviation(m) -= statusM.ActivePowerInjection();
					if (statusM.Type() == NodeType::PQNode)
						PowerInjectionDeviation(subM) -= statusM.ReactivePowerInjection();
				}
			}
			_PS_TRACE("平衡节点 ======");
			_PS_TRACE("有功注入：" << PSolution->NodeStatus(NodeCount - 1).ActivePowerInjection());
			_PS_TRACE("无功注入：" << PSolution->NodeStatus(NodeCount - 1).ReactivePowerInjection());
			_PS_TRACE("偏差 deltaY ==========");
			_PS_TRACE(PowerInjectionDeviation);
		}

		void NRSolver::GenerateJacobian()
		{
			//PQ节点：[dP dQ] = J * [V theta]
			//PV节点：[dP] = J * [theta]
			//       PQ/PV        PQ
			//Δy = [dP ... dP dQ ... dQ]
			//Δx = [theta ... theta V ... V]
			//    / H | N \
			//J = | --+-- |
			//    \ M | L /
			// 起始索引：（注意排除平衡节点）
			// H[0,0], N[NodeCount - 1, 0]
			// M[NodeCount - 1, 0], L[..., ...]
			// 对非对角元有
			// H =  L
			// N = -M
			auto& Admittance = PNetwork->Admittance;
			// 雅可比矩阵非零元。
			vector<Triplet<double>> values;
			values.reserve(JocobianReservedValuesCount);
			// 建议使用 Triplet 构造矩阵，可以取得不错的性能，而且生成的矩阵是压缩过的。
			for (int outer = 0; outer < Admittance.outerSize(); outer++)
				for (SparseMatrix<complexd>::InnerIterator it(Admittance, outer); it; ++it)
				{
					auto m = it.row();
					auto n = it.col();
					// 注意排除平衡节点
					if (m == NodeCount - 1 || n == NodeCount - 1)
						continue;
					auto& nodeM = PSolution->NodeStatus(m);
					auto Um = NodeVoltage(m);
					auto thetaM = NodeAngle(m);
					auto subM = Block1EquationCount() + nodeM.SubIndex();
					//计算非对角元素
					if (n > m)
					{
						complexd Ymn = it.value();
						complexd Ynm = it.value();
						double UmUn = Um * NodeVoltage(n);
						double thetaMn = thetaM - NodeAngle(n);
						double sinMn = sin(thetaMn);
						double cosMn = cos(thetaMn);
						//H(m,n)
						double H = -UmUn * (Ymn.real() * sinMn - Ymn.imag() * cosMn);
						values.emplace_back(m, n, H);
						//H(n,m)
						double Hp = UmUn * (Ynm.real() * sinMn + Ynm.imag() * cosMn);
						values.emplace_back(n, m, Hp);
						//i.e. H(n,m) = -UmUn * (-Y.real() * sinMn - Y.imag() * cosMn)
						//N
						double N = -UmUn * (Ymn.real() * cosMn + Ymn.imag() * sinMn);
						double Np = -UmUn * (Ynm.real() * cosMn - Ynm.imag() * sinMn);
						//cout << "::" << m << "," << n << " = " << Jocobian.insert(1, 0) << endl;
						auto& nodeN = PSolution->NodeStatus(n);
						int subN = Block1EquationCount() + nodeN.SubIndex();
						if (nodeM.Type() == NodeType::PQNode)
						{
							if (nodeN.Type() == NodeType::PQNode)
							{
								//PQ-PQ，子阵具有对称性
								// N
								values.emplace_back(m, subN, N);
								values.emplace_back(n, subM, Np);
								// M = -N
								values.emplace_back(subM, n, -N);
								values.emplace_back(subN, m, -Np);
								// L = H
								values.emplace_back(subM, subN, H);
								values.emplace_back(subN, subM, Hp);
							}
							else {
								//PQ-PV（m,n）
								// M = -N
								values.emplace_back(subM, n, -N);
								//PV-PQ（n,m）
								// N
								values.emplace_back(n, subM, Np);
							}
						}
						else {
							if (nodeN.Type() == NodeType::PQNode)
							{
								//PV-PQ（m,n）
								// N
								values.emplace_back(m, subN, N);
								//PQ-PV（n,m）
								// M
								values.emplace_back(subN, m, -Np);
							}
						}
					}
					else if (n == m)
					{
						complexd Y = it.value();
						double UmSqr = Um * Um;
						//计算对角元素
						//H
						values.emplace_back(m, m, UmSqr * Y.imag() + nodeM.ReactivePowerInjection());
						if (nodeM.Type() == NodeType::PQNode)
						{
							//N
							values.emplace_back(m, subM, -UmSqr * Y.real() - nodeM.ActivePowerInjection());
							//M
							values.emplace_back(subM, m, UmSqr * Y.real() - nodeM.ActivePowerInjection());
							//L
							values.emplace_back(subM, subM, UmSqr * Y.imag() - nodeM.ReactivePowerInjection());
						}
					}
				}
			Jocobian.setFromTriplets(values.begin(), values.end());
			_PS_TRACE("雅可比矩阵 ==========");
			_PS_TRACE(Jocobian);
		}

		bool NRSolver::GenerateNextAnswer()
		{
			SparseLU<SparseMatrix<double>> solver;
			//求解矩阵方程
			Jocobian.makeCompressed();
			solver.compute(Jocobian);
			if (solver.info() != Eigen::Success) return false;
			CorrectionAnswer = solver.solve(PowerInjectionDeviation);
			if (solver.info() != Eigen::Success) return false;
			//计算新的结果
			//注意到 Δy = -J Δx
			//而此处实际解的方程组为 Δy' = J Δx
			//也就是说，Δy = -Δy'
			//Δx = [theta ... theta V ... V]
			CurrentAnswer.head(Block1EquationCount()) -= CorrectionAnswer.head(Block1EquationCount());
			CurrentAnswer.tail(Block2EquationCount()) -= CurrentAnswer.tail(Block2EquationCount()).cwiseProduct(CorrectionAnswer.tail(Block2EquationCount()));
			_PS_TRACE("当前解向量 ==========");
			_PS_TRACE(CurrentAnswer);
			return true;
		}
	}
}