/* JYFL Filament driven ion source for target sputtering
 * Ar+ beam
 * Zoomed to extraction to get closer to Debye length
 */

#include <fstream>
#include <iostream>
#include <ibsimu.hpp>
#include <error.hpp>
#include <geometry.hpp>
#include <dxf_solid.hpp>
#include <mydxffile.hpp>
#include <epot_field.hpp>
#include <epot_efield.hpp>
#include <meshvectorfield.hpp>
#include <epot_bicgstabsolver.hpp>
#include <gtkplotter.hpp>
#include <trajectorydiagnostics.hpp>



using namespace std;


/* Simulation parameters
 */
double Vgnd = -10e3;       // Gnd voltage (plasma at 0 V)
double Vpuller = -10e3;    // Puller voltage
double Veinzel = -2e3;     // Einzel voltage
double h = 0.2e-3;        // Mesh spacing (>5 mesh cells in plasma electrode radius)
double nperh = 1000;        // Number of particles per mesh cell
double r0 = 5e-3;          // Radius of emission
double rplasma = 2e-3;     // Radius of plasma electrode aperture
double Npart = r0/h*nperh; // Number of particles absolute
double Q = 1;              // Charge (e)
double M = 40;             // Mass (u)
double I = 0.4e-3;           // Extracted current (A)
double A = rplasma*rplasma*M_PI; // Plasma electrode area
double J = I/A;            // Emission current density
double Te = 2.0;           // Electron temperature (eV)
double E0 = 1.1*0.5*Te;    // Initial energy of ions at sheath edge (eV), >Bohm velocity, 10 % margin
double Tt = 0.5;           // Ion transverse temperature (eV)
double Up = 9.36;          // Plasma potential (V)
double xstart = -2.0e-3;   // Emission plane coordinate
double sc_alpha = 1.0;     // Space charge under-relaxation
uint32_t Niter = 30;

double m = M*MASS_U;
double q = Q*CHARGE_E;
double n_e = J/(q*sqrt(Te*CHARGE_E/m));
double lamb_d = sqrt(EPSILON0*Te*CHARGE_E/(n_e*q*q));

void simu( int *argc, char ***argv )
{
    ibsimu.message(1) << "Simulation parameters\n";
    ibsimu.message(1) << "Vgnd = " << Vgnd << " V\n";
    ibsimu.message(1) << "Vpuller = " << Vpuller << " V\n";
    ibsimu.message(1) << "Veinzel = " << Veinzel << " V\n";
    ibsimu.message(1) << "h = " << h << " m\n";
    ibsimu.message(1) << "rplasma = " << rplasma << " m\n";
    ibsimu.message(1) << "r0 = " << r0 << " m\n";
    ibsimu.message(1) << "J = " << J << " A/m2\n";
    ibsimu.message(1) << "M = " << M << " u\n";
    ibsimu.message(1) << "Q = " << Q << " e\n";
    ibsimu.message(1) << "n_e = " << n_e << " 1/m3\n";
    ibsimu.message(1) << "lamb_d = " << lamb_d << " m\n";
    ibsimu.message(1) << "h = " << h << " m\n";
    ibsimu.message(1) << "Te = " << Te << " eV\n";
    ibsimu.message(1) << "Tt = " << Tt << " eV\n";
    ibsimu.message(1) << "Up = " << Up << " V\n";
    ibsimu.message(1) << "E0 = " << E0 << " eV\n";
    ibsimu.message(1) << "xstart = " << xstart << " m\n";
    ibsimu.message(1) << "Niter = " << Niter << "\n";
    ibsimu.message(1) << "\n";

    /// Define simulation box
    Vec3D origin( xstart, 0, 0 );
    Vec3D sizereq( 320e-3-xstart, 50e-3, 0 );
    Int3D size( floor(sizereq[0]/h)+1,
		floor(sizereq[1]/h)+1,
		1 );
    Geometry geom( MODE_CYL, size, origin, h );

    // Read in DXF file
    MyDXFFile *dxffile = new MyDXFFile;
    dxffile->set_warning_level( 1 );
    dxffile->read( "../sputter.dxf" );

    // Set solids based on DXF layers
    DXFSolid *s1 = new DXFSolid( dxffile, "plasma" );
    s1->scale( 1e-3 );
    geom.set_solid( 7, s1 );
    DXFSolid *s2 = new DXFSolid( dxffile, "puller" );
    s2->scale( 1e-3 );
    geom.set_solid( 8, s2 );
    DXFSolid *s3 = new DXFSolid( dxffile, "gnd" );
    s3->scale( 1e-3 );
    geom.set_solid( 9, s3 );
    DXFSolid *s4 = new DXFSolid( dxffile, "einzel" );
    s4->scale( 1e-3 );
    geom.set_solid( 10, s4 );

    // Define boundary conditions
    geom.set_boundary(  1, Bound(BOUND_NEUMANN, 0.0) ); // xmin
    geom.set_boundary(  2, Bound(BOUND_NEUMANN, 0.0) ); // xmax
    geom.set_boundary(  3, Bound(BOUND_NEUMANN, 0.0) ); // rmin
    geom.set_boundary(  4, Bound(BOUND_NEUMANN, 0.0) ); // rmax

    geom.set_boundary(  7, Bound(BOUND_DIRICHLET, 0.0) );
    geom.set_boundary(  8, Bound(BOUND_DIRICHLET, Vpuller) );
    geom.set_boundary(  9, Bound(BOUND_DIRICHLET, Vgnd) );
    geom.set_boundary( 10, Bound(BOUND_DIRICHLET, Veinzel) );

    // Build geometry
    geom.build_mesh();

    // Initialize objects for calculation
    EpotField epot( geom );
    MeshScalarField scharge( geom );
    MeshScalarField scharge_ave( geom );
    MeshVectorField bfield;
    EpotEfield efield( epot );
    field_extrpl_e extrapl[6] = { FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE,
	FIELD_SYMMETRIC_POTENTIAL, FIELD_EXTRAPOLATE,
	FIELD_EXTRAPOLATE, FIELD_EXTRAPOLATE };
    efield.set_extrapolation( extrapl );

    // Initialize solver (default)
    EpotBiCGSTABSolver solver( geom );

    // Initial guess for plasma boundary at x=0.0
    InitialPlasma init_plasma( AXIS_X, 0.0 );
    solver.set_initial_plasma( Up, &init_plasma );

    // Initialize particle database
    ParticleDataBaseCyl pdb( geom );
    bool pmirror[6] = {false, false,
	false, false,
	false, false};
    pdb.set_mirror( pmirror );

    // Major loop
    for( uint32_t a = 0; a < Niter; a++ ) {

	ibsimu.message(1) << "\n";
	ibsimu.message(1) << "Major cycle " << a << " / " << Niter << "\n";
	ibsimu.message(1) << "-----------------------\n";

	// Set plasma model
	if( a == 1 ) {
	    double rhoe = pdb.get_rhosum();
	    solver.set_pexp_plasma( rhoe, Te, Up );
	}

	// Solve Poisson
	solver.solve( epot, scharge_ave );
	if( solver.get_iter() == 0 ) {
	    ibsimu.message(1) << "No iterations, breaking major cycle\n";
	    break;
	}
	efield.recalculate();
	
	// Calculate particles
	pdb.clear();
	pdb.add_2d_beam_with_energy( Npart, J, Q, M, E0, 0, Tt, 
				     geom.origo(0), 0,
				     geom.origo(0), r0 );
	pdb.iterate_trajectories( scharge, efield, bfield );
	
	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);
	    }
	}
	
	// Do trajectory diagnostics
	TrajectoryDiagnosticData tdata;
	vector<trajectory_diagnostic_e> diag;
	diag.push_back( DIAG_R );
	diag.push_back( DIAG_RP );
	diag.push_back( DIAG_CURR );
	pdb.trajectories_at_plane( tdata, AXIS_X, geom.max(0)-geom.h(), diag );
	tdata.mirror( AXIS_R, geom.origo(1) );
	Emittance emit( tdata(0).data(), tdata(1).data(), tdata(2).data() );

	// Output diagnostics into log
	ofstream dataout( "convergence.txt", ios_base::app );
	dataout << emit.alpha() << " "
		<< emit.beta() << " "
		<< emit.epsilon() << " "
		<< tdata.traj_size() << "\n";
	dataout.close();
    }

    // Save binary data
    geom.save("geom.dat");
    epot.save("epot.dat");
    pdb.save("pdb.dat");

    // Plot emittance as a function of distance
    ofstream emitout( "emit_vs_x.txt" );
    double step = (geom.max(0)-geom.origo(0))/1000.0;
    for( double x = 0.0; x < geom.max(0)-geom.h(); x += step ) {
	TrajectoryDiagnosticData tdata;
	vector<trajectory_diagnostic_e> diag;
	diag.push_back( DIAG_R );
	diag.push_back( DIAG_RP );
	diag.push_back( DIAG_CURR );
	pdb.trajectories_at_plane( tdata, AXIS_X, x, diag );
	tdata.mirror( AXIS_R, geom.origo(1) );
	Emittance emit( tdata(0).data(), tdata(1).data(), tdata(2).data() );
	emitout << x << " "
		<< emit.epsilon() << " "
		<< epot(Vec3D(x,0,0)) << "\n";
    }
    
    // Start plotter
    if( true ) {
	MeshScalarField tdens( geom );
	pdb.build_trajectory_density_field( tdens );

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


int main( int argc, char **argv )
{
    remove( "convergence.txt" );

    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) );
    }

    return( 0 );
}