/**
 *  The definition of an abstraction class.
 *
 *  Created by Yinan Li on Aug. 30, 2017.
 *
 *  Hybrid Systems Group, University of Waterloo.
 */

#ifndef _abstraction_h_
#define _abstraction_h_

#include <armadillo>
#include <boost/dynamic_bitset.hpp>
#include <functional>
#include <set>

#include "transition.hpp"
#include "system.hpp"

namespace rocs {

    // /**
    //  * Boolean function defining a set.
    //  */
    // typedef bool (*gset)(const ivec &x);

    /**
     * \brief Weighting function (of edges) of transitions.
     */
    typedef double (*WGT)(std::vector<double> &x0,
    			  std::vector<double> &x1,
    			  std::vector<double> &u0);

    /**
     * \brief An abstraction class of a dynamical system.
     *
     * This is a template class where the typename %S will be replaced by the actual dynamical system type.
     */
    template<typename S>
    class abstraction {

    public:
	grid _x;  /**< A grid of states */
	S *_ptrsys; /**< A pointer to the system object */
	std::vector<int> _labels; /**< A vector of labels for the state grid _x */
	WGT _wf;  /**< Weighting callback function */
	fts _ts;  /**< The finite transition system */

	/**
	 * \brief A constructor.
	 *
	 * Assign user defined system dynamics to an abstraction.
	 *
	 * @param[in] sys The pointer to a system object
	 */
	abstraction(S *sys) : _x(sys->_xdim), _ptrsys(sys), _wf(NULL) {}
	/**
	 * \brief A constructor.
	 *
	 * Assign user defined system dynamics as well as a weighting function to an abstraction.
	 *
	 * @param[in] sys The pointer to a system object
	 * @param[in] wf A weighting function
	 */
	abstraction(S *sys, WGT wf) : _x(sys->_xdim), _ptrsys(sys), _wf(wf) {}

	/**
	 * \brief Initialize the states in the abstraction.
	 *
	 * This includes
	 * - assigning a grid of states by the given bounds and grid size.
	 * - assigning labels
	 *
	 * @param[in] eta An array of grid size
	 * @param[in] xlb An array of lower bounds
	 * @param[in] xub An array of upper bounds
	 */
	void init_state(const double eta[], const double xlb[], const double xub[]) {
	    _x.gridding(eta, xlb, xub);
	    _labels.resize(_x._nv+1, 0); //the extra one is an out-of-domain node
	}

	/**
	 * \brief Initialize the finite transition system.
	 *
	 * Assign the number of predecessing (fts._npre) and successing (fts._npost) states.
	 */
	bool init_transitions() {
	    if (_x._nv > 0 && _ptrsys->_ugrid._nv > 0) {
		_ts.init(_x._nv+1, _ptrsys->_ugrid._nv);
	    } else {
		std::cout << "Transition initialization failed: gridding problem.\n";
		return false;
	    }
	    return true;
	}

	/**
	 * \brief Get state indicies of a given hyper-rectangle area.
	 *
	 * The area is given by its lower and upper bounds.
	 * It is allowed to be out of domain, but only the inside domain part is considered.
	 *
	 * @param xlb The lower bound of the area
	 * @param xub The upper bound of the area
	 * @return a list of indicies.
	 */
	std::vector<size_t> get_discrete_states(const double xlb[], const double xub[]) {
	    ivec states(_x._dim);
	    for (int i = 0; i < _x._dim; ++i) {
		states.setval(i, interval(xlb[i], xub[i]));
	    }

	    return _x.subset(states, false, false);  // allow area out of domain, and collect grids intersect the area
	}

	/**
	 * Get discrete state indicies of a given area:
	 * \f$\{x|g(x)=\text{True}\}\f$.
	 *
	 * @param[in] g The function that determines if x is inside the set
	 * @return a list of indicies.
	 */
	std::vector<size_t> get_discrete_states(std::function<bool(const ivec &)> g) {
	    std::vector<size_t> r;
	    if (_x._nv>0 && !_x._data.empty()) {
		ivec x(_x._dim);
		for (int i = 0; i < _x._nv; ++i) {
		    /* set interval x */
		    for (int j = 0; j < _x._dim; ++j) {
			x.setval(j, interval(_x._data[i][j]-_x._gw[j]/2, _x._data[i][j]+_x._gw[j]/2));
		    }
		    /* test */
		    if (g(x)) {
			r.push_back(i);
		    }
		}

	    } else {
		std::cout << "get_discrete_states: a grid of state space hasn't been iniitlized.\n";
	    }
	    return r;
	}

	/**
	 * \brief A template function that assigns labels to the state grid.
	 *
	 * @param[in] labeling A function that returns the label of a grid id
	 */
	template<typename F>
	void assign_labels(F labeling) {
	    for (size_t i = 0; i < _x._nv; ++i) {
		if(_labels[i] > -1)
		    _labels[i] = labeling(i);
	    }
	}

	/**
	 * \brief Assign the label to the out-of-domain part.
	 *
	 * @param[in] label The label given to the out-of-domain part
	 */
	void assign_label_outofdomain(int label) {
	    _labels[_x._nv] = label;
	}

	/**
	 * \brief Assign transitions with robustness margins e1,e2.
	 *
	 * This function constructs an fts according to the given system.
	 *
	 * @param[in] e1 the robustness margin at the initial point.
	 * @param[in] e2 the robustness margin at the end point.
	 * @return whether construction is successful.
	 */
	bool assign_transitions(const double e1[], const double e2[]) {
	    /* Initialize _ts (the transition system) */
	    if (!init_transitions())
		return false;
	    int n = _x._dim;
	    size_t nx = _x._nv;
	    size_t nu = _ptrsys->_ugrid._nv;
	    // std::cout << "dim=" << n << ',' << "#states=" << nx << ",#inputs=" << nu <<'\n';

	    ivec ie2(n);
	    for (int j = 0; j < n; ++j)
	    	ie2.setval(j, interval(-e2[j],e2[j]));

	    std::vector<double> xmin(n), xmax(n);
	    for (int k = 0; k < n; ++k) {
	    	xmin[k] = _x._valmin[k]-_x._gw[k]/2.0;
	    	xmax[k] = xmin[k] + _x._gw[k]*_x._size[k];
	    }
	    ivec xbds(n);
	    for (int k = 0; k < n; ++k)
		xbds.setval(k, interval(xmin[k], xmax[k]));

	    /* Out-of-domain indicator:
	     * 0:inside, 1:fully outside, 2:partially outside
	     */
	    std::vector<int> out_of_domain(nx*nu,0);

	    ivec y0(n);
	    std::vector<ivec> yt(nu, ivec(n));
	    std::vector<double> ytl(n), ytu(n);
	    size_t r, postid;
	    double idu;
	    std::vector<size_t> il(n), iu(n);
	    std::vector<size_t> stpost(nx*nu, 0), rngpost(nx*nu*n, 0);

	    /* loop state grids */
	    for (size_t row = 0; row < nx; ++row) {
		/* the avoid points (labeled by -1) has no outgoing edges */
		if (_labels[row] == -1)
		    continue;
		/* compute reachable set */
		for (int k = 0; k < n; ++k) {
		    y0.setval(k, interval(_x._data[row][k]-_x._gw[k]/2.0,
		    			  _x._data[row][k]+_x._gw[k]/2.0));
		    // y0.setval(k, interval(_x._data[row][k]-_x._gw[k]/2.0-e1[k],
		    // 			  _x._data[row][k]+_x._gw[k]/2.0+e1[k]));
		}
		_ptrsys->get_reach_set(yt, y0);
		if (!yt.empty()) {
		    /* loop inputs */
		    for (size_t col = 0; col < nu; ++col) {
			// yt[col] += ie2;
			// /* Skip the yt[col] that is out of the domain _bds.
			//  * No need to delete the transitions to avoid area, since they are sinks.
			//  * */
			// if (!xbds.isin(yt[col]))
			//     continue;
			// /** Using this line will probably result in failure in synthesis **/
			// // if (!_x._bds.isin(yt[col]))
			// //     continue;

			if(xbds.isout(yt[col])) { //fully out of domain
			    out_of_domain[row*nu+col] = 1;
			    _ts._npost[row*nu+col] = 1;
			    _ts._ptrpost[row*nu+col] = _ts._ntrans;
			    ++_ts._ntrans;
			    continue;
			}
			if(!xbds.isin(yt[col])) {//partially inside domain
			    out_of_domain[row*nu+col] = 2;
			    /* Take intersection of xbds and yt[col] */
			    for(int k = 0; k < n; ++k) {
				ytl[k] = yt[col][k].getinf()<xmin[k] ? xmin[k] : yt[col][k].getinf();
				ytu[k] = yt[col][k].getsup()>xmax[k] ? xmax[k] : yt[col][k].getsup();
			    }
			} else { //fully inside domain
			    yt[col].getinf(ytl);
			    yt[col].getsup(ytu);
			}
			/* compute the indices of the bounds of the yt[col] */
			_ts._npost[row*nu + col] = 1;
			for (int k = 0; k < n; ++k) {
			    il[k] = static_cast<size_t> ((ytl[k]-_x._valmin[k])/_x._gw[k]+0.5);
			    idu = (ytu[k]-_x._valmin[k])/_x._gw[k] + 0.5;
			    iu[k] = static_cast<size_t> (idu);
			    if (std::fabs(idu - iu[k]) < EPSIVAL)
			    	iu[k] = iu[k] - 1;
			    // if (std::fabs(fmod(ytu[k]-_x._valmin[k], _x._gw[k]) - _x._gw[k]/2.) < EPSIVAL) /* fmod is slow */
			    // 	iu[k] = static_cast<size_t> ((ytu[k]-_x._valmin[k])/_x._gw[k]);
			    // else
			    // 	iu[k] = static_cast<size_t> ((ytu[k]-_x._valmin[k]+_x._gw[k]/2.0)/_x._gw[k]);
			    stpost[row*nu+col] += _x._base[k] * il[k]; /* compute the smallest post ID of all the post grid points */
			    rngpost[n*(row*nu+col)+k] = iu[k] - il[k] + 1;
			    _ts._npost[row*nu+col] *= rngpost[n*(row*nu+col)+k];
			}
			_ts._ptrpost[row*nu+col] = _ts._ntrans;
			_ts._ntrans += _ts._npost[row*nu+col];

			/* Add an out-of-domain transition */
			if(out_of_domain[row*nu+col]) {
			    ++_ts._npost[row*nu+col];
			    ++_ts._ntrans;
			}
			// /********** logging **********/
			// if(row == 0 && col == 16) {
			//     std::cout << "y0=" << y0 << '\n'
			// 	      << "yt=" << yt[col] << '\n'
			// 	      << "out_of_domain=" << out_of_domain[row*nu+col] << '\n';
			//     if(out_of_domain[row*nu+col]>1) {
			// 	std::cout << "Intersection with domain: ";
			// 	for(int k = 0; k < n; ++k)
			// 	    std::cout << '[' << ytl[k] << ',' << ytu[k] << "] ";
			// 	std::cout << '\n';
			//     }
			//     std::cout << "starting address of post transitions: "
			// 	      << _ts._ptrpost[row*nu+col] << '\n';
			//     std::cout << "# of post transitions: "
			// 	      << _ts._npost[row*nu+col] << '\n';
			// }
			// /********** logging **********/
		    } // end input loop
		} else {
		    yt.resize(nu, ivec(n));
		}// end yt empty check
	    }  //end loop state grids
	    /* Assign out-of-domain posts */
	    if(_labels[nx] >= 0) {//assign a self-loop if the out-of-domain node is not avoided
		for(size_t col = 0; col < nu; ++col) {
		    _ts._npost[nx*nu+col] = 1;
		    _ts._ptrpost[nx*nu+col] = _ts._ntrans;
		    ++_ts._ntrans;
		}
	    }

	    /* assign _idpost, _cost and _npre */
	    _ts._idpost.resize(_ts._ntrans);
	    double w;
	    // for (size_t row = 0; row < nx; ++row) {
	    // 	for (size_t col = 0; col < nu; ++col) {
	    // 	    /* assign post grid point IDs to _idpost */
	    // 	    if (_ts._npost[row*nu+col] == 0)
	    // 		continue;
	    // 	    for (int l = 0; l < _ts._npost[row*nu+col]; ++l) {
	    // 		postid = stpost[row*nu+col];
	    // 		r = l;
	    // 		for (int k = 0; k < n; ++k) {
	    // 		    postid += (r % rngpost[n*(row*nu+col)+k])*_x._base[k];
	    // 		    r = r / rngpost[n*(row*nu+col)+k];
	    // 		}
	    // 		_ts._idpost[_ts._ptrpost[row*nu+col]+l] = postid;
	    // 		/* assign cost (worst case): maximum from all posts */
	    // 		if (_wf) {
	    // 		    w = (*_wf)(_x._data[row], _x._data[postid], _ptrsys->_ugrid._data[col]);
	    // 		    _ts._cost[row*nu+col] = _ts._cost[row*nu+col]<w ? w : _ts._cost[row*nu+col];
	    // 		}

	    // 		_ts._npre[postid*nu+col] ++;
	    // 	    }
	    // 	} /* end for col */
	    // } /* end for row */
	    int num_post;
	    size_t ptrout;
	    for (size_t row = 0; row < nx; ++row) {
		for (size_t col = 0; col < nu; ++col) {
		    // /********** logging **********/
		    // if(row == 0 && col == 16) {
		    // 	std::cout << "post nodes: ";
		    // }
		    // /********** logging **********/
		    /* assign post grid point IDs to _idpost */
		    if(out_of_domain[row*nu+col] != 1) {//intersect with domain
			if(out_of_domain[row*nu+col]) {
			    num_post = _ts._npost[row*nu+col] - 1;
			} else {
			    num_post = _ts._npost[row*nu+col];
			}
			for (int l = 0; l < num_post; ++l) {
			    postid = stpost[row*nu+col];
			    r = l;
			    for (int k = 0; k < n; ++k) {
				postid += (r % rngpost[n*(row*nu+col)+k])*_x._base[k];
				r = r / rngpost[n*(row*nu+col)+k];
			    }
			    _ts._idpost[_ts._ptrpost[row*nu+col]+l] = postid;
			    // /********** logging **********/
			    // if(row == 0 && col == 16) {
			    // 	std::cout << "ptrpost[" << _ts._ptrpost[row*nu+col]+l << "]="
			    // 		  << postid << '\n';
			    // }
			    // /********** logging **********/
			    /* assign cost (worst case): maximum from all posts */
			    if (_wf) {
				w = (*_wf)(_x._data[row], _x._data[postid], _ptrsys->_ugrid._data[col]);
				_ts._cost[row*nu+col] = _ts._cost[row*nu+col]<w ? w : _ts._cost[row*nu+col];
			    }

			    ++_ts._npre[postid*nu+col];
			    // /********** logging **********/
			    // if(postid == 0 && col == 16)
			    // 	std::cout << "current # of predecessor of node x=0 under u=16: "
			    // 		  << _ts._npre[postid*nu+col] << '\n';
			    // /********** logging **********/
			}
		    }
		    if(out_of_domain[row*nu+col]) {//fully or partially out of domain
			ptrout = _ts._ptrpost[row*nu+col]+_ts._npost[row*nu+col]-1;
			_ts._idpost[ptrout] = nx; //post node is xout
			++_ts._npre[nx*nu+col];
			// /********** logging **********/
			// if(row == 0 && col == 16) {
			//     std::cout << "ptrpost[" << ptrout << "]="
			// 	      << nx << '\n';
			// }
			// /********** logging **********/
		    }
		} /* end for col */
	    } /* end for row */
	    if(_labels[nx] >= 0) {//assign a self-loop if the out-of-domain node is not avoided
		for(size_t col = 0; col < nu; ++col) {
		    _ts._idpost[_ts._ptrpost[nx*nu+col]] = nx; //xout is a sink
		    ++_ts._npre[nx*nu+col];
		}
	    }

	    /* Determine pre's by post's: loop _npost and _idpost */
	    _ts._idpre.resize(_ts._ntrans);  // initialize the size of pre's
	    /* assign _ptrpre by _npre */
	    size_t sum = 0;
	    for (size_t row = 0; row < nx+1; ++row) {
		for (size_t col = 0; col < nu; ++col) {
		    _ts._ptrpre[row*nu+col] = sum;
		    sum += _ts._npre[row*nu+col];
		}
	    }
	    assert(sum == _ts._ntrans);

	    /* assign _idpre */
	    std::vector<size_t> precount(nu*(nx+1), 0);
	    size_t idtspre;
	    for (size_t row = 0; row < nx+1; ++row) {
		for (size_t col = 0; col < nu; ++col) {
		    // /********** logging **********/
		    // if (row == 0 && col == 16) {
		    // 	std::cout << "Assign pre transitions:\n";
		    // }
		    // /********** logging **********/
		    for (int ip = 0; ip < _ts._npost[row*nu+col]; ++ip) {
			postid = _ts._idpost[_ts._ptrpost[row*nu+col]+ip];
			idtspre = postid * nu + col;
			_ts._idpre[_ts._ptrpre[idtspre]+precount[idtspre]++] = row;
			// precount[idtspre]++;
			// /********** logging **********/
			// if (row == 0 && col == 16) {
			//     std::cout << postid << ": idpre[" << _ts._ptrpre[idtspre]
			// 	      << '+' << precount[idtspre]-1 << "]="
			// 	      << row << '\n';
			// }
			// if(_ts._ptrpre[idtspre]+precount[idtspre]-1 == 46) {
			//     std::cout << "idpre[46] is filled at: "
			// 	      << row << "->(" << col << ")->" << postid
			// 	      << '\n';
			// }
			// /********** logging **********/
		    }
		}
	    }
	    return true;
	}

	/**
	 * \brief Assign transitions by subgridding.
	 * see assign_transitions()
	 *
	 * @param[in] rp[] the pointer to an array of relative subgridding size
	 * @return whether construction is successful.
	 */
	bool assign_transitions_subgridding(const double rp[]) {
	    /*********** logging ***********/
	    // std::fstream logfile;
	    // logfile.open("y.log", std::ios::out | std::ios::ate);
	    // std::fstream logfile2;
	    // logfile2.open("post946.log", std::ios::out | std::ios::ate);
	    /*********** logging ***********/
	    if (!init_transitions())
		return false;
	    int n = _x._dim;
	    size_t nx = _x._nv;
	    size_t nu = _ptrsys->_ugrid._nv;

	    /* compute the number of sub grid points */
	    size_t subnv = 1;
	    std::vector<double> subgw(n);
	    std::vector<size_t> number(n);
	    for (int k = 0; k < n; ++k) {
		number[k] = ceil(1.0 / rp[k]);
		subgw[k] = _x._gw[k] / number[k];
		subnv *= number[k];
	    }
	    // std::cout << "number of subgrids: " << subnv << '\n';

	    /* transition computation by interval subgridding: loop states */
	    std::vector<double> xmin(n);
	    // std::vector<double> xc(n);
	    ivec v(n);
	    std::vector<size_t> subposts;
	    std::vector<std::vector<double>> sub(subnv, std::vector<double> (n));
	    std::vector<size_t>::iterator iter;
	    int np = 0;
	    /* loop state grids */
	    for (size_t row = 0; row < nx; ++row) {
		if (_labels[row] == -1) /* skip the avoid grid points (labeled by -1) */
		    continue;
		/* compute reachable set by interval subgridding */
		if (subnv == 1) {  // no subgridding
		    sub[0] = _x._data[row];
		} else {  // subnv > 1
		    for (int k = 0; k < n; ++k) {
			// xmin[k] = xc[k] - _gw[k]/2. + _rp[k]*_gw[k]/2.;
			xmin[k] = _x._data[row][k] - _x._gw[k]/2. + subgw[k]/2.;
		    }
		    _x.griddingHelper(sub, xmin, subgw, number, subnv);
		}
		std::vector< std::vector<ivec> > ys(subnv, std::vector<ivec> (nu));  // ys[vi][ui]
		for (int vi = 0; vi < subnv; ++vi) {
		    for (int k = 0; k < n; ++k) {
			v.setval(k, interval(sub[vi][k]-subgw[k]/2, sub[vi][k]+subgw[k]/2));
		    } // assign v
		    /* get a list of post intervals w.r.t. different inputs */
		    _ptrsys->get_reach_set(ys[vi], v);
		} // end for loop (subgrid)

		/* loop inputs */
		for (size_t col = 0; col < nu; ++col) {
		    std::set<size_t> posts;
		    /*********** logging ***********/
		    // logfile << col << ":\n";
		    /*********** logging ***********/

		    /* loop subgrids: collect all unique posts */
		    for (int vi = 0; vi < subnv; ++vi) {
			subposts = _x.subset(ys[vi][col], true, false);
			if (subposts.empty()) {  // out of domain
			    np = 0;
			    posts.clear();
			    break;  // jump out of the subgrid loop
			} else {
			    posts.insert(subposts.begin(), subposts.end());
			} // end if
			/*********** logging ***********/
			// logfile << ys[vi][col] << "(";
			// for (int i = 0; i < subposts.size(); ++i)
			//     logfile << subposts[i] << ',';
			// logfile << ")\n";
			/*********** logging ***********/
		    } // end collecting posts

		    /* assign posts */
		    if (!posts.empty()) {
			/*********** logging ***********/
			// logfile2 << col << '(' << posts.size() << "): ";
			/*********** logging ***********/
			_ts._npost[row*nu + col] = posts.size();
			_ts._ptrpost[row*nu + col] = _ts._ntrans;
			_ts._ntrans += posts.size();

			for (std::set<size_t>::iterator it = posts.begin(); it != posts.end(); ++it) {
			    _ts._idpost.push_back(*it);

			    /* record the number of pres for state (*it) */
			    _ts._npre[(*it) * nu + col] ++;
			    /*********** logging ***********/
			    // logfile2 << *it << ", ";
			    /*********** logging ***********/
			}
			// logfile2 << '\n';
		    }  // end transition assignment
		}  // end input loop
		// /*********** logging ***********/
		// logfile << '\n';
		// /*********** logging ***********/

	    }  // end for loop states

	    std::cout << "# of transitions: " << _ts._ntrans << '\n';
	    std::cout << "length of _idpost: " << _ts._idpost.size() << '\n';
	    assert(_ts._idpost.size() == _ts._ntrans);

	    /* assign _cost */
	    double w;
	    size_t postid;
	    for (size_t row = 0; row < nx; ++row) {
		for (size_t col = 0; col < nu; ++col) {
		    /* assign post grid point IDs to _idpost */
		    if (_ts._npost[row*nu+col] == 0)
			continue;
		    for (int l = 0; l < _ts._npost[row*nu+col]; ++l) {
			/* assign cost (worst case): maximum from all posts */
			postid = _ts._ptrpost[row*nu+col]+l;
			if (_wf) {
			    w = (*_wf)(_x._data[row], _x._data[_ts._idpost[postid]], _ptrsys->_ugrid._data[col]);
			    _ts._cost[row*nu+col] = _ts._cost[row*nu+col]<w ? w : _ts._cost[row*nu+col];
			}
		    }
		} /* end for col */
	    } /* end for row */

	    /* determine pre's by post's: loop _ts._npost and _idpost */
	    _ts._idpre.resize(_ts._ntrans);  // initialize the size of pre's
	    size_t sum = 0;
	    for (size_t row = 0; row < nx; ++row) {
		for (size_t col = 0; col < nu; ++col) {
		    _ts._ptrpre[row*nu + col] = sum;
		    sum += _ts._npre[row*nu + col];
		}
	    }
	    /* assign _idpre */
	    std::vector<size_t> precount(_ptrsys->_ugrid._nv*_x._nv, 0);
	    size_t idtspre;
	    for (size_t row = 0; row < nx; ++row) {
		for (size_t col = 0; col < nu; ++col) {
		    for (int ip = 0; ip < _ts._npost[row*nu + col]; ++ip) {

			idtspre = _ts._idpost[_ts._ptrpost[row*nu+col] + ip]*nu + col;
			_ts._idpre[_ts._ptrpre[idtspre] + precount[idtspre]] = row;
			precount[idtspre] ++;
		    }
		}
	    }

	    std::cout << "length of _idpre: " << _ts._idpre.size() << '\n';
	    assert(_ts._idpre.size() == _ts._ntrans);

	    // logfile.close();
	    // logfile2.close();

	    return true;
	}

    }; /* the abstraction class */

} // namespace rocs

#endif


// Boost.Polygon library voronoi_basic_tutorial.cpp file

//          Copyright Andrii Sydorchuk 2010-2012.
// Distributed under the Boost Software License, Version 1.0.
//    (See accompanying file LICENSE_1_0.txt or copy at
//          http://www.boost.org/LICENSE_1_0.txt)

// See http://www.boost.org for updates, documentation, and revision history.

#include <cstdio>
#include <vector>

#include <boost/polygon/voronoi.hpp>
using boost::polygon::voronoi_builder;
using boost::polygon::voronoi_diagram;
using boost::polygon::x;
using boost::polygon::y;
using boost::polygon::low;
using boost::polygon::high;

#include "voronoi_visual_utils.hpp"

struct Point {
  int a;
  int b;
  Point(int x, int y) : a(x), b(y) {}
};

struct Segment {
  Point p0;
  Point p1;
  Segment(int x1, int y1, int x2, int y2) : p0(x1, y1), p1(x2, y2) {}
};

namespace boost {
namespace polygon {

template <>
struct geometry_concept<Point> {
  typedef point_concept type;
};

template <>
struct point_traits<Point> {
  typedef int coordinate_type;

  static inline coordinate_type get(
      const Point& point, orientation_2d orient) {
    return (orient == HORIZONTAL) ? point.a : point.b;
  }
};

template <>
struct geometry_concept<Segment> {
  typedef segment_concept type;
};

template <>
struct segment_traits<Segment> {
  typedef int coordinate_type;
  typedef Point point_type;

  static inline point_type get(const Segment& segment, direction_1d dir) {
    return dir.to_int() ? segment.p1 : segment.p0;
  }
};
}  // polygon
}  // boost

// Traversing Voronoi edges using edge iterator.
int iterate_primary_edges1(const voronoi_diagram<double>& vd) {
  int result = 0;
  for (voronoi_diagram<double>::const_edge_iterator it = vd.edges().begin();
       it != vd.edges().end(); ++it) {
    if (it->is_primary())
      ++result;
  }
  return result;
}

// Traversing Voronoi edges using cell iterator.
int iterate_primary_edges2(const voronoi_diagram<double> &vd) {
  int result = 0;
  for (voronoi_diagram<double>::const_cell_iterator it = vd.cells().begin();
       it != vd.cells().end(); ++it) {
    const voronoi_diagram<double>::cell_type& cell = *it;
    const voronoi_diagram<double>::edge_type* edge = cell.incident_edge();
    // This is convenient way to iterate edges around Voronoi cell.
    do {
      if (edge->is_primary())
        ++result;
      edge = edge->next();
    } while (edge != cell.incident_edge());
  }
  return result;
}

// Traversing Voronoi edges using vertex iterator.
// As opposite to the above two functions this one will not iterate through
// edges without finite endpoints and will iterate only once through edges
// with single finite endpoint.
int iterate_primary_edges3(const voronoi_diagram<double> &vd) {
  int result = 0;
  for (voronoi_diagram<double>::const_vertex_iterator it =
       vd.vertices().begin(); it != vd.vertices().end(); ++it) {
    const voronoi_diagram<double>::vertex_type& vertex = *it;
    const voronoi_diagram<double>::edge_type* edge = vertex.incident_edge();
    // This is convenient way to iterate edges around Voronoi vertex.
    do {
      if (edge->is_primary())
        ++result;
      edge = edge->rot_next();
    } while (edge != vertex.incident_edge());
  }
  return result;
}

int main() {
  // Preparing Input Geometries.
  std::vector<Point> points;
  points.push_back(Point(0, 0));
  points.push_back(Point(1, 6));
  std::vector<Segment> segments;
  segments.push_back(Segment(-4, 5, 5, -1));
  segments.push_back(Segment(3, -11, 13, -1));

  // Construction of the Voronoi Diagram.
  voronoi_diagram<double> vd;
  construct_voronoi(points.begin(), points.end(),
                    segments.begin(), segments.end(),
                    &vd);

  // Traversing Voronoi Graph.
  {
    printf("Traversing Voronoi graph.\n");
    printf("Number of visited primary edges using edge iterator: %d\n",
        iterate_primary_edges1(vd));
    printf("Number of visited primary edges using cell iterator: %d\n",
        iterate_primary_edges2(vd));
    printf("Number of visited primary edges using vertex iterator: %d\n",
        iterate_primary_edges3(vd));
    printf("\n");
  }

  // Using color member of the Voronoi primitives to store the average number
  // of edges around each cell (including secondary edges).
  {
    printf("Number of edges (including secondary) around the Voronoi cells:\n");
    for (voronoi_diagram<double>::const_edge_iterator it = vd.edges().begin();
         it != vd.edges().end(); ++it) {
      std::size_t cnt = it->cell()->color();
      it->cell()->color(cnt + 1);
    }
    for (voronoi_diagram<double>::const_cell_iterator it = vd.cells().begin();
         it != vd.cells().end(); ++it) {
      printf("%lu ", it->color());
    }
    printf("\n");
    printf("\n");
  }

  // Linking Voronoi cells with input geometries.
  {
    unsigned int cell_index = 0;
    for (voronoi_diagram<double>::const_cell_iterator it = vd.cells().begin();
         it != vd.cells().end(); ++it) {
      if (it->contains_point()) {
        if (it->source_category() ==
            boost::polygon::SOURCE_CATEGORY_SINGLE_POINT) {
          std::size_t index = it->source_index();
          Point p = points[index];
          printf("Cell #%u contains a point: (%d, %d).\n",
                 cell_index, x(p), y(p));
        } else if (it->source_category() ==
                   boost::polygon::SOURCE_CATEGORY_SEGMENT_START_POINT) {
          std::size_t index = it->source_index() - points.size();
          Point p0 = low(segments[index]);
          printf("Cell #%u contains segment start point: (%d, %d).\n",
                 cell_index, x(p0), y(p0));
        } else if (it->source_category() ==
                   boost::polygon::SOURCE_CATEGORY_SEGMENT_END_POINT) {
          std::size_t index = it->source_index() - points.size();
          Point p1 = high(segments[index]);
          printf("Cell #%u contains segment end point: (%d, %d).\n",
                 cell_index, x(p1), y(p1));
        }
      } else {
        std::size_t index = it->source_index() - points.size();
        Point p0 = low(segments[index]);
        Point p1 = high(segments[index]);
        printf("Cell #%u contains a segment: ((%d, %d), (%d, %d)). \n",
               cell_index, x(p0), y(p0), x(p1), y(p1));
      }
      ++cell_index;
    }
  }
  return 0;
}


#ifndef DATASTRUCTURES__PROBLEM_H
#define DATASTRUCTURES__PROBLEM_H

#include <util/id.hpp>
#include <util/my_hash.hpp>
#include <util/print_printable.hpp>

#include <unordered_map>
#include <vector>

#include <boost/optional.hpp>

namespace simplex {

struct VariableIDTag { static const std::ptrdiff_t DEFAULT_VALUE = -1; };
typedef util::ID<std::ptrdiff_t, VariableIDTag> VariableID;

class RandomProblemSpecification : public util::print_printable {
public:
	using FloatType = double;

	RandomProblemSpecification(int num_variables, int num_constraints)
		: num_variables(num_variables)
		, num_constraints(num_constraints)
	{ }

	template<typename STREAM>
	void print(STREAM& os) const {
		os
			<< "{ num_variables = " << num_variables
			<< ", num_constraints = " << num_constraints
			<< ", random_seed = "; if (random_seed) os << *random_seed; else os << "none"; os
			<< ", density = " << density
			<< ", constr_coeff_range = {" << constr_coeff_range.first << " - " << constr_coeff_range.second << "}"
			<< ", constr_rhs_coeff_range = {" << constr_rhs_coeff_range.first << " - " << constr_rhs_coeff_range.second << "}"
			<< ", objfunc_coeff_range = {" << objfunc_coeff_range.first << " - " << objfunc_coeff_range.second << "}"
			<< " }"
		;
	}

	int num_variables;
	int num_constraints;

	boost::optional<unsigned long> random_seed = {};

	FloatType density = 0.8;
	std::pair<FloatType, FloatType> constr_coeff_range     = {-100.0, 100.0};
	std::pair<FloatType, FloatType> constr_rhs_coeff_range = {-100.0, 100.0};
	std::pair<FloatType, FloatType> objfunc_coeff_range    = {-100.0, 100.0};
};

class Problem : public util::print_printable {
public:
	using FloatType = double;

	Problem(boost::optional<RandomProblemSpecification> spec_for_gen = {})
		: m_spec_for_gen(spec_for_gen)
	{ }

	std::ptrdiff_t num_variables() const { if(m_max_var_id) return m_max_var_id->getValue() + 1; else return 0; }
	std::ptrdiff_t num_constraints() const { return static_cast<std::ptrdiff_t>(m_constraints.size()); }

	auto& constraints() const { return m_constraints; }
	auto& variables()   const { return m_variables; }
	auto& spec_for_gen() const { return m_spec_for_gen; }

	template<typename CONSTRAINTS>
	void add_constraint(CONSTRAINTS&& c, FloatType rhs) {
		m_constraints.emplace_back(rhs);
		auto& new_constr = m_constraints.back();
		for (const auto& r : c) {
			m_variables[r.first];
			m_max_var_id = m_max_var_id.value_or(r.first);
			m_max_var_id = r.first.getValue() > m_max_var_id->getValue() ? r.first : m_max_var_id;
			new_constr.m_coeffs.emplace_back(r.first, r.second);
		}
	}

	void set_cost(VariableID var, FloatType coeff) {
		m_variables.at(var).m_coeff = coeff;
	}

	bool has_variable(VariableID var) {
		return m_variables.find(var) != end(m_variables);
	}

	/* Prints the file in LP format */
	template<typename STREAM>
	void print(STREAM& os) const {
		const auto& print_term = [&](auto& coeff, auto&& vid) {
			os << (coeff >= 0 ? "+" : "") << coeff << " x" << vid.getValue() << ' ';
		};

		if (spec_for_gen()) {
			// only link comments supported. Will break if spec print is multiline.
			os << "\\ Generated with: " << spec_for_gen() << "\n";
		}

		os << "Maximize\n\t";
		for (const auto vid_and_vinf : variables()) {
			print_term(vid_and_vinf.second.m_coeff, vid_and_vinf.first);
		}

		os << "\n\nSubject to\n";
		for (const auto& c : constraints()) {
			os << '\t';
			for (const auto vid_and_vinf : c.m_coeffs) {
				print_term(vid_and_vinf.second, vid_and_vinf.first);
			}
			os << "<= " << c.m_rhs << '\n';
		}

		os << "\nBounds\n";
		os << "\nEnd";
	}

private:
	struct VariableProperties {
		struct hasher { std::size_t operator()(const VariableProperties& vp) const { return std::hash<FloatType>()(vp.m_coeff); } };
		FloatType m_coeff = 0;
	};

	struct Constraint {
		Constraint(FloatType rhs)
			: m_coeffs()
			, m_rhs(rhs)
		{ }

		std::vector<std::pair<VariableID, FloatType>> m_coeffs;
		FloatType m_rhs;
	};

	std::unordered_map<VariableID, VariableProperties, typename util::MyHash<VariableID>::type> m_variables = {};
	std::vector<Constraint> m_constraints = {};
	boost::optional<VariableID> m_max_var_id = {};
	boost::optional<RandomProblemSpecification> m_spec_for_gen;
};

class Assignments {

};

Problem generate_random_problem(const RandomProblemSpecification& prob_spec);

Problem make_small_sample_problem();

Problem pad_with_zeroes_modulo(Problem&& p, std::ptrdiff_t height_modulus, std::ptrdiff_t width_modulus);

} // end namespace simplex

#endif /* DATASTRUCTURES__PROBLEM_H */


/**
 * @file 	SteinerTreeHeuristic.cpp
 * @author 	Fabian Wegscheider
 * @date 	Jul 14, 2017
 */

#include <sstream>
#include <boost/heap/priority_queue.hpp>
#include "SteinerTreeHeuristic.h"

using std::pair;
typename graph_traits<Graph>::out_edge_iterator it, itEnd;

/**
 * Data that is stored in one node of the heap. Contains an integer and a double.
 * Comparisons are made by the double, smaller has higher priority
 */
struct heapData
{
	pair<int,double> data;

    heapData(pair<int,double> p):
        data(p)
    {}

    bool operator<(heapData const & rhs) const {
        return data.second > rhs.data.second;
    }
};

/*
 * Implementation of the improved Shortest-Path-Heuristic. A standard binary heap
 * from boost is used and keys are never decreased, but instead new nodes are pushed
 * whenever a better path is found. The resulting tree is stored implicitly in
 * treePredecessors.
 */
double SteinerTreeHeuristic::computeSteinerTree(Graph& g, int numVertices,
		int treePredecessors[], vector<int>& terminals, int root) {

	assert(NULL != treePredecessors);

	//terminals are scanned
	int numTerminals = terminals.size();
	bool* isTerminal = new bool[numVertices]();
	for (int i = 0; i < numTerminals; ++i) {
		assert(terminals[i] < numVertices);
		isTerminal[terminals[i]] = true;
	}

	if (!isTerminal[root]) {
		delete[] isTerminal;
		return -1;
	}

	heap::priority_queue<heapData> heap;

	double* distances = new double[numVertices];
	int* predecessors = new int[numVertices];
	predecessors[root] = root;
	double treeCost = 0.;

	//initialization of heap, distance array and predecessors array. note: a -1
	//in predecessors means that the vertex has not been connected to the tree yet
	for (int i = 0; i < numVertices; ++i) {
		if (i == root) {
			heap.push(std::make_pair(i,0.));
			distances[i] = 0.;
			treePredecessors[i] = root;
		} else {
			heap.push(std::make_pair(i, std::numeric_limits<double>::infinity()));
			distances[i] = std::numeric_limits<double>::infinity();
			treePredecessors[i] = -1;
		}
	}

	property_map<Graph, edge_weight_t>::type weights = get(edge_weight, g);
	property_map<Graph, vertex_index_t>::type index = get(vertex_index, g);

	int connectedTerminals = 1;

	//here the actual algorithm starts
	while (connectedTerminals < numTerminals) {
		assert(!heap.empty());	//heap should never be empty before tree contains all terminals

		pair<int,double> min = heap.top().data;
		heap.pop();

		//if we scan a terminal, all vertices on shortest path to subtree
		//are added to heap with weight 0 and included into the tree
		if (min.first != root && isTerminal[min.first] && treePredecessors[min.first] == -1) {
			++connectedTerminals;

			min.second = 0.;
			int nextVertex = predecessors[min.first];
			treePredecessors[min.first] = predecessors[min.first];

			auto edgeInfo = edge(predecessors[min.first], min.first, g);
			assert(edgeInfo.second);
			treeCost += weights[edgeInfo.first];

			//here all vertices on path from new terminal to tree are added again with key=0
			while (treePredecessors[nextVertex] == -1) {
				heap.push(std::make_pair(nextVertex, 0.));
				treePredecessors[nextVertex] = predecessors[nextVertex];
				treeCost += weights[edge(predecessors[nextVertex], nextVertex, g).first];
				nextVertex = predecessors[nextVertex];
			}
		}

		//this is basicalles the usual dijkstra with pushes instead of decreasekey operations
		for (tie(it, itEnd) = out_edges(*(vertices(g).first+min.first), g);
				it != itEnd; ++it) {
			double tmp = min.second + weights[*it];
			int targetIndex = index[target(*it, g)];
			//we only consider vertices which have not been added to the tree yet
			if (treePredecessors[targetIndex] == -1 && tmp < distances[targetIndex]) {
				distances[targetIndex] = tmp;
				predecessors[targetIndex] = min.first;
				heap.push(std::make_pair(targetIndex, tmp));
			}
		}
	}

	delete[] isTerminal;
	delete[] distances;
	delete[] predecessors;

	return treeCost;
}

bool SteinerTreeHeuristic::testAndPrintTree(Graph& g, int numVertices, int treePredecessors[],
		vector<int>& terminals, int root, string& edgeString) {

	std::stringstream stream;
	bool* visited = new bool[numVertices]();

	for (unsigned int i = 0; i < terminals.size(); ++i) {
		int curr = terminals[i];
		int vertexCount = 1;
		int edgeCount = 0;
		while (curr != root && vertexCount != numVertices && curr != -1) {
			if (!visited[curr]) {
				++edgeCount;
				//edges are added to string and every 50 lines a linebreak is
				//added for the sake of readability
				if (edgeCount % 50 == 0) stream << std::endl;
				stream << "(" << curr << "," << treePredecessors[curr] << ") ";
			}
			curr = treePredecessors[curr];
			++vertexCount;
		}

		//if a cycle is found or not all terminals are connected, result is false
		if (vertexCount == numVertices || curr == -1){
			delete[] visited;
			return false;
		}
	}

	edgeString = stream.str();

	delete[] visited;
	return true;
}

bool SteinerTreeHeuristic::testTree(Graph& g, int numVertices, int treePredecessors[],
		vector<int>& terminals, int root) {

		bool* visited = new bool[numVertices]();

		for (unsigned int i = 0; i < terminals.size(); ++i) {
			int curr = terminals[i];
			int vertexCount = 1;
			int edgeCount = 0;
			while (curr != root && vertexCount != numVertices && curr != -1) {
				if (!visited[curr]) ++edgeCount;
				curr = treePredecessors[curr];
				++vertexCount;
			}

			//if a cycle is found or not all terminals are connected, result is false
			if (vertexCount == numVertices || curr == -1){
				delete[] visited;
				return false;
			}
		}

		delete[] visited;
		return true;
}


/**
 * This file contains an example of graphlab used for gibbs sampling
 * in a pairwise markov random field to denoise a synthetic noisy
 * image.
 *
 *  \author Joseph Gonzalez
 */

// INCLUDES ===================================================================>

// Including Standard Libraries
#include <cstdlib>
#include <cassert>
#include <cstring>
#include <fstream>
#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
#include <limits>
#include <cmath>

#include <boost/program_options.hpp>
#include <boost/bind.hpp>

#include <graphlab.hpp>

// ============================================================================>
//  Support code for loopy belief propagation
#include "image.hpp"

// Include the macro for the for each operation
#include <graphlab/macros_def.hpp>

// Determine the shared ata location of the shared edge factor and the
// number of samples
enum constants {EDGE_FACTOR_ID, NUM_SAMPLES_ID};

// STRUCTS (Edge and Vertex data) =============================================>
/**
 * The data associated with each variable in the pairwise markov
 * random field
 */
struct vertex_data {
  size_t asg;
  size_t updates;
  unsigned char color;
  graphlab::unary_factor potential;
  //! store the Roa-Blackwell conditional belief estimate
  graphlab::unary_factor belief;
  std::vector<double> counts;
  vertex_data() :
    asg(0), updates(0), color(0) { }
}; // End of vertex data

/**
 * The data associated with each directed edge in the pairwise markov
 * random field
 */
struct edge_data { };

typedef graphlab::graph< vertex_data, edge_data> graph_type;
typedef graphlab::types<graph_type> gl;

// GraphLab Update Function ===================================================>

/** Construct denoising ising model based on the image */
void construct_graph(image& img,
                     size_t num_rings,
                     double sigma,
                     gl::graph& graph);

/**
 * The core belief propagation update function.  This update satisfies
 * the graphlab update_function interface.
 */
template<bool UseCallback>
void gibbs_update(gl::iscope& scope,
                  gl::icallback& scheduler,
                  gl::ishared_data* shared_data);

/**
 *  The colore schedule is used by the set scheduler to describe an
 *  execution that updates all vertices of the same color (given by a
 *  checkerboard) pattern) in parallel.
 */
void color_schedule(gl::set_scheduler &sched);
// This ugly global variable is needed for the color_schedule
size_t nsamples;

// Command Line Parsing =======================================================>
struct options {
  size_t ncpus;
  size_t samples;
  size_t num_rings;
  size_t rows;
  size_t cols;
  double sigma;
  double lambda;
  std::string smoothing;
  std::string engine;
  std::string scope;
  std::string scheduler;
  std::string orig_fn;
  std::string noisy_fn;
  std::string rb_pred_fn;
  std::string counts_pred_fn;
  std::string pred_type;
};

/**
 * Parse the command line arguments.  Returns false if there was a
 * problem in parsing command line arguments
 */
bool parse_command_line(int argc, char** argv, options& opts);

/**
 * Display the program options
 */
void display_options(options& opts);

// MAIN =======================================================================>
int main(int argc, char** argv) {
  std::cout << "This program uses gibbs sampling to denoise a synthetic image."
            << std::endl;

  // set the global logger
  global_logger().set_log_level(LOG_WARNING);
  global_logger().set_log_to_console(true);

  // Parse command line arguments --------------------------------------------->
  options opts;
  bool success = parse_command_line(argc, argv, opts);
  if(!success)  return EXIT_FAILURE;
  display_options(opts);

  // Create synthetic images -------------------------------------------------->
  // Creating image for denoising
  std::cout << "Creating a synethic image." << std::endl;
  image img(opts.rows, opts.cols);
  img.paint_sunset(opts.num_rings);
  std::cout << "Saving image. " << std::endl;
  img.save(opts.orig_fn.c_str());
  std::cout << "Corrupting Image. " << std::endl;
  img.corrupt(opts.sigma);
  std::cout << "Saving corrupted image. " << std::endl;
  img.save(opts.noisy_fn.c_str());
  img.save_vec("corrupted.tsv");

  // Create the graph --------------------------------------------------------->
  gl::graph graph;
  std::cout << "Constructing pairwise Markov Random Field. " << std::endl;
  construct_graph(img, opts.num_rings, opts.sigma, graph);
  image coloring_img(opts.rows, opts.cols);
  for(size_t i = 0; i < graph.num_vertices(); ++i)
    coloring_img.pixel(i) = graph.vertex_data(i).color;
  coloring_img.save("coloring.pgm");

  // Setup global shared variables -------------------------------------------->
  gl::thread_shared_data sdm;
  // Initialize the edge agreement factor
  std::cout << "Initializing shared edge agreement factor. " << std::endl;
  // dummy variables 0 and 1 and num_rings by num_rings
  graphlab::binary_factor edge_potential(0, opts.num_rings,
                                         0, opts.num_rings);

  // Set the smoothing type
  if(opts.smoothing == "square") {
    edge_potential.set_as_agreement(opts.lambda);
  } else if (opts.smoothing == "laplace") {
    edge_potential.set_as_laplace(opts.lambda);
  } else {
    std::cout << "Invalid smoothing stype!" << std::endl;
    return EXIT_FAILURE;
  }
  std::cout << edge_potential << std::endl;
  sdm.set_constant(EDGE_FACTOR_ID, graphlab::any(edge_potential));
  sdm.set_constant(NUM_SAMPLES_ID, graphlab::any(opts.samples));

  // Create the engine -------------------------------------------------------->
  gl::iengine* engine = NULL;
  if(opts.scheduler == "set") {
    nsamples = opts.samples;

    // Here we use the set scheduler which has some additional special setup
    engine =
      graphlab::engine_factory::new_engine(opts.engine,
                                           "set",
                                           "null",
                                           graph,
                                           opts.ncpus);
    if(engine == NULL) {
      std::cout << "Unable to construct engine!" << std::endl;
      return EXIT_FAILURE;
    }
    // Set the shared data manager for the engine
    engine->set_shared_data_manager(&sdm);
    // Attach the correct schedule
    engine->get_scheduler().set_option(gl::scheduler_options::SCHEDULING_FUNCTION,
                                       (void*)color_schedule);

  } else {
    // Here we use the other schedulers which use a more default setup
    engine =
      graphlab::engine_factory::new_engine(opts.engine,
                                           opts.scheduler,
                                           opts.scope,
                                           graph,
                                           opts.ncpus);
    if(engine == NULL) {
      std::cout << "Unable to construct engine!" << std::endl;
      return EXIT_FAILURE;
    }
    // Set the shared data manager for the engine
    engine->set_shared_data_manager(&sdm);

    // Create a shuffled set of vertices
    std::vector<gl::vertex_id_t> vertex_ids(opts.rows * opts.cols);
    for(size_t i = 0; i < vertex_ids.size(); ++i) vertex_ids[i] = i;
    std::random_shuffle(vertex_ids.begin(), vertex_ids.end());

    // Add the tasks in shuffled order (helps with fifo scheduler)
    const bool use_callback = true;
    double residual = 1.0;
    engine->get_scheduler().add_tasks(vertex_ids,
                                      gibbs_update<use_callback>,
                                      residual);
  }

  // Running the engine ------------------------------------------------------->
  std::cout << "Running the engine. " << std::endl;
  graphlab::timer timer; timer.start();
  engine->start();
  double runtime = timer.current_time();
  size_t update_count = engine->last_update_count();
  std::cout << "Finished Running engine in " << runtime
            << " seconds." << std::endl
            << "Total updates: " << update_count << std::endl
            << "Efficiency: " << (double(update_count) / runtime)
            << " updates per second "
            << std::endl;

  // Saving the output -------------------------------------------------------->
  std::cout << "Rendering the cleaned image. " << std::endl;
  image rb_img(opts.rows, opts.cols);
  image counts_img(opts.rows, opts.cols);
  if(opts.pred_type == "map") {
    for(size_t v = 0; v < graph.num_vertices(); ++v) {
      const vertex_data& vdata = graph.vertex_data(v);
      rb_img.pixel(v) = vdata.belief.max_asg();
      counts_img.pixel(v) =
        std::max_element(vdata.counts.begin(), vdata.counts.end()) -
        vdata.counts.begin();

    }
  } else if(opts.pred_type == "exp") {
    for(size_t v = 0; v < graph.num_vertices(); ++v) {
      const vertex_data& vdata = graph.vertex_data(v);
      rb_img.pixel(v) = vdata.belief.expectation();
      double expectation = 0;
      for(size_t i = 0; i < vdata.counts.size(); ++i)
        expectation += i * vdata.counts[i];
      counts_img.pixel(v) = expectation;
    }
  } else {
    std::cout << "Invalid prediction type! : " << opts.pred_type
              << std::endl;
    return EXIT_FAILURE;
  }

  std::cout << "Saving cleaned image. " << std::endl;
  rb_img.save(opts.rb_pred_fn.c_str());
  counts_img.save(opts.counts_pred_fn.c_str());

  std::cout << "Done!" << std::endl;
  return EXIT_SUCCESS;
} // End of main

// Implementations ============================================================>
template<bool UseCallback>
void gibbs_update(gl::iscope& scope,
                  gl::icallback& scheduler,
                  gl::ishared_data* shared_data) {

  assert(shared_data != NULL);
  // Get the shared data
  size_t samples = shared_data->get_constant(NUM_SAMPLES_ID).as<size_t>();

  // Grab the state from the scope -------------------------------------------->
  // Get the vertex data
  vertex_data& vdata = scope.vertex_data();

  // If this vertex has been updated sufficiently often then we don't
  // update again
  if(vdata.updates >= samples) return;

  // Construct the conditional------------------------------------------------->
  // Initialize the conditional as the node potential
  graphlab::unary_factor conditional(vdata.potential);

  const graphlab::binary_factor& edge_factor =
    shared_data->get_constant(EDGE_FACTOR_ID).as<graphlab::binary_factor>();

  // Condition on all neighbor assignments
  foreach(graphlab::edge_id_t edgeid, scope.in_edge_ids()) {
    graphlab::vertex_id_t source = scope.source(edgeid);
    // Get the neighboring assignment
    const vertex_data& neighbor = scope.neighbor_vertex_data(source);
    // Get the edge factor
    conditional.condition(edge_factor, neighbor.asg);
    conditional.normalize();
  }

  // Generate a sample -------------------------------------------------------->
  size_t sample = conditional.sample();
  vdata.asg = sample;  // Record the sample
  vdata.updates++;     // Increment the updates
  vdata.counts[vdata.asg]++; // increment the counts

  // Rao-Blackwell belief estimate
  vdata.belief.plus(conditional);

  // Finalize belief on last update ------------------------------------------->
  if(vdata.updates == samples ) { // If this is the last update
    vdata.belief.normalize();
    double Z = 0;
    for(size_t i = 0; i < vdata.counts.size(); ++i)
      Z += vdata.counts[i];
    for(size_t i = 0; i < vdata.counts.size(); ++i)
      vdata.counts[i] /= Z;
  } else if(UseCallback) {
    // Otherwise reschedule self
    double residual = 1.0;
    gl::update_task task(scope.vertex(), gibbs_update<UseCallback> );
    scheduler.add_task(task, residual);
  }
} // end of BP_update

bool parse_command_line(int argc, char** argv, options& opts) {
  // Because typing is painful
  namespace boost_po = boost::program_options;
  // Create a description for this program
  boost_po::options_description
    desc("Denoise a randomly generated image using Gibbs Sampling.");
  // Set the program options
  desc.add_options()
    ("help",   "produce this help message")
    ("ncpus",  boost_po::value<size_t>(&(opts.ncpus))->default_value(2),
     "Number of cpus to use.")
    ("samples",  boost_po::value<size_t>(&(opts.samples))->default_value(100),
     "Number of samples to generate")
    ("rings",  boost_po::value<size_t>(&(opts.num_rings))->default_value(5),
     "Number of rings in the noisy image")
    ("rows",  boost_po::value<size_t>(&(opts.rows))->default_value(200),
     "Number of rows in the noisy image")
    ("cols",  boost_po::value<size_t>(&(opts.cols))->default_value(200),
     "Number of columns in the noisy image")
    ("sigma",  boost_po::value<double>(&(opts.sigma))->default_value(1.2),
     "Standard deviation of noise.")
    ("lambda",  boost_po::value<double>(&(opts.lambda))->default_value(3),
     "Smoothness parameter (larger => smoother).")
    ("smoothing",
     boost_po::value<std::string>(&(opts.smoothing))->default_value("square"),
     "Options are {square, laplace}")
    ("engine",
     boost_po::value<std::string>(&(opts.engine))->default_value("threaded"),
     "Options are {threaded, sequential}")
    ("scope",
     boost_po::value<std::string>(&(opts.scope))->default_value("edge"),
     "Options are {vertex, edge, full}")
    ("scheduler",
     boost_po::value<std::string>(&(opts.scheduler))->default_value("fifo"),
     "Options are {fifo, priority, sampling}")
    ("orig",
     boost_po::value<std::string>(&(opts.orig_fn))->default_value("source_img.pgm"),
     "Original image file name.")
    ("noisy",
     boost_po::value<std::string>(&(opts.noisy_fn))->default_value("noisy_img.pgm"),
     "Noisy image file name.")
    ("rb_pred",
     boost_po::value<std::string>(&(opts.rb_pred_fn))->
     default_value("rb_pred_img.pgm"),
     "Predicted image file name for the Rao-Blackwell estimator.")
    ("counts_pred",
     boost_po::value<std::string>(&(opts.counts_pred_fn))->
     default_value("counts_pred_img.pgm"),
     "Predicted image file name for the raw counts estimator.")
    ("pred_type",
     boost_po::value<std::string>(&(opts.pred_type))->default_value("map"),
     "Predicted image type {map, exp}");
  // Parse the arguments
  boost_po::variables_map vm;
  boost_po::store(boost_po::parse_command_line(argc, argv, desc), vm);
  boost_po::notify(vm);
  if(vm.count("help")) {
    std::cout << desc << std::endl;
    return false;
  }
  return true;
} // end of parse command line arguments

void display_options(options& opts) {
  std::cout << "ncpus:          " << opts.ncpus << std::endl
            << "samples:        " << opts.samples << std::endl
            << "num_rings:      " << opts.num_rings << std::endl
            << "rows:           " << opts.rows << std::endl
            << "cols:           " << opts.cols << std::endl
            << "sigma:          " << opts.sigma << std::endl
            << "lambda:         " << opts.lambda << std::endl
            << "smoothing:      " << opts.smoothing << std::endl
            << "engine:         " << opts.engine << std::endl
            << "scope:          " << opts.scope << std::endl
            << "scheduler:      " << opts.scheduler << std::endl
            << "orig_fn:        " << opts.orig_fn << std::endl
            << "noisy_fn:       " << opts.noisy_fn << std::endl
            << "rb_pred_fn:     " << opts.rb_pred_fn << std::endl
            << "counts_pred_fn: " << opts.counts_pred_fn << std::endl
            << "pred_type:      " << opts.pred_type << std::endl;
}

// STRUCTS (Edge and Vertex data) =============================================>

void construct_graph(image& img,
                     size_t num_rings,
                     double sigma,
                     gl::graph& graph) {

  // Initialize the vertex data to somethine sensible
  vertex_data vdata;

  vdata.belief.resize(num_rings);
  vdata.belief.uniform(-std::numeric_limits<double>::max());
  vdata.potential.resize(num_rings);
  vdata.potential.uniform();
  vdata.potential.normalize();
  vdata.counts.resize(num_rings, 0);
  for(size_t i = 0; i < num_rings; ++i)
    vdata.counts[i] = 0;

  // Add all the vertices
  double sigmaSq = sigma*sigma;
  for(size_t i = 0; i < img.rows(); ++i) {
    for(size_t j = 0; j < img.cols(); ++j) {
      // initialize the potential and belief
      uint32_t pixel_id = img.vertid(i, j);
      vdata.potential.var() = vdata.belief.var() = pixel_id;
      // Determine the color (in checkerboard form)
      vdata.color = (i%2 == 0) ^ (j%2 == 0)? 1 : 0;
      // Set the node potential
      double obs = img.pixel(i, j);
      for(size_t pred = 0; pred < num_rings; ++pred) {
        vdata.potential.logP(pred) =
          -(obs - pred)*(obs - pred) / (2.0 * sigmaSq);
      }
      vdata.potential.normalize();
      // Set the initial assignment
      vdata.asg = vdata.potential.sample();
      // Store the actual data in the graph
      size_t vert_id = graph.add_vertex(vdata);
      // Ensure that we are using a consistent numbering
      assert(vert_id == pixel_id);
    } // end of for j in cols
  } // end of for i in rows
  // Construct an edge blob

  // Add all the edges
  for(size_t i = 0; i < img.rows(); ++i) {
    for(size_t j = 0; j < img.cols(); ++j) {
      size_t vertid = img.vertid(i,j);
      if(i-1 < img.rows())
        graph.add_edge(vertid, img.vertid(i-1, j));
      if(i+1 < img.rows())
        graph.add_edge(vertid, img.vertid(i+1, j));
      if(j-1 < img.cols())
        graph.add_edge(vertid, img.vertid(i, j-1));
      if(j+1 < img.cols())
        graph.add_edge(vertid, img.vertid(i, j+1));
    } // end of for j in cols
  } // end of for i in rows
} // End of construct graph

/** Color selection function */
template<size_t color>
bool select_color(graphlab::vertex_id_t v,
                  const vertex_data& vdata) {
  return vdata.color == color;
}

/**
 * The scheduling function used by the set scheduler to execute in
 * colored order
 */
void color_schedule(gl::set_scheduler &sched) {
  // There is a two coloring of this model
  // All sets must be created before scheduling calls
  std::vector<gl::ivertex_set*> colorsets(2);
  // Colors do not change during execution
  bool fixed_selection = true;
  gl::selector_function select_color_0( select_color<0> );
  colorsets[0] = &sched.attach(gl::rvset(select_color_0, fixed_selection),
                               sched.root_set());
  gl::selector_function select_color_1( select_color<1> );
  colorsets[1] = &sched.attach(gl::rvset(select_color_1, fixed_selection),
                               sched.root_set());
  // Actually construct the color sets
  sched.init();

  // Build the execution plan
  gl::execution_plan eplan;
  for(size_t i = 0; i < colorsets.size(); ++i) {
    std::cout << "Set " << i << " has "
              << colorsets[i]->size() << " vertices."
              << std::endl;
    const bool use_callback = false;
    eplan.execute(*colorsets[i], gibbs_update<use_callback>);
  }

  // Compile the execution plan
  graphlab::timer timer; timer.start();
  eplan.generate_plan(sched.get_graph(), sched.num_cpus());
  std::cout << "Execution plan compiled in "
            << timer.current_time() << " seconds" << std::endl;

  // Run the plan nsamples times
  for(size_t iter = 0; iter < nsamples; ++iter)
    sched.execute_plan(eplan);
}


#include <fstream>
#include <algorithm>
#include <iterator>

#include <boost/functional/value_factory.hpp>
#include <boost/array.hpp>

#include <CGAL/algorithm.h>
#include <CGAL/point_generators_3.h>
#include <CGAL/Exact_predicates_inexact_constructions_kernel.h>

#include <CGAL/AABB_tree.h>
#include <CGAL/AABB_traits.h>
#include <CGAL/Surface_mesh.h>
#include <CGAL/AABB_face_graph_triangle_primitive.h>
#include <CGAL/AABB_halfedge_graph_segment_primitive.h>
#include <CGAL/Timer.h>

typedef CGAL::Epick K;
typedef K::FT FT;
typedef K::Point_3 Point;
typedef K::Vector_3 Vector;
typedef K::Segment_3 Segment;
typedef K::Ray_3 Ray;
typedef CGAL::Surface_mesh<CGAL::Point_3<CGAL::Epick> > Mesh;
typedef CGAL::AABB_face_graph_triangle_primitive<Mesh,
CGAL::Default,
CGAL::Tag_false> T_Primitive;
typedef CGAL::AABB_traits<K, T_Primitive> T_Traits;
typedef CGAL::AABB_tree<T_Traits> T_Tree;
typedef T_Tree::Primitive_id T_Primitive_id;

typedef CGAL::AABB_halfedge_graph_segment_primitive<Mesh,
CGAL::Default,
CGAL::Tag_false> E_Primitive;
typedef CGAL::AABB_traits<K, E_Primitive> E_Traits;
typedef CGAL::AABB_tree<E_Traits> E_Tree;
typedef E_Tree::Primitive_id E_Primitive_id;

int main()
{
  CGAL::Surface_mesh<CGAL::Point_3<CGAL::Epick> > m1, m2;
  std::ifstream in("data/cube.off");
  if(in)
    in >> m1;
  else{
    std::cout << "error reading cube" << std::endl;
    return 1;
  }
  in.close();
  in.open(CGAL::data_file_path("meshes/tetrahedron.off"));
  if(in)
    in >> m2;
  else{
    std::cout << "error reading tetrahedron" << std::endl;
    return 1;
  }
  in.close();
  T_Tree tree(faces(m1).first, faces(m1).second, m1);
  tree.insert(faces(m2).first, faces(m2).second, m2);
  tree.build();
  T_Tree::Bounding_box bbox = tree.bbox();
  Point bbox_center((bbox.xmin() + bbox.xmax()) / 2,
                     (bbox.ymin() + bbox.ymax()) / 2,
                     (bbox.zmin() + bbox.zmax()) / 2);
  std::vector< T_Primitive_id > intersections;
  Ray ray(bbox_center+Vector(3,-0.25,0),bbox_center+Vector(-3,+0.25,0));
  tree.all_intersected_primitives(ray,
                                  std::back_inserter(intersections));
  E_Tree e_tree(edges(m1).first, edges(m1).second, m1);
  e_tree.insert(edges(m2).first, edges(m2).second, m2);
  e_tree.build();
  std::vector< E_Primitive_id > e_intersections;
  Ray e_ray(Point(0,0,0),Point(0,1,1));
  e_tree.all_intersected_primitives(e_ray,
                                  std::back_inserter(e_intersections));

  return 0;
}


/********************************************************
  Stanford Driving Software
  Copyright (c) 2011 Stanford University
  All rights reserved.

  Redistribution and use in source and binary forms, with
  or without modification, are permitted provided that the
  following conditions are met:

* Redistributions of source code must retain the above
  copyright notice, this list of conditions and the
  following disclaimer.
* Redistributions in binary form must reproduce the above
  copyright notice, this list of conditions and the
  following disclaimer in the documentation and/or other
  materials provided with the distribution.
* The names of the contributors may not be used to endorse
  or promote products derived from this software
  without specific prior written permission.

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
  CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
  WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
  PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
  INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
  OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
  DAMAGE.
 ********************************************************/

#include <cmath>
#include <assert.h>
#include <set>
#include <map>
#include <sstream>
#include <Eigen/Dense>

#include <vlrException.h>
#include <aw_geometry_2d.h>

#include <aw_match_to_graph.hpp>
#include <aw_Route.h>
#include <aw_Vehicle.h>

using std::set;
using std::map;
using namespace Eigen;

namespace vlr {

#define TRACE(str) std::cout << "[M2G] "<< str << std::endl

RoutePlanner::RndfEdge* match_2_graph( const double x, const double y,
               const double yaw, const std::map<int, RoutePlanner::RndfEdge*>& edges,
               double& d,
               double& cp_x,
               double& cp_y,
               double& offset_to_first_point,
               int& sign,
               const double penalty_meter_per_degree  )
{
  //  std::cout << "penalty_meter_per_degree " << penalty_meter_per_degree << "\n";

	if (edges.empty()) {return NULL;}

  double yaw_r = yaw * ( 180. / M_PI );

  RoutePlanner::RndfEdge* best_edge = NULL;
  d = std::numeric_limits<double>::max();
  offset_to_first_point = -1;
  double yaw_r_cropped = yaw_r;

  while( yaw_r_cropped < 0. ) yaw_r_cropped += 360.;
  while( yaw_r_cropped > 360. ) yaw_r_cropped -= 360.;
  // BETTER would be  yaw_r_cropped -= floor( yaw_r_cropped/360 ) * 360;

  double best_angle = 123;

//  std::cout << "edges.size == " << edges.size() << "\n";

  for( std::map<int, RoutePlanner::RndfEdge*>::const_iterator it = edges.begin(); it != edges.end(); it++ )
    {
//      assert( it->second != 0 );

      if(!it->second) {
        char* a=NULL; *a=1;
      }
      double x1 = it->second->fromVertex()->x();
      double y1 = it->second->fromVertex()->y();

      double x2 = it->second->toVertex()->x();
      double y2 = it->second->toVertex()->y();

      double pc_x, pc_y, offset;

      closest_point_on_segment( x1, y1, x2, y2, x, y, pc_x, pc_y, offset );

      double delta_x =  x2 - x1;
      double delta_y =  y2 - y1;

      double dist = sqrt( (pc_y - y) * (pc_y - y) + (pc_x - x) * (pc_x - x) );

      // std::cout << "signed dist " << dist << " vs " << dist_signded << " vs " << dist_signded2 << "\n";

      double angle = atan2( delta_y, delta_x ) * 180./M_PI;

      double angle_diff = fabs( angle - yaw_r_cropped );

      while( angle_diff > 360. ) angle_diff -= 360;
      while( angle_diff < 0. ) angle_diff += 360;

//       assert( angle_diff < 360. );
//       assert( angle_diff > 0. );

      if( angle_diff > 180. ) angle_diff = 360. - angle_diff;

      dist += penalty_meter_per_degree * angle_diff;

      if( fabs( dist ) < fabs( d ) )
  {
    d = dist;
    best_edge = it->second;
    offset_to_first_point = offset;
    cp_x = pc_x;
    cp_y = pc_y;
    best_angle = angle;
  }
    }

//  assert(best_edge != 0);
  if(!best_edge) {
    std::stringstream s;
    s << "Object at " << x << ", " << y << " (" << yaw << ") could not be matched to graph.";
    throw VLRException(s.str());
  }
  double x1 = best_edge->fromVertex()->x();
  double y1 = best_edge->fromVertex()->y();

  double x2 = best_edge->toVertex()->x();
  double y2 = best_edge->toVertex()->y();

  double delta_x =  x2 - x1;
  double delta_y =  y2 - y1;

  double norm = sqrt( delta_x*delta_x + delta_y*delta_y );
  if (norm < 0.0001) {
//  	std::cerr << " norm very small! " << norm << std::endl;
  }
  double nx = -delta_y / norm;
  double ny = delta_x / norm;

  double dist_signded = nx * ( x - cp_x ) + ny * ( y - cp_y );

  if( dist_signded < 0 ) sign = -1;
  else sign = 1;

  //  std::cout << "signed distance " << sign * d << "\n";

  return best_edge;
}

std::map< RoutePlanner::RndfEdge*, double > multi_match_2_graph( const Vehicle& veh, const set< RndfEdge* >& edges )
{
	map< RndfEdge*, double > matched_edges;

	Vector2d veh_m( veh.xMatchedFrom(), veh.yMatchedFrom() );
	Point_2 veh_p( veh.xMatchedFrom(), veh.yMatchedFrom() );

	for (set< RndfEdge* >::iterator it = edges.begin(); it != edges.end(); ++it)
	{
		RndfEdge* edge = *it;
		RndfVertex* v1 = edge->fromVertex();
		RndfVertex* v2 = edge->toVertex();
		Point_2 p1( v1->x(), v1->y() );
		Point_2 p2( v2->x(), v2->y() );
		Vector_2 vec = p2 - p1;
		Point_2 m = p1 + vec/2.;

		if ( rect_rect_X( Vector2d( m.x(), m.y() ), angle(vec), edge->width(), edge->length(),
							veh_m, veh.yawMatchedFrom(), veh.width(), veh.length()) )
		{
			//TRACE(edge->name() << " intersects with " << veh.id);
			Line_2 l(p1, p2);
			Point_2 pp = l.projection( veh_p );
			double offset = std::sqrt( (pp - p1).squared_length() );
			if ( offset < 0. ) offset = 0.;
			if ( offset > edge->length() ) offset = edge->length();

			matched_edges.insert( make_pair( edge, offset ) );

		} else {
			//TRACE(edge->name() << " NOT intersects with " << veh.id);
		}
	}

	return matched_edges;

//	************ Alte Version die auch die Ausrichtungsdiffernz berücksichtigt hat
//
//	Point_2 pos(veh.xMatchedFrom(), veh.yMatchedFrom());
//	Vector_2 vec_y( cos(veh.yawMatchedFrom()         ) * veh.length / 2. + 1.0, sin(veh.yawMatchedFrom()         ) * veh.length / 2. + 1.0) ;
//	Vector_2 vec_x( cos(veh.yawMatchedFrom() + M_PI_2) * veh.width  / 2. + 1.0, sin(veh.yawMatchedFrom() + M_PI_2) * veh.width  / 2. + 1.0) ;
//	vector< Point_2 > coords(5, pos);
//	coords[1] = pos + vec_x + vec_y;
//	coords[2] = pos + vec_x - vec_y;
//	coords[3] = pos - vec_x + vec_y;
//	coords[4] = pos - vec_x - vec_y;
//	map< RndfEdge*, double > matched_edges;
//
//	for (vector< Point_2 >::iterator pos_it = coords.begin(); pos_it != coords.end(); ++pos_it) {
//		Point_2 pos = *pos_it;
//		for (set<RndfEdge*>::iterator it = edges.begin(); it != edges.end(); ++it)
//		{
//			RndfEdge* edge = *it;
//			RndfVertex* v1 = edge->fromVertex();
//			RndfVertex* v2 = edge->toVertex();
//			Segment_2 seg( Point_2( v1->x(), v1->y() ), Point_2( v2->x(), v2->y() ) );
//
//			// Abstand überprüfen
//			if ( squared_distance(seg, pos) > sqr( edge->laneWidth()/2. ) ) continue;
//
//			// Winkelabweichung überprüfen (max Abweichung 60°)
//			//		if ( deltaAngle( yaw, angle( seg.to_vector() ) ) > M_PI_4 * 0.66 ) continue;
//
//			// Offset bestimmen
//			Line_2 lin =  seg.supporting_line();
//			Point_2 pp = lin.projection( pos );
//			//		if ( ! seg.has_on( pp ) ) continue;
//			if ( squared_distance( seg, pp ) > 0.0001) continue;
//			double offset = sqrt( (pp - seg.source()).squared_length() );
////			if (offset < 0.) offset = 0.;
////			if (offset > edge->length()) offset = edge->length();
//			if (matched_edges.find(edge) == matched_edges.end() || matched_edges[ edge ] > offset) {
//				matched_edges[ edge ] = offset;
//			}
//		}
//	}
//
//	return matched_edges;
}

map<double, RndfEdge*> getBestMatchings( const Vehicle& veh, const map<int, RndfEdge*>& edges, const double penalty_meter_per_degree, double min_score, double max_score )
{
	Point_2 pos( veh.xMatchedFrom(), veh.yMatchedFrom() );
	map<double, RndfEdge*> result;

	for (map<int, RndfEdge*>::const_iterator it = edges.begin(); it != edges.end(); ++it) {
		RndfEdge* edge = it->second;
		RndfVertex* v1 = edge->fromVertex();
		RndfVertex* v2 = edge->toVertex();
		Segment_2 seg( Point_2( v1->x(), v1->y() ), Point_2( v2->x(), v2->y() ) );

		// Score berechen und Edge einsortieren
		double score = std::sqrt( squared_distance( pos, seg) ) + deltaAngle( angle( seg.to_vector() ), veh.yawMatchedFrom() ) * penalty_meter_per_degree;
		if (score <= min_score || score > max_score) continue;
		result[ score ] = edge;
	}

	return result;
}

double calcDeltaAngle(const Vehicle& veh, const RndfEdge* edge)
{
	const RndfVertex* v1 = edge->fromVertex();
	const RndfVertex* v2 = edge->toVertex();
	Vector_2 vec( Point_2( v1->x(), v1->y() ), Point_2( v2->x(), v2->y() ) );
	return deltaAngle( angle( vec ), veh.yawMatchedFrom() );
}

void closest_point_on_segment( double p1_x, double p1_y, double p2_x, double p2_y, double p_x, double p_y, double& pc_x, double& pc_y, double& offset_to_first_point )
{
  double vx = p2_x - p1_x;
  double vy = p2_y - p1_y;

  double wx = p_x - p1_x;
  double wy = p_y - p1_y;

  double c1 = dot_product( wx, wy, vx, vy );
  if( c1 <= 0 )
    {
      pc_x = p1_x;
      pc_y = p1_y;
      offset_to_first_point = 0;
      return;
    }

  double c2 = dot_product( vx, vy, vx, vy );
  if( c2 <= c1 )
    {
      pc_x = p2_x;
      pc_y = p2_y;
      offset_to_first_point = sqrt( c2 );
      return;
    }

  double b = c1/c2;

  offset_to_first_point = b * sqrt( c2 );

  //assert(offset_to_first_point >= 0.0);
  //assert(offset_to_first_point <= sqrt(vx*vx+vy*vy));

  pc_x = p1_x + b * vx;
  pc_y = p1_y + b * vy;

}

double dot_product( double p1_x, double p1_y, double p2_x, double p2_y )
{
  return p1_x * p2_x + p1_y * p2_y;
}

} // namespace vlr


//==================================================================================================
/*
  EVE - Expressive Vector Engine
  Copyright : EVE Contributors & Maintainers
  SPDX-License-Identifier: MIT
*/
//==================================================================================================
#include <cmath>
#include <boost/math/special_functions/bessel.hpp>
#include <eve/function/cyl_bessel_j1.hpp>
#include <eve/constant/valmin.hpp>
#include <eve/constant/valmax.hpp>
#include "producers.hpp"
#include <cmath>

TTS_CASE_TPL("wide random check on cyl_bessel_j1", EVE_TYPE)
{
  auto bj1 = [](auto x) -> EVE_VALUE {return boost::math::cyl_bessel_j(EVE_VALUE(1), x); };

  eve::uniform_prng<EVE_VALUE> p(0, 4);
  TTS_RANGE_CHECK(p, bj1, eve::cyl_bessel_j1);
}


#include "world.hpp"

#define _USE_MATH_DEFINES
#include <cmath>

#include <Eigen/Core>

#include "fox/counter.hpp"
#include "bullet.hpp"
#include "asteroid.hpp"
#include "ship.hpp"

//------------------------------------------------------------------------------
world::world(events::manager_interface *emi, const std::string &data_root,
	int width, int height)
{
	this->emi = emi;
	this->data_root = data_root;
	this->width = width;
	this->height = height;
	this->generator = std::mt19937_64(std::random_device{}());

	player_ship = nullptr;
	main_counter = nullptr;
}

//------------------------------------------------------------------------------
world::~world()
{
	deinit();
}

//------------------------------------------------------------------------------
void world::update()
{
	float dt = main_counter->update();

	for(asteroid *a : asteroids)
	{
		a->pos = a->pos + dt * a->vel;
		if(a->pos[0] > (float)width)
			a->pos[0] -= (float)width;
		if(a->pos[1] > (float)height)
			a->pos[1] -= (float)height;
	}
}

//------------------------------------------------------------------------------
void world::init()
{
	std::uniform_real_distribution<float> dist(asteroid::speed_min, asteroid::speed_max);
	asteroid *a = new asteroid();
	a->speed = dist(generator);
	dist = std::uniform_real_distribution<float>(0.0f, (float)M_PI);
	a->direction = dist(generator);
	a->vel[0] = cosf(a->direction);
	a->vel[1] = sinf(a->direction);
	a->vel = a->vel * a->speed;
	dist = std::uniform_real_distribution<float>(0.0f, (float)width);
	a->pos[0] = dist(generator);
	dist = std::uniform_real_distribution<float>(0.0f, (float)height);
	a->pos[1] = dist(generator);
	a->radius = a->init_radius;
	asteroids.push_back(a);

	player_ship = new ship();
	player_ship->radius = 16.0f;
	for(int i = 0; i < 100; i++)
	{
		dist = std::uniform_real_distribution<float>(0.0f, (float)width);
		player_ship->pos[0] = dist(generator);
		dist = std::uniform_real_distribution<float>(0.0f, (float)height);
		player_ship->pos[1] = dist(generator);

		Eigen::Vector2f d = player_ship->pos - a->pos;

		if(d.norm() > 2.0f * a->radius)
			break;
	}

	main_counter = new fox::counter();
}

//------------------------------------------------------------------------------
void world::deinit()
{
	for(asteroid *a : asteroids)
	{
		delete a;
	}
	asteroids.clear();

	for(bullet *b : bullets)
	{
		delete b;
	}
	bullets.clear();

	delete player_ship;
	delete main_counter;
}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_GENERIC_FUNCTION_REMQUO_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_GENERIC_FUNCTION_REMQUO_HPP_INCLUDED

#include <boost/simd/detail/brigand.hpp>
#include <boost/simd/function/divides.hpp>
#include <boost/simd/function/if_allbits_else.hpp>
#include <boost/simd/function/is_eqz.hpp>
#include <boost/simd/function/fms.hpp>
#include <boost/simd/function/is_invalid.hpp>
#include <boost/simd/function/logical_or.hpp>
#include <boost/simd/function/minus.hpp>
#include <boost/simd/function/multiplies.hpp>
#include <boost/simd/function/nearbyint.hpp>
#include <boost/simd/function/toint.hpp>
#include <boost/simd/detail/dispatch/function/overload.hpp>
#include <boost/simd/detail/dispatch/meta/as_integer.hpp>
#include <boost/config.hpp>
#include <utility>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;

  BOOST_DISPATCH_OVERLOAD ( remquo_
                          , (typename A0)
                          , bd::cpu_
                          , bd::generic_<bd::floating_<A0> >
                          , bd::generic_<bd::floating_<A0> >
                          )
  {
    using quo_t = bd::as_integer_t<A0, signed>;
    using result_type = std::pair<A0,quo_t>;

    BOOST_FORCEINLINE std::pair<A0, quo_t> operator() ( A0 const& a0, A0 const& a1
                                      ) const BOOST_NOEXCEPT
    {
      A0 const d = nearbyint(a0/a1);

#if defined(BOOST_SIMD_NO_INVALIDS)
      A0  a2 = if_allbits_else(is_eqz(a1),  fma(-d, a1, a0));
#else
      A0  a2 = if_allbits_else(logical_or(is_invalid(a0), is_eqz(a1)),  fma(-d, a1, a0));
#endif
      return {a2, toint(d)};
    }
  };
} } }

#endif


#include "config/scalar.h"
#include "controls/foc.h"
#include "controls/pi_control.h"
#include "controls/six_step.h"
#include "controls/space_vector_modulation.h"
#include "gui.h"
#include "motor.h"
#include "util/clarke_transform.h"
#include "util/conversions.h"
#include "util/math_constants.h"
#include "util/rotation.h"
#include "util/time.h"
#include "wrappers/sdl_context.h"
#include "wrappers/sdl_imgui.h"
#include "wrappers/sdl_imgui_context.h"
#include <Eigen/Dense>
#include <absl/strings/str_format.h>
#include <array>
#include <glad/glad.h>
#include <implot.h>
#include <iostream>

using namespace biro;

int main(int argc, char* argv[]) {
    SimState state;
    init_sim_state(&state);
    state.foc.i_controller_params = make_motor_pi_params(
        /*bandwidth=*/10000,
        /*resistance=*/state.motor.params.phase_resistance,
        /*inductance=*/state.motor.params.phase_inductance);

    VizData viz_data;
    init_viz_data(&viz_data);

    VizOptions viz_options;

    wrappers::SdlContext sdl_context("Biro Motor Simulator",
                                     /*width=*/1920 / 2,
                                     /*height=*/1080 / 2);

    if (sdl_context.status_ != 0) {
        return -1;
    }

    if (gladLoadGL() == 0) {
        printf("Failed to initialize OpenGL loader!\n");
        exit(-1);
    }

    wrappers::SdlImguiContext imgui_context(sdl_context.window_,
                                            sdl_context.gl_context_);

    bool quit_flag = false;
    while (!quit_flag) {
        quit_flag = wrappers::process_sdl_imgui_events(sdl_context.window_);
        if (quit_flag) {
            break;
        }

        wrappers::sdl_imgui_newframe(sdl_context.window_);

        if (!state.paused) {
            update_rolling_buffers(state.time, state.board, state.motor,
                                   state.foc, &viz_data.rolling_buffers);
        }
        run_gui(viz_data, &viz_options, &state);

        if (!state.paused) {
            for (int i = 0; i < state.step_multiplier; ++i) {
                const bool new_pwm_cycle =
                    step_pwm_state(state.dt, &state.board.pwm);

                std::array<bool, 3> gate_command = {};

                // update relevant commutation modes
                if (state.commutation_mode == kCommutationModeManual) {
                    // maintain the existing the gate command
                    // which has probably been set from the GUI
                    gate_command = state.board.gate.commanded;
                }

                if (state.commutation_mode == kCommutationModeSixStep) {
                    gate_command = six_step_commutate(
                        get_electrical_angle(state.motor.params.num_pole_pairs,
                                             state.motor.kinematic.rotor_angle),
                        state.six_step_phase_advance);
                }

                if (state.commutation_mode == kCommutationModeFOC) {
                    if (periodic_timer(state.foc.period, state.dt,
                                       &state.foc.timer)) {
                        Scalar desired_torque = state.foc_desired_torque;
                        if (state.foc_use_cogging_compensation) {
                            desired_torque -= interp_cogging_torque(
                                state.motor.kinematic.rotor_angle,
                                state.motor.params.cogging_torque_map);
                        }

                        std::complex<Scalar> desired_current_qd;
                        if (state.foc_non_sinusoidal_drive_mode) {
                            desired_current_qd =
                                get_desired_current_qd_non_sinusoidal(
                                    desired_torque, state.motor);
                        } else {
                            desired_current_qd = get_desired_current_qd(
                                desired_torque,
                                state.motor.params.normed_bEmf_coeffs(0));
                        }

                        step_foc_current_controller(desired_current_qd,
                                                    state.motor, &state.foc);

                        // anti-windup
                        if (state.foc_pi_anti_windup) {
                            const Scalar voltage_qd_norm =
                                std::abs(state.foc.voltage_qd);
                            if (voltage_qd_norm >
                                state.board.bus_voltage * kClarkeScale) {
                                const std::complex<Scalar>
                                    voltage_qd_saturation =
                                        state.foc.voltage_qd *
                                        (state.board.bus_voltage *
                                         kClarkeScale / voltage_qd_norm);
                                // unwind integral terms
                                pi_unwind(state.foc.i_controller_params,
                                          voltage_qd_saturation.real(),
                                          &state.foc.iq_controller);
                                pi_unwind(state.foc.i_controller_params,
                                          voltage_qd_saturation.imag(),
                                          &state.foc.id_controller);
                            }
                        }
                    }

                    // assert the requested qd voltage with PWM
                    if (new_pwm_cycle) {
                        const std::complex<Scalar> inv_park_transform =
                            get_rotation(get_q_axis_electrical_angle(
                                state.motor.params.num_pole_pairs,
                                state.motor.kinematic.rotor_angle));

                        std::complex<Scalar> voltage_ab =
                            inv_park_transform * state.foc.voltage_qd;

                        if (state.foc_use_qd_decoupling) {
                            const std::complex<Scalar> existing_back_emf_ab =
                                clarke_transform(
                                    state.motor.electrical.normed_bEmfs) *
                                state.motor.kinematic.rotor_angular_vel;
                            voltage_ab += existing_back_emf_ab;
                        }

                        state.board.pwm.duties =
                            get_pwm_duties(state.board.bus_voltage, voltage_ab);
                    }

                    gate_command = get_pwm_gate_command(state.board.pwm);
                }

                state.board.gate.commanded = gate_command;
                update_gate_state(state.dt, &state.board.gate);

                const auto pole_voltages = get_pole_voltages(
                    state.board.bus_voltage,
                    state.motor.electrical.phase_currents, state.board.gate);

                step_motor(state.dt, state.load_torque, pole_voltages,
                           &state.motor);

                state.time += state.dt;
            }
        }

        ImGui::Render();
        int display_w, display_h;
        SDL_GetWindowSize(sdl_context.window_, &display_w, &display_h);
        glViewport(0, 0, display_w, display_h);
        glClearColor(0.0, 0.0, 0.0, 0.0);
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
        glEnable(GL_DEPTH_TEST);
        ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
        SDL_GL_SwapWindow(sdl_context.window_);
    }

    return 0;
}


/**
 * \file libsanm/tensor_svd.cpp
 * This file is part of SANM, a symbolic asymptotic numerical solver.
 */

#include "libsanm/tensor_svd.h"
#include "libsanm/tensor_impl_helper.h"

#include <Eigen/Dense>
#include <Eigen/SVD>

using namespace sanm;

namespace {
void assign_check(size_t& dst, size_t v) {
    if (!dst) {
        dst = v;
    } else {
        sanm_assert(dst == v);
    }
}

fp_t* get_full_ptr(TensorND& dst, const StSparseLinearTrans& lt) {
    dst = lt.check_batched(true).as_full();
    return const_cast<fp_t*>(dst.ptr());
}

fp_t clip_div(fp_t x, fp_t y) {
    constexpr fp_t eps = 1e-12;
    return x * y / (y * y + eps);
}

fp_t* ptr_offset(fp_t* p, size_t off) {
    return p ? p + off : nullptr;
}

#define FOR4_BEGIN(i, j, k, l, n)              \
    do {                                       \
        for (size_t i = 0; i < n; ++i)         \
            for (size_t j = 0; j < n; ++j)     \
                for (size_t k = 0; k < n; ++k) \
                    for (size_t l = 0; l < n; ++l)
#define FOR4_END() \
    }              \
    while (0)
}  // anonymous namespace

const TensorND& TensorND::compute_batched_svd_w(TensorND& u, TensorND& s,
                                                TensorND& w,
                                                bool require_rotation) const {
    SANM_SCOPED_PROFILER("batched_svd_w");
    sanm_assert(rank() == 3);
    const size_t batch = shape(0), n = shape(1);
    sanm_assert(n >= 2);
    sanm_assert(shape(2) == n, "not square matrices: %s",
                shape().str().c_str());
    u.set_shape({batch, n, n});
    s.set_shape({batch, n});
    w.set_shape({batch, n, n});

    auto mptr = ptr();
    auto uptr = u.woptr(), sptr = s.woptr(), wptr = w.woptr();
    auto run = [mptr, uptr, sptr, wptr, batch, n,
                require_rotation]<Eigen::Index sn>() {
        ScopedAllowMalloc scoped_allow_malloc;
        Eigen::JacobiSVD<Eigen::Matrix<fp_t, sn, sn>> svd(
                n, n, Eigen::ComputeFullU | Eigen::ComputeFullV);
        Eigen::Matrix<fp_t, sn, sn> msrcT(n, n);
        // note: with EIGEN_USE_LAPACKE, it JacobiSVD seems to fallback to MKL
        // which would not be so slow
        sanm_assert(n <= 16, "TODO: switch to BDCSVD for large matrices");
        if (sn != Eigen::Dynamic) {
            if (static_cast<Eigen::Index>(n) != sn) {
                // hint the compiler that n equals to sn
                __builtin_unreachable();
            }
        }
        for (size_t i = 0; i < batch; ++i) {
            EigenMat<sn, sn> msrc(const_cast<fp_t*>(mptr) + i * n * n, n, n);
            msrcT = msrc;  // eigen svd class does not take Map inputs
            svd.compute(msrcT);
            EigenMat<sn, 1> ms(sptr + i * n, n, 1);
            EigenMat<sn, sn> muT(uptr + i * n * n, n, n);
            EigenMat<sn, sn> mwT(wptr + i * n * n, n, n);
            // eigen uses column-major and computes M' = VSU'
            ms = svd.singularValues();
            muT = svd.matrixV().transpose();
            if (require_rotation && ((svd.matrixU().determinant() < 0) !=
                                     (svd.matrixV().determinant() < 0))) {
#if 1
                // negate some singular values so that det(w) = 1
                constexpr fp_t EPS = 1e-3;
                int best_idx = -1, best_idx_nr = n + 1;
                for (size_t i = 0; i < n; ++i) {
                    size_t j = i + 1;
                    // ms already sorted
                    while (j < n && std::fabs(ms(i) - ms(j)) < EPS) {
                        ++j;
                    }
                    int nr = j - i;
                    // best case is to negate an odd number of smallest singular
                    // values (so si+sj != 0 in the hessian);
                    // otherwise negate one value whose has the least
                    // repetitionss
                    if (nr <= best_idx_nr ||
                        (nr == best_idx_nr + 1 && nr % 2 == 1)) {
                        best_idx = i;
                        best_idx_nr = nr;
                        if (nr == 1) {
                            break;
                        }
                    }
                    i = j;
                }
                if (best_idx_nr == 1 || best_idx_nr % 2 == 0) {
                    ms(best_idx) = -ms(best_idx);
                    muT.row(best_idx) = -muT.row(best_idx);
                } else {
                    for (int i = best_idx; i < best_idx + best_idx_nr; ++i) {
                        ms(i) = -ms(i);
                        muT.row(i) = -muT.row(i);
                    }
                }
#else
                ms(n - 1) = -ms(n - 1);
                muT.row(n - 1) = -muT.row(n - 1);
#endif
            }
            // w=uv', w'=vu'
            mwT.noalias() = svd.matrixU() * muT;
        }
    };
#define CASE(sn)                  \
    do {                          \
        if (n == sn) {            \
            run.operator()<sn>(); \
            return *this;         \
        }                         \
    } while (0)
    CASE(2);
    CASE(3);
#undef CASE
    run.operator()<Eigen::Dynamic>();
    return *this;
}

void sanm::svd_w_grad_revmode(StSparseLinearTrans& grad, const TensorND& mU,
                              const TensorND& mS, const TensorND& mW,
                              const StSparseLinearTrans& mdU,
                              const StSparseLinearTrans& mdS,
                              const StSparseLinearTrans& mdW) {
    SANM_SCOPED_PROFILER("batched_svd_w_grad");
    auto run = [&]<Eigen::Index sn>() {
        // extract sizes and pointers
        const size_t batch = mU.shape(0), n = mU.shape(1);
        if (sn != Eigen::Dynamic && static_cast<Eigen::Index>(n) != sn) {
            // hint the compiler that n equals to sn
            __builtin_unreachable();
        }
        TensorND mdU_full, mdS_full, mdW_full, mdM;
        fp_t *mdU_ptr = nullptr, *mdS_ptr = nullptr, *mdW_ptr = nullptr;
        size_t out_dim = 0;
        if (mdU.valid()) {
            mdU_ptr = get_full_ptr(mdU_full, mdU);
            assign_check(out_dim, mdU.out_dim());
        }
        if (mdS.valid()) {
            mdS_ptr = get_full_ptr(mdS_full, mdS);
            assign_check(out_dim, mdS.out_dim());
        }
        if (mdW.valid()) {
            mdW_ptr = get_full_ptr(mdW_full, mdW);
            assign_check(out_dim, mdW.out_dim());
        }
        sanm_assert(out_dim, "no output grad");
        if (grad.valid()) {
            mdM = grad.check_batched(true).as_full();
            sanm_assert(mdM.shape() == (TensorShape{batch, out_dim, n * n}));
            grad = {};
        } else {
            mdM.set_shape({batch, out_dim, n * n}).fill_with_inplace(0);
        }
        auto mdMptr = mdM.rwptr();
        auto mU_ptr = const_cast<fp_t*>(mU.ptr()),
             mS_ptr = const_cast<fp_t*>(mS.ptr()),
             mW_ptr = const_cast<fp_t*>(mW.ptr());

        // allocate temporaries
        ScopedAllowMalloc scoped_allow_malloc;
        constexpr Eigen::Index snsqr = sn == Eigen::Dynamic ? sn : sn * sn;
        Eigen::Matrix<fp_t, sn, sn> cV(n, n);
        Eigen::Matrix<fp_t, snsqr, sn> dsdmT(n * n, n);
        Eigen::Matrix<fp_t, snsqr, snsqr> dxdmT(n * n, n * n),
                dydmT(n * n, n * n), dwdyT(n * n, n * n), dudxT(n * n, n * n),
                tmpn2(n * n, n * n);
        scoped_allow_malloc.disallow();

        // matrix variables:
        // prefix c means current (this batch)
        // prefix d means derivative/jacobian
        // suffix t/T means transposed (eigen uses col-major)

        for (size_t ib = 0; ib < batch; ++ib) {
            size_t ib_off = ib * n * n;
            EigenMat<sn, sn> cUt(mU_ptr + ib_off, n, n),
                    cWt(mW_ptr + ib_off, n, n);
            EigenMat<sn, 1> cS(mS_ptr + ib * n, n, 1);
            EigenMat<snsqr, Eigen::Dynamic> cdMt(mdMptr + ib_off * out_dim,
                                                 n * n, out_dim);
            cV.noalias() = cWt * cUt.transpose();
            if (mdS_ptr) {
                for (size_t si = 0; si < n; ++si) {
                    for (size_t mi = 0; mi < n; ++mi) {
                        for (size_t mj = 0; mj < n; ++mj) {
                            dsdmT(mi * n + mj, si) = cUt(si, mi) * cV(mj, si);
                        }
                    }
                }
                EigenMat<sn, Eigen::Dynamic> cdSt(mdS_ptr + ib * n * out_dim, n,
                                                  out_dim);
                cdMt.noalias() += dsdmT * cdSt;
            }
            if (mdW_ptr || mdU_ptr) {
                FOR4_BEGIN(i, j, k, l, n) {
                    fp_t cij = cUt(i, k) * cV(l, j), cji = cUt(j, k) * cV(l, i),
                         si = cS(i), sj = cS(j);
                    if (mdW_ptr) {
                        dydmT(k * n + l, i * n + j) =
                                i == j ? 0 : clip_div(cij - cji, si + sj);
                        dwdyT(k * n + l, i * n + j) = cUt(k, i) * cV(j, l);
                    }
                    if (mdU_ptr) {
                        dudxT(k * n + l, i * n + j) = l == j ? cUt(k, i) : 0;
                        dxdmT(k * n + l, i * n + j) =
                                i == j ? 0
                                       : clip_div(cij * sj + cji * si,
                                                  sj * sj - si * si);
                    }
                }
                FOR4_END();

                if (mdW_ptr) {
                    tmpn2.noalias() = dydmT * dwdyT;
                    EigenMat<snsqr, Eigen::Dynamic> cdWt(
                            mdW_ptr + ib_off * out_dim, n * n, out_dim);
                    cdMt.noalias() += tmpn2 * cdWt;
                }
                if (mdU_ptr) {
                    tmpn2.noalias() = dxdmT * dudxT;
                    EigenMat<snsqr, Eigen::Dynamic> cdUt(
                            mdU_ptr + ib_off * out_dim, n * n, out_dim);
                    cdMt.noalias() += tmpn2 * cdUt;
                }
            }
        }

        grad.reset(StSparseLinearTrans::FULL, true, mdM);
    };

    size_t n = mU.shape(1);
    sanm_assert(n >= 2);
#define CASE(sn)                  \
    do {                          \
        if (n == sn) {            \
            run.operator()<sn>(); \
            return;               \
        }                         \
    } while (0)
    CASE(2);
    CASE(3);
#undef CASE
    run.operator()<Eigen::Dynamic>();
}

void sanm::svd_w_taylor_fwd(TensorND& mUk, TensorND& mSk, TensorND& mWk,
                            const TensorND& mMk, const TensorND& mMbiask,
                            const TensorND& mU0, const TensorND& mS0,
                            const TensorND& mW0, const TensorND* mBu,
                            const TensorND& mBw) {
    SANM_SCOPED_PROFILER("svd_w_taylor_fwd");
    auto run = [&]<Eigen::Index sn>() {
        const size_t batch = mMk.shape(0), n = mMk.shape(1);
        if (sn != Eigen::Dynamic && static_cast<Eigen::Index>(n) != sn) {
            // hint the compiler that n equals to sn
            __builtin_unreachable();
        }
        fp_t *mUkptr = mBu ? mUk.woptr() : nullptr,
             *mSkptr = mBu ? mSk.woptr() : nullptr, *mWkptr = mWk.woptr();
#define DEF(x) auto x##ptr = const_cast<fp_t*>(x.ptr())
        DEF(mMk);
        DEF(mMbiask);
        DEF(mU0);
        DEF(mS0);
        DEF(mW0);
        DEF(mBw);
#undef DEF
        fp_t* mBuptr = mBu ? const_cast<fp_t*>(mBu->ptr()) : nullptr;

        // matrix variables:
        // prefix c means current (this batch)
        // suffix t/T means transposed (eigen uses col-major)

        // allocate temporaries
        ScopedAllowMalloc scoped_allow_malloc;
        Eigen::Matrix<fp_t, sn, sn> cV0(n, n), eqbT(n, n), tmp0(n, n),
                tmp1(n, n);
        MatrixMultiProduct<Eigen::Matrix<fp_t, sn, sn>> mprod(n, n);
        scoped_allow_malloc.disallow();

        for (size_t ib = 0; ib < batch; ++ib) {
            size_t ib_off = ib * n * n;
            EigenMat<sn, 1> cS0(mS0ptr + ib * n, n, 1),
                    cSk(mSkptr + ib * n, n, 1);
#define DEF(x) EigenMat<sn, sn> c##x##T(ptr_offset(m##x##ptr, ib_off), n, n)
            DEF(Uk);
            DEF(Wk);
            DEF(Mk);
            DEF(Mbiask);
            DEF(U0);
            DEF(W0);
            DEF(Bu);
            DEF(Bw);
#undef DEF
            cV0.noalias() = cW0T * cU0T.transpose();
            tmp0 = cMkT - cMbiaskT;
            mprod.init(cV0.transpose(), tmp0).mul_r_to(eqbT, cU0T.transpose());
            {
                // solve Wk
                auto& rhs = tmp0;  // right hand side of the equation
                tmp0 = eqbT.transpose();
                rhs = tmp0 - eqbT;
                rhs -= mprod.init(cV0.transpose(), cBwT.transpose())
                               .mul_r(cV0)
                               .mul_r(cS0.asDiagonal())
                               .get();
                auto& x = rhs;
                for (size_t j = 0; j < n; ++j) {
                    for (size_t i = 0; i < n; ++i) {
                        x(i, j) = clip_div(rhs(i, j), cS0(i) + cS0(j));
                    }
                }
                // update eqbT for future solves
                if (mBu) {
                    eqbT -= mprod.init(x.transpose(), cS0.asDiagonal()).get();
                }
                // solve cWkT
                mprod.init(cV0, x.transpose()).mul_r_to(cWkT, cU0T);
            }
            if (!mBu) {
                continue;
            }
            eqbT.noalias() += cBuT * cS0.asDiagonal();
            cSk = eqbT.diagonal();
            {
                // solve Uk
                auto& cUkTU0 = eqbT;  // Uk.T * U0
                for (size_t j = 0; j < n; ++j) {
                    for (size_t i = 0; i < j; ++i) {
                        fp_t v = clip_div(eqbT(i, j), cS0(i) - cS0(j));
                        cUkTU0(i, j) = v;
                        cUkTU0(j, i) = -cBuT(j, i) - v;
                    }
                    cUkTU0(j, j) = -cBuT(j, j) / 2;
                }
                cUkT.noalias() = cUkTU0 * cU0T;
            }
        }
    };
    const size_t batch = mMk.shape(0), n = mMk.shape(1);
    sanm_assert(n >= 2);
    if (mBu) {
        mUk.set_shape({batch, n, n});
        mSk.set_shape({batch, n});
    }
    mWk.set_shape({batch, n, n});
#define CASE(sn)                  \
    do {                          \
        if (n == sn) {            \
            run.operator()<sn>(); \
            return;               \
        }                         \
    } while (0)
    CASE(2);
    CASE(3);
#undef CASE
    run.operator()<Eigen::Dynamic>();
}

void sanm::svd_w_taylor_fwd_p(TensorND& mPk, TensorND& mWk, const TensorND& mMk,
                              const TensorND& mU0, const TensorND& mS0,
                              const TensorND& mW0, const TensorND& mBm,
                              const TensorND& mBp, const TensorND& mBpw) {
    SANM_SCOPED_PROFILER("svd_w_taylor_fwd_p");
    auto run = [&]<Eigen::Index sn>() {
        const size_t batch = mMk.shape(0), n = mMk.shape(1);
        if (sn != Eigen::Dynamic && static_cast<Eigen::Index>(n) != sn) {
            // hint the compiler that n equals to sn
            __builtin_unreachable();
        }
        fp_t *mPkptr = mPk.woptr(), *mWkptr = mWk.woptr();
#define DEF(x) auto x##ptr = const_cast<fp_t*>(x.ptr())
        DEF(mMk);
        DEF(mU0);
        DEF(mS0);
        DEF(mW0);
        DEF(mBm);
        DEF(mBp);
        DEF(mBpw);
#undef DEF

        // matrix variables:
        // prefix c means current (this batch)
        // suffix t/T means transposed (eigen uses col-major)

        // allocate temporaries
        ScopedAllowMalloc scoped_allow_malloc;
        Eigen::Matrix<fp_t, sn, sn> cV0(n, n), eqbT(n, n);
        Eigen::Matrix<fp_t, sn, 1> cS0inv(n, 1);
        MatrixMultiProduct<Eigen::Matrix<fp_t, sn, sn>> mprod(n, n);
        scoped_allow_malloc.disallow();

        for (size_t ib = 0; ib < batch; ++ib) {
            size_t ib_off = ib * n * n;
            EigenMat<sn, 1> cS0(mS0ptr + ib * n, n, 1);
#define DEF(x) EigenMat<sn, sn> c##x##T(m##x##ptr + ib_off, n, n)
            DEF(Pk);
            DEF(Wk);
            DEF(Mk);
            DEF(U0);
            DEF(W0);
            DEF(Bm);
            DEF(Bp);
            DEF(Bpw);
#undef DEF
            cV0.noalias() = cW0T * cU0T.transpose();
            eqbT = cBmT - cBpT;
            mprod.init(cU0T, eqbT).mul_r_to(eqbT, cU0T.transpose());
            mprod.init(cS0.asDiagonal(), cV0.transpose())
                    .mul_r(cMkT)
                    .mul_r(cU0T.transpose());
            eqbT += mprod.get();
            eqbT += mprod.get().transpose();
            auto& x = eqbT;
            for (size_t j = 0; j < n; ++j) {
                for (size_t i = 0; i < n; ++i) {
                    x(i, j) = clip_div(eqbT(i, j), cS0(i) + cS0(j));
                }
            }
            for (size_t i = 0; i < n; ++i) {
                cS0inv(i) = clip_div(1, cS0(i));
            }
            mprod.init(cU0T.transpose(), x).mul_r_to(cPkT, cU0T);
            eqbT = cMkT - cBpwT;
            eqbT.noalias() -= cW0T * cPkT;
            mprod.init(eqbT, cU0T.transpose())
                    .mul_r(cS0inv.asDiagonal())
                    .mul_r_to(cWkT, cU0T);
        }
    };
    const size_t batch = mMk.shape(0), n = mMk.shape(1);
    sanm_assert(n >= 2);
    mPk.set_shape({batch, n, n});
    mWk.set_shape({batch, n, n});
#define CASE(sn)                  \
    do {                          \
        if (n == sn) {            \
            run.operator()<sn>(); \
            return;               \
        }                         \
    } while (0)
    CASE(2);
    CASE(3);
#undef CASE
    run.operator()<Eigen::Dynamic>();
}