import pandas as pd
import numpy as np
import pymc as pm
import arviz as az

# Load data (adapt the filename/columns to your actual dataset)
df = pd.read_excel("Dataset_S1_FieldData.xlsx")
df.columns = df.columns.str.strip()

# Ensure plot codes are 0,...,n_plots-1
df["Plot_idx"] = df["Plot_code"].astype("category").cat.codes
n_plots = df["Plot_idx"].nunique()

def fit_gamma_glmm(response_var, df):
    """
    Fit a Gamma GLMM with log link:
      log(mu_ij) = beta0 + beta_T*T_i + beta_D*D_i + beta_R*R_i
                   + beta_TD*T_i*D_i + beta_TR*T_i*R_i + u_j

    Priors:
      beta ~ Normal(0, 1)
      sigma_plot ~ HalfNormal(1)
      u_j ~ Normal(0, sigma_plot)
      phi ~ Gamma(2, 0.5)
    """
    y = df[response_var].values

    with pm.Model() as model:
        # Random intercept for plots
        sigma_plot = pm.HalfNormal("sigma_plot", sigma=1)
        u_plot = pm.Normal("u_plot", mu=0, sigma=sigma_plot, shape=n_plots)

        # Fixed effects (weakly informative priors)
        beta0 = pm.Normal("beta0", mu=0, sigma=1)
        beta_treat = pm.Normal("beta_treat", mu=0, sigma=1)
        beta_density = pm.Normal("beta_density", mu=0, sigma=1)
        beta_rain = pm.Normal("beta_rain", mu=0, sigma=1)
        beta_treat_density = pm.Normal("beta_treat_density", mu=0, sigma=1)
        beta_treat_rain = pm.Normal("beta_treat_rain", mu=0, sigma=1)

        # Linear predictor (log-link)
        eta = (
            beta0
            + beta_treat * df["Treatments_code"].values
            + beta_density * df["Density_Plants"].values
            + beta_rain * df["Rainfall_code"].values
            + beta_treat_density * df["Treatments_code"].values * df["Density_Plants"].values
            + beta_treat_rain * df["Treatments_code"].values * df["Rainfall_code"].values
            + u_plot[df["Plot_idx"].values]
        )
        mu = pm.math.exp(eta)

        # Gamma likelihood with shape–rate parameterisation
        # phi is shape here; rate is phi/mu so that E[Y] = mu
        phi = pm.Gamma("phi", alpha=2, beta=0.5)
        y_obs = pm.Gamma("y_obs", alpha=phi, beta=phi / mu, observed=y)

        # Sampling (as in the manuscript: 4 chains, 4000 tune, 4000 draws)
        trace = pm.sample(
            draws=4000,
            tune=4000,
            chains=4,
            target_accept=0.95,
            cores=4
        )

    print(f"\nSummary for {response_var}:")
    print(
        az.summary(
            trace,
            var_names=["beta0", "beta_treat", "beta_density",
                       "beta_rain", "beta_treat_density", "beta_treat_rain"],
            round_to=2
        )
    )

    return trace

# Fit models for all four responses
trace_runoff = fit_gamma_glmm("Runoff", df)
trace_soil   = fit_gamma_glmm("Soil_loss", df)
trace_TN     = fit_gamma_glmm("TN", df)
trace_TP     = fit_gamma_glmm("TP", df)
