# Copyright (c) 2026 Jaime Yan. See LICENSE and CITATION.cff.
# Independent scheduled-grid computation from raw input tables.
source('R/clinical_extensions.R')

prepare_visit_matrix <- function(frames, score_unit='points_0_100', diameter_unit='mm') {
  validate_extensions(frames,score_unit,diameter_unit)
  s<-frames$subjects;d<-frames$tumor;weeks<-c(0,4,8,12);out<-list()
  for(stratum in extension_strata) for(arm in extension_arms) {
    ids<-sort(s$subject[s$arm==arm & (stratum=='All' | s$sex==stratum)])
    for(subject in ids) {
      subject_rows<-d[d$subject==subject,];subject_rows<-subject_rows[match(weeks,subject_rows$week),]
      observed<-!is.na(subject_rows$diameter)
      pattern<-if(all(observed)) 'Complete' else if(!any(observed)) 'All missing' else {
        last<-max(which(observed))
        if(last<length(weeks) && all(observed[seq_len(last)])) 'Trailing missing' else 'Intermittent or early missing'
      }
      for(i in seq_along(weeks)) {
        week<-weeks[i];n<-sum(!is.na(d$diameter[d$subject %in% ids & d$week==week]));expected<-length(ids)
        out[[length(out)+1]]<-data.frame(stratum,arm,subject,week,status=if(observed[i]) 'Observed' else 'Missing',pattern,n,expected,missing=expected-n)
      }
    }
  }
  if(length(out)==0) return(data.frame(stratum=character(),arm=character(),subject=character(),week=numeric(),status=character(),pattern=character(),n=integer(),expected=integer(),missing=integer()))
  ans<-do.call(rbind,out);rownames(ans)<-NULL;ans
}

draw_visit_matrix <- function(rows,stratum='All') {
  if(!stratum %in% extension_strata) stop('Unknown stratum')
  library(ggplot2)
  h<-rows[rows$stratum==stratum,];h$arm<-factor(h$arm,levels=extension_arms)
  h$subject<-factor(h$subject,levels=rev(sort(unique(h$subject))))
  h$mark<-ifelse(h$status=='Observed','O','X');h$text_color<-ifelse(h$status=='Observed','white','#333333')
  counts<-unique(h[c('arm','week','n','expected')])
  count_notes<-vapply(extension_arms,function(a) {
    z<-counts[as.character(counts$arm)==a,]
    paste(a,paste(paste0('W',z$week,': ',z$n,'/',z$expected),collapse='; '))
  },character(1))
  p<-ggplot(h,aes(x=factor(week,levels=c(0,4,8,12)),y=subject))+
    geom_tile(aes(fill=status),width=.96,height=.92,color='#555555',linewidth=.3)+
    geom_text(aes(label=mark,color=text_color),size=3)+scale_color_identity()+
    scale_fill_manual(values=c(Observed='#0072B2',Missing='white'),limits=c('Observed','Missing'),drop=FALSE)+
    facet_wrap(~arm,scales='free_y',drop=FALSE)+
    labs(title='Which scheduled observations are missing?',subtitle=paste('SYNTHETIC TEACHING DATA | Stratum:',stratum,'| Independent R computation'),
      x='Scheduled week',y='Fictional subject',fill='Observation',
      caption=paste(c(count_notes,'O = observed; X = missing. Fixed subject order; no clustering.',
        'Trailing missingness is not a withdrawal or censoring event.',
        'Jaime Yan | Personal noncommercial use | Attribution required'),collapse='\n'))+
    theme_minimal(base_size=11)+theme(panel.grid=element_blank(),plot.title=element_text(face='bold',size=18),plot.caption=element_text(hjust=0),legend.position='bottom',axis.text.y=element_text(size=9))
  if(nrow(h)==0) {
    p<-ggplot()+annotate('text',x=0,y=0,label='No participants')+theme_void()+labs(title='Scheduled observation matrix',subtitle=paste('Empty stratum:',stratum))
  }
  p
}