#include <cstdio>
#include <numeric>
#include "black_holes.h"

/* BEGIN legacy inclusion */
// I'm not going to touch this C file
#define SQR(x) ((x)*(x))
double L[3]; // needed in pn_bh_spin.c
#include "pn_bh_spin.c"
#undef SQR
/* END legacy inclusion */

void two_body_gravity(
    const double m1, const double3 x1, const double3 v1,
    const double m2, const double3 x2, const double3 v2,
    const double eps,
    double& pot1, double3& acc1, double3& jrk1,
    double& pot2, double3& acc2, double3& jrk2)
{
    double3 dx = x1 - x2;
    double3 dv = v1 - v2;
    double r2 = dx.norm2() + eps*eps;
    double r  = sqrt(r2);
    double r3 = r2*r;
    double r4 = r2*r2;
    double RP = 3*(dx*dv)/r;

    pot1 = -m2/r;
    pot2 = -m1/r;
    acc1 = -m2*dx/r3;
    acc2 =  m1*dx/r3;
    jrk1 = -m2*(dv/r3 - RP*dx/r4);
    jrk2 =  m1*(dv/r3 - RP*dx/r4);
}

void Black_hole_physics::adjust_softening(
    double& pot1, double& pot2,
    double3& acc1, double3& acc2,
    double3& jrk1, double3& jrk2)
{
    if (eps_new < 0) return;
    // calculate and "minus" the BH <-> BH softened pot, acc & jerk
    two_body_gravity(
        m1, x1, v1,
        m2, x2, v2,
        eps_old,
        bbh_grav.pot1, bbh_grav.a1, bbh_grav.adot1,
        bbh_grav.pot2, bbh_grav.a2, bbh_grav.adot2);
    pot1 -= bbh_grav.pot1;
    pot2 -= bbh_grav.pot2;
    acc1 -= bbh_grav.a1;
    acc2 -= bbh_grav.a2;
    jrk1 -= bbh_grav.adot1;
    jrk2 -= bbh_grav.adot2;

    // calculate and "plus" the new BH <-> BH unsoftened pot, acc, jerk
    two_body_gravity(
        m1, x1, v1,
        m2, x2, v2,
        eps_new,
        bbh_grav.pot1, bbh_grav.a1, bbh_grav.adot1,
        bbh_grav.pot2, bbh_grav.a2, bbh_grav.adot2);
    pot1 += bbh_grav.pot1;
    pot2 += bbh_grav.pot2;
    acc1 += bbh_grav.a1;
    acc2 += bbh_grav.a2;
    jrk1 += bbh_grav.adot1;
    jrk2 += bbh_grav.adot2;
}

void Black_hole_physics::adjust_post_newtonian(
    const double dt_bh, // pn_usage should be const
    double3& acc1, double3& acc2,
    double3& jrk1, double3& jrk2)
{
    // calculate and "plus" the new BH <-> BH : PN1, PN2, PN2.5, PN3, PN3.5 : acc, jerk
    // TODO maybe have the PN terms as local variables here?
    int tmp;
    tmp = calc_force_pn_BH(m1, x1, v1, bbh_grav.spin1,
                           m2, x2, v2, bbh_grav.spin2,
                           c,  dt_bh, pn_usage,
                           (double(*)[3])bbh_grav.a_pn1, (double(*)[3])bbh_grav.adot_pn1,
                           (double(*)[3])bbh_grav.a_pn2, (double(*)[3])bbh_grav.adot_pn2, myRank, rootRank);
    if (tmp == 505) exit(-1); // Very ugly way to terminate

    // NOTE we have these _corr variables accumulating the corrections before
    // applying it. It's almost the same but different from applying each
    // correction term in a loop.
    double3 acc1_corr(0,0,0), acc2_corr(0,0,0), jrk1_corr(0,0,0), jrk2_corr(0,0,0);
    for (int i=1; i<7; i++) {
        acc1_corr += bbh_grav.a_pn1[i];
        acc2_corr += bbh_grav.a_pn2[i];
        jrk1_corr += bbh_grav.adot_pn1[i];
        jrk2_corr += bbh_grav.adot_pn2[i];
    }
    acc1 += acc1_corr;
    acc2 += acc2_corr;
    jrk1 += jrk1_corr;
    jrk2 += jrk2_corr;
}

void Black_hole_physics::write_bh_data(double time_cur, double m[], double3 x[], double3 v[], double pot[], double3 a[], double3 adot[], double dt[])
{
    // This function logs data on the black hole(s). It uses both external data
    // (the arguments to this function) and optionall internal data to this
    // object (the most recently calculated force and post-Newtonian terms).
    auto out = fopen("bh.dat", "a");
    if (count == 2) {
        for (int i=0; i < 2; i++) {
            double3 *a_pn, *adot_pn, a_bh, adot_bh;
            double pot_bh;
            if (i==0) {
                a_pn    = bbh_grav.a_pn1;
                adot_pn = bbh_grav.adot_pn1;
                pot_bh  = bbh_grav.pot1;
                a_bh    = bbh_grav.a1;
                adot_bh = bbh_grav.adot1;
            } else {
                a_pn    = bbh_grav.a_pn2;
                adot_pn = bbh_grav.adot_pn2;
                pot_bh  = bbh_grav.pot2;
                a_bh    = bbh_grav.a2;
                adot_bh = bbh_grav.adot2;
            }
            if (eps_new >= 0) {
                fprintf(out, "%.16E \t %.8E \t % .16E % .16E % .16E \t %.16E \t % .16E % .16E % .16E \t %.16E \t % .8E \t % .8E % .8E % .8E \t %.8E \t % .8E % .8E % .8E \t %.8E \t %.8E \t\t % .8E \t % .8E % .8E % .8E \t %.8E \t % .8E % .8E % .8E \t %.8E ",
                        time_cur, m[i],
                        x[i][0], x[i][1], x[i][2], x[i].norm(),
                        v[i][0], v[i][1], v[i][2], v[i].norm(),
                        pot[i],
                        a[i][0],    a[i][1],    a[i][2],    a[i].norm(),
                        adot[i][0], adot[i][1], adot[i][2], adot[i].norm(),
                        dt[i],
                        pot_bh,
                        a_bh[0],    a_bh[1],    a_bh[2],    a_bh.norm(),
                        adot_bh[0], adot_bh[1], adot_bh[2], adot_bh.norm());
                if (c > 0) {
                    fprintf(out, "\t");
                    for (int pn_idx=0; pn_idx < 7; pn_idx++) {
                        fprintf(out, "\t % .8E % .8E % .8E \t %.8E \t % .8E % .8E % .8E \t %.8E ", a_pn[pn_idx][0], a_pn[pn_idx][1], a_pn[pn_idx][2], a_pn[pn_idx].norm(), adot_pn[pn_idx][0], adot_pn[pn_idx][1], adot_pn[pn_idx][2], adot_pn[pn_idx].norm());
                    }
                }
                fprintf(out, "\n");
            } else {
                fprintf(out,"%.16E \t %.8E \t % .16E % .16E % .16E \t %.16E \t % .16E % .16E % .16E \t %.16E \t % .8E \t % .8E % .8E % .8E \t %.8E \t % .8E % .8E % .8E \t %.8E \t %.8E \n",
                        time_cur, m[i],
                        x[i][0], x[i][1], x[i][2], x[i].norm(),
                        v[i][0], v[i][1], v[i][2], v[i].norm(),
                        pot[i],
                        a[i][0],    a[i][1],    a[i][2],    a[i].norm(),
                        adot[i][0], adot[i][1], adot[i][2], adot[i].norm(),
                        dt[i]);
            }
        }
        fprintf(out, "\n");
    } else if (count == 1) {
        fprintf(out,"%.16E \t %.8E \t % .16E % .16E % .16E \t %.16E \t % .16E % .16E % .16E \t %.16E \t % .8E \t % .8E % .8E % .8E \t %.8E \t % .8E % .8E % .8E \t %.8E \t %.8E \n",
                time_cur, m[0],
                x[0][0], x[0][1], x[0][2], x[0].norm(),
                v[0][0], v[0][1], v[0][2], v[0].norm(),
                pot[0],
                a[0][0],    a[0][1],    a[0][2],    a[0].norm(),
                adot[0][0], adot[0][1], adot[0][2], adot[0].norm(),
                dt[0]);
        fprintf(out,"\n");
    }
    fclose(out);
}

void Write_bh_nb_data::operator()(double time_cur)
{
    for (int i_bh=0; i_bh < smbh_count; i_bh++) {
        for (int i=0; i<N; i++) var_sort[i] = (x[i]-x[i_bh]).norm();
        std::iota(ind_sort.begin(), ind_sort.end(), 0);
        std::partial_sort(ind_sort.begin(), ind_sort.begin() + nb, ind_sort.begin() + N, [&](int i, int j) {return var_sort[i] < var_sort[j];});

        fprintf(out,"%.16E \t %07d \t %.8E \t % .8E % .8E % .8E \t % .8E % .8E % .8E \t",
                time_cur,
                i_bh,
                m[i_bh],
                x[i_bh][0], x[i_bh][1], x[i_bh][2],
                v[i_bh][0], v[i_bh][1], v[i_bh][2]);

        for (int j=1; j < nb; j++) {
            int i = ind_sort[j];
            fprintf(out,"%02d %07d   %.8E   % .8E % .8E % .8E   %.8E   % .8E % .8E % .8E   %.8E \t",
                    j, i,
                    m[i],
                    x[i][0], x[i][1], x[i][2], (x[i]-x[i_bh]).norm(),
                    v[i][0], v[i][1], v[i][2], (v[i]-v[i_bh]).norm());
        }
        fprintf(out,"\n");
    }
    fprintf(out, "\n"); // this is redundant
    fflush(out);
}

void Binary_smbh_influence_sphere_output::operator()(int ind_act[], int n_act, double timesteps, double time_cur)
{
    double m_bh1  = m[0];
    double m_bh2  = m[1];
    double3 x_bh1 = x[0];
    double3 x_bh2 = x[1];
    double3 v_bh1 = v[0];
    double3 v_bh2 = v[1];

    double3 x_bbhc = (m_bh1*x_bh1 + m_bh2*x_bh2)/(m_bh1 + m_bh2);
    double3 v_bbhc = (m_bh1*v_bh1 + m_bh2*v_bh2)/(m_bh1 + m_bh2);

    double DR2 = (x_bh1 - x_bh2).norm2();
    double DV2 = (v_bh1 - v_bh2).norm2();
    double EB = -(m_bh1 + m_bh2) / sqrt(DR2) + 0.5 * DV2;
    double SEMI_a = -0.5 * (m_bh1 + m_bh2)/EB;
    double SEMI_a2 = SEMI_a*SEMI_a;

    for (int i=0; i<n_act; i++) {
        int j_act = ind_act[i];
        if (j_act<2) continue;
        double& pot_bh1 = pot[0];
        double& pot_bh2 = pot[1];
        double& m_act   = m[j_act];
        double3& x_act  = x[j_act];
        double3& v_act  = v[j_act];
        double& dt_act  = dt[j_act];
        double& pot_act = pot[j_act];
        double tmp_r2 = (x_act - x_bbhc).norm2();
        if (tmp_r2 < SEMI_a2*factor*factor) {
            if (inf_event[j_act] == 0) {
                fprintf(out,"INF1 %.6E %.16E %07d %07d   %.6E % .6E % .6E % .6E % .6E % .6E % .6E   %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E  %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E   %.6E   %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E  %.6E \n",
                    timesteps, time_cur, i, j_act,
                    sqrt(DR2), x_bbhc[0], x_bbhc[1], x_bbhc[2], v_bbhc[0], v_bbhc[1], v_bbhc[2],
                    m_bh1, x_bh1[0], x_bh1[1], x_bh1[2], v_bh1[0], v_bh1[1], v_bh1[2], pot_bh1,
                    m_bh2, x_bh2[0], x_bh2[1], x_bh2[2], v_bh2[0], v_bh2[1], v_bh2[2], pot_bh2,
                    sqrt(tmp_r2),
                    m_act, x_act[0], x_act[1], x_act[2], v_act[0], v_act[1], v_act[2], pot_act,
                    dt_act);

            inf_event[j_act] = 1;
            }
        } else {
            if (inf_event[j_act] == 1) {
                fprintf(out,"INF2 %.6E %.16E %07d %07d   %.6E % .6E % .6E % .6E % .6E % .6E % .6E   %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E  %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E   %.6E   %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E  %.6E \n",
                    timesteps, time_cur, i, ind_act[i],
                    sqrt(DR2), x_bbhc[0], x_bbhc[1], x_bbhc[2], v_bbhc[0], v_bbhc[1], v_bbhc[2],
                    m_bh1, x_bh1[0], x_bh1[1], x_bh1[2], v_bh1[0], v_bh1[1], v_bh1[2], pot_bh1,
                    m_bh2, x_bh2[0], x_bh2[1], x_bh2[2], v_bh2[0], v_bh2[1], v_bh2[2], pot_bh2,
                    sqrt(tmp_r2),
                    m_act, x_act[0], x_act[1], x_act[2], v_act[0], v_act[1], v_act[2], pot_act,
                    dt_act);
            }
            inf_event[j_act] = 0;
        } /* if (tmp_r2 < DR2*R_INF2) */
    } /* i */
    fflush(out);
}