#include <iostream>
#include <map>
#include <fstream>

#include <unsupported/Eigen/KroneckerProduct>
#include <Eigen/Geometry>

#include <polyvec/curve-tracer/bezier_merging.hpp>
#include <polyvec/utils/num.hpp>
#include <polyvec/geometry/angle.hpp>
#include <polyvec/geometry/line.hpp>
#include <polyvec/utils/string.hpp>
#include <polyvec/geometry/winding_number.hpp>
#include <polyvec/curve-tracer/curve_bezier.hpp>
#include <polyvec/shortest-path/dijkstra.hpp>
#include <polyvec/io/vtk_curve_writer.hpp>

NAMESPACE_BEGIN ( polyfit )
NAMESPACE_BEGIN ( BezierMerging )

// ===========================================
// Control points parameterization
// ===========================================

void
PointParameters::is_potracable ( bool& result, Eigen::Vector2d& oo, double& alpha, double& beta ) const {

#define undoable() do{ result = false; return; }while(0)

    // const double norm_cap = 10000;

    // if ( control_points.norm() > norm_cap ) {
    //     undoable();
    // }

    bool is_inflection_okay =  !AngleUtils::have_opposite_convexity_with_tol (
                                   control_points.col ( 0 ),
                                   control_points.col ( 1 ),
                                   control_points.col ( 2 ),
                                   control_points.col ( 3 )
                               );

    if ( !is_inflection_okay ) {
        undoable();
    }

    double lhs_at, rhs_at;
    bool do_intersect = ::polyvec::LineUtils::intersect (
                            control_points.col ( 0 ),  control_points.col ( 1 ),
                            control_points.col ( 3 ),  control_points.col ( 2 ),
                            lhs_at,
                            rhs_at );

    bool is_intersection_okay = do_intersect && ( lhs_at >= 1. )  && ( rhs_at >= 1. );

    if ( !is_intersection_okay ) {
        undoable();
    }

    oo = ::polyvec::LineUtils::line_at ( control_points.col ( 0 ),  control_points.col ( 1 ), lhs_at );
    alpha = 1. / lhs_at;
    beta = 1. / rhs_at;
    result = true;

#undef undoable
}

bool
PointParameters::is_potracable () const {
    Eigen::Vector2d oy;
    double alpha, beta;
    bool potracable;
    is_potracable ( potracable, oy, alpha, beta );
    return potracable;
}

PotraceParameters
PointParameters::as_potrace_params() const {
    bool potracable;
    Eigen::Vector2d oy;
    double alpha, beta;
    PotraceParameters ans;

    is_potracable ( potracable, oy, alpha, beta );

    if ( potracable ) {
        ans.alpha = alpha;
        ans.beta = beta;
        Mat23 p0, p1;
        //
        p0.col ( 0 ) << -1, 0;
        p0.col ( 1 ) << 0, 1;
        p0.col ( 2 ) << 1, 0;
        //
        p1.col ( 0 ) << control_points.col ( 0 );
        p1.col ( 1 ) << oy;
        p1.col ( 2 ) << control_points.col ( 3 );
        //
        find_affine_transformation ( p0, p1, ans.A, ans.b );
    } else {
        assert_break ( 0 );
    }

    return ans;
}

// ===========================================
// Potrace parameterization
// ===========================================

double
PotraceParameters::get_areaz() const {
    return 3./10* ( 2*alpha + 2*beta - alpha*beta );
}

double
PotraceParameters::get_areay() const {
    return std::abs ( polyvec::Num::determinant ( A ) ) *get_areaz();
}

PointParameters
PotraceParameters::get_zz() const {
    Mat24 ans;
    ans.col ( 0 ) << -1, 0;
    ans.col ( 1 ) << -1+alpha, alpha;
    ans.col ( 2 ) << 1-beta, beta;
    ans.col ( 3 ) << 1, 0;

    return {ans};
}

Eigen::Vector2d
PotraceParameters::get_oz() {
    return Eigen::Vector2d ( 0, 1 );
}

PointParameters
PotraceParameters::get_yy() const {
    return { ( A*get_zz().control_points ).colwise() + b};
}

Eigen::Vector2d
PotraceParameters::get_oy() const {
    return A*get_oz() + b;
}

PotraceParameters
PotraceParameters::equiparamed() const {
    PotraceParameters ans;
    ans = *this;
    ans.alpha = ans.beta = 2 - std::sqrt ( ( 2-alpha ) * ( 2-beta ) );
    return ans;
}

bool
PotraceParameters::is_ccw() const {
    return polyvec::Num::determinant ( A ) < 0;
}

// stolen from potrace <3
// calculate (p1-p0)x(p3-p2)
namespace {
	double
		_cprod(const Eigen::Vector2d& p0, const Eigen::Vector2d& p1, const Eigen::Vector2d& p2, const Eigen::Vector2d& p3) {
		double x1, y1, x2, y2;

		x1 = p1.x() - p0.x();
		y1 = p1.y() - p0.y();
		x2 = p3.x() - p2.x();
		y2 = p3.y() - p2.y();

		return x1 * y2 - x2 * y1;
	}

	// calculate the point t in [0..1] on the (convex) bezier curve
	//   (p0,p1,p2,p3) which is tangent to q1-q0. Return -1.0 if there is no
	//   solution in [0..1].
	double
		_tangent(const Eigen::Vector2d& p0, const Eigen::Vector2d& p1, const Eigen::Vector2d& p2, const Eigen::Vector2d& p3, const Eigen::Vector2d& q0,
			const Eigen::Vector2d& q1) {
		double A, B, C;   /* (1-t)^2 A + 2(1-t)t B + t^2 C = 0 */
		double a, b, c;   /* a t^2 + b t + c = 0 */
		double d, s, r1, r2;

		A = _cprod(p0, p1, q0, q1);
		B = _cprod(p1, p2, q0, q1);
		C = _cprod(p2, p3, q0, q1);

		a = A - 2 * B + C;
		b = -2 * A + 2 * B;
		c = A;

		d = b * b - 4 * a * c;

		if ((std::abs(a) < 1e-12)
			&& (std::abs(b) > 1e-12)
			&& (-c / b >= 0)
			&& (-c / b <= 1.)) {
			return -c / b;
		}

		if ((std::abs(a) < 1e-12) || (d < 1e-12)) {
			return -1.0;
		}

		s = sqrt(d);

		r1 = (-b + s) / (2 * a);
		r2 = (-b - s) / (2 * a);

		if (r1 >= 0 && r1 <= 1) {
			return r1;
		}
		else if (r2 >= 0 && r2 <= 1) {
			return r2;
		}
		else {
			return -1.0;
		}
	}
}

void
PotraceParameters::t_for_tangent ( const Eigen::Vector2d& tangent_y, double& t, bool& success ) const {
    const Mat24 yy = get_yy().control_points;
    t = _tangent ( yy.col ( 0 ), yy.col ( 1 ), yy.col ( 2 ), yy.col ( 3 ), Eigen::Vector2d::Zero(), tangent_y.normalized() );

    if ( t < 0 ) {
        success = false;
    } else {
        success = true;
    }

# if 0 // my own attempt. Potrace seems nicer.
    const double pi = ::polyvec::constants::PI;
    // Find the tangent in the reference coordiantes
    Eigen::Vector2d tangent_z = Num::solve_linear_system ( A, tangent_y );
    const Eigen::Vector2d zz = get_zz().control_points;
    assert_break ( tangent_z.norm() > 1e-10 );
    // Now find the angle of the tangent with the x axis
    // also find the begin and end angle of the bezier
    const double angle = std::atan2 ( tangent_z.y(), tangent_z.x() );
    const double begin_angle = std::atan2 ( zz.col ( 1 ).y() - zz.col ( 0 ).y(), zz.col ( 1 ).x() - zz.col ( 0 ).x() );
    const double end_angle = std::atan2 ( zz.col ( 3 ).y() - zz.col ( 2 ).y(), zz.col ( 3 ).x() - zz.col ( 2 ).x() );
    assert_break ( std::abs ( begin_angle-pi/4 ) < 1e-5 );
    assert_break ( std::abs ( end_angle+pi/4 ) < 1e-5 );
#endif
}

// ===========================================
// Functions
// ===========================================

void
find_affine_transformation ( const Mat23& points0, const Mat23& points1, Eigen::Matrix2d& A, Eigen::Vector2d& b ) {
    Eigen::Matrix<double, 6, 6> LHS;
    Eigen::Matrix<double, 6, 1> RHS;
    Eigen::Matrix<double, 6, 1> unknowns;

    Eigen::Matrix2d eye =  Eigen::Matrix2d::Identity();

    LHS.block<2, 4> ( 0, 0 ) = Eigen::kroneckerProduct ( points0.col ( 0 ).transpose(), eye );
    LHS.block<2, 4> ( 2, 0 ) = Eigen::kroneckerProduct ( points0.col ( 1 ).transpose(), eye );
    LHS.block<2, 4> ( 4, 0 ) = Eigen::kroneckerProduct ( points0.col ( 2 ).transpose(), eye );
    //
    LHS.block<2, 2> ( 0, 4 ) = eye;
    LHS.block<2, 2> ( 2, 4 ) = eye;
    LHS.block<2, 2> ( 4, 4 ) = eye;
    //
    RHS.segment<2> ( 0 ) = points1.col ( 0 );
    RHS.segment<2> ( 2 ) = points1.col ( 1 );
    RHS.segment<2> ( 4 ) = points1.col ( 2 );

    // Debugging
    //std::cout << LHS << std::endl;
    //std::cout << Num::determinant(LHS) << std::endl;

    unknowns = polyvec::Num::solve_linear_system ( LHS, RHS );

    A.col ( 0 ) = unknowns.segment<2> ( 0 );
    A.col ( 1 ) = unknowns.segment<2> ( 2 );
    b = unknowns.segment<2> ( 4 );

    // Make sure it is working
    const double error = ( ( A*points0 ).colwise() + b - points1 ).norm() ;
    assert_break ( error < 1e-10 );
}

namespace {
bool
skip_comments ( FILE* fl ) {
    char tmp[4096];
    bool is_there_more_to_read = false;
    assert_break ( fl );

    for ( ;; ) {
        const int c = fgetc ( fl );

        if ( c ==EOF ) {
            is_there_more_to_read = false;
            break;
        } else if ( c == '#' ) {
            fgets ( tmp, 4096, fl );
            continue;
        } else if ( ( c != ' ' ) && ( c != '\n' ) && ( c != '\r' ) ) {
            int success = fseek ( fl, -1, SEEK_CUR );
            assert ( success == 0 );
            is_there_more_to_read = true;
            break;
        }
    }

    return is_there_more_to_read;
}
}

void
dump_curve_sequence ( const std::vector<Eigen::Matrix2Xd>& curves, FILE* file ) {
    fprintf ( file, "# N curves \n %d \n", ( int ) curves.size() );

    for ( int i = 0 ; i < curves.size() ; ++ i ) {
        fprintf ( file, "# Curve %d, n points \n %d \n", i, ( int ) curves[i].cols() );

        for ( int j = 0 ; j < curves[i].cols() ; ++ j ) {
            fprintf ( file, "%.12g ", curves[i] ( 0, j ) );
        }

        fprintf ( file, "\n" );

        for ( int j = 0 ; j < curves[i].cols() ; ++ j ) {
            fprintf ( file, "%.12g ", curves[i] ( 1, j ) );
        }

        fprintf ( file, "\n" );
    }
}

std::vector<Eigen::Matrix2Xd>
read_curve_sequence ( FILE* file ) {
#define ENSURE(X) if(!(X)) {assert_break(0);}

    int n_curves;
    ENSURE ( skip_comments ( file ) );
    ENSURE ( fscanf ( file, "%d", &n_curves ) == 1 );

    std::vector<Eigen::Matrix2Xd> curves ( n_curves );

    for ( int i = 0 ; i < curves.size() ; ++ i ) {
        int n_points;
        ENSURE ( skip_comments ( file ) );
        ENSURE ( fscanf ( file, "%d", &n_points ) == 1 );
        curves[i].resize ( 2, n_points );

        ENSURE ( skip_comments ( file ) );

        for ( int j = 0 ; j < curves[i].cols() ; ++ j ) {
            ENSURE ( fscanf ( file, "%lf", &curves[i] ( 0, j ) ) == 1 );
        }

        ENSURE ( skip_comments ( file ) );

        for ( int j = 0 ; j < curves[i].cols() ; ++ j ) {
            ENSURE ( fscanf ( file, "%lf", &curves[i] ( 1, j ) ) == 1 );
        }
    }

    return curves;

#undef ENSURE
}

void merged_curve (
    const std::vector<Eigen::Matrix2d>& line_begin,
    const std::vector<Mat24>& beziers,
    const std::vector<Eigen::Matrix2d>& line_end,
    FILE* verbose,
    bool& doable,
    Mat24& ans_mat ) {

    PotraceParameters ans;

    auto eprintf = [verbose] ( const std::string strr ) {
        if ( verbose ) {
            fprintf ( verbose, "%s", strr.c_str() );
        }
    };

#define undoable() do{doable =false; eprintf(""); return;}while(0)

    assert_break ( line_begin.size() <= 1 );
    assert_break ( line_end.size() <= 1 );
    assert_break ( beziers.size() >= 1 );

    // ============
    // First check that each bezier should be potracable
    // ============
    for ( int i = 0 ; i < ( int ) beziers.size() ; ++i ) {
        if ( !PointParameters ( {beziers[i]} ).is_potracable() ) {
            eprintf ( StringUtils::fmt ( "curve[%d] is not potracable \n", i ) );
            undoable();
        }
    }

    // ============
    // Now check that transitions from each curve also does not
    // introduce an inflection
    // The assumption is that the tangents of all the transitions match
    // ============

    // line begin
    if ( line_begin.size() ) {
        if ( AngleUtils::have_opposite_convexity (
                    line_begin[0].col ( 0 ),
                    beziers[0].col ( 1 ),
                    beziers[0].col ( 2 ),
                    beziers[0].col ( 3 ) ) ) {
            eprintf ( StringUtils::fmt ( "line0 to bezier 1 inflected \n" ) );
            undoable();
        }
    }

    // line end
    if ( line_end.size() ) {
        if ( AngleUtils::have_opposite_convexity (
                    beziers.back().col ( 0 ),
                    beziers.back().col ( 1 ),
                    beziers.back().col ( 2 ),
                    line_end[0].col ( 1 ) ) ) {
            eprintf ( StringUtils::fmt ( "line_end to bezier end inflected \n" ) );
            undoable();
        }
    }

    // Beziers
    for ( int i = 0 ; i < ( int ) beziers.size()-1 ; ++i ) {
        if ( AngleUtils::have_opposite_convexity (
                    beziers[i].col ( 1 ),
                    beziers[i].col ( 2 ),
                    beziers[i+1].col ( 1 ),
                    beziers[i+1].col ( 2 ) ) ) {
            eprintf ( StringUtils::fmt ( "bezier[%d] to bezier[%d] inflected \n", i, i +1 ) );
            undoable();
        }
    }

    // ============
    //   Check that the first and last line intersection each other
    // ============
    Eigen::Vector2d y0, y3, oy;
    {
        Eigen::Matrix2d ray_begin;
        Eigen::Matrix2d ray_end;

        if ( line_begin.size() ) {
            ray_begin = line_begin.front();
        } else {
            ray_begin = beziers.front().leftCols<2>();
        }

        if ( line_end.size() ) {
            ray_end = line_end.front();
        } else {
            ray_end = beziers.back().rightCols<2>();
        }

        ray_end.rowwise().reverseInPlace();

        double lhs_at, rhs_at;
        const bool do_intersect = ::polyvec::LineUtils::intersect ( ray_begin.col ( 0 ), ray_begin.col ( 1 ), ray_end.col ( 0 ), ray_end.col ( 1 ), lhs_at, rhs_at );

        if ( ( !do_intersect )
                || ( lhs_at < 1. )
                || ( rhs_at < 1. ) ) {
            eprintf ( StringUtils::fmt ( "Intersection Problem \n" ) );
            undoable();
        }

        y0 = ray_begin.col ( 0 );
        y3 = ray_end.col ( 0 );
        oy = ::polyvec::LineUtils::line_at ( ray_begin.col ( 0 ), ray_begin.col ( 1 ), lhs_at );

        // This is intersecting on the other side of the universe
        if ( oy.norm() > 10000 ) {
            //polyvec::VtkCurveWriter writer;
            //writer.add_polyline( ray_begin );
            //writer.add_polyline( ray_end   );
            //writer.dump("DAMN.vtk");
            undoable();
        }
    }

    // ============
    //  Now find the area of the bezier sequence
    // ============
    double area_shape = 0;
    {
        auto find_carpet  = [&line_begin, &line_end, &beziers] {

            Eigen::Matrix2Xd carpet ( 2, beziers.size() + line_begin.size() + line_end.size() + 1 );
            int colid = 0;

            if ( line_begin.size() ) {
                carpet.col ( colid ) = line_begin.front().col ( 0 );
                ++colid;
            }

            for ( int i = 0 ; i < ( int ) beziers.size() ; ++i ) {
                carpet.col ( colid ) = beziers[i].col ( 0 );
                ++colid;
            }

            if ( line_end.size() ) {
                carpet.col ( colid ) = line_end.front().col ( 0 );
                ++colid;
                carpet.col ( colid ) = line_end.front().col ( 1 );
                ++colid;
            } else {
                carpet.col ( colid ) = beziers.back().col ( 3 );
                ++colid;
            }

            assert_break ( colid == ( int ) carpet.cols() );

            return carpet;
        }; // find carpet

        bool is_carpet_ccw;
        double area_carpet;
        Eigen::Matrix2Xd carpet = find_carpet ();
        ::polyvec::WindingNumber::compute_orientation ( carpet, is_carpet_ccw, area_carpet );

        double area_beziers = 0;

        for ( int i = 0 ; i < ( int ) beziers.size() ; ++i ) {
            area_beziers += PointParameters ( {beziers[i]} ).as_potrace_params().get_areay();
        }

        area_shape = area_carpet + area_beziers;
    }

    //
    // If we are here, we can finally create the merged bezier
    //
    {
        Mat23 p0, p1;
        //
        p0.col ( 0 ) << -1, 0;
        p0.col ( 1 ) << 0, 1;
        p0.col ( 2 ) << 1, 0;
        //
        p1.col ( 0 ) << y0;
        p1.col ( 1 ) << oy;
        p1.col ( 2 ) << y3;
        //
        find_affine_transformation ( p0, p1, ans.A, ans.b );

        const double area_ref  =  area_shape / std::abs ( polyvec::Num::determinant ( ans.A ) );

        assert_break( !std::isnan( area_ref ) );
        if ( ( area_ref > 1.19 ) ) {
            eprintf ( StringUtils::fmt ( "Area Problem \n" ) );
            undoable();
        }

        ans.alpha =   ans.beta  = 2. - sqrt ( 4-10*area_ref/3. ) ;
        assert_break (  ans.beta > 0 );

        if ( ans.beta > 1 ) {
            // DEBUGGING - TEMP, REMOVE
            //for (auto bz : beziers ) {
            //    std::cout << bz.row(0) << " " << bz.row(1) << std::endl;
            // }
            undoable();
        }

        assert_break ( area_ref > 0 );
        assert_break ( std::abs ( area_ref - ans.get_areaz() ) < 1e-4 );
        ans_mat = ans.get_yy().control_points;
        doable = true;
    }

#undef undoable
}

void
is_merge_acceptable (
    const std::vector<Mat24>& beziers,
    const Mat24& merged_bezier,
    const double tolerance,
    const double alphamax,
    double& penalty,
    bool& is_acceptable ) {

    //
    // HELPERS
    //

    auto get_distance = [] ( const PotraceParameters &bzp, ::polyvec::BezierCurve &bz, const Eigen::Vector2d& p0, const Eigen::Vector2d& p1,
    bool &proj_success, double &dist ) {
        double t;
        bzp.t_for_tangent ( p1-p0, t, proj_success );

        if ( !proj_success ) {
            return;
        }

        dist = ::polyvec::LineUtils::distance_from_point ( p0, p1, bz.pos ( t ) );
    };

    auto get_distance_signed = [] ( const PotraceParameters &bzp,  ::polyvec::BezierCurve &bz, const Eigen::Vector2d& p0, const Eigen::Vector2d& p1,
    const Eigen::Vector2d pforsign,  bool &proj_success, double &dist ) {
        double t;
        bzp.t_for_tangent ( p1-p0, t, proj_success );

        if ( !proj_success ) {
            return;
        }

        const double dist_pref =  ::polyvec::LineUtils::signed_distance_from_point ( p0, p1, pforsign );
        dist =  ::polyvec::LineUtils::signed_distance_from_point ( p0, p1, bz.pos ( t ) );

        if ( dist*dist_pref < 0 ) {
            dist = -std::abs ( dist );
        } else {
            dist = std::abs ( dist );
        }
    };

    auto check_distance_a = [tolerance] ( const double dist_a ) {
        assert_break ( dist_a >= 0 );

        if ( dist_a > tolerance ) {
            return false;
        }

        return true;
    };

    auto check_distance_b = [tolerance] ( const double dist_b ) {
        if ( dist_b < -tolerance ) {
            return false;
        }

        return true;
    };

    //
    // Loop over all beziers. Calculate the dist_a and dist_b distances
    // from the potrace paper.
    //
    penalty = 0;
    is_acceptable = true;
    ::polyvec::BezierCurve eval_ctx;
    eval_ctx.set_control_points ( merged_bezier );
    PotraceParameters project_ctx;
    project_ctx = PointParameters ( {merged_bezier} ).as_potrace_params();

    //
    // Check alpha
    //
    if ( project_ctx.alpha > alphamax ) {
        is_acceptable = false;
        return;
    }

    //
    // Check distace
    //
    for ( int i = 0 ; i < ( int ) beziers.size() ; ++i ) {

        // dist a
        const Eigen::Vector2d a = PointParameters ( {beziers[i]} ).as_potrace_params().get_oy();

        if ( i < ( int ) beziers.size() -1 ) {
            const Eigen::Vector2d anext = PointParameters ( {beziers[i+1]} ).as_potrace_params().get_oy();
            bool can_project;
            double dist_a;
            get_distance ( project_ctx, eval_ctx, a, anext, can_project, dist_a );

            if ( ( !can_project ) || ( !check_distance_a ( dist_a ) ) ) {
                is_acceptable = false;
                return;
            }

            penalty += dist_a * dist_a;
        } // end of dist a

        // dist b
        {
            const Eigen::Vector2d b = beziers[i].col ( 0 );
            const Eigen::Vector2d bnext =  beziers[i].col ( 3 );
            bool can_project;
            double dist_b;
            get_distance_signed ( project_ctx, eval_ctx, b, bnext, a, can_project, dist_b );

            if ( ( !can_project ) || ( !check_distance_b ( dist_b ) ) ) {
                is_acceptable = false;
                return;
            }

            penalty += dist_b * dist_b;
        } // end of dist b
    } // end of beziers

} //  merge_penalty()

void merge_penalty (
    const std::vector<Mat24>& beziers,
    const Mat24& merged_bezier,
    double& penalty ) {

    // Do a one sided projection of the points
    constexpr int n_samples = 10;

    penalty = 0;
    ::polyvec::BezierCurve big_eval_ctx;
    big_eval_ctx.set_control_points ( merged_bezier );

    for ( int i = 0 ; i < ( int ) beziers.size() ; ++i ) {
        ::polyvec::BezierCurve sub_eval_ctx;
        sub_eval_ctx.set_control_points ( beziers[i] );
        // This is a question . Should we multiply by length?
        const double weight = 1. / n_samples *  sub_eval_ctx.length();
        // Perhaps not
        // const double weight = 1. / n_samples;

        for ( int j = 1 ; j < ( n_samples+1 ) ; ++j ) {
            const double t_sub = 1. / ( n_samples+1 ) * j ;
            const Eigen::Vector2d point_sub = sub_eval_ctx.pos ( t_sub );
            const double t_big = big_eval_ctx.project ( point_sub );
            const Eigen::Vector2d point_big = big_eval_ctx.pos ( t_big );
            penalty += weight * ( point_big-point_sub ).norm();
        } // end of points
    }  // end of bezeirs
}

void
merge_curves (
    const std::vector<Eigen::Matrix2Xd>& curves_in, // lines of bezier control points
    const bool is_circular,
    const MergeRecursivelyOptions options,
	const std::set<Regularity::SymmetryPair>& symmetric_curves,
    std::vector<Eigen::Matrix2Xd>& curves_out,
    std::vector<std::vector<int>>& out2in  ) {

    const int n_curves = ( int ) curves_in.size();

    //
    // Check if a curve is right after a corner
    //
    auto is_line_after_corner = [&n_curves, &curves_in, &is_circular] ( const int i ) {
        if ( ( i == 0 ) && ( !is_circular ) ) {
            return false;
        } else {
            const int iprev = ( ( i-1 )+n_curves ) %n_curves;

            if ( ( curves_in[iprev].cols() == 2 )
                    && ( curves_in[i].cols() == 2 ) ) {
                return true;
            } else {
                return false;
            }
        }
    };

    //
    // Prepare the data for a merge curve
    //
    auto get_merge_data = [&n_curves, &curves_in, &is_circular]
                          ( const int i,
                            const int n_merge,
                            std::vector<Mat22>& line_begin,
                            std::vector<Mat24>& beziers,
                            std::vector<Mat22>& line_end,
    std::vector<int>& parent_ids ) {

        assert_break ( is_circular || ( n_merge+i<=n_curves ) );
        line_begin.resize ( 0 );
        line_end.resize ( 0 );
        beziers.resize ( 0 );
        parent_ids.resize ( 0 );

        for ( int offset = 0 ; offset < n_merge ; ++offset ) {
            const int curve_id = ( i + offset ) %n_curves;

            if ( ( curves_in[curve_id].cols() == 2 ) && ( offset==0 ) ) {
                line_begin.push_back ( curves_in[curve_id] );
            } else if ( ( curves_in[curve_id].cols() == 2 ) && ( offset==n_merge-1 ) ) {
                line_end.push_back ( curves_in[curve_id] );
            }
            else {
                assert_break ( curves_in[curve_id].cols() == 4 );
                beziers.push_back ( curves_in[curve_id] );
            }

            parent_ids.push_back ( curve_id );
        }

    };

    //
    // Convert a curves end or begin point into node ids
    //
    auto n_graph_nodes = [&is_circular, &curves_in, &n_curves]() {
        if ( is_circular  ) {
            return n_curves;
        } else {
            return n_curves+1;
        }
    };
    auto curve_end_node = [&is_circular, &curves_in, &n_curves] ( const int i ) {
        if ( is_circular && ( i == n_curves-1 ) ) {
            return 0;
        } else {
            return i+1;
        }
    };
    auto curve_begin_node = [&is_circular, &curves_in, &n_curves] ( const int i ) {
        return i;
    };

    //
    // All the possible merge situations
    //
    std::vector<ShortestPath::Node> graph_nodes ( n_graph_nodes() );
    std::map<std::pair<int, int>, int> mc_map;
    std::vector<double> mc_penalty;
    std::vector<std::vector<int>> mc_parent_ids;
    std::vector<Eigen::Matrix2Xd> mc_pts;

    //
    // Create all the merge candidates
    //
    {
        std::vector<Mat22> line_begin;
        std::vector<Mat24> beziers;
        std::vector<Mat22> line_end;
        ShortestPath::Node node;
        std::vector<int> parent_ids;
        Mat24 merged_bezier;
        bool is_merge_doable;
        double potrace_penalty;
        double distace_penalty;

        for ( int cbeginid = 0 ; cbeginid < n_curves ; ++cbeginid ) {
            const int max_n_merge = is_circular ? n_curves : n_curves - cbeginid;

			// find the next symmetry pair after the current curve
			auto next_symmetry = std::lower_bound(symmetric_curves.begin(), symmetric_curves.end(), Regularity::SymmetryPair(cbeginid, cbeginid, -1));
			if (next_symmetry == symmetric_curves.end())
				next_symmetry = symmetric_curves.begin(); //wrap around

			// We employ the following strategy to promote symmetries:
			// A merge is only valid if all its symmetries are completely within the
			// merge or none is (i.e., it is ok if only one curve of symmetry pairs
			// is within the merge). This is evaluated independently for each
			// participating symmetric region

			struct SymmetryInfo
			{
				// stores the number of symmetric pairs that are completely
				// contained
				int n_complete_symmetries = 0;

				// stores the number of symmetric paris where only one part is
				// contained
				int n_partial_symmetries = 0;
			};

			std::map<int, SymmetryInfo> symmetry_info;

            for ( int n_merge = 1 ; n_merge <= max_n_merge ; ++n_merge ) {

				const int cendid = (cbeginid + n_merge - 1) % n_curves;

				// check if we have new symmetries
				while (next_symmetry != symmetric_curves.end() && next_symmetry->first == cendid)
				{
					auto& info = symmetry_info[next_symmetry->region];
					// check if the other curve is already in the merge
					if (PathUtils::contains_closed(curves_in.size(), cbeginid, cendid, next_symmetry->second, is_circular))
					{
						++info.n_complete_symmetries;
						--info.n_partial_symmetries;
					}
					else
						++info.n_partial_symmetries;

					++next_symmetry;
					if (next_symmetry == symmetric_curves.end())
						next_symmetry = symmetric_curves.begin(); //wrap around
				}

                // If n_merge is just 1 then we don't have to do much
                // The fit is accepted and the error is 0
                if ( n_merge==1 ) {
                    mc_pts.push_back ( curves_in[cbeginid] );
                    mc_penalty.push_back ( 0 );
                    mc_parent_ids.push_back ( {cbeginid} );
                } else {

                    // Check that we do not pass a corner
                    const int cendidm1 = ( cbeginid+n_merge-2 ) %n_curves;

                    // Handle lines
                    if ( options.allow_merging_lines ) {
                        // If were were fitting to potrace data this should be enough
                        if (  is_line_after_corner ( cendid ) ) {
                            break;
                        }
                        // For ed's assymetric case, we also need this
                        if ( ( curves_in[cendidm1].cols() == 2 ) && ( cendidm1 != cbeginid) ) {
                            break;
                        }
                    }

                    // New strategy, should work for Ed too
                    // Simply don't let cbeginid or cendid to become line
                    // basically we are not attempting to merge lines with anything
                    if ( !options.allow_merging_lines ) {
                        if ( ( curves_in[cbeginid].cols() == 2 ) || ( curves_in[cendid].cols() == 2 ) ) {
                            break;
                        }
                    }

					bool symmetry_valid = true;
					for(auto& info : symmetry_info)
						if (info.second.n_complete_symmetries > 0 && info.second.n_partial_symmetries > 0)
							symmetry_valid = false; // violates symmetry
					if (!symmetry_valid)
						continue;

                    // Get info needed for merging
                    get_merge_data ( cbeginid, n_merge, line_begin, beziers, line_end, parent_ids );

                    // Merge the curve
                    merged_curve ( line_begin, beziers, line_end, nullptr, is_merge_doable, merged_bezier );

                    if ( !is_merge_doable ) {
                        continue;
                    }

                    // Check if acceptable
                    is_merge_acceptable ( beziers, merged_bezier, options.distance_tolerance, options.alpha_max, potrace_penalty, is_merge_doable );

                    if ( !is_merge_doable ) {
                        continue;
                    }

                    // Find the penalty
                    merge_penalty ( beziers, merged_bezier, distace_penalty );

                    // Add to the curve to nodes
                    mc_pts.push_back ( merged_bezier );
                    mc_penalty.push_back ( distace_penalty );
                    mc_parent_ids.push_back ( parent_ids );

                } // end of more than one curve

            } // end of n_merge
        } // end of begin id
    } // end of creating the merge candidates

    //
    // Now create a graph
    //
    for ( int node_id = 0 ; node_id < n_graph_nodes() ; ++node_id  ) {
        graph_nodes[node_id].v = node_id;
    }

    for ( int mcid = 0 ; mcid < ( int ) mc_parent_ids.size() ; ++mcid ) {
        const int node_from = curve_begin_node ( mc_parent_ids[mcid].front() );
        const int node_to = curve_end_node ( mc_parent_ids[mcid].back() );
        graph_nodes[node_from].add_neighbor(node_to, mc_penalty[mcid] + options.per_bezier_const );
        mc_map.insert ( std::make_pair ( std::make_pair ( node_from, node_to ), mcid ) );
    }

    // Find the shortest path
    {
        std::vector<Eigen::Index> P;
        ShortestPath::State state;

        if ( is_circular ) {
            ShortestPath::find_cycle ( state, graph_nodes, {0}, P );
            // Debugging
            // ShortestPath::print_graph ( graph_nodes, stderr );
        } else {
            ShortestPath::find ( state, graph_nodes, 0, n_graph_nodes()-1, P );
        }

        const int max_node_offset = is_circular ? ( int ) P.size() : ( int ) P.size()-1;

        curves_out.resize ( 0 );
        out2in.resize ( 0 );

        for ( int nodeoffset = 0 ; nodeoffset <  max_node_offset ; ++nodeoffset ) {
            const int node_from = (int)P[nodeoffset];
            const int node_to = (int)P[ ( nodeoffset+1 ) %P.size()];
            auto mc_id_it = mc_map.find ( std::make_pair ( node_from, node_to ) );
            assert_break ( mc_id_it != mc_map.end() );
            const int mc_id = (int)mc_id_it->second;

            curves_out.push_back ( mc_pts[ mc_id ] );
            out2in.push_back ( mc_parent_ids[ mc_id ] );
        }

    }

} // merge_recursively()

NAMESPACE_END ( BezierMerging )
NAMESPACE_END ( polyfit )


//
// Implements Tree-iLQR.
//
// Arun Venkatraman (arunvenk@cs.cmu.edu)
// December 2016
//

#pragma once

#include <ilqr/tree.hh>
#include <ilqr/ilqr_node.hh>

#include <Eigen/Dense>

#include <memory>
#include <list>
#include <vector>

namespace ilqr
{

// State-Action-Cost-Value(CostToGo).
struct SACV
{
    Eigen::VectorXd x;
    Eigen::VectorXd u;
    double c;
    double probability;
    double value;
};

// Shared pointer to a Node in the underlying Tree structure of the iLQR-Tree. Calling ->item()
// returns a shared pointer to the iLQRNode that holds the state/control/dynamics/etc.
// information.
using TreeNodePtr = std::shared_ptr<data::Node<iLQRNode>>;

class iLQRTree
{
public:
    iLQRTree(int state_dim, int control_dim);
    virtual ~iLQRTree() = default;

    // Construct a iLQRNode which represents the dynamics, cost functions as
    // well as the state and control policy at a specific time step along the
    // tree. Nominal state and control inputs are passed and used for the
    // initial Taylor expansions of the cost and dynamics functions.
    std::shared_ptr<iLQRNode> make_ilqr_node(const Eigen::VectorXd &x_star,
            const Eigen::VectorXd &u_star, const DynamicsFunc &dynamics, const
            CostFunc &cost, const double probablity);

    // Add the root node to the iLQR Tree with similar arguments to the
    // make_ilqr_node function.
    TreeNodePtr add_root(const Eigen::VectorXd &x_star, const Eigen::VectorXd
            &u_star, const DynamicsFunc &dynamics, const CostFunc &cost);

    // Add a iLQRNode as the root node to the iLQR Tree. Requires the
    // probability to be 1.
    TreeNodePtr add_root(const std::shared_ptr<iLQRNode> &ilqr_node);

    // Add a list of iLQRNodes as children under a parent. The probabilities
    // must sum to 1.
    std::vector<TreeNodePtr> add_nodes(const
            std::vector<std::shared_ptr<iLQRNode>> &ilqr_nodes, TreeNodePtr
            &parent);

    // Get the root node of the Tree.
    TreeNodePtr root();

    // Forward pass to generate new nominal points for iLQR.
    // The step-size is controlled by alpha (alpha * K).
    void forward_tree_update(const double alpha);

    // Do a full bellman backup on the tree.
    void bellman_tree_backup();

    // Returns a Tree of states-action-costs-value with expected cost-to-gos.
    data::Tree<SACV> forward_pass(const Eigen::VectorXd &x0, const TreeNodePtr &top = nullptr, const double top_alpha = 1.0);

private:
    int state_dim_ = 0;
    int control_dim_ = 0;

    data::Tree<iLQRNode> tree_;

    // Zeros matrix of size [state_dim +1] x [state_dim +1]. Used as
    // initialization for the zeros matrix.
    const QuadraticValue ZERO_VALUE_;

    // Special case for just the leaves of the tree. We can compute this by
    // giving the leaves synthetic children with $V_{T+1} = 0$.
    void control_and_value_for_leaves();

    // Backups the value matrix and control gains matrix from the children of
    // the node to get the value and control policies for the parents of the
    // children. Returns a list of all the parents.
    std::list<TreeNodePtr> backup_to_parents(const std::list<TreeNodePtr>
            &all_children);

};

} // namespace ilqr


/*
 Copyright (C) 2016 Quaternion Risk Management Ltd
 All rights reserved.

 This file is part of ORE, a free-software/open-source library
 for transparent pricing and risk analysis - http://opensourcerisk.org

 ORE is free software: you can redistribute it and/or modify it
 under the terms of the Modified BSD License.  You should have received a
 copy of the license along with this program.
 The license is also available online at <http://opensourcerisk.org>

 This program is distributed on the basis that it will form a useful
 contribution to risk analytics and model standardisation, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 FITNESS FOR A PARTICULAR PURPOSE. See the license for more details.
*/

/*! \file interpolateddiscountcurve.hpp
    \brief interpolated discount term structure
    \ingroup termstructures
*/

#ifndef quantext_interpolated_discount_curve_hpp
#define quantext_interpolated_discount_curve_hpp

#include <boost/make_shared.hpp>
#include <ql/termstructures/yieldtermstructure.hpp>
#include <qle/quotes/logquote.hpp>

namespace QuantExt {
using namespace QuantLib;

//! InterpolatedDiscountCurve based on loglinear interpolation of DiscountFactors
/*! InterpolatedDiscountCurve based on loglinear interpolation of DiscountFactors,
    flat fwd extrapolation is always enabled, the term structure has always a
    floating reference date

        \ingroup termstructures
    */
class InterpolatedDiscountCurve : public YieldTermStructure {
public:
    enum class Interpolation { logLinear, linearZero };
    enum class Extrapolation { flatFwd, flatZero };
    //! \name Constructors
    //@{
    //! default constructor
    InterpolatedDiscountCurve(const std::vector<Time>& times, const std::vector<Handle<Quote>>& quotes,
                              const Natural settlementDays, const Calendar& cal, const DayCounter& dc,
                              const Interpolation interpolation = Interpolation::logLinear,
                              const Extrapolation extrapolation = Extrapolation::flatFwd)
        : YieldTermStructure(settlementDays, cal, dc), times_(times), interpolation_(interpolation),
          extrapolation_(extrapolation) {
        initalise(quotes);
    }

    //! constructor that takes a vector of dates
    InterpolatedDiscountCurve(const std::vector<Date>& dates, const std::vector<Handle<Quote>>& quotes,
                              const Natural settlementDays, const Calendar& cal, const DayCounter& dc,
                              const Interpolation interpolation = Interpolation::logLinear,
                              const Extrapolation extrapolation = Extrapolation::flatFwd)
        : YieldTermStructure(settlementDays, cal, dc), times_(dates.size()), interpolation_(interpolation),
          extrapolation_(extrapolation) {
        for (Size i = 0; i < dates.size(); ++i)
            times_[i] = timeFromReference(dates[i]);
        initalise(quotes);
    }
    //@}

private:
    void initalise(const std::vector<Handle<Quote>>& quotes) {
        QL_REQUIRE(times_.size() > 1, "at least two times required");
        QL_REQUIRE(times_[0] == 0.0, "First time must be 0, got " << times_[0]); // or date=asof
        QL_REQUIRE(times_.size() == quotes.size(), "size of time and quote vectors do not match");
        for (Size i = 0; i < quotes.size(); ++i) {
            quotes_.push_back(boost::make_shared<LogQuote>(quotes[i]));
        }
        for (Size i = 0; i < times_.size() - 1; ++i)
            timeDiffs_.push_back(times_[i + 1] - times_[i]);
    }

    //! \name TermStructure interface
    //@{
    Date maxDate() const { return Date::maxDate(); } // flat fwd extrapolation
    //@}

protected:
    DiscountFactor discountImpl(Time t) const {
        if (t > this->times_.back() && extrapolation_ == Extrapolation::flatZero) {
            Real tMax = this->times_.back();
            Real dMax = std::exp(quotes_.back()->value());
            return std::pow(dMax, t / tMax);
        }
        std::vector<Time>::const_iterator it = std::upper_bound(times_.begin(), times_.end(), t);
        Size i = std::min<Size>(it - times_.begin(), times_.size() - 1);
        Real weight = (times_[i] - t) / timeDiffs_[i - 1];
        if (interpolation_ == Interpolation::logLinear || t > this->times_.back()) {
            // this handles flat fwd extrapolation (t > times.back()) as well
            Real value = (1.0 - weight) * quotes_[i]->value() + weight * quotes_[i - 1]->value();
            return ::exp(value);
        } else {
            Real value =
                (1.0 - weight) * quotes_[i]->value() / times_[i] + weight * quotes_[i - 1]->value() / times_[i - 1];
            return ::exp(t * value);
        }
    }

private:
    std::vector<Time> times_;
    std::vector<Time> timeDiffs_;
    std::vector<boost::shared_ptr<Quote>> quotes_;
    Interpolation interpolation_;
    Extrapolation extrapolation_;
};

} // namespace QuantExt

#endif


//=======================================================================
// Copyright (c)
//
// 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)
//=======================================================================
/**
 * @file bounded_degree_mst_oracle.hpp
 * @brief
 * @author Piotr Godlewski
 * @version 1.0
 * @date 2013-06-05
 */
#ifndef PAAL_BOUNDED_DEGREE_MST_ORACLE_HPP
#define PAAL_BOUNDED_DEGREE_MST_ORACLE_HPP

#include "paal/iterative_rounding/min_cut.hpp"
#include "paal/lp/lp_base.hpp"

#include <boost/optional.hpp>
#include <boost/range/as_array.hpp>

#include <vector>

namespace paal {
namespace ir {

/**
 * @class bdmst_violation_checker
 * @brief Violations checker for the separation oracle
 *      in the bounded degree minimum spanning tree problem.
 */
class bdmst_violation_checker {
    using AuxEdge = min_cut_finder::Edge;
    using AuxVertex = min_cut_finder::Vertex;
    using AuxEdgeList = std::vector<AuxEdge>;
    using Violation = boost::optional<double>;

  public:
    using Candidate = std::pair<AuxVertex, AuxVertex>;
    using CandidateList = std::vector<Candidate>;

    /**
     * Returns an iterator range of violated constraint candidates.
     */
    template <typename Problem, typename LP>
    const CandidateList &get_violation_candidates(const Problem &problem,
                                                  const LP &lp) {
        fill_auxiliary_digraph(problem, lp);
        initialize_candidates(problem);
        return m_candidate_list;
    }

    /**
     * Checks if the given constraint candidate is violated an if it is,
     * returns the violation value and violated constraint ID.
     */
    template <typename Problem>
    Violation check_violation(Candidate candidate, const Problem &problem) {
        double violation = find_violation(candidate.first, candidate.second);
        if (problem.get_compare().g(violation, 0)) {
            return violation;
        } else {
            return Violation{};
        }
    }

    /**
     * Adds a violated constraint to the LP.
     */
    template <typename Problem, typename LP>
    void add_violated_constraint(Candidate violating_pair,
                                 const Problem &problem, LP &lp) {
        if (violating_pair != m_min_cut.get_last_cut()) {
            find_violation(violating_pair.first, violating_pair.second);
        }

        auto const &g = problem.get_graph();
        auto const &index = problem.get_index();

        lp::linear_expression expr;
        for (auto const &e : problem.get_edge_map().right) {
            auto u = get(index, source(e.second, g));
            auto v = get(index, target(e.second, g));
            if (m_min_cut.is_in_source_set(u) &&
                m_min_cut.is_in_source_set(v)) {
                expr += e.first;
            }
        }
        lp.add_row(std::move(expr) <= m_min_cut.source_set_size() - 2);
    }

  private:

    /**
     * Creates the auxiliary directed graph used for feasibility testing.
     */
    template <typename Problem, typename LP>
    void fill_auxiliary_digraph(const Problem &problem, const LP &lp) {
        auto const &g = problem.get_graph();
        auto const &index = problem.get_index();
        m_vertices_num = num_vertices(g);
        m_min_cut.init(m_vertices_num);
        m_src_to_v.resize(m_vertices_num);
        m_v_to_trg.resize(m_vertices_num);

        for (auto const &e : problem.get_edge_map().right) {
            lp::col_id col_idx = e.first;
            double col_val = lp.get_col_value(col_idx) / 2;

            if (!problem.get_compare().e(col_val, 0)) {
                auto u = get(index, source(e.second, g));
                auto v = get(index, target(e.second, g));
                m_min_cut.add_edge_to_graph(u, v, col_val, col_val);
            }
        }

        m_src = m_min_cut.add_vertex_to_graph();
        m_trg = m_min_cut.add_vertex_to_graph();

        for (auto v : boost::as_array(vertices(g))) {
            auto aux_v = get(index, v);
            m_src_to_v[aux_v] = m_min_cut
                .add_edge_to_graph(m_src, aux_v, degree_of(problem, v, lp) / 2)
                .first;
            m_v_to_trg[aux_v] =
                m_min_cut.add_edge_to_graph(aux_v, m_trg, 1).first;
        }
    }

    /**
     * Initializes the list of cut candidates.
     */
    template <typename Problem>
    void initialize_candidates(const Problem &problem) {
        auto const &g = problem.get_graph();
        auto const &index = problem.get_index();
        auto src = *(std::next(vertices(g).first, rand() % m_vertices_num));
        auto aux_src = get(index, src);
        m_candidate_list.clear();
        for (auto v : boost::as_array(vertices(g))) {
            if (v != src) {
                auto aux_v = get(index, v);
                m_candidate_list.push_back(std::make_pair(aux_src, aux_v));
                m_candidate_list.push_back(std::make_pair(aux_v, aux_src));
            }
        }
    }

    /**
     * Calculates the sum of the variables for edges incident with a given
     * vertex.
     */
    template <typename Problem, typename LP, typename Vertex>
    double degree_of(const Problem &problem, const Vertex &v, const LP &lp) {
        double res = 0;

        for (auto e : boost::as_array(out_edges(v, problem.get_graph()))) {
            auto col_id = problem.edge_to_col(e);
            if (col_id) {
                res += lp.get_col_value(*col_id);
            }
        }
        return res;
    }

    /**
     * Finds the most violated set of vertices containing \c src and not
     * containing \c trg and returns its violation value.
     * @param src vertex to be contained in the violating set
     * @param trg vertex not to be contained in the violating set
     * @return violation of the found set
     */
    double find_violation(AuxVertex src, AuxVertex trg) {
        double orig_cap = m_min_cut.get_capacity(m_src_to_v[src]);

        m_min_cut.set_capacity(m_src_to_v[src], m_vertices_num);
        // capacity of m_src_to_v[trg] does not change
        m_min_cut.set_capacity(m_v_to_trg[src], 0);
        m_min_cut.set_capacity(m_v_to_trg[trg], m_vertices_num);

        double min_cut_weight = m_min_cut.find_min_cut(m_src, m_trg);
        double violation = m_vertices_num - 1 - min_cut_weight;

        // reset the original values for the capacities
        m_min_cut.set_capacity(m_src_to_v[src], orig_cap);
        // capacity of m_src_to_v[trg] does not change
        m_min_cut.set_capacity(m_v_to_trg[src], 1);
        m_min_cut.set_capacity(m_v_to_trg[trg], 1);

        return violation;
    }

    int m_vertices_num;

    AuxVertex m_src;
    AuxVertex m_trg;

    AuxEdgeList m_src_to_v;
    AuxEdgeList m_v_to_trg;

    CandidateList m_candidate_list;

    min_cut_finder m_min_cut;
};

} //! ir
} //! paal
#endif // PAAL_BOUNDED_DEGREE_MST_ORACLE_HPP


/* -*- mode: c++; c-basic-offset: 2; indent-tabs-mode: nil; tab-width: 8 -*-
 *
 * $Id: ext.cpp 21225 2014-11-28 21:06:02Z phyy-nx $
 */
#include <cctbx/boost_python/flex_fwd.h>
#include <boost/python.hpp>
#include <boost/python/class.hpp>
#include <scitbx/array_family/flex_types.h>
#include <scitbx/array_family/shared.h>
#include <scitbx/math/basic_statistics.h>
#include <cctbx/miller.h>
#include <cctbx/uctbx.h>

/*

This file contains an implmentation of calc_avg_I from
postrefine/mod_util.py, which averages observed intensities,
does outlier rejection and computes various statistics.

*/

using namespace boost::python;

namespace prime {
  enum Average_Mode {
    Average,  ///< normal avearaging
    Weighted, ///< weighted averaging
    Final     ///< averaging done when the postrefinement is finishing
  };

  typedef
    scitbx::af::shared<cctbx::miller::index<int> > shared_miller;

  struct average_result_store {
    // class for returning results
    shared_miller miller_index;
    scitbx::af::shared<double> I_avg;
    scitbx::af::shared<double> sigI_avg;
    scitbx::af::shared<double> r_meas_w_top;
    scitbx::af::shared<double> r_meas_w_btm;
    scitbx::af::shared<double> r_meas_top;
    scitbx::af::shared<double> r_meas_btm;
    scitbx::af::shared<int> multiplicity;
    scitbx::af::shared<double> I_avg_even;
    scitbx::af::shared<double> I_avg_odd;
    scitbx::af::shared<double> I_avg_even_h;
    scitbx::af::shared<double> I_avg_odd_h;
    scitbx::af::shared<double> I_avg_even_k;
    scitbx::af::shared<double> I_avg_odd_k;
    scitbx::af::shared<double> I_avg_even_l;
    scitbx::af::shared<double> I_avg_odd_l;
    std::string txt_obs_out;
    std::string txt_reject_out;
  };

  class averaging_engine {
    // interface for computing averages
    const int group_no_;
    const scitbx::af::shared<int> group_id_list_;
    const shared_miller miller_index_;
    const shared_miller miller_index_ori_;
    const scitbx::af::shared<double> I_;
    const scitbx::af::shared<double> sigI_;
    const scitbx::af::shared<double> G_;
    const scitbx::af::shared<double> B_;
    const scitbx::af::shared<double> p_set_;
    const scitbx::af::shared<double> rs_set_;
    const scitbx::af::shared<double> wavelength_set_;
    const scitbx::af::shared<double> sin_theta_over_lambda_sq_;
    const scitbx::af::shared<double> SE_;
    const scitbx::af::shared<std::string> pickle_filename_set_;
    public:
    Average_Mode avg_mode_;
    double sigma_max_;
    bool flag_volume_correction_;
    int n_rejection_cycle_;
    bool flag_output_verbose_;

    public: averaging_engine(
      // main constructor, used for passing in arrays. other parmaters exposed
      // as properties to python
      int group_no,
      scitbx::af::shared<int> group_id_list,
      const shared_miller& miller_index,
      const shared_miller& miller_index_ori,
      const scitbx::af::shared<double>& I,
      const scitbx::af::shared<double>& sigI,
      const scitbx::af::shared<double>& G,
      const scitbx::af::shared<double>& B,
      const scitbx::af::shared<double>& p_set,
      const scitbx::af::shared<double>& rs_set,
      const scitbx::af::shared<double>& wavelength_set,
      const scitbx::af::shared<double>& sin_theta_over_lambda_sq,
      const scitbx::af::shared<double>& SE,
      const scitbx::af::shared<std::string>& pickle_filename_set
      ):
        group_no_(group_no),
        group_id_list_(group_id_list),
        miller_index_(miller_index),
        miller_index_ori_(miller_index_ori),
        I_(I),sigI_(sigI),G_(G),B_(B),p_set_(p_set),
        rs_set_(rs_set),
        wavelength_set_(wavelength_set),
        sin_theta_over_lambda_sq_(sin_theta_over_lambda_sq),
        SE_(SE),
        pickle_filename_set_(pickle_filename_set)
    {
      avg_mode_ = Average;
      sigma_max_ = 99.0;
      flag_volume_correction_ = true;
      n_rejection_cycle_ = 1;
      flag_output_verbose_ = false;
    }

    void calc_avg_two_halves(
      const scitbx::af::shared<double>& I_full_group,
      const scitbx::af::shared<double>& SE_norm,
      const Average_Mode& avg_mode_,
      double& I_avg_even,
      double& I_avg_odd
      )
    {
      double I_even_sum = 0;
      double I_odd_sum = 0;
      double I_even_weighted_sum = 0;
      double I_odd_weighted_sum = 0;
      double SE_norm_even_sum = 0;
      double SE_norm_odd_sum = 0;
      I_avg_even = 0;
      I_avg_odd = 0;
      if (I_full_group.size() > 2) {
        for (int i = 0; i < I_full_group.size(); i++) {
          if (i % 2 == 0) {
            I_even_sum += I_full_group[i];
            I_even_weighted_sum += I_full_group[i] * SE_norm[i];
            SE_norm_even_sum += SE_norm[i];
          }
          else {
            I_odd_sum += I_full_group[i];
            I_odd_weighted_sum += I_full_group[i] * SE_norm[i];
            SE_norm_odd_sum += SE_norm[i];
          }
        }

        if (I_full_group.size() % 2 == 1) {
          I_odd_sum += I_full_group[I_full_group.size()-1];
          I_odd_weighted_sum += I_full_group[I_full_group.size()-1] * SE_norm[SE_norm.size()-1];
          SE_norm_odd_sum += SE_norm[SE_norm.size()-1];
        }

        if (avg_mode_ == Weighted || avg_mode_ == Final) {
          I_avg_even = I_even_weighted_sum/SE_norm_even_sum;
          I_avg_odd = I_odd_weighted_sum/SE_norm_odd_sum;
        }
        else {
          SCITBX_ASSERT(avg_mode_ == Average);
          int size;
          if(I_full_group.size() % 2 == 1)
            size = (I_full_group.size()+1)/2;
          else
            size = I_full_group.size()/2;
          I_avg_even = I_even_sum/size;
          I_avg_odd = I_odd_sum/size;
        }
      }
    }

    static const double CONST_SE_MIN_WEIGHT;
    static const double CONST_SE_MAX_WEIGHT;
    static const double CONST_SIG_I_FACTOR;

    public: average_result_store
    calc_avg_I() {
      /*
      Average the intensites.

      I_ is an array with N observations.  miller_indices is an array with M indices.
      M should always be <= N. Observations are arranged into groups where the members
      of the same group have the same miller index.  group_id_list is N long and indexes
      each observation to a group and its miller index. group_id_list should be sorted
      (ascending) and that sort order applied to G, B, etc.
      */

      average_result_store results;
      std::ostringstream txt_obs_out;
      std::ostringstream txt_reject_out;

      // convert to the full intensity and calculate mosaic spread (currently only logged)
      scitbx::af::shared<double> I_full;
      scitbx::af::shared<double> sigI_full;
      scitbx::af::shared<double> mosaic_radian_set;
      for(int x = 0; x < G_.size(); x++) {
        double tmp1 = G_[x] *std::exp(-2*B_[x]*sin_theta_over_lambda_sq_[x]);
        double tmp2 = I_[x];
        double tmp3 = p_set_[x];
        if (x==0)
          printf("c++ check overflow: %70.70f\n",tmp1*tmp2); // Without this printf, I_full comes out differently between
                                                               // python and cpp in the highest significant digits. No idea why.
        I_full.push_back(tmp2/(tmp1*tmp3));
        sigI_full.push_back(sigI_[x]/(G_[x] * std::exp(-2*B_[x]*sin_theta_over_lambda_sq_[x]) * p_set_[x]));
        mosaic_radian_set.push_back(2 * rs_set_[x] * wavelength_set_[x]);
      }

      if (flag_volume_correction_)
        for (int x = 0; x < I_full.size(); x++){
          I_full[x] *= (4.0/3.0) * (rs_set_[x]);
          sigI_full[x] *= (4.0/3.0) * (rs_set_[x]);
        }

      // Iterate over each group of intensites. They will match a single miller_index each
      int obs_ptr = 0; // this will track along the intensites array as each group is processed
      for (int g = 0; g < group_no_; g++) {
        SCITBX_ASSERT(group_id_list_[obs_ptr] == g); // this verifies that group_id_list is sorted

        int obs_ptr_start = obs_ptr;

        // gather intensites and errors for this group
        double max_w = CONST_SE_MAX_WEIGHT;
        double min_w = std::sqrt(CONST_SE_MIN_WEIGHT);
        scitbx::af::shared<double> I_group;
        scitbx::af::shared<double> sigI_group;
        scitbx::af::shared<double> I_full_group;
        scitbx::af::shared<double> sigI_full_group;
        scitbx::af::shared<double> SE_group;
        cctbx::miller::index<int> current_index;
        shared_miller current_index_ori;
        scitbx::af::shared<std::string>pickle_filename_set_group;
        std::ostringstream txt_reject_out_group;

        scitbx::af::shared<int> valid_ptrs;

        bool found_one = false;
        while (obs_ptr < I_.size()) {
          if (group_id_list_[obs_ptr] != g) {
            break; // found them all
          }

          if (found_one) {
            SCITBX_ASSERT(current_index == miller_index_[obs_ptr]);
          }
          else {
            found_one = true;
            current_index = miller_index_[obs_ptr];
          }
          I_group.push_back(I_[obs_ptr]);
          sigI_group.push_back(sigI_[obs_ptr]);
          I_full_group.push_back(I_full[obs_ptr]);
          sigI_full_group.push_back(sigI_full[obs_ptr]);
          SE_group.push_back(SE_[obs_ptr]);
          current_index_ori.push_back(miller_index_ori_[obs_ptr]);
          pickle_filename_set_group.push_back(pickle_filename_set_[obs_ptr]);
          valid_ptrs.push_back(obs_ptr);
          obs_ptr++;
        }
        SCITBX_ASSERT(found_one);

        // log
        char buf[512];
        sprintf(buf, "Reflection: %d,%d,%d\nmeanI    medI  sigI_est sigI_true delta_sigI   n_refl\n",current_index[0],current_index[1],current_index[2]);
        txt_obs_out << buf;
        scitbx::af::shared<double> I_full_group_copy;
        for (int i = 0; i < I_full_group.size(); i++)
          I_full_group_copy.push_back(I_full_group[i]);
        scitbx::math::basic_statistics<double> basic_stat(I_full_group_copy.const_ref());
        scitbx::math::median_functor mf;

        double median_I = mf(I_full_group_copy.ref());
        double mean_I = basic_stat.mean;
        double std_I = 0;
        //if (I_full_group.size() > 1)
        std_I = basic_stat.biased_standard_deviation; // based on what I think numpy is doing compared to basic_statistics

        sprintf(buf, "%6.2f %6.2f %8.2f %8.0f\n", mean_I, median_I, std_I, double(I_full_group.size()));
        txt_obs_out << buf;

        //reject outliers
        if (I_full_group.size() > 2) {
          for (int i_rejection = 0; i_rejection < n_rejection_cycle_; i_rejection++) {
            scitbx::af::shared<double> I_full_group_copy;
            for (int i = 0; i < I_full_group.size(); i++)
              I_full_group_copy.push_back(I_full_group[i]);
            scitbx::math::basic_statistics<double> basic_stat(I_full_group_copy.const_ref());
            scitbx::math::median_functor mf;

            scitbx::af::shared<double> I_group_filtered;
            scitbx::af::shared<double> sigI_group_filtered;
            scitbx::af::shared<double> I_full_group_filtered;
            scitbx::af::shared<double> sigI_full_group_filtered;
            scitbx::af::shared<double> SE_group_filtered;
            shared_miller current_index_ori_filtered;
            scitbx::af::shared<std::string> pickle_filename_set_group_filtered;

            scitbx::af::shared<int> valid_ptrs_filtered;

            double median_I = mf(I_full_group_copy.ref());
            double mean_I = basic_stat.mean;
            double std_I = basic_stat.biased_standard_deviation; // based on what I think numpy is doing compared to basic_statistics

            for (int i = 0; i < I_full_group.size(); i++) {
              double I_full_as_sigma = (I_full_group[i] - median_I) / std_I;
              if (std::abs(I_full_as_sigma) > sigma_max_) {
                char buf[512];
                sprintf(buf, "%s %3.0f %3.0f %3.0f %10.2f %10.2f\n", pickle_filename_set_group[i].c_str(),
                  double(current_index_ori[i][0]), double(current_index_ori[i][1]), double(current_index_ori[i][2]), I_group[i], sigI_group[i]);
                txt_reject_out_group << buf;
              }
              else {
                I_group_filtered.push_back(I_group[i]);
                sigI_group_filtered.push_back(sigI_group[i]);
                I_full_group_filtered.push_back(I_full_group[i]);
                sigI_full_group_filtered.push_back(sigI_full_group[i]);
                SE_group_filtered.push_back(SE_group[i]);
                current_index_ori_filtered.push_back(current_index_ori[i]);
                pickle_filename_set_group_filtered.push_back(pickle_filename_set_group[i]);
                valid_ptrs_filtered.push_back(valid_ptrs[i]);
              }
            }
            I_group = I_group_filtered;
            sigI_group = sigI_group_filtered;
            I_full_group = I_full_group_filtered;
            sigI_full_group = sigI_full_group_filtered;
            SE_group = SE_group_filtered;
            current_index_ori = current_index_ori_filtered;
            pickle_filename_set_group = pickle_filename_set_group_filtered;
            valid_ptrs = valid_ptrs_filtered;

            char buf[512];
            sprintf(buf, "%6.2f %6.2f %8.2f %8.0f\n", mean_I, median_I, std_I, double(I_full_group.size()));
            txt_obs_out << buf;

            if (I_full_group.size() <= 3)
              break;
          }
          if (I_full_group.size() == 0) {
            printf("miller_index (%d, %d, %d) rejected at calc_avg", current_index[0], current_index[1], current_index[2]);
            continue;
          }
        }
        // normalize the SE
        scitbx::af::shared<double> SE_norm;
        double se_max = scitbx::af::max(SE_group.ref());
        double se_min = scitbx::af::min(SE_group.ref());
        double SE_norm_val;
        if (SE_group.size() == 1 || ((se_max-se_min) < 0.1) || avg_mode_ == Average) {
          SE_norm_val = 1;
          for (int i = 0; i < SE_group.size(); i++) {
            SE_norm.push_back(SE_norm_val);
          }
        }
        else {
          double m = (max_w - min_w)/(se_min-se_max);
          double b = max_w - (m*se_min);

          for (int i = 0; i < SE_group.size(); i++) {
            SE_norm.push_back( (m*SE_group[i]) + b);
          }
        }

        double SE_norm_sum = scitbx::af::sum(SE_norm.const_ref());
        SCITBX_ASSERT(SE_norm_sum != 0);
        scitbx::af::shared<double> avg_tmp;
        for (int i = 0; i < SE_norm.size(); i++) {
          avg_tmp.push_back(SE_norm[i] * I_full_group[i]);
        }

        double I_avg = scitbx::af::sum(avg_tmp.const_ref())/SE_norm_sum;
        double sigI_avg = scitbx::af::mean(sigI_full_group.const_ref());

        //Rmeas, Rmeas_w, multiplicity
        int multiplicity = I_full_group.size();
        double r_meas_w_top = 0;
        double r_meas_w_btm = 0;
        double r_meas_top = 0;
        double r_meas_btm = 0;
        double r_meas = 0;
        double r_meas_w = 0;
        if (multiplicity > 1) {
          for (int i = 0; i < multiplicity; i++) {
            r_meas_w_top += std::pow(((I_full_group[i] - I_avg)*SE_norm[i]),2);
            r_meas_w_btm += std::pow(I_full_group[i]*SE_norm[i],2);
            r_meas_top += std::abs(((I_full_group[i] - I_avg)*SE_norm[i]));
            r_meas_btm += std::abs(I_full_group[i]*SE_norm[i]);
          }
          r_meas_w = r_meas_w_top/r_meas_w_btm;
          r_meas = r_meas_top/r_meas_btm;
        }

        //for calculation of cc1/2
        //separate the observations into two groups
        double I_avg_even = 0;
        double I_avg_odd = 0;
        double I_avg_even_h = 0;
        double I_avg_odd_h = 0;
        double I_avg_even_k = 0;
        double I_avg_odd_k = 0;
        double I_avg_even_l = 0;
        double I_avg_odd_l = 0;

        calc_avg_two_halves(I_full_group, SE_norm, avg_mode_, I_avg_even, I_avg_odd);

        //select reflections on h axis
        scitbx::af::shared<double> I_full_group_h;
        scitbx::af::shared<double> SE_norm_h;
        scitbx::af::shared<double> I_full_group_k;
        scitbx::af::shared<double> SE_norm_k;
        scitbx::af::shared<double> I_full_group_l;
        scitbx::af::shared<double> SE_norm_l;
        for (int i = 0; i < I_full_group.size(); i++) {
          if (current_index_ori[i][0] == 0) {
            I_full_group_h.push_back(I_full_group[i]);
            SE_norm_h.push_back(SE_norm[i]);
          }

          if (current_index_ori[i][1] == 0) {
            I_full_group_k.push_back(I_full_group[i]);
            SE_norm_k.push_back(SE_norm[i]);
          }

          if (current_index_ori[i][2] == 0) {
            I_full_group_l.push_back(I_full_group[i]);
            SE_norm_l.push_back(SE_norm[i]);
          }
        }

        calc_avg_two_halves(I_full_group_h, SE_norm_h, avg_mode_, I_avg_even_h, I_avg_odd_h);
        calc_avg_two_halves(I_full_group_k, SE_norm_k, avg_mode_, I_avg_even_k, I_avg_odd_k);
        calc_avg_two_halves(I_full_group_l, SE_norm_l, avg_mode_, I_avg_even_l, I_avg_odd_l);

        // save the results for this group
        results.miller_index.push_back(current_index);
        results.I_avg.push_back(I_avg);
        results.sigI_avg.push_back(sigI_avg);
        results.r_meas_w_top.push_back(r_meas_w_top);
        results.r_meas_w_btm.push_back(r_meas_w_btm);
        results.r_meas_top.push_back(r_meas_top);
        results.r_meas_btm.push_back(r_meas_btm);
        results.multiplicity.push_back(multiplicity);
        results.I_avg_even.push_back(I_avg_even);
        results.I_avg_odd.push_back(I_avg_odd);
        results.I_avg_even_h.push_back(I_avg_even_h);
        results.I_avg_odd_h.push_back(I_avg_odd_h);
        results.I_avg_even_k.push_back(I_avg_even_k);
        results.I_avg_odd_k.push_back(I_avg_odd_k);
        results.I_avg_even_l.push_back(I_avg_even_l);
        results.I_avg_odd_l.push_back(I_avg_odd_l);

        if (flag_output_verbose_) {
          txt_obs_out << "    I_o        sigI_o    G      B     Eoc      rs    lambda rocking(deg) W     I_full     sigI_full\n";
          for (int i = 0; i < I_full_group.size(); i++) {
            char buf[512];
            sprintf(buf, "%10.2f %10.2f %6.2f %6.2f %6.2f %8.5f %8.5f %8.5f %6.2f %10.2f %10.2f\n",
              I_group[i],sigI_group[i],1/G_[valid_ptrs[i]],B_[valid_ptrs[i]],p_set_[valid_ptrs[i]],rs_set_[valid_ptrs[i]],
              wavelength_set_[valid_ptrs[i]],mosaic_radian_set[valid_ptrs[i]]*180/scitbx::constants::pi,SE_norm[i],
              I_full_group[i],sigI_full[i]);
            txt_obs_out << buf;
          }
          char buf[512];
          sprintf(buf, "Merged I, sigI: %6.2f, %6.2f\n",I_avg,sigI_avg);
          txt_obs_out << buf;
          sprintf(buf, "Rmeas: %6.2f Qw: %6.2f\n",r_meas,r_meas_w);
          txt_obs_out << buf;
          sprintf(buf, "No. total observed: %4.0f No. after rejection: %4.0f\n", double(obs_ptr-obs_ptr_start), double(I_full_group.size()));
          txt_obs_out << buf;
          txt_obs_out << "List of rejected observations:\n";
          txt_obs_out << txt_reject_out_group.str();
        }
        txt_reject_out << txt_reject_out_group.str();
      }
      results.txt_obs_out = txt_obs_out.str();
      results.txt_reject_out = txt_reject_out.str();

      return results;
    }
  };

const double averaging_engine::CONST_SE_MIN_WEIGHT = 0.17;
const double averaging_engine::CONST_SE_MAX_WEIGHT = 1.0;
const double averaging_engine::CONST_SIG_I_FACTOR = 1.5;

namespace boost_python { namespace {
  void
  init_module() {
    using namespace boost::python;
    typedef return_value_policy<return_by_value> rbv;
    typedef default_call_policies dcp;
    typedef averaging_engine w_t;

    class_<w_t>("averaging_engine", no_init)
      .def(init<int,scitbx::af::shared<int>,
        const shared_miller&,const shared_miller&,
        const scitbx::af::shared<double>&,
        const scitbx::af::shared<double>&,
        const scitbx::af::shared<double>&,
        const scitbx::af::shared<double>&,
        const scitbx::af::shared<double>&,
        const scitbx::af::shared<double>&,
        const scitbx::af::shared<double>&,
        const scitbx::af::shared<double>&,
        const scitbx::af::shared<double>&,
        const scitbx::af::shared<std::string>&
        >(
        (arg("group_no"),arg("group_id_list"),
        arg("miller_list"),arg("miller_list_ori"),arg("I"),arg("sigI"),
        arg("G"),arg("B"),arg("p_set"),arg("rs_set"),
        arg("wavelength_set"),arg("sin_theta_over_lambda_sq"),arg("SE"),
        arg("pickle_filename_set")
        )))
      .def("calc_avg_I", &averaging_engine::calc_avg_I)
      .add_property("avg_mode",
        make_getter(&averaging_engine::avg_mode_, rbv()),
        make_setter(&averaging_engine::avg_mode_, dcp()))
      .add_property("sigma_max",
        make_getter(&averaging_engine::sigma_max_, rbv()),
        make_setter(&averaging_engine::sigma_max_, dcp()))
      .add_property("flag_volume_correction",
        make_getter(&averaging_engine::flag_volume_correction_, rbv()),
        make_setter(&averaging_engine::flag_volume_correction_, dcp()))
      .add_property("n_rejection_cycle",
        make_getter(&averaging_engine::n_rejection_cycle_, rbv()),
        make_setter(&averaging_engine::n_rejection_cycle_, dcp()))
      .add_property("flag_output_verbose",
        make_getter(&averaging_engine::flag_output_verbose_, rbv()),
        make_setter(&averaging_engine::flag_output_verbose_, dcp()))
      ;

    class_<average_result_store>("average_result_store",init<>())
      .add_property("miller_index",
        make_getter(&average_result_store::miller_index, rbv()),
        make_setter(&average_result_store::miller_index, dcp()))
      .add_property("I_avg",
        make_getter(&average_result_store::I_avg, rbv()),
        make_setter(&average_result_store::I_avg, dcp()))
      .add_property("sigI_avg",
        make_getter(&average_result_store::sigI_avg, rbv()),
        make_setter(&average_result_store::sigI_avg, dcp()))
      .add_property("r_meas_w_top",
        make_getter(&average_result_store::r_meas_w_top, rbv()),
        make_setter(&average_result_store::r_meas_w_top, dcp()))
      .add_property("r_meas_w_btm",
        make_getter(&average_result_store::r_meas_w_btm, rbv()),
        make_setter(&average_result_store::r_meas_w_btm, dcp()))
      .add_property("r_meas_top",
        make_getter(&average_result_store::r_meas_top, rbv()),
        make_setter(&average_result_store::r_meas_top, dcp()))
      .add_property("r_meas_btm",
        make_getter(&average_result_store::r_meas_btm, rbv()),
        make_setter(&average_result_store::r_meas_btm, dcp()))
      .add_property("multiplicity",
        make_getter(&average_result_store::multiplicity, rbv()),
        make_setter(&average_result_store::multiplicity, dcp()))
      .add_property("I_avg_even",
        make_getter(&average_result_store::I_avg_even, rbv()),
        make_setter(&average_result_store::I_avg_even, dcp()))
      .add_property("I_avg_odd",
        make_getter(&average_result_store::I_avg_odd, rbv()),
        make_setter(&average_result_store::I_avg_odd, dcp()))
      .add_property("I_avg_even_h",
        make_getter(&average_result_store::I_avg_even_h, rbv()),
        make_setter(&average_result_store::I_avg_even_h, dcp()))
      .add_property("I_avg_odd_h",
        make_getter(&average_result_store::I_avg_odd_h, rbv()),
        make_setter(&average_result_store::I_avg_odd_h, dcp()))
      .add_property("I_avg_even_k",
        make_getter(&average_result_store::I_avg_even_k, rbv()),
        make_setter(&average_result_store::I_avg_even_k, dcp()))
      .add_property("I_avg_odd_k",
        make_getter(&average_result_store::I_avg_odd_k, rbv()),
        make_setter(&average_result_store::I_avg_odd_k, dcp()))
      .add_property("I_avg_even_l",
        make_getter(&average_result_store::I_avg_even_l, rbv()),
        make_setter(&average_result_store::I_avg_even_l, dcp()))
      .add_property("I_avg_odd_l",
        make_getter(&average_result_store::I_avg_odd_l, rbv()),
        make_setter(&average_result_store::I_avg_odd_l, dcp()))
      .add_property("txt_obs_out",
        make_getter(&average_result_store::txt_obs_out, rbv()),
        make_setter(&average_result_store::txt_obs_out, dcp()))
      .add_property("txt_reject_out",
        make_getter(&average_result_store::txt_reject_out, rbv()),
        make_setter(&average_result_store::txt_reject_out, dcp()))
    ;
  };

  using namespace boost::python;

  void export_average_mode()
  {
    enum_<Average_Mode>("Average_Mode")
      .value("Average", Average)
      .value("Weighted", Weighted)
      .value("Final", Final);
  }

}}} // namespace prime::boost_python::<anonymous>

BOOST_PYTHON_MODULE(prime_ext)
{
  prime::boost_python::init_module();
  prime::boost_python::export_average_mode();

}


//-----------------------------------------------------------------------------
// Created on: 03 December 2016
//-----------------------------------------------------------------------------
// Copyright (c) 2017, Sergey Slyadnev
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
//    * Redistributions of source code must retain the above copyright
//      notice, this list of conditions and the following disclaimer.
//    * Redistributions in binary form must reproduce the above copyright
//      notice, this list of conditions and the following disclaimer in the
//      documentation and/or other materials provided with the distribution.
//    * Neither the name of the copyright holder(s) nor the
//      names of all contributors may be used to endorse or promote products
//      derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE AUTHORS 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.
//-----------------------------------------------------------------------------

// Own include
#include <asiAlgo_PlaneOnPoints.h>

// Eigen includes
#pragma warning(disable : 4701 4702)
#include <Eigen/Dense>
#pragma warning(default : 4701 4702)

#undef COUT_DEBUG
#if defined COUT_DEBUG
  #pragma message("===== warning: COUT_DEBUG is enabled")
#endif

//-----------------------------------------------------------------------------

bool compare(const std::pair<double, int>& p1, const std::pair<double, int>& p2)
{
  return p1.first > p2.first;
}

//-----------------------------------------------------------------------------

asiAlgo_PlaneOnPoints::asiAlgo_PlaneOnPoints(ActAPI_ProgressEntry progress,
                                             ActAPI_PlotterEntry  plotter)
: ActAPI_IAlgorithm(progress, plotter)
{}

//-----------------------------------------------------------------------------

bool asiAlgo_PlaneOnPoints::Build(const std::vector<gp_XYZ>& points,
                                  gp_Pln&                    result) const
{
  return this->internalBuild(points, result);
}

//-----------------------------------------------------------------------------

bool asiAlgo_PlaneOnPoints::Build(const Handle(asiAlgo_BaseCloud<double>)& points,
                                  gp_Pln&                                  result) const
{
  // Repack point cloud to a vector.
  std::vector<gp_XYZ> pointsVec;
  //
  for ( int k = 0; k < points->GetNumberOfElements(); ++k )
    pointsVec.push_back( points->GetElement(k) );

  return this->internalBuild(pointsVec, result);
}

//-----------------------------------------------------------------------------

bool asiAlgo_PlaneOnPoints::internalBuild(const std::vector<gp_XYZ>& points,
                                          gp_Pln&                    result) const
{
  const int nPts = (int) points.size();

  /* ======================
   *  Calculate mean point
   * ====================== */

  gp_XYZ mu;
  for ( size_t i = 0; i < points.size(); ++i )
  {
    mu += points[i];
  }
  mu /= nPts;

  /* =========================
   *  Build covariance matrix
   * ========================= */

  Eigen::Matrix3d C;
  for ( int j = 1; j <= 3; ++j )
  {
    for ( int k = 1; k <= 3; ++k )
    {
      C(j-1, k-1) = 0.0; // TODO: is that necessary?
    }
  }

  for ( size_t i = 0; i < points.size(); ++i )
  {
    const gp_XYZ& p      = points[i];
    gp_XYZ        p_dash = p - mu;

    for ( int j = 1; j <= 3; ++j )
    {
      for ( int k = 1; k <= 3; ++k )
      {
        C(j-1, k-1) += ( p_dash.Coord(j)*p_dash.Coord(k) );
      }
    }
  }

  for ( int j = 1; j <= 3; ++j )
  {
    for ( int k = 1; k <= 3; ++k )
    {
      C(j-1, k-1) /= nPts;
    }
  }

  Eigen::EigenSolver<Eigen::Matrix3d> EigenSolver(C);

#if defined COUT_DEBUG
  std::cout << "\tCovariance matrix: " << std::endl << C << std::endl;
  std::cout << "\tThe eigen values of C are:" << std::endl << EigenSolver.eigenvalues() << std::endl;
  std::cout << "\tThe matrix of eigenvectors, V, is:" << std::endl << EigenSolver.eigenvectors() << std::endl << std::endl;
#endif

  Eigen::Vector3cd v1 = EigenSolver.eigenvectors().col(0);
  Eigen::Vector3cd v2 = EigenSolver.eigenvectors().col(1);
  Eigen::Vector3cd v3 = EigenSolver.eigenvectors().col(2);

  gp_Vec V[3] = { gp_Vec( v1.x().real(), v1.y().real(), v1.z().real() ),
                  gp_Vec( v2.x().real(), v2.y().real(), v2.z().real() ),
                  gp_Vec( v3.x().real(), v3.y().real(), v3.z().real() ) };
  //
  std::vector< std::pair<double, int> >
    lambda { std::pair<double, int>( EigenSolver.eigenvalues()(0).real(), 0 ),
             std::pair<double, int>( EigenSolver.eigenvalues()(1).real(), 1 ),
             std::pair<double, int>( EigenSolver.eigenvalues()(2).real(), 2 ) };
  //
  std::sort(lambda.begin(), lambda.end(), compare);
  //
  gp_Ax1 ax_X(mu, V[lambda[0].second]);
  gp_Ax1 ax_Y(mu, V[lambda[1].second]);
  gp_Ax1 ax_Z(mu, V[lambda[2].second]);
  //
  gp_Vec vec_X( ax_X.Direction() );
  gp_Vec vec_Y( ax_Y.Direction() );
  gp_Vec vec_Z( ax_Z.Direction() );
  //
  if ( (vec_X ^ vec_Y).Magnitude() < gp::Resolution() ||
       (vec_X ^ vec_Z).Magnitude() < gp::Resolution() ||
       (vec_Y ^ vec_Z).Magnitude() < gp::Resolution() )
  {
    std::cout << "Warning: degenerated normal" << std::endl;
    return false; // Degenerated normal
  }

  // Check if the system is right-handed
  const double ang = ax_X.Direction().AngleWithRef( ax_Y.Direction(), ax_Z.Direction() );
  if ( ang < 0 )
  {
    gp_Ax1 tmp = ax_X;
    ax_X = ax_Y;
    ax_Y = tmp;
  }

  // Store results
  gp_Ax3 ax3( gp_Pnt(mu), ax_Z.Direction(), ax_X.Direction() );
  result.SetPosition(ax3);
  //
  return true;
}


/**
 * @file TriangleElement.cpp
 * @author Adam Li (adamli@umich.edu)
 * @date 2021-11-12
 *
 * @brief Implements TriangleElement from meshGenerator.py.
 *
 */

#include "TriangleElement.hpp"
#include <cstring>
#include <cstdlib>
#include <Eigen/Dense>
#include "TriangularMesh.hpp"

using sel_map::core::ElemArray;
using sel_map::core::ElemArrayBase;

#include "core/EigenMathUtil.hpp"
namespace EigenMathUtil = sel_map::core::EigenMathUtil;

//Avoid rewriting the whole thing everytime
using sel_map::mesh::TriangularMesh;
using sel_map::mesh::TriangleElement;
using sel_map::mesh::Element;

// TriangleElement::TriangleElement(const std::reference_wrapper<TriangularMesh>& parentMesh, unsigned int index, double elementLength, unsigned int pointLimit, float heightPartition,
//                     double heightSafetyCheck, unsigned int classes)
//                       : Element<TriangularMesh>(parentMesh, index, elementLength, pointLimit, heightPartition, heightSafetyCheck, classes)
//                     //   , lambdas()
TriangleElement::TriangleElement(const std::reference_wrapper<TriangularMesh>& parentMesh, unsigned int index)
                      : Element<TriangularMesh>(parentMesh, index)
{
    // Initializer
    active = false;
}

TriangleElement::~TriangleElement()
{
    // Destructor
}

void TriangleElement::reset()
{
    active = false;
    dirichletParamsForTerrainClass.setZero();
}

/**
 * @brief Computes the Barycentric coordinates of all points within the element from Cartesian points.
 */
// void TriangleElement::cartesianToBarycentric()
// {
//     // Auto keywords in eigen are really dangerous, oops
//     // These are all effectively Map<vector2d, unaligned> types
//     auto vertices = parentMesh.get().vertexVector(simplex(), Eigen::all);
//     auto v1 = vertices(0, Eigen::seqN(Eigen::fix<0>, Eigen::fix<2>));
//     auto v2 = vertices(0, Eigen::seqN(Eigen::fix<0>, Eigen::fix<2>));
//     auto v3 = vertices(0, Eigen::seqN(Eigen::fix<0>, Eigen::fix<2>));
//     // These two variables should never be touched again! they're just intermediates
//     auto area_iexp1 = v2 - v1;
//     auto area_iexp2 = v3 - v1;
//     double area = EigenMathUtil::crossProduct2D(area_iexp1, area_iexp2);

//     // This is effectively Map<ArrayXd, aligned> type
//     auto eigenPointsArray = points.getEigenArray();
//     auto eigenLambdaArray = lambdas.reserve(points.length(), true).getEigenArray();

//     // The following is simplified to keep eveything in scalar math
//     // alpha = abs(np.cross(p2-self.points[:,0:2], p3-self.points[:,0:2]) / area)
//     // beta  = abs(np.cross(p1-self.points[:,0:2], p3-self.points[:,0:2]) / area)
//     // gamma = abs(np.cross(p1-self.points[:,0:2], p2-self.points[:,0:2]) / area)
//     // self.lambdas = np.array([alpha, beta, gamma])

//     double v1_x_v2 = EigenMathUtil::crossProduct2D(v1, v2);
//     double v1_x_v3 = EigenMathUtil::crossProduct2D(v1, v3);
//     double v2_x_v3 = EigenMathUtil::crossProduct2D(v2, v3);

//     // update alpha
//     eigenLambdaArray.col(0) = ((v2_x_v3
//                                  + (v2(1) - v3(1)) * eigenPointsArray.col(0)
//                                  + (v3(0) - v2(0)) * eigenPointsArray.col(1)
//                               ) / area).abs();

//     // update beta
//     eigenLambdaArray.col(1) = ((v1_x_v3
//                                  + (v1(1) - v3(1)) * eigenPointsArray.col(0)
//                                  + (v3(0) - v1(0)) * eigenPointsArray.col(1)
//                               ) / area).abs();

//     // update gamma
//     eigenLambdaArray.col(2) = ((v1_x_v2
//                                  + (v1(1) - v2(1)) * eigenPointsArray.col(0)
//                                  + (v2(0) - v1(0)) * eigenPointsArray.col(1)
//                               ) / area).abs();
// }

void TriangleElement::computeElementNormal()
{
    // auto vertices = parentMesh.get().vertexVector(simplex(), Eigen::all);
    // auto a = vertices(2, Eigen::seqN(Eigen::fix<0>(), Eigen::fix<3>()));
    // auto b = vertices(1, Eigen::seqN(Eigen::fix<0>(), Eigen::fix<3>()));
    // auto c = vertices(0, Eigen::seqN(Eigen::fix<0>(), Eigen::fix<3>()));

    // auto v1 = a - b;
    // auto v2 = c - b;
    // normal = v1.matrix().cross(v2.matrix());
}

Eigen::Block<Eigen::Array<unsigned int, -1, 3, 1>, 1, 3, true> TriangleElement::simplex()
{
    return parentMesh.get().simplices.row(index);
}


#include <cassert>
#include <NTL/RR.h>
#include "ring.h"

/**********************************************************************/
ostream& operator<< (ostream& s, const ring& rg) {
  s << rg.label << ": type = " << rg.type << ", mR = " << rg.mR << ", pR = " << rg.pR
    << ", gR = " << rg.gR << ", tR = " << rg.tR << ", depth = " << rg.depth << endl;

  return s;
}

ostream& operator<< (ostream& s, const param& prm) {
  s << endl << "<" << prm.label << ">" << endl;
  s << "l = " << prm.l << ", ";
  s << "r = " << prm.r << ", ";
  s << "s = " << prm.s << ", ";
  s << "lw = " << prm.lw << endl;
  s << prm.rg << prm.ns;

  return s;
}

/**********************************************************************/
void ring::to_eta (ZZ_pX& w, vec_ZZ_p& a) {
  if (type != prime) {
    cout << "Error: to_eta" << endl;
    cout << "- not yet implemented for " << type << "-type" << endl;
  }

  for (long i = 0; i < gR; ++i) w += a[i] * conv<ZZ_pX>(eta(i));
}

void ring::to_delta (ZZ_pX& w, vec_ZZ_p& a) {
  if (type != prime) {
    cout << "Error: to_eta" << endl;
    cout << "- not yet implemented for " << type << "-type" << endl;
  }

  for (long i = 0; i < gR; ++i) w += a[i] * delta(i);
}

void mult (vec_ZZ_p& x, const vec_ZZ_p& a, const vec_ZZ_p& b) {
  for (long i = 0; i < a.length(); ++i) x[i] = a[i] * b[i];
}

void sqr (vec_ZZ_p& x, const vec_ZZ_p& a) {
  for (long i = 0; i < a.length(); ++i) x[i] = sqr(a[i]);
}

void add_scalar (vec_ZZ_p& x, const vec_ZZ_p& a, const ZZ_p& b) {
  for (long i = 0; i < a.length(); ++i) x[i] = a[i] + b;
}

void add_scalar (vec_ZZ_p& x, const ZZ_p& a, const vec_ZZ_p& b) {
  for (long i = 0; i < b.length(); ++i) x[i] = a + b[i];
}

void zzx_convolution (vec_ZZ_p& c, const vec_ZZ_p& a, const vec_ZZ_p& b) {
  ZZ_pX aX, bX;
  long g = a.length();

  aX.SetLength(g);
  bX.SetLength(g);
  c.SetLength(g);

  aX[0] = a[0];
  for (long i = 1; i < g; ++i) aX[i] = a[g - i];
  for (long i = 0; i < g; ++i) bX[i] = b[i];

  ZZ_pX w = aX * bX;
  w.SetLength(2*g-1);

  // c = w mod m, m = x^gR - 1
  for (long i = 0; i < g-1; i++) c[i] = w[i] + w[i + g];
  c[g-1] = w[g-1];
}

// Transpose a (m*n)-dim. vector 'a'
// into a m x n matrix (m-dim vector of n-dim vector)'w'
void mat (Vec<vec_ZZ_p>& aa, const vec_ZZ_p& a, const long n, const long m) {
  aa.SetLength(m);
  for (long j = 0; j < m; ++j) aa[j].SetLength(n);

  for (long j = 0; j < m; ++j)
    for (long i = 0; i < n; ++i)
      aa[j][i] = a[i + n*j];
}

// returns m-dim vector of n-dim vectors
void mat (Vec<vec_ZZ>& aa, const vec_ZZ& a, const long n, const long m) {
  aa.SetLength(m);
  for (long j = 0; j < m; ++j) aa[j].SetLength(n);

  for (long j = 0; j < m; ++j)
    for (long i = 0; i < n; ++i)
      aa[j][i] = a[i + n*j];
}

void transpose (Vec<vec_ZZ_p>& b, const Vec<vec_ZZ_p>& a) {
  long m = a.length();
  long n = a[0].length();

  b.SetLength(n);
  for (long j = 0; j < n; ++j) b[j].SetLength(m);

  for (long j = 0; j < n; ++j)
    for (long i = 0; i < m; ++i)
      b[j][i] = a[i][j];
}

void transpose (Vec<vec_ZZ>& b, const Vec<vec_ZZ>& a) {
  long m = a.length();
  long n = a[0].length();

  b.SetLength(n);
  for (long j = 0; j < n; ++j) b[j].SetLength(m);

  for (long i = 0; i < m; ++i)
    for (long j = 0; j < n; ++j)
      b[j][i] = a[i][j];
}

void vec (vec_ZZ_p& b, const Vec<vec_ZZ_p>& a) {
  long m = a.length();
  long n = a[0].length();

  b.SetLength(m*n);

  for (long i = 0; i < m; ++i)
    for (long j = 0; j < n; ++j)
      b[i*n+j] = a[i][j];
}

void transpose_vec (vec_ZZ_p& b, const Vec<vec_ZZ_p>& a) {
  long m = a.length();
  long n = a[0].length();

  b.SetLength(m*n);

  for (long i = 0; i < n; ++i)
    for (long j = 0; j < m; ++j)
      b[i*m+j] = a[j][i];
}

void transpose_vec (vec_ZZ& b, const Vec<vec_ZZ>& a) {
  long m = a.length();
  long n = a[0].length();

  b.SetLength(m*n);

  for (long i = 0; i < n; ++i)
    for (long j = 0; j < m; ++j)
      b[i*m+j] = a[j][i];
}

void diagonal (vec_ZZ& MD, const long k, const vec_ZZ& M) {
  long dim = M.length();
  MD.SetLength(dim);

  for (long i = 0; i < dim; ++i) MD[i] = M[(i+(i+k)) % dim];
}

void diagonal (vec_ZZ_p& MD, const long k, const vec_ZZ_p& M) {
  long dim = M.length();
  MD.SetLength(dim);

  for (long i = 0; i < dim; ++i) MD[i] = M[(i+(i+k)) % dim];
}

void naiveApplyMatrix (vec_ZZ_p& b, const vec_ZZ_p& m, const vec_ZZ_p& a) {
  long dim = m.length();
  vec_ZZ_p S; S.SetLength(dim);

  vec_ZZ_p aa; //aa.SetLength(dim);
  VectorCopy(aa, a, dim);

  for (long i=0; i < dim; ++i) {
    vec_ZZ_p md;
    diagonal(md, i, m);

    vec_ZZ_p w; w.SetLength(dim);
    mult(w, md, aa);

    shift(aa);
    S += w;
  }

  b = S;
}

void applyMatrix (vec_ZZ_p& b, const vec_ZZ_p& m, const vec_ZZ_p& a) {
  zzx_convolution(b, a, m);
  //naiveApplyMatrix(b, m, a);  // slow, for debug only
}

void ring::applyOmega (vec_ZZ_p& b, const vec_ZZ_p& a) {
  if (type == prime) {
    vec_ZZ_p w = conv<vec_ZZ_p>(omega);
    applyMatrix(b, w, a);
  }
  else {  // type == composite
    long gl = left->gR;
    long gr = right->gR;

    Vec<vec_ZZ_p> w;
    mat(w, a, gr, gl);
    for (long i = 0; i < gl; ++i) right->applyOmega(w[i], w[i]);

    Vec<vec_ZZ_p> wt;
    transpose(wt, w);
    for (long j = 0; j < gr; ++j) left->applyOmega(wt[j], wt[j]);

    transpose_vec(b, wt);
  }
}

void ring::applyOmega (vec_ZZ& b, const vec_ZZ& a) {
  vec_ZZ_p _b;
  applyOmega(_b, conv<vec_ZZ_p>(a));
  b = conv<vec_ZZ>(_b);
}

void ring::applyOmegaInv (vec_ZZ_p& b, const vec_ZZ_p& a) {
  if (type == prime) {
    vec_ZZ_p w = conv<vec_ZZ_p>(omega_inv);
    applyMatrix(b, w, a);
  }
  else {  // type == composite
    long gl = left->gR;
    long gr = right->gR;

    Vec<vec_ZZ_p> w;
    mat(w, a, gr, gl);
    for (long i = 0; i < gl; ++i) right->applyOmegaInv(w[i], w[i]);

    Vec<vec_ZZ_p> wt;
    transpose(wt, w);
    for (long j = 0; j < gr; ++j) left->applyOmegaInv(wt[j], wt[j]);

    transpose_vec(b, wt);
  }
}

void ring::applyOmegaInv (vec_ZZ& b, const vec_ZZ& a) {
  vec_ZZ_p _b;
  applyOmegaInv(_b, conv<vec_ZZ_p>(a));
  b = conv<vec_ZZ>(_b);
}

/**********************************************************************/
ZZ inf_norm (const vec_ZZ& a) {
  ZZ w(0);
  for (long i = 0; i < a.length(); ++i)
    if (w < abs(a[i])) w = abs(a[i]);
  return w;
}

void center_lift (vec_ZZ& b, const vec_ZZ_p& a) {
  b.SetLength(a.length());
  ZZ q = ZZ_p::modulus();
  for (long i=0; i < a.length(); ++i) {
    if (conv<ZZ>(a[i]) >= q/2) b[i] = conv<ZZ>(a[i]) - q;
    else b[i] = conv<ZZ>(a[i]);
  }
}

void center_lift (vec_ZZ& b, const vec_ZZ& a, ZZ& p) {
  ZZ_pContext context;
  context.save();

  ZZ_p::init(p);
  center_lift(b, conv<vec_ZZ_p>(a));

  context.restore();
}

vec_ZZ center_lift (const vec_ZZ_p& a) {
  vec_ZZ b;
  center_lift(b, a);
  return b;
}

vec_ZZ center_lift (const vec_ZZ& a, ZZ& p) {
  vec_ZZ b;
  center_lift(b, a, p);
  return b;
}

// q -> q1
void rescale (vec_ZZ &b, const vec_ZZ &a, const ZZ& q, const ZZ& q1) {
  b.SetLength(a.length());
  for (long i = 0; i < a.length(); ++i) {
    b[i] = (a[i] * q1 + q/2) / q;
    if (b[i] >= q1) b[i] -= q1;
    //b[i] = b[i] % q1;  /*@@*/
  }
}

void lsd_rescale (vec_ZZ &b, const vec_ZZ &a, const ZZ& P, const ZZ& t) {
  long dim = a.length();
  b.SetLength(dim);
  vec_ZZ delta;
  delta.SetLength(dim);

  ZZ _g, u, v;
  XGCD(_g, u, v, P, t);  // 1 = w = u*P + v*t
  ZZ w1 = v*t; // = 1 mod P, = 0 mod t

  for (long i = 0; i < dim; ++i) {
    delta[i] = (((- a[i]) % P) * w1) % (P*t);
    b[i] = (a[i] + delta[i]) / P;
  }
}

/***********************************************************************/
/* returns the power decomp of x mod 2^n  */
void power_decomp (Vec<vec_ZZ>& z, long n, const vec_ZZ& x) {
  long dim = x.length();
  vec_ZZ zero; zero.SetLength(dim);

  z.SetLength(n);
  for (long j = 0; j < n; ++j) z[j] = zero;

  Vec<Vec<vec_ZZ>> y;
  y.SetLength(n);
  for (long i = 0; i < n; ++i) {
    y[i].SetLength(n-i);
    for (long j = 0; j < n-i; ++j) y[i][j] = zero;
  }

  for (long i = 0; i < n; ++i) {
    // y[i][0] = (x - sum_{j=0}^{i-1} 2^j * y[j][i-j]) / 2^i
    y[i][0] = x;
    for (long j = 0; j < i; ++j) {
      for (long k = 0; k < dim; ++k) {
	y[i][0][k] = (y[i][0][k] - y[j][i-j][k]) % power(ZZ(2), n);
	y[i][0][k] /= ZZ(2);
      }
    }

    // y[i][j] = y[i][j-1]^2  (j = 1..(n-1-i))
    for (long j = 1; j < n-i; ++j) {
      for (long k = 0; k < dim; ++k) y[i][j][k] = (y[i][j-1][k] * y[i][j-1][k]) % power(ZZ(2), n);
    }
  }

  for (long i = 0; i < n; ++i) {
    z[i] = y[i][n-1-i];
    for (long k = 0; k < dim; ++k) z[i][k] = z[i][k] % power(ZZ(2), n-i);
  }
}

// returns (x >> m) mod 2^(n-m) given (x mod 2^n)
void right_shift (vec_ZZ& z, const vec_ZZ& x, long m, long n) {
  long dim = x.length();
  vec_ZZ zero; zero.SetLength(dim);

  long l = n - m;

  Vec<vec_ZZ> y; y.SetLength(n);
  for (long j = 0; j < n; ++j) y[j] = zero;
  power_decomp(y, n, x);

  // Drop the least m bits (LowerClear)
  for (long i = 0; i < l; ++i) y[i] = y[m+i];

  // z = y[0] + 2*y[1] + ... + 2^(l-1)*y[l-1]
  z = zero;
  for (long i = 0; i < l; ++i) {
    // z = (z + (2^i)*y[i]) % 2^l
    vec_ZZ w = zero;
    for (long j = 0; j < dim; ++j) {
      w[j] = (power(ZZ(2),i) * y[i][j]) % power(ZZ(2),l);
      z[j] = (z[j] + w[j]) % power(ZZ(2),l);
    }
  }
}

void right_shift (vec_ZZ_p& z, const vec_ZZ_p& x, long m, long n) {
  vec_ZZ _z; _z.SetLength(x.length());
  right_shift(_z, conv<vec_ZZ>(x), m, n);
  z = conv<vec_ZZ_p>(_z);
}

// returns the (i,j)-th component of the cyclic matrix with the first row of 'm'
ZZ_p to_mat (const vec_ZZ_p& m, long i, long j) {
  long dim = m.length();
  assert(0 <= i && i < dim && 0 <= j && j < dim);
  return m[(i + j) % dim];
}

/***********************************************************************/

/* returns the tensor product of a and a1 */
void tensor_prod (vec_ZZ& A, const vec_ZZ& a, const vec_ZZ& a1, const ZZ& q) {
  long gout = a.length();
  long gin = a1.length();

  for (long i = 0; i < gout; ++i) {
    for (long j = 0; j < gin; ++j) {
      A[i*gin + j] = (a[i] * a1[j]) % q;
    }
  }
}

// left shift of vectors
void shift (vec_ZZ& b, const long i, const vec_ZZ& a) {
  long dim = a.length();
  vec_ZZ w;
  w.SetLength(dim);
  for (long j = 0; j < dim; ++j) w[j] = a[(j + i) % dim];
  b = w;
}

void shift (vec_ZZ_p& b, const long i, const vec_ZZ_p& a) {
  long dim = a.length();
  vec_ZZ_p w;
  w.SetLength(dim);
  for (long j = 0; j < dim; ++j) w[j] = a[(j + i) % dim];
  b = w;
}

vec_ZZ shift (const long i, const vec_ZZ& a) {
  vec_ZZ w;
  shift(w, i, a);
  return w;
}

// one-position left shift
void shift (vec_ZZ_p& a) {
  long dim = a.length();
  ZZ_p w = a[0];
  for (long i = 1; i < dim; ++i) a[i-1] = a[i];
  a[dim-1] = w;
}

void shift_right (vec_ZZ& b, const long i, const vec_ZZ& a) {
  shift(b, a.length()-i, a);
}

void shift_right (vec_ZZ_p& b, const long i, const vec_ZZ_p& a) {
  shift(b, a.length()-i, a);
}

// one-position right shift
void right_shift (vec_ZZ_p& a) {
  long dim = a.length();
  ZZ_p w = a[dim-1];
  for (long i = dim-1; i > 0; --i) a[i] = a[i-1];
  a[0] = w;
}

void ring::deep_shift (string _label, vec_ZZ& v, long i) {
  if ((type == prime) && (label == _label)) {
    shift(v, (i<0)?(i+gR):i, v);
  }
  else if ((type == prime) && (label != _label)) {
    return;
  }
  else {  // composite ring
    long gl = left->gR;
    long gr = right->gR;

    Vec<vec_ZZ> w;
    mat(w, v, gr, gl);
    for (long j = 0; j < gl; ++j) right->deep_shift(_label, w[j], i);

    Vec<vec_ZZ> wt;
    transpose(wt, w);
    for (long j = 0; j < gr; ++j) left->deep_shift(_label, wt[j], i);

    transpose_vec(v, wt);
  }
}

void ring::deep_shift_right (string _label, vec_ZZ& v, long i) {
  deep_shift(_label, v, -i);
}

/***********************************************************************/
void get_subrings (ring& rg, Vec<ring>& subrings) {
  if (rg.type == prime) {
    append(subrings, rg);
  }
  else {  // type == composite
    get_subrings(*(rg.right), subrings);
    get_subrings(*(rg.left), subrings);
  }
}

long ring::get_subrings (Vec<ring>& subrings) {
  ::get_subrings(*this, subrings);
  return subrings.length();
}

long param::get_subrings (Vec<ring>& subrings) {
  return rg.get_subrings(subrings);
}

long ring::get_subring_index (string label) {
  Vec<ring> subrings;
  long depth = get_subrings(subrings);

    for (long i = 0; i < depth; ++i) {
      if (subrings[i].label == label) return i;
    }

    return -1;
}

vec_ZZ ring::index_convert (long j) {
  vec_ZZ tensor_index;
  tensor_index.SetLength(depth);

  long jj = j;

  for (long i = 0; i < depth; ++i) {
    long a = jj % subrings[i].gR;
    tensor_index[i] = a;
    jj = (jj - a) / subrings[i].gR;
  }

  return tensor_index;
}

// y = (x, x, ..., x) where x is assumed to be an element of the right subring
void ring::embed_1 (vec_ZZ& y, vec_ZZ& x) {
  y.SetLength(gR);
  for (long i = 0; i < left->gR; ++i) {
    for (long j = 0; j < right->gR; ++j) {
      y[i*right->gR + j] = x[j];
    }
  }
}

vec_ZZ ring::embed_1 (vec_ZZ& x) {
  vec_ZZ y;
  embed_1(y, x);
  return y;
}

// y = (x, 0, 0, ..., 0) where x is assumed to be an element of the right subring
void ring::embed_0 (vec_ZZ& y, vec_ZZ& x) {
  y.SetLength(gR);
  for (long j = 0; j < right->gR; ++j) {
    y[j] = x[j];
  }
}

vec_ZZ ring::embed_0 (vec_ZZ& x) {
  vec_ZZ y;
  embed_0(y, x);
  return y;
}

// y = (x, 0, 0, ..., 0) where x is assumed to be an element of the prime subring with label subring_label
void ring::embed_0 (vec_ZZ& y, string subring_label, vec_ZZ& x) {
  y.SetLength(gR);
  long d = get_subring_index(subring_label);

  for (long j = 0; j < gR; ++j) {
    vec_ZZ tensor_index = index_convert(j);
    long skip = 0;
    for (long k = 0; k < tensor_index.length(); ++k) {
      if (k != d && tensor_index[k] != 0) {
        skip = 1;
        break;
      }
    }
    if (!skip) y[j] = x[conv<long>(tensor_index[d])];
  }
}

vec_ZZ ring::embed_0 (string subring_label, vec_ZZ& x) {
  vec_ZZ y;
  embed_0(y, subring_label, x);
  return y;
}

// y = (x, x, ..., x) where x is assumed to be an element of the prime subring with label subring_label
void ring::embed_1 (vec_ZZ& y, string subring_label, vec_ZZ& x) {
  y.SetLength(gR);
  long d = get_subring_index(subring_label);

  for (long j = 0; j < gR; ++j) {
    vec_ZZ tensor_index = index_convert(j);
    y[j] = x[conv<long>(tensor_index[d])];
  }
}

vec_ZZ ring::embed_1 (string subring_label, vec_ZZ& x) {
  vec_ZZ y;
  embed_1(y, subring_label, x);
  return y;
}

// y = (-x, -x, ..., -x) where x is assumed to be an element of the prime subring with label subring_label
void ring::embed_11 (vec_ZZ& y, string subring_label, vec_ZZ& x) {
  y.SetLength(gR);
  long d = get_subring_index(subring_label);

  for (long j = 0; j < gR; ++j) {
    vec_ZZ tensor_index = index_convert(j);
    y[j] = -x[conv<long>(tensor_index[d])];
  }
}

vec_ZZ ring::embed_11 (string subring_label, vec_ZZ& x) {
  vec_ZZ y;
  embed_11(y, subring_label, x);
  return y;
}

// trace to the right ring
void ring::trace (vec_ZZ& y, vec_ZZ& x, ZZ& q) {
  y.SetLength(right->gR);
  for (long j = 0; j < x.length(); ++j) {
    long i = j % right->gR;
    y[i] = (y[i] + x[j]) % q;
  }
}

// trace to the subring of label 'subring_label'
// -- the subring must be prime
void ring::trace (vec_ZZ& y, string subring_label, vec_ZZ& x, ZZ& q) {
  //y.SetLength(gR);
  long d = get_subring_index(subring_label);
  y.SetLength(subrings[d].gR);

  for (long j = 0; j < x.length(); ++j) {
    vec_ZZ tensor_index = index_convert(j);
    long i = conv<long>(tensor_index[d]);
    y[i] = (y[i] + x[j]) % q;
  }
}

// returns the i-th diagonal vector of the extended Omega matrix,
// corresponding to the subring with label '_label'
void ring::deep_omega (vec_ZZ& extended_omega, string _label, long i) {
  Vec<ring> subrings;
  get_subrings(subrings);
  long d = get_subring_index(_label);

  vec_ZZ omg;
  diagonal(omg, i, subrings[d].omega);

  //embed_0(extended_omega, _label, omg); /*@@*/
  embed_1(extended_omega, _label, omg); /*@@*/
}

// returns the i-th diagonal vector of the extended Omega_inv matrix,
// corresponding to the subring with label '_label'
void ring::deep_omega_inv (vec_ZZ& extended_omega_inv, string _label, long i) {
  Vec<ring> subrings;
  get_subrings(subrings);
  long d = get_subring_index(_label);

  vec_ZZ omg_inv;
  diagonal(omg_inv, i, subrings[d].omega_inv);

  //embed_0(extended_omega_inv, _label, omg_inv); /*@@*/
  embed_1(extended_omega_inv, _label, omg_inv); /*@@*/
}