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.*

/* Return pointer to matplot::line.
   Type for x, y and z: Float64 */
foreign func c_stemplot3(
    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_stemplot3_z(
    z: CPointer<Unit>, z_len: Int64, line_spec: CString): CPointer<Unit>
foreign func c_stemplot3_zs(
    z: CPointer<Unit>, row: Int64, col: Int64,
    line_spec: CString): CPointer<Unit>
foreign func c_axes_stemplot3(
    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 stem3(x: Vector<Float64>, y: Vector<Float64>, 
    z: Vector<Float64>, line_spec!: String = "-o"): Line {
    if (!(x.size() == y.size()) || !(y.size() == z.size())) {
        throw IllegalArgumentException("stem plot 3d: input are not vectors of the same size.")
    }

    let size = x.size()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe{ c_stemplot3(x.ptr, size, y.ptr, size, z.ptr, size, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Line(handle)
}

public func stem3(z: Vector<Float64>, line_spec!: String = "-o"): Line {
    let size = z.size()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe{ c_stemplot3_z(z.ptr, size, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Line(handle)
}

public func stem3(z: Matrix<Float64>, line_spec!: String = "-o"): Line {
    let row = z.getRows()
    let col = z.getCols()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe { c_stemplot3_zs(z.ptr, row, col, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Line(handle)
}

public func stem3(axes: AxesType, x: Vector<Float64>, y: Vector<Float64>, 
                  z: Vector<Float64>, line_spec!: String = "-o"): Line {
    if (!(x.size() == y.size()) || !(y.size() == z.size())) {
        throw IllegalArgumentException("stem plot 3d: input are not vectors of the same size.")
    }
    let size = x.size()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe { c_axes_stemplot3(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 testStemPlot3d1() {
    let x = linspace<Float64>(-Float64.getPI() / 2.0, Float64.getPI() / 2.0, num: 40)
    let z = x.apply({t => cos(t)})

    stem3(z)
    save("./tests/imgs/stem_plot_3d/stem_plot_3d_1.svg", "svg")
    clear()
}

public func testStemPlot3d2() {
    let x = vector<Float64>(40, 1.0)
    let y = linspace<Float64>(-Float64.getPI() / 2.0, Float64.getPI() / 2.0, num: 40)
    let z = y.apply({t => cos(t)})

    stem3(x, y, z)
    save("./tests/imgs/stem_plot_3d/stem_plot_3d_2.svg", "svg")
    clear()
}

public func testStemPlot3d3() {
    let x = linspace<Float64>(-Float64.getPI() / 2.0, Float64.getPI() / 2.0, num: 40)

    let Z = empty<Float64>(2, 40)
    Z[0] = x.apply({t => sin(t)});
    Z[1] = x.apply({t => cos(t)});

    stem3(Z)
    save("./tests/imgs/stem_plot_3d/stem_plot_3d_3.svg", "svg")
    clear()
}

public func testStemPlot3d4() {
    let x = linspace<Float64>(-5.0, 5.0, num: 60)
    let y = x.apply({t => cos(t)})
    let z = x.apply({x => pow(x, 2.0)})
    
    stem3(x, y, z)
    view(-8.0, 30.0)
    save("./tests/imgs/stem_plot_3d/stem_plot_3d_4.svg", "svg")
    clear()
}

public func testStemPlot3d5() {
    let grid_ticks = linspace<Float64>(0.0, 1.0, num: 11)
    let size = 11
    let x = vector<Float64>(size * size, 1.0)
    let y = vector<Float64>(size * size, 1.0)
    let z = vector<Float64>(size * size, 1.0)
    var k = 0
    for (i in 0..size) {
        for (j in 0..size) {
            x[k] = grid_ticks[i]
            y[k] = grid_ticks[j]
            z[k] = exp(grid_ticks[i] + grid_ticks[j])
            k = k + 1
        }
    }

    stem3(x, y, z)
    save("./tests/imgs/stem_plot_3d/stem_plot_3d_5.svg", "svg")
    clear()
}

public func testStemPlot3d6() {
    let x = linspace<Float64>(-Float64.getPI(), Float64.getPI(), num: 40)
    let z = x.apply({t => cos(t)})

    stem3(z, line_spec: "filled")
    save("./tests/imgs/stem_plot_3d/stem_plot_3d_6.svg", "svg")
    clear()
}

public func testStemPlot3d7() {
    let x = linspace<Float64>(-Float64.getPI(), Float64.getPI(), num: 40)
    let z = x.apply({t => cos(t)})

    stem3(z, line_spec: "--*m")
    save("./tests/imgs/stem_plot_3d/stem_plot_3d_7.svg", "svg")
    clear()
}

public func testStemPlot3d8() {
    let theta = linspace<Float64>(0.0, 2.0 * Float64.getPI())
    let x = theta.apply({t => cos(t)})
    let y = theta.apply({t => sin(t)})
    let z = theta

    stem3(x, y, z, line_spec: ":*m")
    save("./tests/imgs/stem_plot_3d/stem_plot_3d_8.svg", "svg")
    clear()
}

public func testStemPlot3d9() {
    let x = linspace<Float64>(-Float64.getPI(), Float64.getPI(), num: 40)
    let z = x.apply({t => cos(t)})

    stem3(z).marker("s").marker_color("m").marker_face_color("g")
    save("./tests/imgs/stem_plot_3d/stem_plot_3d_9.svg", "svg")
    clear()
}

public func testStemPlot3d10() {
    let x = linspace<Float64>(-2.0, 2.0, num: 50)
    let y = x.apply({x => pow(x, 3.0)})
    let z = x.apply(exp)

    tiledlayout(2, 1)

    let ax1 = nexttile()
    stem(ax1, x, z)

    let ax2 = nexttile()
    stem3(ax2, x, y, z)
    
    save("./tests/imgs/stem_plot_3d/stem_plot_3d_10.svg", "svg")
    clear()
}

public func testStemPlot3d11() {
    let x = linspace<Float64>(0.0, 2.0)
    let y = x.apply({x => pow(x, 3.0)})
    let z = apply2(x, y, {x: Float64, y: Float64 => exp(x) * cos(y)})
    
    stem3(x, y, z, line_spec: "filled").color("m").marker_face_color("y")
    view(-10.0, 35.0)

    save("./tests/imgs/stem_plot_3d/stem_plot_3d_11.svg", "svg")
    clear()
}

public func testStemPlot3d() {
    testStemPlot3d1()
    testStemPlot3d2()
    testStemPlot3d3()
    testStemPlot3d4()
    testStemPlot3d5()
    testStemPlot3d6()
    testStemPlot3d7()
    testStemPlot3d8()
    testStemPlot3d9()
    testStemPlot3d10()
    testStemPlot3d11()
}