#include "FEMSpace.h"
#include "Mesh.h"
#include "Element.h"
#include "Equation.h"
#include <Eigen/Sparse>
#include "Error.h"
#include "Eigen/IterativeLinearSolvers"
#include "MultigridSolver.h"
#include "Surface.h"
#include <cmath>
#define pi 4*atan(1.0)
//===========右端项================//
double f(double *p,double t)
{
    double x = p[0];
    double y = p[1];
    //return  2 * (t - 1) * std::exp(t * t + x + y);
    return (1 + 2 * t * 4 * pi * pi) * sin(2 * pi * x) * sin(2 * pi * y);
    //return (1 + 2 * t * pi * pi) * cos(pi * x) * cos(pi * y);
}
//===========初始条件==============//
double u0(double *p)
{
    double x = p[0];
    double y = p[1];
    //return std::exp(x + y);
    //return 0;
    return 0;
}
//====真解　＋边界条件 ==============//
double u(double *p,double t)
{
    double x = p[0];
    double y = p[1];
    //return std::exp(t * t + x + y);
    return t * sin(2 * pi * x) *sin(2 * pi * y); 
    //return t * cos(pi * x) *cos(pi * y);
}
// =======终止函数：用于计算误差 =====//
double u1(double *p)
{
    double x = p[0];
    double y = p[1];
    //return std::exp(1 + x + y);
    return sin(2 * pi * x) *sin(2 * pi * y);
    //return cos(pi * x) *cos(pi * y);
}

// =======扩散系数 =============//
double c(double *p,double t)
{
    return 1;
}
typedef Eigen::SparseMatrix<Real> SpMat;
typedef Eigen::Triplet<Real> Tri;
typedef Eigen::VectorXd VectorXd;
//for windows
//g++ -o main .\HeatEquation.cpp -std=c++11 -I \\wsl$\Ubuntu-18.04\usr\include\eigen3 .\include\tinyxml2.cpp
//for linux
//g++ -o main HeatEquation.cpp -std=c++11 -I /usr/include/eigen3/ ./include/tinyxml2.cpp -g
int main(int argv,char *argc[])
{
    //======设置参数，网格大小，时间步数划分，单元选择，网格选择 =========//
    double theta = 0.5;
    int n = 6;
    double T = 1.0;
    int n_t = 16;
    double d_t = T / double(n_t);
    RectangleDomain * Domain= new RectangleDomain({{0,0},{1,0},{1,1},{0,1}});
    Mesh<2> * mesh = new Q1Mesh(Domain,{POW2(n),POW2(n)});
    Element<2> * element = new Quadrilateral_1_Element();
    //============组装过程中需要的矩阵和向量========================//
    int n_dofs_point = mesh ->n_dofs();
    SpMat M(n_dofs_point,n_dofs_point);
    SpMat A(n_dofs_point,n_dofs_point);
    std::vector<VectorXd> solution_list;
    std::vector<VectorXd> rhs_list;
    int n_Dofs = element->n_Dofs();
    std::vector<Tri> TriList((mesh->n_element()) * n_Dofs * n_Dofs);
    std::vector<Tri> TriListM((mesh->n_element()) * n_Dofs * n_Dofs);
    std::vector<Tri>::iterator it = TriList.begin();
    std::vector<Tri>::iterator it_m = TriListM.begin();
    //============组装刚度矩阵和Mass矩阵 ========================//
    for (int k = 0; k < mesh -> n_element(); k++)
    {
        std::vector<int> Idx = mesh->NodeofEle(k);
        std::vector<Dofs<2> > temp(n_Dofs);
        for(int i = 0;i < n_Dofs; i++)
        {
	        temp[i] = mesh->DofsofIndex(Idx[i]);
        }
        element->SetDofsList(temp);
        for (int i = 1; i <= n_Dofs; i++)
	    for (int j = 1; j <= n_Dofs; j++)
	    {
            Real det = element -> det_Jacobi(0,0);
            Real a = 0;
            Real m = 0;
            for (int h = 0; h < element->n_GaussPnt(); h++)
            {
	            Real xi = element->GaussionPoint(h)[0];
	            Real eta = element->GaussionPoint(h)[1];
                a = a + element-> det_Jacobi(xi, eta) * element->GaussionWeight(h) * InnerProuduct(element->gradient(xi,eta,i),element->gradient(xi,eta,j));
                m = m + element-> det_Jacobi(xi, eta) * element ->GaussionWeight(h) * element ->phi(xi,eta,i) * element -> phi(xi , eta , j);
            }
		    *it = Tri(element->NdIdx(i), element->NdIdx(j), a);
            *it_m = Tri(element->NdIdx(i), element->NdIdx(j), m);
		    it++;
            it_m++;
	    }
    }
    A.setFromTriplets(TriList.begin(), TriList.end());
    A.makeCompressed();
    M.setFromTriplets(TriListM.begin(), TriListM.end());
    M.makeCompressed();
    //std::cout << A << std::endl;
    //std::cout << M << std::endl;
    //=============组装所有右端项 ==========================//
    for(int m = 0;m <= n_t;m++)
    {
        VectorXd bm = Eigen::MatrixXd::Zero(n_dofs_point,1);
        int n_GaussPnt = element->n_GaussPnt();
        for (int k = 0; k < mesh -> n_element(); k++)
        {
            std::vector<int> Idx = mesh->NodeofEle(k);
            std::vector<Dofs<2>> temp(n_Dofs);
            for(int i = 0; i < n_Dofs; i++)
            {
	            temp[i] = mesh ->DofsofIndex(Idx[i]);
            }
            element->SetDofsList(temp);
            for(int i = 1; i <= n_Dofs; i++)
            {
                Real a = 0.0;
                for (int j = 0; j < n_GaussPnt; j++)
                {
		            double xi = element->GaussionPoint(j)[0];
		            double eta = element->GaussionPoint(j)[1];
		            double xi_ = element->Global_x(xi, eta);
		            double eta_ = element->Global_y(xi, eta);
                    Dofs<2> temp({xi_,eta_});
		            a = a + element->det_Jacobi(xi, eta) * element->GaussionWeight(j) * element->phi(xi, eta, i)* f(*temp,double(d_t * m));
                }
                bm[element->NdIdx(i)]+= a;
            }
        }
        rhs_list.push_back(bm);
    }
    //======================初始条件离散=============================//
    VectorXd x0 = VectorXd(n_dofs_point);
    for(int j = 0; j <= POW2(n);j++)
    {
        for(int i = 0;i <= POW2(n);i++)
        {
            double point[2] = {i * mesh -> x_h(),j * mesh->y_h()};
            x0[i + (POW2(n) + 1) * j] = u0(point);
        }
    }
    solution_list.push_back(x0);
    //================全离散格式并迭代求解 ====================================//
    SpMat A_c(n_dofs_point,n_dofs_point);
    VectorXd b_c(n_dofs_point);
    Eigen::ConjugateGradient<Eigen::SparseMatrix<double> > Solver_sparse;
    for(int m = 1;m <= n_t ;m++)
    {
        A_c = M *(1.0 /  d_t) + theta * A;
        
        b_c = theta * rhs_list[m] + (1 - theta) * rhs_list[m - 1] + (M *(1.0 /  d_t) - (1 - theta) * A) * solution_list[m - 1];
        
        for(auto k:mesh->Boundary())
        {
            Dofs<2> bnd_point = mesh ->DofsofIndex(k);
            Real bnd_value = u(*bnd_point,double(d_t * m));
            b_c[k] = bnd_value * A_c.coeffRef(k,k);
            for(Eigen::SparseMatrix<Real>::InnerIterator it(A_c,k);it;++it)
            {
                int row = it.row();
                if(row == k)
                    continue;
                b_c[row] -= A_c.coeffRef(k,row)* bnd_value;
                A_c.coeffRef(k,row) = 0.0;
                A_c.coeffRef(row,k) = 0.0;
            }
        }
        Solver_sparse.setTolerance(1e-12);
	    Solver_sparse.compute(A_c);
        solution_list.push_back(Solver_sparse.solve(b_c));
    }
    //===============输出成VTK等格式进行可视化===============//
    for(int m = 0; m <=n_t;m++)
    {
        Surface<2> S(mesh,element,solution_list[m]);
        std::string s("surface_heat");
        s+= std::to_string(m);
        S.WriteVTKData(s);
    }
    //==============计算最终误差 ==========================//
    Error<2> E(mesh,element,u1,solution_list[n_t]);
    std::cout << "L2 error is :"  << E.L2Norm() << std::endl;
    //===========d_t = d_h 时，有二阶精度 ==================//
    delete mesh;
    delete element;
    delete Domain;
    return 0;
}