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

/* x, y and errors have type Float64 */
foreign func c_errorbar(
    x: CPointer<Unit>, y: CPointer<Unit>, error: CPointer<Unit>,
    len: Int64, errorbar_type: CString, line_spec: CString): CPointer<Unit>

foreign func c_errorbar2(
    x: CPointer<Unit>, y: CPointer<Unit>,
    y_neg_delta: CPointer<Unit>, y_pos_delta: CPointer<Unit>,
    x_neg_delta: CPointer<Unit>, x_pos_delta: CPointer<Unit>, len: Int64,
    line_spec: CString): CPointer<Unit>

/* Accept and return pointers to matplot::error_bar. */
foreign func c_filled_curve(errorbar: CPointer<Unit>, filled_curve: Bool): CPointer<Unit>
foreign func c_cap_size(errorbar: CPointer<Unit>, cap_size: Float32): CPointer<Unit>

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

    public func filled_curve(filled_curve: Bool): ErrorBar {
        this.ptr = unsafe { c_filled_curve(this.ptr, filled_curve) }
        return this
    }

    public func cap_size(cap_size: Float32): ErrorBar {
        this.ptr = unsafe { c_cap_size(this.ptr, cap_size) }
        return this
    }

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

    public override func marker_color(c: String): ErrorBar {
        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_color(c: String): ErrorBar {
        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(c: String): ErrorBar {
        var cstr = unsafe { LibC.mallocCString(c) }
        this.ptr = unsafe { c_marker_str(this.ptr, cstr) }
        unsafe { LibC.free(cstr) }
        return this
    }

    public override func color(c: String): ErrorBar {
        var cstr = unsafe { LibC.mallocCString(c) }
        this.ptr = unsafe { c_color_str(this.ptr, cstr) }
        unsafe { LibC.free(cstr) }
        return this
    }
}

public enum ErrorBarType {
    | Vertical
    | Horizontal
    | Both

    func toString(): String {
        return match (this) {
            case Vertical => "vertical"
            case Horizontal => "horizontal"
            case Both => "both"
        }
    }
}

public func errorbar(x: Vector<Float64>, y: Vector<Float64>, error: Vector<Float64>,
                     errorbar_type!:ErrorBarType = Vertical, line_spec!:String = ""): ErrorBar {
    let len = x.size()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    var cstr_errorbar_type = unsafe { LibC.mallocCString(errorbar_type.toString()) }
    let handle = unsafe {
        c_errorbar(x.ptr, y.ptr, error.ptr, len,
                   cstr_errorbar_type, cstr_line_spec)
    }
    unsafe { LibC.free(cstr_line_spec) }
    unsafe { LibC.free(cstr_errorbar_type) }
    return ErrorBar(handle)
}

public func errorbar(x: Vector<Float64>, y: Vector<Float64>,
                     yneg: Vector<Float64>, ypos: Vector<Float64>,
                     xneg: Vector<Float64>, xpos: Vector<Float64>, line_spec!:String = ""): ErrorBar {
    let len = x.size()
    var cstr_line_spec = unsafe { LibC.mallocCString(line_spec) }
    let handle = unsafe {
        c_errorbar2(x.ptr, y.ptr, yneg.ptr, ypos.ptr, xneg.ptr, xpos.ptr, len, cstr_line_spec)
    }
    unsafe { LibC.free(cstr_line_spec) }
    return ErrorBar(handle)
}

public func testErrorBar1() {
    let x = arange<Float64>(0.0, 100.0, 10.0)
    let y = vector<Float64>([20.0, 30.0, 45.0, 40.0, 60.0, 65.0, 80.0, 75.0, 95.0, 90.0])
    let err = vector<Float64>(y.size(), { x => 10.0 })
    errorbar(x, y, err)
    axis(0.0, 100.0, 0.0, 110.0)
    save("./tests/imgs/errorbar/errorbar_1.svg", "svg")
    clear()
}

public func testErrorBar2() {
    let x = arange<Float64>(0.0, 100.0, 10.0)
    let y = vector<Float64>([20.0, 30.0, 45.0, 40.0, 60.0, 65.0, 80.0, 75.0, 95.0, 90.0])
    let err = vector<Float64>(y.size(), { x => 10.0 })
    errorbar(x, y, err).filled_curve(true)
    axis(0.0, 100.0, 0.0, 110.0)
    save("./tests/imgs/errorbar/errorbar_2.svg", "svg")
    clear()
}

public func testErrorBar3() {
    let x = arange<Float64>(0.0, 100.0, 10.0)
    let y = vector<Float64>([20.0, 30.0, 45.0, 40.0, 60.0, 65.0, 80.0, 75.0, 95.0, 90.0])
    let err = vector<Float64>([5.0, 8.0, 2.0, 9.0, 3.0, 3.0, 8.0, 3.0, 9.0, 3.0])
    errorbar(x, y, err)
    axis(0.0, 100.0, 0.0, 110.0)
    save("./tests/imgs/errorbar/errorbar_3.svg", "svg")
    clear()
}

public func testErrorBar4() {
    let x = arange<Float64>(0.0, 100.0, 10.0)
    let y = vector<Float64>([20.0, 30.0, 45.0, 40.0, 60.0, 65.0, 80.0, 75.0, 95.0, 90.0])
    let err = vector<Float64>([1.0, 3.0, 5.0, 3.0, 5.0, 3.0, 6.0, 4.0, 3.0, 3.0])
    errorbar(x, y, err, errorbar_type:ErrorBarType.Horizontal)
    axis(0.0, 100.0, 0.0, 110.0)
    save("./tests/imgs/errorbar/errorbar_4.svg", "svg")
    clear()
}

public func testErrorBar5() {
    let x = arange<Float64>(0.0, 100.0, 10.0)
    let y = vector<Float64>([20.0, 30.0, 45.0, 40.0, 60.0, 65.0, 80.0, 75.0, 95.0, 90.0])
    let err = vector<Float64>([4.0, 3.0, 5.0, 3.0, 5.0, 3.0, 6.0, 4.0, 3.0, 3.0])
    errorbar(x, y, err, errorbar_type:ErrorBarType.Both)
    axis(0.0, 100.0, 0.0, 110.0)
    axis(AxisStyle.Equal)
    save("./tests/imgs/errorbar/errorbar_5.svg", "svg")
    clear()
}

public func testErrorBar6() {
    let x = arange<Float64>(0.0, 100.0, 10.0)
    let y = vector<Float64>([20.0, 30.0, 45.0, 40.0, 60.0, 65.0, 80.0, 75.0, 95.0, 90.0])
    let err = vector<Float64>([4.0, 3.0, 5.0, 3.0, 5.0, 3.0, 6.0, 4.0, 3.0, 3.0])
    errorbar(x, y, err, errorbar_type:ErrorBarType.Both, line_spec:"o")
    axis(0.0, 100.0, 0.0, 100.0)
    axis(AxisStyle.Equal)
    save("./tests/imgs/errorbar/errorbar_6.svg", "svg")
    clear()
}

public func testErrorBar7() {
    let x = arange<Float64>(0.0, 100.0, 10.0)
    let y = vector<Float64>([20.0, 30.0, 45.0, 40.0, 60.0, 65.0, 80.0, 75.0, 95.0, 90.0])
    let yneg = vector<Float64>([1.0, 3.0, 5.0, 3.0, 5.0, 3.0, 6.0, 4.0, 3.0, 3.0])
    let ypos = vector<Float64>([2.0, 5.0, 3.0, 5.0, 2.0, 5.0, 2.0, 2.0, 5.0, 5.0])
    let xneg = vector<Float64>([1.0, 3.0, 5.0, 3.0, 5.0, 3.0, 6.0, 4.0, 3.0, 3.0])
    let xpos = vector<Float64>([2.0, 5.0, 3.0, 5.0, 2.0, 5.0, 2.0, 2.0, 5.0, 5.0])
    errorbar(x, y, yneg, ypos, xneg, xpos, line_spec:"o")
    axis(0.0, 100.0, 0.0, 100.0)
    save("./tests/imgs/errorbar/errorbar_7.svg", "svg")
    clear()
}

public func testErrorBar8() {
    let x = linspace<Float64>(0.0, 10.0, num:15)
    let y = x.apply({ x => sin(x / 2.0) })
    let err = vector<Float64>(y.size(), { x => 0.3 })
    var e:ErrorBar = errorbar(x, y, err, line_spec:"-s").marker_size(10.0)
                        .marker_color("red")
                        .marker_face_color("red")
    save("./tests/imgs/errorbar/errorbar_8.svg", "svg")
    clear()
}

public func testErrorBar9() {
    let x = linspace<Float64>(0.0, 2.0, num:15)
    let y = x.apply({ x => exp(x) })
    let err = vector<Float64>(y.size(), { x => 0.3 })
    var e:ErrorBar = errorbar(x, y, err).cap_size(18.0);
    save("./tests/imgs/errorbar/errorbar_9.svg", "svg")
    clear()
}

public func testErrorBar10() {
    let x = linspace<Float64>(0.0, 10.0, num:10)
    let y = x.apply({ x => sin(x / 2.0) })
    let err = vector<Float64>(y.size(), { x => 0.3 })
    var e:ErrorBar = errorbar(x, y, err).cap_size(15.0).marker("*")
                        .marker_size(10.0).color("red")
    save("./tests/imgs/errorbar/errorbar_10.svg", "svg")
    clear()
}

public func testErrorBar() {
    print("  + testErrorBar\n")
    testErrorBar1()
    testErrorBar2()
    testErrorBar3()
    testErrorBar4()
    testErrorBar5()
    testErrorBar6()
    testErrorBar7()
    testErrorBar8()
    testErrorBar9()
    testErrorBar10()
}