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

/* Note: The same api and data generate different figures with original docs,
   maybe 'heatmap' api changed. */

/* Type for x: Float64 */
foreign func c_heatmap(x: CPointer<Unit>, x_row: Int64, x_col: Int64): CPointer<Unit>
foreign func c_normalization(ptr: CPointer<Unit>, str_color_normalization: CString): Unit

public func heatmap(data: Matrix<Float64>): MatplotMatrix {
    let handle = unsafe { c_heatmap(data.ptr, data.getRows(), data.getCols()) }
    return MatplotMatrix(handle)
}

public enum color_normalization {
    | none
    | rows
    | columns

    func toString(): String {
        return match (this) {
            case none => "none"
            case rows => "rows"
            case columns => "columns"
        }
    }
}

public class MatplotMatrix {
    var ptr: CPointer<Unit>

    /* Input is a pointer to matplot::matrix. */
    init(ptr: CPointer<Unit>) {
        this.ptr = ptr
    }

    public func normalization (color_normalization: color_normalization): Unit {
        var cstr_color_normalization = unsafe { LibC.mallocCString(color_normalization.toString()) }
        unsafe { c_normalization(this.ptr, cstr_color_normalization) }
        unsafe { LibC.free(cstr_color_normalization) }
    }
}

public func testHeatMap1() {
    let data = matrix([
        [45.0, 60.0, 32.0],
        [43.0, 54.0, 76.0],
        [32.0, 94.0, 68.0],
        [23.0, 95.0, 58.0]
    ])
    
    heatmap(data)
    save("./tests/imgs/heatmap/heatmap_1.png", "png")
    clear()
}

public func testHeatMap2() {
    let data = matrix([
        [24.0, 10.0],
        [10.0, 5.0],
        [24.0, 16.0],
        [8.0, 3.0]
    ])

    heatmap(data)
    title("Conut of SelfAssessedHealthStatus vs. Smoker")
    xticklabels(["false", "true"])
    yticklabels(["Excellent", "Fair", "Good", "Poor"])
    xlabel("Smoker")
    ylabel("SelfAssessedHealthStatus")

    save("./tests/imgs/heatmap/heatmap_2.png", "png")
    clear()
}

public func testHeatMap3() {
    let data = matrix([
        [38.46, 39.0],
        [39.7, 36.2],
        [38.13, 38.88],
        [33.88, 43.0]
    ])

    heatmap(data)
    title("Mean of Age")
    xticklabels(["false", "true"])
    yticklabels(["Excellent", "Fair", "Good", "Poor"])
    xlabel("Smoker")
    ylabel("SelfAssessedHealthStatus")

    save("./tests/imgs/heatmap/heatmap_3.png", "png")
    clear()
}

public func testHeatMap4() {
    let data = matrix([
        [45.0, 60.0, 32.0],
        [43.0, 54.0, 76.0],
        [32.0, 94.0, 68.0],
        [23.0, 95.0, 58.0]
    ])

    heatmap(data)
    title("T-Shirt Orders")
    xticklabels(["Small", "Medium", "Large"])
    yticklabels(["Green", "Red", "Blue", "Gray"])
    xlabel("Sizes")
    ylabel("Colors")

    save("./tests/imgs/heatmap/heatmap_4.png", "png")
    clear()
}

public func testHeatMap5() {
    let data = matrix([
        [12.0, 135.0, 20.0, 0.0, 127.0],
        [0.0, 1.0, 0.0, 0.0, 1.0],
        [19.0, 31.0, 81.0, 8.0, 49.0],
        [9.0, 18.0, 42.0, 2.0, 85.0],
        [0.0, 5.0, 3.0, 0.0, 17.0],
        [31.0, 143.0, 135.0, 6.0, 23.0],
        [32.0, 102.0, 54.0, 6.0, 7.0],
        [5.0, 11.0, 4.0, 0.0, 4.0],
        [16.0, 41.0, 13.0, 3.0, 22.0],
        [18.0, 70.0, 37.0, 1.0, 19.0]
    ])

    heatmap(data)
    title("Count of Cause vs. Region")
    let ax = gca()
    xticklabels(["MidWest", "NorthEast", "SouthEast", "SouthWest", "West"])
    yticklabels(["Attack", "Earthquake", "Energy emergency", "Equipment fault", "Fire",
         "Severe Storm", "Thunder Storm", "Unknown", "Wind", "Winter Storm"])
    xlabel("Region")
    ylabel("Cause")
    var w = ax.width()
    ax.width(w * 0.85)
    ax.x_origin(ax.x_origin() + w * 0.1)
    save("./tests/imgs/heatmap/heatmap_5.png", "png")
    clear()
}

public func testHeatMap6() {
    let data = matrix([
        [12.0, 135.0, 20.0, 0.0, 127.0],
        [0.0, 1.0, 0.0, 0.0, 1.0],
        [19.0, 31.0, 81.0, 8.0, 49.0],
        [9.0, 18.0, 42.0, 2.0, 85.0],
        [0.0, 5.0, 3.0, 0.0, 17.0],
        [31.0, 143.0, 135.0, 6.0, 23.0],
        [32.0, 102.0, 54.0, 6.0, 7.0],
        [5.0, 11.0, 4.0, 0.0, 4.0],
        [16.0, 41.0, 13.0, 3.0, 22.0],
        [18.0, 70.0, 37.0, 1.0, 19.0]
    ])

    heatmap(data).normalization(color_normalization.columns)
    title("Count of Cause vs. Region")
    let ax = gca()
    xticklabels(["MidWest", "NorthEast", "SouthEast", "SouthWest", "West"])
    yticklabels(["Attack", "Earthquake", "Energy emergency", "Equipment fault", "Fire",
         "Severe Storm", "Thunder Storm", "Unknown", "Wind", "Winter Storm"])
    xlabel("Region")
    ylabel("Cause")
    var w = ax.width()
    ax.width(w * 0.85)
    ax.x_origin(ax.x_origin() + w * 0.1)
    save("./tests/imgs/heatmap/heatmap_6.png", "png")
    clear()
}

public func testHeatMap7() {
    let data = matrix([
        [12.0, 135.0, 20.0, 0.0, 127.0],
        [0.0, 1.0, 0.0, 0.0, 1.0],
        [19.0, 31.0, 81.0, 8.0, 49.0],
        [9.0, 18.0, 42.0, 2.0, 85.0],
        [0.0, 5.0, 3.0, 0.0, 17.0],
        [31.0, 143.0, 135.0, 6.0, 23.0],
        [32.0, 102.0, 54.0, 6.0, 7.0],
        [5.0, 11.0, 4.0, 0.0, 4.0],
        [16.0, 41.0, 13.0, 3.0, 22.0],
        [18.0, 70.0, 37.0, 1.0, 19.0]
    ])

    heatmap(data).normalization(color_normalization.rows)
    title("Count of Cause vs. Region")
    let ax = gca()
    xticklabels(["MidWest", "NorthEast", "SouthEast", "SouthWest", "West"])
    yticklabels(["Attack", "Earthquake", "Energy emergency", "Equipment fault", "Fire",
         "Severe Storm", "Thunder Storm", "Unknown", "Wind", "Winter Storm"])
    xlabel("Region")
    ylabel("Cause")
    var w = ax.width()
    ax.width(w * 0.85)
    ax.x_origin(ax.x_origin() + w * 0.1)
    save("./tests/imgs/heatmap/heatmap_7.png", "png")
    clear()
}

public func testHeatMap() {
    print("  + testHeatMap\n")
    testHeatMap1()
    testHeatMap2()
    testHeatMap3()
    testHeatMap4()
    testHeatMap5()
    testHeatMap6()
    testHeatMap7()
}