"""
Copyright (c) 2026 AI First Engineering LLC
Author: Taner Onsay
Date: September 2026

Educational Use Notice
----------------------
Permission is granted to students and educators to use, copy, and modify this
software for non-commercial educational purposes, provided that this copyright
and attribution notice is retained. Commercial use, redistribution for profit,
publication as part of a commercial product or service, or removal of attribution
requires prior written permission from AI First Engineering LLC.

This software is provided for educational use without warranty of any kind.

Fundamentals of AI-Augmented Engineering - Lab 6
Program 2: Vibration Decay, Interpolation & Validation — Version 2

Workflow
--------
1) Import the phyphox Excel file and display the measured Z acceleration.
2) Display positive and negative experimental peak envelopes.
3) Convert the envelope to |peak acceleration| and create TWO missing-data
   interpolation regions, near 2 s and near 15 s.
4) Students enter straight-line interpolation predictions.
5) Reveal withheld experimental measurements and validate.
6) Fit A(t) = A0 exp(-lambda t), report the decay constant, measured
   damped frequency, and an estimated damping ratio.

Required packages:
    pandas, numpy, matplotlib, openpyxl
"""

from pathlib import Path
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

SHEET_NAME = "Raw Data"
TIME_COLUMN = "Time (s)"
Z_COLUMN = "Linear Acceleration z (m/s^2)"

EARLY_TARGET = 2.0
LATE_TARGET = 15.0
HIDDEN_PER_REGION = 3


def display_usage_notice():
    """Display ownership and permitted-use information at program launch."""
    print("Copyright (c) 2026 AI First Engineering LLC | Author: Taner Onsay")
    print("Educational use permitted with attribution; commercial use requires written permission.")
    print("-" * 78)


def stop(message):
    print("\nERROR:")
    print(message)
    input("\nPress ENTER to close the program...")
    sys.exit(1)


def get_file():
    print("=" * 72)
    print("ENGR 200 - LAB 6")
    print("PROGRAM 2: VIBRATION DECAY, INTERPOLATION & VALIDATION")
    print("=" * 72)
    print("\nPlace your phyphox Excel (.xlsx) file in the SAME folder")
    print("as this Python program.\n")

    name = input("Enter the phyphox Excel filename: ").strip().strip('"').strip("'")
    if not name.lower().endswith(".xlsx"):
        name += ".xlsx"

    path = Path(name)
    if not path.exists():
        stop(
            f'File "{path.name}" was not found.\n'
            "Check the spelling and make sure the Excel file is in the same folder."
        )
    return path


def load_data(path):
    try:
        df = pd.read_excel(path, sheet_name=SHEET_NAME)
    except ImportError:
        stop("Install Excel support with:\n\n    pip install openpyxl")
    except Exception as exc:
        stop(f'Could not read worksheet "{SHEET_NAME}".\nDetails: {exc}')

    missing = [c for c in (TIME_COLUMN, Z_COLUMN) if c not in df.columns]
    if missing:
        stop("Missing expected column(s): " + ", ".join(missing))

    df = df[[TIME_COLUMN, Z_COLUMN]].copy()
    for c in df.columns:
        df[c] = pd.to_numeric(df[c], errors="coerce")
    df = df.dropna().sort_values(TIME_COLUMN).reset_index(drop=True)

    if len(df) < 100:
        stop("Not enough valid measurements were found.")

    return df


def estimate_frequency(t, z):
    dt = np.median(np.diff(t))
    lo = t[0] + 0.20 * (t[-1] - t[0])
    hi = t[0] + 0.80 * (t[-1] - t[0])
    mask = (t >= lo) & (t <= hi)

    zz = z[mask] - np.mean(z[mask])
    freqs = np.fft.rfftfreq(len(zz), d=dt)
    spec = np.abs(np.fft.rfft(zz))
    valid = (freqs >= 0.5) & (freqs <= 20.0)

    if not np.any(valid):
        stop("Could not estimate the vibration frequency.")

    return freqs[valid][np.argmax(spec[valid])], dt


def spaced_local_peaks(signal, min_spacing_samples, positive=True):
    if positive:
        cand = np.where(
            (signal[1:-1] > signal[:-2]) &
            (signal[1:-1] >= signal[2:])
        )[0] + 1
        strength = signal[cand]
    else:
        cand = np.where(
            (signal[1:-1] < signal[:-2]) &
            (signal[1:-1] <= signal[2:])
        )[0] + 1
        strength = -signal[cand]

    order = cand[np.argsort(strength)[::-1]]
    kept = []
    for idx in order:
        if all(abs(idx - k) >= min_spacing_samples for k in kept):
            kept.append(idx)

    return np.array(sorted(kept), dtype=int)


def detect_signed_peaks(t, z, frequency, dt):
    period = 1.0 / frequency
    min_spacing = max(2, int(round(0.60 * period / dt)))

    pos = spaced_local_peaks(z, min_spacing, positive=True)
    neg = spaced_local_peaks(z, min_spacing, positive=False)

    # Mild noise rejection based on the strong part of the measured response.
    reference = np.percentile(np.abs(z), 95)
    threshold = max(0.05, 0.05 * reference)
    pos = pos[z[pos] >= threshold]
    neg = neg[-z[neg] >= threshold]

    if len(pos) < 8 or len(neg) < 8:
        stop("Too few positive/negative vibration peaks were detected.")

    return pos, neg


def build_absolute_peak_series(t, z, pos, neg):
    """
    Combine positive and negative extrema into one |peak| series.
    Since |sin| has a peak every half cycle, this gives a dense decay record.
    """
    idx = np.sort(np.concatenate([pos, neg]))
    return idx, t[idx], np.abs(z[idx])


def show_raw_z(t, z):
    plt.figure(figsize=(11, 6))
    plt.plot(t, z, color="goldenrod", linewidth=1.1,
             label="Linear Acceleration z")
    plt.axhline(0, linewidth=0.8)
    plt.xlabel("Time (s)")
    plt.ylabel("a (m/s²)")
    plt.title("Linear Acceleration z - Raw Experimental Data")
    plt.grid(True, alpha=0.25)
    plt.legend(loc="lower left")
    plt.tight_layout()
    plt.show(block=False)
    plt.pause(0.5)


def show_signed_envelope(t, z, pos, neg):
    plt.figure(figsize=(11, 6))
    plt.plot(t, z, color="goldenrod", linewidth=0.9,
             label="Measured Z acceleration")
    plt.plot(t[pos], z[pos], marker="o", markersize=4, linewidth=1.2,
             label="Positive peak envelope")
    plt.plot(t[neg], z[neg], marker="o", markersize=4, linewidth=1.2,
             label="Negative peak envelope")
    plt.axhline(0, linewidth=0.8)
    plt.xlabel("Time (s)")
    plt.ylabel("Z acceleration (m/s²)")
    plt.title("Experimental Vibration and Decay Envelope")
    plt.grid(True, alpha=0.25)
    plt.legend(loc="lower left")
    plt.tight_layout()
    plt.show(block=False)
    plt.pause(0.5)


def choose_region(peak_idx, t, target):
    """
    Select five consecutive |peak| measurements near target:
        anchor, hidden, hidden, hidden, anchor
    """
    block_size = HIDDEN_PER_REGION + 2
    best = None
    best_dist = np.inf

    for start in range(len(peak_idx) - block_size + 1):
        block = peak_idx[start:start + block_size]
        center = np.mean(t[block])
        dist = abs(center - target)
        if dist < best_dist:
            best_dist = dist
            best = block

    if best is None:
        stop(f"Could not create interpolation region near {target:.1f} s.")

    return best, best[1:-1]


def show_absolute_challenge(t, z, peak_idx, early_block, early_hidden,
                            late_block, late_hidden):
    hidden = set(np.concatenate([early_hidden, late_hidden]).tolist())
    visible = np.array([i for i in peak_idx if i not in hidden], dtype=int)

    plt.figure(figsize=(11, 6))
    plt.scatter(t[visible], np.abs(z[visible]), s=32,
                label="Available experimental peak measurements")

    plt.axvspan(t[early_block[0]], t[early_block[-1]], alpha=0.12,
                label="Early missing-data region")
    plt.axvspan(t[late_block[0]], t[late_block[-1]], alpha=0.12,
                label="Late missing-data region")

    plt.xlabel("Time (s)")
    plt.ylabel("|Z acceleration| peak amplitude (m/s²)")
    plt.title("Vibration Decay - Experimental Peak Data with Missing Measurements")
    plt.grid(True, alpha=0.25)
    plt.legend(loc="lower left")
    plt.tight_layout()
    plt.show(block=False)
    plt.pause(0.5)


def print_region(label, t, z, block, hidden):
    hidden_set = set(hidden.tolist())

    print("\n" + "-" * 72)
    print(label)
    print("-" * 72)
    print(f"{'Time (s)':>12}    {'Peak |Z acceleration| (m/s^2)':>30}")
    print("-" * 48)

    for idx in block:
        val = "MISSING" if idx in hidden_set else f"{abs(z[idx]):.6f}"
        print(f"{t[idx]:12.5f}    {val:>30}")

    print("\nUse the two available bounding measurements to perform")
    print("straight-line interpolation at the missing times.")


def enter_predictions(label, t, hidden):
    print(f"\nEnter your interpolated values for {label}.")
    print("Enter |Z acceleration| in m/s^2.")
    print("The experimental values remain hidden.\n")

    answers = {}
    for idx in hidden:
        while True:
            try:
                answers[idx] = float(
                    input(f"Prediction at t = {t[idx]:.5f} s: ").strip()
                )
                break
            except ValueError:
                print("Please enter a numerical value.")
    return answers


def print_validation(label, t, z, hidden, answers):
    print("\n" + "=" * 82)
    print(label.upper() + " - VALIDATION")
    print("=" * 82)
    print(
        f"{'Time (s)':>10} {'Prediction':>14} {'Measured':>14} "
        f"{'Abs. Error':>14} {'% Difference':>15}"
    )
    print("-" * 82)

    errors = []
    for idx in hidden:
        pred = answers[idx]
        actual = abs(z[idx])
        err = abs(pred - actual)
        errors.append(err)
        pct = 100 * err / actual if actual > 1e-12 else np.nan
        pct_text = f"{pct:.2f}%" if np.isfinite(pct) else "N/A"
        print(
            f"{t[idx]:10.5f} {pred:14.6f} {actual:14.6f} "
            f"{err:14.6f} {pct_text:>15}"
        )

    print("-" * 82)
    print(f"Mean absolute error: {np.mean(errors):.6f} m/s^2")


def validation_plot(filename, title, t, z, peak_idx, block, hidden, answers):
    hidden_set = set(hidden.tolist())
    peak_list = peak_idx.tolist()
    p0 = peak_list.index(block[0])
    p1 = peak_list.index(block[-1])
    lo = max(0, p0 - 2)
    hi = min(len(peak_idx), p1 + 3)
    region = peak_idx[lo:hi]
    available = np.array([i for i in region if i not in hidden_set], dtype=int)

    pred = np.array([answers[i] for i in hidden])

    plt.figure(figsize=(9.5, 5.8))
    plt.scatter(t[available], np.abs(z[available]), s=48,
                label="Available experimental peak measurements")
    plt.scatter(t[hidden], pred, s=85, marker="x", linewidths=2,
                label="Your interpolated predictions")
    plt.scatter(t[hidden], np.abs(z[hidden]), s=48, marker="o",
                color="red", zorder=5,
                label="Withheld experimental peak measurements")

    line_t = np.concatenate(([t[block[0]]], t[hidden], [t[block[-1]]]))
    line_a = np.concatenate(([abs(z[block[0]])], pred, [abs(z[block[-1]])]))
    plt.plot(line_t, line_a, linewidth=1.3, label="Your linear interpolation")

    plt.xlabel("Time (s)")
    plt.ylabel("|Z acceleration| peak amplitude (m/s²)")
    plt.title(title)
    plt.grid(True, alpha=0.25)
    plt.legend(loc="lower left")
    plt.tight_layout()
    plt.savefig(filename, dpi=200, bbox_inches="tight")
    plt.show(block=False)
    plt.pause(0.5)


def fit_decay(t, z, peak_idx):
    """
    Fit ln(A) = ln(A0) - lambda*t using the measured |peak| amplitudes.
    """
    tp = t[peak_idx]
    ap = np.abs(z[peak_idx])

    # Exclude very tiny amplitudes if present.
    threshold = max(0.05, 0.05 * np.percentile(ap, 90))
    mask = ap > threshold
    tf = tp[mask]
    af = ap[mask]

    slope, intercept = np.polyfit(tf, np.log(af), 1)
    decay = -slope
    A0 = np.exp(intercept)

    predicted = A0 * np.exp(-decay * tf)
    ss_res = np.sum((af - predicted) ** 2)
    ss_tot = np.sum((af - np.mean(af)) ** 2)
    r2 = 1 - ss_res / ss_tot if ss_tot > 0 else np.nan

    return A0, decay, r2


def show_decay_fit(t, z, peak_idx, A0, decay, frequency, r2):
    tp = t[peak_idx]
    ap = np.abs(z[peak_idx])
    curve_t = np.linspace(tp[0], tp[-1], 500)
    curve_a = A0 * np.exp(-decay * curve_t)

    # For a lightly damped oscillator:
    # wd = 2*pi*fd; lambda = zeta*wn; wd = wn*sqrt(1-zeta^2)
    # Solving gives zeta = lambda/sqrt(lambda^2 + wd^2).
    wd = 2 * np.pi * frequency
    zeta = decay / np.sqrt(decay**2 + wd**2)
    damping_percent = 100 * zeta

    plt.figure(figsize=(11, 6))
    plt.scatter(tp, ap, s=30, label="Experimental |peak| measurements")
    plt.plot(curve_t, curve_a, linewidth=2.0, color="red",
             label="Exponential decay model")
    plt.xlabel("Time (s)")
    plt.ylabel("|Z acceleration| peak amplitude (m/s²)")
    plt.title("Experimental Vibration Decay and Exponential Model")

    # Display the fitted decay equation directly on the graph.
    equation_text = (
        r"$A(t) = A_0 e^{-\lambda t}$" + "\n"
        + rf"$A(t) = {A0:.3f}e^{{-{decay:.5f}t}}$"
        + "\n"
        + rf"$\lambda = {decay:.5f}\ \mathrm{{s^{{-1}}}}$"
    )
    plt.text(
        0.52, 0.58, equation_text,
        transform=plt.gca().transAxes,
        fontsize=13,
        ha="center", va="center",
        bbox=dict(boxstyle="round,pad=0.45", facecolor="white",
                  edgecolor="gray", alpha=0.90)
    )

    plt.grid(True, alpha=0.25)
    plt.legend(loc="lower left")
    plt.tight_layout()
    plt.savefig("Lab6_Decay_Exponential_Model.png", dpi=200,
                bbox_inches="tight")
    plt.show(block=False)
    plt.pause(0.5)

    print("\n" + "=" * 72)
    print("ENGINEERING INTERPRETATION - VIBRATION DECAY")
    print("=" * 72)
    print(f"\nExponential model:  A(t) = {A0:.4f} e^(-{decay:.5f} t)")
    print(f"Decay constant, lambda:             {decay:.5f} 1/s")
    print(f"Measured damped frequency, f_d:     {frequency:.3f} Hz")
    print(f"Estimated damping ratio, zeta:      {zeta:.5f}")
    print(f"Estimated percent damping:          {damping_percent:.3f}%")
    print(f"Decay-model R^2:                    {r2:.4f}")
    print("\nThe decay constant describes how rapidly the vibration amplitude")
    print("decreases with time. A larger decay constant means faster decay.")
    print("\nThe damping ratio is a dimensionless measure of how strongly the")
    print("system dissipates vibration. The value reported here is an estimate")
    print("based on a lightly damped vibration model.")


def main():
    display_usage_notice()
    path = get_file()
    df = load_data(path)
    t = df[TIME_COLUMN].to_numpy(float)
    z = df[Z_COLUMN].to_numpy(float)

    frequency, dt = estimate_frequency(t, z)
    pos, neg = detect_signed_peaks(t, z, frequency, dt)
    peak_idx, _, _ = build_absolute_peak_series(t, z, pos, neg)

    print("\nFile loaded successfully.")
    print(f"File: {path.name}")
    print(f"Valid measurements: {len(df)}")
    print(f"Recording duration: {t[-1] - t[0]:.2f} s")

    # STEP 1
    show_raw_z(t, z)
    print("\nFigure 1 shows the Z acceleration measurements")
    print("from your phyphox experiment.")
    print("\nCompare this figure with the Z acceleration data displayed in phyphox.")
    input("\nPress ENTER on your computer to confirm.")

    # STEP 2
    show_signed_envelope(t, z, pos, neg)
    print("\nFigure 2 shows the measured vibration together with")
    print("the positive and negative experimental peak envelopes.")
    input("\nPress ENTER on your computer to continue.")

    # STEP 3
    early_block, early_hidden = choose_region(peak_idx, t, EARLY_TARGET)
    late_block, late_hidden = choose_region(peak_idx, t, LATE_TARGET)

    if set(early_block.tolist()) & set(late_block.tolist()):
        stop("The early and late interpolation regions overlap.")

    show_absolute_challenge(
        t, z, peak_idx,
        early_block, early_hidden,
        late_block, late_hidden
    )

    print_region(
        "CHALLENGE A - EARLY DECAY REGION",
        t, z, early_block, early_hidden
    )
    print_region(
        "CHALLENGE B - LATE DECAY REGION",
        t, z, late_block, late_hidden
    )

    print("\nThe graph will remain open while you perform your calculations.")
    input("\nPress ENTER on your computer to proceed to data entry...")

    early_answers = enter_predictions(
        "Challenge A - Early Decay Region", t, early_hidden
    )
    late_answers = enter_predictions(
        "Challenge B - Late Decay Region", t, late_hidden
    )

    print("\nPredictions recorded.")
    input("Press ENTER to REVEAL the withheld experimental measurements...")

    print_validation(
        "Challenge A - Early Decay Region",
        t, z, early_hidden, early_answers
    )
    print_validation(
        "Challenge B - Late Decay Region",
        t, z, late_hidden, late_answers
    )

    validation_plot(
        "Lab6_Decay_Early_Validation.png",
        "Validation - Interpolation in the Early Decay Region",
        t, z, peak_idx, early_block, early_hidden, early_answers
    )
    validation_plot(
        "Lab6_Decay_Late_Validation.png",
        "Validation - Interpolation in the Late Decay Region",
        t, z, peak_idx, late_block, late_hidden, late_answers
    )

    input("\nPress ENTER on your computer to view the vibration decay model...")

    # STEP 6
    A0, decay, r2 = fit_decay(t, z, peak_idx)
    show_decay_fit(t, z, peak_idx, A0, decay, frequency, r2)

    print("\nThree report-ready figures were saved in this folder:")
    print("  Lab6_Decay_Early_Validation.png")
    print("  Lab6_Decay_Late_Validation.png")
    print("  Lab6_Decay_Exponential_Model.png")

    input("\nPress ENTER to finish...")
    plt.close("all")


if __name__ == "__main__":
    main()