#include <fstream>
#include <iomanip>
#include <limits>

#include <multimeshvectorfield.hpp>
#include <dxf_solid.hpp>
#include <gtkplotter.hpp>
#include <geomplotter.hpp>
#include <geometry.hpp>
#include <func_solid.hpp>
#include <epot_bicgstabsolver.hpp>
#include <epot_gssolver.hpp>
#include <epot_efield.hpp>
#include <error.hpp>
#include <ibsimu.hpp>
#include <mydxffile.hpp>
#include <trajectorydiagnostics.hpp>
#include <particledatabase.hpp>
#include <particlediagplotter.hpp>
#include <transformation.hpp>
#include <compmath.hpp>

#include "scalarfield2d.hpp"


using namespace std;


double Vpull = 0.0;
double Veinzel = 0.0;
double Vbias = 0.0;
double J = 0.0;
double eir = 0.0;
double Ef = 0.0;
double Et = 0.0;
double E0 = 0.0;
double Tt = 0.0;
double rhoff = 0.0;
double epfac = 0.0;
double comp = 0.0;
std::string outdir = "";

double sc_alpha = 0.50;
int32_t Nrounds = 10;
double Nperh = 500;
const double r0 = 2.5e-3;


const char *file1 = "bfield.dat";
const char *file2 = "bfield_fine.dat";


MeshScalarField *simu_rough( int argc, char **argv )
{
    ibsimu.message(1) << "----------------------------\n";
    ibsimu.message(1) << "ROUGH SIMULATION\n";
    ibsimu.message(1) << "----------------------------\n";

    double start = -1.0e-3;
    double h = 2e-4;
    double sizereq[3] = {  24.0e-3, 
                           24.0e-3, 
                          168.5e-3-start };
    Int3D meshsize( (int)floor(sizereq[0]/h)+1,
                    (int)floor(sizereq[1]/h)+1,
                    (int)floor(sizereq[2]/h)+1 );
    Vec3D origo( -0.5*(meshsize[0]-1)*h,
                 -0.5*(meshsize[1]-1)*h,
                 start );
    Geometry geom( MODE_3D, meshsize, origo, h );

    MyDXFFile *dxffile = new MyDXFFile( "geom.dxf" );
    MyDXFEntities *e = dxffile->get_entities();
    MyDXFEntitySelection *sel = e->selection_all();
    e->scale( sel, dxffile, 1.0e-3 );

    DXFSolid *s1 = new DXFSolid( dxffile, "plasmae" );
    s1->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 7, s1 );
    DXFSolid *s3 = new DXFSolid( dxffile, "puller" );
    s3->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 8, s3 );
    DXFSolid *s4 = new DXFSolid( dxffile, "einzel1" );
    s4->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 9, s4 );
    DXFSolid *s5 = new DXFSolid( dxffile, "gnd" );
    s5->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 10, s5 );

    geom.set_boundary(  1,  Bound(BOUND_NEUMANN,     0.0) );
    geom.set_boundary(  2,  Bound(BOUND_NEUMANN,     0.0) );
    geom.set_boundary(  3,  Bound(BOUND_NEUMANN,     0.0) );
    geom.set_boundary(  4,  Bound(BOUND_NEUMANN,     0.0) );
    geom.set_boundary(  5,  Bound(BOUND_DIRICHLET,   0.0) );
    geom.set_boundary(  6,  Bound(BOUND_DIRICHLET, 1.0e3*Vbias) );

    geom.set_boundary(  7,  Bound(BOUND_DIRICHLET,   0.0) );
    geom.set_boundary(  8,  Bound(BOUND_DIRICHLET, 1.0e3*Vpull) );
    geom.set_boundary(  9,  Bound(BOUND_DIRICHLET, 1.0e3*Vbias-1.0e3*Veinzel) );
    geom.set_boundary( 10,  Bound(BOUND_DIRICHLET, 1.0e3*Vbias) );
    geom.build_mesh();

    EpotField epot( geom );
    MeshScalarField scharge_ele( geom );
    MeshScalarField scharge( geom );
    MeshScalarField scharge_ave( geom );

    EpotBiCGSTABSolver solver( geom );
    InitialPlasma initp( AXIS_Z, 0.1e-3 );
    solver.set_nsimp_initial_plasma( &initp );

    // Define magnetic field
    bool fout[3] = {true, true, true};
    MultiMeshVectorField bfield( MODE_3D, fout, 1.0e-3, 1.0, file1 );
    bfield.add_mesh( 1.0e-3, 1.0, file2 );
    field_extrpl_e bfldextrpl[6] = { FIELD_ZERO, FIELD_ZERO, 
				     FIELD_ZERO, FIELD_ZERO, 
				     FIELD_ZERO, FIELD_ZERO };
    bfield.set_extrapolation( bfldextrpl );

    EpotEfield efield( epot );
    field_extrpl_e efldextrpl[6] = { FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE, 
				     FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE, 
				     FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE };
    efield.set_extrapolation( efldextrpl );

    ParticleDataBase3D pdb( geom );
    pdb.set_max_steps( 1000 );
    bool pmirror[6] = { false, false, false, false, false, false };
    pdb.set_mirror( pmirror );
    pdb.set_polyint( true );
    
    NPlasmaBfieldSuppression psup( epot, 10.0 );
    pdb.set_bfield_suppression( &psup );

    double rhotot = 0.0;
    for( int32_t a = 0; a < Nrounds; a++ ) {

	if( a == 1 ) {
            std::vector<double> Ei, rhoi;
	    Ei.push_back( Et );
            rhoi.push_back( (1.0-rhoff)*rhotot );
            solver.set_nsimp_plasma( rhoff*rhotot, Ef, rhoi, Ei );
        } 

	solver.solve( epot, scharge_ave );
	efield.recalculate();

	uint32_t N = Nperh*M_PI*r0*r0/(h*h);
        pdb.clear(); 
	pdb.add_cylindrical_beam_with_energy( N, J*eir, -1.0, 1.0/1836.15, 
					      E0, 0.0, Tt, 
					      Vec3D(0,0,start), // center
					      Vec3D(1,0,0), // dir1
                                              Vec3D(0,1,0), // dir2
					      r0 );
        pdb.iterate_trajectories( scharge_ele, efield, bfield );
	rhotot = epfac*pdb.get_rhosum();

        pdb.clear(); 
	pdb.add_cylindrical_beam_with_energy( N, J, -1.0, 1.0, 
					      E0, 0.0, Tt, 
					      Vec3D(0,0,start), // center
					      Vec3D(1,0,0), // dir1
                                              Vec3D(0,1,0), // dir2
					      r0 );
        pdb.iterate_trajectories( scharge, efield, bfield );
	rhotot += pdb.get_rhosum();

	// Space charge compensation and sum
        for( uint32_t k = 0; k < scharge.size(2); k++ ) {       
            double z = scharge.origo(2)+k*scharge.h();
            for( uint32_t j = 0; j < scharge.size(1); j++ ) {   
                for( uint32_t i = 0; i < scharge.size(0); i++ ) {       
                    if( z >= 90.0e-3 )
                        scharge(i,j,k) = comp*scharge(i,j,k) + scharge_ele(i,j,k);
                    else if( z <= 5.0e-3 )
                        scharge(i,j,k) = scharge(i,j,k) + epfac*scharge_ele(i,j,k);
		    else
			scharge(i,j,k) = scharge(i,j,k) + scharge_ele(i,j,k);
                }
            }
        }

	// Average space charge
        if( a == 0 ) {
            scharge_ave = scharge;
        } else {
            double sc_beta = 1.0-sc_alpha;
            uint32_t nodecount = scharge.nodecount();
            for( uint32_t b = 0; b < nodecount; b++ ) {
                scharge_ave(b) = sc_alpha*scharge(b) + sc_beta*scharge_ave(b);
            }
        }

	// Trajectory diagnostics
	TrajectoryDiagnosticData tdata;
	std::vector<trajectory_diagnostic_e> diagnostics;
	diagnostics.push_back( DIAG_X );
	diagnostics.push_back( DIAG_XP );
	pdb.trajectories_at_plane( tdata, AXIS_Z, geom.max(2)-geom.h(), diagnostics );
	Emittance emit( tdata(0).data(), tdata(1).data() );	

	// Output
	string dout_fn = outdir + "dout_rough.dat";
	ofstream dout( dout_fn.c_str(), ios_base::app );
	dout << emit.alpha() << " "
	     << emit.beta() << " "
	     << emit.epsilon() << "\n";
	dout.close();

	if( false ) {
	    MeshScalarField tdens( geom );
	    pdb.build_trajectory_density_field( tdens );

	    GTKPlotter plotter( &argc, &argv );
	    plotter.set_geometry( &geom );
	    plotter.set_epot( &epot );
	    plotter.set_bfield( &bfield );
	    plotter.set_efield( &efield );
	    plotter.set_scharge( &scharge );
	    plotter.set_trajdens( &tdens );
	    plotter.set_particledatabase( &pdb );
	    plotter.new_geometry_plot_window();
	    plotter.run();
	}
    }

    if( false ) {
	MeshScalarField tdens( geom );
	pdb.build_trajectory_density_field( tdens );

	GTKPlotter plotter( &argc, &argv );
	plotter.set_geometry( &geom );
	plotter.set_epot( &epot );
	plotter.set_bfield( &bfield );
	plotter.set_efield( &efield );
	plotter.set_scharge( &scharge );
	plotter.set_trajdens( &tdens );
	plotter.set_particledatabase( &pdb );
	plotter.new_geometry_plot_window();
	plotter.run();
    }

    geom.save( outdir + "rough_geom.dat" );
    epot.save( outdir + "rough_epot.dat" );
    pdb.save( outdir + "rough_pdb.dat" );

    GeomPlotter geomplotter( geom );
    geomplotter.set_epot( &epot );
    geomplotter.set_particle_database( &pdb );
    geomplotter.set_particle_div( 11 );
    geomplotter.set_size( 1024, 768 );
    geomplotter.set_font_size( 20 );
    std::vector<double> eqlines;
    eqlines.push_back( 1.0 );
    eqlines.push_back( 2.0 );
    eqlines.push_back( 4.0 );
    geomplotter.set_eqlines_manual( eqlines );
    geomplotter.set_view( VIEW_ZX, -1 );
    geomplotter.plot_png( outdir + "geom_rough_zx.png" );
    geomplotter.set_view( VIEW_ZY, -1 );
    geomplotter.plot_png( outdir + "geom_rough_zy.png" );

    ParticleDiagPlotter pplotter1( geom, pdb, AXIS_Z, geom.max(2)-geom.h(), 
                                   PARTICLE_DIAG_PLOT_SCATTER,
                                   DIAG_X, DIAG_XP );
    pplotter1.set_size( 800, 600 );
    pplotter1.set_font_size( 16 );
    pplotter1.set_ranges( -0.03, -0.060, 0.03, 0.060 );
    pplotter1.set_emittance_ellipse( true );
    pplotter1.set_histogram_n( 151 );
    pplotter1.set_histogram_m( 151 );
    pplotter1.set_colormap_interpolation( INTERPOLATION_CLOSEST );
    pplotter1.plot_png( outdir + "emit_rough_xxp.png" );
    pplotter1.export_data( outdir + "emit_rough_xxp.txt" );

    ParticleDiagPlotter pplotter2( geom, pdb, AXIS_Z, geom.max(2)-geom.h(), 
                                   PARTICLE_DIAG_PLOT_SCATTER,
                                   DIAG_Y, DIAG_YP );
    pplotter2.set_size( 800, 600 );
    pplotter2.set_font_size( 16 );
    pplotter2.set_ranges( -0.03, -0.060, 0.03, 0.060 );
    pplotter2.set_emittance_ellipse( true );
    pplotter2.set_histogram_n( 151 );
    pplotter2.set_histogram_m( 151 );
    pplotter2.set_colormap_interpolation( INTERPOLATION_CLOSEST );
    pplotter2.plot_png( outdir + "emit_rough_yyp.png" );
    pplotter2.export_data( outdir + "emit_rough_yyp.txt" );

    // Return a copy of epot
    MeshScalarField *epot2 = new MeshScalarField( epot );
    return( epot2 );
}


class BoundPot : public CallbackFunctorD_V {

    const MeshScalarField &_field;

public:

    BoundPot( const MeshScalarField &field ) : _field(field) {}

    ~BoundPot() {}

    virtual double operator()( const Vec3D &x ) const {
	double pot = _field( x );
	return( pot );
    }

};


void simu_fine( int argc, char **argv, MeshScalarField &roughpot,
		std::vector<Particle3D> &part, uint32_t range[6] )
{
    ibsimu.message(1) << "----------------------------\n";
    ibsimu.message(1) << "FINE SIMULATION\n";
    ibsimu.message(1) << "----------------------------\n";

    // Define fine geometry, rough grid planes are on fine planes
    int coef = 3;
    double h = roughpot.h()/((double)coef);
    double start = -1.0e-3;
    Vec3D sizereq( 12.0e-3, 
		   12.0e-3, 
		   10.0e-3-start );
    Vec3D origoreq( -0.5*sizereq(0),
		    -0.5*sizereq(1),
		    start );
    Int3D ogrid( floor( (origoreq(0)-roughpot.origo(0))*roughpot.div_h() + 0.5 ),
		 floor( (origoreq(1)-roughpot.origo(1))*roughpot.div_h() + 0.5 ),
		 floor( (origoreq(2)-roughpot.origo(2))*roughpot.div_h() + 0.5 ) );
    Vec3D origo( roughpot.origo(0)+roughpot.h()*ogrid(0),
		 roughpot.origo(1)+roughpot.h()*ogrid(1),
		 roughpot.origo(2)+roughpot.h()*ogrid(2) );
    Int3D cellsize_rough( floor(sizereq(0)/roughpot.h()),
			  floor(sizereq(1)/roughpot.h()),
			  floor(sizereq(2)/roughpot.h()) );
    Int3D gridsize( coef*cellsize_rough(0)+1,
		    coef*cellsize_rough(1)+1,
		    coef*cellsize_rough(2)+1 );
    Geometry geom( MODE_3D, gridsize, origo, h );

    MyDXFFile *dxffile = new MyDXFFile( "V4_pellis.dxf" );
    MyDXFEntities *e = dxffile->get_entities();
    MyDXFEntitySelection *sel = e->selection_all();
    e->scale( sel, dxffile, 1.0e-3 );

    DXFSolid *s1 = new DXFSolid( dxffile, "plasmae" );
    s1->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 7, s1 );
    DXFSolid *s3 = new DXFSolid( dxffile, "puller" );
    s3->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 8, s3 );
    DXFSolid *s4 = new DXFSolid( dxffile, "einzel1" );
    s4->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 9, s4 );
    DXFSolid *s5 = new DXFSolid( dxffile, "gnd" );
    s5->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 10, s5 );

    BoundPot bound( roughpot );
    geom.set_boundary(  1,  Bound(BOUND_DIRICHLET, &bound) );
    geom.set_boundary(  2,  Bound(BOUND_DIRICHLET, &bound) );
    geom.set_boundary(  3,  Bound(BOUND_DIRICHLET, &bound) );
    geom.set_boundary(  4,  Bound(BOUND_DIRICHLET, &bound) );
    geom.set_boundary(  5,  Bound(BOUND_DIRICHLET,    0.0) );
    geom.set_boundary(  6,  Bound(BOUND_DIRICHLET, &bound) );

    geom.set_boundary(  7,  Bound(BOUND_DIRICHLET,    0.0) );
    geom.set_boundary(  8,  Bound(BOUND_DIRICHLET,  1.0e3*Vpull) );
    geom.set_boundary(  9,  Bound(BOUND_DIRICHLET, 1.0e3*Vbias-1.0e3*Veinzel) );
    geom.set_boundary( 10,  Bound(BOUND_DIRICHLET,  1.0e3*Vbias) );
    geom.build_mesh();

    // Use rough solution as a starting guess
    EpotField epot( geom );
    for( uint32_t k = 0; k < epot.size(2); k++ ) {
	double z = epot.origo(2) + k*epot.h();
	for( uint32_t j = 0; j < epot.size(1); j++ ) {
	    double y = epot.origo(1) + j*epot.h();
	    for( uint32_t i = 0; i < epot.size(0); i++ ) {
		double x = epot.origo(0) + i*epot.h();
		epot(i,j,k) = roughpot( Vec3D(x,y,z) );
	    }
	}
    }

    MeshScalarField scharge( geom );
    MeshScalarField scharge_ele( geom );
    MeshScalarField scharge_ave( geom );

    EpotBiCGSTABSolver solver( geom );

    // Define magnetic field
    bool fout[3] = {true, true, true};
    MultiMeshVectorField bfield( MODE_3D, fout, 1.0e-3, 1.0, file1 );
    bfield.add_mesh( 1.0e-3, 1.0, file2 );
    field_extrpl_e bfldextrpl[6] = { FIELD_ZERO, FIELD_ZERO, 
				     FIELD_ZERO, FIELD_ZERO, 
				     FIELD_ZERO, FIELD_ZERO };
    bfield.set_extrapolation( bfldextrpl );

    EpotEfield efield( epot );
    field_extrpl_e efldextrpl[6] = { FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE, 
				     FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE, 
				     FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE };
    efield.set_extrapolation( efldextrpl );

    ParticleDataBase3D pdb( geom );
    pdb.set_max_steps( 1000 );
    bool pmirror[6] = { false, false, false, false, false, false };
    pdb.set_mirror( pmirror );
    pdb.set_polyint( true );
    
    NPlasmaBfieldSuppression psup( epot, 10.0 );
    pdb.set_bfield_suppression( &psup );

    double rhotot = 0.0;
    for( int32_t a = 0; a < Nrounds; a++ ) {
	
	if( a == 1 ) {
            std::vector<double> Ei, rhoi;
	    Ei.push_back( Et );
            rhoi.push_back( (1.0-rhoff)*rhotot );
            solver.set_nsimp_plasma( rhoff*rhotot, Ef, rhoi, Ei );
	}
	if( a > 0 ) {
	    solver.solve( epot, scharge_ave );
	    efield.recalculate();
	}

	uint32_t N = Nperh*M_PI*r0*r0/(h*h);
        pdb.clear(); 
	pdb.add_cylindrical_beam_with_energy( N, J*eir, -1.0, 1.0/1836.15, 
					      E0, 0.0, Tt, 
					      Vec3D(0,0,geom.origo(2)), // center
					      Vec3D(1,0,0), // dir1
                                              Vec3D(0,1,0), // dir2
					      r0 );
        pdb.iterate_trajectories( scharge_ele, efield, bfield );
	rhotot = epfac*pdb.get_rhosum();

        pdb.clear();
	pdb.add_cylindrical_beam_with_energy( N, J, -1.0, 1.0, 
					      E0, 0.0, Tt, 
					      Vec3D(0,0,geom.origo(2)), // center
					      Vec3D(1,0,0), // dir1
                                              Vec3D(0,1,0), // dir2
					      r0 );
        pdb.iterate_trajectories( scharge, efield, bfield );
	rhotot += pdb.get_rhosum();

	// Space charge compensation and sum
        for( uint32_t k = 0; k < scharge.size(2); k++ ) {       
            double z = scharge.origo(2)+k*scharge.h();
            for( uint32_t j = 0; j < scharge.size(1); j++ ) {   
                for( uint32_t i = 0; i < scharge.size(0); i++ ) {       
                    if( z >= 90.0e-3 )
                        scharge(i,j,k) = comp*scharge(i,j,k) + scharge_ele(i,j,k);
                    else if( z <= 5.0e-3 )
                        scharge(i,j,k) = scharge(i,j,k) + epfac*scharge_ele(i,j,k);
		    else
                        scharge(i,j,k) = scharge(i,j,k) + scharge_ele(i,j,k);
                }
            }
        }

	// Average space charge
        if( a == 0 ) {
            scharge_ave = scharge;
        } else {
            double sc_beta = 1.0-sc_alpha;
            uint32_t nodecount = scharge.nodecount();
            for( uint32_t b = 0; b < nodecount; b++ ) {
                scharge_ave(b) = sc_alpha*scharge(b) + sc_beta*scharge_ave(b);
            }
        }

	// Trajectory diagnostics
	TrajectoryDiagnosticData tdata;
	std::vector<trajectory_diagnostic_e> diagnostics;
	diagnostics.push_back( DIAG_X );
	diagnostics.push_back( DIAG_XP );
	pdb.trajectories_at_plane( tdata, AXIS_Z, geom.max(2)-geom.h(), diagnostics );
	Emittance emit( tdata(0).data(), tdata(1).data() );	

	// Output
	string dout_fn = outdir + "dout_fine.dat";
	ofstream dout( dout_fn.c_str(), ios_base::app );
	dout << emit.alpha() << " "
	     << emit.beta() << " "
	     << emit.epsilon() << "\n";
	dout.close();

	if( false ) {
	    MeshScalarField tdens( geom );
	    pdb.build_trajectory_density_field( tdens );

	    GTKPlotter plotter( &argc, &argv );
	    plotter.set_geometry( &geom );
	    plotter.set_epot( &epot );
	    plotter.set_bfield( &bfield );
	    plotter.set_efield( &efield );
	    plotter.set_scharge( &scharge );
	    plotter.set_trajdens( &tdens );
	    plotter.set_particledatabase( &pdb );
	    plotter.new_geometry_plot_window();
	    plotter.run();
	}
    }

    if( false ) {
	MeshScalarField tdens( geom );
	pdb.build_trajectory_density_field( tdens );

	GTKPlotter plotter( &argc, &argv );
	plotter.set_geometry( &geom );
	plotter.set_epot( &epot );
	plotter.set_bfield( &bfield );
	plotter.set_efield( &efield );
	plotter.set_scharge( &scharge );
	plotter.set_trajdens( &tdens );
	plotter.set_particledatabase( &pdb );
	plotter.new_geometry_plot_window();
	plotter.run();
    }

    geom.save( outdir + "fine_geom.dat" );
    epot.save( outdir + "fine_epot.dat" );
    pdb.save( outdir + "fine_pdb.dat" );

    GeomPlotter geomplotter( geom );
    geomplotter.set_epot( &epot );
    geomplotter.set_particle_database( &pdb );
    geomplotter.set_particle_div( 11 );
    geomplotter.set_size( 1024, 768 );
    geomplotter.set_font_size( 20 );
    std::vector<double> eqlines;
    eqlines.push_back( 1.0 );
    eqlines.push_back( 2.0 );
    eqlines.push_back( 4.0 );
    geomplotter.set_eqlines_manual( eqlines );
    geomplotter.set_view( VIEW_ZX, -1 );
    geomplotter.plot_png( outdir + "geom_fine_zx.png" );
    geomplotter.set_view( VIEW_ZY, -1 );
    geomplotter.plot_png( outdir + "geom_fine_zy.png" );

    ParticleDiagPlotter pplotter1( geom, pdb, AXIS_Z, geom.max(2)-geom.h(), 
                                   PARTICLE_DIAG_PLOT_SCATTER,
                                   DIAG_X, DIAG_XP );
    pplotter1.set_size( 800, 600 );
    pplotter1.set_font_size( 16 );
    pplotter1.set_ranges( -0.03, -0.060, 0.03, 0.060 );
    pplotter1.set_emittance_ellipse( true );
    pplotter1.set_histogram_n( 151 );
    pplotter1.set_histogram_m( 151 );
    pplotter1.set_colormap_interpolation( INTERPOLATION_CLOSEST );
    pplotter1.plot_png( outdir + "emit_fine_xxp.png" );
    pplotter1.export_data( outdir + "emit_fine_xxp.txt" );

    ParticleDiagPlotter pplotter2( geom, pdb, AXIS_Z, geom.max(2)-geom.h(), 
                                   PARTICLE_DIAG_PLOT_SCATTER,
                                   DIAG_Y, DIAG_YP );
    pplotter2.set_size( 800, 600 );
    pplotter2.set_font_size( 16 );
    pplotter2.set_ranges( -0.03, -0.060, 0.03, 0.060 );
    pplotter2.set_emittance_ellipse( true );
    pplotter2.set_histogram_n( 151 );
    pplotter2.set_histogram_m( 151 );
    pplotter2.set_colormap_interpolation( INTERPOLATION_CLOSEST );
    pplotter2.plot_png( outdir + "emit_fine_yyp.png" );
    pplotter2.export_data( outdir + "emit_fine_yyp.txt" );

    // Export particles on a mesh level of rough simulation

    // Load epot to roughpot. For grid nodes of roughpot inside 
    // the fine simulation - no extrapolation. 
    range[0] = (uint32_t)ceil( (epot.origo(0)-roughpot.origo(0))*roughpot.div_h() );
    range[1] = (uint32_t)floor( (epot.max(0)-roughpot.origo(0))*roughpot.div_h() );
    range[2] = (uint32_t)ceil( (epot.origo(1)-roughpot.origo(1))*roughpot.div_h() );
    range[3] = (uint32_t)floor( (epot.max(1)-roughpot.origo(1))*roughpot.div_h() );
    range[4] = (uint32_t)ceil( (epot.origo(2)-roughpot.origo(2))*roughpot.div_h() );
    range[5] = (uint32_t)floor( (epot.max(2)-roughpot.origo(2))*roughpot.div_h() );

    for( uint32_t k = range[4]; k < range[5]; k++ ) {
	double z = roughpot.origo(2) + k*roughpot.h();
	for( uint32_t j = range[2]; j < range[3]; j++ ) {
	    double y = roughpot.origo(1) + j*roughpot.h();
	    for( uint32_t i = range[0]; i < range[1]; i++ ) {
		double x = roughpot.origo(0) + i*roughpot.h();
		roughpot(i,j,k) = epot( Vec3D(x,y,z) );
	    }
	}
    }

    // Limits for export in rough grid
    range[0] += 3;
    range[1] -= 3;
    range[2] += 3;
    range[3] -= 3;
    range[5] -= 10;

    // Particles
    double zmax = roughpot.origo(2)+range[5]*roughpot.h();
    pdb.trajectories_at_plane( part, AXIS_Z, zmax );
}


class ForceFine : public CallbackFunctorD_V {
    
    const MeshScalarField &_pot;
    double                 _frange[6];

public:

    ForceFine( const MeshScalarField &roughpot, uint32_t range[6] ) 
	: _pot(roughpot) {

	// Extend by half grid cell to ensure that callback coordinate
	// fits in.
	_frange[0] = roughpot.origo(0) + (range[0]-0.5)*roughpot.h();
	_frange[1] = roughpot.origo(0) + (range[1]+0.5)*roughpot.h();
	_frange[2] = roughpot.origo(1) + (range[2]-0.5)*roughpot.h();
	_frange[3] = roughpot.origo(1) + (range[3]+0.5)*roughpot.h();
	_frange[4] = roughpot.origo(2) + (range[4]-0.5)*roughpot.h();
	_frange[5] = roughpot.origo(2) + (range[5]+0.5)*roughpot.h();
    }

    ~ForceFine() {}

    virtual double operator()( const Vec3D &x ) const {
	
	for( uint32_t a = 0; a < 3; a++ ) {
	    if( x[a] < _frange[2*a+0] || x[a] > _frange[2*a+1] )
		return( std::numeric_limits<double>::quiet_NaN() );
	}

	return( _pot(x) );
    }
};


void simu_rough2( int argc, char **argv, MeshScalarField &roughpot,
		  std::vector<Particle3D> &part, uint32_t range[6] )
{
    ibsimu.message(1) << "----------------------------\n";
    ibsimu.message(1) << "ROUGH II SIMULATION\n";
    ibsimu.message(1) << "----------------------------\n";

    double start = -1.0e-3;
    double h = 2e-4;
    double sizereq[3] = {  24.0e-3, 
                           24.0e-3, 
                          168.5e-3-start };
    Int3D meshsize( (int)floor(sizereq[0]/h)+1,
                    (int)floor(sizereq[1]/h)+1,
                    (int)floor(sizereq[2]/h)+1 );
    Vec3D origo( -0.5*(meshsize[0]-1)*h,
                 -0.5*(meshsize[1]-1)*h,
                 start );
    Geometry geom( MODE_3D, meshsize, origo, h );

    MyDXFFile *dxffile = new MyDXFFile( "V4_pellis.dxf" );
    MyDXFEntities *e = dxffile->get_entities();
    MyDXFEntitySelection *sel = e->selection_all();
    e->scale( sel, dxffile, 1.0e-3 );

    DXFSolid *s1 = new DXFSolid( dxffile, "plasmae" );
    s1->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 7, s1 );
    DXFSolid *s3 = new DXFSolid( dxffile, "puller" );
    s3->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 8, s3 );
    DXFSolid *s4 = new DXFSolid( dxffile, "einzel1" );
    s4->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 9, s4 );
    DXFSolid *s5 = new DXFSolid( dxffile, "gnd" );
    s5->define_2x3_mapping( DXFSolid::rotz );
    geom.set_solid( 10, s5 );

    geom.set_boundary(  1,  Bound(BOUND_NEUMANN,     0.0) );
    geom.set_boundary(  2,  Bound(BOUND_NEUMANN,     0.0) );
    geom.set_boundary(  3,  Bound(BOUND_NEUMANN,     0.0) );
    geom.set_boundary(  4,  Bound(BOUND_NEUMANN,     0.0) );
    geom.set_boundary(  5,  Bound(BOUND_DIRICHLET,   0.0) );
    geom.set_boundary(  6,  Bound(BOUND_DIRICHLET, 1.0e3*Vbias) );

    geom.set_boundary(  7,  Bound(BOUND_DIRICHLET,   0.0) );
    geom.set_boundary(  8,  Bound(BOUND_DIRICHLET, 1.0e3*Vpull) );
    geom.set_boundary(  9,  Bound(BOUND_DIRICHLET, 1.0e3*Vbias-1.0e3*Veinzel) );
    geom.set_boundary( 10,  Bound(BOUND_DIRICHLET, 1.0e3*Vbias) );
    geom.build_mesh();

    EpotField epot( geom );
    // Starting guess from first rough solution
    for( uint32_t k = 0; k < epot.size(2); k++ ) {
	for( uint32_t j = 0; j < epot.size(1); j++ ) {
	    for( uint32_t i = 0; i < epot.size(0); i++ ) {
		epot(i,j,k) = roughpot(i,j,k);
	    }
	}
    }

    MeshScalarField scharge( geom );
    MeshScalarField scharge_ave( geom );

    EpotBiCGSTABSolver solver( geom );
    ForceFine force_fine( roughpot, range );
    solver.set_forced_potential_volume( &force_fine );

    // Define magnetic field
    bool fout[3] = {true, true, true};
    MultiMeshVectorField bfield( MODE_3D, fout, 1.0e-3, 1.0, file1 );
    bfield.add_mesh( 1.0e-3, 1.0, file2 );
    field_extrpl_e bfldextrpl[6] = { FIELD_ZERO, FIELD_ZERO, 
				     FIELD_ZERO, FIELD_ZERO, 
				     FIELD_ZERO, FIELD_ZERO };
    bfield.set_extrapolation( bfldextrpl );

    EpotEfield efield( epot );
    field_extrpl_e efldextrpl[6] = { FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE, 
				     FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE, 
				     FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE };
    efield.set_extrapolation( efldextrpl );

    ParticleDataBase3D pdb( geom );
    pdb.set_max_steps( 1000 );
    bool pmirror[6] = { false, false, false, false, false, false };
    pdb.set_mirror( pmirror );
    pdb.set_polyint( true );
    
    NPlasmaBfieldSuppression psup( epot, 10.0 );
    pdb.set_bfield_suppression( &psup );

    for( int32_t a = 0; a < Nrounds; a++ ) {

	solver.solve( epot, scharge_ave );
	efield.recalculate();

        pdb.clear(); 
	// Add particles from fine simulation
	for( size_t b = 0; b < part.size(); b++ )
            pdb.add_particle( part[b] );
        pdb.iterate_trajectories( scharge, efield, bfield );

	// Space charge compensation
        for( uint32_t k = 0; k < scharge.size(2); k++ ) {       
            double z = scharge.origo(2)+k*scharge.h();
            for( uint32_t j = 0; j < scharge.size(1); j++ ) {   
                //double y = scharge.origo(1)+j*scharge.h();
                for( uint32_t i = 0; i < scharge.size(0); i++ ) {       
                    //double x = scharge.origo(0)+i*scharge.h();
                    if( z >= 90.0e-3 )
                        scharge(i,j,k) *= comp;
                }
            }
        }

	// Average space charge
        if( a == 0 ) {
            scharge_ave = scharge;
        } else {
            double sc_beta = 1.0-sc_alpha;
            uint32_t nodecount = scharge.nodecount();
            for( uint32_t b = 0; b < nodecount; b++ ) {
                scharge_ave(b) = sc_alpha*scharge(b) + sc_beta*scharge_ave(b);
            }
        }

	// Trajectory diagnostics
	TrajectoryDiagnosticData tdata;
	std::vector<trajectory_diagnostic_e> diagnostics;
	diagnostics.push_back( DIAG_X );
	diagnostics.push_back( DIAG_XP );
	pdb.trajectories_at_plane( tdata, AXIS_Z, geom.max(2)-geom.h(), diagnostics );
	Emittance emit( tdata(0).data(), tdata(1).data() );	

	// Output
	string dout_fn = outdir + "dout_rough2.dat";
	ofstream dout( dout_fn.c_str(), ios_base::app );
	dout << emit.alpha() << " "
	     << emit.beta() << " "
	     << emit.epsilon() << "\n";
	dout.close();

	if( false ) {
	    MeshScalarField tdens( geom );
	    pdb.build_trajectory_density_field( tdens );

	    GTKPlotter plotter( &argc, &argv );
	    plotter.set_geometry( &geom );
	    plotter.set_epot( &epot );
	    plotter.set_bfield( &bfield );
	    plotter.set_efield( &efield );
	    plotter.set_scharge( &scharge );
	    plotter.set_trajdens( &tdens );
	    plotter.set_particledatabase( &pdb );
	    plotter.new_geometry_plot_window();
	    plotter.run();
	}
    }

    if( false ) {
	MeshScalarField tdens( geom );
	pdb.build_trajectory_density_field( tdens );

	GTKPlotter plotter( &argc, &argv );
	plotter.set_geometry( &geom );
	plotter.set_epot( &epot );
	plotter.set_bfield( &bfield );
	plotter.set_efield( &efield );
	plotter.set_scharge( &scharge );
	plotter.set_trajdens( &tdens );
	plotter.set_particledatabase( &pdb );
	plotter.new_geometry_plot_window();
	plotter.run();
    }

    geom.save( outdir + "rough2_geom.dat" );
    epot.save( outdir + "rough2_epot.dat" );
    pdb.save( outdir + "rough2_pdb.dat" );

    GeomPlotter geomplotter( geom );
    geomplotter.set_epot( &epot );
    geomplotter.set_particle_database( &pdb );
    geomplotter.set_particle_div( 11 );
    geomplotter.set_size( 1024, 768 );
    geomplotter.set_font_size( 20 );
    std::vector<double> eqlines;
    eqlines.push_back( 1.0 );
    eqlines.push_back( 2.0 );
    eqlines.push_back( 4.0 );
    geomplotter.set_eqlines_manual( eqlines );
    geomplotter.set_view( VIEW_ZX, -1 );
    geomplotter.plot_png( outdir + "geom_rough2_zx.png" );
    geomplotter.set_view( VIEW_ZY, -1 );
    geomplotter.plot_png( outdir + "geom_rough2_zy.png" );

    ParticleDiagPlotter pplotter1( geom, pdb, AXIS_Z, geom.max(2)-geom.h(), 
                                   PARTICLE_DIAG_PLOT_SCATTER,
                                   DIAG_X, DIAG_XP );
    pplotter1.set_size( 800, 600 );
    pplotter1.set_font_size( 16 );
    pplotter1.set_ranges( -0.03, -0.060, 0.03, 0.060 );
    pplotter1.set_emittance_ellipse( true );
    pplotter1.set_histogram_n( 151 );
    pplotter1.set_histogram_m( 151 );
    pplotter1.set_colormap_interpolation( INTERPOLATION_CLOSEST );
    pplotter1.plot_png( outdir + "emit_rough2_xxp.png" );
    pplotter1.export_data( outdir + "emit_rough2_xxp.txt" );

    ParticleDiagPlotter pplotter2( geom, pdb, AXIS_Z, geom.max(2)-geom.h(), 
                                   PARTICLE_DIAG_PLOT_SCATTER,
                                   DIAG_Y, DIAG_YP );
    pplotter2.set_size( 800, 600 );
    pplotter2.set_font_size( 16 );
    pplotter2.set_ranges( -0.03, -0.060, 0.03, 0.060 );
    pplotter2.set_emittance_ellipse( true );
    pplotter2.set_histogram_n( 151 );
    pplotter2.set_histogram_m( 151 );
    pplotter2.set_colormap_interpolation( INTERPOLATION_CLOSEST );
    pplotter2.plot_png( outdir + "emit_rough2_yyp.png" );
    pplotter2.export_data( outdir + "emit_rough2_yyp.txt" );
}


int main( int argc, char **argv )
{
    if( argc <= 12 ) {
	std::cout << "Usage: simu Vpull Veinzel Vbias J eir Ef Et E0 Tt rhoff epfac comp\n";
	exit( 1 );
    }

    Vpull   = atof( argv[1] );
    Veinzel = atof( argv[2] );
    Vbias   = atof( argv[3] );
    J       = atof( argv[4] );
    eir     = atof( argv[5] );
    Ef      = atof( argv[6] );
    Et      = atof( argv[7] );
    E0      = atof( argv[8] );
    Tt      = atof( argv[9] );
    rhoff   = atof( argv[10] );
    epfac   = atof( argv[11] );
    comp    = atof( argv[12] );
    outdir  =
	to_string(Vpull) + "_" + 
	to_string(Veinzel) + "_" +
	to_string(Vbias) + "_" +
	to_string(J) + "_" +
	to_string(eir) + "_" +
	to_string(Ef) + "_" +
	to_string(Et) + "_" +
	to_string(E0) + "_" +
	to_string(Tt) + "_" +
	to_string(rhoff) + "_" +
	to_string(epfac) + "_" +
	to_string(comp);

    // Make directory
    std::string cmd = "mkdir " + outdir;
    FILE *fp = popen( cmd.c_str(), "r" );
    pclose( fp );

    outdir += "/";
    
    try {
        ibsimu.set_message_output( outdir + "ibsimu.txt" );
	ibsimu.set_message_threshold( MSG_VERBOSE, 1 );
	ibsimu.set_thread_count( 4 );

	//ScalarFieldOnBox *finebox = simu_rough( argc, argv );
	MeshScalarField *roughpot = simu_rough( argc, argv );
	std::vector<Particle3D> part;
	uint32_t range[6];
	simu_fine( argc, argv, *roughpot, part, range );
	simu_rough2( argc, argv, *roughpot, part, range );
    } catch( Error e ) {
        e.print_error_message( ibsimu.message(0) );
        exit( 1 );
    }

    return( 0 );
}