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 and z: Float64 */
foreign func c_plot3(
    x: CPointer<Unit>, x_len: Int64, y: CPointer<Unit>, y_len: Int64,
    z: CPointer<Unit>, z_len: Int64, line_spec: CString): CPointer<Unit>

foreign func c_axes_plot3(
    axes: CPointer<Unit>, x: CPointer<Unit>, x_len: Int64, y: CPointer<Unit>, y_len: Int64,
    z: CPointer<Unit>, z_len: Int64, line_spec: CString): CPointer<Unit>


public func plot3(x: Vector<Float64>, y: Vector<Float64>, z: Vector<Float64>, line_spec!:String = ""): Line {
    let size = x.size()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe { c_plot3(x.ptr, size, y.ptr, size, z.ptr, size, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Line(handle)
}

public func plot3(axes: AxesType, x: Vector<Float64>, y: Vector<Float64>, z: Vector<Float64>, 
                  line_spec!:String = ""): Line {
    let size = x.size()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe { c_axes_plot3(axes.ptr, x.ptr, size, y.ptr, size, z.ptr, size, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Line(handle)
}

public func testPlot3d1() {
    let t = linspace<Float64>(0.0, 10.0 * Float64.getPI(), num: 500)
    let st = t.apply({t => sin(t)})
    let ct = t.apply({t => cos(t)})
    plot3(st, ct, t)
    save("./tests/imgs/line_plot_3d/plot3d_1.svg", "svg")
    clear()
}

public func testPlot3d2() {
    let t = linspace<Float64>(0.0, Float64.getPI(), num: 500)
    let xt1 = t.apply({x: Float64 => sin(x) * cos(10.0 * x)})
    let yt1 = t.apply({x: Float64 => sin(x) * sin(10.0 * x)})
    let zt1 = t.apply({t => cos(t)})
    plot3(xt1, yt1, zt1)
    hold(true)
    let xt2 = t.apply({x: Float64 => sin(x) * cos(12.0 * x)})
    let yt2 = t.apply({x: Float64 => sin(x) * sin(12.0 * x)})
    let zt2 = t.apply({t => cos(t)})
    plot3(xt2, yt2, zt2)
    save("./tests/imgs/line_plot_3d/plot3d_2.svg", "svg")
    clear()   
}

public func testPlot3d3() {
    let t = linspace<Float64>(0.0, Float64.getPI(), num: 500)

    let X = empty<Float64>(3, 500)
    X[0] = t.apply({x: Float64 => sin(x) * cos(10.0 * x)})
    X[1] = t.apply({x: Float64 => sin(x) * cos(12.0 * x)})
    X[2] = t.apply({x: Float64 => sin(x) * cos(20.0 * x)})

    let Y = empty<Float64>(3, 500)
    Y[0] = t.apply({x: Float64 => sin(x) * sin(10.0 * x)})
    Y[1] = t.apply({x: Float64 => sin(x) * sin(12.0 * x)})
    Y[2] = t.apply({x: Float64 => sin(x) * sin(20.0 * x)})

    let z = t.apply({t => cos(t)})

    plot3(X[0], Y[0], z)
    hold(true)
    plot3(X[1], Y[1], z)
    plot3(X[2], Y[2], z)
    save("./tests/imgs/line_plot_3d/plot3d_3.svg", "svg")
    clear()
}

public func testPlot3d4() {
    let t = linspace<Float64>(0.0, 40.0 * Float64.getPI(), num: 20000)
    let xt = t.apply({x: Float64 => (3.0 + cos(sqrt(32.0) * x)) * cos(x)})
    let yt = t.apply({x: Float64 => sin(sqrt(32.0) * x)})
    let zt = t.apply({x: Float64 => (3.0 + cos(sqrt(32.0) * x)) * sin(x)})
    plot3(xt, yt, zt)
    axis(AxisStyle.Equal)
    xlabel("x(t)")
    ylabel("y(t)")
    zlabel("z(t)")
    save("./tests/imgs/line_plot_3d/plot3d_4.svg", "svg")
    clear() 
}

public func testPlot3d5() {
    let t = linspace<Float64>(0.0, 10.0 * Float64.getPI(), num: 200)
    let xt = t.apply({t => sin(t)})
    let yt = t.apply({t => cos(t)})
    plot3(xt, yt, t, line_spec: "-ob").marker_size(10.0).marker_face_color("#D9FFFF")
    xlabel("x(t)")
    ylabel("y(t)")
    zlabel("t")
    save("./tests/imgs/line_plot_3d/plot3d_5.svg", "svg")
    clear()    
}

public func testPlot3d6() {
    let t = linspace<Float64>(0.0, 10.0 * Float64.getPI(), num: 200)
    let xt1 = t.apply({t => sin(t)})
    let yt1 = t.apply({t => cos(t)})
    let xt2 = t.apply({t => sin(2.0 * t)})
    let yt2 = t.apply({t => cos(2.0 * t)})
    plot3(xt1, yt1, t, line_spec: "--")
    hold(true)
    plot3(xt2, yt2, t, line_spec: "--")
    save("./tests/imgs/line_plot_3d/plot3d_6.svg", "svg")
    clear()
}

public func testPlot3d7() {
    let t = linspace<Float64>(-10.0, 10.0, num: 1000)
    let xt = t.apply({x: Float64 => exp(-x / 10.0) * sin(5.0 * x)})
    let yt = t.apply({x: Float64 => exp(-x / 10.0) * cos(5.0 * x)})
    plot3(xt, yt, t).line_width(3.0)
    save("./tests/imgs/line_plot_3d/plot3d_7.svg", "svg")
    clear()
}

public func testPlot3d8() {
    tiledlayout(1, 2)
    let ax1 = nexttile()

    let t1 = linspace<Float64>(0.0, 10.0 * Float64.getPI(), num: 200)
    let xt1 = t1.apply({t => sin(t)})
    let yt1 = t1.apply({t => cos(t)})
    plot3(ax1, xt1, yt1, t1)
    title(ax1, "Helix with 5 Turns")

    let ax2 = nexttile();
    let t2 = linspace<Float64>(0.0, 10.0 * Float64.getPI(), num: 200)
    let xt2 = t2.apply({x: Float64 => sin(2.0 * x)})
    let yt2 = t2.apply({x: Float64 => cos(2.0 * x)})
    plot3(ax2, xt2, yt2, t2);
    ax2.box(false);
    title(ax2, "Helix with 10 Turns");
    save("./tests/imgs/line_plot_3d/plot3d_8.svg", "svg")
    clear()
}

public func testPlot3d9() {
    let xt = linspace<Float64>(0.0, 10.0, num: 10)
    let yt = xt
    let zt = xt
    plot3(xt, yt, zt, line_spec: "o")

    xlabel("X")
    ylabel("Y")
    zlabel("Duration")
    grid(true)
    save("./tests/imgs/line_plot_3d/plot3d_9.svg", "svg")
    clear()
}

public func testPlot3d10() {
    let t = linspace<Float64>(0.0, Float64.getPI(), num: 500)
    let xt = t.apply({x: Float64 => sin(x) * cos(10.0 * x)})
    let yt = t.apply({x: Float64 => sin(x) * sin(10.0 * x)})
    let zt = t.apply({t => cos(t)})
    let p = plot3(xt, yt, zt, line_spec: "-o")
    p.marker_indices(vector<Int64>([200]))
    save("./tests/imgs/line_plot_3d/plot3d_10.svg", "svg")
    clear()
}


public func testPlot3d() {
    testPlot3d1()
    testPlot3d2()
    testPlot3d3()
    testPlot3d4()
    testPlot3d5()
    testPlot3d6()
    testPlot3d7()
    testPlot3d8()
    testPlot3d9()
    testPlot3d10()
}