import * as se3 from '../se3';
import * as linalg from './linalg';
import {Vec3} from './linalg';

interface Orbit {
    excentricity: number,
    semimajor_axis: number,
    inclination: number,
    ascending_node_longitude: number,
    periapsis_argument: number,
    t0: number,

    lastE: number | undefined,
    tf: number[] | undefined,
}

interface Body {
    position: Vec3,
    velocity: Vec3,
    orientation: Vec3,
    children: Body[],
    mass: number,
    orbit: Orbit,
    spin: Vec3,
    name: string,
}

export function updateBodyPhysics(time: number, body: Body, parentBody : Body | undefined = undefined) {
    if (parentBody !== undefined) {
        const mu = parentBody.mass;
        const {position, velocity} = getCartesianState(body.orbit, mu, time);
        body.position = [
            parentBody.position[0] + position[0],
            parentBody.position[1] + position[1],
            parentBody.position[2] + position[2],
        ];
        body.velocity = [
            parentBody.velocity[0] + velocity[0],
            parentBody.velocity[1] + velocity[1],
            parentBody.velocity[2] + velocity[2],
        ];
    } else {
        body.position = [0, 0, 0];
        body.velocity = [0, 0, 0];
    }
    body.orientation = getOrientation(body, time);

    if (body.children !== undefined) {
        for (const child of body.children) {
            updateBodyPhysics(time, child, body);
        }
    }
}

export function findSoi(rootBody: Body, position: number[]) : Body {
    const bodies = [rootBody];
    let body : Body;
    while (bodies.length > 0) {
        body = bodies.shift()!;
        
        if (body.children === undefined) {
            return body;
        }

        for (const child of body.children) {
            const soi = child.orbit.semimajor_axis * Math.pow(child.mass / body.mass, 2/5);
            const pos = position;
            const bod = child.position;
            const dr = [pos[0] - bod[0], pos[1] - bod[1], pos[2] - bod[2]];
            if (dr[0]**2 + dr[1]**2 + dr[2]**2 < soi**2) {
                bodies.push(child);
            }
        }
    }

    return body!;
}

export function computeOrbit(player: any, body: Body, time: number) {
    const {cross, diff, norm, dot, scale} = linalg;
    const rvec = diff(player.position, body.position);
    const r = norm(rvec);

    if (norm(player.velocity) < 1e-6) {
        // cheating
        console.log('cheating');
        player.velocity = scale(cross([1, 1, 1], rvec), 0.01/(r**2));
    }
    const vvec = diff(player.velocity, body.velocity);
    const v = norm(vvec);
    const Hvec = cross(rvec, vvec);
    const H = norm(Hvec);

    const mu = body.mass;
    const p = H**2 / mu;
    const resinnu = Math.sqrt(p/mu) * dot(vvec, rvec)
    const recosnu = p - r;

    const e = Math.sqrt(resinnu**2 + recosnu**2) / r;

    // should also work for hyperbolic orbits
    const a = p/(1-e**2);

    const x = scale(rvec, 1/r);
    const yy = cross(Hvec, rvec);
    const y = scale(yy, 1/norm(yy));
    const z = scale(Hvec, 1/H);

    // Om i and w can be skipped when we just give tf...
    let Om = Math.atan2(Hvec[0], -Hvec[1]);
    if (Hvec[0] === 0 && Hvec[1] === 0) {
        Om = 0;
    }

    let i = Math.atan2(Math.sqrt(Hvec[0]**2 + Hvec[1]**2), Hvec[2]);
    if (i * Math.sign((Hvec[0] || 1) * Math.sin(Om)) < 0) {
        i *= -1;
    }

    const nu = Math.atan2(resinnu, recosnu);

    let w = Math.atan2(rvec[2] / r / Math.sin(i), y[2] / Math.sin(i)) || 0 - nu;
    
    let t0;
    let E;

    if (a < 0) {
        E = 2 * Math.atanh(Math.sqrt((e-1)/(e+1)) * Math.tan(nu/2));
        const n = Math.sqrt(-mu / (a**3));
        t0 = time - (e*Math.sinh(E) - E) / n;

        t0 %= 2 * Math.PI / n;
    } else {
        E = 2 * Math.atan(Math.sqrt((1-e)/(1+e)) * Math.tan(nu/2));
        const n = Math.sqrt(mu / (a**3));
        t0 = time - (1/n)*(E - e*Math.sin(E));

        t0 %= 2 * Math.PI / n;
    }

    // column-major... see se3.js
    const tf = se3.product([
        x[0], x[1], x[2], 0,
        y[0], y[1], y[2], 0,
        z[0], z[1], z[2], 0,
        0,    0,    0,    1,
    ], se3.rotz(-nu));

    const orbit = {
        excentricity: e,
        semimajor_axis: a,
        inclination: i,
        ascending_node_longitude: Om,
        periapsis_argument: w,
        t0,
        tf,
        lastE: E,
    };

    return orbit;
}

/** Let's be honest I should clean this up.
 * 
 * This is the part that solves Kepler's equation using Newton's method.
 * For circular-ish orbits, one or two iterations are usually enough.
 * More excentric orbits can take more (6 or 7?).
 * 
 * For near-parabolic orbits (and some others?) it often fails to converge...
 */
export function getCartesianState(orbit: Orbit, mu: number, time: number) {
    const {
        excentricity: e,
        semimajor_axis: a,
        inclination: i,
        ascending_node_longitude: Om,
        periapsis_argument: w,
        t0,
    } = orbit;
    let n = Math.sqrt(mu/(a**3));
    if (a < 0) {
        n = Math.sqrt(mu/-(a**3));  // mean motion
    }
    const M = (n * (time - t0)) % (2 * Math.PI);  // mean anomaly

    // Newton's method
    var E2 = 0;
    var E = M;
    let iterations = 0;
    // a clever guess? https://link.springer.com/article/10.1023/A:1008200607490
    // doesn't work at all.

    while (Math.abs(E - E2) > 1e-10) {
        if (e < 0.001) {
            break;
        }
        E = E2;
        if (e < 1) {
            E2 = E - (E - e * Math.sin(E) - M) / (1 - e * Math.cos(E));
        } else if (e > 1) {
            E2 = E - (-E + e * Math.sinh(E) - M) / (e * Math.cosh(E) - 1);
        } else {
            E2 = E - (E + E*E*E/3 - M) / (1 + E*E);
        }

        iterations++;
        if (iterations > 100) {
            console.log('numerical instability');
            return {};
        }
    }
    orbit.lastE = E;
    let nu;
    if (e > 1) {
        nu = 2 * Math.atan(Math.sqrt((e+1) / (e-1)) * Math.tanh(E/2));
    } else {
        nu = 2 * Math.atan(Math.sqrt((1+e) / (1-e)) * Math.tan(E/2));
    }
    const p = a * (1 - e**2);
    const r = p / (1 + e * Math.cos(nu));// * ((a < 0) ? -1 : 1);
    const rd = e * Math.sqrt(mu / p) * Math.sin(nu);

    if (orbit.tf === undefined) {
        // FIXME: this is actually borken. :/
        orbit.tf = [se3.rotz(Om), se3.rotx(i), se3.rotz(w)].reduce(se3.product);
    }

    const tf = se3.product(orbit.tf, se3.rotz(nu));
    const pos = se3.apply(tf, [r, 0, 0, 1]);
    const vel = se3.apply(tf, [rd, Math.sqrt(p * mu) / r, 0, 1]);

    return {
        position: pos.slice(0, 3),
        velocity: vel.slice(0, 3),
    };
}

function getOrientation(body: Body, time: number) {
    return se3.rotxyz(
        body.spin[0] * time,
        body.spin[1] * time,
        body.spin[2] * time,
    );
}

export function makeOrbitObject(gl: WebGLRenderingContext, glContext: any, orbit: Orbit, parentPosition: number[]) {
    const position = parentPosition;
    
    // FIXME: currently borken.
    // const orientation = [
    //     se3.rotz(orbit.ascendingNodeLongitude),
    //     se3.rotx(orbit.inclination),
    //     se3.rotz(orbit.periapsisArgument),
    // ].reduce(se3.product);
    const orientation = orbit.tf;

    const buffer = gl.createBuffer();
    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);

    const a = orbit.semimajor_axis;
    const b = a * Math.sqrt(1 - orbit.excentricity**2);

    const x = - orbit.semimajor_axis * orbit.excentricity;

    gl.bufferData(gl.ARRAY_BUFFER, new Float32Array([
        x-a, -b, 0, -1, -1,
        x-a, +b, 0, -1, +1,
        x+a, -b, 0, +1, -1,
        x+a, +b, 0, +1, +1,
    ]), gl.STATIC_DRAW);

    const geometry = {
        glBuffer: buffer,
        numVertices: 4,
        delete: () => gl.deleteBuffer(buffer),
    };

    return {
        geometry,
        orientation,
        position,
        glContext,
    };
}