#include "epot_bicgstabsolver.hpp"
#include "geometry.hpp"
#include "func_solid.hpp"
#include "epot_efield.hpp"
#include "meshvectorfield.hpp"
#include "particledatabase.hpp"
#include "geomplotter.hpp"
#include "ibsimu.hpp"
#include "error.hpp"
#include "gtkplotter.hpp"



double h = 1e-3;

double r0 = 10e-3;
double fac = 1.2;
double len = 15e-3;
double gap = 5e-3;
double thickness = 5e-3;

double rbeam = 2e-3;
double nperh2 = 100;
double Npart = rbeam*rbeam/h/h*nperh2; // Number of particles absolute
double Q = 1;              // Charge (e)
double M = 40;             // Mass (u)
double I = 0.4e-3;           // Beam current (A)
double A = rbeam*rbeam*M_PI; // Plasma electrode area
double J = I/A;            // Current density
double E0 = 5e3;           // Beam energy (eV)

bool quad_ends( double x, double y, double z )
{
    double r = sqrt(x*x + y*y);
    return( r > r0 && ((z < -len-gap && z > -len-gap-thickness) ||
		       (z > len+gap && z < len+gap+thickness) ) );
}


bool quad_plus( double x, double y, double z )
{
    return( (sqrt(pow(x-(1+fac)*r0,2) + pow(y,2)) < r0 ||
	     sqrt(pow(x+(1+fac)*r0,2) + pow(y,2)) < r0) &&
	    z < len && z > -len );
}


bool quad_minus( double x, double y, double z )
{
    return( (sqrt(pow(y-(1+fac)*r0,2) + pow(x,2)) < r0 ||
	     sqrt(pow(y+(1+fac)*r0,2) + pow(x,2)) < r0) &&
	    z < len && z > -len );
}


void simu( int argc, char **argv )
{
    Vec3D origin( -4*r0, -4*r0, -(len+2*gap+thickness) );
    Vec3D sizereq( 8*r0, 8*r0, 2*(len+2*gap+thickness) );
    Int3D size( floor(sizereq[0]/h)+1,
                floor(sizereq[1]/h)+1,
                floor(sizereq[2]/h)+1 );
    Geometry geom( MODE_3D, size, origin, h );

    Solid *s1 = new FuncSolid( quad_ends );
    geom.set_solid( 7, s1 );
    Solid *s2 = new FuncSolid( quad_plus );
    geom.set_solid( 8, s2 );
    Solid *s3 = new FuncSolid( quad_minus );
    geom.set_solid( 9, s3 );
    
    geom.set_boundary( 1, Bound(BOUND_DIRICHLET,     0.0) );
    geom.set_boundary( 2, Bound(BOUND_DIRICHLET,     0.0) );
    geom.set_boundary( 3, Bound(BOUND_DIRICHLET,     0.0) );
    geom.set_boundary( 4, Bound(BOUND_DIRICHLET,     0.0) );
    
    geom.set_boundary( 7, Bound(BOUND_DIRICHLET,     0.0) );
    geom.set_boundary( 8, Bound(BOUND_DIRICHLET,   1.0e3) );
    geom.set_boundary( 9, Bound(BOUND_DIRICHLET,  -1.0e3) );
    geom.build_mesh();
    geom.build_surface();

    EpotField epot( geom );
    MeshScalarField scharge( geom );
    MeshVectorField bfield;
    EpotEfield efield( epot );
    field_extrpl_e efldextrpl[6] = { FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE, 
	                             FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE,
				     FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE };
    efield.set_extrapolation( efldextrpl );

    EpotBiCGSTABSolver solver( geom );

    ParticleDataBase3D pdb( geom );
    bool pmirror[6] = { false, false, false, false, false, false };
    pdb.set_mirror( pmirror );

    for( size_t i = 0; i < 10; i++ ) {
	solver.solve( epot, scharge );
	efield.recalculate();
	pdb.clear();
	pdb.add_cylindrical_beam_with_energy( Npart, J, Q, M, E0,
					      0.0, 0.0, 
					      Vec3D(0,0,-2*len), 
					      Vec3D(1,0,0),
					      Vec3D(0,1,0),
					      rbeam );
	pdb.iterate_trajectories( scharge, efield, bfield );
    }

    MeshScalarField tdens(geom);
    pdb.build_trajectory_density_field( tdens );
    
#ifdef GTK3
    GTKPlotter plotter( &argc, &argv );
    plotter.set_geometry( &geom );
    plotter.set_trajdens( &tdens );
    plotter.set_epot( &epot );
    plotter.set_scharge( &scharge );
    plotter.set_particledatabase( &pdb );
    plotter.new_geometry_plot_window();
    plotter.run();
#endif    
}


int main( int argc, char **argv )
{
    try {
	ibsimu.set_message_threshold( MSG_VERBOSE, 1 );
	ibsimu.set_thread_count( 4 );
        simu( argc, argv );
    } catch ( Error e ) {
        e.print_error_message( ibsimu.message( 0 ) );
        exit( 1 );
    }

    return( 0 );
}