import numpy as np
import matplotlib.pyplot as plt
from mh_sampler import metropolis_hastings as mh
from corner import corner

# Prepare data
data = np.loadtxt("gForce_2026-07-09_11-28-47.csv", delimiter=",", skiprows=1)
timedata = data[:, 0]
acc_x_data = data[:, 1]
acc_y_data = data[:, 2]
acc_z_data = data[:, 3]


# Cut right time window for analysis
start_time = 7.5
end_time = 19.2
start_index = np.searchsorted(timedata, start_time)
end_index = np.searchsorted(timedata, end_time)

# Clean up the data by subtracting the mean of the z-acceleration in the selected time window
mean_z_data = np.mean(acc_z_data[start_index:end_index])
acc_z_data -= mean_z_data

# Define model
def acc_z_model(t, theta):
    A, phi0, L, g, gamma = theta
    omega = np.sqrt(g / L)
    return A * omega**2 * np.cos(2 * omega * (t - start_time) + phi0) * np.exp(-gamma * (t - start_time)) / g

# Define log prior
def log_prior(theta):
    A0, phi0, L, g, gamma = theta
    sum = 0.0
    if not (0 < A0 < 10):
        return -np.inf
    if not (0 <= phi0 <= 2 * np.pi):
        return -np.inf
    if not (0 < L < 10):
        return -np.inf
    #if not (9.7 < g < 9.9):
    #    return -np.inf
    sum -= (g - 9.81)**2 / (2 * 0.05**2) + np.log(0.05 * np.sqrt(2 * np.pi))
    if not (0 <= gamma < 2):
        return -np.inf
    return sum

# Define log likelihood
def log_like(theta):
    """Log-Likelihood for a 2D Gaussian with known covariance."""
    
    A0, phi0, L, g, gamma = theta
    
    model = acc_z_model(timedata[start_index:end_index], theta)
    residuals = acc_z_data[start_index:end_index] - model
    sigma = 0.1 # Assumed measurement error in units of g
    return -0.5 * np.sum((residuals / sigma) ** 2) - len(residuals) * np.log(sigma * np.sqrt(2 * np.pi))

def log_post(theta):
    lp = log_prior(theta)
    if not np.isfinite(lp):
        return -np.inf
    ll = log_like(theta)
    return lp + ll


n_steps = 1000000
burn_in = 10000
thin = 10

chain, logp, acc = mh(
    log_post, x0=[0.0, 0.0, 0, 9.81, 0], step_sizes=[0.1, 0.05, 0.1, 0.005, 0.05],
    n_steps=n_steps, seed=1)

samples = chain[burn_in:]
logp_samples = logp[burn_in:]

print(f"Akzeptanzrate: {acc:.2f}  (sollte ungefaehr zwischen 0.2 und 0.5 liegen)")
print(f"Sample-Mittelwert: {samples.mean(axis=0)}")
print(f"Sample-Kovarianz:\n{np.cov(samples.T)}")

fig, axes = plt.subplots(1, 3, figsize=(12, 4))

axes[0].plot(chain[::thin, 0], lw=0.5, label="$A_0 \, / \,  g$")
axes[0].plot(chain[::thin, 1], lw=0.5, label="$\phi_0 \, / \, \mathrm{rad}$")
axes[0].plot(chain[::thin, 2], lw=0.5, label="$L \, / \, \mathrm{m}$")
axes[0].plot(chain[::thin, 3], lw=0.5, label="$g \, / \, \mathrm{m/s^2}$")
axes[0].plot(chain[::thin, 4], lw=0.5, label="$\gamma \, / \, \mathrm{s^{-1}}$")
axes[0].axvline(burn_in, color="k", ls="--", label="Ende burn-in")
axes[0].set_xlabel("Schritt")
axes[0].set_title("Trace-Plot")
axes[0].legend(loc="upper right", fontsize=8)

axes[1].plot(samples[::thin, 2], samples[::thin, 3], '.', ms=1, alpha=0.3)
axes[1].set_xlabel(r"$L \,  /  \, \mathrm{m}$")
axes[1].set_ylabel(r"$g \,  / \, \mathrm{m/s^2}$")
axes[1].set_title("Samples nach dem burn-in")

lin_t_space = np.linspace(timedata[start_index], timedata[end_index], 1000)
best_params = samples[abs(logp_samples - np.max(logp_samples)) < 1e-3][0]
print("Best-Fit Parameter:\n")
print("A0:", best_params[0], "+-", np.std(samples[:, 0]))
print("phi0:", best_params[1], "+-", np.std(samples[:, 1]))
print("L:", best_params[2], "+-", np.std(samples[:, 2]))
print("g:", best_params[3], "+-", np.std(samples[:, 3]))
print("gamma:", best_params[4], "+-", np.std(samples[:, 4]))
fun_values = acc_z_model(lin_t_space, best_params)
axes[2].plot(timedata[start_index:end_index], acc_z_data[start_index:end_index], label="Daten")
axes[2].plot(lin_t_space, fun_values, label="Model")
axes[2].set_xlabel("Zeit / s")
axes[2].set_ylabel("Rel. Beschleunigung in Einheiten von $g$")
axes[2].set_title("Model vs. Daten")
axes[2].legend()

fig.tight_layout()
fig.savefig("pendel_fit.pdf", dpi=150, bbox_inches="tight")
print("\nPlot gespeichert als pendel_fit.pdf")
plt.close(fig)

labels = [R"$A_0 / g$",
          R"$\phi_0\,(\mathrm{rad})$",
          R"$L\,(\mathrm{m})$",
          R"$g\,(\mathrm{m/s^2})$",
          R"$\gamma\,(\mathrm{s^{-1}})$"]

corner(samples, labels=labels, show_titles=True, title_fmt=".3f")
plt.savefig("pendel_corner.pdf")