// // Here are useful functions that can speed-up coding of your problem
#include "../extensions/wslda_utils.h"

/**
 * THIS FUNCTION IS CALLED AT THE BEGINNING OF SIMULATION.
 * After loading params array from input file, the parameters are processed by this routine.
 * @param params array of size MAX_USER_PARAMS with parameters from input file.
 * @param kF typical Fermi momentum scale of the problem.
 *           kF=referencekF if the referencekF tag is indicated in the input file,
 *           otherwise to kF value is assigned according formula kF=(3*pi^2*n)^{1/3}, where n corresponds to density in the box center.
 *           Note that it is passed via a pointer, to access/modify it use kF[0] or (*kF).
 * @param mu array with chemical potentials: mu[SPINA] and mu[SPINB].
 * @param extra_data_size size of extra_data in bytes, if extra_data size=0 the optional data is not uploaded
 * @param extra_data optional set of data uploaded by load_extra_data()
 * For more info see: Wiki->User defined parameters
 * */
extern "C" void process_params(double *params, double *kF, double *mu, size_t extra_data_size, void *extra_data)
{
    // if akF is active, then overwrite the value of input->sclgth
    // value of kF is generated by the function referencekF(...)
    if(input->akF!=0.0) input->sclgth = input->akF/kF[0];

    // PROCESS INPUT FILE PARAMETERS
    // Definition of R_TF:  0.5*m*omega^2 * R_TF^2 = mu
    //                      => omega = sqrt(2*mu/m)/R_TF
    params[11] = sqrt(2.*mu[SPINA])/ (params[1]*LX*0.5); // omega_x
    double eF = kF[0]*kF[0]/2.0;
    params[4] = params[4]*eF; // convert barrier height to internal units
    params[5] = params[5]/kF[0]; // convert barrier width to internal units
}

/** 
 * EXTERNAL POTENTIAL V_ext
 * @param ix x-coordinate from range [0,NX), to convert to Cartesian use: x = DX*(ix-NX/2)
 * @param iy y-coordinate from range [0,NY), to convert to Cartesian use: y = DY*(iy-NY/2),
 *           NOTE: in case of 1d code iy=0
 * @param iz z-coordinate from range [0,NZ), to convert to Cartesian use: z = DZ*(iz-NZ/2)
 *           NOTE: in case of 1d and 2d codes iz=0
 * @param it time iteration, use dc_t0 + dc_dt*it to compute corresponding time 
 * @param spin spin indicator, value from set {SPINA,SPINB}
 * @param params array of input parameters, before call of this routine the params array is processed by process_params() routine
 * @param extra_data_size size of extra_data in bytes, if extra_data size=0 the optional data is not uploaded
 * @param extra_data optional set of data uploaded by load_extra_data()
 * @return value of the external potential V_spin(x,y,z)
 * */
__device__ double v_ext(int ix, int iy, int iz, int it, int spin, double *params, size_t extra_data_size, void *extra_data)
{
    double x = DX*(ix-NX/2);

    double V_ho = harmonic_oscillator_smooth_edges(x, params[11], params[2]*LX*0.5, params[3]*LX*0.5, 1.0);
    double V_barrier = params[4] * exp(-2.0*pow(x/params[5],2));
    double V_tilt = params[6]*x/(LX/2);
    return V_ho + V_barrier + V_tilt;
}
 
/** 
 * EXTERNAL PAIRING POTENTIAL Delta_ext
 * @param ix x-coordinate from range [0,NX), to convert to Cartesian use: x = DX*(ix-NX/2)
 * @param iy y-coordinate from range [0,NY), to convert to Cartesian use: y = DY*(iy-NY/2)
 *           NOTE: in case of 1d code iy=0
 * @param iz z-coordinate from range [0,NZ), to convert to Cartesian use: z = DZ*(iz-NZ/2)
 *           NOTE: in case of 1d and 2d codes iz=0
 * @param it time iteration, use dc_t0 + dc_dt*it to compute corresponding time 
 * @param delta - value of delta computed self-consistently for given iteration it. 
 * @param params array of input parameters, before call of this routine the params array is processed by process_params() routine
 * @param extra_data_size size of extra_data in bytes, if extra_data size=0 the optional data is not uploaded
 * @param extra_data optional set of data uploaded by load_extra_data()
 * @return value of external pairing potential Delta_{ext}(x,y,z)
 * Complex type is equivalent to thrust::complex<double>
 * for more info see: https://thrust.github.io/doc/structthrust_1_1complex.html
 * */
__device__ Complex delta_ext(int ix, int iy, int iz, int it, Complex delta, double *params, size_t extra_data_size, void *extra_data)
{
//     double x = DX*(ix-NX/2);
//     double y = DY*(iy-NY/2);     // for 1d code iy will be always 0
//     double z = DZ*(iz-NZ/2);     // for 1d and 2d codes iz will be always 0
//     double t = dc_t0 + dc_dt*it; // time

    // ADD HERE FORMULA FOR Delta_ext(r)
    Complex D_ext = Complex(0.0,0.0);

    return D_ext; 
}

/** 
 * EXTERNAL VELOCITY FIELD vec[v]_ext
 * @param ix x-coordinate from range [0,NX), to convert to Cartesian use: x = DX*(ix-NX/2)
 * @param iy y-coordinate from range [0,NY), to convert to Cartesian use: y = DY*(iy-NY/2)
 *           NOTE: in case of 1d code iy=0
 * @param iz z-coordinate from range [0,NZ), to convert to Cartesian use: z = DZ*(iz-NZ/2)
 *           NOTE: in case of 1d and 2d codes iz=0
 * @param it time iteration, use dc_t0 + dc_dt*it to compute corresponding time 
 * @param spin spin indicator, value from set {SPINA,SPINB}
 * @param coordinate - Cartesian coordinate of the external velocity vector that should be computed, value from set {XAXIS, YAXIS, ZAXIS}
 *                     NOTE: for 1d code only XAXIS is requested, for 2d code XAXIS and YAXIS are requested.
 * @param params array of input parameters, before call of this routine the params array is processed by process_params() routine
 * @param extra_data_size size of extra_data in bytes, if extra_data size=0 the optional data is not uploaded
 * @param extra_data optional set of data uploaded by load_extra_data()
 * @return value of the external velocity vector v_ext(x,y,z)
 * */
__device__ double velocity_ext(int ix, int iy, int iz, int it, int spin, int coordinate, double *params, size_t extra_data_size, void *extra_data)
{
//     double x = DX*(ix-NX/2);
//     double y = DY*(iy-NY/2);     // for 1d code iy will be always 0
//     double z = DZ*(iz-NZ/2);     // for 1d and 2d codes iz will be always 0
//     double t = dc_t0 + dc_dt*it; // time

    // ADD HERE FORMULAS FOR vec{v}_ext=(vx, vy, vz)
    double v_ext;
    if(coordinate==XAXIS) v_ext=0.0;
    if(coordinate==YAXIS) v_ext=0.0;
    if(coordinate==ZAXIS) v_ext=0.0;

    return v_ext; 
}

/**
 * THIS FUNCTION IS CALLED AFTER THE DENSITIES ARE CONSTRUCTED.
 * IT IS CALLED ONLY IF ENABLE_MODIFY_DENSITIES IS DEFINED (in predefines.h)
 * Before each computation of potentials, the user can modify densities arbitrarily.
 * The function is called by HOST. You need to copy data from device to host, modify and copy it back.
 * @param it iteration number
 * (HOST POINTERS)
 * @param h_densities structure with densities, see (wiki) documentation for list of fields (HOST)
 *                    NOTE: some of densities can be destroyed.
 *                    To avoid any problems copy densities from DEVICE to this (HOST) buffer.
 * @param params array of input parameters (HOST).
 *               NOTE: If you modify this array, you must copy it to the DEVICE.
 *               Use: cudaMemcpyToSymbol(dc_params, params, MAX_USER_PARAMS*sizeof(double))
 * @param extra_data_size size of extra_data in bytes, if extra_data size=0 the optional data is not uploaded (HOST)
 * @param h_extra_data optional set of data uploaded by load_extra_data() (HOST)
 * (DEVICE POINTERS)
 * @param d_densities structure with densities, see (wiki) documentation for list of fields (DEVICE)
 * @param d_extra_data optional set of data uploaded by load_extra_data() (DEVICE)
 *                     NOTE: further computation process uses only d_extra_data
 * */
#include "tdwslda_functionals_framework_enable.h" // DO NOT REMOVE!
extern "C" void modify_densities(int it, wslda_density h_densities, double *params, size_t extra_data_size, void *h_extra_data,
                                         wslda_density d_densities,                                         void *d_extra_data)
{
    // // SNIPPETS
    // // To copy density from DEVICE to HOST use
    // memcopy_gpu2host(d_densities.rho_a, h_densities.rho_a, NUMBER_ELEMENT*sizeof(double));
    // // To copy density from HOST to DEVICE use
    // memcopy_host2gpu(h_densities.rho_a, d_densities.rho_a, NUMBER_ELEMENT*sizeof(double));
    // // To update array of DEVICE params use
    // memcopy_const_params(params);
    // // To update array of DEVICE d_extra_data use
    // memcopy_host2gpu(h_extra_data, d_extra_data, extra_data_size);
    // // To extract time use
    // double time = hc_t0+hc_dt*it;

    // ... add here your code ...
}
#include "tdwslda_functionals_framework_disable.h" // DO NOT REMOVE!

/**
 * THIS FUNCTION IS CALLED AFTER EACH COMPUTATION OF POTENTIALS.
 * IT IS CALLED ONLY IF ENABLE_MODIFY_POTENTIALS IS DEFINED (in predefines.h)
 * Before each application of the Hamiltonian, the user can arbitrarily modify potentials.
 * The function is called by DEVICE. All pointers are DEVICE pointers.
 * @param it iteration number
 * @param h_densities structure with densities, see (wiki) documentation for list of fields
 *                    NOTE: densities structure is processed by modify_densities(...) function before call ot this function.
 * @param h_potentials struture with potentials, see (wiki) documentation for list of fields
 * Global variabls deliver by tdwslda_functionals_framework_enable.h are:
 * @param params array of input parameters, before call of this routine the params array is processed by process_params() routine
 * @param extra_data_size size of extra_data in bytes, if extra_data size=0 the optional data is not uploaded
 * @param extra_data optional set of data uploaded by load_extra_data()
 * */
#include "tdwslda_functionals_framework_enable.h" // DO NOT REMOVE!
__global__ void modify_potentials(int it, wslda_density h_densities, wslda_potential h_potentials)
{    
    size_t ixyz= threadIdx.x + blockIdx.x * blockDim.x; // compute for this point
    int ix, iy, iz, i;
    
    if(ixyz<NUMBER_ELEMENT)
    {
        // decode_ixyz2ixiyiz(ixyz,ix,iy,iz, i);
        // // Now ix, iy, iz keeps lattice coordinate. 
        
        // h_potentials.V_a[ixyz] stores value of spin-up particles mean-field potential for coordinate (x,y,z)
        // and similarly for other potentials
        // ... below you can modify them your wish ...
        // 
    }
}
#include "tdwslda_functionals_framework_disable.h" // DO NOT REMOVE!

/**
 * THIS FUNCTION IS CALLED AFTER EACH COMPUTATION OF POTENTIALS
 * IT IS CALLED ONLY IF ENABLE_MODIFY_ENERGIES IS DEFINED (in predefines.h)
 * Users can modify arbitrary expressions for the energy computation.
 * The function is called by DEVICE. All pointers are DEVICE pointers.
 * @param it iteration number
 * @param h_densities structure with densities, see (wiki) documentation for list of fields
 *                    NOTE: densities structure is processed by modify_densities(...) function before call of this function.
 * @param h_potentials struture with potentials, see (wiki) documentation for list of fields
 *                    NOTE: potential structure is processed by modify_potentials(...) function before call of this function.
 * @param energy these entries user can modify.
 *                    Contributions are stored in energy[ETAG], where TAG in {EKIN, EPOT, EPAIR, ECURRENT, EPOTEXT, EPAIREXT, EVELEXT}
 * GLOBAL VARIABLES ACCESSIBLE WITHIN THIS ROUTINE
 * @param params array of input parameters, before call of this routine the params array is processed by process_params() routine
 * @param extra_data_size size of extra_data in bytes, if extra_data size=0 the optional data is not uploaded
 * @param extra_data optional set of data uploaded by load_extra_data()
 * */
#include "tdwslda_functionals_framework_enable.h" // DO NOT REMOVE!
__global__ void modify_energies(int it, wslda_density h_densities, wslda_potential h_potentials, double *energy)
{
    size_t ixyz= threadIdx.x + blockIdx.x * blockDim.x; // compute for this point
    int ix, iy, iz, i;

    if(ixyz<NUMBER_ELEMENT)
    {
        // decode_ixyz2ixiyiz(ixyz,ix,iy,iz, i);
        // // Now ix, iy, iz keeps lattice coordinate.

        // compute your contribution to the energy
        // double myE_contrib = (...)*VOLUME_ELEMENT;

        // Add contribution to desired tag, for example
        // energy[EPOT*NUMBER_ELEMENT+ixyz]+=myE_contrib;
    }
}
#include "tdwslda_functionals_framework_disable.h" // DO NOT REMOVE!

/**
 * This function provides size of extra_data array, in bytes.
 * The extra_data of specified size will be allocated by the main process.
 * This function is thread-safe.
 * @param params with input file parameters. 
 *              NOTE: the array contains bare input file values, not processed by process_params()!
 * @return size of the extra_data array that needs to be allocated, if 0 then extra_data will not be allocated.
 * */
extern "C" size_t get_extra_data_size(double *params)
{
    return 0;
}

/**
 * This function loads data into extra_data array.
 * This function is thread-safe.
 * @param extra_data_size size of array computed using function get_extra_data_size()
 * @param extra_data pointer to array that should be filled with data
 * @param params with input file parameters. 
 *               NOTE: the array contains bare input file values, not processed by process_params()!
 * @return 0 if load is successful, otherwise return error code. If nonzero value is returned the main code will terminate.
 * */
extern "C" int load_extra_data(size_t extra_data_size, void *extra_data, double *params)
{
    return 0;
}

/**
 * Scattering length, in code units.
 * This function is meaningful only in the case of BDG or SLDAE functionals.
 * For SLDA and ASLDA the scattering length is assumed to be infinite, and the function is ignored.
 * The function is called by DEVICE. All pointers are DEVICE pointers.
 * @param ix x-coordinate from range [0,NX), to convert to Cartesian use: x = DX*(ix-NX/2)
 * @param iy y-coordinate from range [0,NY), to convert to Cartesian use: y = DY*(iy-NY/2),
 *           NOTE: in case of 1d code iy=0
 * @param iz z-coordinate from range [0,NZ), to convert to Cartesian use: z = DZ*(iz-NZ/2)
 *           NOTE: in case of 1d and 2d codes iz=0
 * @param it time iteration, use dc_t0 + dc_dt*it to compute corresponding time 
 * @param params array of input parameters, before call of this routine the params array is processed by process_params() routine
 * @param extra_data_size size of extra_data in bytes, if extra_data size=0 the optional data is not uploaded
 * @param extra_data optional set of data uploaded by load_extra_data()
 * @return value of the scattering length a(x,y,z,t).
 * */
__device__ double scattering_length(int ix, int iy, int iz, int it, double *params, size_t extra_data_size, void *extra_data)
{
    return dc_sclgth; // by default return value from input file.
    
    // however, here you can define your own prescription
    // it can be time and position dependent
    // ...
}

/**
 * ------------------------ FOR FUNCTIONAL == CUSTOMEDF ------------------------
 * @see hpc-engine/tdwslda_functionals.h for more details
 * The function is called by DEVICE. All pointers are DEVICE pointers.
 * GLOBAL VARIABLES ACCESSIBLE WITHIN ROUTINES BELOW
 * @param params array of input parameters, before call of this routine the params array is processed by process_params() routine
 * @param extra_data_size size of extra_data in bytes, if extra_data size=0 the optional data is not uploaded
 * @param extra_data optional set of data uploaded by load_extra_data()
 * */

#if FUNCTIONAL==CUSTOMEDF
#include "tdwslda_functionals_framework_enable.h" // DO NOT REMOVE!
__global__ void tdwslda_compute_potentials(int it, wslda_density h_densities, wslda_potential h_potentials, double cccoeff)
{
    size_t ixyz= threadIdx.x + blockIdx.x * blockDim.x; // compute for this point
    int ix, iy, iz, i;
    
//     double na,nb;
//     double Va_ext,Vb_ext;
    
    if(ixyz<NUMBER_ELEMENT)
    {
        decode_ixyz2ixiyiz(ixyz,ix,iy,iz, i); // decode cartesian coordinates
        
//         // external potential
//         Va_ext=u_ext(ix,iy,iz,it,SPINA);
//         Vb_ext=u_ext(ix,iy,iz,it,SPINB);
//         
//         // read densities
//         na=h_densities.rho_a[ixyz];
//         nb=h_densities.rho_b[ixyz];
        
//         // save potential to global memory
//         h_potentials.V_a[ixyz]=...;  // <-- mean field + external potential
//         h_potentials.V_b[ixyz]=...;  // <-- mean field + external potential 
//         h_potentials.delta[ixyz]=...;  // <-- pairing field
//         ...
    }
}


__global__ void tdwslda_compute_energy(int it, wslda_density h_densities, wslda_potential h_potentials, // <-- INPUT
                                       double *E_kin, double *E_pot, double *E_pair, double *E_curr)    // <- OUTPUT
{
    size_t ixyz= threadIdx.x + blockIdx.x * blockDim.x; // compute for this point
    
    double na, nb;
    double taua, taub;
    double jxa, jya, jza, jxb, jyb, jzb;
    
    if(ixyz<NUMBER_ELEMENT)
    {
        // densities
        na=h_densities.rho_a[ixyz];
        nb=h_densities.rho_b[ixyz];
        
        // kinetic energy
        taua=h_densities.tau_a[ixyz]; 
        taub=h_densities.tau_b[ixyz];  
        
        // currents
#if CODEDIM>=1
        jxa=h_densities.j_a_x[ixyz];
        jxb=h_densities.j_b_x[ixyz];
#endif
#if CODEDIM>=2
        jya=h_densities.j_a_y[ixyz];
        jyb=h_densities.j_b_y[ixyz];
#else        
        jya=0.0;
        jyb=0.0;
#endif
#if CODEDIM>=3
        jza=h_densities.j_a_z[ixyz];
        jzb=h_densities.j_b_z[ixyz];
#else
        jza=0.0;
        jzb=0.0;
#endif
        
//         taua-=p_regularization(na)*(jxa*jxa+jya*jya+jza*jza)/na; // -ja^2/na: correction for tilde{tau}_a
//         taub-=p_regularization(nb)*(jxb*jxb+jyb*jyb+jzb*jzb)/nb; // -jb^2/nb: correction for tilde{tau}_b
// 
//         // galilean invariant contribution
//         E_kin[ixyz]=0.5*(h_potentials.alpha_a[ixyz]*taua + h_potentials.alpha_b[ixyz]*taub)*VOLUME_ELEMENT;
//         
//         // potential energy
//         E_pot[ixyz]=0.0;
//         
//         // pairing energy
//         E_pair[ixyz]=(h_potentials.delta[ixyz]*thrust::conj(h_densities.nu[ixyz])).real()*(-1.0)*VOLUME_ELEMENT;
//         
//         // flow energy
//         E_curr[ixyz]= p_regularization(na)*(jxa*jxa + jya*jya + jza*jza)/(2.*na)*VOLUME_ELEMENT  
//                     + p_regularization(nb)*(jxb*jxb + jyb*jyb + jzb*jzb)/(2.*nb)*VOLUME_ELEMENT;

    }
}
#include "tdwslda_functionals_framework_disable.h" // DO NOT REMOVE!
#endif