package scientific.matplot

import std.collection.*
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 and y: Float64 */
foreign func c_stemplot(
    x: CPointer<Unit>, x_len: Int64, 
    y: CPointer<Unit>, y_len: Int64, line_spec: CString): CPointer<Unit>

foreign func c_stemplot_ys(
    y: CPointer<Unit>, row: Int64, col: Int64,
    line_spec: CString): CPointer<CPointer<Unit>>

foreign func c_stemplot_x_ys(
    x: CPointer<Unit>, x_len: Int64,
    y: CPointer<Unit>, row: Int64, col: Int64,
    line_spec: CString): CPointer<CPointer<Unit>>

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

public func stem(x: Vector<Float64>, y: Vector<Float64>, line_spec!: String = "-o"): Line {
    if (!(x.size() == y.size())) {
        throw IllegalArgumentException("stem plot: 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_stemplot(x.ptr, size, y.ptr, size, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Line(handle)
}

public func stem(y: Matrix<Float64>, line_spec!: String = "-o"): Array<Line> {
    let row = y.getRows()
    let col = y.getCols()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handles = unsafe { c_stemplot_ys(y.ptr, row, col, cstr_line_spec) }
    let res = ArrayList<Line>()
    for (i in 0..row) {
        res.add(Line(unsafe { handles.read(i) }))
    }
    unsafe { LibC.free(cstr_line_spec) }
    return res.toArray()
}

public func stem(x: Vector<Float64>, y: Matrix<Float64>, line_spec!: String = "-o"): Array<Line> {
    let size = x.size()
    let row = y.getRows()
    let col = y.getCols()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handles = unsafe { c_stemplot_x_ys(x.ptr, size, y.ptr, row, col, cstr_line_spec) }
    let res = ArrayList<Line>()
    for (i in 0..row) {
        res.add(Line(unsafe {handles.read(i)}))
    }
    unsafe { LibC.free(cstr_line_spec) }
    return res.toArray()
}

public func stem(axes: AxesType, x: Vector<Float64>, y: Vector<Float64>, line_spec!: String = "-o"): Line {
    if (!(x.size() == y.size())) {
        throw IllegalArgumentException("stem plot: 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_stemplot(axes.ptr, x.ptr, size, y.ptr, size, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Line(handle)
}

public func testStemPlot1() {
    let x = linspace<Float64>(0.0, 2.0 * Float64.getPI(), num: 50)

    let Y = empty<Float64>(2, 50);
    Y[0] = x.apply({t => cos(t)});
    Y[1] = x.apply({x => 0.5 * sin(x)})

    stem(Y, line_spec: "-o")
    save("./tests/imgs/stem_plot/stem_plot_1.svg", "svg")
    clear()
}

public func testStemPlot2() {
    let x = linspace(0.0, 2.0 * Float64.getPI(), num: 50)
    let y = x.apply({t => cos(t)})

    stem(x, y)
    save("./tests/imgs/stem_plot/stem_plot_2.svg", "svg")
    clear()
}

public func testStemPlot3() {
    let x = linspace<Float64>(0.0, 2.0 * Float64.getPI(), num: 50)

    let Y = empty<Float64>(2, 50);
    Y[0] = x.apply({t => cos(t)});
    Y[1] = x.apply({x => 0.5 * sin(x)})

    stem(x, Y)
    save("./tests/imgs/stem_plot/stem_plot_3.svg", "svg")
    clear()
}

public func testStemPlot4() {
    let x1 = linspace(0.0, 2.0 * Float64.getPI(), num: 50)
    let y1 = x1.apply({t => cos(t)})

    let x2 = linspace(Float64.getPI(), 3.0 * Float64.getPI(), num: 50)
    let y2 = x2.apply({t => sin(t)}) * 0.5

    stem(x1, y1);
    hold(true);
    stem(x2, y2);

    save("./tests/imgs/stem_plot/stem_plot_4.svg", "svg")
    clear()
}

public func testStemPlot5() {
    let x = linspace(0.0, 10.0, num: 20)
    let y = x.apply({x: Float64 => exp(0.25 * x)})

    stem(x, y, line_spec: "filled")
    save("./tests/imgs/stem_plot/stem_plot_5.svg", "svg")
    clear()
}

public func testStemPlot6() {
    let x = linspace(0.0, 2.0 * Float64.getPI(), num: 50)
    let y = x.apply({x: Float64 => exp(x) * sin(x)})

    stem(x, y, line_spec: ":dr");
    save("./tests/imgs/stem_plot/stem_plot_6.svg", "svg")
    clear()
}

public func testStemPlot7() {
    let x = linspace(0.0, 2.0 * Float64.getPI(), num: 25)
    let y = x.apply({x: Float64 => cos(2.0 * x)})

    stem(x, y).line_style("-.").marker_face_color("red").marker_color("green")

    save("./tests/imgs/stem_plot/stem_plot_7.svg", "svg")
    clear()
}

public func testStemPlot8() {
    let x = linspace(0.0, 25.0, num: 26)
    let y1 = x.apply({x: Float64 => exp(0.1 * x)})
    let y2 = x.apply({x: Float64 => -exp(0.05 * x)})

    tiledlayout(2, 1)
    let ax1 = nexttile()
    stem(ax1, x, y1)

    let ax2 = nexttile()
    stem(ax2, x, y2)

    save("./tests/imgs/stem_plot/stem_plot_8.svg", "svg")
    clear()
}

public func testStemPlot9() {
    let x = linspace(0.0, 2.0 * Float64.getPI(), num: 50)
    let y = x.apply({x: Float64 => exp(0.3 * x) * sin(3.0 * x)})

    stem(x, y)
    gca().x_axis().zero_axis(false)

    save("./tests/imgs/stem_plot/stem_plot_9.svg", "svg")
    clear()
}

public func testStemPlot(){
    testStemPlot1()
    testStemPlot2()
    testStemPlot3()
    testStemPlot4()
    testStemPlot5()
    testStemPlot6()
    testStemPlot7()
    testStemPlot8()
    testStemPlot9()
}