optimization.py

import numpy as np
import matplotlib.pyplot as plt
from dsp import AccelerationData, SegmentType, FeatureExtractor

from plot_utils import FinalPlot

def get_centroids(features):
    """Calculates centroids (mean of re_x, re_y) for each class."""
    classes = {SegmentType.FLAT: [], SegmentType.DOWN: [], SegmentType.UP: []}

    for f in features:
        if f.segment_type in classes:
            classes[f.segment_type].append([f.re_x, f.re_z])

    centroids = {}
    for cls, points in classes.items():
        if points:
            centroids[cls] = np.mean(np.array(points), axis=0) * 100
        else:
            # Handle empty class if necessary
            centroids[cls] = np.array([0.0, 0.0])

    return centroids


def calculate_separability(features):
    """
        Calculates separability considering both the distance between centroids
        and the variance of the data (distance / spread).
        """
    classes = {SegmentType.FLAT: [], SegmentType.DOWN: [], SegmentType.UP: []}

    for f in features:
        if f.segment_type in classes:
            classes[f.segment_type].append([f.re_x, f.re_z])

    means = {}
    stds = {}

    for cls, points in classes.items():
        if points:
            # Scale by 100 to match original scale
            pts = np.array(points) * 100
            means[cls] = np.mean(pts, axis=0)
            # Calculate spread: sqrt(sum of variances) as a measure of cluster radius
            stds[cls] = np.sqrt(np.sum(np.var(pts, axis=0)))
        else:
            means[cls] = np.zeros(2)
            stds[cls] = 1.0  # Avoid division by zero

    c_flat = means[SegmentType.FLAT]
    c_down = means[SegmentType.DOWN]
    c_up = means[SegmentType.UP]

    s_flat = stds[SegmentType.FLAT]
    s_down = stds[SegmentType.DOWN]
    s_up = stds[SegmentType.UP]

    epsilon = 1e-6

    sep_fd = np.linalg.norm(c_flat - c_down) / (s_flat + s_down + epsilon)
    sep_fu = np.linalg.norm(c_flat - c_up) / (s_flat + s_up + epsilon)
    sep_du = np.linalg.norm(c_down - c_up) / (s_down + s_up + epsilon)

    return sep_fd + sep_fu + sep_du

def main():
    # 1. Load Data
    filename = "data/walking_2.csv"
    print(f"Loading {filename}...")
    try:
        data = AccelerationData.from_csv(filename)
    except Exception as e:
        print(f"Error loading file: {e}")
        return

    # 2. Resample to 100Hz
    resampled_data = data.resample(100.0)

    # 3. Segment Data
    resampled_data.threshold_percentile = 80.0
    segments, segments_mask = resampled_data.segment(min_period=1.5)

    # 4. Select and Label Segments
    try:
        segments_flat = [segments[0], segments[17]]
        segments_down = [segments[1], segments[3], segments[5], segments[7]]
        segments_up = [segments[9], segments[12], segments[14], segments[16]]

        for segment in segments_flat:
            segment.type = SegmentType.FLAT
        for segment in segments_down:
            segment.type = SegmentType.DOWN
        for segment in segments_up:
            segment.type = SegmentType.UP

        # Combine valid segments for processing
        labeled_segments = segments_flat + segments_down + segments_up

    except IndexError:
        print("Error: Not enough segments found for the hardcoded indices. Check threshold/min_period.")
        return

    # 5. Optimization Loop
    window_lengths = [1.5, 2.0, 2.5, 3.0]

    bandwidths = range(4, 10)
    freq_step = 0.5
    f_starts = np.arange(1.0, 50.0, freq_step)

    print(f"\nStarting optimization for windows: {window_lengths}")

    results = []

    for W in window_lengths:
        print(f"Processing Window Length W = {W}s...")

        best_sep = -1.0
        best_band = (0.0, 0.0)
        best_features = []

        # Create windows for labeled segments
        windows = []
        for seg in labeled_segments:
            windows.extend(seg.create_windows(window_length=W))

        if not windows:
            print(f"  No windows generated for W={W}s")
            continue

        # Sweep frequency bands
        for f_low in f_starts:
            for bandwidth in bandwidths:
                f_high = f_low + bandwidth

                extractor = FeatureExtractor(windows, f_low=float(f_low), f_high=float(f_high))
                sep = calculate_separability(extractor.features)

                if sep > best_sep:
                    best_sep = sep
                    best_band = (f_low, f_high)
                    best_features = extractor.features

        results.append({
            'W': W,
            'band': best_band,
            'separability': best_sep,
            'features': best_features
        })

        print(f"  -> Best Band: {best_band[0]:.1f}-{best_band[1]:.1f} Hz, Separability: {best_sep:.4f}")

    # 6. Final Visualization and Reporting
    print("\nOptimization Results:")
    print(f"{'Window (s)':<10} | {'Freq Band (Hz)':<15} | {'Separability':<12}")
    print("-" * 45)

    for res in results:
        band_str = f"{res['band'][0]:.1f}-{res['band'][1]:.1f}"
        print(f"{res['W']:<10.1f} | {band_str:<15} | {res['separability']:<12.4f}")

        # Plot the best result for this window using FinalPlot if available
        try:
            final_plot = FinalPlot(res['features'])
            fig, ax = final_plot.plot_features(fig_size=(8, 7))

            title = f"Best Result: W={res['W']}s, Band={band_str}Hz (Sep={res['separability']:.2f})"
            ax.set_title(title)

            # Add centroids for visual verification
            centroids = get_centroids(res['features'])
            # for c_pos in centroids.values():
            #     ax.plot(c_pos[0], c_pos[1], 'kx', markersize=10, markeredgewidth=2)

            try:
                final_plot.add_confidence_ellipses(ax)
            except AttributeError:
                pass

            plt.show()
            fig.savefig(
                f"export/optimization/features_w[{float(res['W'])}]_f[{band_str}]_sep[{res['separability']:.2f}].png",
                dpi=300, bbox_inches="tight", pad_inches=0.1)
        except NameError:
            pass


if __name__ == "__main__":
    main()