# usage: R -e "shiny::runApp(port = 8080)"
# usage: R -e "shiny::runApp(host = '127.0.0.1', port = 8080)"

library(shiny)
library(mvbernoulli)
library(tensorPredictors)

# configuration
color.palet <- hcl.colors(64, "YlOrRd", rev = TRUE)

# GMLM parameters
n <- 250
p <- c(4, 4)
q <- c(2, 2)
r <- 2


eta1 <- 0                                                       # intercept
linspace <- seq(-1, 1, length.out = 4)

# 270 deg (90 deg clockwise) rotation of matrix layout
#
# Used to get proper ploted matrices cause `image` interprets the `z` matrix as
# a table of `f(x[i], y[j])` values, so that the `x` axis corresponds to row
# number and the `y` axis to column number, with column 1 at the bottom,
# i.e. a 90 degree counter-clockwise rotation of the conventional printed layout
# of a matrix. By first calling `rot270` on a matrix before passing it to
# `image` the plotted matrix layout now matches the conventional printed layout.
rot270 <- function(A) {
    t(A)[, rev(seq_len(nrow(A))), drop = FALSE]
}
plot.mat <- function(mat, add.values = FALSE, zlim = range(mat)) {
    par(oma = rep(0, 4), mar = rep(0, 4))
    img <- rot270(mat)
    image(x = seq_len(nrow(img)), y = seq_len(ncol(img)), z = img,
          zlim = zlim, col = color.palet, xaxt = "n", yaxt = "n", bty = "n")
    if (add.values) {
        text(x = rep(seq_len(nrow(img)), ncol(img)),
             y = rep(seq_len(ncol(img)), each = nrow(img)),
             round(img, 2), adj = 0.5, col = "black")
    }
}

AR <- function(rho, dim) {
    rho^abs(outer(seq_len(dim), seq_len(dim), `-`))
}
AR.inv <- function(rho, dim) {
    A <- diag(c(1, rep(rho^2 + 1, dim - 2), 1))
    A[abs(.row(dim(A)) - .col(dim(A))) == 1] <- -rho
    A / (1 - rho^2)
}

# User Interface (page layout)
ui <- fluidPage(
    tags$head(HTML("
        <script type='text/x-mathjax-config'>
            MathJax.Hub.Config({
                tex2jax: { inlineMath: [['$','$']] }
            });
        </script>
    ")),
    withMathJax(),

    titlePanel("Ising Model Simulation Data Generation"),
    sidebarLayout(
        sidebarPanel(
            h2("Settings"),
            h4("c1 (influence of $\\eta_1$)"),
            sliderInput("c1", "", min = 0, max = 1, value = 1, step = 0.01),
            h4("c2 (influence of $\\eta_2$)"),
            sliderInput("c2", "", min = 0, max = 1, value = 1, step = 0.01),
            sliderInput("y", "y", min = -1, max = 1, value = 0, step = 0.01),
            fluidRow(
                column(6,
                    radioButtons("alphaType", "Type: $\\boldsymbol{\\alpha}_k$",
                        choices = list(
                            "linspace" = "linspace", "squared" = "squared", "QR" = "QR"
                        ),
                        selected = "linspace"
                    )
                ),
                column(6,
                    radioButtons("OmegaType", "Type: $\\boldsymbol{\\Omega}_k$",
                        choices = list(
                            "Identity" = "identity", "AR$(\\rho)$" = "AR", "AR$(\\rho)^{-1}$" = "AR.inv"
                        ),
                        selected = "AR"
                    )
                )
            ),
            sliderInput("rho", "rho", min = -1, max = 1, value = -0.55, step = 0.01),
            actionButton("reset", "Reset")
        ),
        mainPanel(
            fluidRow(
                column(4,
                    h3("$\\eta_{y,1}$"),
                    plotOutput("eta_y1"),
                ),
                column(4,
                    h3("$\\eta_{y,2}$"),
                    plotOutput("eta_y2"),
                ),
                column(4,
                    h3("$\\boldsymbol{\\Theta}_y$"),
                    plotOutput("Theta_y")
                )
            ),
            fluidRow(
                column(4, offset = 2,
                    h3("Expectation $\\mathbb{E}[\\mathcal{X}\\mid Y = y]$"),
                    plotOutput("expectationPlot"),
                ),
                column(4,
                    h3("Covariance $\\operatorname{Cov}(\\text{vec}(\\mathcal{X})\\mid Y = y)$"),
                    plotOutput("covariancePlot"),
                    textOutput("covRange"),
                )
            ),
            fluidRow(
                column(8, offset = 4,
                    h3("iid samples $(X_i, y_i)$ with $y_i \\sim U[-1, 1]$ sorted")
                ),
                column(4,
                    "Conditional Expectations",
                    plotOutput("cond_expectations")
                ),
                column(4,
                    "observations sorted by $y_i$",
                    plotOutput("sample_sorted_y")
                ),
                column(4,
                    "observations sorted (lexicographic order) by $\\mathcal{X}_i$",
                    plotOutput("sample_sorted_X")
                ),
            ),
            fluidRow(
                column(6,
                    h3("Sample Mean"),
                    plotOutput("sampleMean")
                ),
                column(6,
                    h3("Sample Cov"),
                    plotOutput("sampleCov")
                )
            ),
            h2("Explanation"),
            "
            The response $y$ follows a continuous uniform distributed
            $y\\sim U[-1, 1]$ from which $\\mathcal{F}_y$ is computed as

            $$\\mathcal{F}_y = \\begin{pmatrix}
                \\cos(\\pi y) & -\\sin(\\pi y) \\\\
                \\sin(\\pi y) & \\cos(\\pi y)
            \\end{pmatrix}.$$

            Next are the GMLM parameters (for 'linspace' or 'squared' type $\\boldsymbol{\\alpha}_k$
            with the 'QR' type being random semi-orthogonal matrices) which are set to be
            $$\\overline{\\eta}_1 = 0$$
            $$\\boldsymbol{\\alpha}_k^{\\text{linspace}} = \\begin{pmatrix}
                -1 & 1 \\\\
                -1/3 & 1/3 \\\\
                1/3 & -1/3 \\\\
                1 & -1
            \\end{pmatrix},\\qquad\\boldsymbol{\\alpha}_k^{\\text{squared}} = \\begin{pmatrix}
                -1 & 1 \\\\
                -1/3 & 1/9 \\\\
                1/3 & 1/9 \\\\
                1 & 1
            \\end{pmatrix}$$
            for $k = 1,2$. The two-way interactions are modeled via the $\\boldsymbol{\\Omega}_k$ which
            are the identity $\\boldsymbol{I}_4$ or one of
            $$\\operatorname{AR}(\\rho) = \\begin{pmatrix}
                {\\color{gray}1} & \\rho^1 & \\rho^2 & \\rho^3 \\\\
                \\rho^1 & {\\color{gray}1} & \\rho^1 & \\rho^2 \\\\
                \\rho^2 & \\rho^1 & {\\color{gray}1} & \\rho^1 \\\\
                \\rho^3 & \\rho^2 & \\rho^1 & {\\color{gray}1}
            \\end{pmatrix},\\qquad
            \\operatorname{AR}(\\rho)^{-1} = \\frac{1}{1 - \\rho^2}\\begin{pmatrix}
                {\\color{gray}1} & -\\rho & 0 & 0 \\\\
                -\\rho & {\\color{gray}1+\\rho^2} & -\\rho & 0 \\\\
                0 & -\\rho & {\\color{gray}1+\\rho^2} & -\\rho \\\\
                0 & 0 & -\\rho & {\\color{gray}1}
            \\end{pmatrix}.$$

            The natural parameters given $y$ are then
            $$\\boldsymbol{\\eta}_{y,1} \\equiv \\overline{\\boldsymbol{\\eta}}_1
                + \\mathcal{F}_y\\times_{k\\in[2]}\\boldsymbol{\\alpha}_k,$$
            $$\\boldsymbol{\\eta}_{y,2} \\equiv \\bigotimes_{k = 2}^{1}\\boldsymbol{\\Omega}_k.$$
            With that the conditional Ising model parameters are
            $$\\boldsymbol{\\theta}_y = \\operatorname{vech}(
                \\operatorname{diag}(\\boldsymbol{\\eta}_{y,1})
                    + (\\boldsymbol{1}_p \\boldsymbol{1}_p' - \\mathbf{I}_p) \\odot \\boldsymbol{\\eta}_{y,2}
            )$$
            which to sample the predictors via the conditional distribution
            $$\\operatorname{vec}(\\mathcal{X})\\mid Y = y \\sim \\text{Ising}(\\boldsymbol{\\theta}_y).$$
            "
        )
    )
)

# Server logic
server <- function(input, output, session) {

    Fy <- reactive({
        phi <- pi * input$y
        matrix(c(
            cos(phi), -sin(phi),
            sin(phi), cos(phi)
        ), 2, 2, byrow = TRUE)
    })
    alphas <- reactive({
        switch(input$alphaType,
            "linspace" = list(
                             matrix(c(linspace, rev(linspace)), length(linspace), 2),
                             matrix(c(linspace, rev(linspace)), length(linspace), 2)
                         ),
            "squared"  = list(
                             matrix(c(linspace, linspace^2), length(linspace), 2),
                             matrix(c(linspace, linspace^2), length(linspace), 2)
                         ),
            "QR"       = Map(function(pj, qj) {
                             qr.Q(qr(matrix(rnorm(pj * qj), pj, qj)))
                         }, p, q)
        )
    })
    Omegas <- reactive({
        switch(input$OmegaType,
            "identity" = Map(diag, p),
            "AR"       = Map(AR, list(input$rho), dim = p),
            "AR.inv"   = Map(AR.inv, list(input$rho), dim = p)
        )
    })

    eta_y1 <- reactive({
        input$c1 * (mlm(Fy(), alphas()) + c(eta1))
    })
    eta_y2 <- reactive({
        input$c2 * Reduce(`%x%`, rev(Omegas()))
    })

    # compute Ising model parameters from GMLM parameters given single `Fy`
    theta_y <- reactive({
        vech(diag(c(eta_y1())) + (1 - diag(nrow(eta_y2()))) * eta_y2())
    })

    E_y <- reactive({
        mvbernoulli::ising_expectation(theta_y())
    })
    Cov_y <- reactive({
        mvbernoulli::ising_cov(theta_y())
    })

    random_sample <- reactive({
        c1 <- input$c1
        c2 <- input$c2
        eta_y_i2 <- eta_y2()

        y <- sort(runif(n, -1, 1))
        X <- sapply(y, function(y_i) {
            phi <- pi * y_i
            Fy_i <- matrix(c(
                cos(phi), -sin(phi),
                sin(phi), cos(phi)
            ), 2, 2)

            eta_y_i1 <- c1 * (mlm(Fy_i, alphas()) + c(eta1))

            theta_y_i <- vech(diag(c(eta_y_i1)) + (1 - diag(nrow(eta_y_i2))) * eta_y_i2)

            ising_sample(1, theta_y_i)
        })
        attr(X, "p") <- prod(p)

        as.mvbmatrix(X)
    })

    cond_expectations <- reactive({
        c1 <- input$c1
        c2 <- input$c2
        eta_y_i2 <- eta_y2()

        y <- seq(-1, 1, length.out = 50)
        t(sapply(y, function(y_i) {
            phi <- pi * y_i
            Fy_i <- matrix(c(
                cos(phi), -sin(phi),
                sin(phi), cos(phi)
            ), 2, 2)

            eta_y_i1 <- c1 * (mlm(Fy_i, alphas()) + c(eta1))

            theta_y_i <- vech(diag(c(eta_y_i1)) + (1 - diag(nrow(eta_y_i2))) * eta_y_i2)

            ising_expectation(theta_y_i)
        }))
    })

    output$eta_y1 <- renderPlot({
        plot.mat(eta_y1(), add.values = TRUE, zlim = c(-2, 2))
    }, res = 108)
    output$eta_y2 <- renderPlot({
        plot.mat(eta_y2())
    })
    output$Theta_y <- renderPlot({
        plot.mat(vech.pinv(theta_y()))
    })

    output$expectationPlot <- renderPlot({
        plot.mat(matrix(E_y(), p[1], p[2]), add.values = TRUE, zlim = c(0, 1))
    }, res = 108)
    output$covariancePlot <- renderPlot({
        plot.mat(Cov_y())
    })
    output$covRange <- renderText({
        paste(round(range(Cov_y()), 3), collapse = " - ")
    })
    output$cond_expectations <- renderPlot({
        plot.mat(cond_expectations(), zlim = 0:1)
    })
    output$sample_sorted_y <- renderPlot({
        plot.mat(random_sample())
    })
    output$sample_sorted_X <- renderPlot({
        X <- random_sample()
        plot.mat(X[do.call(order, as.data.frame(X)), ])
    })
    output$sampleMean <- renderPlot({
        Xmean <- matrix(colMeans(random_sample()), p[1], p[2])
        plot.mat(Xmean, add.values = TRUE, zlim = c(0, 1))
    }, res = 108)
    output$sampleCov <- renderPlot({
        plot.mat(cov(random_sample()))
    })

    observeEvent(input$reset, {
        updateNumericInput(session, "c1",  value = 1)
        updateNumericInput(session, "c2",  value = 1)
        updateNumericInput(session, "y",   value = 0)
        updateNumericInput(session, "rho", value = -0.55)
        updateRadioButtons(session, "OmegaType", selected = "AR")
        updateRadioButtons(session, "alphaType", selected = "squared")
    })
}

# launch Shiny Application (start server)
shinyApp(ui = ui, server = server)