#!/bin/env python3

import pandas as pd
import matplotlib.pyplot as plt # Import matplotlib
import re
import argparse # Import the argparse module
import os # Import os module for path manipulation

# Set up argument parsing
parser = argparse.ArgumentParser(description='Plot monitor_mode_test.py results.')
parser.add_argument(
    'log_file',
    type=str,
    help='Path to the log file containing mode data.'
    )
parser.add_argument(
    '--show',
    action=argparse.BooleanOptionalAction,
    default=True,
    help='Display the plot window.'
    )
parser.add_argument(
    '--plot-duration',
    action=argparse.BooleanOptionalAction,
    default=False,
    help='Include the duration subplot.'
    )
args = parser.parse_args()
base_name = os.path.splitext(os.path.basename(args.log_file))[0]

# Read log data from the specified file
try:
    with open(args.log_file, 'r') as f:
        log_data = f.read()
except FileNotFoundError:
    print(f"Error: The file '{args.log_file}' was not found.")
    exit()
except Exception as e:
    print(f"Error reading file '{args.log_file}': {e}")
    exit()

sleep_time = 1 # Default value if not found
data = []
gps_seconds = None
last_gps_seconds = 0
for line in log_data.splitlines():
    mode = None # Initialize mode for each line

    if "404 Client Error: File not found for url" in line:
        mode = 1 # Monitor
    elif "HTTPConnectionPool" in line and "Read timed out." in line:
        mode = 2 # Failed Conn
    elif "Error fetching data: ('Connection aborted.', ConnectionResetError(104, 'Connection reset by peer'))" in line:
        mode = 2 # Failed Conn
    elif "Temperature temp=" in line:
        mode = 0 # Normal
    elif line.startswith("fname_log: ") or \
         line.startswith("ip: ") or \
         line.startswith("port: ") or \
         line.startswith("quiet_mode: "):
        continue # Skip header lines
    elif line.startswith("sleep_time: "):
        sleep_time_match = re.match(r"sleep_time: (\d+)", line)
        if sleep_time_match:
            sleep_time = int(sleep_time_match.group(1))
        continue
    else:
        # Unknown line type, assign a mode if the line isn't just whitespace
        if line.strip():
            mode = 3 # Unknown
        else:
            continue

    # Extract gps_seconds if present in the line
    gps_seconds_match = re.search(r"gps_seconds:(\d+)", line)
    if gps_seconds_match:
        gps_seconds = int(gps_seconds_match.group(1))
    else:
        if last_gps_seconds is not None:
            gps_seconds = last_gps_seconds + sleep_time
        else:
            # If it's the very first line and no gps_seconds, we can't determine it.
            pass

    if mode is not None and gps_seconds is not None:
        data.append({"gps_seconds": gps_seconds, "mode": mode})
        last_gps_seconds = gps_seconds # Update last_gps_seconds for the next iteration

# Create the DataFrame for raw log data (now named df)
df = pd.DataFrame(data)
print(f'{sleep_time=}')

# Define mode_map for both raw data plotting and duration calculation
# This map is used for y-ticks on the first plot and for simplifying modes for duration
mode_map = {0: 'Normal', 1: 'Monitor/Failed', 2: 'Monitor/Failed', 3: 'Unknown'}

if args.plot_duration and not df.empty:
    state_durations = [] # Temporary list to build the new df_state_durations

    # Add simplified_mode to df for duration calculation
    df['simplified_mode'] = df['mode'].map(mode_map)

    current_state_start_time = df.iloc[0]['gps_seconds']
    current_state_mode = df.iloc[0]['simplified_mode']

    for i in range(1, len(df)):
        if df.iloc[i]['simplified_mode'] != current_state_mode:
            # State change detected
            duration = df.iloc[i]['gps_seconds'] - current_state_start_time
            state_durations.append({
                'gps_seconds': current_state_start_time, # Start time of the state
                'duration_sec': duration,
                'mode': current_state_mode
            })
            # Reset for the new state
            current_state_start_time = df.iloc[i]['gps_seconds']
            current_state_mode = df.iloc[i]['simplified_mode']

    # Add the last state's duration (from its start to the end of the log)
    if len(df) > 0:
        last_duration = df.iloc[-1]['gps_seconds'] - current_state_start_time
        state_durations.append({
            'gps_seconds': current_state_start_time,
            'duration_sec': last_duration,
            'mode': current_state_mode
        })

    # Assign the summarized state duration data to 'df_state_durations'
    df_state_durations = pd.DataFrame(state_durations)
    print("\ndf_state_durations:")
    print(df_state_durations)
else:
    # If plot_duration is False, df_state_durations will be an empty DataFrame
    df_state_durations = pd.DataFrame()

# Save the raw data ('df') to a CSV file as it was before the last change
output_csv_file = f"{base_name}.csv"
print(f'Saving {output_csv_file }')
df.to_csv(output_csv_file, index=False)

# Create a figure with subplots based on whether the duration plot is requested
if args.plot_duration:
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(20, 12), sharex=True) # sharex=True aligns x-axes
else:
    fig, ax1 = plt.subplots(1, 1, figsize=(20, 6)) # Single plot if no duration subplot

# Plot 1: Mode vs GPS Seconds (uses df, the raw log data)
ax1.scatter(df['gps_seconds'], df['mode'], marker='+')
ax1.set_ylabel('Mode')
ax1.set_title('Mode vs GPS Seconds (Raw Log Data)')
ax1.set_yticks(list(mode_map.keys()), list(mode_map.values()))
ax1.grid(True)

# Plot 2: Duration vs GPS Seconds (Scatter plot) - uses df_state_durations
if args.plot_duration:
    # Get unique modes and assign colors from df_state_durations
    unique_modes = df_state_durations['mode'].unique()
    colors = plt.cm.get_cmap('tab10', len(unique_modes)) # Using a colormap for distinct colors

    for i, mode_name in enumerate(unique_modes):
        subset = df_state_durations[df_state_durations['mode'] == mode_name]
        ax2.scatter(subset['gps_seconds'], subset['duration_sec'],
                    marker='o',
                    color=colors(i),
                    label=mode_name)

    ax2.set_xlabel('GPS Seconds (State Start Time)')
    ax2.set_ylabel('Duration (Seconds)')
    ax2.set_title('State Duration vs GPS Seconds (Summarized Data)')
    ax2.grid(True)
    ax2.legend(title='Mode') # Add a legend for the modes

plt.tight_layout() # Adjust layout to prevent overlapping titles/labels

# Derive output PNG filename
output_png_file = f"{base_name}.png"

# Save the chart as a PNG file
print(f'Saving {output_png_file }')
plt.savefig(output_png_file)

# Display the plot
if args.show:
    plt.show()
