"""
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 1: Oscillatory Interpolation & Validation
--------------------------------------------------
Reads a phyphox "Acceleration (without g)" Excel export, selects a clean
3-4 cycle window near the middle of the recording, and creates two
interpolation challenges:
    1) missing data around a vibration peak
    2) missing data around a zero crossing

Students see the gaps first, calculate estimates, enter their predictions,
and only then are the withheld experimental measurements revealed.

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


# ----------------------------- SETTINGS ---------------------------------
SHEET_NAME = "Raw Data"
TIME_COLUMN = "Time (s)"
X_COLUMN = "Linear Acceleration x (m/s^2)"
Y_COLUMN = "Linear Acceleration y (m/s^2)"
Z_COLUMN = "Linear Acceleration z (m/s^2)"

# Number of consecutive measurements hidden in each challenge.
# Keep this odd so the target measurement is centered in the gap.
N_HIDDEN = 5

# Number of vibration cycles displayed in the challenge graph.
N_CYCLES_TO_SHOW = 4

# Search for the challenge near the middle of the recording.
# This makes the program adapt to recordings that are not exactly 20 s.
# ------------------------------------------------------------------------


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_with_message(message):
    print("\nERROR:")
    print(message)
    input("\nPress Enter to close the program...")
    sys.exit(1)


def find_excel_file():
    print("=" * 66)
    print("ENGR 200 - LAB 6")
    print("PROGRAM 1: OSCILLATORY INTERPOLATION & VALIDATION")
    print("=" * 66)
    print("\nPlace your phyphox Excel (.xlsx) file in the SAME folder")
    print("as this Python program.\n")

    filename = input("Enter the phyphox Excel filename: ").strip().strip('"').strip("'")

    if not filename.lower().endswith(".xlsx"):
        filename += ".xlsx"

    path = Path(filename)

    if not path.exists():
        stop_with_message(
            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_phyphox(path):
    try:
        data = pd.read_excel(path, sheet_name=SHEET_NAME)
    except ImportError:
        stop_with_message(
            "Python needs the openpyxl package to read Excel files.\n"
            "Install it with:\n\n    pip install openpyxl"
        )
    except Exception as exc:
        stop_with_message(
            f'Could not read worksheet "{SHEET_NAME}".\n'
            f"Details: {exc}"
        )

    missing = [c for c in (TIME_COLUMN, X_COLUMN, Y_COLUMN, Z_COLUMN) if c not in data.columns]
    if missing:
        stop_with_message(
            "This does not appear to be the expected phyphox export.\n"
            "Missing column(s): " + ", ".join(missing)
        )

    clean = data[[TIME_COLUMN, X_COLUMN, Y_COLUMN, Z_COLUMN]].copy()
    clean[TIME_COLUMN] = pd.to_numeric(clean[TIME_COLUMN], errors="coerce")
    clean[X_COLUMN] = pd.to_numeric(clean[X_COLUMN], errors="coerce")
    clean[Y_COLUMN] = pd.to_numeric(clean[Y_COLUMN], errors="coerce")
    clean[Z_COLUMN] = pd.to_numeric(clean[Z_COLUMN], errors="coerce")
    clean = clean.dropna().sort_values(TIME_COLUMN).reset_index(drop=True)

    if len(clean) < 100:
        stop_with_message("Not enough valid acceleration measurements were found.")

    return clean



def show_raw_acceleration_data(t, x, y, z):
    """Display the imported X, Y, and Z data in a phyphox-like stacked view."""
    fig, axes = plt.subplots(3, 1, figsize=(11, 8), sharex=True)
    fig.suptitle("Raw Experimental Acceleration Data", fontsize=16)

    axes[0].plot(t, x, color="limegreen", linewidth=1.0)
    axes[0].set_title("Linear Acceleration x")
    axes[0].set_ylabel("a (m/s²)")
    axes[0].grid(True, alpha=0.25)

    axes[1].plot(t, y, color="dodgerblue", linewidth=1.0)
    axes[1].set_title("Linear Acceleration y")
    axes[1].set_ylabel("a (m/s²)")
    axes[1].grid(True, alpha=0.25)

    axes[2].plot(t, z, color="goldenrod", linewidth=1.0)
    axes[2].set_title("Linear Acceleration z")
    axes[2].set_ylabel("a (m/s²)")
    axes[2].set_xlabel("Time (s)")
    axes[2].grid(True, alpha=0.25)

    plt.tight_layout(rect=[0, 0, 1, 0.96])
    plt.show(block=False)
    plt.pause(0.5)


def estimate_frequency(t, z):
    """Estimate dominant vibration frequency from a central portion of the data."""
    dt = np.median(np.diff(t))
    if not np.isfinite(dt) or dt <= 0:
        stop_with_message("The time data are not valid.")

    # Use the middle 60% to reduce influence of the beginning/end.
    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])
    n = len(zz)

    spectrum = np.abs(np.fft.rfft(zz))
    frequencies = np.fft.rfftfreq(n, d=dt)

    # Cantilever phone experiment: ignore very slow drift and implausibly
    # high-frequency sensor noise.
    valid = (frequencies >= 0.5) & (frequencies <= 20.0)
    if not np.any(valid):
        stop_with_message("Could not determine the vibration frequency.")

    f = frequencies[valid][np.argmax(spectrum[valid])]
    return f, dt


def local_peak_indices(z):
    """Return indices of simple local extrema (positive peaks and negative valleys)."""
    positive = np.where((z[1:-1] > z[:-2]) & (z[1:-1] >= z[2:]))[0] + 1
    negative = np.where((z[1:-1] < z[:-2]) & (z[1:-1] <= z[2:]))[0] + 1
    return positive, negative


def choose_peak(t, z, center_time, period):
    """
    Choose a strong local extremum near the center of the recording.
    Using absolute amplitude allows either a positive peak or negative valley.
    """
    pos, neg = local_peak_indices(z)
    candidates = np.concatenate([pos, neg])

    radius = 0.55 * period
    nearby = candidates[np.abs(t[candidates] - center_time) <= radius]

    if len(nearby) == 0:
        # Fall back to strongest sample in a one-period neighborhood.
        region = np.where(np.abs(t - center_time) <= period)[0]
        if len(region) == 0:
            stop_with_message("Could not locate a vibration peak.")
        return region[np.argmax(np.abs(z[region]))]

    # Prefer strong extrema while staying reasonably close to center.
    amplitude = np.abs(z[nearby])
    distance_penalty = 0.15 * np.max(amplitude) * (
        np.abs(t[nearby] - center_time) / radius
    )
    score = amplitude - distance_penalty
    return nearby[np.argmax(score)]


def choose_zero_crossing(t, z, peak_idx, period):
    """
    Find a zero crossing roughly a quarter-cycle away from the selected peak.
    This keeps the peak and zero-crossing challenges separate.
    """
    sign_change = np.where(z[:-1] * z[1:] <= 0)[0]
    crossing_times = (t[sign_change] + t[sign_change + 1]) / 2

    # Prefer a crossing about 1/4 period after the chosen peak.
    target = t[peak_idx] + 0.25 * period
    valid = np.abs(crossing_times - target) <= 0.30 * period

    if np.any(valid):
        cand = sign_change[valid]
        times = crossing_times[valid]
        j = np.argmin(np.abs(times - target))
        left = cand[j]
    else:
        # Try before the peak.
        target = t[peak_idx] - 0.25 * period
        j = np.argmin(np.abs(crossing_times - target))
        left = sign_change[j]

    # Select whichever sample is closest to zero.
    return left if abs(z[left]) <= abs(z[left + 1]) else left + 1


def hidden_indices(center_idx, n_points, n_total):
    half = n_points // 2
    start = center_idx - half
    end = center_idx + half

    if start < 2 or end >= n_total - 2:
        stop_with_message("The selected gap is too close to the end of the dataset.")

    return np.arange(start, end + 1)


def interpolation_inputs(t, z, hidden):
    """
    Use the immediately adjacent measured values as the two endpoints for
    straight-line interpolation across the missing interval.
    """
    left = hidden[0] - 1
    right = hidden[-1] + 1

    t1, a1 = t[left], z[left]
    t2, a2 = t[right], z[right]

    predicted = a1 + (z * 0)  # placeholder array shape
    pred_hidden = a1 + (a2 - a1) * (t[hidden] - t1) / (t2 - t1)

    return left, right, pred_hidden


def show_challenge_plot(t, z, window_mask, peak_hidden, zero_hidden):
    visible = window_mask.copy()
    visible[peak_hidden] = False
    visible[zero_hidden] = False

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

    # Mark the gap locations without revealing values.
    plt.axvspan(t[peak_hidden[0]], t[peak_hidden[-1]], alpha=0.12,
                label="Peak-region missing data")
    plt.axvspan(t[zero_hidden[0]], t[zero_hidden[-1]], alpha=0.12,
                label="Zero-crossing missing data")

    plt.axhline(0, linewidth=0.8)
    plt.xlabel("Time (s)")
    plt.ylabel("Z acceleration (m/s²)")
    plt.title("Interpolation Challenge - Experimental 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_challenge_table(name, t, z, hidden):
    left = hidden[0] - 1
    right = hidden[-1] + 1

    print("\n" + "-" * 66)
    print(name)
    print("-" * 66)
    print("The measurements marked MISSING were recorded by phyphox,")
    print("but they are being withheld until you submit your prediction.\n")
    print(f"{'Time (s)':>12}    {'Z acceleration (m/s^2)':>24}")
    print("-" * 42)

    print(f"{t[left]:12.5f}    {z[left]:24.6f}")
    for idx in hidden:
        print(f"{t[idx]:12.5f}    {'MISSING':>24}")
    print(f"{t[right]:12.5f}    {z[right]:24.6f}")

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


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

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


def validation_plot(title, filename, t, z, hidden, student):
    left = hidden[0] - 3
    right = hidden[-1] + 3
    left = max(0, left)
    right = min(len(t) - 1, right)
    region = np.arange(left, right + 1)

    available = np.array([i for i in region if i not in set(hidden)])

    plt.figure(figsize=(9, 5.5))
    plt.scatter(t[available], z[available], s=45,
                label="Available experimental measurements")
    plt.scatter(t[hidden], student, s=85, marker="x",
                linewidths=2.0, label="Your interpolated predictions")
    plt.scatter(t[hidden], z[hidden], s=45, marker="o",
                color="red", zorder=5,
                label="Withheld experimental measurements")

    # Connect the student prediction through the two interpolation endpoints.
    interp_t = np.concatenate(([t[hidden[0] - 1]], t[hidden], [t[hidden[-1] + 1]]))
    interp_a = np.concatenate(([z[hidden[0] - 1]], student, [z[hidden[-1] + 1]]))
    plt.plot(interp_t, interp_a, linewidth=1.2, label="Your interpolation")

    plt.axhline(0, linewidth=0.8)
    plt.xlabel("Time (s)")
    plt.ylabel("Z acceleration (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 print_validation(name, t, z, hidden, student):
    print("\n" + "=" * 76)
    print(f"{name.upper()} - VALIDATION")
    print("=" * 76)
    print(
        f"{'Time (s)':>10} {'Prediction':>14} {'Measured':>14} "
        f"{'Abs. Error':>14} {'% Difference':>15}"
    )
    print("-" * 76)

    for idx, prediction in zip(hidden, student):
        measured = z[idx]
        error = abs(prediction - measured)
        if abs(measured) > 1e-12:
            pct = 100.0 * error / abs(measured)
            pct_text = f"{pct:.2f}%"
        else:
            pct_text = "N/A"

        print(
            f"{t[idx]:10.5f} {prediction:14.6f} {measured:14.6f} "
            f"{error:14.6f} {pct_text:>15}"
        )

    mae = np.mean(np.abs(student - z[hidden]))
    print("-" * 76)
    print(f"Mean absolute error for this gap: {mae:.6f} m/s^2")


def main():
    display_usage_notice()
    path = find_excel_file()
    data = load_phyphox(path)

    t = data[TIME_COLUMN].to_numpy(dtype=float)
    x = data[X_COLUMN].to_numpy(dtype=float)
    y = data[Y_COLUMN].to_numpy(dtype=float)
    z = data[Z_COLUMN].to_numpy(dtype=float)

    duration = t[-1] - t[0]

    print("\nFile loaded successfully.")
    print(f"File: {path.name}")
    print(f"Valid measurements: {len(data)}")
    print(f"Recording duration: {duration:.2f} s")

    # First show the imported data in the same X-Y-Z order used by phyphox.
    show_raw_acceleration_data(t, x, y, z)

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

    frequency, dt = estimate_frequency(t, z)
    period = 1.0 / frequency

    print("\nThe program will now create two interpolation challenges")
    print("from a clean region near the middle of your recording.")

    center_time = 0.5 * (t[0] + t[-1])

    # Build a four-cycle display window centered near the middle of the run.
    half_window = 0.5 * N_CYCLES_TO_SHOW * period
    window_start = center_time - half_window
    window_end = center_time + half_window
    window_mask = (t >= window_start) & (t <= window_end)

    # CHALLENGE A:
    # Select the SECOND positive vibration peak visible from left to right
    # in the displayed window.
    positive_peaks, _ = local_peak_indices(z)
    display_peaks = positive_peaks[
        (t[positive_peaks] >= window_start)
        & (t[positive_peaks] <= window_end)
    ]

    # Reject tiny local wiggles by keeping peaks with substantial amplitude.
    if len(display_peaks) > 0:
        threshold = 0.50 * np.max(z[display_peaks])
        display_peaks = display_peaks[z[display_peaks] >= threshold]

    if len(display_peaks) < 2:
        stop_with_message(
            "Could not identify two clear positive vibration peaks "
            "in the displayed region."
        )

    peak_idx = display_peaks[1]

    # CHALLENGE B:
    # Count zero crossings from left to right in the SAME displayed window
    # and select the SIXTH crossing.
    sign_change = np.where(z[:-1] * z[1:] <= 0)[0]
    crossing_times = (t[sign_change] + t[sign_change + 1]) / 2.0
    display_crossings = sign_change[
        (crossing_times >= window_start)
        & (crossing_times <= window_end)
    ]

    if len(display_crossings) < 6:
        stop_with_message(
            "Could not identify six zero crossings in the displayed region."
        )

    crossing_left = display_crossings[5]
    zero_idx = (
        crossing_left
        if abs(z[crossing_left]) <= abs(z[crossing_left + 1])
        else crossing_left + 1
    )

    peak_hidden = hidden_indices(peak_idx, N_HIDDEN, len(t))
    zero_hidden = hidden_indices(zero_idx, N_HIDDEN, len(t))

    # Make sure the gaps do not overlap.
    if len(set(peak_hidden).intersection(set(zero_hidden))) > 0:
        stop_with_message(
            "The selected peak and zero-crossing gaps overlap.\n"
            "Please show this message to your instructor."
        )

    show_challenge_plot(t, z, window_mask, peak_hidden, zero_hidden)

    print_challenge_table(
        "CHALLENGE A - MISSING DATA NEAR A VIBRATION PEAK",
        t, z, peak_hidden
    )
    print_challenge_table(
        "CHALLENGE B - MISSING DATA NEAR A ZERO CROSSING",
        t, z, zero_hidden
    )

    print("\nThe challenge graph will remain open while you calculate.")
    print("Do NOT close the graph window yet.")

    input("\nWhen you have completed your interpolation calculations, "
          "press ENTER on your computer to proceed to data entry...")

    peak_student = get_student_predictions(
        "Challenge A - Peak Region", t, peak_hidden
    )
    zero_student = get_student_predictions(
        "Challenge B - Zero-Crossing Region", t, zero_hidden
    )

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

    print_validation(
        "Challenge A - Peak Region",
        t, z, peak_hidden, peak_student
    )
    print_validation(
        "Challenge B - Zero-Crossing Region",
        t, z, zero_hidden, zero_student
    )

    validation_plot(
        "Validation - Interpolation Near a Vibration Peak",
        "Lab6_Peak_Interpolation_Validation.png",
        t, z, peak_hidden, peak_student
    )
    validation_plot(
        "Validation - Interpolation Near a Zero Crossing",
        "Lab6_Zero_Crossing_Validation.png",
        t, z, zero_hidden, zero_student
    )

    print("\n" + "=" * 66)
    print("VALIDATION COMPLETE")
    print("=" * 66)
    print("\nTwo report-ready figures were saved in this folder:")
    print("  Lab6_Peak_Interpolation_Validation.png")
    print("  Lab6_Zero_Crossing_Validation.png")
    print("\nCompare the two validation results.")
    print("Consider how the local shape/curvature of the experimental")
    print("signal affected the performance of straight-line interpolation.")

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


if __name__ == "__main__":
    main()
