mmwifi_process.py

from dataclasses import dataclass, field
from typing import Dict, List, Tuple, Optional
import shutil
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import Lasso, OrthogonalMatchingPursuit, ElasticNet
from util import get_cleaned_measurement, get_filtered_measurement_matrix, preprocess_features, scale_to_01
from omp import omp
from spgl1 import spg_bp
from spgl1.spgl1 import norm_l1nn_primal, norm_l1nn_dual, norm_l1nn_project
from clean_data import Point
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from scipy.signal import find_peaks
import os
from datetime import datetime
import pandas as pd
from pathlib import Path
import argparse
from joblib import Parallel, delayed
import traceback
import scipy
import tikzplotlib

class MultiScaleLasso:
    def __init__(self, n_scales=5, lambda_0=0.1, beta=1.0):
        self.n_scales = n_scales
        self.lambda_0 = lambda_0
        self.beta = beta
        self.coef_ = None
        
    def compute_weights(self):
        """Compute scale integration weights"""
        weights = np.exp(-self.beta * np.arange(self.n_scales))
        return weights / weights.sum()
        
    def downsample(self, X, scale):
        """Downsample signal by given scale"""
        return scipy.signal.decimate(X, scale, axis=0)
        
    def upsample(self, X, target_size):
        """Upsample signal to target size"""
        return scipy.signal.resample(X, target_size)
        
    def solve_scale(self, Y, A, scale_level):
        """Solve Lasso at specific scale"""
        lambda_l = self.lambda_0 * (2 ** scale_level)
        lasso = Lasso(alpha=lambda_l, max_iter=2000, positive=True)
        lasso.fit(Y, A)
        return lasso.coef_
        
    def fit(self, Y, A):
        """Main fitting procedure"""
        weights = self.compute_weights()
        H_combined = np.zeros_like(Y[0])
        
        for l in range(self.n_scales):
            # Downsample
            scale = 2 ** l
            Y_l = self.downsample(Y, scale)
            A_l = self.downsample(A, scale)
            
            # Solve at current scale
            H_l = self.solve_scale(Y_l, A_l, l)
            
            # Upsample and combine
            H_l_up = self.upsample(H_l, len(H_combined))
            H_combined += weights[l] * H_l_up
        
        self.coef_ = H_combined
            
        return H_combined

@dataclass
class TransceiverConfig_16:
    tx_position: Point = field(default_factory=lambda: Point(0.75, 0))
    rx_position: Point = field(default_factory=lambda: Point(-0.75, 0))

    def to_dict(self) -> Dict:
        return {
            "tx_position": self.tx_position.to_dict(),
            "rx_position": self.rx_position.to_dict()
        }
        
@dataclass
class TransceiverConfig_18:
    tx_position: Point = field(default_factory=lambda: Point(-0.75, 0))
    rx_position: Point = field(default_factory=lambda: Point(0.75, 0))

    def to_dict(self) -> Dict:
        return {
            "tx_position": self.tx_position.to_dict(),
            "rx_position": self.rx_position.to_dict()
        }
@dataclass
class MMWiFiConfig:
    x_min: float = -2.5
    x_max: float = 2.5
    y_min: float = -0.2
    y_max: float = 3.5
    debug: bool = False

class MMWiFiProcessor:
    def __init__(self, folder_path: str, config: MMWiFiConfig = None):
        self.folder_path = folder_path
        self.config = config or MMWiFiConfig()
        self.setup_initial_data()

    def setup_initial_data(self):
        """Initialize measured sector patterns and angle grid."""
        self.rads, self.measurement_matrix_ori, self.has_sectors = get_filtered_measurement_matrix()
        self.figures_dir = None
        self.config_18 = TransceiverConfig_16()
        self.TX_pos = self.config_18.tx_position.to_numpy()
        self.RX_pos = self.config_18.rx_position.to_numpy()

    def clear_figures_directory(self, confirm: bool = True) -> None:
        """
        Clear all contents of the figures directory.
        
        Args:
            confirm (bool): If True, asks for user confirmation before deleting.
                          Set to False for automatic clearing.
        """
        if not self.figures_dir:
            print("No figures directory set.")
            return

        if not os.path.exists(self.figures_dir):
            print(f"Figures directory {self.figures_dir} does not exist.")
            return

        if confirm:
            response = input(f"Are you sure you want to clear all contents in {self.figures_dir}? (y/n): ")
            if response.lower() != 'y':
                print("Operation cancelled.")
                return

        try:
            # Remove all files and subdirectories
            for item in Path(self.figures_dir).glob('*'):
                if item.is_file():
                    item.unlink()
                elif item.is_dir():
                    shutil.rmtree(item)
            print(f"Successfully cleared {self.figures_dir}")
        except Exception as e:
            print(f"Error clearing figures directory: {str(e)}")

    def process_device_data(self, file_name: str) -> List[Tuple[datetime, Dict]]:
        """Process data from a single device file"""
        real_target_vectors, timestamps, sectors, direction = get_cleaned_measurement(file_name)
        real_target_vectors = preprocess_features(real_target_vectors.T).T
        
        results = []
        for i, (real_vector, timestamp, sector, dir) in enumerate(zip(
            real_target_vectors, timestamps, sectors, direction)):
            coeffs_dict = self._process_single_measurement(real_vector, i)
            results.append((timestamp, coeffs_dict))
        
        return results

    def _process_single_measurement(self, real_vector: np.ndarray, index: int) -> Dict:
        """Process a single measurement vector"""
        vector_raw = np.power(10, real_vector/10)
        X = self.measurement_matrix_ori.T
        y = vector_raw

        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(X)

        coeffs_dict = {
            'lasso': self._get_lasso_coeffs(X_scaled, y),
            'omp': self._get_omp_coeffs(X_scaled, y),
            'enet': self._get_enet_coeffs(X_scaled, y),
            'bpdn': self._get_bpdn_coeffs(X, y),
            'relasso': self._get_re_lasso_coeffs(X_scaled, y)
        }

        if self.config.debug:
            self._plot_debug_info(coeffs_dict, real_vector, index)

        return coeffs_dict

    def _get_coeffs_methods(self) -> Dict:
        """Get all coefficient calculation methods"""
        return {
            'lasso': self._get_lasso_coeffs,
            'omp': self._get_omp_coeffs,
            'enet': self._get_enet_coeffs,
            'bpdn': self._get_bpdn_coeffs,
            'relasso': self._get_re_lasso_coeffs,
        }
    
    def clean_timestamps(self, tx_data: list, rx_data: list) -> tuple:
        all_timestamps = [ts for ts, _ in tx_data + rx_data]
        earliest_time = min(all_timestamps)
        
        cleaned_tx_data = [((ts - earliest_time).total_seconds(), coeffs) for ts, coeffs in tx_data]
        cleaned_rx_data = [((ts - earliest_time).total_seconds(), coeffs) for ts, coeffs in rx_data]
        
        tx_earliest_time = min(t for t, _ in cleaned_tx_data)
        rx_earliest_time = min(t for t, _ in cleaned_rx_data)
        
        cleaned_tx_data = [(t - tx_earliest_time, coeffs) for t, coeffs in cleaned_tx_data]
        cleaned_rx_data = [(t - rx_earliest_time, coeffs) for t, coeffs in cleaned_rx_data]
        
        return cleaned_tx_data, cleaned_rx_data
    
    def find_peak_angles(self, rads, coeffs, height=0.1, distance=10):
        peaks, _ = find_peaks(coeffs, height=height, distance=distance)
        return rads[peaks], coeffs[peaks]

    def get_intersection(self, point1, angle1, point2, angle2):
        # Unpack points
        x1, y1 = point1
        x2, y2 = point2

        # Calculate direction vectors (angles are relative to y-axis)
        dx1, dy1 = np.sin(angle1), np.cos(angle1)
        dx2, dy2 = np.sin(angle2), np.cos(angle2)

        # Calculate the intersection point
        det = dx1 * dy2 - dy1 * dx2
        if np.isclose(det, 0):
            return None  # Lines are parallel, no intersection

        t1 = ((x2 - x1) * dy2 - (y2 - y1) * dx2) / det
        intersection = np.array([x1 + t1 * dx1, y1 + t1 * dy1])

        return intersection

    def estimate_positions(self, tx_data: list, rx_data: list) -> list:
        positions = []
        for tx_timestamp, tx_coeffs in tx_data:
            closest_rx = min(rx_data, key=lambda x: abs(x[0] - tx_timestamp))
            
            tx_angles, _ = self.find_peak_angles(self.rads, tx_coeffs)
            rx_angles, _ = self.find_peak_angles(self.rads, closest_rx[1])
            
            for tx_angle in tx_angles:
                for rx_angle in rx_angles:
                    pos = self.get_intersection(self.TX_pos, tx_angle, self.RX_pos, rx_angle)
                    if pos is not None:
                        positions.append((tx_timestamp, pos))
        
        return positions

    def _get_lasso_coeffs(self, X: np.ndarray, y: np.ndarray) -> np.ndarray:
        """Calculate Lasso coefficients"""
        lasso = Lasso(alpha=0.1, max_iter=2000, positive=True)
        lasso.fit(X, y)
        return scale_to_01(np.abs(lasso.coef_))
    
    def _get_re_lasso_coeffs(self, X: np.ndarray, y: np.ndarray) -> np.ndarray:
        """Calculate Lasso coefficients"""
        lasso = MultiScaleLasso()
        lasso.fit(X, y)
        return scale_to_01(np.abs(lasso.coef_))

    def _get_omp_coeffs(self, X: np.ndarray, y: np.ndarray) -> np.ndarray:
        """Calculate OMP coefficients"""
        omp_result = omp(X, y, ncoef=5, tol=1e-5, maxit=2000)
        return scale_to_01(np.abs(omp_result.coef))

    def _get_enet_coeffs(self, X: np.ndarray, y: np.ndarray) -> np.ndarray:
        """Calculate Elastic Net coefficients"""
        enet = ElasticNet(alpha=0.1, max_iter=2000, positive=True)
        enet.fit(X, y)
        return scale_to_01(np.abs(enet.coef_))

    def _get_bpdn_coeffs(self, X: np.ndarray, y: np.ndarray) -> np.ndarray:
        """Calculate BPDN coefficients"""
        bpdn_coeffs, _, _, _ = spg_bp(X, y, iter_lim=2000, verbosity=1,
                                    project=norm_l1nn_project,
                                    primal_norm=norm_l1nn_primal,
                                    dual_norm=norm_l1nn_dual)
        return scale_to_01(np.abs(bpdn_coeffs))
    
    def _process_file_pair(self, tx_file: str, rx_file: str) -> dict:
        print(f"\nProcessing pair:\nTX: {tx_file}\nRX: {rx_file}")
        
        output_dir = os.path.dirname(tx_file)
        timestamp_str = datetime.now().strftime('%Y%m%d_%H%M%S')
        
        self.figures_dir = os.path.join(output_dir, 'figures')
        os.makedirs(self.figures_dir, exist_ok=True)

        # return self.process_with_tracking(tx_file, rx_file, timestamp_str) # need some refinement
        return self._generate_results(tx_file, rx_file, timestamp_str)

    def _generate_results(self, tx_file: str, rx_file: str, timestamp_str: str) -> dict:
        """Generate results for a pair of TX and RX files"""
        tx_results = self.process_device_data(tx_file)
        rx_results = self.process_device_data(rx_file)
        
        cleaned_tx_results, cleaned_rx_results = self.clean_timestamps(tx_results, rx_results)

        results_dir = os.path.join(os.path.dirname(tx_file), 'results')
        os.makedirs(results_dir, exist_ok=True)
        tx_file_prefix = os.path.splitext(os.path.split(tx_file)[1])[0]
        
        positions = {}
        for method in ['lasso', 'enet', 'omp', 'bpdn', 'relasso']:
            tx_data = [(ts, res[method]) for ts, res in cleaned_tx_results]
            rx_data = [(ts, res[method]) for ts, res in cleaned_rx_results]
            positions[method] = self.estimate_positions(tx_data, rx_data)
            self.save_positions_data(positions[method],
                                     os.path.join(results_dir, f'{method}_position_{tx_file_prefix}.csv'))

        return {
            'tx_file': tx_file,
            'rx_file': rx_file,
            **positions,
            'timestamp': timestamp_str
        }
    
    def _plot_debug_info(self, coeffs_dict: Dict, real_vector: np.ndarray, index: int):
        """Plot debug information for coefficients"""
        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(12, 6))
        
        # Plot Lasso and Elastic Net
        ax1.plot(np.degrees(self.rads), coeffs_dict['lasso'], label='Lasso')
        ax1.plot(np.degrees(self.rads), coeffs_dict['enet'], label='enet')
        ax1.set_xlabel('Angle (degrees)')
        ax1.set_ylabel('Coefficient')
        ax1.set_title('Target Vector Inference')
        ax1.legend()
        ax1.grid(True)

        # Plot OMP and BPDN
        ax2.plot(np.degrees(self.rads), coeffs_dict['omp'], label='OMP')
        ax2.plot(np.degrees(self.rads), coeffs_dict['bpdn'], label='bpdn')
        ax2.set_xlabel('Angle (degrees)')
        ax2.set_ylabel('Coefficient')
        ax2.legend()
        ax2.grid(True)

        # Plot raw data
        ax3.plot(real_vector, label='raw')
        ax3.legend()
        
        plt.tight_layout()
        plt.savefig(os.path.join(self.figures_dir, f"Test_{index}.jpg"))
        plt.close('all')

    def process_folder(self) -> List[Dict]:
        """Process all files in the folder"""
        tx_files, rx_files = self._get_file_pairs()
        return [self._process_file_pair(tx_file, rx_file) 
                for tx_file, rx_file in zip(tx_files, rx_files)]

    def _get_file_pairs(self) -> Tuple[List[str], List[str]]:
        """Get pairs of TX and RX files from the folder"""
        tx_files, rx_files = [], []
        
        for root, _, files in os.walk(self.folder_path):
            for file in files:
                if file.endswith('_cleaned.csv'):
                    if '192.168.3.16' in file and 'position' not in file:
                        tx_files.append(os.path.join(root, file))
                    elif '192.168.3.18' in file and 'position' not in file:
                        rx_files.append(os.path.join(root, file))
        
        return sorted(tx_files), sorted(rx_files)

    def save_positions_data(self, positions: List[Tuple], output_file: str):
        """Save position data to CSV file"""
        positions_data = [
            {'timestamp': ts, 'x': pos[0], 'y': pos[1]}
            for ts, pos in positions
        ]
        
        df = pd.DataFrame(positions_data)
        df.to_csv(output_file, index=False)

def process_single_folder(folder_path, config):
    """Process a single folder and return its results"""
    print(f"Processing folder: {folder_path}")
    try:
        processor = MMWiFiProcessor(folder_path, config)
        results = processor.process_folder()
        print(f"Successfully processed {folder_path}")
        return folder_path, results, True
    except Exception as e:
        print(f"Error processing {folder_path}: {str(e)}")
        return folder_path, None, False

def main(args):
    folders = [
        "E:/mmwave_recording_auto/Single_Square",
        "E:/mmwave_recording_auto/Single_Diamond",
        "E:/mmwave_recording_auto/Single_Hourglass",
        "E:/mmwave_recording_auto/Single_Large_Hourglass",
        "E:/mmwave_recording_auto/Single_Large_Rectangle",
        "E:/mmwave_recording_auto/Single_Z_Shape"
        # "E:/mmwave_recording_auto/Single_Diamond_dis",
    ] if args.folders is None else args.folders

    config = MMWiFiConfig(debug=False)
    start_time = datetime.now()
    print(f"Starting processing of {len(folders)} folders...")
    
    try:
        if args.n_jobs == 1:
            results = [process_single_folder(folder, config) for folder in folders]
        else:
            results = Parallel(n_jobs=args.n_jobs, verbose=10)(
                delayed(process_single_folder)(folder, config)
                for folder in folders
            )
        
        # Separate successful and failed results
        successful_results = {}
        failed_folders = []
        
        for folder_path, result, success in results:
            if success:
                successful_results[folder_path] = result
            else:
                failed_folders.append(folder_path)
        
        # Print detailed results
        print("\nDetailed Results:")
        for folder_path, folder_results in successful_results.items():
            print(f"\nFolder: {folder_path}")
            print(f"Processed {len(folder_results)} pairs of files:")
            for result in folder_results:
                print(f"\nTX: {os.path.basename(result['tx_file'])}")
                print(f"RX: {os.path.basename(result['rx_file'])}")
                print(f"Output directory: {os.path.dirname(result['tx_file'])}")
        
        # Print summary
        duration = datetime.now() - start_time
        print("\nProcessing Summary:")
        print(f"Successfully processed: {len(successful_results)} folders")
        print(f"Failed: {len(failed_folders)} folders")
        if failed_folders:
            print("Failed folders:")
            for folder in failed_folders:
                print(f"  - {folder}")
        print(f"Total processing time: {duration}")
        
    except Exception as e:
        print(f"An error occurred: {str(e)}")
        traceback.print_exc()
    
    print("Processing completed!")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        prog='mmwifi_process',
        description='mmwifi parallel processing'
    )
    
    parser.add_argument(
        '-f', '--folders',
        nargs='+',
        help='List of folder paths to process'
    )
    
    parser.add_argument(
        '-j', '--n_jobs',
        type=int,
        default=12,
        help='Number of parallel jobs (-1 for all cores)'
    )

    args = parser.parse_args()
    main(args)