#!/bin/bash

# Script to fix Tensor API usage across all preprocessing modules

cd "$(dirname "$0")"

# Function to fix a single file
fix_file() {
    local file="$1"
    echo "Fixing $file..."
    
    # Fix shape.len() -> shape.ndim()
    sed -i 's/shape\.len()/shape.ndim()/g' "$file"
    
    # Fix shape[n] -> dims[n] (first add dims = shape.dims())
    sed -i 's/let shape = data\.shape();/let shape = data.shape();\n        let dims = shape.dims();/g' "$file"
    sed -i 's/shape\[\([0-9]\)\]/dims[\1]/g' "$file"
    
    # Fix to_vec calls
    sed -i 's/data\.to_vec::<f64>()?/data.to_cpu()?.iter().map(|&x| x as f64).collect::<Vec<f64>>()/g' "$file"
    sed -i 's/\.to_vec::<f64>()?/.to_cpu()?.iter().map(|&x| x as f64).collect::<Vec<f64>>()/g' "$file"
    
    # Fix from_slice calls - basic pattern
    sed -i 's/Tensor::from_slice(&\([^,]*\), shape)/Ok(Tensor::from_slice(\&\1.iter().map(|\&x| x as f32).collect::<Vec<f32>>(), dims, \&data.device())?)/g' "$file"
    
    # Fix specific patterns for different variable names
    sed -i 's/Tensor::from_slice(&\([^,]*\), &\[\([^]]*\)\])/Ok(Tensor::from_slice(\&\1.iter().map(|\&x| x as f32).collect::<Vec<f32>>(), \&[\2], \&data.device())?)/g' "$file"
}

# Fix all Rust files in the preprocessing crate
find src -name "*.rs" -type f | while read -r file; do
    fix_file "$file"
done

echo "Basic fixes applied. Manual review and additional fixes may be needed."