# Copyright (c) 2026 Jaime Yan. See LICENSE and CITATION.cff.
# Independent raw-data calculations; no Python result is read.
extension_arms <- c('Reference','Investigational')
extension_domains <- c('Fatigue','Pain','Sleep','Appetite','Mobility')
extension_kinds <- c('ecdf','completion','domain-intervals')
extension_strata <- c('All','F','M')

validate_extensions <- function(frames, score_unit='points_0_100', diameter_unit='mm') {
  if(score_unit!='points_0_100'||diameter_unit!='mm') stop('Expected points_0_100 and mm')
  req <- list(subjects=c('subject','arm','sex'),domains=c('subject','domain','week','score'),tumor=c('subject','week','diameter'))
  for(k in names(req)) if(is.null(frames[[k]])||!all(req[[k]] %in% names(frames[[k]]))) stop('Missing table/columns: ',k)
  s<-frames$subjects; q<-frames$domains; d<-frames$tumor
  if(nrow(s)==0) stop('Empty subject roster')
  if(anyNA(s[c('subject','arm','sex')])||any(trimws(s$subject)=='')||!all(s$arm %in% extension_arms)||!all(s$sex %in% c('F','M'))) stop('Invalid roster')
  for(k in names(req)) {
    key<-switch(k,subjects='subject',domains=c('subject','domain','week'),tumor=c('subject','week'))
    if(anyNA(frames[[k]][key])||anyDuplicated(frames[[k]][key])) stop('Missing or duplicate key')
  }
  for(k in c('domains','tumor')) {
    x<-frames[[k]][[if(k=='domains') 'score' else 'diameter']]
    if(!is.numeric(x)||any(!is.finite(x[!is.na(x)]))) stop('Nonnumeric or nonfinite measurement')
  }
  if(any(q$score<0|q$score>100,na.rm=TRUE)||any(d$diameter<=0,na.rm=TRUE)) stop('Measurement outside range')
  eq<-expand.grid(subject=s$subject,domain=extension_domains,week=c(0,12),stringsAsFactors=FALSE)
  ed<-expand.grid(subject=s$subject,week=c(0,4,8,12),stringsAsFactors=FALSE)
  keys<-function(x,cols) do.call(paste,c(x[cols],sep='\r'))
  if(!setequal(keys(q,names(eq)),keys(eq,names(eq)))) stop('Domain grid mismatch')
  if(!setequal(keys(d,names(ed)),keys(ed,names(ed)))) stop('Visit grid mismatch')
}

prepare_extension <- function(kind,frames,score_unit='points_0_100',diameter_unit='mm') {
  validate_extensions(frames,score_unit,diameter_unit)
  if(!kind %in% extension_kinds) stop('Unknown template')
  s<-frames$subjects; q<-frames$domains; d<-frames$tumor
  b<-q[q$week==0,c('subject','domain','score')];names(b)[3]<-'baseline'
  f<-q[q$week==12,c('subject','domain','score')];names(f)[3]<-'followup'
  pairs<-merge(b,f,by=c('subject','domain'));pairs$change<-pairs$followup-pairs$baseline
  rows<-list()
  add<-function(x) rows[[length(rows)+1]] <<- x
  for(stratum in extension_strata) for(arm in extension_arms) {
    ids<-s$subject[s$arm==arm & (stratum=='All' | s$sex==stratum)];expected<-length(ids)
    if(kind=='completion') {
      for(week in c(0,4,8,12)) {
        n<-sum(!is.na(d$diameter[d$subject %in% ids & d$week==week]))
        add(data.frame(stratum,arm,week,n,expected,missing=expected-n,percent=if(expected>0) 100*n/expected else NA_real_))
      }
    } else for(domain in if(kind=='ecdf') 'Fatigue' else extension_domains) {
      x<-pairs$change[pairs$subject %in% ids & pairs$domain==domain];x<-x[!is.na(x)];n<-length(x)
      meta<-data.frame(stratum,arm,domain,n,expected,missing=expected-n)
      if(kind=='ecdf') {
        if(n>0) { v<-sort(unique(x));add(cbind(meta[rep(1,length(v)),],change=v,probability=stats::ecdf(x)(v))) }
        else add(cbind(meta,change=NA_real_,probability=NA_real_))
      } else {
        estimate<-if(n>0) mean(x) else NA_real_;sd<-if(n>=2) stats::sd(x) else NA_real_
        width<-if(n>=2) stats::qt(.975,n-1)*sd/sqrt(n) else NA_real_
        add(cbind(meta,estimate,sd,lower=estimate-width,upper=estimate+width))
      }
    }
  }
  out<-do.call(rbind,rows);rownames(out)<-NULL;out
}

draw_extension <- function(kind,rows,stratum='All') {
  if(!kind %in% extension_kinds||!stratum %in% extension_strata) stop('Unknown template or stratum')
  g<-rows[rows$stratum==stratum,];if(nrow(g)==0) stop('Selected stratum absent')
  g$arm<-factor(g$arm,levels=extension_arms)
  cols<-setNames(c('#0072B2','#D55E00'),extension_arms)
  library(ggplot2)
  title<-switch(kind,ecdf='How broadly is improvement distributed?',completion='Who contributed at each scheduled visit?',`domain-intervals`='Symptom changes, with their precision.')
  note<-switch(kind,ecdf='Complete-pair distribution; no confidence band. Missing pairs are excluded, not censored.',completion='Observed / full scheduled roster. A missing measurement is not a withdrawal.',`domain-intervals`='Pointwise 95% t intervals; no multiplicity adjustment. These are not treatment contrasts.')
  if(kind=='ecdf') {
    lines<-list()
    for(a in extension_arms) {
      h<-g[as.character(g$arm)==a & !is.na(g$change),]
      if(nrow(h)>0) {
        first<-h[1,];first$change<--100;first$probability<-0
        last<-h[nrow(h),];last$change<-100;last$probability<-1
        lines[[a]]<-rbind(if(min(h$change)>-100) first else NULL,h,last)
      }
    }
    z<-if(length(lines)) do.call(rbind,lines) else g
    limits<-range(c(-40,15,g$change-2,g$change+2),na.rm=TRUE)
    labels<-sapply(extension_arms,function(a){h<-g[as.character(g$arm)==a,][1,];paste0(a,': ',h$n,'/',h$expected,' complete')})
    p<-ggplot(z,aes(change,probability,color=arm,linetype=arm))+geom_step(linewidth=.8,na.rm=TRUE)+geom_vline(xintercept=0,linetype=3,color='grey40')+
      coord_cartesian(xlim=limits,ylim=c(0,1.04))+labs(x='Week-12 minus baseline Fatigue (points); negative = improvement',y='Proportion with change <= x')+
      scale_color_manual(values=cols,labels=labels,drop=FALSE)+scale_linetype_manual(values=c('solid','dashed'),labels=labels,drop=FALSE)
  } else if(kind=='completion') {
    g$label<-paste0(g$n,'/',g$expected)
    g$label_y<-vapply(seq_len(nrow(g)),function(i) {
      other<-g$percent[g$week==g$week[i] & g$arm!=g$arm[i]]
      above<-is.na(other)||(!is.na(g$percent[i])&&(g$percent[i]>other||(g$percent[i]==other&&g$arm[i]=='Reference')))
      g$percent[i]+if(above) 5 else -6
    },numeric(1))
    p<-ggplot(g,aes(week,percent,color=arm,shape=arm,linetype=arm))+geom_line(linewidth=.8,na.rm=TRUE)+geom_point(size=2.5,na.rm=TRUE)+
      geom_text(aes(label=label,y=label_y),show.legend=FALSE,size=3.5,na.rm=TRUE)+
      scale_x_continuous(breaks=c(0,4,8,12))+scale_y_continuous(breaks=c(0,20,40,60,80,100))+coord_cartesian(ylim=c(-5,110))+labs(x='Scheduled week',y='Observed measurements (%)')
  } else {
    g$y<-match(g$domain,extension_domains)+ifelse(g$arm=='Reference',-.125,.125)
    p<-ggplot(g,aes(estimate,y,color=arm,shape=arm))+geom_vline(xintercept=0,linetype=3,color='grey40')+
      geom_segment(aes(x=lower,xend=upper,yend=y),linewidth=.8,na.rm=TRUE)+geom_point(size=2.5,na.rm=TRUE)+
      scale_y_reverse(breaks=1:5,labels=extension_domains)+labs(x='Mean within-person change (points); negative = improvement',y=NULL)
    hi<-max(c(0,g$upper,g$estimate),na.rm=TRUE);lo<-min(c(0,g$lower,g$estimate),na.rm=TRUE);span<-max(hi-lo,1)
    p<-p+geom_text(aes(x=hi+.18*span,label=paste0(n,'/',expected)),show.legend=FALSE,size=3.2)+
      annotate('text',x=hi+.18*span,y=.45,label='Pairs / roster',size=3.2)+scale_x_continuous(limits=c(lo-.05*span,hi+.32*span))
  }
  if(kind!='ecdf') p<-p+scale_color_manual(values=cols,drop=FALSE)+scale_shape_manual(values=c(16,15),drop=FALSE)
  p+labs(title=title,subtitle=paste('SYNTHETIC TEACHING DATA | Stratum:',stratum,'| Independent R computation'),
      caption=paste0(note,'\nJaime Yan | Clinical Figure Library | Personal noncommercial use | Attribution and citation required'),color=NULL,linetype=NULL,shape=NULL)+
    theme_minimal(base_size=12)+theme(legend.position='bottom',panel.grid.minor=element_blank(),plot.title=element_text(size=20,face='bold'),
      plot.subtitle=element_text(margin=margin(b=20)),plot.caption=element_text(hjust=0,size=9,margin=margin(t=20)),plot.margin=margin(20,24,16,20))
}