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_stairs_y(
    y: CPointer<Unit>, y_len: Int64, line_spec: CString): CPointer<Unit>
foreign func c_stairs(
    x: CPointer<Unit>, x_len: Int64, y: CPointer<Unit>, y_len: Int64,
    line_spec: CString): CPointer<Unit>
foreign func c_stairs_ys(
    y: CPointer<Unit>, row: Int64, col: Int64,
    line_spec: CString): CPointer<CPointer<Unit>>
foreign func c_stairs_x_ys(
    x: CPointer<Unit>, x_len: Int64, y: CPointer<Unit>, row: Int64, col: Int64,
    line_spec: CString): CPointer<CPointer<Unit>>
foreign func c_stairs_xs_ys(
    x: CPointer<Unit>, y: CPointer<Unit>, row: Int64, col: Int64,
    line_spec: CString): CPointer<CPointer<Unit>>
foreign func c_axes_stairs(
    axes: CPointer<Unit>, x: CPointer<Unit>, x_len: Int64, y: CPointer<Unit>, y_len: Int64,
    line_spec: CString): CPointer<Unit>
foreign func c_stair_style(
    stair: CPointer<Unit>, c_style: CString): CPointer<Unit>

public class Stair <: Line {
    /* Input is a pointer to matplot::line. */
    init(ptr: CPointer<Unit>) {
        super(ptr)
    }

    public func stair_style(style: StairStyle): Stair {
        var cstr_style = unsafe { LibC.mallocCString(style.toString()) }
        this.ptr = unsafe { c_stair_style(this.ptr, cstr_style) }
        unsafe { LibC.free(cstr_style) }
        return this
    }

    public override func line_width(line_width: Float32): Stair {
        this.ptr = unsafe { c_line_width(this.ptr, line_width) }
        return this
    }

    public override func marker(c: String): Stair {
        var cstr = unsafe { LibC.mallocCString(c) }
        this.ptr = unsafe { c_marker_str(this.ptr, cstr) }
        unsafe { LibC.free(cstr) }
        return this
    }

    public override func marker_face_color(c: String): Stair {
        var cstr = unsafe { LibC.mallocCString(c) }
        this.ptr = unsafe { c_marker_face_color_str(this.ptr, cstr) }
        unsafe { LibC.free(cstr) }
        return this
    }

    public override func marker(style: MarkerStyle): Stair {
        return this.marker(style.toString())
    }

    public override func marker_size(marker_size: Float32): Stair {
        this.ptr = unsafe { c_marker_size(this.ptr, marker_size) }
        return this
    }

    public override func marker_color(c: String): Stair {
        var cstr = unsafe { LibC.mallocCString(c) }
        this.ptr = unsafe { c_marker_color_str(this.ptr, cstr) }
        unsafe { LibC.free(cstr) }
        return this
    }

    public override func marker_face(v: Bool): Stair {
        this.ptr = unsafe { c_marker_face(this.ptr, v) }
        return this
    }
}

public enum StairStyle {
    | TraceXFirst
    | TraceYFirst
    | TraceHistogram
    | Fill

    func toString(): String {
        return match (this) {
            case TraceXFirst => "trace_x_first"
            case TraceYFirst => "trace_y_first"
            case TraceHistogram => "histogram"
            case Fill => "fill"
        }
    }
}

public func stairs(y: Vector<Float64>, line_spec!:String = ""): Stair {
    let size = y.size()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe { c_stairs_y(y.ptr, size, cstr_line_spec) }
    unsafe { LibC.free(cstr_line_spec) }
    return Stair(handle)
}

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

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

public func stairs(x: Vector<Float64>, y: Matrix<Float64>, line_spec!:String = ""): Array<Stair> {
    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_stairs_x_ys(x.ptr, size, y.ptr, row, col, cstr_line_spec) }
    let res = ArrayList<Stair>()
    for (i in 0..row) {
        res.add(Stair(unsafe {handles.read(i)}))
    }
    unsafe { LibC.free(cstr_line_spec) }
    return res.toArray()
}

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

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

public func testStairs1() {
    let x = linspace(0.0, 4.0 * Float64.getPI(), num: 40);
    let y = x.apply({t => sin(t)})
    figure();
    stairs(y);

    save("./tests/imgs/stairs/stairs_1.svg", "svg")
    clear()
}

public func testStairs2() {
    let x = linspace(0.0, 4.0 * Float64.getPI(), num: 40)
    let Y = empty<Float64>(2, 40)
    Y[0] = x.apply({x => 0.5 * cos(x)})
    Y[1] = x.apply({x => 2.0 * cos(x)})

    figure()
    stairs(Y)
    save("./tests/imgs/stairs/stairs_2.svg", "svg")
    clear()
}

public func testStairs3() {
    let x = linspace(0.0, 4.0 * Float64.getPI(), num: 40)
    let y = x.apply({t => sin(t)})

    figure()
    stairs(x, y)
    save("./tests/imgs/stairs/stairs_3.svg", "svg")
    clear()
}

public func testStairs4() {
    let x = linspace(0.0, 4.0 * Float64.getPI(), num: 50)
    let Y = empty<Float64>(2, 50)
    Y[0] = x.apply({x => 0.5 * cos(x)})
    Y[1] = x.apply({x => 2.0 * cos(x)})

    figure()
    stairs(x, Y)
    save("./tests/imgs/stairs/stairs_4.svg", "svg")
    clear()
}


public func testStairs5() {
    let X = empty<Float64>(2, 100)
    X[0] = linspace(0.0, 2.0 * Float64.getPI())
    X[1] = linspace(0.0, Float64.getPI())

    let Y = empty<Float64>(2, 100)
    Y[0] = X[0].apply({x => sin(5.0 * x)})
    Y[1] = X[1].apply({x => exp(x) * sin(5.0 * x)})

    figure()
    stairs(X, Y)
    save("./tests/imgs/stairs/stairs_5.svg", "svg")
    clear()
}

public func testStairs6() {
    let x = linspace(0.0, 4.0 * Float64.getPI(), num: 20)
    let y = x.apply({t => sin(t)})

    figure()
    stairs(x, y, line_spec: "-.or")
    save("./tests/imgs/stairs/stairs_6.svg", "svg")
    clear()
}

public func testStairs7() {
    let x = linspace(0.0, 4.0 * Float64.getPI(), num: 20)
    let y = x.apply({t => sin(t)})
    stairs(x, y).line_width(2.0).marker("d").marker_face_color("c")
    save("./tests/imgs/stairs/stairs_7.svg", "svg")
    clear()
}

public func testStairs8() {
    let x = linspace(0.0, 2.0 * Float64.getPI())
    let y1 = x.apply({ x => 5.0 * sin(x) })
    let y2 = x.apply({ x => sin(5.0 * x) })
    tiledlayout(2, 1)

    let ax1 = nexttile()
    stairs(ax1, x, y1)
    let ax2 = nexttile()
    stairs(ax2, x, y2)
    save("./tests/imgs/stairs/stairs_8.svg", "svg")
    clear()
}

public func testStairs9() {
    let x = linspace(0.0, 1.0, num:30)
    let Y = empty<Float64>(2, 30);
    Y[0] = x.apply({ x => cos(10.0 * x) })
    Y[1] = x.apply({ x => exp(x) * sin(10.0 * x) })
    let h: Array<Stair> = stairs(x, Y)
    h[0].marker(MarkerStyle.Circle).marker_size(4.0)
    h[1].marker(MarkerStyle.Circle).marker_face_color("m")
    save("./tests/imgs/stairs/stairs_9.svg", "svg")
    clear()
}

public func testStairs10() {
    let x = vector<Float64>([1.0, 3.0, 5.0, 7.0, 10.0])
    let y = vector<Float64>([2.0, 5.0, 6.0, 7.0, 11.0])
    let h1: Stair = stairs(x, y)
    let h2: Stair = stairs(x, y)
    let h3: Stair = stairs(x, y)
    let h4: Stair = stairs(x, y)
    h1.stair_style(StairStyle.Fill)
    h2.stair_style(StairStyle.TraceXFirst).line_width(4.0)
    h3.stair_style(StairStyle.TraceYFirst).line_width(2.0)
    h4.stair_style(StairStyle.TraceHistogram).marker(MarkerStyle.Circle)
        .line_width(1.0).marker_color("m").marker_face(true)
        .marker_size(10.0)
    save("./tests/imgs/stairs/stairs_10.svg", "svg")
    clear()
}

public func testStairs() {
    testStairs1()
    testStairs2()
    testStairs3()
    testStairs4()
    testStairs5()
    testStairs6()
    testStairs7()
    testStairs8()
    testStairs9()
    testStairs10()
}