#Vectorized code with correct dispersal array and correct order of events#
#setwd("~/Documents/PhD-project Laura/Spatial Model/Code/")
#rm(list=ls())

#correct the dt things#


###Run model#####
total_time = 0
tol <- dt/100
for(y in 1:year) {
  print(paste("year:", y))
  index_dt = 1
  index_t = 1
  for(t in seq(from = 1, to = num_days, by = dt)) {
    #print(t)
    #print(max(Ad_ill_new))
    ####Functions#####
    # Calculate the current day for output
    current_day <- floor(t)
    
    #FUNCTIONS
    #TEMPERATURE change within a year
    Temp <- Ta*cos((t - td) * 2 * pi / 365) + Tc
    #Temperature correction for rates
    tc <- (Temp-4)/(22-4)
    
    #FLORAL RESOURCES  #*
    #Functions flower food availability per habitat
    B1 <- B1l + B1m * dweibull((t-T1s)/T1d, k1, 1) 
    B4 <- B4l + B4m * dweibull((t-T4s)/T4d, ks, 1) #annual
    #B4 <- B4l + B4m * dweibull((t-T4s)/T4d, ks, 1) #perennial
    
    ###Update flowers that vary over time####
    B[woody_habitat] <- B1
    B[flower_margin] <- B4
    
    #APHID DYNAMICS 
    # Aphid resource availability and carrying capacity (Weibull, reverse Weibull)
    n1 <- n1s * dweibull((t-A1s)/A1d,ks,1) + n1a*dweibull((t-A1as)/A1ad,ks,1)
    K2 <- Km * dweibull((A2f-t)/A2d,ks,1)+phi
    K3 <- Km * dweibull((A3f-t)/A3d,ks,1)+phi
    n2 <- ifelse(K2>0.01*Km, 1, 0)
    n3 <- ifelse(K3>0.01*Km, 1, 0)
    
    #Aphid population growth rate
    r1 <- n1 * rN1 * tc * dt
    r2 <- n2 * rN2 * tc * dt
    r3 <- n3 * rN3 * tc * dt
    
    #time delay representing inactive stages
    tau <- tauh/tc  #tau niet delen door dt, want t is in eenheden van 1 dag  
    #needs to be rounded because array only every day. 
    
    #HOVERFLY HIBERNATION
    #Release from hibernation (around day H0=30) h0
    h0 <- dnorm(t, H0, sigma0) * dt
    
    #Induction into hibernation (around day H1=183) h1  
    h1 <- dnorm(t, H1, sigma1) * dt
    
    
    ####POPULATION DYNAMICS####
    
    ###Aphids####
    
    #aphids infest in other habitats 
    if(abs(t - T2i) < tol) {
      N_prev[early_crop] <- N2i
    }
    
    if(abs(t - T3i)< tol) {
      N_prev[late_crop] <- N3i
    }
    
    
    #Woody habitat
    #note that the function f() is already taking care of temperature (tc) and timestep (dt)

    if(t<=Wt1){
    ###Aphids
    N_susceptible <- pmax(N_prev[woody_habitat] - Nr, 0) #individuals that can die 
    N_save <- pmin(Nr, N_prev[woody_habitat]) #individuals that are hidden and cannot die.
    
      N_new[woody_habitat] <- 
      N_prev[woody_habitat] * r1 * (1-N_prev[woody_habitat]/K1) + #growth
    (N_susceptible - pmin(Juv_prev[woody_habitat]*f(N_susceptible), N_susceptible))*(1-mu1*tc*dt) + #mortality (predation followed by background)
        N_save #hidden aphids. 
        
    

    
    ###Juveniles hoverflies####
    Juv_new[woody_habitat] <- Juv_prev[woody_habitat] * (1-muj(N_susceptible)) * 
                              (1-e(N_susceptible)) +
                              G1(N_susceptible) * Ad_well_prev[woody_habitat] * (1-muB * tc * dt)
      if(min(Juv_new[woody_habitat]) < 0) {print(paste("Stop at year ", y, " and day ", t))} 
      }
    
    #winter mortality
    if(abs(t - Wt1) < tol) {
      N_new[woody_habitat] <- WN * N_new[woody_habitat]
      Juv_new[woody_habitat] <- 0
    } 


    #Early crop
    if(T2i <= t & t <= Wt2){
    #Aphids H2
    N_new[early_crop] <- N_prev[early_crop] * r2 * (1-N_prev[early_crop]/K2) + 
    (N_prev[early_crop] - pmin(Juv_prev[early_crop]*f(N_prev[early_crop]), N_prev[early_crop]))*(1-mu2*tc*dt)
     
    
    Juv_new[early_crop] <- Juv_prev[early_crop] * (1-muj(N_prev[early_crop])) * 
                              (1-e(N_prev[early_crop])) +
                              G2(N_prev[early_crop]) * Ad_well_prev[early_crop] * (1-muB * tc * dt)
    }


    #Harvest mortality
    if(abs(t - Wt2) < tol) {
      Juv_new[early_crop] <- 0
      N_new[early_crop] <- 0
    }

    
    #Late crop
    if(T3i <= t & t <= Wt3){
    #Aphids H3
    N_new[late_crop] <- N_prev[late_crop] * r3 * (1-N_prev[late_crop]/K3) + 
    (N_prev[late_crop] - pmin(Juv_prev[late_crop]*f(N_prev[late_crop]),N_prev[late_crop]))*(1-mu3*tc*dt)
    
    

    #Juveniles
    Juv_new[late_crop] <- Juv_prev[late_crop] * (1-muj(N_prev[late_crop])) * 
                              (1-e(N_prev[late_crop])) +
                              G2(N_prev[late_crop]) * Ad_well_prev[late_crop] * (1-muB * tc * dt)
      }

    if(abs(t - Wt3) < tol) { #mortality
      Juv_new[late_crop] <- 0
      N_new[late_crop] <- 0
    }


    
    ###Adult hoverflies new####
    #1) Survival
    Ad_ill_new <- Ad_ill_prev * (1-(muB+muS) * tc * dt)        #survival of ill-fed
    Ad_well_new <- Ad_well_prev * (1-muB * tc * dt)                 #survival of well-fed
    
    #2) Reproduction of adults into juveniles (no change in adults)
    #calculated under juveniles
    
    #3) change of feeding status
    Ad_well_save = Ad_well_new #necessary to save in order to not change status multiple times
    Ad_ill_save = Ad_ill_new
    
    Ad_ill_new  <- Ad_ill_new  + (Ad_well_save * a * tc * dt) - (Ad_ill_save * b * B * tc * dt)
    Ad_well_new <- Ad_well_new + (Ad_ill_save * b * B * tc * dt) - (Ad_well_save * a * tc *dt)
    
    
    #4 Juveniles develop into ill-fed adults
    tau_rounded <- round(tau / dt) * dt
    index <- round((t - tau_rounded - 1) / dt)
    if(t >= (tau + 2)) {
      susceptible <- pmax(N_array_days[, , index] - Nr, 0)
      Ad_ill_new <- Ad_ill_new + Juv_array_days[, , index] *
        (1-muj(susceptible)) *
        e(susceptible)
    }
    
    #5) End of hibernation 
    Ad_hib_new[woody_habitat] <- Ad_hib_prev[woody_habitat]  * (1 - h0) #this could be done without woody_habitat since only in woody habitat anyway
    Ad_ill_new[woody_habitat] <- Ad_ill_new[woody_habitat] + Ad_hib_prev[woody_habitat] * h0 #adult hoverflies leaving hibernation
    
    #6) Start of hibernation 
    Ad_hib_new[woody_habitat] <- Ad_hib_new[woody_habitat] * (1-mu0*dt*tc) + (Ad_well_new[woody_habitat])*h1
    Ad_well_new[woody_habitat] <- Ad_well_new[woody_habitat] * (1 - h1) #well-fed going into hibernation around day h1=183
    
   #7) winter mortality of adults
    #Eerst was Wt1 dag 209, maar ik heb er nu de laatste dag 210 van gemaakt
    if(abs(t - Wt1) < tol) {
      Ad_hib_new <- Ad_hib_new * WP
      Ad_well_new <- 0
      Ad_ill_new <- 0
    } 
    
    ####Dispersal of adults####
    #DISPERSAL####
    # Proportion of (ill-fed) hoverflies that stay in each cell due to flowers (zB)
    zB <- B / (Bd + B)  # Using Michaelis-Menten equation #*
    # Proportion of (well-fed) hoverflies that stay in each cell due to aphids (zN)
    zN <- N_new / (Nd + N_new)  # Using Michaelis-Menten equation
    
    #APPLYING DISPERSAL TO Populations
    #update dispersal not always necessary, maybe only every day and not every dt day
    #dispersal itself should happen each time step I think 
    if (round(t) %in% Update_dispersal && abs(t - round(t)) < tol) {
      #print(paste('At time', t, "dispersal arrays are updated"))
      dispersal_array_ill <- calculate_dispersal_array(B, distance_array, B_localsearch, distance_local, distance_global)
      dispersal_array_well <- calculate_dispersal_array(N_new, distance_array, N_localsearch, distance_local, distance_global)
    }
      Ad_ill_new <- apply_dispersal(Ad_ill_new, dispersal_array_ill, zB)  #*
      Ad_well_new <- apply_dispersal(Ad_well_new, dispersal_array_well, zN)  #*
     
      #print(paste("Arrays at time", t, " and index ", index))
      N_array_days[,,index_dt] <- N_new 
      Juv_array_days[,,index_dt] <- Juv_new
      #print(Juv_array_days[1,1,index])
      index_dt = index_dt + 1
    
    #DATA STORAGE#### 
    #Collect some data#
    
    sumN1 <- sumN1 + sum(N_new[woody_habitat])
    sumJuv1 <- sumJuv1 + sum(Juv_new[woody_habitat])
    sumN2 <- sumN2 + sum(N_new[early_crop])
    sumJuv2 <- sumJuv2 + sum(Juv_new[early_crop])
    sumN3 <- sumN3 + sum(N_new[late_crop])
    sumJuv3 <- sumJuv3 + sum(Juv_new[late_crop])

    if (abs(t - current_day) < tol) {
      total_time = floor(t) + (y - 1) * num_days
      #Creating array for all cells on all days of 1 year
      #N_array_days[,,t] <- N_new 
      #Juv_array_days[,,t] <- Juv_new
      Ad_ill_array_days[,,index_t] <- Ad_ill_new
      Ad_well_array_days[,,index_t] <- Ad_well_new
      Ad_hib_array_days[,,index_t] <- Ad_hib_new
      B_array_days[,,index_t] <- B
      
      # Creating a list with population averages for all days in 1 year
      N1_mean[[total_time]] <- sumN1*dt/ncellsH1
      sumN1<-0
      N2_mean[[total_time]] <- sumN2*dt/ncellsH2
      sumN2<-0
      N3_mean[[total_time]] <- sumN3*dt/ncellsH3
      sumN3<-0
      Juv1_mean[[total_time]] <- sumJuv1*dt/ncellsH1
      sumJuv1<-0
      Juv2_mean[[total_time]] <- sumJuv2*dt/ncellsH2
      sumJuv2<-0  
      Juv3_mean[[total_time]] <- sumJuv3*dt/ncellsH3
      sumJuv3<-0

      
      Ad_ill_mean[[total_time]] <- mean(Ad_ill_new)
      Ad_well_mean[[total_time]] <- mean(Ad_well_new)
      Ad_hib_mean[[total_time]] <- mean(Ad_hib_new)
      index_t = index_t + 1
      ###Visualise data####
      #PLOTTING distributions
      if (y %in% plotyears & t %% 5 == 0 & resplot == T) {
        
        
        
        
        #Aphids
        # Convert matrix to a dataframe for ggplot
        N_df <- melt(t(N_new))
        colnames(N_df) <- c("Row", "Column", "Value")
        
        Aphid_title <- paste0("Aphid density (Year: ", y, ", Day: ", t, ")")
        
        # Plot using ggplot with a gradient color scale
        N_plot = ggplot(N_df, aes(x = Column, y = Row, fill = Value)) +
          geom_tile() +
          scale_fill_gradientn(colors = c("blue", "green", "yellow", "red"),
                               name = "Density",
                               trans = "log",  # Apply log scale transformation
                               limits = c(0.01, 1000),  # Avoid log(0), so set min to 1
                               breaks = c(0.1, 1, 10, 100, 1000),  # Log scale breaks
                               labels = scales::comma) +
          
          theme_minimal() +
          labs(title = Aphid_title, fill = "Value") +
          scale_y_reverse() +  # Flip y-axis so (1,1) is top-left
          theme(axis.text.x = element_blank(),
                axis.text.y = element_blank(),
                axis.ticks = element_blank(),
                axis.title = element_blank(),
                panel.grid = element_blank())
        
        # plot(N_plot)
        
        # Store plot in a list with dynamic naming
        N_plot_list[[paste0("plot_day", t)]] <- N_plot
        
        
        ##Well-fed
        Ad_well_df <- melt(t(Ad_well_new))
        colnames(Ad_well_df) <- c("Row", "Column", "Value")
        
        Ad_well_title <- paste0("Well-fed hoverfly density (Year: ", y, ", Day: ", t, ")")
        
        # Plot using ggplot with a gradient color scale
        Ad_well_plot = ggplot(Ad_well_df, aes(x = Column, y = Row, fill = Value)) +
          geom_tile() +
          scale_fill_gradientn(colors = c("blue", "green", "yellow", "red"),
                               name = "Density",
                               trans = "log",  # Apply log scale transformation
                               limits = c(0.001, 10),  # Avoid log(0), so set min to 1
                               breaks = c(0.001, 0.01, 0.1, 1, 3),  # Log scale breaks
                               labels = scales::comma) +
          
          theme_minimal() +
          labs(title = Ad_well_title, fill = "Value") +
          scale_y_reverse() +  # Flip y-axis so (1,1) is top-left
          theme(axis.text.x = element_blank(),
                axis.text.y = element_blank(),
                axis.ticks = element_blank(),
                axis.title = element_blank(),
                panel.grid = element_blank())
        
        # plot(Ad_well_plot)
        
        Ad_well_plot_list[[paste0("plot_day", t)]] <- Ad_well_plot
        
        ##Ill-fed
        Ad_ill_df <- melt(t(Ad_ill_new))
        colnames(Ad_ill_df) <- c("Row", "Column", "Value")
        
        Ad_ill_title <- paste0("ill-fed hoverfly density (Year: ", y, ", Day: ", t, ")")
        
        # Plot using ggplot with a gradient color scale
        Ad_ill_plot = ggplot(Ad_ill_df, aes(x = Column, y = Row, fill = Value)) +
          geom_tile() +
          scale_fill_gradientn(colors = c("blue", "green", "yellow", "red"),
                               name = "Density",
                               trans = "log",  # Apply log scale transformation
                               limits = c(0.001, 10),  # Avoid log(0), so set min to 1
                               breaks = c(0.001, 0.01, 0.1, 1, 3),  # Log scale breaks
                               labels = scales::comma) +
          
          theme_minimal() +
          labs(title = Ad_ill_title, fill = "Value") +
          scale_y_reverse() +  # Flip y-axis so (1,1) is top-left
          theme(axis.text.x = element_blank(),
                axis.text.y = element_blank(),
                axis.ticks = element_blank(),
                axis.title = element_blank(),
                panel.grid = element_blank())
        
        # plot(Ad_ill_plot)
        
        Ad_ill_plot_list[[paste0("plot_day", t)]] <- Ad_ill_plot
  
      }
      
      
      
      
      
    } #end t<-current-day loop
    
    ####End, set new values to old values for the next time step####
    #and make sure there are no zero's and negatives
    N_new <- pmax(N_new, MinVal)
    N_prev <- N_new
    Juv_new <- pmax(Juv_new, MinVal)
    Juv_prev <- Juv_new
    Ad_ill_new <- pmax(Ad_ill_new, MinVal)
    Ad_ill_prev <- Ad_ill_new
    Ad_well_new <- pmax(Ad_well_new, MinVal)
    Ad_well_prev <- Ad_well_new
    Ad_hib_new  <- pmax(Ad_hib_new, MinVal)
    Ad_hib_prev <- Ad_hib_new
    
  } # end TIME-loop
}

###Save data####
time = seq(1, num_days * year)


NewData <- data.frame(time = time, 
                      well = Ad_well_mean[time], 
                      ill = Ad_ill_mean[time], 
                      hib = Ad_hib_mean[time],
                      Aphids_N1 = N1_mean[time], 
                      Aphids_N2 = N2_mean[time],
                      Aphids_N3 = N3_mean[time],
                      Juv_N1 = Juv1_mean[time],
                      Juv_N2 = Juv2_mean[time],
                      Juv_N3 = Juv3_mean[time]
                      )
NewData$Day = (NewData$time - 1) %% num_days + 1
NewData$Year = (NewData$time -1) %/% num_days + 1


