#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Dispersion Curve Inversion Manager

Backend for multi-mode surface wave inversion using linearized approach.
Supports single position, batch processing, and parallel execution.
"""

import numpy as np
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, as_completed

try:
    import pygimli as pg
    from disba import PhaseDispersion
    DISBA_AVAILABLE = True
except ImportError as e:
    print(f"Warning: Could not import required modules: {e}")
    print("Make sure pygimli and disba are properly installed")
    print("Install disba with: pip install disba")
    PhaseDispersion = None
    DISBA_AVAILABLE = False


def run_batch_inversion(positions_data, params, progress_callback=None, 
                        parallel=True, n_workers=None):
    """
    Run inversion for multiple positions with optional parallelization
    
    Parameters:
    -----------
    positions_data : dict
        {xmid_1: {modes: {...}, selected_for_inversion: True}, ...}
    params : dict
        Inversion parameters (shared across all positions)
    progress_callback : callable, optional
        Function(xmid, idx, total, status, result) called after each position
    parallel : bool
        If True, use multiprocessing; if False, sequential
    n_workers : int, optional
        Number of parallel workers (default: CPU count - 1)
    
    Returns:
    --------
    results : dict
        {xmid_1: inversion_result_1, xmid_2: inversion_result_2, ...}
    """
    
    # Filter selected positions
    selected_positions = {
        xmid: pos_data 
        for xmid, pos_data in positions_data.items()
        if pos_data.get('selected_for_inversion', True)
    }
    
    if not selected_positions:
        raise ValueError("No positions selected for inversion")
    
    n_positions = len(selected_positions)
    results = {}
    
    # Sequential mode (easier for debugging)
    if not parallel or n_positions == 1:
        for idx, (xmid, pos_data) in enumerate(selected_positions.items()):
            if progress_callback:
                progress_callback(xmid, idx, n_positions, 'running', None)
            
            try:
                result = run_multimode_sw_inversion(pos_data, params)
                result['xmid'] = xmid
                result['success'] = True
                results[xmid] = result
                
                if progress_callback:
                    progress_callback(xmid, idx + 1, n_positions, 'completed', result)
            except Exception as e:
                import traceback
                results[xmid] = {
                    'xmid': xmid,
                    'error': str(e), 
                    'traceback': traceback.format_exc(),
                    'success': False
                }
                
                if progress_callback:
                    progress_callback(xmid, idx + 1, n_positions, 'failed', str(e))
    
    # Parallel mode using ProcessPoolExecutor
    else:
        if n_workers is None:
            n_workers = max(1, mp.cpu_count() - 1)  # Leave one core free
        
        # Prepare work items
        work_items = list(selected_positions.items())
        
        with ProcessPoolExecutor(max_workers=n_workers) as executor:
            # Submit all jobs
            future_to_xmid = {
                executor.submit(_inversion_worker, xmid, pos_data, params): xmid
                for xmid, pos_data in work_items
            }
            
            # Collect results as they complete
            completed = 0
            for future in as_completed(future_to_xmid):
                xmid = future_to_xmid[future]
                
                try:
                    result = future.result()
                    results[xmid] = result
                    
                    if progress_callback:
                        progress_callback(xmid, completed + 1, n_positions, 'completed', result)
                except Exception as e:
                    import traceback
                    results[xmid] = {
                        'xmid': xmid,
                        'error': str(e),
                        'traceback': traceback.format_exc(),
                        'success': False
                    }
                    
                    if progress_callback:
                        progress_callback(xmid, completed + 1, n_positions, 'failed', str(e))
                
                completed += 1
    
    return results


def _inversion_worker(xmid, pos_data, params):
    """
    Worker function for parallel processing (must be picklable)
    """
    try:
        result = run_multimode_sw_inversion(pos_data, params)
        result['xmid'] = xmid
        result['success'] = True
        return result
    except Exception as e:
        import traceback
        error_msg = str(e)
        tb = traceback.format_exc()
        print(f"Inversion failed for xmid={xmid}: {error_msg}")
        print(f"Traceback:\n{tb}")
        return {
            'xmid': xmid,
            'success': False,
            'error': error_msg,
            'traceback': tb
        }


def run_multimode_sw_inversion(position_data, params):
    """
    Execute joint multi-mode surface wave inversion for single position
    
    Inverts a SINGLE Vs profile using data from MULTIPLE modes simultaneously.
    
    Parameters:
    -----------
    position_data : dict
        {
            'xmid': float,
            'modes': {
                'M0': {'resampled': {'frequencies': array, 'velocities': array, 'errors': array},
                       'selected_for_inversion': True},
                'M1': {...}
            }
        }
    params : dict
        {
            'n_layers': int,
            'vs_min': float, 'vs_max': float,
            'nu': float (Poisson's ratio),
            'rho': float (density),
            'lambda': float,
            'max_iter': int,
            ...
        }
    
    Returns:
    --------
    result : dict
        {
            'vs_model': array,
            'thickness': array,
            'vp_model': array,
            'density': array,
            'mode_fits': {mode_name: {observed, calculated, residuals, rms, chi2}},
            'total_chi2': float,
            'total_rms': float,
            'rrms': float,
            'iterations': int,
            'convergence': bool,
            'lambda_final': float
        }
    """
    
    # Check if disba is available
    if not DISBA_AVAILABLE:
        raise ImportError("disba library is required for dispersion inversion. "
                        "Please install it with: pip install disba")
    
    # 1. Collect selected modes
    selected_modes = []
    for mode_name, mode_data in position_data.get('modes', {}).items():
        if mode_data.get('selected_for_inversion', True):
            # Use resampled data if available, otherwise original
            data_source = mode_data.get('resampled', mode_data.get('original', None))
            if data_source and len(data_source.get('frequencies', [])) > 0:
                selected_modes.append((mode_name, data_source))
    
    if len(selected_modes) == 0:
        raise ValueError("No modes selected for inversion")
    
    # Sort modes by mode number (M0, M1, M2, ...) to ensure fundamental mode is first
    selected_modes = sorted(selected_modes, key=lambda x: int(x[0].replace('M', '')))
    
    print(f"  Starting inversion with {len(selected_modes)} mode(s): {[m[0] for m in selected_modes]}")
    
    # 2. Concatenate data from all modes
    all_frequencies = []
    all_velocities = []
    all_errors = []
    mode_indices = []  # Track which data belongs to which mode
    mode_start_end = {}  # Track start/end indices for each mode
    
    current_idx = 0
    for mode_name, data in selected_modes:
        n_points = len(data['frequencies'])
        
        all_frequencies.extend(data['frequencies'])
        all_velocities.extend(data['velocities'])
        
        # Handle errors - do NOT generate defaults, user should provide errors
        if 'errors' in data and data['errors'] is not None and len(data['errors']) > 0:
            all_errors.extend(data['errors'])
        else:
            # No errors - warning (this should not happen)
            print(f"Warning: No errors provided for mode {mode_name}")
            print(f"  Errors should be present from main Pyckster (Lorentzian) or PVC file column.")
            # Use zeros to maintain array structure
            all_errors.extend([0.0] * len(data['velocities']))
        
        mode_indices.extend([mode_name] * n_points)
        mode_start_end[mode_name] = (current_idx, current_idx + n_points)
        
        current_idx += n_points
    
    all_frequencies = np.array(all_frequencies)
    all_velocities = np.array(all_velocities)
    all_errors = np.array(all_errors)
    
    # 3. Create start model
    # Use only first available mode for starting model
    first_mode_name, first_mode_data = selected_modes[0]
    first_mode_freqs = np.array(first_mode_data['frequencies'])
    first_mode_vels = np.array(first_mode_data['velocities'])
    
    print(f"  Using {first_mode_name} to build starting model")
    
    # Adjust parameters if higher modes are present
    has_higher_modes = any(mode_name != 'M0' for mode_name, _ in selected_modes)
    if has_higher_modes and 'n_layers' not in params:
        # Use more layers for higher modes (they need more complex structures)
        params['n_layers'] = max(12, len(selected_modes) * 6)
        print(f"  Adjusted n_layers to {params['n_layers']} for multi-mode inversion")
    if has_higher_modes and 'max_depth' not in params:
        # Use greater depth for higher modes
        max_wavelength = np.max(all_velocities / all_frequencies)
        params['max_depth'] = max(50, max_wavelength * 1.5)
        print(f"  Adjusted max_depth to {params['max_depth']:.1f}m based on wavelengths")
    
    vs_start, thickness = create_start_model(params, first_mode_vels, first_mode_freqs)
    
    # 4. Setup forward operator for multi-mode
    fop = MultiModeDispersionModelling(
        frequencies=all_frequencies,
        mode_indices=mode_indices,
        thickness=thickness,
        nu=params.get('nu', 0.25),
        rho=params.get('rho', 2000),
        algorithm=params.get('algorithm', 'dunkin')
    )
    
    # Set model transformation and constraints
    vs_min_constraint = params.get('vs_min', np.min(all_velocities) * 0.3)
    vs_max_constraint = params.get('vs_max', np.max(all_velocities) * 2.0)
    
    # Use bounded log transform to keep velocities physical
    fop.modelTrans = pg.trans.TransLogLU(vs_min_constraint, vs_max_constraint)
    
    # 5. Setup inversion
    inv = pg.Inversion(fop)
    
    # 6. Run inversion
    model = pg.Vector(vs_start)
    data = pg.Vector(all_velocities)
    
    # Use relative error (larger for multi-mode to be more forgiving)
    rel_error = params.get('rel_error', 0.15 if len(selected_modes) > 1 else 0.1)
    
    # Run inversion with start model
    final_model = inv.run(
        data,
        lam=params.get('lambda', 100),
        errorVals=rel_error,
        startModel=model,
        maxIter=params.get('max_iter', 15),
        verbose=params.get('verbose', False)
    )
    
    # 7. Get response
    response_full = fop.response(final_model)
    
    # 8. Compute Vp and density
    nu = params.get('nu', 0.25)
    vp_model = vs_to_vp(np.array(final_model), nu)
    density = np.ones_like(final_model) * params.get('rho', 2000)
    
    # 9. Split response back into per-mode and compute statistics
    mode_fits = {}
    total_sse = 0
    total_n = 0
    
    for mode_name, (start_idx, end_idx) in mode_start_end.items():
        observed = all_velocities[start_idx:end_idx]
        calculated = np.array(response_full[start_idx:end_idx])
        residuals = calculated - observed
        
        sse = np.sum(residuals**2)
        n = len(residuals)
        rms = np.sqrt(sse / n)
        
        # Chi-squared (normalized by errors)
        errors_mode = all_errors[start_idx:end_idx]
        chi2_mode = np.sum((residuals / errors_mode)**2) / n
        
        mode_fits[mode_name] = {
            'observed': observed,
            'calculated': calculated,
            'residuals': residuals,
            'rms': rms,
            'chi2': chi2_mode,
            'sse': sse,
            'n': n
        }
        
        total_sse += sse
        total_n += n
    
    # Total statistics
    total_rms = np.sqrt(total_sse / total_n)
    total_rrms = total_rms / np.mean(all_velocities) * 100  # Relative RMS in %
    
    # Overall chi-squared
    total_chi2 = np.sum(((np.array(response_full) - all_velocities) / all_errors)**2) / total_n
    
    return {
        'vs_model': np.array(final_model),
        'vs_start': vs_start,  # Starting model for visualization
        'thickness': thickness,
        'vp_model': vp_model,
        'vp_start': vs_to_vp(vs_start, nu),  # Starting Vp model
        'density': density,
        'mode_fits': mode_fits,
        'total_chi2': total_chi2,
        'total_rms': total_rms,
        'rrms': total_rrms,
        'iterations': params.get('max_iter', 20),  # Actual iterations not easily accessible
        'convergence': total_chi2 <= params.get('target_chi2', 1.5),
        'lambda_final': params.get('lambda', 100),
        'success': True
    }


def create_start_model(params, observed_velocities, observed_frequencies=None):
    """
    Create layered start model based on wavelength-depth relationship
    
    Uses physical relationship:
    - VR (phase velocity) ≈ 0.9-0.95 * Vs, so Vs ≈ VR / 0.92
    - Investigation depth ≈ lambda/3 to lambda/4 (wavelength = VR/frequency)
    
    For wavelength-based models:
    1. Converts wavelength to investigation depth using depth_factor
    2. Creates layer structure based on n_layers and thickness_strategy
    3. For each layer, computes mean velocity of all samples within that depth range
    4. Scales phase velocity to Vs using vs_vr_ratio
    
    Parameters:
    -----------
    params : dict
        Must contain 'n_layers', optionally 'max_depth', 'thickness_strategy',
        'vs_vr_ratio', 'depth_factor'
    observed_velocities : array
        Observed phase velocities (typically from M0 or lowest available mode)
    observed_frequencies : array, optional
        Observed frequencies (for wavelength-based model)
    
    Returns:
    --------
    vs_start : array
        Initial Vs values for each layer
    thickness : array
        Thickness of each layer (n_layers - 1 values, last layer is half-space)
    """
    n_layers = params.get('n_layers', 8)
    thickness_strategy = params.get('thickness_strategy', 'equal')
    
    # Use wavelength-based model if frequencies are provided
    if observed_frequencies is not None and len(observed_frequencies) > 0:
        # Calculate wavelengths: lambda = VR / frequency
        wavelengths = observed_velocities / observed_frequencies
        
        # Investigation depth factor (lambda/3 or lambda/4)
        depth_factor = params.get('depth_factor', 3.5)  # Default: lambda/3.5
        investigation_depths = wavelengths / depth_factor
        
        # Vs-VR relationship: VR ≈ 0.92*Vs for Rayleigh wave fundamental mode
        # Therefore: Vs ≈ VR / 0.92
        vs_vr_ratio = params.get('vs_vr_ratio', 0.92)
        
        # Sort by investigation depth to create depth profile
        sort_idx = np.argsort(investigation_depths)
        sorted_depths = investigation_depths[sort_idx]
        sorted_velocities = observed_velocities[sort_idx]
        
        # Maximum investigation depth
        max_depth = params.get('max_depth', np.max(sorted_depths) * 1.2)
        
        # Create layer structure based on thickness strategy
        if thickness_strategy == 'logarithmic':
            # Logarithmic thickness distribution
            layer_boundaries = np.zeros(n_layers + 1)
            layer_boundaries[0] = 0
            layer_boundaries[-1] = max_depth
            
            # Create logarithmic distribution for intermediate boundaries
            # Need (n_layers - 1) intermediate boundaries
            log_depths = np.logspace(0, 1, n_layers - 1)
            log_depths = log_depths / log_depths[-1] * max_depth
            layer_boundaries[1:-1] = log_depths
        else:
            # Equal thickness distribution
            layer_boundaries = np.linspace(0, max_depth, n_layers + 1)
        
        # Thickness of each layer (except half-space)
        thickness = np.diff(layer_boundaries[:-1])  # n_layers - 1 values
        
        # For each layer, compute mean velocity from samples within that depth range
        vs_start = np.zeros(n_layers)
        
        for i in range(n_layers):
            layer_top = layer_boundaries[i]
            layer_bottom = layer_boundaries[i + 1]
            
            # Find all depth samples within this layer
            mask = (sorted_depths >= layer_top) & (sorted_depths < layer_bottom)
            
            if np.sum(mask) > 0:
                # Mean velocity of samples in this layer
                layer_mean_velocity = np.mean(sorted_velocities[mask])
                vs_start[i] = layer_mean_velocity / vs_vr_ratio
            else:
                # No samples in this layer - interpolate from nearest samples
                layer_center = (layer_top + layer_bottom) / 2
                vs_start[i] = np.interp(layer_center, sorted_depths, 
                                       sorted_velocities / vs_vr_ratio,
                                       left=sorted_velocities[0] / vs_vr_ratio * 0.8,
                                       right=sorted_velocities[-1] / vs_vr_ratio * 1.1)
        
        print(f"  Wavelength-based start model: depths up to {max_depth:.1f}m, "
              f"Vs range {vs_start.min():.0f}-{vs_start.max():.0f} m/s, "
              f"thickness={thickness_strategy}")
    
    else:
        # Fallback to generic model if no frequencies provided
        max_depth = params.get('max_depth', 40)
        
        # Logarithmic thickness distribution
        thickness = np.logspace(0, 1, n_layers - 1)
        thickness = thickness / thickness.sum() * max_depth
        
        # Simple gradient based on observed velocities
        vs_vr_ratio = params.get('vs_vr_ratio', 0.92)
        vs_min = max(50, np.min(observed_velocities) / vs_vr_ratio * 0.7)
        vs_max = np.max(observed_velocities) / vs_vr_ratio * 1.1
        vs_start = np.linspace(vs_min, vs_max, n_layers)
        
        print(f"  Generic start model: {max_depth:.1f}m depth, "
              f"Vs range {vs_start.min():.0f}-{vs_start.max():.0f} m/s")
    
    return vs_start, thickness


def vs_to_vp(vs, nu):
    """
    Convert Vs to Vp using Poisson's ratio
    
    Vp = Vs * sqrt(2(1-nu)/(1-2nu))
    """
    return vs * np.sqrt(2 * (1 - nu) / (1 - 2 * nu))


class MultiModeDispersionModelling(pg.core.ModellingBase):
    """
    Forward operator for joint multi-mode dispersion calculation
    
    Compatible with pygimli's inversion framework.
    Uses disba's surf96 for dispersion calculation.
    """
    
    def __init__(self, frequencies, mode_indices, thickness, nu=0.25, rho=2000, algorithm='dunkin'):
        """
        Parameters:
        -----------
        frequencies : array
            All frequencies (concatenated from all modes)
        mode_indices : list
            Mode label for each frequency (e.g., ['M0', 'M0', 'M1', ...])
        thickness : array
            Layer thicknesses (n_layers-1 values, last layer is half-space)
        nu : float
            Poisson's ratio
        rho : float or array
            Density (kg/m3)
        algorithm : str
            Dispersion algorithm ('dunkin' or 'fast-delta')
        """
        super().__init__(verbose=False)
        
        self.frequencies = np.array(frequencies)
        self.mode_indices = mode_indices
        self.thickness = np.array(thickness)
        self.nu = nu
        self.rho = rho if isinstance(rho, (list, np.ndarray)) else np.array([rho])
        self.algorithm = algorithm
        
        # Determine unique modes and their numbers
        self.unique_modes = sorted(set(mode_indices), key=lambda x: int(x.replace('M', '')))
        self.mode_numbers = [int(m.replace('M', '')) for m in self.unique_modes]
        self.max_mode = max(self.mode_numbers)
        
        # Create mesh (1D with layers) - following SRT-SWA-SJI pattern
        n_layers = len(thickness) + 1
        self.mesh_ = pg.meshtools.createMesh1D(n_layers)
        self.setMesh(self.mesh_)
        
        # Set region properties - required by pygimli
        self.regionManager().setParameterCount(n_layers)
    
    @property
    def parameterCount(self):
        """Return number of model parameters (layers)"""
        return len(self.thickness) + 1
    
    def response(self, model):
        """
        Calculate dispersion for all modes given Vs model
        
        Parameters:
        -----------
        model : array
            Vs values for each layer
        
        Returns:
        --------
        pg.Vector
            Concatenated phase velocities for all data points
        """
        vs_model = np.array(model)
        n_layers = len(vs_model)
        
        # Ensure thickness matches (may have been updated)
        if len(self.thickness) != n_layers - 1:
            # Adjust thickness if needed
            pass
        
        # Compute Vp from Vs
        vp_model = vs_to_vp(vs_model, self.nu)
        
        # Prepare density (constant or per-layer)
        if isinstance(self.rho, (int, float)):
            rho_model = np.ones(n_layers) * self.rho
        elif len(self.rho) == 1:
            rho_model = np.ones(n_layers) * self.rho[0]
        else:
            rho_model = self.rho
        
        # Calculate dispersion using disba PhaseDispersion (following user's example)
        try:
            if not DISBA_AVAILABLE:
                raise ImportError("disba not available")
            
            # Prepare thickness array with large value for half-space
            thk = self.thickness.copy()
            if len(thk) == n_layers - 1:
                thk = np.hstack((thk, 1000))  # Large value for half-space
            else:
                thk[-1] = 1000
            
            # Create velocity model arrays and convert to km units
            thk_km = thk / 1e3
            vp_km = vp_model / 1e3
            vs_km = vs_model / 1e3
            rho_g = rho_model / 1e3  # kg/m³ -> g/cm³
            
            # Create PhaseDispersion object with algorithm parameter
            pd = PhaseDispersion(thk_km, vp_km, vs_km, rho_g, algorithm=self.algorithm, dc=0.005)
            
            # Get unique frequencies per mode and calculate
            mode_responses = {}
            for mode_name in self.unique_modes:
                mode_num = int(mode_name.replace('M', ''))
                mask = np.array([m == mode_name for m in self.mode_indices])
                freqs_this_mode = self.frequencies[mask]
                
                if len(freqs_this_mode) > 0:
                    # Convert to periods (reversed like in user's example)
                    periods = 1.0 / freqs_this_mode[::-1]
                    
                    # Calculate phase velocity for this mode
                    try:
                        phase_disp = pd(periods, mode=mode_num, wave="rayleigh")
                        vr = phase_disp.velocity
                        
                        # Handle cases where dispersion curve is too short
                        # (happens when mode doesn't exist at certain periods)
                        if phase_disp.period.shape[0] < periods.shape[0]:
                            # Use extrapolation for missing values instead of NaN
                            n_missing = periods.shape[0] - phase_disp.period.shape[0]
                            if len(vr) > 0:
                                # Linear extrapolation from last two points
                                if len(vr) >= 2:
                                    slope = (vr[-1] - vr[-2]) / (phase_disp.period[-1] - phase_disp.period[-2])
                                    extrapolated = [vr[-1] + slope * (i+1) * (phase_disp.period[-1] - phase_disp.period[-2]) 
                                                   for i in range(n_missing)]
                                else:
                                    extrapolated = [vr[-1]] * n_missing
                                vr = np.append(vr, extrapolated)
                            else:
                                # No valid points - use rough estimate
                                vr = np.mean(vs_km) * 1.1 * 1000 * np.ones(periods.shape[0])
                        
                        # Reverse back and convert to m/s
                        vr = vr[::-1] * 1000
                        
                        # Check for NaN values and replace with reasonable estimates
                        if np.any(np.isnan(vr)):
                            # Replace NaN with interpolated or median values
                            nan_mask = np.isnan(vr)
                            if not np.all(nan_mask):
                                # Interpolate from non-NaN values
                                good_indices = np.where(~nan_mask)[0]
                                bad_indices = np.where(nan_mask)[0]
                                vr[bad_indices] = np.interp(bad_indices, good_indices, vr[good_indices])
                            else:
                                # All NaN - use rough estimate based on Vs
                                vr = np.mean(vs_km) * 1.1 * 1000 * np.ones_like(vr)
                        
                        mode_responses[mode_name] = vr
                    except Exception as e:
                        print(f"Warning: Failed to calculate mode {mode_name}: {e}")
                        # Use rough estimate instead of skipping
                        vr = np.mean(vs_km) * 1.1 * 1000 * np.ones(len(freqs_this_mode))
                        mode_responses[mode_name] = vr
            
            # Concatenate responses in same order as input using boolean indexing
            response = np.zeros(len(self.frequencies))
            for mode_name in self.unique_modes:
                if mode_name in mode_responses:
                    mask = np.array([m == mode_name for m in self.mode_indices])
                    response[mask] = mode_responses[mode_name]
                else:
                    # Mode calculation failed - use rough estimate
                    print(f"Warning: Using fallback for mode {mode_name}")
                    mask = np.array([m == mode_name for m in self.mode_indices])
                    response[mask] = np.mean(vs_model) * 1.1  # Rough estimate
            
            # Final NaN check - critical to prevent PyGIMLi errors
            if np.any(np.isnan(response)):
                print(f"Warning: NaN detected in response, replacing with estimates")
                nan_mask = np.isnan(response)
                if not np.all(nan_mask):
                    # Interpolate
                    good_indices = np.where(~nan_mask)[0]
                    bad_indices = np.where(nan_mask)[0]
                    response[bad_indices] = np.interp(bad_indices, good_indices, response[good_indices])
                else:
                    # All NaN - use model estimate
                    response = np.mean(vs_model) * 1.1 * np.ones_like(response)
            
        except Exception as e:
            import traceback
            print(f"Error in dispersion calculation: {e}")
            print(f"Traceback: {traceback.format_exc()}")
            print(f"Model: Vs={vs_model}, Vp={vp_model}")
            # Don't re-raise - return reasonable estimate to allow inversion to continue
            print(f"Returning fallback response to allow inversion to recover")
            response = np.mean(vs_model) * 1.1 * np.ones(len(self.frequencies))
        
        # Ensure no NaN or inf values
        response = np.nan_to_num(response, nan=np.mean(vs_model)*1.1, 
                                posinf=np.max(vs_model)*2, neginf=np.min(vs_model)*0.5)
        
        return pg.Vector(response)
    
    def createStartModel(self, dataVals):
        """
        Create default start model based on data values
        
        Returns a model with Vs values estimated from observed velocities
        """
        # Use median velocity as starting estimate for all layers
        median_velocity = np.median(dataVals)
        n_params = len(self.thickness) + 1  # number of layers
        return pg.Vector(n_params, median_velocity)
