(***************************************************************************************************
*            F# sample for the usage of TopoART-R (classes TopoART_R and Fast_TopoART_R)           *
*                                 (approximation of the Mandelbrot set)                            *
****************************************************************************************************
*                              Created by Marko Tscherepanow, 9 July 2019                          *
***************************************************************************************************)

// Compile and run from the console: dotnet run --project TopoART-R_sample3.fsproj

/// <summary>
/// Pixel-wise regression analysis of an image using TopoART-R. [F#]
/// <para>
/// This sample trains a TopoART-R network with a colour image depicting the Mandelbrot set. The 
/// pixel coordinates are used as independent variables and the corresponding colour values as 
/// dependent variables. The accuracy of the regression function can be controlled by means of the 
/// vigilance parameter <c>rho_a</c>.
/// </para>
/// <para>The training image <c>Mandelbrot_reference.png</c> and the predicted image 
/// <c>TopoART-R_Mandelbrot_regression.png</c> are written into the subfolder 
/// <c>results/regression</c>.
/// </para>
/// </summary>
module LibTopoART_samples.TopoART_R_sample3

open LibTopoART
open System
open SixLabors.ImageSharp
open SixLabors.ImageSharp.PixelFormats
open System.IO
open System.Reflection

type Selector = USE_TopoART_R | USE_Fast_TopoART_R

// Choose TopoART-R implementation (see enum Selector above)
let sel = USE_Fast_TopoART_R

// Destination directory for the regression results
let resultPath = "../../../../../results/regression/"
let networkPath = "../../../../../results/networks/"

// Input and output dimensions
let iLen = 2; // Number of independent variables (input dimensions)
let dLen = 3; // Number of dependent variables (output dimensions)

// TopoART-R parameters
let rho_a = 0.95m    // Most relevant parameter (decrease and/or increase to visualise effects)
let beta_sbm = 0.5m  // Default
let phi = 1L         // Default (increase to enable noise reduction)
let tau = 100L       // Default

// Application parameters
let imgWidth = 800
let imgHeight = 600
let maxIter = 1023

(*--------------------------------------------------------------------------------------------------
-                                           Helper functions                                       -
--------------------------------------------------------------------------------------------------*)

let ColorFromIteration iter =
    if iter = maxIter then
        Rgb24(0uy, 0uy, 0uy)
    else
        let b_tmp = float iter / float maxIter * 100.0
        let b = int (b_tmp / ( 1.0 + b_tmp ) * 255.0)
        let g_tmp = float iter / float maxIter * 25.0
        let g = int (g_tmp / ( 1.0 + g_tmp ) * 255.0)
        let r_tmp = float iter / float maxIter * 5.0
        let r = int ((1.0 - (r_tmp / ( 1.0 + r_tmp ))) * 31.0)
        Rgb24(byte r, byte g, byte b)

let ColorFromPrediction (pred : decimal[]) =
    let r = Math.Round(pred[0] * 255.0m)
    let g = Math.Round(pred[1] * 255.0m)
    let b = Math.Round(pred[2] * 255.0m)
    Rgb24(byte r, byte g, byte b)

let GetPixel x y =
    let x0 = float x / float imgWidth * 3.0 - 2.0
    let y0 = float y / float imgHeight * 2.0 - 1.0

    let mutable x = 0.0
    let mutable y = 0.0
    let mutable i = 0

    while x * x + y * y <= 4.0 && i < maxIter do 
        let tmp = x * x - y * y + x0
        y <- 2.0 * x * y + y0
        x <- tmp
        i <- i + 1
    i

(*--------------------------------------------------------------------------------------------------
-                                             Main program                                         -
--------------------------------------------------------------------------------------------------*)

[<EntryPoint>]
let main args =

    // Set working directory to assembly directory
    let cb = Uri(Assembly.GetEntryAssembly().Location)
    Directory.SetCurrentDirectory(Path.GetDirectoryName(cb.LocalPath))

    let sampleNum = imgWidth * imgHeight
    let i = Array.create sampleNum <| Array.create iLen 0.0m
    let d = Array.create sampleNum <| Array.create dLen 0.0m
    
    printf "Generate training data\n"

    // Generate data (used for training and testing)
    let mutable sampleIndex = 0
    use refImage = new Image<Rgb24>(imgWidth, imgHeight)
    for y = 0 to imgHeight - 1 do
        for x = 0 to imgWidth - 1 do
            let iter = GetPixel x y
            let c = ColorFromIteration iter
            refImage[x, y] <- c
            i[sampleIndex] <- [| (decimal x) / (decimal (imgWidth - 1)); (decimal y) / (decimal (imgHeight - 1)) |]
            d[sampleIndex] <- [| decimal c.R / 255.0m; decimal c.G / 255.0m; decimal c.B / 255.0m |]
            sampleIndex <- sampleIndex + 1

    // Save reference image
    refImage.SaveAsJpeg(resultPath + "Mandelbrot_reference.png")

    // Create TopoART-R network
    let tar = 
        if sel = USE_TopoART_R then 
            new TopoART_R(int64 iLen, int64 dLen, 2L, rho_a) :> ITopoART_R
        else
            new Fast_TopoART_R(int64 iLen, int64 dLen, 2L, rho_a) :> ITopoART_R

    let filePrefix = if sel = USE_TopoART_R then "TopoART-R" else "Fast_TopoART-R"
    let fileSuffix = if sel = USE_TopoART_R then "tar" else "ftar"

    tar.Beta_sbm <- beta_sbm
    tar.Tau <- tau
    tar.Phi <- phi
    tar.SkipEdgeLearning <- true
    
    // Start time measuring
    let trainingStart = DateTime.Now

    // Train
    for trainIndex = 0 to sampleNum - 1 do
        tar.Learn(i[trainIndex], d[trainIndex])
        if trainIndex % 10000 = 0 then printf "."
    printf "\n"

    // Stop time measuring
    let trainingEnd = DateTime.Now

    // Output the required time
    let trainingTime = trainingEnd - trainingStart
    printf "Time for training: %s\n" <| trainingTime.ToString()

    use resultImage = new Image<Rgb24>(imgWidth, imgHeight)

    // Start time measuring
    let predictionStart = DateTime.Now

    // Predict
    for predictIndex = 0 to sampleNum - 1 do
        let x = predictIndex % imgWidth
        let y = predictIndex / imgWidth
        let prediction = tar.Predict(i[predictIndex])
        resultImage[x, y] <- ColorFromPrediction prediction
        if predictIndex % 10000 = 0 then printf "."
    printf "\n"

    // Stop time measuring
    let predictionEnd = DateTime.Now

    // Output the required time
    let predictionTime = predictionEnd - predictionStart
    printf "Time for prediction: %s\n" <| predictionTime.ToString()

    // Save reference image
    resultImage.SaveAsJpeg(resultPath + filePrefix + "_Mandelbrot_regression.png")

    printf "Save network\n"

    // Save network in human-readable form
    // tar.SaveText(networkPath + filePrefix + "_Mandelbrot_regression.txt");

    // Save network in binary form
    tar.Save(networkPath + filePrefix + "_Mandelbrot_regression." + fileSuffix);

    0
