import os

# Create the 'output' folder if it doesn't exist
if not os.path.exists('output'):
    os.makedirs('output')

# Import libraries
import numpy as np
import pandas as pd
from obspy import read
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import os
cat_directory = 'lunar/training/catalogs/'
cat_file = cat_directory + 'apollo12_catalog_GradeA_final.csv'
cat = pd.read_csv(cat_file)
for cat_len in range(len(cat)):
    row = cat.iloc[cat_len]
    arrival_time = datetime.strptime(row['time_abs(%Y-%m-%dT%H:%M:%S.%f)'],'%Y-%m-%dT%H:%M:%S.%f')
    # If we want the value of relative time, we don't need to use datetime
    arrival_time_rel = row['time_rel(sec)']
    # Let's also get the name of the file
    test_filename = row.filename
    data_directory = 'lunar/training/data/S12_GradeA/'
    csv_file = f'{data_directory}{test_filename}.csv'
    #csv_file = '/content/drive/MyDrive/lunar/test/data/S15_GradeA/xa.s15.00.mhz.1973-04-04HR00_evid00098.csv'
    data_cat = pd.read_csv(csv_file)
    # Read a miniseed file
    data_directory = 'lunar/training/data/S12_GradeA/'
    mseed_file = f'{data_directory}{test_filename}.mseed'
    #mseed_file = '/content/drive/MyDrive/lunar/test/data/S15_GradeA/xa.s15.00.mhz.1973-04-04HR00_evid00098.mseed'
    st = read(mseed_file)
    print(st)
    # The stream file also contains some useful header information
    print(st[0].stats)

    tr = st.traces[0].copy()
    tr_times = tr.times()
    tr_data = tr.data

    # Apply a bandpass filter
    st_filtered = st.copy()
    st_filtered.filter("bandpass", freqmin=0.1, freqmax=10.0)
    #st_filtered = st_p.copy()
    # Access the first trace in the filtered stream
    tr_filtered = st_filtered[0]

    # Get the filtered data (dependent variable)
    filtered_data = tr_filtered.data

    sq_filtered_data = filtered_data*filtered_data
    plt.plot(tr_times, sq_filtered_data)
    plt.xlabel('Time (s)')
    plt.ylabel('Velocity (m/s)')
    plt.title('Velocity Data with Detected Peaks')
    plt.legend()
    plt.show()

    # LEAST Squares ----------------------------------------------------------------
    # Manually construct the design matrix X for a quadratic fit
    # X = [[t1^2, t1, 1], [t2^2, t2, 1], ..., [tn^2, tn, 1]]
    X = np.vstack([tr_times**2, tr_times, np.ones(len(tr_times))]).T

    # Compute the Least Squares solution: p = (X^T X)^(-1) X^T y
    # p will contain the coefficients [a, b, c] where y = a * t^2 + b * t + c
    p = np.linalg.inv(X.T @ X) @ X.T @ sq_filtered_data

    # Extract the coefficients
    a, b, c = p

    # Generate the fitted signal for the quadratic model
    fitted_signal = a * tr_times**2 + b * tr_times + c

    # Calculate residuals
    residuals = sq_filtered_data - (fitted_signal**2)
    # Set a threshold to detect seismic events (e.g., 3 times the standard deviation)
    threshold = 25 * np.std(residuals)

    # Identify points where residuals exceed the threshold (potential events)
    event_indices = np.where(np.abs(residuals) > threshold)[0]

    # Convert indices to relative times
    event_times = tr_times[event_indices]
    event_residuals = residuals[event_indices]
    # Print detected event times
    print("Detected event times (seconds):", event_times[11])
    print("Detected event residuals:", event_residuals[11])
    # Plot the filtered data, fitted signal, and detected events
    plt.figure(figsize=(16, 6))

    # Plot filtered data
    plt.plot(tr_times, sq_filtered_data, label='Filtered Seismic Data (0.1 - 5.0 Hz)', color='black')

    # Plot fitted signal (Least Squares Fit)
    plt.plot(tr_times, (fitted_signal**2), label='Fitted Signal (Manual Least Squares)', color='blue')

    # Mark detected events
    plt.scatter(event_times, sq_filtered_data[event_indices], color='red', label='Detected Events')

    # Add labels and title
    plt.xlabel('Time (seconds since start)')
    plt.ylabel('Amplitude')
    plt.title('Seismic Event Detection using Manual Least Squares Method')
    plt.legend()

    # Show plot
    plt.show()

    #Least Square -------------------------------------------------------------------

    #Moving Average -------------------------------------------------------------------------
    # Assuming tr_times is in seconds and sq_filtered_data contains the seismic data

    # Step 1: Define window sizes in terms of data points
    # To get the number of points corresponding to 10 minutes and 1 hour,
    # we calculate based on the time differences in tr_times (assuming uniform sampling).

    sampling_interval = np.mean(np.diff(tr_times))  # Calculate the average time step (seconds)

    # Calculate window sizes based on sampling interval
    window_10min = int(600 / sampling_interval)  # 600 seconds = 10 minutes
    window_1hr = int(3600 / sampling_interval)   # 3600 seconds = 1 hour

    # Step 2: Calculate moving averages using a rolling window
    # Use np.convolve to calculate the moving averages for both window sizes

    # Moving average for 10 minutes
    moving_avg_10min = np.convolve(sq_filtered_data, np.ones(window_10min)/window_10min, mode='valid')

    # Moving average for 1 hour
    moving_avg_1hr = np.convolve(sq_filtered_data, np.ones(window_1hr)/window_1hr, mode='valid')

    # Step 3: Plot the original data and the moving averages

    # Plot the original seismic data
    plt.figure(figsize=(16, 6))
    plt.plot(tr_times, sq_filtered_data, label='Original Seismic Data', color='black')

    # Plot the 10-minute moving average (aligned with time)
    plt.plot(tr_times[window_10min-1:], moving_avg_10min, label='10-min Moving Average', color='red')

    # Plot the 1-hour moving average (aligned with time)
    plt.plot(tr_times[window_1hr-1:], moving_avg_1hr, label='1-hr Moving Average', color='blue')

    # Add labels and title
    plt.xlabel('Time (seconds since start)')
    plt.ylabel('Amplitude')
    plt.title('Seismic Data with 10-min and 1-hr Moving Averages')

    # Add legend
    plt.legend()

    # Show plot
    plt.show()

    #plotting highest point in moving average
    # Step 1: Define window size for 10 minutes in terms of data points
    window_10min = int(600 / sampling_interval)  # 600 seconds = 10 minutes

    # Step 2: Calculate the 10-minute moving average
    moving_avg_10min = np.convolve(sq_filtered_data, np.ones(window_10min)/window_10min, mode='valid')

    # Step 3: Find the index of the maximum point in the moving average
    max_index_in_avg = np.argmax(moving_avg_10min)  # Index within the moving average

    # Since the moving average is shorter than the original data (due to 'valid' mode),
    # we need to shift the index to match the original time series.
    max_index_in_original = max_index_in_avg + (window_10min - 1)

    # Get the corresponding time for this index
    max_time = tr_times[max_index_in_original]

    # Print the index and time of the highest point in the 10-min moving average
    print("Index of highest point in 10-min moving average:", max_index_in_original)
    print("Time of highest point in 10-min moving average (seconds):", max_time)

    # Step 4: Plot the original data and the 10-minute moving average
    plt.figure(figsize=(16, 6))

    # Plot the original seismic data
    plt.plot(tr_times, sq_filtered_data, label='Original Seismic Data', color='black')

    # Plot the 10-minute moving average
    plt.plot(tr_times[window_10min-1:], moving_avg_10min, label='10-min Moving Average', color='red')

    # Mark the highest point in the moving average
    plt.scatter(tr_times[max_index_in_original], moving_avg_10min[max_index_in_avg], color='green', label='Max Point in 10-min Avg', zorder=5)

    # Add labels and title
    plt.xlabel('Time (seconds since start)')
    plt.ylabel('Amplitude')
    plt.title('Seismic Data with 10-min Moving Average and Max Point')

    # Add legend
    plt.legend()

    # Show plot
    plt.show()

    # plotting highest of data -----------------------------------------------

    # Step 1: Define window size for 10 minutes in terms of data points (from the previous code)
    window_10min = int(600 / sampling_interval)

    # Calculate the 10-minute moving average
    moving_avg_10min = np.convolve(sq_filtered_data, np.ones(window_10min) / window_10min, mode='valid')

    # Calculate slopes to detect sharp upward slope
    slopes = np.diff(moving_avg_10min)
    slope_threshold = 1.5 * np.std(slopes)

    # Find sharp upward slope index (from previous code)
    sharp_slope_index_in_avg = max_index_in_avg
    for i in range(max_index_in_avg - 1, 0, -1):
      if slopes[i-1] > slope_threshold:
          sharp_slope_index_in_avg = i
          break

    # Convert to the original index
    sharp_slope_index_in_original = sharp_slope_index_in_avg + (window_10min - 1)

    # Step 2: Find the highest point within the -5000 index range
    # Define the range to search for the highest point
    start_index_range = max(sharp_slope_index_in_original - 5000, 0)  # Ensure we don't go below index 0
    end_index_range = sharp_slope_index_in_original

    # Extract the portion of the original data in this range
    data_in_range = sq_filtered_data[start_index_range:end_index_range]

    # Find the index of the highest point in the range
    max_in_range_index = np.argmax(data_in_range)

    # Get the actual index of the highest point in the original data
    highest_point_index = start_index_range + max_in_range_index

    # Get the time and amplitude of the highest point
    highest_point_time = tr_times[highest_point_index]
    highest_point_amplitude = sq_filtered_data[highest_point_index]

    # Print the details of the highest point
    print("Highest point in the -5000 index range:")
    print("Index:", highest_point_index)
    print("Time (seconds):", highest_point_time)
    print("Amplitude:", highest_point_amplitude)

    # Step 3: Plot the seismic data and highlight the highest point
    plt.figure(figsize=(16, 6))

    # Plot the original seismic data
    plt.plot(tr_times, sq_filtered_data, label='Original Seismic Data', color='black')

    # Plot the 10-minute moving average
    plt.plot(tr_times[window_10min-1:], moving_avg_10min, label='10-min Moving Average', color='red')

    # Mark the point with sharp upward slope
    plt.scatter(tr_times[sharp_slope_index_in_original], moving_avg_10min[sharp_slope_index_in_avg], color='green', label='Sharp Upward Slope', zorder=5)

    # Mark the highest point in the -5000 range
    plt.scatter(highest_point_time, highest_point_amplitude, color='blue', label='Highest Point in -5000 Range', zorder=6)

    # Add labels and title
    plt.xlabel('Time (seconds since start)')
    plt.ylabel('Amplitude')
    plt.title('Seismic Data with Detected Points')

    # Add legend
    plt.legend()

    # Show plot
    plt.show()

    #Backwards Standard Deviation-----------------------------------
    # Set a threshold to detect seismic events (e.g., 3 times the standard deviation)
    threshold = 0.1*np.std(residuals)

    # Identify points where residuals exceed the threshold (potential events)
    event_indices = np.where(np.abs(residuals) > threshold)[0]

    # Convert indices to relative times
    event_times = tr_times[event_indices]
    event_residuals = residuals[event_indices]
    # Step 2: Find the maximum detected event point in terms of index
    max_event_index = highest_point_index  # The last event detected

    # Step 3: Define the range to plot, going backward 5000 indices from the max event point
    start_index_plot = max(0, max_event_index - 5000)  # Ensure index doesn't go below 0
    end_index_plot = max_event_index

    # Extract the event indices within this range
    event_indices_in_range = event_indices[(event_indices >= start_index_plot) & (event_indices <= end_index_plot)]

    # Get the corresponding times and residuals for events in this range
    event_times_in_range = tr_times[event_indices_in_range]
    event_residuals_in_range = residuals[event_indices_in_range]
    print("orginal time: ",arrival_time_rel)

    # Print the last detected event time and residual within the range
    if len(event_times_in_range) > 0:
      print("Detected event time (seconds):", event_times_in_range[1])
      print("Detected event residual:", event_residuals_in_range[1])
    else:
      print("No events detected in the -5000 index range.")

    # Step 4: Plot the filtered data, fitted signal, and only the detected events in the backward range
    plt.figure(figsize=(16, 6))

    # Plot the filtered seismic data
    plt.plot(tr_times, sq_filtered_data, label='Filtered Seismic Data (0.1 - 5.0 Hz)', color='black')

    # Plot the fitted signal (Least Squares Fit)
    plt.plot(tr_times, (fitted_signal**2), label='Fitted Signal (Manual Least Squares)', color='blue')

    # Mark the detected events only within the -5000 index range
    #plt.scatter(event_times_in_range[1], sq_filtered_data[1], color='red', label='Detected Events (-5000 range)', zorder=5)
    plt.axvline(x=event_times_in_range[1], color='red', linestyle='--', label='Event Time (-5000 range)')
    # Add labels and title
    plt.xlabel('Time (seconds since start)')
    plt.ylabel('Amplitude')
    plt.title('Seismic Event Detection (Backward 5000 indices from last event)')

    # Add legend
    plt.legend()
    # Save the plot to the output folder
    plt.savefig(f'output/{cat_len}.png')
    # Show plot
    plt.show()

C:\Users\tanuj\anaconda3\lib\site-packages\obspy\signal\filter.py:62: UserWarning: Selected high corner frequency (10.0) of bandpass is at or above Nyquist (3.3125). Applying a high-pass instead.
  warnings.warn(msg)
No artists with labels found to put in legend.  Note that artists whose label start with an underscore are ignored when legend() is called with no argument.

1 Trace(s) in Stream:
XA.S12.00.MHZ | 1970-01-19T00:00:00.665000Z - 1970-01-20T00:00:02.778208Z | 6.6 Hz, 572415 samples
         network: XA
         station: S12
        location: 00
         channel: MHZ
       starttime: 1970-01-19T00:00:00.665000Z
         endtime: 1970-01-20T00:00:02.778208Z
   sampling_rate: 6.625
           delta: 0.1509433962264151
            npts: 572415
           calib: 1.0
         _format: MSEED
           mseed: AttribDict({'dataquality': 'D', 'number_of_records': 1136, 'encoding': 'FLOAT64', 'byteorder': '>', 'record_length': 4096, 'filesize': 4653056})

Detected event times (seconds): 73769.05660377358
Detected event residuals: 3.294661726926621e-17

Index of highest point in 10-min moving average: 492039
Time of highest point in 10-min moving average (seconds): 74270.03773584905

Highest point in the -5000 index range:
Index: 489650
Time (seconds): 73909.43396226416
Amplitude: 6.868709837768176e-17

orginal time:  73500.0
Detected event time (seconds): 73430.64150943396
Detected event residual: 1.9808949195720213e-19

C:\Users\tanuj\anaconda3\lib\site-packages\obspy\signal\filter.py:62: UserWarning: Selected high corner frequency (10.0) of bandpass is at or above Nyquist (3.3125). Applying a high-pass instead.
  warnings.warn(msg)
No artists with labels found to put in legend.  Note that artists whose label start with an underscore are ignored when legend() is called with no argument.

1 Trace(s) in Stream:
XA.S12.00.MHZ | 1970-03-25T00:00:00.440000Z - 1970-03-26T00:00:01.949434Z | 6.6 Hz, 572411 samples
         network: XA
         station: S12
        location: 00
         channel: MHZ
       starttime: 1970-03-25T00:00:00.440000Z
         endtime: 1970-03-26T00:00:01.949434Z
   sampling_rate: 6.625
           delta: 0.1509433962264151
            npts: 572411
           calib: 1.0
         _format: MSEED
           mseed: AttribDict({'dataquality': 'D', 'number_of_records': 1136, 'encoding': 'FLOAT64', 'byteorder': '>', 'record_length': 4096, 'filesize': 4653056})

Detected event times (seconds): 13332.981132075472
Detected event residuals: 9.888539495735593e-18

Index of highest point in 10-min moving average: 91298
Time of highest point in 10-min moving average (seconds): 13780.830188679245