#' Calculate Genotype Probabilities and Impute Missing Data
#' 
#' Calculates genotype probabilities and imputes missing genotype data for genetic mapping studies.
#' Supports both Association Mapping (AM) and Linkage Mapping (LM) methods with comprehensive
#' missing data handling and virtual marker creation capabilities.
#' 
#' @param GeneticMap A data frame containing genetic map information with columns:
#'   \itemize{
#'     \item \code{marker}: Marker names (character)
#'     \item \code{chr}: Chromosome numbers (numeric or factor)
#'     \item \code{pos}: Genetic positions in centimorgans (numeric)
#'   }
#' @param GenoData A matrix or data frame containing genotype data with individuals as rows 
#'   and markers as columns. Genotype codes:
#'   \itemize{
#'     \item \code{0}: Homozygous for reference allele
#'     \item \code{1}: Heterozygous
#'     \item \code{2}: Homozygous for alternative allele
#'     \item \code{9} or \code{NA}: Missing data
#'   }
#' @param method Character string specifying the mapping method:
#'   \itemize{
#'     \item \code{"AM"}: Association Mapping - combines map with genotype data
#'     \item \code{"LM"}: Linkage Mapping - performs missing data imputation
#'   }
#' @param croType Character string specifying the cross type for LM method. Required when \code{method="LM"}:
#'   \itemize{
#'     \item \code{"RIL"}: Recombinant Inbred Lines
#'     \item \code{"DH"}: Doubled Haploids
#'     \item \code{"F2"}: F2 population
#'     \item \code{"BCP1"}: Backcross to Parent 1
#'     \item \code{"BCP2"}: Backcross to Parent 2
#'     \item \code{"Fn"}: Advanced generation (n > 2)
#'   }
#' @param steps Numeric value specifying the step size (in cM) for virtual marker creation in LM.
#'   If 0, no virtual markers are created. Default is 0.
#' @param Gn Numeric value specifying the generation number for advanced populations.
#'   Must be greater than 1. Default is 2.
#' 
#' @return A data frame containing the processed genotype data with genetic map information
#'   and calculated genotype probabilities. The structure depends on the method:
#'   \itemize{
#'     \item \strong{AM method}: Original data combined with genetic map
#'     \item \strong{LM method}: Imputed genotype data with probabilities and optional virtual markers
#'   }
#' 
#' @details This function provides comprehensive genotype data processing for genetic mapping:
#' 
#' \strong{Association Mapping (AM):}
#' \itemize{
#'   \item Simply combines genetic map with genotype data
#'   \item Converts missing values to code 9
#'   \item No imputation performed
#' }
#' 
#' \strong{Linkage Mapping (LM):}
#' \itemize{
#'   \item Validates genotype codes based on cross type expectations
#'   \item Imputes missing genotypes using flanking marker information
#'   \item Creates virtual markers at specified intervals when \code{steps > 0}
#'   \item Uses Haldane mapping function for recombination calculations
#'   \item Applies cross-type specific genotype probability calculations
#' }
#' 
#' \strong{Cross Type Genotype Expectations:}
#' \itemize{
#'   \item \code{RIL/DH}: Only codes 0 and 2 allowed
#'   \item \code{F2}: Codes 0, 1, and 2 allowed
#'   \item \code{BCP1/BCP2}: Codes 0, 1, and 2 with specific constraints
#'   \item \code{Fn}: Codes 0, 1, and 2 for advanced generations
#' }
#' 
#' \strong{Virtual Marker Creation:}
#' When \code{steps > 0}, virtual markers are created at regular intervals between
#' existing markers to improve mapping resolution and handle large gaps in the genetic map.
#' 
#' @examples
#' # Example genetic map
#' genetic_map <- data.frame(
#'   marker = c("M1", "M2", "M3", "M4", "M5"),
#'   chr = c(1, 1, 1, 1, 1),
#'   pos = c(0, 10, 25, 40, 50)
#' )
#' 
#' # Example genotype data with missing values
#' geno_data <- matrix(
#'   c(2, 0, NA, 1, 2,
#'     1, 0, 2, NA, 0,
#'     2, 1, 0, 2, 1,
#'     0, NA, 1, 0, 2),
#'   nrow = 4, ncol = 5,
#'   dimnames = list(c("Ind1", "Ind2", "Ind3", "Ind4"), 
#'                   c("M1", "M2", "M3", "M4", "M5"))
#' )
#' 
#' # Association mapping (no imputation)
#' result_am <- genotype_prob(genetic_map, geno_data, method = "AM")
#' result_am
#' 
#' # Linkage mapping with virtual marker creation
#' result_lm_vm <- genotype_prob(genetic_map, geno_data, method = "LM", 
#'                            croType = "F2", steps = 5)
#' result_lm_vm
#' 
#' @seealso \code{\link{haldane_map}} for recombination fraction calculations,
#'          \code{\link{expected_genotype_dist}} for expected genotype probabilities
#' 
#' @references
#' Haldane, J.B.S. (1919). The combination of linkage values and the calculation of 
#' distances between the loci of linked factors. Journal of Genetics, 8(3), 299-309.
#' 
#' @export
genotype_prob <- function(GeneticMap,GenoData,method,croType=NULL,steps=0,Gn=2){
    if(nrow(GeneticMap) != ncol(GenoData)){
      stop("Marker No. in geneticMap and GenoData must be equal\n")
    }
    if(method=="AM"){
      GenoData[is.na(GenoData)] <- 9
      tGeneticMapt <- t(GeneticMap[,-1])
      colnames(tGeneticMapt) <- colnames(GenoData)
      RetValue <- rbind(tGeneticMapt,GenoData)
    }else if(method=="LM"){
      if(steps < 0){
        stop("Parameter steps must be postive\n")
      }
      if(croType=="RIL" | croType=="DH") {
        if(!all(is.na(GenoData[!GenoData==2 & !GenoData==0] ))){
          warning("GenoData numeric should be 0 or 2 for RIL and DH population, other numeric has been transformed to NA\n")
          GenoData[!GenoData==2 & !GenoData==0] <- NA
        }
      }
      if(croType=="F2") {
        if(!all(is.na(GenoData[!GenoData==2 & !GenoData==1 & !GenoData==0] ))){
          warning("GenoData numeric should be 0 1 2 for F2 population, other numeric has been transformed to NA\n")
          GenoData[!GenoData==2 & !GenoData==1 & !GenoData==0] <- NA
        }
      }
      if(croType=="BCP1" | croType=="BCP2") {
        if(!all(is.na(GenoData[!GenoData==2 & !GenoData==1 & !GenoData==0] ))){
          warning("GenoData numeric should be 1 2 and 1 0 for BCP1 and BCP2 population respectively, other numeric has been transformed to NA\n")
          GenoData[!GenoData==2 & !GenoData==1 & !GenoData==0] <- NA
        }
      }
      NoChr <- nlevels(as.factor(GeneticMap$chr))
      ConChr <- levels(as.factor(GeneticMap$chr))
      rownames(GeneticMap) <- GeneticMap[,1]
      #RMapList
      RMapList <- NULL
      for (i in 1:NoChr){
        tempV <- subset(GeneticMap,GeneticMap$chr==ConChr[i])
        RMapList[[i]] <- tempV$pos
        names(RMapList[[i]]) <- tempV$marker
      }
      if(steps==0){
        VMmatrix <- GenoData
        nObs <- nrow(GenoData)
        for(i in 1:NoChr){
          Chr_R <- RMapList[[i]]
          for(nS in 1:nObs){
            for(j in 1:length(Chr_R)){
              if(GenoData[nS,names(Chr_R)[j]] %in% NA){
                jL <- j
                jR <- j
                while(jR < length(Chr_R) & GenoData[nS,names(Chr_R)[jR]] %in% NA){
                  jR <- jR+1
                }
                while(jL > 1 & GenoData[nS,names(Chr_R)[jL]] %in% NA){
                  jL <- jL-1
                }
                if(jL==1 & GenoData[nS,names(Chr_R)[jL]] %in% NA){
                  xb <- as.numeric(Chr_R[jR]-Chr_R[j])
                  rb <- haldane_mapping_function(xb/100)
                  if(croType != "BCP2" & GenoData[nS,names(Chr_R)[jR]] %in% 2){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("N2",croType,Gn,rb),digits=3)+1
                  }
                  if(croType != "DH" & croType != "RIL"& GenoData[nS,names(Chr_R)[jR]] %in% 1){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("N1",croType,Gn,rb),digits=3)+1
                  }
                  if(croType != "BCP1"& GenoData[nS,names(Chr_R)[jR]] %in% 0){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("N0",croType,Gn,rb),digits=3)+1
                  }
                }
                if(jR==length(Chr_R) & GenoData[nS,names(Chr_R)[jR]] %in% NA){
                  xa <- as.numeric(Chr_R[j]-Chr_R[jL])
                  ra <- haldane_mapping_function(xa/100)
                  if (croType != "BCP2" & GenoData[nS,names(Chr_R)[jL]] %in% 2){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("2N",croType,Gn,ra),digits=3)+1
                  }
                  if (croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 1){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("1N",croType,Gn,ra),digits=3)+1
                  }
                  if (croType != "BCP1" & GenoData[nS,names(Chr_R)[jL]] %in% 0){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("0N",croType,Gn,ra),digits=3)+1
                  }
                }
                #No 1st and last missing marker
                if(!GenoData[nS,names(Chr_R)[jL]] %in% NA & !GenoData[nS,names(Chr_R)[jR]] %in% NA){
                  #else{
                  xa <- as.numeric(Chr_R[j]-Chr_R[jL])
                  xb <- as.numeric(Chr_R[jR]-Chr_R[j])
                  ra <- haldane_mapping_function(xa/100)
                  rb <- haldane_mapping_function(xb/100)
                  if (croType != "BCP2" & GenoData[nS,names(Chr_R)[jL]] %in% 2 & GenoData[nS,names(Chr_R)[jR]] %in% 2){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("22",croType,Gn,ra,rb),digits=3)+1
                  }
                  if (croType != "BCP2" & croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 2 & GenoData[nS,names(Chr_R)[jR]] %in% 1){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("21",croType,Gn,ra,rb),digits=3)+1
                  }
                  if (croType != "BCP2" & croType != "BCP1" & GenoData[nS,names(Chr_R)[jL]] %in% 2 & GenoData[nS,names(Chr_R)[jR]] %in% 0){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("20",croType,Gn,ra,rb),digits=3)+1
                  }
                  if (croType != "BCP2" & croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 1 & GenoData[nS,names(Chr_R)[jR]] %in% 2){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("12",croType,Gn,ra,rb),digits=3)+1
                  }
                  if (croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 1 & GenoData[nS,names(Chr_R)[jR]] %in% 1){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("11",croType,Gn,ra,rb),digits=3)+1
                  }
                  if (croType != "BCP1" & croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 1 & GenoData[nS,names(Chr_R)[jR]] %in% 0){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("10",croType,Gn,ra,rb),digits=3)+1
                  }
                  if (croType != "BCP2" & croType != "BCP1" & GenoData[nS,names(Chr_R)[jL]] %in% 0 & GenoData[nS,names(Chr_R)[jR]] %in% 2){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("02",croType,Gn,ra,rb),digits=3)+1
                  }
                  if (croType != "BCP1" & croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 0 & GenoData[nS,names(Chr_R)[jR]] %in% 1){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("01",croType,Gn,ra,rb),digits=3)+1
                  }
                  if (croType != "BCP1" & GenoData[nS,names(Chr_R)[jL]] %in% 0 & GenoData[nS,names(Chr_R)[jR]] %in% 0){
                    VMmatrix[nS,names(Chr_R)[j]] <- round(calculate_expected_genotype_distribution("00",croType,Gn,ra,rb),digits=3)+1
                  }
                } 
              }
            }#j
          }#nObs
        }#NoChr
        VMmatrix <- round(VMmatrix,digits = 3)
        RetValue <- rbind(t(GeneticMap[,-1]),VMmatrix)
      }else{
        #VMapList
        VMapList <- NULL
        geneticMap <- NULL
        for(i in 1:NoChr){
          #i=1
          chi <- RMapList[[i]]
          newCh <- chi[1]
          n <- 0
          for(j in 2:length(chi)){
            #j=8
            gap <- chi[j]-chi[j-1]
            cushion <- chi[j-1]
            if(gap > steps){
              nVM <- trunc(gap/steps)
              perdist <- gap/(nVM+1)
              vecVM <- round(cushion + perdist*c(1:nVM),digits=3)
              names(vecVM) <- paste("V",i,"_",c((n+j):(n+j+nVM-1)),sep="") 
              newCh <- append(newCh,vecVM)
              n <- n+nVM
            }#gap
            newCh <- append(newCh,chi[j])
          }#j
          subMap <- matrix(NA,2,length(newCh))
          subMap <- as.data.frame(subMap)
          colnames(subMap) <- names(newCh)
          #subMap <- as.data.frame(subMap)
          rownames(subMap) <- c("chr","pos")
          subMap[1,] <- ConChr[i]
          subMap[2,] <- newCh
          if(i==1){
            geneticMap <- subMap
          }else{
            geneticMap <- cbind(geneticMap,subMap)
          }
          VMapList[i] <- list(newCh)
        }#i
        #VMapLine
        NoMM <- 0      #No. of Marker on the Map
        LineDis <- 0   #Line distance
        VMapLine <- NULL     #mapline including Virtual markers
        FMCh <- NULL   #vector of First marker of each chromosome 
        EMCh <- NULL   #vector of End marker of each chromosome 
        for (i in 1:NoChr) {
          chi <- VMapList[[i]]
          StartP <- chi[1]
          for (j in 1:length(chi)) {
            NoMM <- NoMM + 1
            VMapLine <- append(VMapLine, (chi[j]+LineDis))
            names(VMapLine)[NoMM] <- names(chi)[j]
            if (j == 1) {
              FMCh <- append(FMCh,names(chi)[j])
            }else if (j == length(chi)) {
              EMCh <- append(EMCh,names(chi)[j])
              LineDis <- VMapLine[NoMM]
            }
          }
          #LineDis <- LineDis + 20
        }
        nM_V <- NoMM          #nM_V: No. of Markers
        nObs <- nrow(GenoData)
        VMmatrix <- matrix(NA,nObs,nM_V)
        colnames(VMmatrix) <- names(VMapLine)
        rownames(VMmatrix) <- rownames(GenoData)
        for(j in 1:ncol(GenoData)){
          VMmatrix[,colnames(GenoData)[j]] <- GenoData[,j]
        }
        for(i in 1:NoChr){
          #i=1
          Chr_R <- RMapList[[i]]
          Chr_V <- VMapList[[i]]
          for(jr in 2:length(Chr_R)){
            #jr=8
            nL <- which(names(Chr_V) %in% names(Chr_R)[jr-1])
            nR <- which(names(Chr_V) %in% names(Chr_R)[jr])
            for(nS in 1:nObs){
              #nS=2
              jR <- jr  #j=1
              jL <- jr-1
              while(jR < length(Chr_R) & GenoData[nS,names(Chr_R)[jR]] %in% NA){
                jR <- jR+1
              }
              while(jL > 1 & GenoData[nS,names(Chr_R)[jL]] %in% NA){
                jL <- jL-1
              }
              for(jv in nL:nR){
                #jv = 66
                if(VMmatrix[nS,names(Chr_V)[jv]] %in% NA){
                  if(jL==1 & jR ==length(Chr_R)){
                    stop("marker data without valid marker")
                  }
                  #the 1st missing marker 
                  if(jL==1 & GenoData[nS,names(Chr_R)[jL]] %in% NA){
                    xb <- as.numeric(Chr_R[jR]-Chr_V[jv])
                    rb <- haldane_mapping_function(xb/100)
                    if(croType != "BCP2" & GenoData[nS,names(Chr_R)[jR]] %in% 2){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("N2",croType,Gn,rb),digits=3)+1
                    }
                    if(croType != "DH" & croType != "RIL"& GenoData[nS,names(Chr_R)[jR]] %in% 1){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("N1",croType,Gn,rb),digits=3)+1
                    }
                    if(croType != "BCP1"& GenoData[nS,names(Chr_R)[jR]] %in% 0){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("N0",croType,Gn,rb),digits=3)+1
                    }
                  }
                  #the last missing marker
                  if(jR==length(Chr_R) & GenoData[nS,names(Chr_R)[jR]] %in% NA){
                    xa <- as.numeric(Chr_V[jv]-Chr_R[jL])
                    ra <- haldane_mapping_function(xa/100)
                    if (croType != "BCP2" & GenoData[nS,names(Chr_R)[jL]] %in% 2){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("2N",croType,Gn,ra),digits=3)+1
                    }
                    if (croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 1){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("1N",croType,Gn,ra),digits=3)+1
                    }
                    if (croType != "BCP1" & GenoData[nS,names(Chr_R)[jL]] %in% 0){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("0N",croType,Gn,ra),digits=3)+1
                    }
                  }
                  #No 1st and last missing marker
                  if(!GenoData[nS,names(Chr_R)[jL]] %in% NA & !GenoData[nS,names(Chr_R)[jR]] %in% NA){
                    #else{
                    xa <- as.numeric(Chr_V[jv]-Chr_R[jL])
                    xb <- as.numeric(Chr_R[jR]-Chr_V[jv])
                    ra <- haldane_mapping_function(xa/100)
                    rb <- haldane_mapping_function(xb/100)
                    if (croType != "BCP2" & GenoData[nS,names(Chr_R)[jL]] %in% 2 & GenoData[nS,names(Chr_R)[jR]] %in% 2){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("22",croType,Gn,ra,rb),digits=3)+1
                    }
                    if (croType != "BCP2" & croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 2 & GenoData[nS,names(Chr_R)[jR]] %in% 1){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("21",croType,Gn,ra,rb),digits=3)+1
                    }
                    if (croType != "BCP2" & croType != "BCP1" & GenoData[nS,names(Chr_R)[jL]] %in% 2 & GenoData[nS,names(Chr_R)[jR]] %in% 0){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("20",croType,Gn,ra,rb),digits=3)+1
                    }
                    if (croType != "BCP2" & croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 1 & GenoData[nS,names(Chr_R)[jR]] %in% 2){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("12",croType,Gn,ra,rb),digits=3)+1
                    }
                    if (croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 1 & GenoData[nS,names(Chr_R)[jR]] %in% 1){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("11",croType,Gn,ra,rb),digits=3)+1
                    }
                    if (croType != "BCP1" & croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 1 & GenoData[nS,names(Chr_R)[jR]] %in% 0){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("10",croType,Gn,ra,rb),digits=3)+1
                    }
                    if (croType != "BCP2" & croType != "BCP1" & GenoData[nS,names(Chr_R)[jL]] %in% 0 & GenoData[nS,names(Chr_R)[jR]] %in% 2){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("02",croType,Gn,ra,rb),digits=3)+1
                    }
                    if (croType != "BCP1" & croType != "DH" & croType != "RIL" & GenoData[nS,names(Chr_R)[jL]] %in% 0 & GenoData[nS,names(Chr_R)[jR]] %in% 1){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("01",croType,Gn,ra,rb),digits=3)+1
                    }
                    if (croType != "BCP1" & GenoData[nS,names(Chr_R)[jL]] %in% 0 & GenoData[nS,names(Chr_R)[jR]] %in% 0){
                      VMmatrix[nS,names(Chr_V)[jv]] <- round(calculate_expected_genotype_distribution("00",croType,Gn,ra,rb),digits=3)+1
                    }
                }
              }
            }#jv
          }#nObs
        }#jr
      }#NoChr
      VMmatrix <- round(VMmatrix,digits = 3)
      RetValue <- rbind(geneticMap,VMmatrix)
    }#steps
  }else{
    stop("method must be AM or LM")
  }
  return(as.data.frame(RetValue))
}
