---
title: "pupil registration"
output:
  html_document: 
    toc: true
    toc_float: true
    code_folding: hide
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = FALSE, warning = FALSE, message = FALSE)

library(tidyverse)
library(qgam)
```


## Data loading 

First, load the cleaned pupil data.

```{r}
# load cleaned pupils data
pupils <- readRDS(file = here::here("data", "cleaned_pupils.rds"))

```


Next, define a function that will be applied to each trajectory to register the data.

```{r}
# dat is a single clean but unregistered pupil trajectory
register_pupil_trajectory = function(dat){
  
  ##### Calculate end of light test
  end_df = dat %>%
    filter(!is.na(major_axis)) %>% select(frame, major_axis)
  
  # calculate derivatives
  der_vec = c(0, diff(end_df$major_axis)) 
  der_vec2 = c(0, diff(end_df$major_axis, 2),0) 
  
  ending_frame = end_df %>%
    mutate(prev = c(0, abs(diff(major_axis))),
           dat = abs(der_vec + der_vec2)/2) %>%
    filter(frame > 300, prev < 5, dat > 10) %>%
    arrange(frame) %>% first() %>% pull(frame)
  
  if(is.na(ending_frame)){ending_frame = max(end_df$frame, na.rm= TRUE)} 
  
  # remove frames after last frame
  dat = dat %>% filter(frame <= ending_frame)
  
  ##### Calculate beginning of light test
  n_knots = if_else(nrow(dat) > 30, 25, nrow(dat) - 5)
  # smooth data to get better derivative calculations
  fit <- qgam(major_axis~ s(frame, k=n_knots, bs="cs"), data=dat, qu = 0.5)
  
  ysmooth = predict(fit, newdata = data.frame(frame = dat$frame))
  dat$ysmooth = as.numeric(ysmooth)
  
  # get first derivative, local min, and local max
  der_vec = c(0, diff(ysmooth)) 
  sign_dif = c(sign(diff(sign(der_vec))), 0)
  der_vec_max = ifelse(sign_dif == -1, 1, 0)
  der_vec_min = ifelse(sign_dif == 1, 1, 0)
  
  ##### calculate beginning of light test
  beginning_frame = dat %>% 
    mutate(extremum = factor(der_vec_max + der_vec_min)) %>%
    filter(extremum == 1, frame < 300) %>% 
    mutate(slope = ysmooth - lead(ysmooth)) %>%
    arrange(-slope) %>% 
    first() %>% pull(frame) %>% round()
  
  dat %>% 
    mutate(beginning_frame = beginning_frame,
           centered = round(frame - beginning_frame)) %>%
    filter(frame >= beginning_frame) %>%
    as_tibble() %>%
    select(timeid, frame, centered, eye, major_axis, ysmooth)

}

```


## Registration

Apply function to register data.

```{r}
pupils_nest = pupils %>%
  select(timeid, frame, eye, major_axis) %>%
  mutate(timeid_eye = str_c(timeid, "_", eye)) %>%
  nest(data = c(-timeid_eye))

pupils_reg = map_dfr(pupils_nest$data, register_pupil_trajectory)

rm(pupils_nest)
```


Create percent change variable.


```{r}

## Creating Percent Change Variable 
pupils_reg = pupils_reg %>% 
  arrange(timeid, eye, centered) %>%
  group_by(timeid, eye) %>%
  mutate(percent_change = (ysmooth - first(ysmooth))/first(ysmooth) * 100) %>%
  ungroup() %>%
  select(timeid, eye, frame, centered, percent_change)



```


## Saving fully cleaned, smoothed, and registered data

```{r}
saveRDS(pupils_reg, 
        file = here::here("data", "cleaned_registered_pupils.rds"))
```
