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

public func fplot3(fx: (Float64)->Float64, fy: (Float64)->Float64, 
                   fz: (Float64)->Float64, line_spec!:String = "") {
    let t = linspace<Float64>(-5.0, 5.0, num: 500)
    let x = t.apply(fx)
    let y = t.apply(fy)
    let z = t.apply(fz)
    let handle = plot3(x, y, z, line_spec: line_spec)
    return handle
}

public func fplot3(fx: (Float64)->Float64, fy: (Float64)->Float64, fz: (Float64)->Float64,
                   t_start: Float64, t_end: Float64, mesh_density!:Int64 = 500, line_spec!:String = "")
{
    let t = linspace<Float64>(t_start, t_end, num: mesh_density)
    let x = t.apply(fx)
    let y = t.apply(fy)
    let z = t.apply(fz)
    let handle = plot3(x, y, z, line_spec: line_spec)
    return handle
}

public func testFuncPlot3d1() {
    fplot3(sin, cos, {t:Float64 => t})
    save("./tests/imgs/func_plot_3d/func_plot3d1.svg", "svg")
    clear()
}

public func testFuncPlot3d2() {
    let xt = { t:Float64 => exp(-t / 10.0) * sin(5.0 * t) }
    let yt = { t:Float64 => exp(-t / 10.0) * cos(5.0 * t) }
    fplot3(xt, yt, {t:Float64 => t}, -10.0, 10.0)
    save("./tests/imgs/func_plot_3d/func_plot3d2.svg", "svg")
    clear()
}

public func testFuncPlot3d3() {
    let f = fplot3(sin, cos, {t:Float64 => t}, 0.0, 2.0 * Float64.getPI())
    f.line_width(2.0)
    hold(true)
    fplot3(sin, cos, {t:Float64 => t}, 2.0 * Float64.getPI(), 4.0 * Float64.getPI(), line_spec: "--or")
    fplot3(sin, cos, {t:Float64 => t}, 4.0 * Float64.getPI(), 6.0 * Float64.getPI(), line_spec: "-.*c")
    hold(false)
    save("./tests/imgs/func_plot_3d/func_plot3d3.svg", "svg")
    clear()
}

public func testFuncPlot3d4() {
    // Note: animation not included
    let xt = { t:Float64 => exp(-abs(t) / 10.0) * sin(5.0 * abs(t)) }
    let yt = { t:Float64 => exp(-abs(t) / 10.0) * cos(5.0 * abs(t)) }
    let zt = { t:Float64 => t }
    fplot3(xt, yt, zt, -10.0, 10.0).color("r")
    xlabel("e^{-|z|/10} sin(2|z|")
    ylabel("e^{-|z|/10} cos(2|z|")
    zlabel("z")
    grid(true)
    save("./tests/imgs/func_plot_3d/func_plot3d4.svg", "svg")
    clear()
}

public func testFuncPlot3d5() {
    let xt = {t:Float64 => t}
    let yt = {t:Float64 => t / 2.0}
    let zt = {t:Float64 => sin(6.0 * t) }
    fplot3(xt, yt, zt, -2.0 * Float64.getPI(), 2.0 * Float64.getPI(), mesh_density: 300)
    .line_width(1.0)
    title("x=t, y=t/2, z=sin(6t) for -2π<t<2π")
    xlabel("x")
    ylabel("y")
    view(52.5, 30.0)
    box(true)
    xrange(-2.0 * Float64.getPI(), 2.0 * Float64.getPI())
    yrange(-Float64.getPI(), Float64.getPI())
    xticks(linspace(-2.0 * Float64.getPI(), 2.0 * Float64.getPI(), num: 9))
    xticklabels(["-2π", "-3π/2", "-π", "-π/2", "0", "π/2", "π", "3π/2", "2π"])
    yticks(linspace(-Float64.getPI(), Float64.getPI(), num: 5))
    yticklabels(["-π", "-π/2", "0", "π/2", "π"])
    save("./tests/imgs/func_plot_3d/func_plot3d5.svg", "svg")
    clear()
}

public func testFunctionPlot3d(){
    testFuncPlot3d1()
    testFuncPlot3d2()
    testFuncPlot3d3()
    testFuncPlot3d4()
    testFuncPlot3d5()
}