package scientific.matplot

import std.math.*
import std.unittest.*
import std.unittest.testmacro.*

import scientific.numbers.*
import scientific.linear.*
import scientific.stats.random.*
import scientific.stats.normal.*

/* Type for x, y, u, v, m, z, w, c: Float64 */
foreign func c_quiver(x: CPointer<Unit>, y: CPointer<Unit>, 
                      u: CPointer<Unit>, v: CPointer<Unit>,
                      rows: Int64, cols: Int64, scale: Float64, line_spec: CString): CPointer<Unit>

foreign func c_quiver_m(x: CPointer<Unit>, y: CPointer<Unit>, 
                        u: CPointer<Unit>, v: CPointer<Unit>, m: CPointer<Unit>,
                        rows: Int64, cols: Int64, scale: Float64, line_spec: CString): CPointer<Unit>

foreign func c_quiver_3d(z: CPointer<Unit>, u: CPointer<Unit>, v: CPointer<Unit>,
                         w: CPointer<Unit>, rows: Int64, cols: Int64, scale: Float64,
                         line_spec: CString): CPointer<Unit>

foreign func c_quiver_3d_m(x: CPointer<Unit>, y: CPointer<Unit>, z: CPointer<Unit>,
                           u: CPointer<Unit>, v: CPointer<Unit>, w: CPointer<Unit>,
                           c: CPointer<Unit>, len: Int64, scale: Float64, line_spec: CString): CPointer<Unit>

foreign func c_vectors_normalize(ptr: CPointer<Unit>, b: Bool): CPointer<Unit>
foreign func c_vectors_line_width(ptr: CPointer<Unit>, width: Float32): CPointer<Unit>

foreign func c_feather(u: CPointer<Unit>, v: CPointer<Unit>, len: Int64): CPointer<Unit>

public class Vectors {
    var ptr: CPointer<Unit> = CPointer<Unit>()

    /* Input is a pointer to matplot::vectors */
    init(ptr: CPointer<Unit>) {
        this.ptr = ptr
    }

    public func normalize(b: Bool) {
        this.ptr = unsafe { c_vectors_normalize(this.ptr, b) }
        return this
    }

    public func line_width(width: Float32) {
        this.ptr = unsafe { c_vectors_line_width(this.ptr, width) }
        return this
    }
}

public func quiver(x: Matrix<Float64>, y: Matrix<Float64>,
                   u: Matrix<Float64>, v: Matrix<Float64>,
                   scale!: Float64 = 1.0, line_spec!: String = ""): Vectors {
    let row = x.getRows()
    let col = x.getCols()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe { c_quiver(x.ptr, y.ptr, u.ptr, v.ptr, row, col, scale, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Vectors(handle)
}

public func quiver(x: Matrix<Float64>, y: Matrix<Float64>,
                   u: Matrix<Float64>, v: Matrix<Float64>, m: Matrix<Float64>,
                   scale!: Float64 = 1.0, line_spec!: String = ""): Vectors {
    let row = x.getRows()
    let col = x.getCols()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe { c_quiver_m(x.ptr, y.ptr, u.ptr, v.ptr, m.ptr, row, col, scale, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Vectors(handle)
}

public func quiver3(z: Matrix<Float64>, u: Matrix<Float64>, v: Matrix<Float64>, w: Matrix<Float64>,
                    scale!: Float64 = 1.0, line_spec!: String = ""): Vectors {
    let row = z.getRows()
    let col = z.getCols()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe { c_quiver_3d(z.ptr, u.ptr, v.ptr, w.ptr, row, col, scale, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Vectors(handle)
}

public func quiver3(x: Vector<Float64>, y: Vector<Float64>, z: Vector<Float64>,
                    u: Vector<Float64>, v: Vector<Float64>, w: Vector<Float64>, c: Vector<Float64>,
                    scale!: Float64 = 1.0, line_spec!: String = ""): Vectors {
    let len = x.size()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe { c_quiver_3d_m(x.ptr, y.ptr, z.ptr, u.ptr, v.ptr, w.ptr, c.ptr, len, scale, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Vectors(handle)
}

public func feather(u: Vector<Float64>, v: Vector<Float64>): Vectors {
    let handle = unsafe { c_feather(u.ptr, v.ptr, u.size()) }
    return Vectors(handle)
}

public func testQuiver1() {
    let data1 = linspace(0.0, 2.0, num:11)
    let (x, y) = meshgrid(data1, data1)

    let u = apply2(x, y, {t1: Float64, t2:Float64 => cos(t1) * t2})
    let v = apply2(x, y, {t1: Float64, t2:Float64 => sin(t1) * t2})

    quiver(x, y, u, v)
    save("./tests/imgs/quiver/quiver_1.svg", "svg")
    clear()
}

public func testQuiver2() {
    let data1 = linspace(0.0, 2.0, num:11)
    let (x, y) = meshgrid(data1, data1)

    let u = apply2(x, y, {t1: Float64, t2:Float64 => cos(t1) * t2})
    let v = apply2(x, y, {t1: Float64, t2:Float64 => sin(t1) * t2})

    quiver(x, y, u, v, scale:2.0)
    save("./tests/imgs/quiver/quiver_2.svg", "svg")
    clear()
}

public func testQuiver3() {
    let data1 = linspace(0.0, 2.0, num:11)
    let (x, y) = meshgrid(data1, data1)

    let u = apply2(x, y, {t1: Float64, t2:Float64 => cos(t1) * t2})
    let v = apply2(x, y, {t1: Float64, t2:Float64 => sin(t1) * t2})

    quiver(x, y, u, v, scale:0.0)
    save("./tests/imgs/quiver/quiver_3.svg", "svg")
    clear()
}

public func testQuiver4() {
    let data1 = linspace(-2.0, 2.0, num:21)
    let (x, y, z) = meshgrid(data1, data1, {
        x:Float64, y:Float64 => x * exp(-pow(x, 2.0) - pow(y, 2.0))})
    let (dx, dy) = gradient(z, spacing_x: 0.2, spacing_y: 0.2)
    contour(x, y, z)
    hold(true)
    quiver(x, y, dx, dy)
    save("./tests/imgs/quiver/quiver_4.svg", "svg")
    clear()
}

public func testQuiver5() {
    let data1 = linspace(-2.5, 2.5, num:21)
    let (x, y, u) = meshgrid(data1, data1, {
        x:Float64, y:Float64 => cos(x) * sin(x + y)})
    let v = apply2(x, y, {x: Float64, y: Float64 => sin(y) * cos(x + y)})
    let m = apply2(u, v, {u: Float64, v: Float64 => sqrt(u * u + v * v)})
    quiver(x, y, u, v, m, scale:0.2).normalize(true).line_width(1.5)
    colormap(Palette.Jet)
    save("./tests/imgs/quiver/quiver_5.svg", "svg")
    clear()
}

func testQuiver3D1(): Unit {
    let U = matrix<Float64>([
        [-0.688247, -0.620174, -0.534522, -0.428571, -0.301511, -0.156174, 0.0,
         0.156174, 0.301511, 0.428571, 0.534522, 0.620174, 0.688247],
        [-0.744208, -0.680414, -0.596285, -0.486664, -0.348155, -0.182574, 0.0,
         0.182574, 0.348155, 0.486664, 0.596285, 0.680414, 0.744208],
        [-0.801784, -0.745356, -0.666667, -0.557086, -0.408248, -0.218218, 0.0,
         0.218218, 0.408248, 0.557086, 0.666667, 0.745356, 0.801784],
        [-0.857143, -0.811107, -0.742781, -0.639602, -0.485071, -0.267261, 0.0,
         0.267261, 0.485071, 0.639602, 0.742781, 0.811107, 0.857143],
        [-0.904534, -0.870388, -0.816497, -0.727607, -0.57735, -0.333333, 0.0,
         0.333333, 0.57735, 0.727607, 0.816497, 0.870388, 0.904534],
        [-0.937043, -0.912871, -0.872872, -0.801784, -0.666667, -0.408248, 0.0,
         0.408248, 0.666667, 0.801784, 0.872872, 0.912871, 0.937043],
        [-0.948683, -0.928477, -0.894427, -0.83205, -0.707107, -0.447214, -0.0,
         0.447214, 0.707107, 0.83205, 0.894427, 0.928477, 0.948683],
        [-0.937043, -0.912871, -0.872872, -0.801784, -0.666667, -0.408248, -0.0,
         0.408248, 0.666667, 0.801784, 0.872872, 0.912871, 0.937043],
        [-0.904534, -0.870388, -0.816497, -0.727607, -0.57735, -0.333333, -0.0,
         0.333333, 0.57735, 0.727607, 0.816497, 0.870388, 0.904534],
        [-0.857143, -0.811107, -0.742781, -0.639602, -0.485071, -0.267261, -0.0,
         0.267261, 0.485071, 0.639602, 0.742781, 0.811107, 0.857143],
        [-0.801784, -0.745356, -0.666667, -0.557086, -0.408248, -0.218218, -0.0,
         0.218218, 0.408248, 0.557086, 0.666667, 0.745356, 0.801784],
        [-0.744208, -0.680414, -0.596285, -0.486664, -0.348155, -0.182574, -0.0,
         0.182574, 0.348155, 0.486664, 0.596285, 0.680414, 0.744208],
        [-0.688247, -0.620174, -0.534522, -0.428571, -0.301511, -0.156174, -0.0,
         0.156174, 0.301511, 0.428571, 0.534522, 0.620174, 0.688247]
    ])

    let V = matrix<Float64>([
        [0.688247, 0.744208, 0.801784, 0.857143, 0.904534, 0.937043, 0.948683,
         0.937043, 0.904534, 0.857143, 0.801784, 0.744208, 0.688247],
        [0.620174, 0.680414, 0.745356, 0.811107, 0.870388, 0.912871, 0.928477,
         0.912871, 0.870388, 0.811107, 0.745356, 0.680414, 0.620174],
        [0.534522, 0.596285, 0.666667, 0.742781, 0.816497, 0.872872, 0.894427,
         0.872872, 0.816497, 0.742781, 0.666667, 0.596285, 0.534522],
        [0.428571, 0.486664, 0.557086, 0.639602, 0.727607, 0.801784, 0.83205,
         0.801784, 0.727607, 0.639602, 0.557086, 0.486664, 0.428571],
        [0.301511, 0.348155, 0.408248, 0.485071, 0.57735, 0.666667, 0.707107,
         0.666667, 0.57735, 0.485071, 0.408248, 0.348155, 0.301511],
        [0.156174, 0.182574, 0.218218, 0.267261, 0.333333, 0.408248, 0.447214,
         0.408248, 0.333333, 0.267261, 0.218218, 0.182574, 0.156174],
        [-0.0, -0.0, -0.0, -0.0, -0.0, -0.0, -0.0, -0.0, -0.0, -0.0, -0.0, -0.0, -0.0],
        [-0.156174, -0.182574, -0.218218, -0.267261, -0.333333, -0.408248,
         -0.447214, -0.408248, -0.333333, -0.267261, -0.218218, -0.182574,
         -0.156174],
        [-0.301511, -0.348155, -0.408248, -0.485071, -0.57735, -0.666667,
         -0.707107, -0.666667, -0.57735, -0.485071, -0.408248, -0.348155,
         -0.301511],
        [-0.428571, -0.486664, -0.557086, -0.639602, -0.727607, -0.801784,
         -0.83205, -0.801784, -0.727607, -0.639602, -0.557086, -0.486664,
         -0.428571],
        [-0.534522, -0.596285, -0.666667, -0.742781, -0.816497, -0.872872,
         -0.894427, -0.872872, -0.816497, -0.742781, -0.666667, -0.596285,
         -0.534522],
        [-0.620174, -0.680414, -0.745356, -0.811107, -0.870388, -0.912871,
         -0.928477, -0.912871, -0.870388, -0.811107, -0.745356, -0.680414,
         -0.620174],
        [-0.688247, -0.744208, -0.801784, -0.857143, -0.904534, -0.937043,
         -0.948683, -0.937043, -0.904534, -0.857143, -0.801784, -0.744208,
         -0.688247]
    ])

    let W = matrix<Float64>([
        [0.229416, 0.248069, 0.267261, 0.285714, 0.301511, 0.312348, 0.316228,
         0.312348, 0.301511, 0.285714, 0.267261, 0.248069, 0.229416],
        [0.248069, 0.272166, 0.298142, 0.324443, 0.348155, 0.365148, 0.371391,
         0.365148, 0.348155, 0.324443, 0.298142, 0.272166, 0.248069],
        [0.267261, 0.298142, 0.333333, 0.371391, 0.408248, 0.436436, 0.447214,
         0.436436, 0.408248, 0.371391, 0.333333, 0.298142, 0.267261],
        [0.285714, 0.324443, 0.371391, 0.426401, 0.485071, 0.534522, 0.5547,
         0.534522, 0.485071, 0.426401, 0.371391, 0.324443, 0.285714],
        [0.301511, 0.348155, 0.408248, 0.485071, 0.57735, 0.666667, 0.707107,
         0.666667, 0.57735, 0.485071, 0.408248, 0.348155, 0.301511],
        [0.312348, 0.365148, 0.436436, 0.534522, 0.666667, 0.816497, 0.894427,
         0.816497, 0.666667, 0.534522, 0.436436, 0.365148, 0.312348],
        [0.316228, 0.371391, 0.447214, 0.5547, 0.707107, 0.894427, 1.0, 0.894427,
         0.707107, 0.5547, 0.447214, 0.371391, 0.316228],
        [0.312348, 0.365148, 0.436436, 0.534522, 0.666667, 0.816497, 0.894427,
         0.816497, 0.666667, 0.534522, 0.436436, 0.365148, 0.312348],
        [0.301511, 0.348155, 0.408248, 0.485071, 0.57735, 0.666667, 0.707107,
         0.666667, 0.57735, 0.485071, 0.408248, 0.348155, 0.301511],
        [0.285714, 0.324443, 0.371391, 0.426401, 0.485071, 0.534522, 0.5547,
         0.534522, 0.485071, 0.426401, 0.371391, 0.324443, 0.285714],
        [0.267261, 0.298142, 0.333333, 0.371391, 0.408248, 0.436436, 0.447214,
         0.436436, 0.408248, 0.371391, 0.333333, 0.298142, 0.267261],
        [0.248069, 0.272166, 0.298142, 0.324443, 0.348155, 0.365148, 0.371391,
         0.365148, 0.348155, 0.324443, 0.298142, 0.272166, 0.248069],
        [0.229416, 0.248069, 0.267261, 0.285714, 0.301511, 0.312348, 0.316228,
         0.312348, 0.301511, 0.285714, 0.267261, 0.248069, 0.229416]
    ])

    let data = linspace(-3.0, 3.0, num:13)
    let (X, Y, Z) = meshgrid(data, data, {x:Float64, y:Float64 => pow(y, 2.0) - pow(x, 2.0)})
    quiver3(Z, U, V, W, scale: 5.0)
    view(-35.0, 45.0)
    save("./tests/imgs/quiver3d/quiver3d_1.png", "png")
    clear()
}

public func testQuiver3D2() {
    let x = vector<Float64>([
        -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0,
        -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0,
        -1.0, -1.0, -1.0, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5,
        -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5,
        -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, 0.0, 0.0, 0.0, 0.0, 0.0,
        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.5,
        0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5,
        0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5,
        0.5, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
        1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
        1.0, 1.0, 1.0, 1.0
    ])

    let y = vector<Float64>([
        -1.0, -1.0, -1.0, -1.0, -1.0, -0.5, -0.5, -0.5, -0.5, -0.5, 0.0, 0.0,
        0.0, 0.0, 0.0, 0.5, 0.5, 0.5, 0.5, 0.5, 1.0, 1.0, 1.0, 1.0,
        1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -0.5, -0.5, -0.5, -0.5, -0.5, 0.0,
        0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.5, 0.5, 0.5, 1.0, 1.0, 1.0,
        1.0, 1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -0.5, -0.5, -0.5, -0.5, -0.5,
        0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.5, 0.5, 0.5, 1.0, 1.0,
        1.0, 1.0, 1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -0.5, -0.5, -0.5, -0.5,
        -0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.5, 0.5, 0.5, 1.0,
        1.0, 1.0, 1.0, 1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -0.5, -0.5, -0.5,
        -0.5, -0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.5, 0.5, 0.5,
        1.0, 1.0, 1.0, 1.0, 1.0
    ])

    let z = vector<Float64>([
        -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0,
        -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0,
        -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0,
        -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0,
        -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0,
        -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0,
        -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0,
        -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0, -1.0, -0.5, 0.0, 0.5, 1.0,
        -1.0, -0.5, 0.0, 0.5, 1.0
    ])

    let sz = x.size()
    let u = empty<Float64>(sz)
    let v = empty<Float64>(sz)
    let w = empty<Float64>(sz)
    let m = empty<Float64>(sz)
    for (i in 0..sz) {
        let t = exp(-pow(x[i], 2.0) - pow(y[i], 2.0) - pow(z[i], 2.0))
        u[i] = x[i] * t
        v[i] = y[i] * t
        w[i] = z[i] * t
        m[i] = sqrt(x[i] * x[i] + y[i] * y[i] + z[i] * z[i])
    }
    quiver3(x, y, z, u, v, w, m).normalize(true).line_width(2.0)
    save("./tests/imgs/quiver3d/quiver3d_2.png", "png")
    clear()
}

public func testFeather1() {
    let theta = linspace(-Float64.getPI()/2.0, Float64.getPI()*2.0, num:17)
    let rho = linspace(2.0, 2.0, num:17)

    let u = apply2(rho, theta, {r:Float64, t:Float64 => r * cos(t)})
    let v = apply2(rho, theta, {r:Float64, t:Float64 => r * sin(t)})

    feather(u, v)
    save("./tests/imgs/feather/feather_1.svg", "svg")
    clear()
}

public func testQuiver() {
    testQuiver1()
    testQuiver2()
    testQuiver3()
    testQuiver4()
    testQuiver5()

    testQuiver3D1()
    testQuiver3D2()

    testFeather1()
}