"""
Inversion Manager Module for Pyckster

This module handles seismic traveltime inversion processes, including data preparation,
inversion setup, and execution.
"""

import os
import numpy as np
import pygimli as pg
from pygimli.physics import TravelTimeManager as refrac

def run_inversion_from_file(sgt_file, params=None):
    """
    Run seismic traveltime inversion from a .sgt file.
    
    Parameters:
    -----------
    sgt_file : str
        Path to the .sgt file containing picks data
    params : dict, optional
        Dictionary containing inversion parameters
    
    Returns:
    --------
    refrac_manager : TravelTimeManager
        PyGIMLI's TravelTimeManager object with inversion results
    """

    # Check if the file exists
    if not os.path.exists(sgt_file):
        print(f"Error: File {sgt_file} does not exist.")
        return None
    
    # Create refraction manager from file
    try:
        refrac_manager = refrac(sgt_file)
    except Exception as e:
        print(f"Error creating refraction manager: {e}")
        return None
    
    return run_inversion(refrac_manager, params)

def run_inversion_from_manager(refrac_manager, params=None):
    """
    Run seismic traveltime inversion from a TravelTimeManager object.

    Parameters:
    -----------
    refrac_manager : TravelTimeManager
        PyGIMLI's TravelTimeManager object with picks data
    params : dict, optional
        Dictionary containing inversion parameters

    Returns:
    --------
    refrac_manager : TravelTimeManager
        PyGIMLI's TravelTimeManager object with inversion results
    """
    if not isinstance(refrac_manager, refrac):
        raise TypeError("refrac_manager must be an instance of TravelTimeManager")
    
    return run_inversion(refrac_manager, params)

def run_inversion(refrac_manager, params=None):
    """ Run seismic traveltime inversion on a TravelTimeManager object.
    This function prepares the inversion parameters, applies time corrections to
    set t=0 at shot positions (which also removes zero-time picks), and runs the inversion.
    
    Parameters:
    -----------
    refrac_manager : TravelTimeManager
        PyGIMLI's TravelTimeManager object with picks data
    params : dict, optional
        Dictionary containing inversion parameters. If None, default parameters are used.
        
        Time correction parameters:
        - apply_time_correction : bool (default: True)
            Apply time correction to set t=0 at shot positions before inversion.
            Zero-time picks are automatically removed after correction.
        - min_velocity : float (default: 500)
            Minimum/fallback velocity for time correction in m/s
        - max_offset_threshold : float or None (default: None, auto = half geophone spacing)
            Maximum offset for reliable surface velocity estimation
        - debug_time_correction : bool (default: False)
            Print detailed debug information about time correction
    
    Returns:
    --------
    refrac_manager : TravelTimeManager
        PyGIMLI's TravelTimeManager object with inversion results
    """
    
    # Initialize parameters dictionary if not provided
    if params is None:
        params = {}
    
    # Apply time corrections to set t=0 at shot positions before inversion
    # This ensures proper zero reference and handles negative times
    apply_correction = params.get('apply_time_correction', True)  # Default: always correct
    
    if apply_correction:
        correct_time_picks(
            refrac_manager, 
            params.get('min_velocity', 500),
            max_offset_threshold=params.get('max_offset_threshold', None),
            debug=params.get('debug_time_correction', False)
        )
    
    # Create inversion parameter dictionary with defaults
    inversion_params = {
        'verbose': params.get('verbose', True),
        'vTop': params.get('vTop', 500),
        'vBottom': params.get('vBottom', 5000),
        'secNodes': params.get('secNodes', 2),
        'lam': params.get('lam', 30),
        'maxIter': params.get('maxIter', 6),
        'zWeight': params.get('zWeight', 0.5),
        'balanceDepth': params.get('balanceDepth', False),
    }
    
    # Add optional parameters only if they're not None
    optional_params = ['paraDX', 'paraDepth', 'paraMaxCellSize']
    for param in optional_params:
        if params.get(param) is not None:
            inversion_params[param] = params.get(param)
    
    # Run the inversion
    refrac_manager.invert(**inversion_params)
    
    # Print inversion QC
    print("chi2 =", round(refrac_manager.inv.chi2(), 2), 
          "- rrms =", round(refrac_manager.inv.relrms(), 2), '%')
    
    return refrac_manager

def correct_time_picks(refrac_manager, min_velocity=500, max_offset_threshold=None, debug=False):
    """
    Correct negative times in the picks, assuming a surface velocity.
    The function resets minimum times for each shot at 0, then corrects 
    the time picks with the minimum offset and the assumed surface velocity.
    After correction, removes picks with t=0 (at shot positions) as PyGIMLI cannot invert them.
    
    Parameters:
    -----------
    refrac_manager : TravelTimeManager
        PyGIMLI's TravelTimeManager object with data
    min_velocity : float or array
        Fallback minimum velocity for time correction in m/s (default: 500)
    max_offset_threshold : float or None
        Maximum offset (in m) to consider picks "close enough" for reliable 
        surface velocity estimation. If None, defaults to half the geophone spacing.
    debug : bool
        If True, print detailed debug information about velocity estimation (default: False)
    """
    offset_max = None
    apply_positive_correction = True

    unique_sources = np.unique(refrac_manager.data['s'])
    sensor_positions = refrac_manager.data.sensors().array()
    
    if not isinstance(min_velocity, np.ndarray):
        min_velocity = np.ones(len(unique_sources)) * min_velocity

    # Calculate default max_offset_threshold as half the geophone spacing
    if max_offset_threshold is None:
        x_positions = sensor_positions[:, 0]
        unique_x = np.unique(x_positions)
        if len(unique_x) > 1:
            geophone_spacing = np.median(np.diff(np.sort(unique_x)))
            max_offset_threshold = geophone_spacing / 2.0
            print(f"Using max_offset_threshold = {max_offset_threshold:.2f} m (half of {geophone_spacing:.2f} m geophone spacing)")
        else:
            # Fallback if we can't determine spacing
            max_offset_threshold = 5.0
            print(f"Warning: Could not determine geophone spacing, using default threshold = {max_offset_threshold} m")

    # Estimate surface velocities from picks using existing get_surf_velocity logic
    surf_vel_data = get_surf_velocity(refrac_manager, smooth=False, show_plot=False)
    # Returns: [x_surf_array, smooth_vel_surf_array, vel_surf_array, min_offset_array, max_offset_array]
    
    estimated_velocities = []
    velocity_per_source = np.copy(min_velocity)
    velocity_source_type = []  # Track how each velocity was determined
    
    if debug:
        print("\n" + "="*80)
        print("VELOCITY ESTIMATION DEBUG")
        print("="*80)
        print(f"{'Shot':>4} {'X_pos':>8} {'Min_off':>8} {'Est_vel':>8} {'Status':^30} {'Used_vel':>8}")
        print("-"*80)
    
    for i in range(len(unique_sources)):
        x_source = sensor_positions[unique_sources[i], 0]
        min_offset = surf_vel_data[i, 3]
        estimated_vel = surf_vel_data[i, 2]  # Use unsmoothed velocity
        
        # Check conditions
        offset_ok = min_offset <= max_offset_threshold
        vel_ok = 100 < estimated_vel < 5000
        
        # Only use estimated velocity if picks are close enough and velocity is reasonable
        if offset_ok and vel_ok:
            velocity_per_source[i] = estimated_vel
            estimated_velocities.append(estimated_vel)
            velocity_source_type.append('estimated')
            status = "✓ ESTIMATED"
        else:
            # Build rejection reason
            reasons = []
            if not offset_ok:
                reasons.append(f"offset>{max_offset_threshold:.1f}m")
            if not vel_ok:
                reasons.append(f"vel_out_of_range")
            status = f"✗ REJECTED: {', '.join(reasons)}"
            velocity_source_type.append('pending')
        
        if debug:
            print(f"{unique_sources[i]:4d} {x_source:8.1f} {min_offset:8.2f} {estimated_vel:8.1f} {status:^30} {velocity_per_source[i]:8.1f}")
    
    # Compute average of estimated velocities as backup
    if len(estimated_velocities) > 0:
        avg_estimated_velocity = np.mean(estimated_velocities)
        
        if debug:
            print("-"*80)
            print(f"\nEstimated surface velocities from {len(estimated_velocities)}/{len(unique_sources)} shots")
            print(f"  Range: {np.min(estimated_velocities):.1f} - {np.max(estimated_velocities):.1f} m/s")
            print(f"  Average: {avg_estimated_velocity:.1f} m/s")
            print(f"\nAssigning average velocity to shots without reliable estimates:")
            print("-"*80)
        else:
            print(f"Estimated surface velocities from {len(estimated_velocities)}/{len(unique_sources)} shots")
            print(f"  Range: {np.min(estimated_velocities):.1f} - {np.max(estimated_velocities):.1f} m/s")
            print(f"  Average: {avg_estimated_velocity:.1f} m/s")
        
        # Use average for shots without reliable estimates
        for i in range(len(unique_sources)):
            if velocity_source_type[i] == 'pending':
                old_vel = velocity_per_source[i]
                velocity_per_source[i] = avg_estimated_velocity
                velocity_source_type[i] = 'average'
                if debug:
                    x_source = sensor_positions[unique_sources[i], 0]
                    print(f"{unique_sources[i]:4d} {x_source:8.1f} → Using AVERAGE: {old_vel:.1f} → {avg_estimated_velocity:.1f} m/s")
        
        if debug:
            print("="*80 + "\n")
    else:
        print(f"Warning: Could not reliably estimate surface velocity from picks")
        print(f"  Using provided min_velocity: {np.mean(min_velocity):.1f} m/s")
        if debug:
            print("="*80 + "\n")

    no_correc = 0
    tneg_correc = 0
    tneg_offset_correc = 0
    tpos_correc = 0
    tpos_offset_correc = 0
    
    t_pick_original = np.copy(refrac_manager.data['t'])

    if debug:
        print("\n" + "="*80)
        print("TIME CORRECTION DEBUG")
        print("="*80)
        print(f"{'Shot':>4} {'X_pos':>8} {'MinOff':>7} {'Vel':>7} {'MinT_orig':>10} {'MinT_corr':>10} {'Correction':^20}")
        print("-"*80)

    # Second pass: apply corrections using computed velocities
    for i in range(len(unique_sources)):
        # Get indices for picks from this source
        shot_indices = np.where(refrac_manager.data['s'] == unique_sources[i])[0]
        
        x_source = sensor_positions[unique_sources[i], 0]
        geophones = refrac_manager.data['g'][shot_indices]
        x_pick = sensor_positions[geophones, 0]
        t_pick = refrac_manager.data['t'][shot_indices]
        offset = np.abs(x_pick - x_source)
        min_picked_offset = np.min(offset)
        min_picked_time = np.min(t_pick)
        t_offset = np.min(offset) / velocity_per_source[i]

        if offset_max is None:
            max_picked_offset = np.max(offset)
        else:
            max_picked_offset = offset_max
        
        correction_type = ""
            
        # Correct negative and 0 time picks
        if min_picked_time <= 0:
            t_pick_corr = t_pick - min_picked_time
            
            # Correct time picks when closest picked geophone is not at the source location
            if min_picked_offset > 0:
                t_pick_corrected = t_pick_corr + t_offset
                tneg_offset_correc += 1
                correction_type = "neg+offset"
            # No further correction if closest geophone is at the source location
            else:
                t_pick_corrected = t_pick_corr
                tneg_correc += 1
                correction_type = "neg"
        # Correct positive time picks
        else:
            if apply_positive_correction:
                if min_picked_offset == 0:
                    # Correct positive time picks when closest picked geophone is at the source location
                    if min_picked_time > 0:
                        t_pick_corrected = t_pick - min_picked_time
                        tpos_correc += 1
                        correction_type = "pos"
                    else:
                        t_pick_corrected = t_pick
                        no_correc += 1
                        correction_type = "none"
                # Correct positive time picks when closest picked geophone is within offset range
                elif min_picked_offset > 0 and min_picked_offset <= max_picked_offset:
                    t_pick_corrected = t_pick + (t_offset - min_picked_time)
                    tpos_offset_correc += 1
                    correction_type = "pos+offset"
                else:
                    t_pick_corrected = t_pick
                    no_correc += 1
                    correction_type = "none"
            else:
                t_pick_corrected = t_pick
                no_correc += 1
                correction_type = "none"

        if debug:
            min_t_corr = np.min(t_pick_corrected)
            print(f"{unique_sources[i]:4d} {x_source:8.1f} {min_picked_offset:7.2f} {velocity_per_source[i]:7.1f} "
                  f"{min_picked_time*1000:10.3f} {min_t_corr*1000:10.3f} {correction_type:^20}")

        # Update the data in-place using the correct indices
        refrac_manager.data['t'][shot_indices] = t_pick_corrected

    if debug:
        print("="*80)
        # Show statistics of time changes
        time_changes = refrac_manager.data['t'] - t_pick_original
        print(f"\nTime change statistics (ms):")
        print(f"  Min change: {np.min(time_changes)*1000:8.3f}")
        print(f"  Max change: {np.max(time_changes)*1000:8.3f}")
        print(f"  Mean change: {np.mean(time_changes)*1000:8.3f}")
        print(f"  Changed picks: {np.sum(time_changes != 0)}/{len(time_changes)}")
        print("="*80 + "\n")
        
    print(f"Time correction summary:")
    print(f"  No correction: {no_correc}")
    print(f"  Negative time correction: {tneg_correc}")
    print(f"  Negative time with offset correction: {tneg_offset_correc}")
    print(f"  Positive time correction: {tpos_correc}")
    print(f"  Positive time with offset correction: {tpos_offset_correc}")
    
    # Remove picks with t=0 (at shot positions) as PyGIMLI can't handle them
    # Also remove any remaining negative times
    t_threshold = 1e-6  # Effectively 0
    
    original_count = len(refrac_manager.data['t'])
    invalid_picks = refrac_manager.data['t'] <= t_threshold
    
    if np.any(invalid_picks):
        removed_count = np.sum(invalid_picks)
        print(f"\nRemoving {removed_count} picks with t <= {t_threshold*1000:.3f} ms (at shot positions)")
        
        # Mark invalid picks
        refrac_manager.data.markInvalid(invalid_picks)
        
        # Remove invalid data
        refrac_manager.data.removeInvalid()
        
        print(f"Remaining picks: {len(refrac_manager.data['t'])}/{original_count}")
    
def get_surf_velocity(refrac_manager, smooth=True, w=9, 
                      method='convolve',s=50, show_plot=False):

    sensor_positions = refrac_manager.data.sensors().array()
    unique_sources = np.unique(refrac_manager.data['s'])
    x_source_all = sensor_positions[unique_sources,0]
    
    x_surf_array = []
    vel_surf_array = []
    min_offset_array = []
    max_offset_array = []

    for i in range(len(unique_sources)):

        source_index = unique_sources[i]
        x_source = sensor_positions[source_index,0]
        geophone_index = refrac_manager.data['g'][refrac_manager.data['s']==source_index]
        x_pick = sensor_positions[geophone_index,0]
        t_pick = refrac_manager.data['t'][refrac_manager.data['s']==source_index]
        offset = np.round(x_pick - x_source,3)

        # plt.plot(x_source,0,'r*')
        # plt.plot(x_pick,t_pick,'b+')

        offset_left = offset[offset<0]
        offset_right = offset[offset>0]
        offset_zero = offset[offset==0]

        left_ok = False
        right_ok = False
        vel_left = []
        vel_right = []
        min_offset_left = []
        max_offset_left = []
        min_offset_right = []
        max_offset_right = []

        if len(offset_zero)>0: # If closest geophone is at the source location
            closest_geophone_left_index = geophone_index[np.where(offset==np.sort(offset_zero)[0])]
            closest_geophone_right_index = geophone_index[np.where(offset==np.sort(offset_zero)[0])]

            if len(offset_left)>0:
                second_geophone_left_index = geophone_index[np.where(offset==np.max(offset_left))]
                left_ok = True
                # print(offset_left)
                min_offset_left.append(0)
                max_offset_left.append(np.abs(np.sort(offset_left)[-1]))
            else:
                second_geophone_left_index = []
            
            if len(offset_right)>0:
                second_geophone_right_index = geophone_index[np.where(offset==np.min(offset_right))]
                right_ok = True
                # print(offset_right)
                min_offset_right.append(0)
                max_offset_right.append(np.abs(np.sort(offset_right)[0]))
            else:  
                second_geophone_right_index = []
                
        else: # If closest geophone is not at the source location
            if len(offset_left)>1:
                closest_geophone_left_index = geophone_index[np.where(offset==np.sort(offset_left)[-1])]
                second_geophone_left_index = geophone_index[np.where(offset==np.sort(offset_left)[-2])]
                left_ok = True
                min_offset_left.append(np.abs(np.sort(offset_left)[-1]))
                max_offset_left.append(np.abs(np.sort(offset_left)[-2]))

            else:
                closest_geophone_left_index = []
                second_geophone_left_index = []
            
            if len(offset_right)>1:
                closest_geophone_right_index = geophone_index[np.where(offset==np.sort(offset_right)[0])]
                second_geophone_right_index = geophone_index[np.where(offset==np.sort(offset_right)[1])]
                right_ok = True
                min_offset_right.append(np.abs(np.sort(offset_right)[0]))
                max_offset_right.append(np.abs(np.sort(offset_right)[1]))

            else:  
                closest_geophone_right_index = []
                second_geophone_right_index = []

        if left_ok:
            x1 = sensor_positions[closest_geophone_left_index,0]
            x2 = sensor_positions[second_geophone_left_index,0]
            indices = np.where(np.isin(geophone_index, closest_geophone_left_index))
            indices_list_closest = indices[0].tolist()
            y1 = np.max(t_pick[indices_list_closest])
            indices = np.where(np.isin(geophone_index, second_geophone_left_index))
            indices_list_second = indices[0].tolist()
            y2 = np.max(t_pick[indices_list_second])

            if y2 == y1:
                left_ok = False
            else:
                vel_left = np.abs((x2-x1)/(y2-y1))

            # plt.plot(sensor_positions[closest_geophone_left_index,0],t_pick[indices_list_closest],'g+')
            # plt.plot(sensor_positions[second_geophone_left_index,0],t_pick[indices_list_second],'g+')
                
        if right_ok:
            x1 = sensor_positions[closest_geophone_right_index,0]
            x2 = sensor_positions[second_geophone_right_index,0]
            indices = np.where(np.isin(geophone_index, closest_geophone_right_index))
            indices_list_closest = indices[0].tolist()
            y1 = np.max(t_pick[indices_list_closest])
            indices = np.where(np.isin(geophone_index, second_geophone_right_index))
            indices_list_second = indices[0].tolist()
            y2 = np.max(t_pick[indices_list_second])

            if y2 == y1:
                right_ok = False
            else:
                vel_right = np.abs((x2-x1)/(y2-y1))

            # plt.plot(sensor_positions[closest_geophone_right_index,0],t_pick[indices_list_closest],'y+')
            # plt.plot(sensor_positions[second_geophone_right_index,0],t_pick[indices_list_second],'y+')

        if left_ok and right_ok:
            vel = np.mean([vel_left,vel_right])
            min_offset = np.mean([min_offset_left,min_offset_right])
            max_offset = np.mean([max_offset_left,max_offset_right])
        elif left_ok:
            vel = vel_left[0]
            min_offset = min_offset_left[0]
            max_offset = max_offset_left[0]
        elif right_ok:
            vel = vel_right[0]
            min_offset = min_offset_right[0]
            max_offset = max_offset_right[0]

        x_surf_array.append(x_source)
        vel_surf_array.append(vel)
        min_offset_array.append(min_offset)
        max_offset_array.append(max_offset)

        # if left_ok:
        #     plt.plot((x_source+offset_left),(offset_left)/(-vel))
        # if right_ok:
        #     plt.plot((x_source+offset_right),(offset_right)/vel)

    # plt.show()

    if smooth:
        smooth_vel_surf_array = smooth_interpolation(x_source_all,x_surf_array,vel_surf_array,
                                                 size=w,method=method,s=s)
    else:
        smooth_vel_surf_array = vel_surf_array
    
    if show_plot:
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots()
        ax.plot(x_surf_array,vel_surf_array,'k')
        ax.plot(x_source_all,smooth_vel_surf_array,'b')
        ax.set_xlabel('X source (m)')
        ax.set_ylabel('Surface velocity (m/s)')
        plt.show()

    return np.column_stack((x_surf_array, smooth_vel_surf_array, 
                            vel_surf_array, min_offset_array, max_offset_array))

def spline_interp(x_interp,x,y,s=1):
    from scipy.interpolate import UnivariateSpline
    spline = UnivariateSpline(x, y, s=s)
    return spline(x_interp)

def median_filter(x, size):
    '''Median filter of x with window w.'''
    from scipy.ndimage import median_filter
    return median_filter(x,size=size)

def moving_average(x, size=5, method='convolve'):
    '''Moving average of x with window size. Method can be 'convolve' or 'sagvol'''
    from scipy.signal import savgol_filter
    from scipy.ndimage import uniform_filter1d

    if method == 'convolve':
        return uniform_filter1d(x,size=size)
    elif method == 'sagvol':
        return savgol_filter(x, size, 2)

def smooth_interpolation(x_interp,x,y,size=9,method='convolve',s=None):
    '''Smooth interpolation of x and y with window size. Method can be 'convolve' or 'sagvol'''
    y_interp = spline_interp(x_interp,x,y,s=s)
    y_interp = median_filter(y_interp,size)
    y_interp = moving_average(y_interp,size,method=method)
    return y_interp