#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Pick I/O utilities for reading and writing seismic picks.

This module provides functions to save and load first-arrival traveltime picks
in multiple formats:

**Supported Formats:**

1. **PyGimLi .sgt format** (read/write)
   - Widely used in seismic refraction processing
   - Contains station coordinates and source-geophone-time picks
   - Functions: read_sgt_file(), save_picks_to_sgt(), match_picks_to_geometry()

2. **Rayfract .LST format** (read only)
   - First break picks from Rayfract software
   - Format: shot_no trace_no position time synthetic
   - Invalid picks marked as -1.000
   - Functions: read_lst_file(), match_lst_picks_to_geometry()

**Typical Usage:**

For SGT files::

    from pick_io import read_sgt_file, match_picks_to_geometry
    
    # Read SGT file
    sgt_data = read_sgt_file('picks.sgt', verbose=True)
    
    # Match to geometry
    matched = match_picks_to_geometry(
        sgt_data['stations'], sgt_data['picks'],
        trace_positions, trace_elevations,
        source_positions, source_elevations
    )

For LST files::

    from pick_io import read_lst_file, match_lst_picks_to_geometry
    
    # Read LST file
    lst_data = read_lst_file('picks.LST', verbose=True)
    
    # Match to geometry
    matched = match_lst_picks_to_geometry(
        lst_data, trace_positions, source_positions
    )

Copyright (C) 2024, 2025, 2026 Sylvain Pasquet
Email: sylvain.pasquet@sorbonne-universite.fr

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
"""

import numpy as np


def save_picks_to_sgt(output_file, trace_positions, trace_elevations, 
                      source_positions, source_elevations, picks, errors, geophone_mapping=None):
    """
    Save picks to PyGimLi .sgt file format.
    
    The .sgt format stores station coordinates and source-geophone-time picks:
    - Header with number of stations and their (x, z) coordinates
    - Pick data with source index, geophone index, time, and error
    
    Parameters
    ----------
    output_file : str
        Path to save the .sgt file
    trace_positions : list of arrays
        List of trace position arrays (one per source)
    trace_elevations : list of arrays
        List of trace elevation arrays (one per source)
    source_positions : list
        List of source positions
    source_elevations : list
        List of source elevations
    picks : list of lists
        Nested list of pick times [source][trace]
    errors : list of lists
        Nested list of pick errors [source][trace]
    geophone_mapping : dict, optional
        Maps (shot_idx, trace_idx) -> sgt_geophone_number for order-matched picks.
        If provided, uses these geophone numbers instead of deriving from positions.
    
    Returns
    -------
    int
        Number of picks saved
    
    Raises
    ------
    ValueError
        If no picks are available to save
    """
    # Get unique traces from list of trace arrays
    trace_pairs = []
    for sublist_position, sublist_elevation in zip(trace_positions, trace_elevations):
        if sublist_position is not None:
            for trace, elevation in zip(sublist_position, sublist_elevation):
                trace_pairs.append((trace, elevation))

    # Get unique sources
    source_pairs = [(source, elevation) 
                   for source, elevation in zip(source_positions, source_elevations) 
                   if source is not None]

    # Convert to numpy structured arrays
    trace_pairs = np.array(trace_pairs, dtype=[('position', float), ('elevation', float)])
    source_pairs = np.array(source_pairs, dtype=[('position', float), ('elevation', float)])

    # Concatenate and get unique stations
    all_pairs = np.concatenate((trace_pairs, source_pairs))
    stations = np.unique(all_pairs)

    # Get trace indices in station list
    trace_indices = [
        np.where((stations['position'] == trace_pair['position']) & 
                (stations['elevation'] == trace_pair['elevation']))[0][0] 
        for trace_pair in trace_pairs
    ]

    # Get source indices in station list
    source_indices = [
        np.where((stations['position'] == source_pair['position']) & 
                (stations['elevation'] == source_pair['elevation']))[0][0] 
        for source_pair in source_pairs
    ]

    # Count non-NaN picks
    picks_list = [pick for sublist in picks if sublist is not None for pick in sublist]
    n_picks = np.sum(~np.isnan(picks_list))

    if n_picks == 0:
        raise ValueError("No picks to save!")
    
    # Write .sgt file
    with open(output_file, 'w') as f:
        # Write number of stations
        f.write(f"{len(stations)} # shot/geophone points\n")
        f.write("# x\ty\n")
        for station in stations:
            x = station[0]
            y = station[1]
            f.write(f"{x}\t{y}\n")
        
        # Write number of picks
        f.write(f"{n_picks} # measurements\n")
        f.write("# s\tg\tt\terr\n")

        # Write pick data
        for i, pick_list in enumerate(picks):
            if pick_list is not None:
                for j, pick in enumerate(pick_list):
                    if not np.isnan(pick):
                        error = errors[i][j]
                        # Determine geophone number to use
                        if geophone_mapping is not None and (i, j) in geophone_mapping:
                            # Use the mapped geophone number from order matching
                            geophone_num = geophone_mapping[(i, j)]
                        else:
                            # Use the trace index-based geophone number
                            geophone_num = trace_indices[j] + 1
                        # Write source index, geophone number, pick time, pick error
                        # Indices are 1-based in .sgt format
                        f.write(f"{source_indices[i] + 1}\t{geophone_num}\t{pick:.5f}\t{error:.5f}\n")
    
    return n_picks


def read_sgt_file(sgt_file, verbose=False):
    """
    Read PyGimLi .sgt file and parse stations and picks.
    
    Automatically detects coordinate format:
    - 2 columns: (x, z)
    - 3 columns: (x, y, z) or (x, z, 0) depending on which columns have values
    
    Automatically detects column order for picks using header comments.
    Supports various synonyms: s/src/source, g/geo/geophone/r/recv, t/time/tt, err/error/unc
    
    Parameters
    ----------
    sgt_file : str
        Path to .sgt file
    verbose : bool, optional
        Print debug information, default False
    
    Returns
    -------
    dict
        Dictionary containing:
        - 'stations': list of (x, y, z) tuples
        - 'picks': list of (source_idx, geophone_idx, time, error) tuples
        - 'n_stations': int
        - 'n_picks': int
    """
    with open(sgt_file, 'r') as f:
        # Read number of stations
        n_stations = int(f.readline().split('#')[0].strip())
        if verbose:
            print(f"Number of stations: {n_stations}")

        # Skip comment lines
        flag_comment = True
        while flag_comment:
            line = f.readline().strip()
            if '#' in line[0]:
                if verbose:
                    print(f"Comment: {line}")
                flag_comment = True
            else:
                flag_comment = False

        # Read station coordinates
        coords_list = []
        for i in range(n_stations):
            if i > 0:
                line = f.readline().strip()

            if verbose and (i < 5 or i > n_stations - 5):
                print(f"Reading station line: {line}")
        
            if line:
                parts = line.split()
                coords = [float(p) for p in parts]
                coords_list.append(coords)
        
        # Determine coordinate format
        stations = []
        if len(coords_list) > 0:
            n_cols = len(coords_list[0])
            
            if n_cols == 2:
                # 2 columns: (x, z) format
                for coords in coords_list:
                    stations.append((coords[0], 0.0, coords[1]))
                if verbose:
                    print(f"Detected 2-column format: (x, z)")
            
            elif n_cols == 3:
                # Check which columns have non-zero values
                col0_nonzero = any(abs(coords[0]) > 1e-10 for coords in coords_list)
                col1_nonzero = any(abs(coords[1]) > 1e-10 for coords in coords_list)
                col2_nonzero = any(abs(coords[2]) > 1e-10 for coords in coords_list)
                
                if verbose:
                    print(f"Column non-zero detection: col0={col0_nonzero}, col1={col1_nonzero}, col2={col2_nonzero}")
                
                if col0_nonzero and col1_nonzero and col2_nonzero:
                    # All 3 columns non-zero: (x, y, z)
                    for coords in coords_list:
                        stations.append((coords[0], coords[1], coords[2]))
                    if verbose:
                        print(f"Detected 3-column format: (x, y, z)")
                elif col0_nonzero and col1_nonzero and not col2_nonzero:
                    # Columns 0 and 1 non-zero: (x, z, 0) -> store as (x, 0, z)
                    for coords in coords_list:
                        stations.append((coords[0], 0.0, coords[1]))
                    if verbose:
                        print(f"Detected 3-column format: (x, z, 0)")

        # Read number of picks
        n_picks = int(f.readline().split('#')[0].strip())
        if verbose:
            print(f"Number of picks: {n_picks}")

        # Initialize default column indices (standard order: s g t err)
        s_ind = 0
        g_ind = 1
        t_ind = 2
        err_ind = 3

        # Read optional header comment lines to infer column order
        comment_lines = []
        while True:
            line = f.readline()
            if not line:
                line = ""
                break
            line = line.strip()
            if line and line[0] == '#':
                comment_lines.append(line)
                if verbose:
                    print(f"Comment: {line}")
                continue
            else:
                break

        # Parse comment lines to detect column order
        if comment_lines:
            synonyms = {
                's': 's', 'src': 's', 'source': 's',
                'g': 'g', 'geo': 'g', 'geophone': 'g', 'r': 'g', 'recv': 'g', 'receiver': 'g',
                't': 't', 'time': 't', 'tt': 't', 'pick': 't', 'tpick': 't',
                'err': 'err', 'error': 'err', 'unc': 'err', 'uncertainty': 'err', 'sigma': 'err'
            }
            best_fields = []
            for cl in comment_lines:
                cl_proc = cl[1:] if cl.startswith('#') else cl
                for sep in [',', ';', '|']:
                    cl_proc = cl_proc.replace(sep, ' ')
                tokens = [tok.strip().lower().strip(':') for tok in cl_proc.split()]
                fields = []
                for tok in tokens:
                    if tok in synonyms:
                        canon = synonyms[tok]
                        if len(fields) == 0 or fields[-1] != canon:
                            fields.append(canon)
                unique_fields = []
                for fcanon in fields:
                    if fcanon not in unique_fields:
                        unique_fields.append(fcanon)
                if len(unique_fields) > len(best_fields):
                    best_fields = unique_fields
            
            if best_fields:
                try:
                    if 's' in best_fields:
                        s_ind = best_fields.index('s')
                    if 'g' in best_fields:
                        g_ind = best_fields.index('g')
                    if 't' in best_fields:
                        t_ind = best_fields.index('t')
                    if 'err' in best_fields:
                        err_ind = best_fields.index('err')
                    if verbose:
                        print(f"Detected column order from header: {best_fields}")
                        print(f"Using indices: s={s_ind}, g={g_ind}, t={t_ind}, err={err_ind}")
                except (ValueError, IndexError) as e:
                    if verbose:
                        print(f"Error parsing column order: {e}, using defaults")

        # Read picks
        picks = []
        for i in range(n_picks):
            if i == 0 and line:
                # Use the line we already read
                pass
            else:
                line = f.readline().strip()
            
            if line:
                parts = line.split()
                if len(parts) >= 4:
                    # Parse based on detected column order
                    source_idx = int(parts[s_ind])
                    geophone_idx = int(parts[g_ind])
                    time = float(parts[t_ind])
                    error = float(parts[err_ind])
                    picks.append((source_idx, geophone_idx, time, error))

    return {
        'stations': stations,
        'picks': picks,
        'n_stations': n_stations,
        'n_picks': n_picks
    }


def match_picks_to_geometry(sgt_stations, sgt_picks, 
                           trace_positions, trace_elevations,
                           source_positions, source_elevations,
                           matching_mode='exact_x', tolerance=0.01,
                           max_distance=50.0, verbose=False):
    """
    Match picks from SGT file to actual trace/source geometry.
    
    Supports multiple matching strategies:
    - 'exact_x': Match by exact X coordinate (within tolerance)
    - 'nearest_x': Match to nearest X coordinate (within max_distance)
    - 'nearest_xz': Match to nearest (X, Z) coordinate (within max_distance)
    
    Parameters
    ----------
    sgt_stations : list of tuples
        Station coordinates [(x, y, z), ...] from SGT file
    sgt_picks : list of tuples
        Pick data [(source_idx, geophone_idx, time, error), ...] from SGT file
        Note: Indices are 1-based
    trace_positions : list of arrays
        Actual trace positions in dataset
    trace_elevations : list of arrays
        Actual trace elevations in dataset
    source_positions : list
        Actual source positions in dataset
    source_elevations : list
        Actual source elevations in dataset
    matching_mode : str, optional
        Matching strategy: 'exact_x', 'nearest_x', or 'nearest_xz'
    tolerance : float, optional
        Tolerance for exact matching in meters, default 0.01
    max_distance : float, optional
        Maximum distance for nearest neighbor matching in meters, default 50.0
    verbose : bool, optional
        Print debug information, default False
    
    Returns
    -------
    dict
        Dictionary containing:
        - 'picks': 2D list [source][trace] of matched pick times (NaN where no pick)
        - 'errors': 2D list [source][trace] of matched errors
        - 'n_matched': Number of successfully matched picks
        - 'n_total': Total picks in SGT file
    """
    n_sources = len(source_positions)
    n_traces_per_source = [len(tp) if tp is not None else 0 for tp in trace_positions]
    
    # Initialize pick and error arrays with NaN
    matched_picks = [[np.nan] * n_traces for n_traces in n_traces_per_source]
    matched_errors = [[np.nan] * n_traces for n_traces in n_traces_per_source]
    
    n_matched = 0
    n_total = len(sgt_picks)
    
    for source_idx_1based, geophone_idx_1based, pick_time, pick_error in sgt_picks:
        # Convert to 0-based indexing
        source_idx = source_idx_1based - 1
        geophone_idx = geophone_idx_1based - 1
        
        # Get station coordinates from SGT file (0-based indexing)
        if source_idx < len(sgt_stations):
            sgt_source_x, sgt_source_y, sgt_source_z = sgt_stations[source_idx]
        else:
            continue
            
        if geophone_idx < len(sgt_stations):
            sgt_geo_x, sgt_geo_y, sgt_geo_z = sgt_stations[geophone_idx]
        else:
            continue
        
        # Find matching source in actual geometry
        matched_source = None
        for i_src, (src_pos, src_elev) in enumerate(zip(source_positions, source_elevations)):
            if src_pos is None:
                continue
            
            if matching_mode == 'exact_x':
                if abs(src_pos - sgt_source_x) < tolerance:
                    matched_source = i_src
                    break
            elif matching_mode == 'nearest_x':
                if matched_source is None or abs(src_pos - sgt_source_x) < abs(source_positions[matched_source] - sgt_source_x):
                    if abs(src_pos - sgt_source_x) <= max_distance:
                        matched_source = i_src
            elif matching_mode == 'nearest_xz':
                dist = np.sqrt((src_pos - sgt_source_x)**2 + (src_elev - sgt_source_z)**2)
                if matched_source is None or dist < min_dist:
                    if dist <= max_distance:
                        matched_source = i_src
                        min_dist = dist
        
        if matched_source is None:
            continue
        
        # Find matching trace in actual geometry
        matched_trace = None
        if trace_positions[matched_source] is not None:
            for i_tr, (tr_pos, tr_elev) in enumerate(zip(trace_positions[matched_source], 
                                                          trace_elevations[matched_source])):
                if matching_mode == 'exact_x':
                    if abs(tr_pos - sgt_geo_x) < tolerance:
                        matched_trace = i_tr
                        break
                elif matching_mode == 'nearest_x':
                    if matched_trace is None or abs(tr_pos - sgt_geo_x) < abs(trace_positions[matched_source][matched_trace] - sgt_geo_x):
                        if abs(tr_pos - sgt_geo_x) <= max_distance:
                            matched_trace = i_tr
                elif matching_mode == 'nearest_xz':
                    dist = np.sqrt((tr_pos - sgt_geo_x)**2 + (tr_elev - sgt_geo_z)**2)
                    if matched_trace is None or dist < min_dist_trace:
                        if dist <= max_distance:
                            matched_trace = i_tr
                            min_dist_trace = dist
        
        # Assign pick if both source and trace matched
        if matched_trace is not None:
            matched_picks[matched_source][matched_trace] = pick_time
            matched_errors[matched_source][matched_trace] = pick_error
            n_matched += 1
    
    if verbose:
        print(f"Matched {n_matched}/{n_total} picks using {matching_mode} mode")
    
    return {
        'picks': matched_picks,
        'errors': matched_errors,
        'n_matched': n_matched,
        'n_total': n_total
    }


def read_vs_file(vs_file, verbose=False):
    """
    Read SeisImager/PickWin .vs file and parse first break picks.
    
    VS file format:
    - Line 1: Header (discarded)
    - Line 2: min_offset n_shots spacing
    - Line 3: Header (discarded)
    - For each shot:
      - Shot header: shot_position n_picks 0
      - n_picks data lines: geophone_offset time quality_flag
    
    Parameters
    ----------
    vs_file : str
        Path to .vs file
    verbose : bool, optional
        Print debug information, default False
    
    Returns
    -------
    dict
        Dictionary containing:
        - 'picks': list of (shot_position, geophone_offset, time) tuples
        - 'n_picks': int (number of valid picks)
        - 'shots': list of shot positions
        - 'n_shots': int (number of shots from header)
        - 'spacing': float (spacing from header)
    """
    picks = []
    shots = []
    n_shots = 0
    spacing = 0.0
    
    with open(vs_file, 'r') as f:
        lines = f.readlines()
    
    if len(lines) < 3:
        if verbose:
            print(f"File too short: {len(lines)} lines")
        return {'picks': [], 'n_picks': 0, 'shots': [], 'n_shots': 0, 'spacing': 0.0}
    
    # Line 1: discard
    # Line 2: 0 n_shots spacing
    try:
        parts = lines[1].strip().split()
        n_shots = int(parts[1])
        spacing = float(parts[2])
    except (ValueError, IndexError) as e:
        if verbose:
            print(f"Error parsing header line 2: {e}")
        return {'picks': [], 'n_picks': 0, 'shots': [], 'n_shots': 0, 'spacing': 0.0}
    
    # Start reading shot data from line 3 (index 2)
    # No third discard line - shot data starts immediately after header
    line_idx = 2
    shot_count = 0
    
    while line_idx < len(lines) and shot_count < n_shots:
        line = lines[line_idx].strip()
        
        if verbose and shot_count < 3:
            print(f"\nProcessing line {line_idx}: '{line}'")
        
        # Skip empty lines
        if not line:
            if verbose and shot_count < 3:
                print(f"  -> Empty line, skipping")
            line_idx += 1
            continue
        
        # Check if this is a shot header line (third value equals 0)
        parts = line.split()
        if verbose and shot_count < 3:
            print(f"  -> Parts: {parts}, length: {len(parts)}")
        
        if len(parts) >= 3:
            try:
                # Check if the third value (index 2) is 0
                third_value = float(parts[2])
                if verbose and shot_count < 3:
                    print(f"  -> Third value: {third_value}, is shot header: {abs(third_value) < 0.0001}")
                
                if abs(third_value) < 0.0001:  # Close to zero (shot header)
                    shot_position = float(parts[0])
                    n_picks_for_shot = int(parts[1])
                    shot_position = float(parts[0])
                    n_picks_for_shot = int(parts[1])
                    shots.append(shot_position)
                    shot_count += 1
                    
                    if verbose:
                        print(f"Shot {shot_count}: position {shot_position}m, {n_picks_for_shot} picks")
                    
                    # Read the pick lines for this shot
                    for i in range(n_picks_for_shot):
                        line_idx += 1
                        if line_idx >= len(lines):
                            break
                        
                        pick_line = lines[line_idx].strip()
                        pick_parts = pick_line.split()
                        
                        if len(pick_parts) >= 3:
                            try:
                                geophone_position = float(pick_parts[0])  # Absolute geophone position
                                time = float(pick_parts[1])
                                # quality_flag = int(pick_parts[2])
                                
                                picks.append((shot_position, geophone_position, time))
                            except (ValueError, IndexError):
                                if verbose:
                                    print(f"Error parsing pick line: {pick_line}")
            except (ValueError, IndexError) as e:
                if verbose:
                    print(f"Error parsing shot header: {line}, error: {e}")
        
        line_idx += 1
    
    if verbose:
        print(f"Read {len(picks)} picks from {len(shots)} shots")
        print(f"Shot positions: {shots[:5]}..." if len(shots) > 5 else f"Shot positions: {shots}")
    
    return {
        'picks': picks,
        'n_picks': len(picks),
        'shots': shots,
        'n_shots': n_shots,
        'spacing': spacing
    }


def match_vs_picks_to_geometry(vs_picks_data, trace_positions, source_positions,
                                offset_tolerance=0.5, position_scale=1.0, verbose=False):
    """
    Match picks from .vs file to actual trace/source geometry.
    
    VS files store shot positions and absolute geophone positions.
    
    Parameters
    ----------
    vs_picks_data : dict
        Dictionary returned by read_vs_file()
    trace_positions : list of arrays
        Actual trace positions in dataset
    source_positions : list
        Actual source positions in dataset
    offset_tolerance : float, optional
        Tolerance for matching offsets in meters, default 0.5
    position_scale : float, optional
        Scaling factor to apply to VS positions (VS_pos * scale = actual_pos), default 1.0
    verbose : bool, optional
        Print debug information, default False
    
    Returns
    -------
    dict
        Dictionary containing:
        - 'picks': 2D list [source][trace] of matched pick times (NaN where no pick)
        - 'errors': 2D list [source][trace] of matched errors (0.001 default)
        - 'n_matched': Number of successfully matched picks
        - 'n_total': Total picks in VS file
        - 'unmatched_shots': List of shot positions that couldn't be matched
    """
    n_sources = len(source_positions)
    n_traces_per_source = [len(tp) if tp is not None else 0 for tp in trace_positions]
    
    # Initialize pick and error arrays with NaN
    matched_picks = [[np.nan] * n_traces for n_traces in n_traces_per_source]
    matched_errors = [[0.001] * n_traces for n_traces in n_traces_per_source]
    
    n_matched = 0
    n_total = vs_picks_data['n_picks']
    unmatched_shots = set()
    unmatched_details = {}
    
    for shot_position, geophone_position, time in vs_picks_data['picks']:
        # Apply position scaling
        scaled_shot_position = shot_position * position_scale
        scaled_geophone_position = geophone_position * position_scale
        
        # Find matching source by position
        matched_source = None
        min_source_dist = float('inf')
        
        for i_src, src_pos in enumerate(source_positions):
            dist = abs(src_pos - scaled_shot_position)
            if dist < min_source_dist:
                min_source_dist = dist
                if dist <= offset_tolerance:
                    matched_source = i_src
        
        if matched_source is None:
            if shot_position not in unmatched_shots:
                unmatched_shots.add(shot_position)  # Store original position
                if shot_position not in unmatched_details:
                    unmatched_details[shot_position] = {
                        'scaled_position': scaled_shot_position,
                        'min_dist': min_source_dist,
                        'closest_source': source_positions[np.argmin(np.abs(source_positions - scaled_shot_position))],
                        'n_picks': 0
                    }
                unmatched_details[shot_position]['n_picks'] += 1
            continue
        
        # Find matching trace by geophone position
        if trace_positions[matched_source] is None:
            continue
        
        matched_trace = None
        min_trace_dist = float('inf')
        
        for i_tr, tr_pos in enumerate(trace_positions[matched_source]):
            dist = abs(tr_pos - scaled_geophone_position)
            if dist < min_trace_dist and dist <= offset_tolerance:
                min_trace_dist = dist
                matched_trace = i_tr
        
        if matched_trace is not None:
            # Convert time to seconds if in milliseconds
            time_seconds = time / 1000.0 if time > 1.0 else time
            matched_picks[matched_source][matched_trace] = time_seconds
            n_matched += 1
    
    if verbose:
        print(f"\nMatched {n_matched}/{n_total} picks ({n_matched*100.0/n_total:.1f}%)")
        if unmatched_shots:
            print(f"Unmatched shots: {len(unmatched_shots)}")
            # Show first 10 unmatched shots with details
            for i, shot_pos in enumerate(sorted(unmatched_shots)[:10]):
                if shot_pos in unmatched_details:
                    details = unmatched_details[shot_pos]
                    print(f"  Shot {shot_pos} (scaled: {details['scaled_position']:.2f}m) - "
                          f"closest source: {details['closest_source']:.2f}m, "
                          f"distance: {details['min_dist']:.2f}m, "
                          f"picks: {details['n_picks']}")
            if len(unmatched_shots) > 10:
                print(f"  ... and {len(unmatched_shots) - 10} more")
    
    return {
        'picks': matched_picks,
        'errors': matched_errors,
        'n_matched': n_matched,
        'n_total': n_total,
        'unmatched_shots': sorted(list(unmatched_shots))
    }


def read_lst_file(lst_file, verbose=False):
    """
    Read Rayfract .LST file and parse first break picks.
    
    LST file format:
    - Header lines (ignored until data starts)
    - Data lines: shot_no trace pos. time synthetic
    - Time of -1.000 indicates no pick
    
    Parameters
    ----------
    lst_file : str
        Path to .LST file
    verbose : bool, optional
        Print debug information, default False
    
    Returns
    -------
    dict
        Dictionary containing:
        - 'picks': list of (shot_no, trace_no, position, time) tuples
        - 'n_picks': int (number of valid picks, excluding -1.000)
        - 'shots': list of unique shot numbers
        - 'traces_per_shot': dict mapping shot_no to list of (trace_no, position, time) tuples
    """
    picks = []
    traces_per_shot = {}
    shots = []
    
    with open(lst_file, 'r') as f:
        # Skip header lines until we find data
        # Data starts when we see lines matching the pattern: integer integer float float
        for line in f:
            line = line.strip()
            
            # Skip empty lines and comments
            if not line or line.startswith('#'):
                continue
            
            # Try to parse as data line
            parts = line.split()
            if len(parts) >= 4:
                try:
                    shot_no = int(parts[0])
                    trace_no = int(parts[1])
                    position = float(parts[2])
                    time = float(parts[3])
                    # synthetic = float(parts[4]) if len(parts) > 4 else None  # Optional synthetic time
                    
                    # Skip picks marked as invalid (-1.000)
                    if time > 0:  # Valid pick
                        picks.append((shot_no, trace_no, position, time))
                        
                        # Organize by shot
                        if shot_no not in traces_per_shot:
                            traces_per_shot[shot_no] = []
                            shots.append(shot_no)
                        traces_per_shot[shot_no].append((trace_no, position, time))
                    
                except (ValueError, IndexError):
                    # Not a data line, skip
                    continue
    
    return {
        'picks': picks,
        'n_picks': len(picks),
        'shots': sorted(shots),
        'traces_per_shot': traces_per_shot
    }


def match_lst_picks_to_geometry(lst_picks_data, trace_positions, source_positions,
                                 shot_number_offset=0, position_tolerance=0.1,
                                 position_scale=1.0, verbose=False):
    """
    Match picks from LST file to actual trace/source geometry.
    
    LST files contain shot number, trace number, and position information.
    This function matches based on:
    1. Shot number -> source index (with optional offset)
    2. Trace position -> actual trace position (within tolerance)
    
    Parameters
    ----------
    lst_picks_data : dict
        Dictionary returned by read_lst_file()
    trace_positions : list of arrays
        Actual trace positions in dataset
    source_positions : list
        Actual source positions in dataset
    shot_number_offset : int, optional
        Offset to apply to shot numbers (LST shot 1 = dataset shot 1+offset), default 0
    position_tolerance : float, optional
        Tolerance for matching trace positions in meters, default 0.1
    position_scale : float, optional
        Scaling factor to apply to LST positions (LST_pos * scale = actual_pos), default 1.0
    verbose : bool, optional
        Print debug information, default False
    
    Returns
    -------
    dict
        Dictionary containing:
        - 'picks': 2D list [source][trace] of matched pick times (NaN where no pick)
        - 'errors': 2D list [source][trace] of matched errors (0.001 default)
        - 'n_matched': Number of successfully matched picks
        - 'n_total': Total picks in LST file
        - 'unmatched_shots': List of shot numbers that couldn't be matched
        - 'unmatched_positions': List of (shot, position) that couldn't be matched
    """
    n_sources = len(source_positions)
    n_traces_per_source = [len(tp) if tp is not None else 0 for tp in trace_positions]
    
    # Initialize pick and error arrays with NaN
    matched_picks = [[np.nan] * n_traces for n_traces in n_traces_per_source]
    matched_errors = [[0.001] * n_traces for n_traces in n_traces_per_source]  # Default error 1ms
    
    n_matched = 0
    n_total = lst_picks_data['n_picks']
    unmatched_shots = set()
    unmatched_positions = []
    unmatched_shot_details = {}
    
    for shot_no, trace_no, position, time in lst_picks_data['picks']:
        # Convert shot number to source index (0-based)
        source_idx = shot_no - 1 + shot_number_offset
        
        # Check if source index is valid
        if source_idx < 0 or source_idx >= n_sources:
            if shot_no not in unmatched_shots:
                unmatched_shots.add(shot_no)
                if shot_no not in unmatched_shot_details:
                    unmatched_shot_details[shot_no] = f"out of range [0, {n_sources-1}]"
            continue
        
        # Find matching trace by position
        if trace_positions[source_idx] is None:
            if shot_no not in unmatched_shot_details:
                unmatched_shot_details[shot_no] = "no trace positions for this shot"
            continue
        
        # Apply position scaling
        scaled_position = position * position_scale
        
        matched_trace = None
        min_distance = float('inf')
        
        for i_tr, tr_pos in enumerate(trace_positions[source_idx]):
            dist = abs(tr_pos - scaled_position)
            if dist < min_distance and dist <= position_tolerance:
                min_distance = dist
                matched_trace = i_tr
        
        if matched_trace is not None:
            # Convert milliseconds to seconds if needed
            time_seconds = time / 1000.0 if time > 1.0 else time
            matched_picks[source_idx][matched_trace] = time_seconds
            n_matched += 1
        else:
            unmatched_positions.append((shot_no, position))
            # Track which traces are available for this shot
            if shot_no not in unmatched_shot_details:
                available_positions = trace_positions[source_idx] if trace_positions[source_idx] is not None else []
                if len(available_positions) > 0:
                    min_avail = min(available_positions)
                    max_avail = max(available_positions)
                    unmatched_shot_details[shot_no] = f"scaled pos {scaled_position:.1f}m (LST: {position:.1f}m) not in [{min_avail:.1f}, {max_avail:.1f}]m"
                else:
                    unmatched_shot_details[shot_no] = "no traces available"
    
    return {
        'picks': matched_picks,
        'errors': matched_errors,
        'n_matched': n_matched,
        'n_total': n_total,
        'unmatched_shots': sorted(unmatched_shots),
        'unmatched_positions': unmatched_positions
    }
