package scientific.matplot

import std.collection.*
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, label_x, label_y: Float64 */
foreign func c_bar_y(y: CPointer<Unit>, y_len: Int64): CPointer<Unit>
foreign func c_bar_x_y(x: CPointer<Unit>, y: CPointer<Unit>, len: Int64): CPointer<Unit>
foreign func c_bar_y_num(y: CPointer<Unit>, y_len: Int64, numbers: Float64): CPointer<Unit>
foreign func c_bar_ys(y: CPointer<Unit>, row: Int64, col: Int64): CPointer<Unit>
foreign func c_bar_x_ys(x: CPointer<Unit>, len: Int64, y: CPointer<Unit>, row: Int64, col: Int64): CPointer<Unit>
foreign func c_axes_bar_ys(axes: CPointer<Unit>, y: CPointer<Unit>, row: Int64, col: Int64): CPointer<Unit>
foreign func c_barstacked_ys(y: CPointer<Unit>, row: Int64, col: Int64): Unit
foreign func c_axes_barstacked_ys(axes: CPointer<Unit>, y: CPointer<Unit>, row: Int64, col: Int64): Unit
foreign func c_barstacked_x_ys(x: CPointer<Unit>, len: Int64, y: CPointer<Unit>, row: Int64, col: Int64): Unit
foreign func c_bar_face_color(bars_handle: CPointer<Unit>, color: CString): CPointer<Unit>
foreign func c_bar_face_color_rgb(bars_handle: CPointer<Unit>, r: Float32, g: Float32, b: Float32): CPointer<Unit>
foreign func c_bar_edge_color_rgb(bars_handle: CPointer<Unit>, r: Float32, g: Float32, b: Float32): CPointer<Unit>
foreign func c_bar_line_width(bars_handle: CPointer<Unit>, width: Float32): CPointer<Unit>
foreign func c_x_end_point(bars_handle: CPointer<Unit>, cluster_index: Int64, index: Int64): Float64
foreign func c_text(
    label_x: CPointer<Unit>, label_y: CPointer<Unit>, labels: CPointer<CString>, len: Int64): Unit
foreign func c_bar_face_colors(
    bars_handle: CPointer<Unit>, index: Int64,
    r: Float32, g: Float32, b: Float32, a: Float32): Unit

public func text(label_x: ArrayList<Float64>, label_y: ArrayList<Float64>,
                 labels: ArrayList<String>): Unit {
    let vec_label_x = vector<Float64>(label_x.toArray())
    let vec_label_y = vector<Float64>(label_y.toArray())

    let size = label_x.size
    var a = unsafe { malloc(UIntNative(8 * size)) }
    var ptr = CPointer<CString>(a)
    for (i in 0..size) {
        var cstr_label = unsafe { LibC.mallocCString(labels[i]) }
        unsafe { ptr.write(i, cstr_label) }
        unsafe { LibC.free(cstr_label) }
    }

    unsafe { c_text(vec_label_x.ptr, vec_label_y.ptr, ptr, size) }
}

public func text(label_x: Vector<Float64>, label_y: Vector<Float64>,
                 labels: ArrayList<String>): Unit {
    var ptr = unsafe { CPointer<CString>(malloc(UIntNative(8 * label_x.size()))) }
    for (i in 0..label_x.size()) {
        var cstr_label = unsafe { LibC.mallocCString(labels[i]) }
        unsafe { ptr.write(i, cstr_label) }
        unsafe { LibC.free(cstr_label) }
    }

    unsafe { c_text(label_x.ptr, label_y.ptr, ptr, label_x.size()) }
}

public func text(label_x: Float64, label_y: Float64, label: String): Unit {
    text(ArrayList([label_x]), ArrayList([label_y]), ArrayList([label]))
}

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

    init(ptr: CPointer<Unit>) {
        this.ptr = ptr
    }

    public func face_color(color: String) {
        var cstr_color = unsafe { LibC.mallocCString(color) }
        this.ptr = unsafe { c_bar_face_color(this.ptr, cstr_color) }
        unsafe { LibC.free(cstr_color) }
        return this
    }

    public func face_color(r: Float32, g: Float32, b: Float32) {
        this.ptr = unsafe { c_bar_face_color_rgb(this.ptr, r, g, b) }
        return this
    }

    public func edge_color(r: Float32, g: Float32, b: Float32) {
        this.ptr = unsafe { c_bar_edge_color_rgb(this.ptr, r, g, b) }
        return this
    }

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

    public func x_end_point(cluster_index: Int64, index: Int64): Float64 {
        return unsafe{ c_x_end_point(this.ptr, cluster_index, index) }
    }

    public func face_colors(index: Int64, r: Float32, g: Float32, b: Float32, a: Float32) {
        unsafe { c_bar_face_colors(this.ptr, index, r, g, b, a) }
    }

}

public func bar(y: Vector<Float64>): Bars {
    let size = y.size()
    let handle = unsafe { c_bar_y(y.ptr, size) }
    return Bars(handle)
}

public func bar(x: Vector<Float64>, y: Vector<Float64>): Bars {
    let size = x.size()
    let handle = unsafe{ c_bar_x_y(x.ptr, y.ptr, size) }
    return Bars(handle)
}

public func bar(y: Vector<Float64>, size: Float64): Bars {
    let handle = unsafe { c_bar_y_num(y.ptr, y.size(), size) }
    return Bars(handle)
}

public func bar(Y: Matrix<Float64>): Bars {
    let handle = unsafe { c_bar_ys(Y.ptr, Y.getRows(), Y.getCols()) }
    return Bars(handle)
}

public func bar(x: Vector<Float64>, Y: Matrix<Float64>): Bars {
    let size = x.size()
    let handle = unsafe { c_bar_x_ys(x.ptr, size, Y.ptr, Y.getRows(), Y.getCols()) }
    return Bars(handle)
}

public func bar(axes: AxesType, Y: Matrix<Float64>): Bars {
    let handle = unsafe { c_axes_bar_ys(axes.ptr, Y.ptr, Y.getRows(), Y.getCols()) }
    return Bars(handle)
}

public func barstacked(Y: Matrix<Float64>) {
    unsafe { c_barstacked_ys(Y.ptr, Y.getRows(), Y.getCols()) }
}

public func barstacked(axes: AxesType, Y: Matrix<Float64>) {
    unsafe { c_axes_barstacked_ys(axes.ptr, Y.ptr, Y.getRows(), Y.getCols()) }
}

public func barstacked(x: Vector<Float64>, Y: Matrix<Float64>) {
    unsafe { c_barstacked_x_ys(x.ptr, x.size(), Y.ptr, Y.getRows(), Y.getCols()) }
}


public func testBar1() {
    let y = vector<Float64>([75.0, 91.0, 105.0, 123.5, 131.0, 150.0, 179.0,
                             203.0, 226.0, 249.0, 281.5])
    bar(y)
    save("./tests/imgs/bar_plot/bar_1.svg", "svg")
    clear()
}

public func testBar2() {
    let x = linspace(1900.0, 2000.0, num: 11)
    let y = vector<Float64>([75.0, 91.0, 105.0, 123.5, 131.0, 150.0, 179.0,
                             203.0, 226.0, 249.0, 281.5])
    bar(x, y)
    save("./tests/imgs/bar_plot/bar_2.svg", "svg")
    clear()
}

public func testBar3() {
    let y = vector<Float64>([75.0, 91.0, 105.0, 123.5, 131.0, 150.0, 179.0,
                             203.0, 226.0, 249.0, 281.5])
    bar(y, 0.4)
    save("./tests/imgs/bar_plot/bar_3.svg", "svg")
    clear()
}

public func testBar4() {
    let Y = matrix([[2.0, 2.0, 2.0, 2.0],
                    [2.0, 5.0, 8.0, 11.0],
                    [3.0, 6.0, 9.0, 12.0]])
    bar(Y)
    save("./tests/imgs/bar_plot/bar_4.svg", "svg")
    clear()
}

public func testBar5() {
    let Y = matrix([[2.0, 2.0, 2.0, 2.0],
                    [2.0, 5.0, 8.0, 11.0],
                    [3.0, 6.0, 9.0, 12.0]])
    barstacked(Y)
    save("./tests/imgs/bar_plot/bar_5.svg", "svg")
    clear()
}

public func testBar6() {
    let x = vector<Float64>([1980.0, 1990.0, 2000.0])
    let Y = matrix([[15.0, 10.0, -10.0],
                    [20.0, -17.0, 5.0],
                    [-5.0, 21.0, 15.0]])
    barstacked(x, Y)
    save("./tests/imgs/bar_plot/bar_6.png", "png")
    clear()
}

public func testBar7() {
    let y = vector<Float64>([10.0, 21.0, 33.0, 52.0])
    bar(y)
    gca().x_axis().ticklabels(["Small", "Medium", "Large", "Extra Large"])
    
    save("./tests/imgs/bar_plot/bar_7.png", "png")
    clear()
}

public func testBar8() {
    let x = vector<Float64>([1.0, 2.0, 3.0])
    let y = matrix([[2.0, 3.0, 6.0],
                    [11.0, 23.0, 26.0]])
    let b = bar(x, y)

    var label_x = ArrayList<Float64>()
    var label_y = ArrayList<Float64>()
    var labels = ArrayList<String>()
    let x_size = x.size()
    let y_size = y.getRows()
    for (i in 0..y_size) {
        for (j in 0..x_size) {
            label_x.add(b.x_end_point(i, j))
            label_y.add(y[i, j] + 1.0)
            labels.add(y[i, j].toString())
        }
    }

    hold(true)
    text(label_x, label_y, labels)

    save("./tests/imgs/bar_plot/bar_8.png", "png")
    clear()
}

public func testBar9() {
    let y = matrix([[1.0, 4.0],
                    [2.0, 5.0],
                    [3.0, 6.0]])

    tiledlayout(2, 1)

    let ax1 = nexttile()
    bar(ax1, y)

    let ax2 = nexttile()
    barstacked(ax2, y)

    save("./tests/imgs/bar_plot/bar_9.png", "png")
    clear()
}

public func testBar10() {
    let y = vector<Float64>([75.0, 91.0, 105.0, 123.5, 131.0, 150.0,
                            179.0, 203.0, 226.0, 249.0, 281.5])

    bar(y).face_color("r")

    save("./tests/imgs/bar_plot/bar_10.png", "png")
    clear()
}

public func testBar11() {
    let y = vector<Float64>([75.0, 91.0, 105.0, 123.5, 131.0, 150.0,
                            179.0, 203.0, 226.0, 249.0, 281.5])

    bar(y).face_color(0.0, 0.5, 0.5).edge_color(0.0, 0.9, 0.9).line_width(1.5)

    save("./tests/imgs/bar_plot/bar_11.png", "png")
    clear()
}

public func testBar12() {
    let x1 = concat(vector<Float64>([1.0]), linspace(3.0, 10.0, num: 8))
    var r: Random = Random(0)
    bar(x1, rand(r, 9, 0.0, 1.0))
    
    hold(true)
    bar(vector<Float64>([2.0]), rand(r, 1, 0.0, 1.0))
    gca().x_axis().tick_values(linspace(1.0, 10.0, num: 10))

    save("./tests/imgs/bar_plot/bar_12.png", "png")
    clear()
}

public func testBar13() {
    let y = matrix(
        [[10.0, 30.0, 50.0],
        [15.0, 35.0, 55.0],
        [20.0, 40.0, 62.0]])

    let b = bar(y)
    b.face_colors(2, 0.0, 0.2, 0.6, 0.5)
    gcf().draw()

    save("./tests/imgs/bar_plot/bar_13.png", "png")
    clear()
}
public func testBarPlot() {
    print("  + testBarPlot\n")
    testBar1()
    testBar2()
    testBar3()
    testBar4()
    testBar5()
    testBar6()
    testBar7()
    testBar8()
    testBar9()
    testBar10()
    testBar11()
    testBar12()
    testBar13()
}