#include "ekf.h"

#include <ros/ros.h>
#include <ros/console.h>
#include <sensor_msgs/Imu.h>
#include <sensor_msgs/Range.h>
#include <nav_msgs/Odometry.h>
#include <Eigen/Eigen>
#include <Eigen/Geometry>
#include <Eigen/Dense>
#include <geometry_msgs/Pose.h>
#include <geometry_msgs/PoseStamped.h>
#include "conversion.h"

using namespace std;
using namespace Eigen;

//X_state: p q v gb ab   with time stamp aligned between imu and img
/*
    EKF model
    prediction:
    xt~ = xt-1 + dt*f(xt-1, ut, 0)
    sigmat~ = Ft*sigmat-1*Ft' + Vt*Qt*Vt'
    Update:
    Kt = sigmat~*Ct'*(Ct*sigmat~*Ct' + Wt*Rt*Wt')^-1
    xt = xt~ + Kt*(zt - g(xt~,0))
    sigmat = sigmat~ - Kt*Ct*sigmat~
*/
/*
   -pi ~ pi crossing problem:
   1. the model prpagation: X_state should be limited to [-pi,pi] after predicting and updating
   2. inovation crossing: (measurement - g(X_state)) should also be limited to [-pi,pi] when getting the inovation.
   z_measurement is normally in [-pi~pi]
*/

// imu frame is imu body frame

//odom: pose px,py pz orientation qw qx qy qz
//imu: acc: x y z gyro: wx wy wz

ros::Publisher odom_pub;
ros::Publisher cam_odom_pub;

//state
geometry_msgs::Pose pose;
Vector3d position, orientation, velocity;

// Now set up the relevant matrices
//states X [p q pdot]  [px,py,pz, wx,wy,wz, vx,vy,vz]
size_t stateSize;                                // x = [p q pdot bg ba]
size_t stateSize_pqv;                                // x = [p q pdot]
size_t measurementSize;                          // z = [p q]
size_t inputSize;                                // u = [w a]
VectorXd X_state(stateSize);                     // x (in most literature)
VectorXd u_input;
VectorXd Z_measurement;                          // z
MatrixXd StateCovariance;                        // sigma
MatrixXd Kt_kalmanGain;                          // Kt
VectorXd X_state_correct(stateSize);                     // x (in most literature)
MatrixXd StateCovariance_correct;                        // sigma
MatrixXd Qt;
MatrixXd Rt;
Vector3d u_gyro;
Vector3d u_acc;
Vector3d gravity(0., 0., -9.8); //need to estimate the bias 9.8099
Vector3d bg_0(0., 0., 0); //need to estimate the bias
Vector3d ba_0(0., 0., 0); //need to estimate the bias  0.1
Vector3d ng(0., 0., 0.);
Vector3d na(0., 0., 0.);
Vector3d nbg(0., 0., 0.);
Vector3d nba(0., 0., 0.);

Vector3d q_last;
Vector3d bg_last;
Vector3d ba_last;

//Qt imu covariance matrix  smaller believe system(imu) more
double gyro_cov = 0.01;
double acc_cov = 0.01;
//Rt visual odomtry covariance smaller believe measurement more
double position_cov = 0.1;
double q_rp_cov = 0.1;
double q_yaw_cov = 0.1;

double dt = 0.005; //second
double t_last, t_now;
bool first_frame_imu = true;
bool first_frame_tag_odom = true;
bool test_odomtag_call = false;
bool odomtag_call = false;

double time_now, time_last;
double time_odom_tag_now;
double diff_time;
//world frame points velocity
vector< pair<VectorXd, sensor_msgs::Imu> > sys_seq;  // keep a sequence of sys imu and X_state(before imu system prediction)
vector<MatrixXd> cov_seq;

//Rotation from the camera frame to the IMU frame
Matrix3d Rc_i;
Vector3d tc_i;  //  cam in imu frame
int cnt = 0;
Vector3d INNOVATION_;
Matrix3d Rr_i;
Vector3d tr_i;  //  rigid body in imu frame

void imu_callback(const sensor_msgs::Imu::ConstPtr &msg)
{
    // first_frame_tag_odom 初始为true，收到vicon后设置为false
    if(!first_frame_tag_odom)
    {
        // 第一次回调：发布初始值
        if(first_frame_imu)
        {
            first_frame_imu = false;
            time_now = msg->header.stamp.toSec();
            time_last = time_now;
            // 发布初始值
            system_pub(msg->header.stamp);
        }
        else
        {
            time_now = msg->header.stamp.toSec();
            // 时间差
            dt = time_now - time_last;

            if(odomtag_call)
            {
                odomtag_call = false;
                diff_time = time_now - time_odom_tag_now;
                if(diff_time<0)
                {
                    cout << "diff time: " << diff_time << endl;  //???!!! exist !!!???
                    cout << "timeimu: " << time_now - 1.60889e9 << " time_odom: " << time_odom_tag_now - 1.60889e9 << endl;
                    // cout << "diff time: " << diff_time << endl;  //about 30ms
                }
            }
            MatrixXd Ft;
            MatrixXd Vt;

            // 将imu存为输入值
            u_gyro(0) = msg->angular_velocity.x;
            u_gyro(1) = msg->angular_velocity.y;
            u_gyro(2) = msg->angular_velocity.z;
            u_acc(0)  = msg->linear_acceleration.x;
            u_acc(1)  = msg->linear_acceleration.y;
            u_acc(2)  = msg->linear_acceleration.z;

            // 上一时刻姿态
            q_last = X_state.segment<3>(3);  // last X2
            bg_last = X_state.segment<3>(9);  //last X4
            ba_last = X_state.segment<3>(12);  //last X5
            Ft = MatrixXd::Identity(stateSize, stateSize) + dt*diff_f_diff_x(q_last, u_gyro, u_acc, bg_last, ba_last);

            Vt = dt*diff_f_diff_n(q_last);

            // 使用输入更新状态
            X_state += dt*F_model(u_gyro, u_acc);
            // 欧拉角限制幅度
            if(X_state(3) > PI)  X_state(3) -= 2*PI;
            if(X_state(3) < -PI) X_state(3) += 2*PI;
            if(X_state(4) > PI)  X_state(4) -= 2*PI;
            if(X_state(4) < -PI) X_state(4) += 2*PI;
            if(X_state(5) > PI)  X_state(5) -= 2*PI;
            if(X_state(5) < -PI) X_state(5) += 2*PI;
            // 更新COV
            StateCovariance = Ft*StateCovariance*Ft.transpose() + Vt*Qt*Vt.transpose();

            time_last = time_now;

            system_pub(msg->header.stamp);
        }
    }

}

VectorXd get_pose_from_mocap(const geometry_msgs::PoseStamped::ConstPtr &msg)
{
    Matrix3d Rr_w;    //rigid body in world
    Vector3d tr_w;
    Matrix3d Ri_w;
    Vector3d ti_w;
    Vector3d p_temp;
    // 位置
    p_temp(0) = msg->pose.position.x;
    p_temp(1) = msg->pose.position.y;
    p_temp(2) = msg->pose.position.z;
    //quaternion2euler:  ZYX  roll pitch yaw
    Quaterniond q;
    // 姿态
    q.w() = msg->pose.orientation.w;
    q.x() = msg->pose.orientation.x;
    q.y() = msg->pose.orientation.y;
    q.z() = msg->pose.orientation.z;

    // 刚体的位置和姿态
    Rr_w = q.toRotationMatrix();
    tr_w = p_temp;
    // imu的位置和姿态
    Ri_w = Rr_w * Rr_i.inverse();
    ti_w = tr_w - Ri_w*tr_i;
    Vector3d euler = mat2euler(Ri_w);

    VectorXd pose = VectorXd::Random(6);
    // imu的位置和姿态
    pose.segment<3>(0) = ti_w;
    pose.segment<3>(3) = euler;

    return pose;
}

void mocap_callback(const geometry_msgs::PoseStamped::ConstPtr &msg)
{
    // 第一次回调
    if(first_frame_tag_odom)
    {
        first_frame_tag_odom = false;
        time_odom_tag_now = msg->header.stamp.toSec();

        VectorXd odom_pose = get_pose_from_mocap(msg);
        X_state.segment<3>(0) = odom_pose.segment<3>(0);
        X_state.segment<3>(3) = odom_pose.segment<3>(3);
    }
    else
    {
        time_odom_tag_now = msg->header.stamp.toSec();
        MatrixXd Ct;
        MatrixXd Wt;

        // 获取得到imu的位置姿态（imu和质心存在偏差）
        VectorXd odom_pose = get_pose_from_mocap(msg);
        // 更新测量值
        Z_measurement.segment<3>(0) = odom_pose.segment<3>(0);
        Z_measurement.segment<3>(3) = odom_pose.segment<3>(3);
        // 发布测量值（即mocap的原始值）
        cam_system_pub(msg->header.stamp);

        Ct = diff_g_diff_x();
        Wt = diff_g_diff_v();

        // 更新kalman增益
        Kt_kalmanGain = StateCovariance*Ct.transpose() * (Ct*StateCovariance*Ct.transpose() + Wt*Rt*Wt.transpose()).inverse();
        VectorXd gg = g_model();
        VectorXd innovation = Z_measurement - gg;

        //Prevent innovation changing suddenly when euler from -Pi to Pi
        if(innovation(3) > 6)  innovation(3) -= 2*PI;
        if(innovation(3) < -6) innovation(3) += 2*PI;
        if(innovation(4) > 6)  innovation(4) -= 2*PI;
        if(innovation(4) < -6) innovation(4) += 2*PI;
        if(innovation(5) > 6)  innovation(5) -= 2*PI;
        if(innovation(5) < -6) innovation(5) += 2*PI;
        INNOVATION_ = innovation.segment<3>(3);
        // 使用测量值更新状态
        X_state += Kt_kalmanGain*(innovation);
        if(X_state(3) > PI)  X_state(3) -= 2*PI;
        if(X_state(3) < -PI) X_state(3) += 2*PI;
        if(X_state(4) > PI)  X_state(4) -= 2*PI;
        if(X_state(4) < -PI) X_state(4) += 2*PI;
        if(X_state(5) > PI)  X_state(5) -= 2*PI;
        if(X_state(5) < -PI) X_state(5) += 2*PI;
        StateCovariance = StateCovariance - Kt_kalmanGain*Ct*StateCovariance;

        test_odomtag_call = true;
        odomtag_call = true;

        if(cnt == 10||cnt==50||cnt==90)
        {
            // cout << "Ct: \n" << Ct << "\nWt:\n" << Wt << endl;
            // cout << "Kt_kalmanGain: \n" << Kt_kalmanGain << endl;
            // cout << "\ninnovation: \n" << Kt_kalmanGain*innovation  << "\ndt:\n" << dt << endl;
            // cout << "\ninnovation: \n" << Kt_kalmanGain*innovation  << endl;
            // cout << "\ninnovation: \n" << INNOVATION_ << endl;
        }
        cnt++;
        if(cnt>100) cnt=101;
    }

}

int main(int argc, char **argv)
{
    ros::init(argc, argv, "mocap_ekf_node");
    ros::NodeHandle n("~");
    // 【订阅】 IMU数据，来自PX4
    ros::Subscriber s1 = n.subscribe("imu", 100, imu_callback, ros::TransportHints().tcpNoDelay());
    // 【订阅】 VICON数据
    ros::Subscriber s2 = n.subscribe("pose", 100, mocap_callback, ros::TransportHints().tcpNoDelay());
    // 【发布】 融合后的odom(发布频率为imu的频率)
    odom_pub = n.advertise<nav_msgs::Odometry>("ekf_odom", 100);
    // 【发布】 相机odom？
    cam_odom_pub = n.advertise<nav_msgs::Odometry>("cam_ekf_odom", 100);

    n.getParam("gyro_cov", gyro_cov);
    n.getParam("acc_cov", acc_cov);
    n.getParam("position_cov", position_cov);
    n.getParam("q_rp_cov", q_rp_cov);
    n.getParam("q_yaw_cov", q_yaw_cov);

    cout << "Q:" << gyro_cov << " " << acc_cov << " R: " << position_cov << " " << q_rp_cov << " " << q_yaw_cov << endl;

    // 初始化
    initsys();
    cout << "initsys" << endl;

    cout << "======================" << endl;
    double r = atan2(1,-100);
    double p = asin(-0.707);
    double y = atan2(-1, -100);
    cout << "r: " << r << " p: " << p << " y: " << y << endl;
    cout << "======================" << endl;

    ros::spin();
}

void system_pub(ros::Time stamp)
{
    nav_msgs::Odometry odom_fusion;
    odom_fusion.header.stamp = stamp;
    odom_fusion.header.frame_id = "world";
    // odom_fusion.header.frame_id = "imu";
    odom_fusion.pose.pose.position.x = X_state(0);
    odom_fusion.pose.pose.position.y = X_state(1);
    odom_fusion.pose.pose.position.z = X_state(2);
    Quaterniond q;
    q = euler2quaternion(X_state.segment<3>(3));
    odom_fusion.pose.pose.orientation.w = q.w();
    odom_fusion.pose.pose.orientation.x = q.x();
    odom_fusion.pose.pose.orientation.y = q.y();
    odom_fusion.pose.pose.orientation.z = q.z();
    odom_fusion.twist.twist.linear.x = X_state(6);
    odom_fusion.twist.twist.linear.y = X_state(7);
    odom_fusion.twist.twist.linear.z = X_state(8);
    odom_pub.publish(odom_fusion);
}

void cam_system_pub(ros::Time stamp)
{
    // 测量值
    nav_msgs::Odometry odom_fusion;
    odom_fusion.header.stamp = stamp;
    odom_fusion.header.frame_id = "world";

    odom_fusion.pose.pose.position.x = Z_measurement(0);
    odom_fusion.pose.pose.position.y = Z_measurement(1);
    odom_fusion.pose.pose.position.z = Z_measurement(2);
    Quaterniond q;
    q = euler2quaternion(Z_measurement.segment<3>(3));
    odom_fusion.pose.pose.orientation.w = q.w();
    odom_fusion.pose.pose.orientation.x = q.x();
    odom_fusion.pose.pose.orientation.y = q.y();
    odom_fusion.pose.pose.orientation.z = q.z();

    odom_fusion.twist.twist.angular.x = diff_time;
    odom_fusion.twist.twist.angular.y = dt;
    cam_odom_pub.publish(odom_fusion);
}

//process model
void initsys()
{
    //  camera position in the IMU frame = (0.05, 0.05, 0)
    // camera orientaion in the IMU frame = Quaternion(0, 1, 0, 0); w x y z, respectively
    //					   RotationMatrix << 1, 0, 0,
    //							             0, -1, 0,
    //                                       0, 0, -1;
    //set the cam2imu params
    Rc_i = Quaterniond(0, 1, 0, 0).toRotationMatrix();
    // cout << "R_cam" << endl << Rc_i << endl;
    tc_i << 0.05, 0.05, 0;

    Rr_i = Quaterniond(1, 0, 0, 0).toRotationMatrix();
    tr_i << 0, 0, 0.055;
    cout << "Rr_i: " << endl << Rr_i << endl;
    cout << "tr_i: " << endl << tr_i << endl;
    //  rigid body position in the IMU frame = (0, 0, 0.04)
    // rigid body orientaion in the IMU frame = Quaternion(1, 0, 0, 0); w x y z, respectively
    //					   RotationMatrix << 1, 0, 0,
    //							             0, 1, 0,
    //                                       0, 0, 1;

    //states X
    // 状态维度 [p q pdot bg ba]  [px,py,pz, wx,wy,wz（欧拉角）, vx,vy,vz bgx,bgy,bgz bax,bay,baz]
    stateSize = 15;                                                                 // x = [p q pdot bg ba]
    stateSize_pqv = 9;                                                              // x = [p q pdot]
    // 测量维度 位置+姿态
    measurementSize = 6;                                                            // z = [p q]
    // 输入维度 w是角速度 a是加速度
    inputSize = 6;                                                                  // u = [w a]
    X_state = VectorXd::Zero(stateSize);                                            // x
    //velocity
    X_state(6) = 0;
    X_state(7) = 0;
    X_state(8) = 0;
    // bias
    X_state.segment<3>(9) = bg_0;
    X_state.segment<3>(12) = ba_0;
    // 输入
    u_input = VectorXd::Zero(inputSize);
    // 测量
    Z_measurement = VectorXd::Zero(measurementSize);                                // z
    // 状态cov
    StateCovariance = MatrixXd::Identity(stateSize, stateSize);                     // sigma
    // kalman增益
    Kt_kalmanGain = MatrixXd::Identity(stateSize, measurementSize);                 // Kt
    // Ct_stateToMeasurement = MatrixXd::Identity(stateSize, measurementSize);         // Ct
    // ？
    X_state_correct = X_state;
    // ？
    StateCovariance_correct = StateCovariance;

    Qt = MatrixXd::Identity(inputSize, inputSize);  //6x6 input [gyro acc]covariance
    Rt = MatrixXd::Identity(measurementSize, measurementSize); //6x6 measurement [p q]covariance

    // You should also tune these parameters
    // Q imu covariance matrix; Rt visual odomtry covariance matrix
    // //Rt visual odomtry covariance smaller believe measurement more
    Qt.topLeftCorner(3, 3) = gyro_cov * Qt.topLeftCorner(3, 3);
    Qt.bottomRightCorner(3, 3) = acc_cov * Qt.bottomRightCorner(3, 3);
    Rt.topLeftCorner(3, 3) = position_cov * Rt.topLeftCorner(3, 3);
    Rt.bottomRightCorner(3, 3) = q_rp_cov * Rt.bottomRightCorner(3, 3);
    Rt.bottomRightCorner(1, 1) = q_yaw_cov * Rt.bottomRightCorner(1, 1);
}

void getState(Vector3d& p, Vector3d& q, Vector3d& v, Vector3d& bg, Vector3d& ba)
{
	p = X_state.segment<3>(0);
	q = X_state.segment<3>(3);
	v = X_state.segment<3>(6);
    bg = X_state.segment<3>(9);
    ba = X_state.segment<3>(12);
}

VectorXd F_model(Vector3d gyro, Vector3d acc)
{
    // IMU is in FLU frame
    // Transform IMU frame into "world" frame whose original point is FLU's original point and the XOY plain is parallel with the ground and z axis is up
    VectorXd f(VectorXd::Zero(stateSize));
    Vector3d p, q, v, bg, ba;
    getState(p, q, v, bg, ba);
    f.segment<3>(0) = v;
    f.segment<3>(3) = w_Body2Euler(q)*(gyro-bg-ng);
    f.segment<3>(6) = gravity + euler2mat(q)*(acc-ba-na);
    f.segment<3>(9) = nbg;
    f.segment<3>(12) = nba;
    return f;
}

VectorXd g_model()
{
    VectorXd g(VectorXd::Zero(measurementSize));

    g.segment<6>(0) = X_state.segment<6>(0);

    return g;
}

//F_model G_model Jocobian
//diff_f()/diff_x (x_t-1  ut  noise=0)   At     Ft = I+dt*At
MatrixXd diff_f_diff_x(Vector3d q_last, Vector3d gyro, Vector3d acc, Vector3d bg_last, Vector3d ba_last)
{
    double cr = cos( q_last(0));
    double sr = sin( q_last(0));
    double cp = cos( q_last(1));
    double sp = sin( q_last(1));
    double cy = cos( q_last(2));
    double sy = sin( q_last(2));

    // ng na = 0 nbg nba = 0
    double Ax = acc(0) - ba_last(0);
    double Ay = acc(1) - ba_last(1);
    double Az = acc(2) - ba_last(2);
    // double Wx = gyro(0) - bg_last(0);
    double Wy = gyro(1) - bg_last(1);
    double Wz = gyro(2) - bg_last(2);

    MatrixXd diff_f_diff_x_jacobian(MatrixXd::Zero(stateSize, stateSize));
    MatrixXd diff_f_diff_x_jacobian_pqv(MatrixXd::Zero(stateSize_pqv, stateSize_pqv));

    diff_f_diff_x_jacobian_pqv <<  0, 0, 0, 0, 0, 0, 1, 0, 0,
      0, 0, 0, 0, 0, 0, 0, 1, 0,
      0, 0, 0, 0, 0, 0, 0, 0, 1,
      0, 0, 0, (sp*(Wy*cr - Wz*sr))/cp, (Wz*cr + Wy*sr)/(cp*cp), 0, 0, 0, 0,
      0, 0, 0, (- Wz*cr - Wy*sr), 0, 0, 0, 0, 0,
      0, 0, 0, (Wy*cr - Wz*sr)/cp, (sp*(Wz*cr + Wy*sr))/(cp*cp), 0, 0, 0, 0,
      0, 0, 0, (Ay*(sr*sy + cr*cy*sp) + Az*(cr*sy - cy*sp*sr)), (Az*cp*cr*cy - Ax*cy*sp + Ay*cp*cy*sr), (Az*(cy*sr - cr*sp*sy) - Ay*(cr*cy + sp*sr*sy) - Ax*cp*sy), 0, 0, 0,
      0, 0, 0, (- Ay*(cy*sr - cr*sp*sy) - Az*(cr*cy + sp*sr*sy)), (Az*cp*cr*sy - Ax*sp*sy + Ay*cp*sr*sy), (Az*(sr*sy + cr*cy*sp) - Ay*(cr*sy - cy*sp*sr) + Ax*cp*cy), 0, 0, 0,
      0, 0, 0, (Ay*cp*cr - Az*cp*sr), (- Ax*cp - Az*cr*sp - Ay*sp*sr), 0, 0, 0, 0 ;

    diff_f_diff_x_jacobian.block<9, 9>(0, 0) = diff_f_diff_x_jacobian_pqv;
    diff_f_diff_x_jacobian.block<3, 3>(3, 9) = -w_Body2Euler(q_last);
    diff_f_diff_x_jacobian.block<3, 3>(6, 12) = -euler2mat(q_last);

    return diff_f_diff_x_jacobian;

    // cp != 0 pitch != 90° !!!!!!!!!
}
//diff_f()/diff_n (x_t-1  ut  noise=0)  Ut    Vt = dt*Ut
MatrixXd diff_f_diff_n(Vector3d q_last)
{
    MatrixXd diff_f_diff_n_jacobian(MatrixXd::Zero(stateSize, inputSize));
    diff_f_diff_n_jacobian.block<3,3>(3,0) = -w_Body2Euler(q_last);
    diff_f_diff_n_jacobian.block<3,3>(6,3) = -euler2mat(q_last);

    return diff_f_diff_n_jacobian;
}
//diff_g()/diff_x  (xt~ noise=0)  Ct
MatrixXd diff_g_diff_x()
{
    MatrixXd diff_g_diff_x_jacobian(MatrixXd::Zero(measurementSize, stateSize));
    diff_g_diff_x_jacobian.block<3,3>(0,0) = MatrixXd::Identity(3,3);
    diff_g_diff_x_jacobian.block<3,3>(3,3) = MatrixXd::Identity(3,3);

    return diff_g_diff_x_jacobian;
}
//diff_g()/diff_v  (xt~ noise=0) Wt
MatrixXd diff_g_diff_v()
{
    MatrixXd diff_g_diff_v_jacobian(MatrixXd::Identity(measurementSize, measurementSize));

    return diff_g_diff_v_jacobian;
}


//  saturating.cpp  ----------------------------------------------------------//

//  Copyright 2008 Howard Hinnant
//  Copyright 2008 Beman Dawes
//  Copyright 2009 Vicente J. Botet Escriba

//  Distributed under the Boost Software License, Version 1.0.
//  See http://www.boost.org/LICENSE_1_0.txt

/*
This code was extracted by Vicente J. Botet Escriba from Beman Dawes time2_demo.cpp which
was derived by Beman Dawes from Howard Hinnant's time2_demo prototype.
Many thanks to Howard for making his code available under the Boost license.
The original code was modified to conform to Boost conventions and to section
20.9 Time utilities [time] of the C++ committee's working paper N2798.
See http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2008/n2798.pdf.

time2_demo contained this comment:

    Much thanks to Andrei Alexandrescu,
                   Walter Brown,
                   Peter Dimov,
                   Jeff Garland,
                   Terry Golubiewski,
                   Daniel Krugler,
                   Anthony Williams.
*/

#define _CRT_SECURE_NO_WARNINGS  // disable VC++ foolishness

#include <boost/chrono/chrono.hpp>
#include <boost/type_traits.hpp>

#include <iostream>

//////////////////////////////////////////////////////////
//////////////////// User2 Example ///////////////////////
//////////////////////////////////////////////////////////

// Demonstrate User2:
// A "saturating" signed integral type  is developed.  This type has +/- infinity and a nan
// (like IEEE floating point) but otherwise obeys signed integral arithmetic.
// This class is subsequently used as the rep in boost::chrono::duration to demonstrate a
// duration class that does not silently ignore overflow.
#include <ostream>
#include <stdexcept>
#include <climits>

namespace User2
{

template <class I>
class saturate
{
public:
    typedef I int_type;

    static const int_type nan = int_type(int_type(1) << (sizeof(int_type) * CHAR_BIT - 1));
    static const int_type neg_inf = nan + 1;
    static const int_type pos_inf = -neg_inf;
private:
    int_type i_;

//     static_assert(std::is_integral<int_type>::value && std::is_signed<int_type>::value,
//                   "saturate only accepts signed integral types");
//     static_assert(nan == -nan && neg_inf < pos_inf,
//                   "saturate assumes two's complement hardware for signed integrals");

public:
    saturate() : i_(nan) {}
    explicit saturate(int_type i) : i_(i) {}
    // explicit
    operator int_type() const;

    saturate& operator+=(saturate x);
    saturate& operator-=(saturate x) {return *this += -x;}
    saturate& operator*=(saturate x);
    saturate& operator/=(saturate x);
    saturate& operator%=(saturate x);

    saturate  operator- () const {return saturate(-i_);}
    saturate& operator++()       {*this += saturate(int_type(1)); return *this;}
    saturate  operator++(int)    {saturate tmp(*this); ++(*this); return tmp;}
    saturate& operator--()       {*this -= saturate(int_type(1)); return *this;}
    saturate  operator--(int)    {saturate tmp(*this); --(*this); return tmp;}

    friend saturate operator+(saturate x, saturate y) {return x += y;}
    friend saturate operator-(saturate x, saturate y) {return x -= y;}
    friend saturate operator*(saturate x, saturate y) {return x *= y;}
    friend saturate operator/(saturate x, saturate y) {return x /= y;}
    friend saturate operator%(saturate x, saturate y) {return x %= y;}

    friend bool operator==(saturate x, saturate y)
    {
        if (x.i_ == nan || y.i_ == nan)
            return false;
        return x.i_ == y.i_;
    }

    friend bool operator!=(saturate x, saturate y) {return !(x == y);}

    friend bool operator<(saturate x, saturate y)
    {
        if (x.i_ == nan || y.i_ == nan)
            return false;
        return x.i_ < y.i_;
    }

    friend bool operator<=(saturate x, saturate y)
    {
        if (x.i_ == nan || y.i_ == nan)
            return false;
        return x.i_ <= y.i_;
    }

    friend bool operator>(saturate x, saturate y)
    {
        if (x.i_ == nan || y.i_ == nan)
            return false;
        return x.i_ > y.i_;
    }

    friend bool operator>=(saturate x, saturate y)
    {
        if (x.i_ == nan || y.i_ == nan)
            return false;
        return x.i_ >= y.i_;
    }

    friend std::ostream& operator<<(std::ostream& os, saturate s)
    {
        switch (s.i_)
        {
        case pos_inf:
            return os << "inf";
        case nan:
            return os << "nan";
        case neg_inf:
            return os << "-inf";
        };
        return os << s.i_;
    }
};

template <class I>
saturate<I>::operator I() const
{
    switch (i_)
    {
    case nan:
    case neg_inf:
    case pos_inf:
        throw std::out_of_range("saturate special value can not convert to int_type");
    }
    return i_;
}

template <class I>
saturate<I>&
saturate<I>::operator+=(saturate x)
{
    switch (i_)
    {
    case pos_inf:
        switch (x.i_)
        {
        case neg_inf:
        case nan:
            i_ = nan;
        }
        return *this;
    case nan:
        return *this;
    case neg_inf:
        switch (x.i_)
        {
        case pos_inf:
        case nan:
            i_ = nan;
        }
        return *this;
    }
    switch (x.i_)
    {
    case pos_inf:
    case neg_inf:
    case nan:
        i_ = x.i_;
        return *this;
    }
    if (x.i_ >= 0)
    {
        if (i_ < pos_inf - x.i_)
            i_ += x.i_;
        else
            i_ = pos_inf;
        return *this;
    }
    if (i_ > neg_inf - x.i_)
        i_ += x.i_;
    else
        i_ = neg_inf;
    return *this;
}

template <class I>
saturate<I>&
saturate<I>::operator*=(saturate x)
{
    switch (i_)
    {
    case 0:
        switch (x.i_)
        {
        case pos_inf:
        case neg_inf:
        case nan:
            i_ = nan;
        }
        return *this;
    case pos_inf:
        switch (x.i_)
        {
        case nan:
        case 0:
            i_ = nan;
            return *this;
        }
        if (x.i_ < 0)
            i_ = neg_inf;
        return *this;
    case nan:
        return *this;
    case neg_inf:
        switch (x.i_)
        {
        case nan:
        case 0:
            i_ = nan;
            return *this;
        }
        if (x.i_ < 0)
            i_ = pos_inf;
        return *this;
    }
    switch (x.i_)
    {
    case 0:
        i_ = 0;
        return *this;
    case nan:
        i_ = nan;
        return *this;
    case pos_inf:
        if (i_ < 0)
            i_ = neg_inf;
        else
            i_ = pos_inf;
        return *this;
    case neg_inf:
        if (i_ < 0)
            i_ = pos_inf;
        else
            i_ = neg_inf;
        return *this;
    }
    int s = (i_ < 0 ? -1 : 1) * (x.i_ < 0 ? -1 : 1);
    i_ = i_ < 0 ? -i_ : i_;
    int_type x_i_ = x.i_ < 0 ? -x.i_ : x.i_;
    if (i_ <= pos_inf / x_i_)
        i_ *= x_i_;
    else
        i_ = pos_inf;
    i_ *= s;
    return *this;
}

template <class I>
saturate<I>&
saturate<I>::operator/=(saturate x)
{
    switch (x.i_)
    {
    case pos_inf:
    case neg_inf:
        switch (i_)
        {
        case pos_inf:
        case neg_inf:
        case nan:
            i_ = nan;
            break;
        default:
            i_ = 0;
            break;
        }
        return *this;
    case nan:
        i_ = nan;
        return *this;
    case 0:
        switch (i_)
        {
        case pos_inf:
        case neg_inf:
        case nan:
            return *this;
        case 0:
            i_ = nan;
            return *this;
        }
        if (i_ > 0)
            i_ = pos_inf;
        else
            i_ = neg_inf;
        return *this;
    }
    switch (i_)
    {
    case 0:
    case nan:
        return *this;
    case pos_inf:
    case neg_inf:
        if (x.i_ < 0)
            i_ = -i_;
        return *this;
    }
    i_ /= x.i_;
    return *this;
}

template <class I>
saturate<I>&
saturate<I>::operator%=(saturate x)
{
//    *this -= *this / x * x;  // definition
    switch (x.i_)
    {
    case nan:
    case neg_inf:
    case 0:
    case pos_inf:
        i_ = nan;
        return *this;
    }
    switch (i_)
    {
    case neg_inf:
    case pos_inf:
        i_ = nan;
    case nan:
        return *this;
    }
    i_ %= x.i_;
    return *this;
}

// Demo overflow-safe integral durations ranging from picoseconds resolution to millennium resolution
typedef boost::chrono::duration<saturate<long long>, boost::pico                 > picoseconds;
typedef boost::chrono::duration<saturate<long long>, boost::nano                 > nanoseconds;
typedef boost::chrono::duration<saturate<long long>, boost::micro                > microseconds;
typedef boost::chrono::duration<saturate<long long>, boost::milli                > milliseconds;
typedef boost::chrono::duration<saturate<long long>                            > seconds;
typedef boost::chrono::duration<saturate<long long>, boost::ratio<         60LL> > minutes;
typedef boost::chrono::duration<saturate<long long>, boost::ratio<       3600LL> > hours;
typedef boost::chrono::duration<saturate<long long>, boost::ratio<      86400LL> > days;
typedef boost::chrono::duration<saturate<long long>, boost::ratio<   31556952LL> > years;
typedef boost::chrono::duration<saturate<long long>, boost::ratio<31556952000LL> > millennium;

}  // User2

// Demonstrate custom promotion rules (needed only if there are no implicit conversions)
namespace User2 { namespace detail {

template <class T1, class T2, bool = boost::is_integral<T1>::value>
struct promote_helper;

template <class T1, class T2>
struct promote_helper<T1, saturate<T2>, true>  // integral
{
    typedef typename boost::common_type<T1, T2>::type rep;
    typedef User2::saturate<rep> type;
};

template <class T1, class T2>
struct promote_helper<T1, saturate<T2>, false>  // floating
{
    typedef T1 type;
};

} }

namespace boost
{

template <class T1, class T2>
struct common_type<User2::saturate<T1>, User2::saturate<T2> >
{
    typedef typename common_type<T1, T2>::type rep;
    typedef User2::saturate<rep> type;
};

template <class T1, class T2>
struct common_type<T1, User2::saturate<T2> >
    : User2::detail::promote_helper<T1, User2::saturate<T2> > {};

template <class T1, class T2>
struct common_type<User2::saturate<T1>, T2>
    : User2::detail::promote_helper<T2, User2::saturate<T1> > {};

// Demonstrate specialization of duration_values:

namespace chrono {

template <class I>
struct duration_values<User2::saturate<I> >
{
    typedef User2::saturate<I> Rep;
public:
    static Rep zero() {return Rep(0);}
    static Rep max BOOST_PREVENT_MACRO_SUBSTITUTION ()  {return Rep(Rep::pos_inf-1);}
    static Rep min BOOST_PREVENT_MACRO_SUBSTITUTION ()  {return -(max)();}
};

}  // namespace chrono

}  // namespace boost

#include <iostream>

void testUser2()
{
    std::cout << "*************\n";
    std::cout << "* testUser2 *\n";
    std::cout << "*************\n";
    using namespace User2;
    typedef seconds::rep sat;
    years yr(sat(100));
    std::cout << "100 years expressed as years = " << yr.count() << '\n';
    nanoseconds ns = yr;
    std::cout << "100 years expressed as nanoseconds = " << ns.count() << '\n';
    ns += yr;
    std::cout << "200 years expressed as nanoseconds = " << ns.count() << '\n';
    ns += yr;
    std::cout << "300 years expressed as nanoseconds = " << ns.count() << '\n';
//    yr = ns;  // does not compile
    std::cout << "yr = ns;  // does not compile\n";
//    picoseconds ps1 = yr;  // does not compile, compile-time overflow in ratio arithmetic
    std::cout << "ps = yr;  // does not compile\n";
    ns = yr;
    picoseconds ps = ns;
    std::cout << "100 years expressed as picoseconds = " << ps.count() << '\n';
    ps = ns / sat(1000);
    std::cout << "0.1 years expressed as picoseconds = " << ps.count() << '\n';
    yr = years(sat(-200000000));
    std::cout << "200 million years ago encoded in years: " << yr.count() << '\n';
    days d = boost::chrono::duration_cast<days>(yr);
    std::cout << "200 million years ago encoded in days: " << d.count() << '\n';
    millennium c = boost::chrono::duration_cast<millennium>(yr);
    std::cout << "200 million years ago encoded in millennium: " << c.count() << '\n';
    std::cout << "Demonstrate \"uninitialized protection\" behavior:\n";
    seconds sec;
    for (++sec; sec < seconds(sat(10)); ++sec)
        ;
    std::cout << sec.count() << '\n';
    std::cout << "\n";
}

void testStdUser()
{
    std::cout << "***************\n";
    std::cout << "* testStdUser *\n";
    std::cout << "***************\n";
    using namespace boost::chrono;
    hours hr = hours(100);
    std::cout << "100 hours expressed as hours = " << hr.count() << '\n';
    nanoseconds ns = hr;
    std::cout << "100 hours expressed as nanoseconds = " << ns.count() << '\n';
    ns += hr;
    std::cout << "200 hours expressed as nanoseconds = " << ns.count() << '\n';
    ns += hr;
    std::cout << "300 hours expressed as nanoseconds = " << ns.count() << '\n';
//    hr = ns;  // does not compile
    std::cout << "hr = ns;  // does not compile\n";
//    hr * ns;  // does not compile
    std::cout << "hr * ns;  // does not compile\n";
    duration<double> fs(2.5);
    std::cout << "duration<double> has count() = " << fs.count() << '\n';
//    seconds sec = fs;  // does not compile
    std::cout << "seconds sec = duration<double> won't compile\n";
    seconds sec = duration_cast<seconds>(fs);
    std::cout << "seconds has count() = " << sec.count() << '\n';
    std::cout << "\n";
}

int main()
{
    testStdUser();
    testUser2();
    return 0;
}


#include <Eigen/Dense>
#include "control.hpp"
#include "defines.hpp"
#include "fad.hpp"
#include "global_residual.hpp"
#include "J2_plane_strain.hpp"
#include "material_params.hpp"

namespace calibr8 {

static ParameterList get_valid_local_residual_params() {
  ParameterList p;
  p.set<std::string>("type", "J2_plane_strain");
  p.set<int>("nonlinear max iters", 0);
  p.set<double>("nonlinear absolute tol", 0.);
  p.set<double>("nonlinear relative tol", 0.);
  p.sublist("materials");
  return p;
}

static ParameterList get_valid_material_params() {
  ParameterList p;
  p.set<double>("E", 0.);
  p.set<double>("nu", 0.);
  p.set<double>("K", 0.);
  p.set<double>("Y", 0.);
  return p;
}

template <typename T>
J2_plane_strain<T>::J2_plane_strain(ParameterList const& inputs, int ndims) {

  this->m_params_list = inputs;
  this->m_params_list.validateParameters(get_valid_local_residual_params(), 0);

  int const num_residuals = 4;

  this->m_num_residuals = num_residuals;
  this->m_num_eqs.resize(num_residuals);
  this->m_var_types.resize(num_residuals);
  this->m_resid_names.resize(num_residuals);

  this->m_resid_names[0] = "zeta";
  this->m_var_types[0] = SYM_TENSOR;
  this->m_num_eqs[0] = get_num_eqs(SYM_TENSOR, ndims);

  this->m_resid_names[1] = "Ie";
  this->m_var_types[1] = SCALAR;
  this->m_num_eqs[1] = get_num_eqs(SCALAR, ndims);

  this->m_resid_names[2] = "alpha";
  this->m_var_types[2] = SCALAR;
  this->m_num_eqs[2] = get_num_eqs(SCALAR, ndims);

  this->m_resid_names[3] = "zeta_zz";
  this->m_var_types[3] = SCALAR;
  this->m_num_eqs[3] = get_num_eqs(SCALAR, ndims);

  m_max_iters = inputs.get<int>("nonlinear max iters");
  m_abs_tol = inputs.get<double>("nonlinear absolute tol");
  m_rel_tol = inputs.get<double>("nonlinear relative tol");

}

template <typename T>
void J2_plane_strain<T>::init_params() {

  int const num_params = 4;
  this->m_params.resize(num_params);
  this->m_param_names.resize(num_params);

  this->m_param_names.resize(num_params);
  this->m_param_names[0] = "E";
  this->m_param_names[1] = "nu";
  this->m_param_names[2] = "K";
  this->m_param_names[3] = "Y";

  int const num_elem_sets = this->m_elem_set_names.size();
  resize(this->m_param_values, num_elem_sets, num_params);

  ParameterList& all_material_params =
      this->m_params_list.sublist("materials", true);

  for (int es = 0; es < num_elem_sets; ++es) {
    std::string const& elem_set_name = this->m_elem_set_names[es];
    ParameterList& material_params =
        all_material_params.sublist(elem_set_name, true);
    material_params.validateParameters(get_valid_material_params(), 0);
    this->m_param_values[es][0] = material_params.get<double>("E");
    this->m_param_values[es][1] = material_params.get<double>("nu");
    this->m_param_values[es][2] = material_params.get<double>("K");
    this->m_param_values[es][3] = material_params.get<double>("Y");
  }

  this->m_active_indices.resize(1);
  this->m_active_indices[0].resize(1);
  this->m_active_indices[0][0] = 0;
}

template <typename T>
J2_plane_strain<T>::~J2_plane_strain() {
}

template <typename T>
void J2_plane_strain<T>::init_variables_impl() {

  int const ndims = this->m_num_dims;
  int const zeta_idx = 0;
  int const Ie_idx = 1;
  int const alpha_idx = 2;
  int const zeta_zz_idx = 3;

  T const Ie = 1.0;
  T const alpha = 0.0;
  Tensor<T> const zeta = minitensor::zero<T>(ndims);
  T const zeta_zz = 0.0;

  this->set_scalar_xi(Ie_idx, Ie);
  this->set_scalar_xi(alpha_idx, alpha);
  this->set_sym_tensor_xi(zeta_idx, zeta);
  this->set_scalar_xi(zeta_zz_idx, zeta_zz);

}

template <typename T>
void eval_be_bar_plane_strain(
    RCP<GlobalResidual<T>> global,
    Tensor<T> const& zeta,
    T const& Ie,
    T const& zeta_zz,
    Tensor<T>& be_bar_2D,
    T& be_bar_zz) {
  int const ndims = global->num_dims();
  Tensor<T> const I = minitensor::eye<T>(ndims);
  Tensor<T> const grad_u = global->grad_vector_x(0);
  Tensor<T> const grad_u_prev = global->grad_vector_x_prev(0);
  Tensor<T> const F = grad_u + I;
  Tensor<T> const F_prev = grad_u_prev + I;
  Tensor<T> const rF = F * minitensor::inverse(F_prev);
  T const det_rF = minitensor::det(rF);
  T const det_rF_13 = cbrt(det_rF);
  Tensor<T> const rF_bar = rF / det_rF_13;
  Tensor<T> const rF_barT = minitensor::transpose(rF_bar);
  be_bar_2D = rF_bar * (zeta + Ie * I) * rF_barT;
  be_bar_zz = (zeta_zz + Ie) / (det_rF_13 * det_rF_13);
}

template <typename T>
T norm_s_3D(Tensor<T> const& s_2D, T const& s_zz) {
  return sqrt(s_2D(0, 0) * s_2D(0, 0) + s_2D(1, 1) * s_2D(1, 1)
      + 2. * s_2D(0, 1) * s_2D(0, 1) + s_zz * s_zz);
}

template <typename T>
T det_be_bar_3D(Tensor<T> const& zeta_2D, T const& zeta_zz, T const& Ie) {
  return ((zeta_2D(0, 0) + Ie) * (zeta_2D(1, 1) + Ie)
      - zeta_2D(0, 1) * zeta_2D(0, 1)) * (zeta_zz + Ie);
}

template <>
int J2_plane_strain<double>::solve_nonlinear(RCP<GlobalResidual<double>>) {
  return 0;
}

template <>
int J2_plane_strain<FADT>::solve_nonlinear(RCP<GlobalResidual<FADT>> global) {

  int path;

  // pick an initial guess for the local variables
  {
    Tensor<FADT> const zeta_old = this->sym_tensor_xi_prev(0);
    FADT const Ie_old = this->scalar_xi_prev(1);
    FADT const alpha_old = this->scalar_xi_prev(2);
    FADT const zeta_zz_old = this->scalar_xi_prev(3);

    int const ndims = global->num_dims();
    Tensor<FADT> be_bar_2D_trial;
    FADT be_bar_zz_trial;
    eval_be_bar_plane_strain(global, zeta_old, Ie_old, zeta_zz_old,
        be_bar_2D_trial, be_bar_zz_trial);
    FADT const Ie_trial = (minitensor::trace(be_bar_2D_trial)
        + be_bar_zz_trial) / 3.;
    Tensor<FADT> const I = minitensor::eye<FADT>(ndims);
    Tensor<FADT> const zeta_trial = be_bar_2D_trial - Ie_trial * I;
    FADT const zeta_zz_trial = be_bar_zz_trial - Ie_trial;
    FADT const alpha_trial = alpha_old;
    this->set_sym_tensor_xi(0, zeta_trial);
    this->set_scalar_xi(1, Ie_trial);
    this->set_scalar_xi(2, alpha_trial);
    this->set_scalar_xi(3, zeta_zz_trial);
    path = ELASTIC;
  }

  // newton iteration until convergence

  int iter = 1;
  double R_norm_0 = 1.;
  bool converged = false;

  while ((iter <= m_max_iters) && (!converged)) {

    path = this->evaluate(global);

    double const R_norm = this->norm_residual();
    if (iter == 1) R_norm_0 = R_norm;
    double const R_norm_rel = R_norm / R_norm_0;
    if ((R_norm_rel < m_rel_tol) || (R_norm < m_abs_tol)) {
      converged = true;
      break;
    }

    EMatrix const J = this->eigen_jacobian();
    EVector const R = this->eigen_residual();
    EVector const dxi = J.fullPivLu().solve(-R);

    this->add_to_sym_tensor_xi(0, dxi);
    this->add_to_scalar_xi(1, dxi);
    this->add_to_scalar_xi(2, dxi);
    this->add_to_scalar_xi(3, dxi);

    iter++;

  }

  // fail if convergence was not achieved
  if ((iter > m_max_iters) && (!converged)) {
    fail("J2_plane_strain:solve_nonlinear failed in %d iterations", m_max_iters);
  }

  return path;

}

template <typename T>
int J2_plane_strain<T>::evaluate(
    RCP<GlobalResidual<T>> global,
    bool force_path,
    int path_in) {

  int path = ELASTIC;
  int const ndims = this->m_num_dims;
  double const sqrt_23 = std::sqrt(2./3.);
  double const sqrt_32 = std::sqrt(3./2.);

  T const E = this->m_params[0];
  T const nu = this->m_params[1];
  T const K = this->m_params[2];
  T const Y = this->m_params[3];
  T const mu = compute_mu(E, nu);

  Tensor<T> const zeta_old = this->sym_tensor_xi_prev(0);
  T const Ie_old = this->scalar_xi_prev(1);
  T const alpha_old = this->scalar_xi_prev(2);
  T const zeta_zz_old = this->scalar_xi_prev(3);

  Tensor<T> const zeta = this->sym_tensor_xi(0);
  T const Ie = this->scalar_xi(1);
  T const alpha = this->scalar_xi(2);
  T const zeta_zz = this->scalar_xi(3);

  Tensor<T> const I = minitensor::eye<T>(ndims);
  Tensor<T> be_bar_trial_2D;
  T be_bar_trial_zz;
  eval_be_bar_plane_strain(global, zeta_old, Ie_old, zeta_zz_old,
      be_bar_trial_2D, be_bar_trial_zz);
  T const Ie_trial = (minitensor::trace(be_bar_trial_2D)
      + be_bar_trial_zz) / 3.;
  Tensor<T> const zeta_trial = be_bar_trial_2D - Ie_trial * I;
  T const zeta_zz_trial = be_bar_trial_zz - Ie_trial;
  Tensor<T> const s_2D = mu * zeta;
  T const s_zz = mu * zeta_zz;
  T const s_mag = norm_s_3D(s_2D, s_zz);
  Tensor<T> const n_2D = s_2D / s_mag;
  T const n_zz = s_zz / s_mag;
  T const sigma_yield = Y + K * alpha;
  T const f = s_mag - sqrt_23 * sigma_yield;

  Tensor<T> R_zeta;
  T R_Ie;
  T R_alpha;
  T R_zeta_zz;

  if (!force_path) {
    // plastic step
    if (f > m_abs_tol || std::abs(f) < m_abs_tol) {
      T const dgam = sqrt_32 * (alpha - alpha_old);
      R_zeta = zeta - zeta_trial + 2. * dgam * Ie * n_2D;
      R_Ie = det_be_bar_3D(zeta, zeta_zz, Ie) - 1.;
      R_alpha = (s_mag - sqrt_23 * sigma_yield) / val(mu);
      R_zeta_zz = zeta_zz - zeta_zz_trial + 2. * dgam * Ie * n_zz;
      path = PLASTIC;
    }
    // elastic step
    else {
      R_zeta = (0. * mu + 1.) * zeta - zeta_trial;
      R_Ie = Ie - Ie_trial + 0. * mu;
      R_alpha = alpha - alpha_old + 0. * mu;
      R_zeta_zz = zeta_zz - zeta_zz_trial + 0. * mu;
      path = ELASTIC;
    }
  }

  // force the path
  else {
    path = path_in;
    // plastic step
    if (path == PLASTIC) {
      T const dgam = sqrt_32 * (alpha - alpha_old);
      R_zeta = zeta - zeta_trial + 2. * dgam * Ie * n_2D;
      R_Ie = det_be_bar_3D(zeta, zeta_zz, Ie) - 1.;
      R_alpha = (s_mag - sqrt_23 * sigma_yield) / val(mu);
      R_zeta_zz = zeta_zz - zeta_zz_trial + 2. * dgam * Ie * n_zz;
    }
    // elastic step
    else {
      R_zeta = (0. * mu + 1.) * zeta - zeta_trial;
      R_Ie = Ie - Ie_trial + 0. * mu;
      R_alpha = alpha - alpha_old + 0. * mu;
      R_zeta_zz = zeta_zz - zeta_zz_trial + 0. * mu;
    }
  }

  this->set_sym_tensor_R(0, R_zeta);
  this->set_scalar_R(1, R_Ie);
  this->set_scalar_R(2, R_alpha);
  this->set_scalar_R(3, R_zeta_zz);

  return path;

}

template <typename T>
Tensor<T> J2_plane_strain<T>::dev_cauchy(RCP<GlobalResidual<T>> global) {
  int const ndims = global->num_dims();
  T const E = this->m_params[0];
  T const nu = this->m_params[1];
  T const mu = E / (2. * (1. + nu));
  Tensor<T> const I = minitensor::eye<T>(ndims);
  Tensor<T> const grad_u = global->grad_vector_x(0);
  Tensor<T> const F = grad_u + I;
  Tensor<T> const zeta = this->sym_tensor_xi(0);
  T const J = minitensor::det(F);
  return mu * zeta / J;
}

template <typename T>
Tensor<T> J2_plane_strain<T>::cauchy(RCP<GlobalResidual<T>> global, T p) {
  int const ndims = global->num_dims();
  Tensor<T> const I = minitensor::eye<T>(ndims);
  Tensor<T> const dev_sigma = this->dev_cauchy(global);
  Tensor<T> const sigma = dev_sigma - p * I;
  return sigma;
}

template class J2_plane_strain<double>;
template class J2_plane_strain<FADT>;

}


// BLAS level 1 (vector) -- complex numbers

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

//#define BOOST_NO_FUNCTION_TEMPLATE_ORDERING

#include <boost/numeric/bindings/ublas/vector.hpp>
#include <boost/numeric/bindings/blas/level1.hpp>
#include "utils.h"

namespace blas = boost::numeric::bindings::blas;
namespace ublas = boost::numeric::ublas;

using std::cout;
using std::endl;

typedef double real_t;
typedef std::complex<real_t> cmplx_t;
typedef ublas::vector<cmplx_t> vct_t;

int main() {

  int n = 6;

  cout << endl;
  vct_t v (n);
  init_v (v, times_plus<cmplx_t> (cmplx_t (1, -1), cmplx_t (0, .1)));
  print_v (v, "v");

  blas::scal (2.0, v);
  print_v (v, "2.0 v");

  blas::scal (cmplx_t (-1, 0), v);
  print_v (v, "(-1, 0) v");

  blas::scal (cmplx_t (0, 1), v);
  print_v (v, "(0, 1) v");

  blas::set (cmplx_t (1, -1), v);
  print_v (v, "v");
  vct_t v1 (n);
  blas::set (cmplx_t (0, -1), v1);
  print_v (v1, "v1");

  cout << endl;
  cout << "v^T v1 = " << blas::dot (v, v1) << " == "
    << blas::dotu (v, v1) << " == "
    << inner_prod (v, v1) << endl;
  cout << "v^T v = " << blas::dot (v, v) << " == "
    << blas::dotu (v, v) << " == "
    << inner_prod (v, v) << endl;
  cout << "v1^T v1 = " << blas::dot (v1, v1) << " == "
    << blas::dotu (v1, v1) << " == "
    << inner_prod (v1, v1) << endl;

  cout << endl;
  cout << "v^H v1 = " << blas::dotc (v, v1) << " == "
    << inner_prod (conj (v), v1) << " == "
    << inner_prod (v1, conj (v)) << " != "
    << inner_prod (v, conj (v1)) << endl;
  cout << "v^H v = " << blas::dotc (v, v) << " == "
    << inner_prod (conj (v), v) << " == "
    << inner_prod (v, conj (v)) << endl;
  cout << "v1^H v1 = " << blas::dotc (v1, v1) << " == "
    << inner_prod (conj (v1), v1) << " == "
    << inner_prod (v1, conj (v1)) << endl;

  cout << endl;
  cout << "||v||_1 = " << blas::asum (v) << endl;
  cout << "||v||_2 = " << blas::nrm2 (v) << " == "
    << norm_2 (v) << endl;

  cout << endl;
}


/* Copyright (C) 2012-2017 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */
#include <cassert>
#include <iostream>
#include <iomanip>
#include <numeric>
#include <cmath>
#include <atomic>
#include <NTL/BasicThreadPool.h>
#include <NTL/ZZ_pX.h>
#include <NTL/xdouble.h>
#include "NumbTh.h"
#include "PAlgebra.h"
#include "DoubleCRT.h"
#include "Context.h"
#include "sample.h"
#include "timing.h"
#include "norms.h"

NTL_CLIENT

void printHistogram(const vector<double>& data,
                    const cx_double& mean, double step)
{
  vector<long> hist(8,0); // histogram: hist[i]= # of x'es with |x|<i*stdev
  double max = 0.0;
  double sum=0.0, sumSqr=0.0;
  for (double x: data) {
    sum += x;
    sumSqr += x*x;
    long j = std::floor(x/step);
    if (j >= lsize(hist)) hist.resize(j+1, 0);
    hist[j]++;
    if (x > max) max = x;
  }
  sum /= data.size();    // E[x]
  sumSqr /= data.size(); // E[x^2]
  double stdev = sqrt(sumSqr - (sum*sum));

  for (long i=hist.size()-1; i>0; i--) // cumulative
    hist[i-1] += hist[i];

  vector<double> dhist(hist.size(), 1.0);
  vector<double> ratio(hist.size(), 0.0);
  for (long i=1; i<lsize(dhist); i++) {
    dhist[i] = hist[i]/double(hist[0]);
    if (dhist[i]>0.0) ratio[i] = dhist[i-1]/dhist[i];
  }
  cout << data.size() << " points, mean="<<mean<<endl;
  cout << "size mean="<<sum<<", stdev="<<stdev
       << ", max="<<max<<".  Step="<<step<<endl;
  cout << "  histogram ="<< hist<<endl
       << " probability="<<dhist<<endl
       << "      ratio ="<<ratio<<endl<<endl;
}

void freshCtxtNoise(zzX& f, const Context& context,
                    double sigma, bool modPhimX)
{
  zzX s, r, e1, e2, e3;
  const PAlgebra& palg = context.zMStar;
  sampleSmallBounded(s, context);
  if (modPhimX) {
    sampleSmall(r, palg.getPhiM()-1);
    sampleGaussian(e1, palg.getPhiM()-1, sigma);
    sampleGaussianBounded(e2, context, sigma);
    sampleGaussian(e3, palg.getPhiM()-1, sigma);
  } else {
    sampleSmall(r, context);
    sampleGaussian(e1, context, sigma);
    sampleGaussianBounded(e2, context, sigma);
    sampleGaussian(e3, context, sigma);
  }
  f = MulMod(s, e1, palg) + MulMod(r, e2, palg) + e3;
}

void roundingNoise(zzX& f, const Context& context,
                   long p2r, bool modPhimX)
{
  zzX s, e1, e2;
  const PAlgebra& palg = context.zMStar;
  sampleSmallBounded(s, context);
  if (modPhimX) {
    sampleUniform(e1, palg.getPhiM()-1, p2r);
    sampleUniform(e2, palg.getPhiM()-1, p2r);
  } else {
    sampleUniform(e1, context, p2r);
    sampleUniform(e2, context, p2r);
  }
  f = MulMod(s, e1, palg) + e2;
}

int main(int argc, char **argv)
{
  FHE_NTIMER_START(init);
  // get parameters from the command line
  ArgMapping amap;

  // long noPrint = 1;
  // amap.arg("noPrint", noPrint, "suppress printouts");

  long m = 15;
  amap.arg("m", m, "the cyclotomic index");
  double sigma = 0.0;
  amap.arg("sigma", sigma, "nomral standard deviation", "heristic");
  long p = 2;
  amap.arg("p", p, "plaintext base");
  long r = 1;
  amap.arg("r", r, "lifting");
  long N = 1000;
  amap.arg("N", N, "# of samples to use");
  long seed=0;
  amap.arg("seed", seed, "PRG seed");
  long nt=4;
  amap.arg("nt", nt, "number of threads");
  amap.parse(argc, argv);

  if (seed != 0) NTL::SetSeed(ZZ(seed));
  if (nt > 1)    NTL::SetNumThreads(nt);
  long p2r = power_long(p, r);
  if (sigma<=0.0) {
    if (m&1) // odd m
      sigma = 3.2*sqrt(m);
    else
      sigma=3.2;
  }

  Context context(m, p, r);
  const PAlgebra& palg = context.zMStar;
  buildModChain(context, /*L=*/5, /*c=*/3);
  long phim = palg.getPhiM();
  FHE_NTIMER_STOP(init);

  cout << "m="<<m<<", phi(m)="<<palg.getPhiM()
       << ", sigma="<<std::setprecision(3)<<sigma
       << ", p^r="<<p2r << endl;

  zzX f;
  vector<double> data;
  cx_double sum, mean;
  double step;

  step = (sigma+0.1)*0.54*(1+((m&1)? sqrt(phim*m): phim));
  // fresh ciphertext, sampling mod X^m-1
  data.resize(0);
  sum = 0;
  for (long i=0; i<N; i++) {
    std::vector<cx_double> cemb;
    freshCtxtNoise(f, context, sigma, false);
    canonicalEmbedding(cemb, f, palg); // Canonical embedding of f
    for (auto& entry : cemb) {
      double sz2 = conv<double>(std::norm<double>(entry));
      data.push_back(std::sqrt(sz2));
      sum += entry/double(lsize(cemb));
    }
  }
  mean = sum / double(N);
  cout << "fresh ctxt noise, sample mod X^m-1: ";
  printHistogram(data, mean, step);
  data.resize(0);
  sum = 0;
  // rounding error, mod Phi_m(X)

  step = (2*p2r+1)*(phim-2)/8.0;
  for (long i=0; i<N; i++) {
    std::vector<cx_double> cemb;
    roundingNoise(f, context, p2r, true);
    canonicalEmbedding(cemb, f, palg); // Canonical embedding of f
    for (auto& entry : cemb) {
      double sz2 = conv<double>(std::norm<double>(entry));
      data.push_back(std::sqrt(sz2));
      sum += entry/double(lsize(cemb));
    }
  }
  mean = sum / double(N);
  cout << "\nrounding noise, sample mod Phi_m(X): ";
  printHistogram(data, mean, step);
  data.resize(0);
  sum = 0;

  cout << endl;

  //  printAllTimers();
  return 0;
}

/********************************************************************/
#if 0 // OLD CODE
  vector<xdouble> l2(8, xdouble(0.0));
  vector<xdouble> ratio(8, xdouble(1.0));

  cout << "*** m="<<m<<", sampling "<<N<<" different f's ***\n";
  for (long i=0; i<N; i++) {
    ZZX f;
    sampleSmall(f, m);
    rem(f, f, phimX);
    if (IsZero(f)) continue;
    xdouble ll = embeddingL2NormSquared(f,m)/phim;
    l2[0] += ll;
    for (long j=1; j<lsize(l2); j++) {
      NTL::SqrMod(f, f, phimX); // fd -> fd^2
      ll *= ll;

      xdouble tt = embeddingL2NormSquared(f, m)/phim;
      l2[j] += tt;

      tt /= ll;
      if (tt>ratio[j]) ratio[j] = tt;
    }
  }
  xdouble base = l2[0] / N;
  long e = 1;
  xdouble factorial(1.0);
  for (long j=0; j<lsize(l2); j++) {
    l2[j] /= N;
    for (long i=e/2 +1; i<=e; i++) factorial *= i;
    cout << "E[|f^"<<e<<"|^2]="<<l2[j]
         << ",\t= "<<(l2[j]/base)<<"*E[|f^2|]^{"<<e<<"} (vs. "
         << e<<"! ="<<factorial<<")\n";
    cout << "\t\t\t max ratio = "<<ratio[j]<<endl;
    base *= base;
    e *= 2;
  }
#endif
#if 0
  step = (1+sqrt(phim*log(phim)))*0.85;
  for (long i=0; i<N; i++) {
    sampleSmall(f, palg);
    data.push_back(embeddingLargestCoeff(f, palg));
  }
  cout << "sampleSmall noise: ";
  printHistogram(data, mean, step);
  data.resize(0);
  step = sigma*(2+((m&1)? sqrt(m*log(phim)) : sqrt(phim*log(phim))))*1.2;
  for (long i=0; i<N; i++) {
    sampleGaussian(f, palg, sigma);
    data.push_back(embeddingLargestCoeff(f, palg));
  }
  cout << "sampleGaussian noise: ";
  printHistogram(data, mean, step);
  data.resize(0);
  exit(0);
#endif
#if 0
  // First test: fresh ciphertext, sampling mod Phi_m(X)
  for (long i=0; i<N; i++) {
    std::vector<cx_double> cemb;
    freshCtxtNoise(f, palg, sigma, true);
    canonicalEmbedding(cemb, f, palg); // Canonical embedding of f
    for (auto& entry : cemb) {
      double sz2 = conv<double>(std::norm<double>(entry));
      data.push_back(std::sqrt(sz2));
      sum += entry/double(lsize(cemb));
    }
  }
  mean = sum / double(N);
  cout << "fresh ctxt noise, sample mod Phi_m(X): ";
  printHistogram(data, mean, step);
#endif
#if 0
  for (long i=0; i<N; i++) {
    std::vector<cx_double> cemb;
    roundingNoise(f, palg, p2r, false);
    canonicalEmbedding(cemb, f, palg); // Canonical embedding of f
    for (auto& entry : cemb) {
      double sz2 = conv<double>(std::norm<double>(entry));
      data.push_back(std::sqrt(sz2));
      sum += entry/double(lsize(cemb));
    }
  }
  mean = sum / double(N);
  cout << "rounding noise, sample mod X^m-1: ";
  printHistogram(data, mean, step);
#endif


/* ---------------------------------------------------------------------
 *
 * Copyright (C) 2009 - 2021 by the deal.II authors
 *
 * This file is part of the deal.II library.
 *
 * The deal.II library is free software; you can use it, redistribute
 * it, and/or modify it under the terms of the GNU Lesser General
 * Public License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 * The full text of the license can be found in the file LICENSE.md at
 * the top level directory of deal.II.
 *
 * ---------------------------------------------------------------------
 *
 * Author: Guido Kanschat, Texas A&M University, 2009
 *         Timo Heister, Clemson University, 2019
 */

// 前面几个文件已经在前面的例子中讲过了，因此不再做进一步的评论。

#include <deal.II/base/quadrature_lib.h>
#include <deal.II/base/function.h>
#include <deal.II/lac/vector.h>
#include <deal.II/lac/dynamic_sparsity_pattern.h>
#include <deal.II/lac/sparse_matrix.h>
#include <deal.II/grid/tria.h>
#include <deal.II/grid/grid_generator.h>
#include <deal.II/grid/grid_out.h>
#include <deal.II/grid/grid_refinement.h>
#include <deal.II/fe/fe_values.h>
#include <deal.II/dofs/dof_handler.h>
#include <deal.II/numerics/vector_tools.h>
#include <deal.II/dofs/dof_tools.h>
#include <deal.II/numerics/data_out.h>
#include <deal.II/fe/mapping_q1.h>

// 这里定义了不连续的有限元。它们的使用方式与所有其他有限元相同，不过--正如你在以前的教程程序中所看到的--用户与有限元类的交互根本不多：它们被传递给 DoFHandler 和 FEValues 对象，仅此而已。

#include <deal.II/fe/fe_dgq.h>

// FEInterfaceValues需要这个头来计算界面上的积分。

#include <deal.II/fe/fe_interface_values.h>

// 我们将使用最简单的求解器，称为Richardson迭代，它代表了一个简单的缺陷修正。这与一个块状SSOR预处理器（定义在precondition_block.h中）相结合，该预处理器使用DG离散产生的系统矩阵的特殊块状结构。

#include <deal.II/lac/solver_richardson.h>
#include <deal.II/lac/precondition_block.h>

// 我们将使用梯度作为细化指标。

#include <deal.II/numerics/derivative_approximation.h>

// 最后，新的包含文件用于使用MeshWorker框架中的Mesh_loop。

#include <deal.II/meshworker/mesh_loop.h>

// 像所有的程序一样，我们在完成这一部分时，要包括所需的C++头文件，并声明我们要使用dealii命名空间中的对象，不含前缀。

#include <iostream>
#include <fstream>

namespace Step12
{
  using namespace dealii;
// @sect3{Equation data}

// 首先，我们定义一个描述不均匀边界数据的类。由于只使用它的值，我们实现value_list()，但不定义Function的所有其他函数。

  template <int dim>
  class BoundaryValues : public Function<dim>
  {
  public:
    BoundaryValues() = default;
    virtual void value_list(const std::vector<Point<dim>> &points,
                            std::vector<double> &          values,
                            const unsigned int component = 0) const override;
  };

// 考虑到流动方向，单位方块 $[0,1]^2$ 的流入边界是右边界和下边界。我们在x轴上规定了不连续的边界值1和0，在右边界上规定了值0。该函数在流出边界上的值将不会在DG方案中使用。

  template <int dim>
  void BoundaryValues<dim>::value_list(const std::vector<Point<dim>> &points,
                                       std::vector<double> &          values,
                                       const unsigned int component) const
  {
    (void)component;
    AssertIndexRange(component, 1);
    Assert(values.size() == points.size(),
           ExcDimensionMismatch(values.size(), points.size()));

    for (unsigned int i = 0; i < values.size(); ++i)
      {
        if (points[i](0) < 0.5)
          values[i] = 1.;
        else
          values[i] = 0.;
      }
  }

// 最后，一个计算并返回风场的函数  $\beta=\beta(\mathbf x)$  。正如在介绍中所解释的，在2D中我们将使用一个围绕原点的旋转场。在3D中，我们只需不设置 $z$ 分量（即为零），而这个函数在目前的实现中不能用于1D。

  template <int dim>
  Tensor<1, dim> beta(const Point<dim> &p)
  {
    Assert(dim >= 2, ExcNotImplemented());

    Tensor<1, dim> wind_field;
    wind_field[0] = -p[1];
    wind_field[1] = p[0];

    if (wind_field.norm() > 1e-10)
      wind_field /= wind_field.norm();

    return wind_field;
  }
// @sect3{The ScratchData and CopyData classes}

// 以下对象是我们在调用 MeshWorker::mesh_loop(). 时使用的抓取和复制对象 新对象是FEInterfaceValues对象，它的工作原理类似于FEValues或FEFacesValues，只是它作用于两个单元格之间的接口，并允许我们以我们的弱形式组装接口条款。

  template <int dim>
  struct ScratchData
  {
    ScratchData(const Mapping<dim> &       mapping,
                const FiniteElement<dim> & fe,
                const Quadrature<dim> &    quadrature,
                const Quadrature<dim - 1> &quadrature_face,
                const UpdateFlags          update_flags = update_values |
                                                 update_gradients |
                                                 update_quadrature_points |
                                                 update_JxW_values,
                const UpdateFlags interface_update_flags =
                  update_values | update_gradients | update_quadrature_points |
                  update_JxW_values | update_normal_vectors)
      : fe_values(mapping, fe, quadrature, update_flags)
      , fe_interface_values(mapping,
                            fe,
                            quadrature_face,
                            interface_update_flags)
    {}

    ScratchData(const ScratchData<dim> &scratch_data)
      : fe_values(scratch_data.fe_values.get_mapping(),
                  scratch_data.fe_values.get_fe(), fas
                  scratch_data.fe_values.get_quadrature(),
                  scratch_data.fe_values.get_update_flags())
      , fe_interface_values(scratch_data.fe_interface_values.get_mapping(),
                            scratch_data.fe_interface_values.get_fe(),
                            scratch_data.fe_interface_values.get_quadrature(),
                            scratch_data.fe_interface_values.get_update_flags())
    {}

    FEValues<dim>          fe_values;
    FEInterfaceValues<dim> fe_interface_values;
  };

  struct CopyDataFace
  {
    FullMatrix<double>                   cell_matrix;
    std::vector<types::global_dof_index> joint_dof_indices;
  };

  struct CopyData
  {
    FullMatrix<double>                   cell_matrix;
    Vector<double>                       cell_rhs;
    std::vector<types::global_dof_index> local_dof_indices;
    std::vector<CopyDataFace>            face_data;

    template <class Iterator>
    void reinit(const Iterator &cell, unsigned int dofs_per_cell)
    {
      cell_matrix.reinit(dofs_per_cell, dofs_per_cell);
      cell_rhs.reinit(dofs_per_cell);

      local_dof_indices.resize(dofs_per_cell);
      cell->get_dof_indices(local_dof_indices);
    }
  };
// @sect3{The AdvectionProblem class}

// 在这个准备工作之后，我们继续进行这个程序的主类，称为AdvectionProblem。

// 这对你来说应该是非常熟悉的。有趣的细节只有在实现集合函数的时候才会出现。

  template <int dim>
  class AdvectionProblem
  {
  public:
    AdvectionProblem();
    void run();

  private:
    void setup_system();
    void assemble_system();
    void solve();
    void refine_grid();
    void output_results(const unsigned int cycle) const;

    Triangulation<dim>   triangulation;
    const MappingQ1<dim> mapping;

// 此外，我们要使用DG元素。

    const FE_DGQ<dim> fe;
    DoFHandler<dim>   dof_handler;

    const QGauss<dim>     quadrature;
    const QGauss<dim - 1> quadrature_face;

// 接下来的四个成员代表要解决的线性系统。  system_matrix and right_hand_side 是由 assemble_system(), the solution 产生的，在 solve(). The sparsity_pattern 中计算，用于确定 system_matrix 中非零元素的位置。

    SparsityPattern      sparsity_pattern;
    SparseMatrix<double> system_matrix;

    Vector<double> solution;
    Vector<double> right_hand_side;
  };

// 我们从构造函数开始。 fe 的构造器调用中的1是多项式的度数。

  template <int dim>
  AdvectionProblem<dim>::AdvectionProblem()
    : mapping()
    , fe(1)
    , dof_handler(triangulation)
    , quadrature(fe.tensor_degree() + 1)
    , quadrature_face(fe.tensor_degree() + 1)
  {}

  template <int dim>
  void AdvectionProblem<dim>::setup_system()
  {

// 在设置通常的有限元数据结构的函数中，我们首先需要分配DoF。

    dof_handler.distribute_dofs(fe);

// 我们从生成稀疏模式开始。为此，我们首先用系统中出现的耦合物填充一个动态稀疏模式（DynamicSparsityPattern）类型的中间对象。在建立模式之后，这个对象被复制到 sparsity_pattern 并可以被丢弃。

// 为了建立DG离散的稀疏模式，我们可以调用类似于 DoFTools::make_sparsity_pattern, 的函数，该函数被称为 DoFTools::make_flux_sparsity_pattern:  。
    DynamicSparsityPattern dsp(dof_handler.n_dofs());
    DoFTools::make_flux_sparsity_pattern(dof_handler, dsp);
    sparsity_pattern.copy_from(dsp);

// 最后，我们设置了线性系统的所有组成部分的结构。

    system_matrix.reinit(sparsity_pattern);
    solution.reinit(dof_handler.n_dofs());
    right_hand_side.reinit(dof_handler.n_dofs());
  }
// @sect4{The assemble_system function}

// 这里我们看到了与手工组装的主要区别。我们不需要在单元格和面上写循环，而是在调用 MeshWorker::mesh_loop() 时包含逻辑，我们只需要指定在每个单元格、每个边界面和每个内部面应该发生什么。这三个任务是由下面的函数里面的lambda函数处理的。

  template <int dim>
  void AdvectionProblem<dim>::assemble_system()
  {
    using Iterator = typename DoFHandler<dim>::active_cell_iterator;
    const BoundaryValues<dim> boundary_function;

// 这是将对每个单元格执行的函数。

    const auto cell_worker = [&](const Iterator &  cell,
                                 ScratchData<dim> &scratch_data,
                                 CopyData &        copy_data) {
      const unsigned int n_dofs =
        scratch_data.fe_values.get_fe().n_dofs_per_cell();
      copy_data.reinit(cell, n_dofs);
      scratch_data.fe_values.reinit(cell);

      const auto &q_points = scratch_data.fe_values.get_quadrature_points();

      const FEValues<dim> &      fe_v = scratch_data.fe_values;
      const std::vector<double> &JxW  = fe_v.get_JxW_values();

// 我们解决的是一个同质方程，因此在单元项中没有显示出右手。 剩下的就是整合矩阵条目。

      for (unsigned int point = 0; point < fe_v.n_quadrature_points; ++point)
        {
          auto beta_q = beta(q_points[point]);
          for (unsigned int i = 0; i < n_dofs; ++i)
            for (unsigned int j = 0; j < n_dofs; ++j)
              {
                copy_data.cell_matrix(i, j) +=
                  -beta_q                      // -\beta
                  * fe_v.shape_grad(i, point)  // \nabla \phi_i
                  * fe_v.shape_value(j, point) // \phi_j
                  * JxW[point];                // dx
              }
        }
    };

// 这是为边界面调用的函数，包括使用FEFaceValues的正常积分。新的逻辑是决定该术语是进入系统矩阵（流出）还是进入右手边（流入）。

    const auto boundary_worker = [&](const Iterator &    cell,
                                     const unsigned int &face_no,
                                     ScratchData<dim> &  scratch_data,
                                     CopyData &          copy_data) {
      scratch_data.fe_interface_values.reinit(cell, face_no);
      const FEFaceValuesBase<dim> &fe_face =
        scratch_data.fe_interface_values.get_fe_face_values(0);

      const auto &q_points = fe_face.get_quadrature_points();

      const unsigned int n_facet_dofs = fe_face.get_fe().n_dofs_per_cell();
      const std::vector<double> &        JxW     = fe_face.get_JxW_values();
      const std::vector<Tensor<1, dim>> &normals = fe_face.get_normal_vectors();

      std::vector<double> g(q_points.size());
      boundary_function.value_list(q_points, g);

      for (unsigned int point = 0; point < q_points.size(); ++point)
        {
          const double beta_dot_n = beta(q_points[point]) * normals[point];

          if (beta_dot_n > 0)
            {
              for (unsigned int i = 0; i < n_facet_dofs; ++i)
                for (unsigned int j = 0; j < n_facet_dofs; ++j)
                  copy_data.cell_matrix(i, j) +=
                    fe_face.shape_value(i, point)   // \phi_i
                    * fe_face.shape_value(j, point) // \phi_j
                    * beta_dot_n                    // \beta . n
                    * JxW[point];                   // dx
            }
          else
            for (unsigned int i = 0; i < n_facet_dofs; ++i)
              copy_data.cell_rhs(i) += -fe_face.shape_value(i, point) // \phi_i
                                       * g[point]                     // g
                                       * beta_dot_n  // \beta . n
                                       * JxW[point]; // dx
        }
    };

// 这是在内部面调用的函数。参数指定了单元格、面和子面的指数（用于自适应细化）。我们只是将它们传递给FEInterfaceValues的reinit()函数。

    const auto face_worker = [&](const Iterator &    cell,
                                 const unsigned int &f,
                                 const unsigned int &sf,
                                 const Iterator &    ncell,
                                 const unsigned int &nf,
                                 const unsigned int &nsf,
                                 ScratchData<dim> &  scratch_data,
                                 CopyData &          copy_data) {
      FEInterfaceValues<dim> &fe_iv = scratch_data.fe_interface_values;
      fe_iv.reinit(cell, f, sf, ncell, nf, nsf);
      const auto &q_points = fe_iv.get_quadrature_points();

      copy_data.face_data.emplace_back();
      CopyDataFace &copy_data_face = copy_data.face_data.back();

      const unsigned int n_dofs        = fe_iv.n_current_interface_dofs();
      copy_data_face.joint_dof_indices = fe_iv.get_interface_dof_indices();

      copy_data_face.cell_matrix.reinit(n_dofs, n_dofs);

      const std::vector<double> &        JxW     = fe_iv.get_JxW_values();
      const std::vector<Tensor<1, dim>> &normals = fe_iv.get_normal_vectors();

      for (unsigned int qpoint = 0; qpoint < q_points.size(); ++qpoint)
        {
          const double beta_dot_n = beta(q_points[qpoint]) * normals[qpoint];
          for (unsigned int i = 0; i < n_dofs; ++i)
            for (unsigned int j = 0; j < n_dofs; ++j)
              copy_data_face.cell_matrix(i, j) +=
                fe_iv.jump(i, qpoint) // [\phi_i]
                *
                fe_iv.shape_value((beta_dot_n > 0), j, qpoint) // phi_j^{upwind}
                * beta_dot_n                                   // (\beta . n)
                * JxW[qpoint];                                 // dx
        }
    };

// 下面的lambda函数将处理从单元格和面组件中复制数据到全局矩阵和右侧的问题。

// 虽然我们不需要AffineConstraints对象，因为在DG离散中没有悬空节点约束，但我们在这里使用一个空对象，因为这允许我们使用其`copy_local_to_global`功能。

    const AffineConstraints<double> constraints;

    const auto copier = [&](const CopyData &c) {
      constraints.distribute_local_to_global(c.cell_matrix,
                                             c.cell_rhs,
                                             c.local_dof_indices,
                                             system_matrix,
                                             right_hand_side);

      for (auto &cdf : c.face_data)
        {
          constraints.distribute_local_to_global(cdf.cell_matrix,
                                                 cdf.joint_dof_indices,
                                                 system_matrix);
        }
    };

    ScratchData<dim> scratch_data(mapping, fe, quadrature, quadrature_face);
    CopyData         copy_data;

// 在这里，我们最终处理了装配问题。我们传入ScratchData和CopyData对象，以及上面的lambda函数，并指定我们要对内部面进行一次装配。

    MeshWorker::mesh_loop(dof_handler.begin_active(),
                          dof_handler.end(),
                          cell_worker,
                          copier,
                          scratch_data,
                          copy_data,
                          MeshWorker::assemble_own_cells |
                            MeshWorker::assemble_boundary_faces |
                            MeshWorker::assemble_own_interior_faces_once,
                          boundary_worker,
                          face_worker);
  }
// @sect3{All the rest}

// 对于这个简单的问题，我们使用了最简单的求解器，称为Richardson迭代，它代表了简单的缺陷修正。这与一个块状SSOR预处理相结合，该预处理使用DG离散化产生的系统矩阵的特殊块状结构。这些块的大小是每个单元的DoF数量。这里，我们使用SSOR预处理，因为我们没有根据流场对DoFs进行重新编号。如果在流的下游方向对DoFs进行重新编号，那么块状的Gauss-Seidel预处理（见PreconditionBlockSOR类，放松=1）会做得更好。

  template <int dim>
  void AdvectionProblem<dim>::solve()
  {
    SolverControl                    solver_control(1000, 1e-12);
    SolverRichardson<Vector<double>> solver(solver_control);

// 这里我们创建了预处理程序。

    PreconditionBlockSSOR<SparseMatrix<double>> preconditioner;

// 然后将矩阵分配给它，并设置正确的块大小。

    preconditioner.initialize(system_matrix, fe.n_dofs_per_cell());

// 做完这些准备工作后，我们就可以启动线性求解器了。

    solver.solve(system_matrix, solution, right_hand_side, preconditioner);

    std::cout << "  Solver converged in " << solver_control.last_step()
              << " iterations." << std::endl;
  }

// 我们根据一个非常简单的细化标准来细化网格，即对解的梯度的近似。由于这里我们考虑的是DG(1)方法（即我们使用片状双线性形状函数），我们可以简单地计算每个单元的梯度。但是我们并不希望我们的细化指标只建立在每个单元的梯度上，而是希望同时建立在相邻单元之间的不连续解函数的跳跃上。最简单的方法是通过差分商计算近似梯度，包括考虑中的单元和其相邻的单元。这是由 DerivativeApproximation 类完成的，它计算近似梯度的方式类似于本教程 step-9 中描述的 GradientEstimation 。事实上， DerivativeApproximation 类是在 step-9 的 GradientEstimation 类之后开发的。与  step-9  中的讨论相关，这里我们考虑  $h^{1+d/2}|\nabla_h u_h|$  。此外，我们注意到，我们不考虑近似的二次导数，因为线性平流方程的解一般不在 $H^2$ 中，而只在 $H^1$ 中（或者，更准确地说：在 $H^1_\beta$ 中，即在方向 $\beta$ 上的导数是可平方整除的函数空间）。

  template <int dim>
  void AdvectionProblem<dim>::refine_grid()
  {

//  DerivativeApproximation 类将梯度计算为浮点精度。这已经足够了，因为它们是近似的，只作为细化指标。

    Vector<float> gradient_indicator(triangulation.n_active_cells());

// 现在，近似梯度被计算出来了

    DerivativeApproximation::approximate_gradient(mapping,
                                                  dof_handler,
                                                  solution,
                                                  gradient_indicator);

//并且它们的单元格按系数 $h^{1+d/2}$ 进行缩放。
    unsigned int cell_no = 0;
    for (const auto &cell : dof_handler.active_cell_iterators())
      gradient_indicator(cell_no++) *=
        std::pow(cell->diameter(), 1 + 1.0 * dim / 2);

// 最后它们作为细化指标。

    GridRefinement::refine_and_coarsen_fixed_number(triangulation,
                                                    gradient_indicator,
                                                    0.3,
                                                    0.1);

    triangulation.execute_coarsening_and_refinement();
  }

// 这个程序的输出包括一个自适应细化网格的vtk文件和数值解。最后，我们还用 VectorTools::integrate_difference(). 计算了解的L-无穷大规范。
  template <int dim>
  void AdvectionProblem<dim>::output_results(const unsigned int cycle) const
  {
    const std::string filename = "solution-" + std::to_string(cycle) + ".vtk";
    std::cout << "  Writing solution to <" << filename << ">" << std::endl;
    std::ofstream output(filename);

    DataOut<dim> data_out;
    data_out.attach_dof_handler(dof_handler);
    data_out.add_data_vector(solution, "u", DataOut<dim>::type_dof_data);

    data_out.build_patches(mapping);

    data_out.write_vtk(output);

    {
      Vector<float> values(triangulation.n_active_cells());
      VectorTools::integrate_difference(mapping,
                                        dof_handler,
                                        solution,
                                        Functions::ZeroFunction<dim>(),
                                        values,
                                        quadrature,
                                        VectorTools::Linfty_norm);
      const double l_infty =
        VectorTools::compute_global_error(triangulation,
                                          values,
                                          VectorTools::Linfty_norm);
      std::cout << "  L-infinity norm: " << l_infty << std::endl;
    }
  }

// 下面的 run 函数与前面的例子类似。

  template <int dim>
  void AdvectionProblem<dim>::run()
  {
    for (unsigned int cycle = 0; cycle < 6; ++cycle)
      {
        std::cout << "Cycle " << cycle << std::endl;

        if (cycle == 0)
          {
            GridGenerator::hyper_cube(triangulation);
            triangulation.refine_global(3);
          }
        else
          refine_grid();

        std::cout << "  Number of active cells:       "
                  << triangulation.n_active_cells() << std::endl;

        setup_system();

        std::cout << "  Number of degrees of freedom: " << dof_handler.n_dofs()
                  << std::endl;

        assemble_system();
        solve();

        output_results(cycle);
      }
  }
} // namespace Step12

// 下面的 main 函数与前面的例子也类似，不需要注释。

int main()
{
  try
    {
      Step12::AdvectionProblem<2> dgmethod;
      dgmethod.run();
    }
  catch (std::exception &exc)
    {
      std::cerr << std::endl
                << std::endl
                << "----------------------------------------------------"
                << std::endl;
      std::cerr << "Exception on processing: " << std::endl
                << exc.what() << std::endl
                << "Aborting!" << std::endl
                << "----------------------------------------------------"
                << std::endl;
      return 1;
    }
  catch (...)
    {
      std::cerr << std::endl
                << std::endl
                << "----------------------------------------------------"
                << std::endl;
      std::cerr << "Unknown exception!" << std::endl
                << "Aborting!" << std::endl
                << "----------------------------------------------------"
                << std::endl;
      return 1;
    }

  return 0;
}


/*
 * Copyright (C) 2019  Rhys Mainwaring
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
*/

#include <cmath>
#include <iostream>
#include <string>

#include <Eigen/Dense>
#include <ignition/common/Console.hh>
#include <ignition/math/Pose3.hh>
#include <ignition/math/Vector2.hh>
#include <ignition/math/Vector3.hh>

#include "Wavefield.hh"

using namespace ignition;
using namespace gazebo;

///////////////////////////////////////////////////////////////////////////////
// Utilities
std::ostream& operator<<(std::ostream &_os, const std::vector<double> &_vec)
{
  for (auto&& v : _vec ) // NOLINT
    _os << v << ", ";
  return _os;
}

///////////////////////////////////////////////////////////////////////////////
/// \brief Private data for the WavefieldParameters.
class ignition::gazebo::WavefieldPrivate
{
  /// \brief Constructor.
  public: WavefieldPrivate():
    size({1000, 1000}),
    cellCount({50, 50}),
    model("PMS"),
    number(1),
    scale(2),
    angle(2.0*M_PI/10.0),
    steepness(1.0),
    amplitude(0.0),
    period(1.0),
    phase(0.0),
    direction(1, 0),
    angularFrequency(2.0*M_PI),
    wavelength(2 * M_PI / this->DeepWaterDispersionToWavenumber(2.0 * M_PI)),
    wavenumber(this->DeepWaterDispersionToWavenumber(2.0 * M_PI)),
    tau(1.0),
    gain(1.0)
  {
  }

  /// \brief The size of the wavefield.
  public: ignition::math::Vector2d size;

  /// \brief The number of grid cells in the wavefield.
  public: ignition::math::Vector2d cellCount;

  /// \brief Name of wavefield model to use - must be "PMS" or "CWR"
  public: std::string model;

  /// \brief The number of component waves.
  public: size_t number;

  /// \brief Set the scale of the largest and smallest waves.
  public: double scale;

  /// \brief Set the angle between component waves and the mean direction.
  public: double angle;

  /// \brief Control the wave steepness. 0 is sine waves, 1 is Gerstner waves.
  public: double steepness;

  /// \brief The mean wave amplitude [m].
  public: double amplitude;

  /// \brief The mean wave period [s]
  public: double period;

  /// \brief The mean wave phase (not currently enabled).
  public: double phase;

  /// \brief The mean wave direction.
  public: ignition::math::Vector2d direction;

  /// \brief The time constant for exponential increasing waves on startup
  public: double tau;

  /// \brief The multiplier applied to PM spectra
  public: double gain;

  /// \brief The mean wave angular frequency (derived).
  public: double angularFrequency;

  /// \brief The mean wavelength (derived).
  public: double wavelength;

  /// \brief The mean wavenumber (derived).
  public: double wavenumber;

  /// \brief The component wave angular frequencies (derived).
  public: std::vector<double> angularFrequencies;

  /// \brief The component wave amplitudes (derived).
  public: std::vector<double> amplitudes;

  /// \brief The component wave phases (derived).
  public: std::vector<double> phases;

  /// \brief The component wave steepness factors (derived).
  public: std::vector<double> steepnesses;

  /// \brief The component wavenumbers (derived).
  public: std::vector<double> wavenumbers;

  /// \brief The component wave dirctions (derived).
  public: std::vector<ignition::math::Vector2d> directions;

  /// \brief Recalculate for constant wavelength-amplitude ratio
  public: void RecalculateCwr()
  {
    // Normalize direction
    this->direction.Normalize();

    // Derived mean values
    this->angularFrequency = 2.0 * M_PI / this->period;
    this->wavenumber = \
      this->DeepWaterDispersionToWavenumber(this->angularFrequency);
    this->wavelength = 2.0 * M_PI / this->wavenumber;

    // Update components
    this->angularFrequencies.clear();
    this->amplitudes.clear();
    this->phases.clear();
    this->wavenumbers.clear();
    this->steepnesses.clear();
    this->directions.clear();

    for (size_t i = 0; i < this->number; ++i)
    {
      const int n = i - this->number / 2;
      const double scaleFactor = std::pow(this->scale, n);
      const double a = scaleFactor * this->amplitude;
      const double k = this->wavenumber / scaleFactor;
      const double omega = this->DeepWaterDispersionToOmega(k);
      const double phi = this->phase;
      double q = 0.0;
      if (!ignition::math::equal(a, 0.0))
      {
        q = std::min(1.0, this->steepness / (a * k * this->number));
      }

      this->amplitudes.push_back(a);
      this->angularFrequencies.push_back(omega);
      this->phases.push_back(phi);
      this->steepnesses.push_back(q);
      this->wavenumbers.push_back(k);

      // Direction
      const double c = std::cos(n * this->angle);
      const double s = std::sin(n * this->angle);

      const ignition::math::Vector2d d(
        c * this->direction.X() - s * this->direction.Y(),
        s * this->direction.X() + c * this->direction.Y());
      directions.push_back(d);
    }
  }

  // \brief Pierson-Moskowitz wave spectrum
  public: double pm(double _omega, double _omegaP)
  {
    double alpha = 0.0081;
    double g = 9.81;
    return alpha * std::pow(g, 2.0) / std::pow(_omega, 5.0) * \
      std::exp(-(5.0 / 4.0) * std::pow(_omegaP / _omega, 4.0));
  }

  /// \brief Recalculate for Pierson-Moskowitz spectrum sampling model
  public: void RecalculatePms()
  {
    // Normalize direction
    this->direction.Normalize();

    // Derived mean values
    this->angularFrequency = 2.0 * M_PI / this->period;
    this->wavenumber = \
      this->DeepWaterDispersionToWavenumber(this->angularFrequency);
    this->wavelength = 2.0 * M_PI / this->wavenumber;

    // Update components
    this->angularFrequencies.clear();
    this->amplitudes.clear();
    this->phases.clear();
    this->wavenumbers.clear();
    this->steepnesses.clear();
    this->directions.clear();

    // Vector for spaceing
    std::vector<double> omegaSpacing;
    omegaSpacing.push_back(this->angularFrequency * (1.0 - 1.0 / this->scale));
    omegaSpacing.push_back(this->angularFrequency * \
                            (this->scale - 1.0 / this->scale) / 2.0);
    omegaSpacing.push_back(this->angularFrequency * (this->scale - 1.0));

    for (size_t i = 0; i < this->number; ++i)
    {
      const int n = i - 1;
      const double scaleFactor = std::pow(this->scale, n);
      const double omega = this->angularFrequency * scaleFactor;
      const double pms = pm(omega, this->angularFrequency);
      const double a = this->gain * std::sqrt(2.0 * pms * omegaSpacing[i]);
      const double k = this->DeepWaterDispersionToWavenumber(omega);
      const double phi = this->phase;
      double q = 0.0;
      if (!ignition::math::equal(a, 0.0))
      {
        q = std::min(1.0, this->steepness / (a * k * this->number));
      }

      this->amplitudes.push_back(a);
      this->angularFrequencies.push_back(omega);
      this->phases.push_back(phi);
      this->steepnesses.push_back(q);
      this->wavenumbers.push_back(k);

      // Direction
      const double c = std::cos(n * this->angle);
      const double s = std::sin(n * this->angle);

      const ignition::math::Vector2d d(
        c * this->direction.X() - s * this->direction.Y(),
        s * this->direction.X() + c * this->direction.Y());
      directions.push_back(d);
    }
  }

  /// \brief Recalculate all derived quantities from inputs.
  public: void Recalculate()
  {
    if (!this->model.compare("PMS"))
    {
      ignmsg << "Using Pierson-Moskowitz spectrum sampling wavefield model "
            << std::endl;
      this->RecalculatePms();
    }
    else if (!this->model.compare("CWR"))
    {
      ignmsg << "Using Constant wavelength-ampltude ratio wavefield model "
            << std::endl;
      this->RecalculateCwr();
    }
    else
    {
      ignwarn<< "Wavefield model specified as <" << this->model
            << "> which is not one of the two supported wavefield models: "
            << "PMS or CWR!!!" << std::endl;
    }
  }

  /////////////////////////////////////////////////
  private: double DeepWaterDispersionToOmega(double _wavenumber)
  {
    const double g = std::fabs(-9.8);
    return std::sqrt(g * _wavenumber);
  }

  /////////////////////////////////////////////////
  private: double DeepWaterDispersionToWavenumber(double _omega)
  {
    const double g = std::fabs(-9.8);
    return _omega * _omega / g;
  }
};

/////////////////////////////////////////////////////////////////////////////
Wavefield::Wavefield()
  : data(std::make_unique<WavefieldPrivate>())
{
  this->data->Recalculate();
}

///////////////////////////////////////////////////////////////////////////////
Wavefield::~Wavefield()
{
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::Load(const std::shared_ptr<const sdf::Element> &_sdf)
{
  if (!_sdf->HasElement("wavefield"))
    return;

  auto ptr = const_cast<sdf::Element *>(_sdf.get());
  auto sdfWavefield = ptr->GetElement("wavefield");

  this->data->size = sdfWavefield->Get<ignition::math::Vector2d>("size",
    this->data->size).first;
  this->data->cellCount = sdfWavefield->Get<ignition::math::Vector2d>("cell_count",
    this->data->cellCount).first;
  if (sdfWavefield->HasElement("wave"))
  {
    auto sdfWave = sdfWavefield->GetElement("wave");

    this->data->model = sdfWave->Get<std::string>("model", "PMS").first;
    this->data->number =
      sdfWave->Get<double>("number", this->data->number).first;
    this->data->amplitude =
      sdfWave->Get<double>("amplitude", this->data->amplitude).first;
    this->data->period =
      sdfWave->Get<double>("period", this->data->period).first;
    this->data->phase = sdfWave->Get<double>("phase", this->data->phase).first;
    this->data->direction =
      sdfWave->Get<ignition::math::Vector2d>("direction",
        this->data->direction).first;
    this->data->scale = sdfWave->Get<double>("scale", this->data->scale).first;
    this->data->angle = sdfWave->Get<double>("angle", this->data->angle).first;
    this->data->steepness =
      sdfWave->Get<double>("steepness", this->data->steepness).first;
    this->data->tau = sdfWave->Get<double>("tau", this->data->tau).first;
    this->data->gain = sdfWave->Get<double>("gain", this->data->gain).first;

    this->data->Recalculate();
  }
  this->DebugPrint();
}

///////////////////////////////////////////////////////////////////////////////
size_t Wavefield::Number() const
{
  return this->data->number;
}

///////////////////////////////////////////////////////////////////////////////
double Wavefield::Angle() const
{
  return this->data->angle;
}

///////////////////////////////////////////////////////////////////////////////
double Wavefield::Scale() const
{
  return this->data->scale;
}

///////////////////////////////////////////////////////////////////////////////
double Wavefield::Steepness() const
{
  return this->data->steepness;
}

///////////////////////////////////////////////////////////////////////////////
double Wavefield::AngularFrequency() const
{
  return this->data->angularFrequency;
}

///////////////////////////////////////////////////////////////////////////////
double Wavefield::Amplitude() const
{
  return this->data->amplitude;
}

///////////////////////////////////////////////////////////////////////////////
double Wavefield::Period() const
{
  return this->data->period;
}

///////////////////////////////////////////////////////////////////////////////
double Wavefield::Phase() const
{
  return this->data->phase;
}

///////////////////////////////////////////////////////////////////////////////
double Wavefield::Wavelength() const
{
  return this->data->wavelength;
}

///////////////////////////////////////////////////////////////////////////////
double Wavefield::Wavenumber() const
{
  return this->data->wavenumber;
}

///////////////////////////////////////////////////////////////////////////////
float Wavefield::Tau() const
{
  return this->data->tau;
}

///////////////////////////////////////////////////////////////////////////////
float Wavefield::Gain() const
{
  return this->data->gain;
}

///////////////////////////////////////////////////////////////////////////////
ignition::math::Vector2d Wavefield::Direction() const
{
  return this->data->direction;
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::SetNumber(size_t _number)
{
  this->data->number = _number;
  this->data->Recalculate();
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::SetAngle(double _angle)
{
  this->data->angle = _angle;
  this->data->Recalculate();
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::SetScale(double _scale)
{
  this->data->scale = _scale;
  this->data->Recalculate();
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::SetSteepness(double _steepness)
{
  this->data->steepness = _steepness;
  this->data->Recalculate();
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::SetAmplitude(double _amplitude)
{
  this->data->amplitude = _amplitude;
  this->data->Recalculate();
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::SetPeriod(double _period)
{
  this->data->period = _period;
  this->data->Recalculate();
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::SetPhase(double _phase)
{
  this->data->phase = _phase;
  this->data->Recalculate();
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::SetTau(double _tau)
{
  this->data->tau = _tau;
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::SetGain(double _gain)
{
  this->data->gain = _gain;
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::SetDirection(const ignition::math::Vector2d &_direction)
{
  this->data->direction = _direction;
  this->data->Recalculate();
}

///////////////////////////////////////////////////////////////////////////////
const std::vector<double> &Wavefield::AngularFrequency_V() const
{
  return this->data->angularFrequencies;
}

///////////////////////////////////////////////////////////////////////////////
const std::vector<double> &Wavefield::Amplitude_V() const
{
  return this->data->amplitudes;
}

///////////////////////////////////////////////////////////////////////////////
const std::vector<double> &Wavefield::Phase_V() const
{
  return this->data->phases;
}

///////////////////////////////////////////////////////////////////////////////
const std::vector<double> &Wavefield::Steepness_V() const
{
  return this->data->steepnesses;
}

///////////////////////////////////////////////////////////////////////////////
const std::vector<double> &Wavefield::Wavenumber_V() const
{
  return this->data->wavenumbers;
}

///////////////////////////////////////////////////////////////////////////////
const std::vector<ignition::math::Vector2d> &Wavefield::Direction_V() const
{
  return this->data->directions;
}

///////////////////////////////////////////////////////////////////////////////
void Wavefield::DebugPrint() const
{
  ignmsg << "Input Parameters:" << std::endl;
  ignmsg << "model:     " << this->data->model << std::endl;
  ignmsg << "number:     " << this->data->number << std::endl;
  ignmsg << "scale:      " << this->data->scale << std::endl;
  ignmsg << "angle:      " << this->data->angle << std::endl;
  ignmsg << "steepness:  " << this->data->steepness << std::endl;
  ignmsg << "amplitude:  " << this->data->amplitude << std::endl;
  ignmsg << "period:     " << this->data->period << std::endl;
  ignmsg << "direction:  " << this->data->direction << std::endl;
  ignmsg << "tau:  " << this->data->tau << std::endl;
  ignmsg << "gain:  " << this->data->gain << std::endl;
  ignmsg << "Derived Parameters:" << std::endl;
  ignmsg << "amplitudes:  " << this->data->amplitudes << std::endl;
  ignmsg << "wavenumbers: " << this->data->wavenumbers << std::endl;
  ignmsg << "omegas:      " << this->data->angularFrequencies << std::endl;
  ignmsg << "periods:     ";
  for (auto&& omega : this->data->angularFrequencies) // NOLINT
  {
    ignmsg << 2.0 * M_PI / omega <<", ";
  }
  ignmsg << std::endl;
  ignmsg << "phases:      " << this->data->phases << std::endl;
  ignmsg << "steepnesses: " << this->data->steepnesses << std::endl;
  ignmsg << "directions:  ";
  for (auto&& d : this->data->directions) // NOLINT
  {
    ignmsg << d << "; ";
  }
  ignmsg << std::endl;
}

///////////////////////////////////////////////////////////////////////////////
double Wavefield::ComputeDepthSimply(const ignition::math::Vector3d &_point,
  double _time, double _timeInit)
{
  double h = 0.0;
  for (std::size_t i = 0; i < this->Number(); ++i)
  {
    double k = this->Wavenumber_V()[i];
    double a = this->Amplitude_V()[i];
    double dx =  this->Direction_V()[i].X();
    double dy =  this->Direction_V()[i].Y();
    double dot = _point.X() * dx + _point.Y() * dy;
    double omega = this->AngularFrequency_V()[i];
    double theta = k * dot - omega * _time;
    double c = cos(theta);
    h += a*c;
  }

  // Exponentially grow the waves
  return h * (1 - exp(-1.0 * (_time - _timeInit) / this->Tau()));
}

///////////////////////////////////////////////////////////////////////////////
double Wavefield::ComputeDepthDirectly(const ignition::math::Vector3d &_point,
  double _time, double _timeInit)
{
  // Struture for passing wave parameters to lambdas
  struct WaveParams
  {
    WaveParams(
      const std::vector<double>& _a,
      const std::vector<double>& _k,
      const std::vector<double>& _omega,
      const std::vector<double>& _phi,
      const std::vector<double>& _q,
      const std::vector<ignition::math::Vector2d>& _dir) :
      a(_a), k(_k), omega(_omega), phi(_phi), q(_q), dir(_dir) {}

    const std::vector<double>& a;
    const std::vector<double>& k;
    const std::vector<double>& omega;
    const std::vector<double>& phi;
    const std::vector<double>& q;
    const std::vector<ignition::math::Vector2d>& dir;
  };

  // Compute the target function and Jacobian. Also calculate pz,
  // the z-component of the Gerstner wave, which we essentially get for free.
  // cppcheck-suppress constParameter
  auto wave_fdf = [=](auto x, auto p, auto t, auto &wp, auto &F, auto &J)
  {
    double pz = 0;
    F(0) = p.x() - x.x();
    F(1) = p.y() - x.y();
    J(0, 0) = -1;
    J(0, 1) =  0;
    J(1, 0) =  0;
    J(1, 1) = -1;
    const size_t n = wp.a.size();
    for (auto&& i = 0; i < n; ++i) // NOLINT
    {
      const double dx = wp.dir[i].X();
      const double dy = wp.dir[i].Y();
      const double q = wp.q[i];
      const double a = wp.a[i];
      const double k = wp.k[i];
      const double dot = x.x() * dx + x.y() * dy;
      const double theta = k * dot - wp.omega[i] * t;
      const double s = std::sin(theta);
      const double c = std::cos(theta);
      const double qakc = q * a * k * c;
      const double df1x = qakc * dx * dx;
      const double df1y = qakc * dx * dy;
      const double df2x = df1y;
      const double df2y = qakc * dy * dy;
      pz += a * c;
      F(0) += a * dx * s;
      F(1) += a * dy * s;
      J(0, 0) += df1x;
      J(0, 1) += df1y;
      J(1, 0) += df2x;
      J(1, 1) += df2y;
    }
    // Exponentially grow the waves
    return pz * (1 - exp(-1.0 * (_time - _timeInit) / this->Tau()));
  };

  // Simple multi-variate Newton solver -
  // this version returns the z-component of the
  // wave field at the desired point p.
  // cppcheck-suppress constParameter
  auto solver = [=](auto& fdfunc, auto x0, auto p, auto t, \
                    auto& wp, auto tol, auto nmax)
  {
    int n = 0;
    double err = 1;
    double pz = 0;
    auto xn = x0;
    Eigen::Vector2d F;
    Eigen::Matrix2d J;
    while (std::abs(err) > tol && n < nmax)
    {
      pz = fdfunc(x0, p, t, wp, F, J);
      xn = x0 - J.inverse() * F;
      x0 = xn;
      err = F.norm();
      n++;
    }
    return pz;
  };

  // Set up parameter references
  WaveParams wp(
    this->Amplitude_V(),
    this->Wavenumber_V(),
    this->AngularFrequency_V(),
    this->Phase_V(),
    this->Steepness_V(),
    this->Direction_V());

  // Tolerances etc.
  const double tol = 1.0E-10;
  const double nmax = 30;

  // Use the target point as the initial guess
  // (this is within sum{amplitudes} of the solution)
  Eigen::Vector2d p2(_point.X(), _point.Y());
  const double pz = solver(wave_fdf, p2, p2, _time, wp, tol, nmax);
  // Removed so that height is reported relative to mean water level
  // const double h = pz - _point.Z();
  const double h = pz;
  return h;
}


/*
 * visualization.cpp
 *
 *  Created on: Jan 18, 2012
 *      Author: eba
 */

#include "mhfpython.h"

#include <boost/foreach.hpp>
#include "visualization.h"
#include <numpy/noprefix.h>

// TODO expand the definition of this function to handle higher dimensional cases
//PyObject * actualDistributionAfterPredict(const GaussianHypothesis<dim> & gh,
//										 transformations<dim,procnoisedim,inputdim> &trans,
//										 const Matrix<double,inputdim,1> & input,
//										 const Matrix<double,procnoisedim,procnoisedim> & noisecov,
//										 int samples, double limit, const int steps) {
//    double * zData = new double[steps * steps];
//    Map<Matrix<double, Dynamic, Dynamic> > z(zData, steps,steps);
//    z.fill(0);
//    double cellArea = pow(2 * limit / steps,2);
//    const Matrix<double,dim,1> & mean = gh.mean;
//    Matrix<double,dim,dim> cholState = gh.cov.llt().matrixL();
//    Matrix<double,procnoisedim,procnoisedim> cholNoise = noisecov.llt().matrixL();
//    boost::mt19937 rng;
//    boost::normal_distribution<> nd(0,1);
//    boost::variate_generator<boost::mt19937&, boost::normal_distribution<> > var_nor(rng, nd);
//
//
//    for (int i=0; i<samples; i++) {
//		Matrix<double,dim,1> sample = mean + cholState * randomMatrix<dim,1>(var_nor);
//		Matrix<double,procnoisedim,1> noise = cholNoise * randomMatrix<procnoisedim,1>(var_nor);
//		sample = trans.statetrans(sample,noise,input);
//		if(limit>sample.maxCoeff() && sample.minCoeff()>-limit) {
//			Matrix<int,dim,1> index = (sample.array() * steps/2 / limit + steps/2).matrix().cast<int>();
//			z(index(0,0),index(1,0)) += 1./(samples*cellArea);
//		}
//	}
//    int N[] = {steps,steps};
//    PyArrayObject * result = (PyArrayObject*) PyArray_SimpleNewFromData(2,N, PyArray_DOUBLE, zData);
//    result->flags = result->flags | OWNDATA;
//    return (PyObject *)result;
//}

void propagateSamples(samplevector & samples,
					  transformations<dim,procnoisedim,inputdim> &trans,
					  const Matrix<double,inputdim,1> & input,
					  const Matrix<double,procnoisedim,procnoisedim> & noisecov
	) {
    Matrix<double,procnoisedim,procnoisedim> cholNoise = noisecov.llt().matrixL();
    boost::mt19937 rng;
    boost::normal_distribution<> nd(0,1);
    boost::variate_generator<boost::mt19937&, boost::normal_distribution<> > var_nor(rng, nd);

    BOOST_FOREACH(Sample & sample, samples) {
		Matrix<double,procnoisedim,1> noise = cholNoise * randomMatrix<procnoisedim,1>(var_nor);
		sample.state = trans.statetrans(sample.state,noise,input);
	}
}

void updateSamples(samplevector & samples,
				   transformations<dim,measdim,measdim> &trans,
				   const Matrix<double,measdim,1> & meas,
				   const Matrix<double,measdim,measdim> & noisecov
	) {

	nullclass nullobj;
	double totalWeight = 0;
    BOOST_FOREACH(Sample & sample, samples) {
		sample.weight = trans.measlikelihood(sample.state,noisecov,meas,nullobj);
		totalWeight += sample.weight;
	}
    BOOST_FOREACH(Sample & sample, samples) {
    	sample.weight /= totalWeight;
    }
}

// TODO expand the definition of this function to handle higher dimensional cases
PyObject * layoutSamples(const samplevector & samples, double limit, const int steps) {
    double * zData = new double[steps * steps];
    Map<Matrix<double, Dynamic, Dynamic> > z(zData, steps,steps);
    z.fill(0);
    double cellArea = pow(2 * limit / steps,2);

    typedef Matrix<double,dim,1> SampleMat;

    BOOST_FOREACH(const Sample & sample, samples) {
		if(limit>sample.state.block<2,1>(0,0).maxCoeff() && sample.state.block<2,1>(0,0).minCoeff()>-limit) { // TODO This line is temporary
			Matrix<int,dim,1> index = (sample.state.array() * steps/2 / limit + steps/2).matrix().cast<int>();
			z(index(0,0),index(1,0)) += sample.weight/cellArea;
		}
	}
    npy_intp N[] = {steps,steps};
    PyArrayObject * result = (PyArrayObject*) PyArray_SimpleNewFromData(2,N, PyArray_DOUBLE, zData);
    result->flags = result->flags | OWNDATA;
    return (PyObject *)result;
}

void resample(samplevector & samples) {
	random_shuffle ( samples.begin(), samples.end() );
	int sampleCount = samples.size();
	double sampleWeight = 1.0/sampleCount;
	samplevector oldsamples = samples;
	double weight = (double(rand()) / double(RAND_MAX)) * sampleWeight;
	int sampleNum = 0;
	BOOST_FOREACH(const Sample & sample, oldsamples) {
		while(sample.weight>=weight) {
			samples[sampleNum].state = sample.state;
			samples[sampleNum].weight = sampleWeight;
			weight += sampleWeight;
			sampleNum++;
		}
		weight -= sample.weight;
	}
	if(sampleNum!=sampleCount) std::cerr<<"problem!\n";
}

void export_visualization() {
	import_array();
	class_<samplevector>("samplevector");
	//def("actualDistributionAfterPredict", actualDistributionAfterPredict);
	def("drawSamples", drawSamples<dim>);
	def("propagateSamples", propagateSamples);
	def("updateSamples", updateSamples);
	def("layoutSamples", layoutSamples);
	def("resample", resample);
}