lb.impala

struct Neighborhood {
    nranks: i32,
    ranks: Buffer,
    naabbs: Buffer,
    offsets: Buffer,
    aabbs: ArrayData,
    rank_capacity: i32,
    aabb_capacity: i32
};

static mut neighborhood : Neighborhood;

fn @get_number_of_neighbor_ranks() -> i32 { neighborhood.nranks }
fn @get_neighborhood_rank(index: i32) -> i32 { bitcast[&[i32]](neighborhood.ranks.data)(index) }
fn @get_rank_number_of_aabbs(index: i32) -> i32 { bitcast[&[i32]](neighborhood.naabbs.data)(index) }
fn @get_rank_offset(index: i32) -> i32 { bitcast[&[i32]](neighborhood.offsets.data)(index) }
fn @get_neighborhood_aabb(nbh: Neighborhood, index: i32) -> AABB {
    let aabbs = get_array_real_ref(array_dev, nbh.aabbs);

    AABB {
        xmin: aabbs(index * 6 + 0),
        xmax: aabbs(index * 6 + 1),
        ymin: aabbs(index * 6 + 2),
        ymax: aabbs(index * 6 + 3),
        zmin: aabbs(index * 6 + 4),
        zmax: aabbs(index * 6 + 5)
    }
}

fn initialize_neighborhood() -> () {
    let rank_capacity = get_initial_maximum_neighbor_ranks();
    let aabb_capacity = rank_capacity as i32 * 8;

    neighborhood = Neighborhood {
        nranks: 0,
        ranks: alloc_cpu(rank_capacity * sizeof[i32]()),
        naabbs: alloc_cpu(rank_capacity * sizeof[i32]()),
        offsets: alloc_cpu(rank_capacity * sizeof[i32]()),
        aabbs: allocate_array(aabb_capacity, 6, sizeof[real_t](), true),
        rank_capacity: rank_capacity as i32,
        aabb_capacity: aabb_capacity
    };
}

fn release_neighborhood() -> () {
    release(neighborhood.ranks);
    release(neighborhood.naabbs);
    release(neighborhood.offsets);
    release_array(neighborhood.aabbs);
}

extern fn md_update_neighborhood(nranks: i32, total_aabbs: i32, ranks: &[i32], naabbs: &[i32], aabbs: &[real_t]) -> () {
    let mut offset = 0;
    neighborhood.nranks = nranks;

    if nranks > neighborhood.rank_capacity {
        let new_capacity = (nranks + 10) as i64;
        release(neighborhood.ranks);
        release(neighborhood.naabbs);
        release(neighborhood.offsets);
        neighborhood.ranks = alloc_cpu(new_capacity * sizeof[i32]());
        neighborhood.naabbs = alloc_cpu(new_capacity * sizeof[i32]());
        neighborhood.offsets = alloc_cpu(new_capacity * sizeof[i32]());
        neighborhood.rank_capacity = new_capacity as i32;
    }

    if total_aabbs > neighborhood.aabb_capacity {
        let new_capacity = total_aabbs + 10;
        reallocate_array(&mut neighborhood.aabbs, new_capacity, 6, sizeof[real_t](), false);
        neighborhood.aabb_capacity = new_capacity;
    }

    range(0, nranks, |r| {
        bitcast[&mut[i32]](neighborhood.ranks.data)(r) = ranks(r);
        bitcast[&mut[i32]](neighborhood.naabbs.data)(r) = naabbs(r);
        bitcast[&mut[i32]](neighborhood.offsets.data)(r) = offset;
        offset += naabbs(r);
    });

    range(0, total_aabbs * 6, |i| {
        get_array_real_ref(array_host, neighborhood.aabbs)(i) = aabbs(i);
    });

    transfer_array_to_device(neighborhood.aabbs);
}

// Get PBC conditions according to particle position
fn @(?pbc_x & ?pbc_y & ?pbc_z) get_pbc_flags_from_position(position: Vector3D, @pbc_x: i32, @pbc_y: i32, @pbc_z: i32, grid: Grid) -> PBCFlags {
    let mut pbc_flags = PBCFlags { x: 0 as i8, y: 0 as i8, z: 0 as i8 };
    let all = pbc_x == 0 && pbc_y == 0 && pbc_z == 0;

    if (all || pbc_x > 0) && position.x < grid.world_aabb.xmin + grid.spacing { pbc_flags.x =  1 as i8; };
    if (all || pbc_x < 0) && position.x > grid.world_aabb.xmax - grid.spacing { pbc_flags.x = -1 as i8; };
    if (all || pbc_y > 0) && position.y < grid.world_aabb.ymin + grid.spacing { pbc_flags.y =  1 as i8; };
    if (all || pbc_y < 0) && position.y > grid.world_aabb.ymax - grid.spacing { pbc_flags.y = -1 as i8; };
    if (all || pbc_z > 0) && position.z < grid.world_aabb.zmin + grid.spacing { pbc_flags.z =  1 as i8; };
    if (all || pbc_z < 0) && position.z > grid.world_aabb.zmax - grid.spacing { pbc_flags.z = -1 as i8; };

    pbc_flags
}

fn is_within_aabb_radius(point: Vector3D, aabb: AABB, radius: real_t) -> bool {
    let mut cond = true;
    if point.x < aabb.xmin + radius { cond = false; }
    if point.x > aabb.xmax - radius { cond = false; }
    if point.y < aabb.ymin + radius { cond = false; }
    if point.y > aabb.ymax - radius { cond = false; }
    if point.z < aabb.zmin + radius { cond = false; }
    if point.z > aabb.zmax - radius { cond = false; }
    cond
}

fn distance_point_line(point: real_t, min: real_t, max: real_t) -> real_t {
    let square = @|x: real_t| { x * x };

    if point < min {
        square(min - point)
    } else if point > max {
        square(point - max)
    } else {
        0.0
    }
}

fn distance_point_aabb(point: Vector3D, aabb: AABB) -> real_t {
    distance_point_line(point.x, aabb.xmin, aabb.xmax) +
    distance_point_line(point.y, aabb.ymin, aabb.ymax) +
    distance_point_line(point.z, aabb.zmin, aabb.zmax)
}

fn distance_point_aabb_periodic(point: Vector3D, aabb: AABB, glob_sizes: Vector3D, body: fn(real_t, Vector3D, PBCFlags) -> ()) -> () {
    let center = Vector3D {
        x: aabb.xmin + (aabb.xmax - aabb.xmin) * 0.5,
        y: aabb.ymin + (aabb.ymax - aabb.ymin) * 0.5,
        z: aabb.zmin + (aabb.zmax - aabb.zmin) * 0.5
    };

    let dis = vector_sub(point, center);
    let half_sizes = vector_scale(0.5, glob_sizes);
    let mut adj_pt = point;
    let mut pbc_flags = PBCFlags { x: 0 as i8, y: 0 as i8, z: 0 as i8 };

    if dis.x < -half_sizes.x { adj_pt.x += glob_sizes.x; pbc_flags.x = +1 as i8; }
    if dis.x > +half_sizes.x { adj_pt.x -= glob_sizes.x; pbc_flags.x = -1 as i8; }
    if dis.y < -half_sizes.y { adj_pt.y += glob_sizes.y; pbc_flags.y = +1 as i8; }
    if dis.y > +half_sizes.y { adj_pt.y -= glob_sizes.y; pbc_flags.y = -1 as i8; }
    if dis.z < -half_sizes.z { adj_pt.z += glob_sizes.z; pbc_flags.z = +1 as i8; }
    if dis.z > +half_sizes.z { adj_pt.z -= glob_sizes.z; pbc_flags.z = -1 as i8; }

    let rdis = distance_point_line(adj_pt.x, aabb.xmin, aabb.xmax) +
               distance_point_line(adj_pt.y, aabb.ymin, aabb.ymax) +
               distance_point_line(adj_pt.z, aabb.zmin, aabb.zmax);

    body(rdis, adj_pt, pbc_flags);
}

fn serialize_and_remove(grid: &mut Grid, comm: &mut Comm, check: fn(Vector3D) -> bool) -> i32 {
    let mut resize = 1;
    grid.nghost = 0;

    while resize > 0 {
        let const_grid = *grid;
        let const_comm = *comm;
        let world_aabb = const_grid.world_aabb;

        set_counter(0, const_grid);
        reset_resize(const_grid);

        particles_scalar(false, const_grid, |i, particle| {
            let send_flags = get_array_i8_ref(array_dev, const_grid.send_flags);
            let buffer = get_array_real_ref(array_dev, const_comm.send_buffer);
            let send_offsets = get_array_i32_ref(array_dev, const_comm.send_offsets);
            let pos = particle.get_position(i);

            if check(pos) {
                let counter = add_counter(const_grid);

                if counter >= const_comm.send_capacity {
                    grow_resize(counter, const_grid);
                } else {
                    let vel = particle.get_velocity(i);
                    let index = counter * 7;

                    buffer(index + 0) = particle.get_mass(i);
                    buffer(index + 1) = pos.x;
                    buffer(index + 2) = pos.y;
                    buffer(index + 3) = pos.z;
                    buffer(index + 4) = vel.x;
                    buffer(index + 5) = vel.y;
                    buffer(index + 6) = vel.z;
                    send_offsets(counter) = i;
                    send_flags(i) = 1 as i8;
                }
            } else {
                send_flags(i) = 0 as i8;
            }
        });

        resize = get_resize(const_grid);

        if resize > 0 {
            resize_send_capacity(comm, resize * 2);
        }
    }

    transfer_array_to_host(comm.send_offsets);
    transfer_array_to_host(grid.send_flags);

    let packed = get_counter(*grid);
    let nparticles = grid.nparticles - packed;
    let send_flags_host = get_array_i8_ref(array_host, grid.send_flags);
    let send_offsets_host = get_array_i32_ref(array_host, comm.send_offsets);
    let copy_list_host = get_array_i32_ref(array_host, comm.copy_list);
    let mut send_pos = grid.nparticles - 1;

    range(0, packed, |i| {
        if send_offsets_host(i) < nparticles {
            while(send_flags_host(send_pos) == 1 as i8) {
                send_pos--;
            }

            copy_list_host(i) = send_pos;
            send_pos--;
        } else {
            copy_list_host(i) = -1;
        }
    });

    transfer_array_to_device(comm.copy_list);

    let const_comm = *comm;
    let const_grid = *grid;
    let copy_list = get_array_i32_ref(array_dev, const_comm.copy_list);
    let particle = make_particle(const_grid, array_dev, ParticleDataLayout(), null_layout());

    device().loop_1d(false, packed, |i| {
        if copy_list(i) > 0 {
            let src = copy_list(i);
            let dst = get_array_i32_ref(array_dev, const_comm.send_offsets)(i);
            particle.set_mass(dst, particle.get_mass(src));
            particle.set_position(dst, particle.get_position(src));
            particle.set_velocity(dst, particle.get_velocity(src));
        }
    });

    transfer_array_to_host(comm.send_buffer);
    grid.nparticles -= packed;
    packed
}

fn deserialize_particles(grid: &mut Grid, comm: &mut Comm, nparticles: i32) -> () {
    grid.nghost = 0;

    let start = add_local_slots(nparticles, grid);
    let const_grid = *grid;
    let const_comm = *comm;
    let particle = make_particle(const_grid, array_dev, ParticleDataLayout(), null_layout());
    let recv_data = get_array_real_ref(array_dev, const_comm.recv_buffer);

    transfer_array_to_device(comm.recv_buffer);

    device().loop_1d(false, nparticles, |i| {
        let index = i * 7;
        let offset = start + i;
        particle.set_mass(offset, recv_data(index));
        particle.set_position(offset, Vector3D { x: recv_data(index + 1), y: recv_data(index + 2), z: recv_data(index + 3) });
        particle.set_velocity(offset, Vector3D { x: recv_data(index + 4), y: recv_data(index + 5), z: recv_data(index + 6) });
    });
}