'''

This program is meant to take the Trace files outputted from ebFRET and determine the FRET efficiency (Viterbi_Mean) of each state in a given HMM
before the state transitions to another state. The program also determines how long each FRET state lasts for, so you can correlate FRET efficiency
with the lifetime of the state.

'''

###################################################################
# 2 state HMM
###################################################################

library(tidyverse)

#Define global variables

path_var <- r"(F:\Backup Plus (Backup)\Analysis\ebFRET_Analysis_4.21.2022\Stop - 25 nM TC - Tig - Lifetime of Final Events\Combo\ebFRET)"
file_name <- "2stateHMM_Stop_IRES_25nMTC_Tigecycline_Lifetime_of_Final_Sampling_Event_T1_to_T2_ChungKennedy_Filtered_Traces.dat"
setwd(path_var)
time_interval <- 0.04 #Time interval in seconds


#Load in the Full Report Trace outputs from ebFRET

file_location <- paste(path_var, file_name, sep = "\\")
TypeAll_traces_2states_ebFRET <- read.table(file_location, quote="\"", comment.char="")


#Name the columns for each dataframe

colnames(TypeAll_traces_2states_ebFRET) <- c("Trace", "Donor", "Acceptor", "FRET", "Viterbi_State", "Viterbi_Mean")


#Subset the dataframe to include only Viterbi_State and Viterbi_Mean columns

test <- TypeAll_traces_2states_ebFRET[, 5:6]


#Generate the index list from test dataset. Convert to a dataframe and rename column
index_test <- data.frame(cumsum(rle(test$Viterbi_Mean)$lengths))
names(index_test) <- c("V1")


#Use index dataframe to extract values at the specified indices
testtest <- test[as.matrix(index_test), ]


#Generate a new index dataframe that has two zeros at the beginning of it (needs to be an odd number)

new <- data.frame(c(0,0))
names(new) <- c("V1")
new_index <- rbind(new, index_test)


#Find the difference between rows

Time <- data.frame(diff(new_index$V1))
names(Time) <- c("Diff")


#Remove the first entry (zero) from the start of the diff_test dataframe

Time <- Time[-1, ]


#Merge the testtest and diff_test dataframes

merged_test <- cbind(testtest, Time)


#Convert the values in the Time column from Frames to Seconds

merged_test <- merged_test %>% mutate(Time = Time*time_interval)


#Filter by Viterbi_State

State1_FRET <- merged_test %>% filter(Viterbi_State == 1)
State2_FRET <- merged_test %>% filter(Viterbi_State == 2)


#Filter out values that are less than zero or greater than one for the Viterbi_Mean

State1_FRET <- State1_FRET %>% filter(Viterbi_Mean >= 0)
State1_FRET <- State1_FRET %>% filter(Viterbi_Mean <= 1)

State2_FRET <- State2_FRET %>% filter(Viterbi_Mean >= 0)
State2_FRET <- State2_FRET %>% filter(Viterbi_Mean <= 1)

write.table(State2_FRET$Viterbi_Mean, "Stop_IRES_15.6nMTC_FRET_of_Initial_Sampling_Event_T1_to_T2_ChungKennedy_Filtered_Traces_State2_2HMM.txt", row.names = FALSE, col.names = FALSE)
#Make a histogram for each FRET population

hist(State1_FRET$Viterbi_Mean, breaks = 30, xlim=c(0,1)) #Baseline FRET (background)
hist(State2_FRET$Viterbi_Mean, breaks = 30, xlim=c(0,1)) #FRET for each sampling event

#Correlation Plots between Viterbi_Mean and Time(s)

plot(State1_FRET$Time, State1_FRET$Viterbi_Mean)


plot(State2_FRET$Time, State2_FRET$Viterbi_Mean)
abline(lm(State2_FRET$Viterbi_Mean ~ State2_FRET$Time), col = "red", lwd = 3)

plot(State2_FRET$Viterbi_Mean, State2_FRET$Time)
abline(lm(State2_FRET$Time ~ State2_FRET$Viterbi_Mean), col = "red", lwd = 3)

lm_xy <- lm(State2_FRET$Viterbi_Mean~State2_FRET$Time)
lm_xy