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

@C
struct ImageInfo {
    let data: CPointer<Unit>  /* Type: UInt8 */
    let nchannel: Int64
    let row: Int64
    let col: Int64

    public init(data: CPointer<Unit>, nchannel: Int64, row: Int64, col: Int64) {
        this.data = data
        this.nchannel = nchannel
        this.row = row
        this.col = col
    }
}

/* Loading and displaying images 
   Type for C, r, g, b: Float64
   Type for data: UInt8 */
foreign func c_imread(filename: CString): ImageInfo
foreign func c_imshow(data: CPointer<Unit>, nchannel: Int64, row: Int64, col: Int64): Unit
foreign func c_imshow_gray(data: CPointer<Unit>, row: Int64, col: Int64): Unit
foreign func c_imshow_file(filename: CString): Unit
foreign func c_imvignette(data: CPointer<Unit>, nchannel: Int64, row: Int64, col: Int64): ImageInfo
foreign func c_imresize(data: CPointer<Unit>, nchannel: Int64, row: Int64, col: Int64, scale: Float64, interpolation: CString): ImageInfo
foreign func c_image(C: CPointer<Unit>, row: Int64, col: Int64): CPointer<Unit>
foreign func c_image_3(r: CPointer<Unit>, g: CPointer<Unit>, b: CPointer<Unit>, row: Int64, col: Int64): CPointer<Unit>
foreign func c_image_v(C: CPointer<Unit>, row: Int64, col: Int64, scaled_colorbar: Bool): CPointer<Unit>
foreign func c_image_range(x_min: Float64, x_max: Float64, y_min: Float64, y_max: Float64, C: CPointer<Unit>, row: Int64, col: Int64): CPointer<Unit>
foreign func c_matrix_alpha(ptr: CPointer<Unit>, alpha: Float64): CPointer<Unit>
foreign func c_imagesc(C: CPointer<Unit>, row: Int64, col: Int64): CPointer<Unit>
foreign func c_imagesc_range(x_min: Float64, x_max: Float64, y_min: Float64, y_max: Float64, C: CPointer<Unit>, row: Int64, col: Int64): CPointer<Unit>

public enum ImageInterpolation {
    | Raw
    | AdditionalSpace
    | Nearest
    | MovingAverage
    | Bilinear
    | Grid
    | Bicubic
    | Lanczos

    func toString(): String {
        return match (this) {
            case Raw => "raw"
            case AdditionalSpace => "additional_space"
            case Nearest => "nearest"
            case MovingAverage => "moving_average"
            case Bilinear => "bilinear"
            case Grid => "grid"
            case Bicubic => "bicubic"
            case Lanczos => "lanczos"
        }
    }
}

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

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

    public func alpha(alpha: Float64): mat_matrix{
        this.ptr = unsafe { c_matrix_alpha(this.ptr, alpha) }
        return this
    }
}

public func imread(filename: String): Tensor<UInt8> {
    var cstr_filename = unsafe { LibC.mallocCString(filename) }
    let info: ImageInfo = unsafe { c_imread(cstr_filename) }
    let res = tensor<UInt8>(info.data, [info.nchannel, info.row, info.col])
    unsafe { LibC.free(cstr_filename) }
    return res
}

public func imshow(img: Tensor<UInt8>): Unit {
    let shape = img.shape()
    if (shape.size != 3) {
        throw IllegalArgumentException("imshow: tensor argument must have 3 dimensions")
    }
    unsafe { c_imshow(img.ptr, shape[0], shape[1], shape[2]) }
}

public func imshow(img: Matrix<UInt8>): Unit {
    unsafe { c_imshow_gray(img.ptr, img.getRows(), img.getCols()) }
}

public func imshow(filename: String): Unit {
    var cstr_filename = unsafe { LibC.mallocCString(filename) }
    unsafe { c_imshow_file(cstr_filename) }
    unsafe { LibC.free(cstr_filename) }
}

public func imvignette(img: Tensor<UInt8>): Tensor<UInt8> {
    let shape = img.shape()
    if (shape.size != 3) {
        throw IllegalArgumentException("imvignette: tensor argument must have 3 dimensions")
    }
    let info: ImageInfo = unsafe { c_imvignette(img.ptr, shape[0], shape[1], shape[2]) }
    let res = tensor<UInt8>(info.data, [info.nchannel, info.row, info.col])
    return res
}

public func imresize(img: Tensor<UInt8>, scale: Float64,
                     interpolation!:ImageInterpolation = ImageInterpolation.Bicubic): Tensor<UInt8> {
    let shape = img.shape()
    if (shape.size != 3) {
        throw IllegalArgumentException("imresize: tensor argument must have 3 dimensions")
    }
    var cstr_interpolation = unsafe { LibC.mallocCString(interpolation.toString()) }
    let info: ImageInfo = unsafe { c_imresize(img.ptr, shape[0], shape[1], shape[2], scale,
                                              cstr_interpolation) }
    let res = tensor<UInt8>(info.data, [info.nchannel, info.row, info.col])
    unsafe { LibC.free(cstr_interpolation) }
    return res
}

public func image(C: Matrix<Float64>) {
    let handle = unsafe { c_image(C.ptr, C.getRows(), C.getCols()) }
    return mat_matrix(handle)
}

public func image(r: Matrix<Float64>, g: Matrix<Float64>, b: Matrix<Float64>) {
    let handle = unsafe { c_image_3(r.ptr, g.ptr, b.ptr, r.getRows(), r.getCols()) }
    return mat_matrix(handle)
}

public func image(C: Matrix<Float64>, scaled_colorbar: Bool) {
    let handle = unsafe { c_image_v(C.ptr, C.getRows(), C.getCols(), scaled_colorbar) }
    return mat_matrix(handle)
}

public func image(x_min: Float64, x_max: Float64, y_min: Float64, y_max: Float64, C: Matrix<Float64>) {
    let handle = unsafe { c_image_range(x_min, x_max, y_min, y_max, C.ptr, C.getRows(), C.getCols()) }
    return mat_matrix(handle)
}

public func imagesc(C: Matrix<Float64>) {
    let handle = unsafe { c_imagesc(C.ptr, C.getRows(), C.getCols()) }
    return mat_matrix(handle)
}

public func imagesc(x_min: Float64, x_max: Float64, y_min: Float64, y_max: Float64, C: Matrix<Float64>) {
    let handle = unsafe { c_imagesc_range(x_min, x_max, y_min, y_max, C.ptr, C.getRows(), C.getCols()) }
    return mat_matrix(handle)
}

public func testImageShow1() {
    let image = imread("./tests/imgs/imshow/lena_gray.tiff")
    imshow(image)
    save("./tests/imgs/imshow/imshow_1.png", "png")
    clear()
}

public func testImageShow2() {
    let image = imread("./tests/imgs/imshow/lena_color.tiff")
    imshow(image)
    save("./tests/imgs/imshow/imshow_2.png", "png")
    clear()    
}

public func testImageShow3() {
    let image = imread("./tests/imgs/imshow/lena_gray.tiff")
    let gray = image[0].asMatrix()
    let h = gray.getRows()
    let w = gray.getCols()
    var mean_density: Float64 = 0.0
    for (i in 0..h) {
        for (j in 0..w) {
            mean_density += Float64(gray[i,j])
        }
    }
    mean_density /= Float64(h * w)
    for (i in 0..h) {
        for (j in 0..w) {
            if (Float64(gray[i,j]) > mean_density) {
                gray[i,j] = 255
            } else {
                gray[i,j] = 0
            }
        }
    }
    imshow(gray)
    save("./tests/imgs/imshow/imshow_3.png", "png")
    clear()
}

public func testImageShow4() {
    imshow("./tests/imgs/imshow/lena_color.tiff")
    save("./tests/imgs/imshow/imshow_4.png", "png")
    clear()
}

func testImageShow5(): Unit {
    let image = imread("./tests/imgs/imshow/lena_gray.tiff")
    imshow(image[0].asMatrix())
    colormap(Palette.Greens)
    save("./tests/imgs/imshow/imshow_5.png", "png")
    clear()
}

func testImageShow6(): Unit {
    let image = imread("./tests/imgs/imshow/lena_gray.tiff")
    imshow(image[0].asMatrix())
    colormap(Palette.DefaultMap)
    save("./tests/imgs/imshow/imshow_6.png", "png")
    clear()
}

func testImageShow7(): Unit {
    let image = imread("./tests/imgs/imshow/lena_color.tiff")
    let image2 = imvignette(image)
    imshow(image2)
    save("./tests/imgs/imshow/imshow_7.png", "png")
    clear()
}

func testImageShow8(): Unit {
    let image = imread("./tests/imgs/imshow/lena_gray.tiff")
    let image2 = imvignette(image)
    imshow(image2)
    save("./tests/imgs/imshow/imshow_8.png", "png")
    clear()
}

func testImageShow9(): Unit {
    let image = imread("./tests/imgs/imshow/lena_color.tiff")
    let image2 = imresize(image, 0.1, interpolation: ImageInterpolation.Bilinear)
    imshow(image2)
    save("./tests/imgs/imshow/imshow_9.png", "png")
    clear()
}

func testImageShow10(): Unit {
    let image = imread("./tests/imgs/imshow/lena_color.tiff")
    let image2 = imvignette(image)
    let image3 = imresize(image2, 0.5)
    imshow(image3)
    save("./tests/imgs/imshow/imshow_10.png", "png")
    clear()
}

func testImageShow11(): Unit {
    let image = imread("./tests/imgs/imshow/lena_gray.tiff")
    imshow(image[0].asMatrix())
    colormap(Palette.Blues)
    save("./tests/imgs/imshow/imshow_11.png", "png")
    clear()
}

public func testImageShow() {
    testImageShow1()
    testImageShow2()
    testImageShow3()
    testImageShow4()
    testImageShow5()
    testImageShow6()
    testImageShow7()
    testImageShow8()
    testImageShow9()
    testImageShow10()
    testImageShow11()
}

public func testImageMatrix1() {
    let C = matrix<Float64>(
        [[0.0, 2.0, 4.0, 6.0],
         [8.0, 10.0, 12.0, 14.0],
         [16.0, 18.0, 20.0, 22.0]])
    image(C)
    colorbar()
    save("./tests/imgs/imagematrix/imagematrix1.png", "png")
    clear()
}

public func testImageMatrix2() {
    let C = matrix<Float64>(
        [[0.0, 2.0, 4.0, 6.0],
         [8.0, 10.0, 12.0, 14.0],
         [16.0, 18.0, 20.0, 22.0]])
    image(C, true)
    colorbar()
    save("./tests/imgs/imagematrix/imagematrix2.png", "png")
    clear()
}

public func testImageMatrix3() {
    let C = matrix<Float64>(
        [[0.0, 2.0, 4.0, 6.0],
         [8.0, 10.0, 12.0, 14.0],
         [16.0, 18.0, 20.0, 22.0]])
    image(5.0, 8.0, 3.0, 6.0, C)
    save("./tests/imgs/imagematrix/imagematrix3.png", "png")
    clear()
}

public func testImageMatrix4() {
    let r = matrix<Float64>(
        [[1.0, 2.0, 3.0],
         [4.0, 5.0, 6.0],
         [7.0, 8.0, 9.0]])
    let r_new = r.apply({t => t * 25.5})
    let g = matrix<Float64>(
        [[0.0, 0.0, 0.0],
         [0.0, 0.0, 0.0],
         [0.0, 0.0, 0.0]])
    let b = matrix<Float64>(
        [[0.0, 0.0, 0.0],
         [0.0, 0.0, 0.0],
         [0.0, 0.0, 0.0]])
    image(r_new, g, b)
    save("./tests/imgs/imagematrix/imagematrix4.png", "png")
    clear()
}

public func testImageMatrix5() {
    let line = linspace(1.0, 3.0)
    plot(line)
    hold(true)
    let C = matrix<Float64>(
        [[1.0, 2.0, 3.0],
         [4.0, 5.0, 6.0],
         [7.0, 8.0, 9.0]])
    let im = image(C)
    im.alpha(0.5)
    gca().y_axis().reverse(false)

    save("./tests/imgs/imagematrix/imagematrix5.png", "png")
    clear()
}

public func testImagesc1() {
    let C = matrix<Float64>(
        [[0.0, 2.0, 4.0, 6.0],
         [8.0, 10.0, 12.0, 14.0],
         [16.0, 18.0, 20.0, 22.0]])
    imagesc(C)
    colorbar()
    save("./tests/imgs/imagesc/imagesc1.png", "png")
    clear()
}

public func testImagesc2() {
    let C = matrix<Float64>(
        [[0.0, 2.0, 4.0, 6.0],
         [8.0, 10.0, 12.0, 14.0],
         [16.0, 18.0, 20.0, 22.0]])
    imagesc(5.0, 8.0, 3.0, 6.0, C)
    save("./tests/imgs/imagesc/imagesc2.png", "png")
    clear()
}

public func testImagesc3() {
    let C = matrix<Float64>(
        [[0.0, 2.0, 4.0, 6.0],
         [8.0, 10.0, 12.0, 14.0],
         [16.0, 18.0, 20.0, 22.0]])
    imagesc(5.0, 8.0, 3.0, 6.0, C)
    gca().color_box_range(4.0, 18.0)
    save("./tests/imgs/imagesc/imagesc3.png", "png")
    clear()
}

public func testImagesc4() {
    let C = matrix<Float64>(
        [[0.0, 2.0, 4.0, 6.0],
         [8.0, 10.0, 12.0, 14.0],
         [16.0, 18.0, 20.0, 22.0]])
    let im = imagesc(5.0, 8.0, 3.0, 6.0, C)
    gca().color_box_range(4.0, 18.0)
    // alpha() can't save as png, but svg will also be different
    im.alpha(0.5)

    save("./tests/imgs/imagesc/imagesc4.svg", "svg")
    clear()
}

public func testImageMatrix() {
    testImageMatrix1()
    testImageMatrix2()
    testImageMatrix3()
    testImageMatrix4()
    testImageMatrix5()
    testImagesc1()
    testImagesc2()
    testImagesc3()
    testImagesc4()
}