#define PY_SSIZE_T_CLEAN
#define _USE_MATH_DEFINES
#include <Python.h>
#include <math.h>


// Class to compute magnetic fields of solenoids.
// Improved version of the original Soleno code.
// Soleno was originally developed at Twente University by Gert Mulder, and upgraded in 1999 by Erik Krooshoop.

// https://gitlab.com/jeroen.van.nugteren/soleno/-/blob/master/soleno_calcB.cpp

static double sign(double val) {
    if(val > 0.0) {
        return 1.0;
    } else if(val < 0.0) {
        return -1.0;
    } else {
        return 0.0;
    }
}


struct field2d {
    double br;
    double bz;
    int valid;
};


struct field2d calc_b(double Zmh, double Rm, double Rp) {
    const unsigned int MAX_ITERATIONS = 50;
    // calculation accuracy
    const double Eps = 1e-12;

    double Hulp = sqrt(Rp * Rp + Zmh * Zmh);
    double Kc = sqrt(Rm * Rm + Zmh * Zmh) / Hulp;
    if(Kc < Eps) Kc = Eps;

    double Mj = 1;
    double Nj = Kc;
    double Aj = -1;
    double Bj = 1;
    double Pj = Rm / Rp;
    double Cj = Pj + 1;
    if(fabs(Rm / Rp) < Eps) Pj = 1;

    double Dj = Cj * sign(Pj);
    Pj = fabs(Pj);
    unsigned int i = 1;
    while(fabs(Mj - Nj) > (Eps * Mj)) {
        if (i > MAX_ITERATIONS) {
            PyErr_SetString(PyExc_StopIteration, "Solution does not converge");
            return (struct field2d){.br = 0, .bz = 0, .valid = 0};
        }

        double Xmn = Mj * Nj;
        double Xmp = Xmn / Pj;
        double D0 = Dj;
        Dj = 2.0 * (Dj + Xmp * Cj);
        Cj = Cj + D0 / Pj;
        double B0 = Bj;
        Bj = 2.0 * (Bj + Nj * Aj);
        Aj = Aj + B0 / Mj;
        Pj = Pj + Xmp;
        Mj = Mj + Nj;
        Nj = 2.0 * sqrt(Xmn);
        i++;
    }

    double Brh = (Aj + Bj / Mj) / (Mj * Hulp);
    double Bzh = (Cj + Dj / Mj) / ((Mj + Pj) * Hulp) * Zmh;

    return (struct field2d){.br = Brh, .bz = Bzh, .valid = 1};
}


struct field2d calc_field_primitive(
    double R,
    double Z,
    double Rin,
    double Rout,
    double Zlow,
    double Zhigh,
    double It,
    int Nlayers) {
    
    double Brd = 0;
    double Bzd = 0;

    if(Rin != Rout && Zlow != Zhigh) {
        double Factor = It / (Zhigh - Zlow) / (double)Nlayers;
        if(Factor != 0.0) {
            double Zm1 = Z - Zlow;
            double Zm2 = Z - Zhigh;
            double dR = (Rout - Rin) / (double)Nlayers;
            double dH = dR / 2.0;
            double dE = dH;
            for(int tt = 0; tt < Nlayers; tt++) {
                double Rs = Rin - dH + dR * (tt + 1);
                if(fabs(R - Rs) >= dE) {
                    // extern
                    double Rp = Rs + R;
                    double Rm = Rs - R;
                    struct field2d B1 = calc_b(Zm1, Rm, Rp);
                    if(!B1.valid) return B1;
                    struct field2d B2 = calc_b(Zm2, Rm, Rp);
                    if(!B2.valid) return B2;
                    Brd = Brd + (B2.br - B1.br) * Factor * Rs;
                    Bzd = Bzd + (B1.bz - B2.bz) * Factor;
                } else {
                    // intern
                    // RIntern = true;
                    double Rp = Rs + Rs - dE;
                    double Rm = dE;
                    double Weeg = (Rs - R + dE) / dE / 2.0;
                    struct field2d B1 = calc_b(Zm1, Rm, Rp);
                    if(!B1.valid) return B1;
                    struct field2d B2 = calc_b(Zm2, Rm, Rp);
                    if(!B2.valid) return B2;
                    Brd = Brd + (B2.br - B1.br) * Weeg * Factor * Rs;
                    Bzd = Bzd + (B1.bz - B2.bz) * Weeg * Factor;
                    Weeg = 1.0 - Weeg;
                    Rp = Rs + Rs + dE;
                    Rm = -dE;
                    B1 = calc_b(Zm1, Rm, Rp);
                    if(!B1.valid) return B1;
                    B2 = calc_b(Zm2, Rm, Rp);
                    if(!B2.valid) return B2;
                    Brd = Brd + (B2.br - B1.br) * Weeg * Factor * Rs;
                    Bzd = Bzd + (B1.bz - B2.bz) * Weeg * Factor;
                }
            }
        }
    }

    double Brt = Brd * M_PI * 1e-7;
    double Bzt = Bzd * M_PI * 1e-7;
    return (struct field2d){.br = Brt, .bz = Bzt, .valid = 1};
}


static PyObject* calc_single_field(PyObject *self, PyObject *args) {
    // This function calculates the field at (R, Z) for a single solenoid element using:
    // `[Br,Bz] = Soleno.calc_single_field(R, Z, Rin, Rout, Zmin, Zmax, I, Nlayers)`
    //
    // Rin, Rout, Zmin, Zmax, I and Nlayers define the coils and must be equal in size
    // R and Z define the position and must be equal in size as well.
    // Output Br and Bz is respectively the field in the radial and axial direction

    double R, Z, Rin, Rout, Zlow, Zhigh, I;
    int Nlayers;
    if(!PyArg_ParseTuple(args, "dddddddi", &R, &Z, &Rin, &Rout, &Zlow, &Zhigh, &I, &Nlayers))
        return NULL;
    
    double Br, Bz;

    if(R >= Rin && R <= Rout) {
        // internal
        double A = (R - Rin) / (Rout - Rin);
        int NL1 = (int)((double)Nlayers * A) + 1;
        int NL2 = (int)((double)Nlayers * (1 - A)) + 1;

        // deel1
        struct field2d B2 = calc_field_primitive(R, Z, Rin, R, Zlow, Zhigh, I * A, NL1);
        if(!B2.valid) return NULL;
        struct field2d Bt = calc_field_primitive(R, Z, Rin, R, Zlow, Zhigh, I * A, NL1 * 2);
        if(!Bt.valid) return NULL;
        double BZtemp = Bt.bz + (Bt.bz - B2.bz) / 3.0;  // Extrapolatie naar 1/Nl/Nl=0
        double BRtemp = Bt.br + (Bt.br - B2.br) / 3.0;

        // deel2
        B2 = calc_field_primitive(R, Z, R, Rout, Zlow, Zhigh, I * (1 - A), NL2);
        if(!B2.valid) return NULL;
        Bt = calc_field_primitive(R, Z, R, Rout, Zlow, Zhigh, I * (1 - A), NL2 * 2);
        if(!Bt.valid) return NULL;

        Bz = BZtemp + Bt.bz + (Bt.bz - B2.bz) / 3.0;  // Extrapolatie naar 1/Nl/Nl=0
        Br = BRtemp + Bt.br + (Bt.br - B2.br) / 3.0;
    } else {
        // external
        struct field2d B2 = calc_field_primitive(R, Z, Rin, Rout, Zlow, Zhigh, I, Nlayers);
        if(!B2.valid) return NULL;
        struct field2d Bt = calc_field_primitive(R, Z, Rin, Rout, Zlow, Zhigh, I, Nlayers * 2);
        if(!Bt.valid) return NULL;
        Bz = Bt.bz + (Bt.bz - B2.bz) / 3.0;  // Extrapolatie naar 1/Nl/Nl=0
        Br = Bt.br + (Bt.br - B2.br) / 3.0;
    }

    return Py_BuildValue("dd", Br, Bz);
}

PyMethodDef hiper_funcs[] = {
	{	"calc_single_field",
		(PyCFunction)calc_single_field,
		METH_VARARGS,
		NULL},
	{	NULL}
};

char hiper_mod_docs[] = "High Perfomance (HiPer) code for NICQS.";

PyModuleDef hiper_mod = {
	PyModuleDef_HEAD_INIT,
	"hiper",
	hiper_mod_docs,
	-1,
	hiper_funcs,
	NULL,
	NULL,
	NULL,
	NULL
};

PyMODINIT_FUNC PyInit_hiper(void) {
	return PyModule_Create(&hiper_mod);
}