# نصب بسته‌ها (اگر نیاز باشد)
if(!require(evd)){install.packages("evd")}
if(!require(openxlsx)){install.packages("openxlsx")}

library(evd)
library(openxlsx)

# تنظیم seed
set.seed(123)

# پارامترهای GEV
mu <- -0.35706
sigma <- 0.99171
xi <- -0.26798

# ماتریس انتقال
transition_matrix <- rbind(
  c(0.0512820512820513, 0.205128205128205, 0.025641026, 0.717948718, 0, 0, 0),
  c(0.183333333, 0.083333333, 0.05, 0.666666667, 0.016666667, 0, 0),
  c(0, 0.031496063, 0.165354331, 0.732283465, 0.05511811, 0.007874016, 0.007874016),
  c(0.026183283, 0.043303122, 0.085599194, 0.666666667, 0.123867069, 0.051359517, 0.003021148),
  c(0, 0, 0.023952096, 0.652694611, 0.137724551, 0.179640719, 0.005988024),
  c(0, 0, 0, 0.135416667, 0.510416667, 0.135416667, 0.21875),
  c(0.368421053, 0, 0, 0.631578947, 0, 0, 0)
)

num_states <- nrow(transition_matrix)

# تابع چگالی احتمال GEV
gev_pdf <- function(x) {
  evd::dgev(x, loc = mu, scale = sigma, shape = xi)
}

# بازه‌های SPEI
spei_ranges <- list(
  c(-3, -2), c(-2, -1.5), c(-1.5, -1), c(-1, 1),
  c(1, 1.5), c(1.5, 2), c(2, 3)
)

# تابع تولید SPEI تصادفی
generate_spei_in_range <- function(state) {
  if (state < 1 || state > length(spei_ranges)) {
    stop(paste("Error: state must be between 1 and", length(spei_ranges)))
  }
  range <- spei_ranges[[state]]
  runif(1, range[1], range[2])
}

# الگوریتم متروپلیس-هستینگز (با تغییرات)
metropolis_hastings_extended <- function(n_samples, initial_state) {
  samples <- numeric(n_samples + 1)
  all_proposed_states <- numeric(n_samples)
  all_proposed_spei <- numeric(n_samples)
  accepted_states <- numeric()
  accepted_spei <- numeric()
  acceptance_indicator <- numeric(n_samples)
  random_numbers <- numeric(n_samples)
  acceptance_ratios <- numeric(n_samples) # اضافه کردن برای ذخیره نسبت پذیرش
  
  current_state <- initial_state
  samples[1] <- current_state
  current_spei <- generate_spei_in_range(current_state)
  
  for (i in 1:n_samples) {
    proposal_probabilities <- transition_matrix[current_state, ]
    proposed_state <- sample(1:num_states, 1, prob = proposal_probabilities)
    proposed_spei <- generate_spei_in_range(proposed_state)
    
    all_proposed_states[i] <- proposed_state
    all_proposed_spei[i] <- proposed_spei
    
    acceptance_ratio <- min(1, (gev_pdf(proposed_spei) * transition_matrix[proposed_state, current_state]) / (gev_pdf(current_spei) * transition_matrix[current_state, proposed_state]))
    acceptance_ratios[i] <- acceptance_ratio # ذخیره نسبت پذیرش
    
    random_number <- runif(1)
    random_numbers[i] <- random_number
    
    if (random_number < acceptance_ratio) {
      current_state <- proposed_state
      current_spei <- proposed_spei
      accepted_states <- c(accepted_states, current_state)
      accepted_spei <- c(accepted_spei, current_spei)
      acceptance_indicator[i] <- 1
    } else {
      acceptance_indicator[i] <- 0
    }
    
    samples[i + 1] <- current_state
  }
  
  list(samples = samples, all_proposed_states = all_proposed_states,
       all_proposed_spei = all_proposed_spei, accepted_states = accepted_states,
       accepted_spei = accepted_spei, acceptance_indicator = acceptance_indicator,
       random_numbers = random_numbers, acceptance_ratios = acceptance_ratios) # اضافه کردن نسبت پذیرش به خروجی
}

# اجرای الگوریتم
n_samples <- 10000
initial_state <- 1

if (initial_state < 1 || initial_state > length(spei_ranges)) {
  stop(paste("Error: initial_state must be between 1 and", length(spei_ranges)))
}

mh_results <- metropolis_hastings_extended(n_samples, initial_state)

all_proposed_states <- mh_results$all_proposed_states
all_proposed_spei <- mh_results$all_proposed_spei
accepted_states <- mh_results$accepted_states
accepted_spei <- mh_results$accepted_spei
acceptance_indicator <- mh_results$acceptance_indicator
random_numbers <- mh_results$random_numbers
acceptance_ratios <- mh_results$acceptance_ratios

# ایجاد دیتافریم‌ها
output_data_all <- data.frame(
  Iteration = 1:length(sampled_states),
  Proposed_State = all_proposed_states, # تغییر نام برای وضوح بیشتر
  Proposed_SPEI = all_proposed_spei,      # تغییر نام برای وضوح بیشتر
  Acceptance_Indicator = acceptance_indicator,
  Random_Number = random_numbers,
  Acceptance_Ratio = acceptance_ratios # اضافه کردن نسبت پذیرش
)

output_data_accepted <- data.frame(
  Accepted_State = accepted_states,
  Accepted_SPEI = accepted_spei
)

# ایجاد فایل اکسل
wb <- createWorkbook()
addWorksheet(wb, "All_Proposals")
addWorksheet(wb, "Accepted")
writeData(wb, sheet = "All_Proposals", x = output_data_all)
writeData(wb, sheet = "Accepted", x = output_data_accepted)
saveWorkbook(wb, "MH_results.xlsx", overwrite = TRUE)

# نمودارها
hist(accepted_spei, main = "Histogram of accepted_spei Values", xlab = "SPEI")
plot(1:length(accepted_states), accepted_states, type = "l", main = "Sampled States over Iterations", xlab = "Iteration", ylab = "State")

# محاسبه و نمایش نرخ پذیرش
acceptance_rate <- sum(acceptance_indicator) / length(acceptance_indicator)
cat("Acceptance Rate:", acceptance_rate, "\n")