import pandas as pd
import glob
import os


# ---------------------------------------------------------
# STEP 1 - Ask user for settings that apply to all springs
# ---------------------------------------------------------

threshold = float(
    input(
        "Enter peak detection threshold (mm): "
    )
)

# This controls whether the larger intermediate files are saved.
# The cleaned and analyzed DataFrames are still created in memory either way
# because they are needed to calculate the peak summary.
generate_cleaned_analyzed = input(
    "Generate cleaned and analyzed files? (y/n): "
).strip().lower()

save_intermediate_files = generate_cleaned_analyzed == "y"


# ---------------------------------------------------------
# STEP 2 - Create list of springs to process
#
# The program will process springs in this order:
#
# 1040_00, 1040_01, ..., 1040_99
# 1090_00, 1090_01, ..., 1090_99
# ---------------------------------------------------------

spring_list = []

for spring_prefix in ["1040", "1090"]:
    for spring_number in range(100):
        spring_name = f"{spring_prefix}_{spring_number:02d}"
        spring_list.append(spring_name)


# ---------------------------------------------------------
# STEP 3 - Function to process one spring
#
# This function does everything your original program did,
# but for one automatically-selected spring instead of one
# manually-entered spring.
# ---------------------------------------------------------

def process_spring(spring_name, threshold, save_intermediate_files):

    print()
    print("---------------------------------------------------------")
    print(f"Checking spring: {spring_name}")
    print("---------------------------------------------------------")

    # ---------------------------------------------------------
    # Find all matching test files for this spring
    # ---------------------------------------------------------

    file_pattern = f"Spring_{spring_name}_Test_*.csv"

    file_list = glob.glob(file_pattern)

    # Make sure files are processed in order.
    # This is important because Test 1 should come before Test 2, etc.
    file_list.sort()

    # If this spring has no matching files, skip it.
    if len(file_list) == 0:
        print(f"No files found matching: {file_pattern}")
        return

    print(f"Found {len(file_list)} file(s).")

    # ---------------------------------------------------------
    # Variables used while combining files
    # ---------------------------------------------------------

    combined_data = []

    time_offset = 0

    # ---------------------------------------------------------
    # Process each test file for this spring
    # ---------------------------------------------------------

    for test_number, file in enumerate(file_list, start=1):

        print(f"Processing: {file}")

        # Read CSV
        df = pd.read_csv(file)

        # Remove rows containing empty data
        df = df.dropna()

        # Time column name
        time_column = "Time (s)"

        # Create Test column.
        # This assigns Test 1, Test 2, Test 3, etc. based on file order.
        df["Test"] = test_number

        # Create continuous time column.
        # This makes the time from multiple files act like one continuous test.
        df["Continuous Time (s)"] = df[time_column] + time_offset

        # Determine ending time for next test.
        # The next file's time starts where this one ended.
        last_time = df[time_column].iloc[-1]

        time_offset += last_time

        # Store cleaned data in memory.
        combined_data.append(df)

    # ---------------------------------------------------------
    # Combine all tests into one DataFrame
    # ---------------------------------------------------------

    final_df = pd.concat(combined_data, ignore_index=True)

    # ---------------------------------------------------------
    # Save cleaned data only if user selected yes
    # ---------------------------------------------------------

    cleaned_output_file = f"Spring_{spring_name}_Cleaned.csv"

    if save_intermediate_files:
        final_df.to_csv(cleaned_output_file, index=False)

        print()
        print("Finished Cleaning!")
        print(f"Cleaned output saved as: {cleaned_output_file}")
    else:
        print()
        print("Finished Cleaning in memory.")
        print("Cleaned file was not saved because user selected no.")

    # ---------------------------------------------------------------------
    # ANALYZING SECTION
    # ---------------------------------------------------------------------

    # Instead of opening the cleaned file from disk, use final_df directly.
    # This allows the analysis to work even when the user chooses not to save
    # the cleaned CSV file.
    df = final_df.copy()

    # ---------------------------------------------------------
    # Define important column names
    # ---------------------------------------------------------

    extension_column = "Extension (mm)"

    # ---------------------------------------------------------
    # Create Above Threshold flag
    #
    # 1 = Extension is above threshold
    # 0 = Extension is below threshold
    # ---------------------------------------------------------

    print("Creating Above Threshold...")

    df["Above Threshold"] = (
        df[extension_column] >= threshold
    ).astype(int)

    # ---------------------------------------------------------
    # Create Previous Above Threshold column
    #
    # Shift the Above Threshold column down by one row.
    # This lets each row compare itself to the row before it.
    # ---------------------------------------------------------

    print("Creating Previous Above Threshold...")

    df["Previous Above Threshold"] = (
        df["Above Threshold"]
        .shift(1)
        .fillna(0)
        .astype(int)
    )

    # ---------------------------------------------------------
    # Create Rising Edge column
    #
    # Rising Edge occurs when:
    #
    # Previous = 0
    # Current  = 1
    #
    # This marks the start of a compression cycle.
    # ---------------------------------------------------------

    print("Creating Rising Edge...")

    df["Rising Edge"] = (
        (df["Above Threshold"] == 1)
        &
        (df["Previous Above Threshold"] == 0)
    ).astype(int)

    # ---------------------------------------------------------
    # Create Falling Edge column
    #
    # Falling Edge occurs when:
    #
    # Previous = 1
    # Current  = 0
    #
    # This marks where the signal drops back below the threshold.
    # ---------------------------------------------------------

    print("Creating Falling Edge...")

    df["Falling Edge"] = (
        (df["Above Threshold"] == 0)
        &
        (df["Previous Above Threshold"] == 1)
    ).astype(int)

    # ---------------------------------------------------------
    # Create Cycle Number
    #
    # Every Rising Edge starts a new cycle.
    #
    # Example:
    #
    # Rising Edge:
    # 0 0 1 0 0 1 0
    #
    # Cycle Number:
    # 0 0 1 1 1 2 2
    #
    # The cumulative sum increases by 1 every time a rising edge occurs.
    # ---------------------------------------------------------

    print("Creating Cycle Number...")

    df["Cycle Number"] = (
        df["Rising Edge"]
        .cumsum()
    )

    # ---------------------------------------------------------
    # Save analyzed file only if user selected yes
    # ---------------------------------------------------------

    analyzed_output_file = f"Spring_{spring_name}_Analyzed.csv"

    if save_intermediate_files:
        print("Saving Analyzed File...")

        df.to_csv(analyzed_output_file, index=False)

        print(f"Analyzed output saved as: {analyzed_output_file}")
    else:
        print("Analyzed file was not saved because user selected no.")

    # ---------------------------------------------------------
    # Create Peak Summary
    # ---------------------------------------------------------

    peak_summary = []

    print("Creating peak summary...")

    # Get all valid cycle numbers.
    # Cycle 0 is ignored because it is before the first rising edge.
    cycle_numbers = sorted(
        df[df["Cycle Number"] > 0]["Cycle Number"].unique()
    )

    for cycle in cycle_numbers:

        # Get only rows for this cycle that are above threshold.
        # This searches the active compression portion of the cycle.
        cycle_data = df[
            (df["Cycle Number"] == cycle)
            &
            (df["Above Threshold"] == 1)
        ]

        # Skip empty cycles just in case.
        if len(cycle_data) == 0:
            continue

        # Find row with maximum load inside this cycle's above-threshold area.
        peak_index = cycle_data["Load (N)"].idxmax()

        peak_row = df.loc[peak_index]

        peak_summary.append({
            "Cycle Number": cycle,
            "Test": peak_row["Test"],
            "Peak Load (N)": peak_row["Load (N)"],
            "Peak Extension (mm)": peak_row[extension_column],
            "Peak Time (s)": peak_row["Continuous Time (s)"],
            "Peak Row": peak_index
        })

    # ---------------------------------------------------------
    # Create peak summary DataFrame
    # ---------------------------------------------------------

    peak_df = pd.DataFrame(peak_summary)

    # ---------------------------------------------------------
    # Save peak summary
    #
    # This file is always saved because it is the main output.
    # ---------------------------------------------------------

    print("Saving peak summary...")

    peak_output_file = f"Spring_{spring_name}_Peaks.csv"

    peak_df.to_csv(
        peak_output_file,
        index=False
    )

    # ---------------------------------------------------------
    # Report cycles per test
    # ---------------------------------------------------------

    if len(peak_df) > 0:

        cycles_per_test = (
            peak_df.groupby("Test").size()
        )

        print("Cycles Per Test:")

        for test, count in cycles_per_test.items():
            print(f"Test {int(test)}: {count}")

    else:
        print("No peaks found for this spring.")

    # ---------------------------------------------------------
    # Report results
    # ---------------------------------------------------------

    total_cycles = int(
        df["Cycle Number"].max()
    )

    print()
    print("Analysis Complete")
    print(f"Spring: {spring_name}")
    print(f"Total Cycles Found: {total_cycles}")
    print(f"Total Peaks Found: {len(peak_summary)}")

    if save_intermediate_files:
        print(f"Cleaned File: {cleaned_output_file}")
        print(f"Analyzed File: {analyzed_output_file}")

    print(f"Peak Summary File: {peak_output_file}")


# ---------------------------------------------------------
# STEP 4 - Run the process for every spring in the list
# ---------------------------------------------------------

for spring_name in spring_list:
    process_spring(
        spring_name,
        threshold,
        save_intermediate_files
    )


# ---------------------------------------------------------
# STEP 5 - Keep command window open
# ---------------------------------------------------------

print()
print("Finished processing all springs.")

input("\nPress Enter to exit")