---
title: "Basic example"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Basic example}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r, include = FALSE}
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)
```

```{r setup}
library(designit)
library(ggplot2)
library(dplyr)
library(tidyr)
```

# Plate layout with two factors

## The samples

Samples of a 2-condition in-vivo experiment are to be
placed on 48 well plates.

These are the conditions

```{r}
# conditions to use
conditions <- data.frame(
  group = c(1, 2, 3, 4, 5),
  treatment = c(
    "vehicle", "TRT1", "TRT2",
    "TRT1", "TRT2"
  ),
  dose = c(0, 25, 25, 50, 50)
)

gt::gt(conditions)
```

We will have 3 animals per groups with 4 replicates each

```{r}
# sample table (2 animals per group with 3 replicates)
n_reps <- 4
n_animals <- 3
animals <- bind_rows(replicate(n_animals, conditions, simplify = FALSE),
  .id = "animal"
)
samples <- bind_rows(replicate(n_reps, animals, simplify = FALSE),
  .id = "replicate"
) |>
  mutate(
    SampleID = paste0(treatment, "_", animal, "_", replicate),
    AnimalID = paste0(treatment, "_", animal)
  ) |>
  mutate(dose = factor(dose))

samples |>
  head(10) |>
  gt::gt()
```
## Plate layout requirements

Corner wells of the plates should be left empty.
This means on a 48 well plate we can place 44 samples.
Since we have `r nrow(samples)` samples, they will fit on
`r ceiling(nrow(samples)/44)` plates

```{r}
n_samp <- nrow(samples)
n_loc_per_plate <- 48 - 4
n_plates <- ceiling(n_samp / n_loc_per_plate)

exclude_wells <- expand.grid(plate = seq(n_plates), column = c(1, 8), row = c(1, 6))
```

## Setting up a Batch container

Create a BatchContainer object that provides all possible locations

```{r}
bc <- BatchContainer$new(
  dimensions = c("plate" = n_plates, "column" = 8, "row" = 6),
  exclude = exclude_wells
)
bc

bc$n_locations
bc$exclude
bc$get_locations() |> head()
```

## Moving samples

Use random assignment function to place samples to plate locations

```{r}
bc <- assign_random(bc, samples)

bc$get_samples()
bc$get_samples(remove_empty_locations = TRUE)
```

Plot of the result using the `plot_plate` function

```{r, fig.width=6, fig.height=3.5}
plot_plate(bc,
  plate = plate, column = column, row = row,
  .color = treatment, .alpha = dose
)
```
To not show empty wells, we can directly plot the sample table as well

```{r, fig.width=6, fig.height=3.5}
plot_plate(bc$get_samples(remove_empty_locations = TRUE),
  plate = plate, column = column, row = row,
  .color = treatment, .alpha = dose
)
```

To move individual samples or manually assigning all locations we can use the 
`batchContainer$move_samples()` method

To swap two or more samples use:

**Warning**: This will change your BatchContainer in-place.

```{r, fig.width=6, fig.height=3.5}
bc$move_samples(src = c(1L, 2L), dst = c(2L, 1L))

plot_plate(bc$get_samples(remove_empty_locations = TRUE),
  plate = plate, column = column, row = row,
  .color = treatment, .alpha = dose
)
```

To assign all samples in one go, use the option `location_assignment`.

**Warning**: This will change your BatchContainer in-place.

The example below orders samples by ID and adds the empty locations afterwards
```{r, fig.width=6, fig.height=3.5}
bc$move_samples(
  location_assignment = c(
    1:nrow(samples),
    rep(NA, (bc$n_locations - nrow(samples)))
  )
)

plot_plate(bc$get_samples(remove_empty_locations = TRUE, include_id = TRUE),
  plate = plate, column = column, row = row,
  .color = .sample_id
)
```

## Run an optimization

The optimization procedure is invoked with e.g. `optimize_design`.
Here we use a simple shuffling schedule: 
swap 10 samples for 100 times, then swap 2 samples for 400 times.

To evaluate how good a layout is, we need a scoring function. 

This function will assess how well treatment 
and dose are balanced across the two plates.

```{r}
bc <- optimize_design(bc,
  scoring = osat_score_generator(
    batch_vars = "plate",
    feature_vars = c("treatment", "dose")
  ),
  # shuffling schedule
  n_shuffle = c(rep(10, 200), rep(2, 400))
)
```

Development of the score can be viewed with

```{r, fig.width=3.5, fig.height=3}
bc$plot_trace()
```

The layout after plate batching looks the following

```{r, fig.width=6, fig.height=3.5}
plot_plate(bc$get_samples(remove_empty_locations = TRUE),
  plate = plate, column = column, row = row,
  .color = treatment, .alpha = dose
)
```

Looking at treatment, we see it's evenly distributed across the plates

```{r, fig.width=4, fig.height=3.5}
ggplot(
  bc$get_samples(remove_empty_locations = TRUE),
  aes(x = treatment, fill = treatment)
) +
  geom_bar() +
  facet_wrap(~plate)
```

## Customizing the plate layout

To properly distinguish between empty and excluded locations one can do the
following.

* Supply the BatchContainer directly
* set `add_excluded = TRUE`, set `rename_empty = TRUE`
* supply a custom color palette
* excluded wells have NA values and can be colored with `na.value`

```{r, fig.width=6, fig.height=3.5}
color_palette <- c(
  TRT1 = "blue", TRT2 = "purple",
  vehicle = "orange", empty = "white"
)

plot_plate(bc,
  plate = plate, column = column, row = row,
  .color = treatment, .alpha = dose,
  add_excluded = TRUE, rename_empty = TRUE
) +
  scale_fill_manual(values = color_palette, na.value = "darkgray")
```

To remove all empty wells from the plot, hand the pruned sample list.
to plot_plate rather than the whole BatchContainer.
You can still assign your own colors.

```{r, fig.width=6, fig.height=3.5}
plot_plate(bc$get_samples(remove_empty_locations = TRUE),
  plate = plate, column = column, row = row,
  .color = treatment, .alpha = dose
) +
  scale_fill_viridis_d()
```

Note: removing all empty and excluded wells will lead to omitting 
completely empty rows or columns!

```{r, fig.width=6, fig.height=3.5}
plot_plate(bc$get_samples(remove_empty_locations = TRUE) |>
  filter(column != 2),
plate = plate, column = column, row = row,
.color = treatment, .alpha = dose
) +
  scale_fill_viridis_d()
```