# G-Force Meter Code v2.2 - Sustained Max G Logic
# IMU (QMI8658) on I2C1 (Adapter Socket: Pins 6,7)
# OLED (SSD1327) on I2C0 (Direct Wire: Pins 4,5)

import machine
import time
import math # Needed for abs(), copysign, pow, ticks_ms, ticks_diff
import ucollections # For deque

# --- Library Imports ---
try:
    import qmi8658; imu_lib_found = True; print("Imported qmi8658")
except ImportError: print("Error: qmi8658.py missing"); imu_lib_found = False; qmi8658 = None
try:
    from ssd1327 import SSD1327_I2C; oled_lib_found = True; print("Imported SSD1327_I2C")
except ImportError: print("Error: ssd1327.py missing"); oled_lib_found = False; SSD1327_I2C = None

# --- Constants ---
IMU_ADDR = 0x6b; OLED_ADDR = 0x3d; OLED_WIDTH = 128; OLED_HEIGHT = 128
FONT_WIDTH = 8; FONT_HEIGHT = 8

# --- Display Configuration ---
TITLE_Y_LAT = 5; MAX_VAL_Y_LAT = 18; BAR_Y_LAT = 38; CURRENT_VAL_Y_LAT = 51
TITLE_Y_ACCEL = 69; MAX_VAL_Y_ACCEL = 82; BAR_Y_ACCEL = 102; CURRENT_VAL_Y_ACCEL = 115
BAR_HALF_WIDTH = 60; BAR_THICKNESS = 8; MAX_G_DISPLAY = 1.0
NON_LINEAR_EXPONENT = 1.5

# --- Filter/Logic Configuration ---
G_DEADZONE_THRESHOLD = 0.05 # For display zeroing
DEADZONE_FILTER_MS = 500    # Time before display zeros out

SUSTAINED_G_DURATION_MS = 500 # Required duration for max G update
SUSTAINED_G_WINDOW = 0.15     # +/- G window for sustained check
SUSTAINED_G_MIN_TRIGGER = 0.2 # Minimum G to start checking for sustained max

LOOP_DELAY_MS = 50            # Target loop time (~20 FPS)
HISTORY_LENGTH = int(SUSTAINED_G_DURATION_MS / LOOP_DELAY_MS) + 2 # Store > 0.5s worth

# Colors (Grayscale 0-15)
GRAY_LEVEL_BG = 0; GRAY_LEVEL_BAR_BG = 4; GRAY_LEVEL_INDICATOR = 15
GRAY_LEVEL_TEXT = 12; GRAY_LEVEL_MAX_VAL = 15; GRAY_LEVEL_CURRENT_VAL = 15

# --- Calibration Offsets ---
ACCEL_X_OFFSET = 0.0; ACCEL_Y_OFFSET = 0.0

# --- Axis Mapping ---
FORWARD_BACKWARD_AXIS = 0; LATERAL_AXIS = 1

# --- High Score Tracking ---
max_lateral_g = 0.0
max_braking_g = 0.0

# --- Filter State Variables ---
g_force_lateral_display = 0.0; g_force_fwd_bwd_display = 0.0
lat_below_thresh_start_time = None; acc_below_thresh_start_time = None

# --- Sustained Max State Variables ---
# Lateral
checking_sustained_lat = False
candidate_lat_start_time = None
candidate_lat_center_g = 0.0
recent_lateral_g = ucollections.deque((), HISTORY_LENGTH) # Store recent values
# Braking (use fwd/bwd axis but only check negative Gs)
checking_sustained_brk = False
candidate_brk_start_time = None
candidate_brk_center_g = 0.0
recent_fwd_bwd_g = ucollections.deque((), HISTORY_LENGTH) # Store recent values

# --- Helper Functions ---
def transform_g(g_value, exponent=NON_LINEAR_EXPONENT):
    if g_value == 0: return 0.0
    return math.copysign(math.pow(abs(g_value), 1.0 / exponent), g_value)

MAX_TRANSFORMED_G = transform_g(MAX_G_DISPLAY) if MAX_G_DISPLAY != 0 else 0.0

def map_transformed_g_to_x(transformed_g_value, transformed_g_min, transformed_g_max, x_min, x_max):
    transformed_g_value = max(transformed_g_min, min(transformed_g_max, transformed_g_value))
    if transformed_g_max == transformed_g_min: return (x_min + x_max) // 2
    norm_g = (transformed_g_value - transformed_g_min) / (transformed_g_max - transformed_g_min)
    x_val = x_min + norm_g * (x_max - x_min)
    return int(max(x_min, min(x_max, x_val)))

def center_text_x(text_string, screen_width=OLED_WIDTH, font_width=FONT_WIDTH):
    return max(0, (screen_width - (len(text_string) * font_width)) // 2)

def calculate_average(dq):
    if not dq: return 0.0
    return sum(dq) / len(dq)

# --- Hardware Setup & Init (Simplified for brevity) ---
i2c0 = None; i2c1 = None; oled = None; imu = None
try: i2c0 = machine.I2C(0, sda=machine.Pin(4), scl=machine.Pin(5), freq=400000); i2c0.scan()
except Exception as e: print(f"I2C0 Error: {e}")
try: i2c1 = machine.I2C(1, sda=machine.Pin(6), scl=machine.Pin(7), freq=100000); i2c1.scan()
except Exception as e: print(f"I2C1 Error: {e}")
if i2c0 and oled_lib_found:
    try: oled = SSD1327_I2C(i2c0, address=OLED_ADDR); print("OLED OK")
    except Exception as e: print(f"OLED Init Error: {e}")
if i2c1 and imu_lib_found:
    try: imu = qmi8658.QMI8658(i2c1, address=IMU_ADDR); print("IMU OK")
    except Exception as e: print(f"IMU Init Error: {e}")

# --- Main Loop ---
print("Starting main loop...")

if not imu or not oled:
    print("Initialization failed. Halting.")
else:
    print("--- G-Force Display Running ---")
    oled.fill(oled.black); oled.show()
    center_x = OLED_WIDTH // 2; bar_start_x = center_x - BAR_HALF_WIDTH; bar_end_x = center_x + BAR_HALF_WIDTH
    title_lat_str = "Lateral G Force"; title_acc_str = "BRK/ACC G Force"
    title_lat_x = center_text_x(title_lat_str); title_acc_x = center_text_x(title_acc_str)

    while True:
        current_ticks = time.ticks_ms()
        loop_start_time = current_ticks # For calculating actual loop duration

        try:
            # 1. Read & Calibrate Actual Gs
            xyz_data = imu.Read_XYZ()
            accel_fwd_bwd_actual = xyz_data[FORWARD_BACKWARD_AXIS] - ACCEL_X_OFFSET
            accel_lateral_actual = xyz_data[LATERAL_AXIS] - ACCEL_Y_OFFSET

            # Append to history (using absolute for lateral, raw for fwd/bwd)
            recent_lateral_g.append(abs(accel_lateral_actual))
            recent_fwd_bwd_g.append(accel_fwd_bwd_actual) # Keep sign for braking check

            # 2. Time-Based Deadzone Filter for Display
            # Lateral
            if abs(accel_lateral_actual) >= G_DEADZONE_THRESHOLD:
                g_force_lateral_display = accel_lateral_actual
                lat_below_thresh_start_time = None
            else:
                if lat_below_thresh_start_time is None: lat_below_thresh_start_time = current_ticks
                elif time.ticks_diff(current_ticks, lat_below_thresh_start_time) >= DEADZONE_FILTER_MS:
                    g_force_lateral_display = 0.0
            # Forward/Backward
            if abs(accel_fwd_bwd_actual) >= G_DEADZONE_THRESHOLD:
                g_force_fwd_bwd_display = accel_fwd_bwd_actual
                acc_below_thresh_start_time = None
            else:
                if acc_below_thresh_start_time is None: acc_below_thresh_start_time = current_ticks
                elif time.ticks_diff(current_ticks, acc_below_thresh_start_time) >= DEADZONE_FILTER_MS:
                    g_force_fwd_bwd_display = 0.0

            # --- 3. Sustained Max G Logic ---
            # Lateral Check
            current_abs_lat = abs(accel_lateral_actual)
            if not checking_sustained_lat:
                if current_abs_lat > max_lateral_g and current_abs_lat >= SUSTAINED_G_MIN_TRIGGER:
                    checking_sustained_lat = True
                    candidate_lat_start_time = current_ticks
                    candidate_lat_center_g = current_abs_lat # Use current abs value as center
                    # Clear history for averaging this specific event
                    recent_lateral_g.clear()
                    recent_lateral_g.append(current_abs_lat)
            else: # Already checking lateral
                if abs(current_abs_lat - candidate_lat_center_g) <= SUSTAINED_G_WINDOW:
                    # Still within window, check duration
                    elapsed_lat = time.ticks_diff(current_ticks, candidate_lat_start_time)
                    if elapsed_lat >= SUSTAINED_G_DURATION_MS:
                        # Sustained! Calculate average and update max if needed
                        avg_sustained_lat = calculate_average(recent_lateral_g)
                        if avg_sustained_lat > max_lateral_g:
                            max_lateral_g = avg_sustained_lat
                            print(f"New Max Lateral: {max_lateral_g:.2f}g (Avg over {elapsed_lat}ms)") # Debug
                        checking_sustained_lat = False # Stop checking
                else:
                    # Fell outside window, reset check
                    checking_sustained_lat = False

            # Braking Check (Only applies if current value is negative)
            current_abs_brk = abs(accel_fwd_bwd_actual) if accel_fwd_bwd_actual < 0 else 0.0
            if not checking_sustained_brk:
                 # Start check if braking is strong enough and potentially new max
                if accel_fwd_bwd_actual < 0 and current_abs_brk > max_braking_g and current_abs_brk >= SUSTAINED_G_MIN_TRIGGER:
                    checking_sustained_brk = True
                    candidate_brk_start_time = current_ticks
                    candidate_brk_center_g = current_abs_brk # Use current abs value
                    # Clear history for averaging
                    recent_fwd_bwd_g.clear()
                    recent_fwd_bwd_g.append(accel_fwd_bwd_actual) # Store signed value
            else: # Already checking braking
                # Check if still braking and within window of candidate center (using abs values for window check)
                if accel_fwd_bwd_actual < 0 and abs(current_abs_brk - candidate_brk_center_g) <= SUSTAINED_G_WINDOW:
                     # Still within window, check duration
                    elapsed_brk = time.ticks_diff(current_ticks, candidate_brk_start_time)
                    if elapsed_brk >= SUSTAINED_G_DURATION_MS:
                        # Sustained! Calculate average braking G and update max if needed
                        # Average the absolute values of the negative readings in the history
                        braking_readings = [abs(g) for g in recent_fwd_bwd_g if g < 0]
                        avg_sustained_brk = calculate_average(braking_readings) if braking_readings else 0.0

                        if avg_sustained_brk > max_braking_g:
                            max_braking_g = avg_sustained_brk
                            print(f"New Max Braking: {max_braking_g:.2f}g (Avg over {elapsed_brk}ms)") # Debug
                        checking_sustained_brk = False # Stop checking
                else:
                     # Fell outside window or stopped braking, reset check
                     checking_sustained_brk = False


            # 4. Transform DISPLAY Gs for bars
            transformed_lat_g = transform_g(g_force_lateral_display)
            transformed_acc_g = transform_g(g_force_fwd_bwd_display)

            # 5. Map TRANSFORMED Gs to Pixel Coordinates
            lat_x = map_transformed_g_to_x(transformed_lat_g, -MAX_TRANSFORMED_G, MAX_TRANSFORMED_G, bar_start_x, bar_end_x)
            acc_x = map_transformed_g_to_x(transformed_acc_g, -MAX_TRANSFORMED_G, MAX_TRANSFORMED_G, bar_start_x, bar_end_x)

            # 6. Prepare dynamic text strings
            current_lat_str = "{:+.2f}g".format(g_force_lateral_display)
            current_lat_x = center_text_x(current_lat_str)
            max_lat_str = "Max:{:.2f}g".format(max_lateral_g)
            max_lat_x = center_text_x(max_lat_str)

            current_acc_str = "{:+.2f}g".format(g_force_fwd_bwd_display)
            current_acc_x = center_text_x(current_acc_str)
            max_acc_str = "Max B:{:.2f}g".format(max_braking_g)
            max_acc_x = center_text_x(max_acc_str)

            # 7. Draw Display
            oled.fill(oled.black)
            # Lateral G Section
            oled.text(title_lat_str, title_lat_x, TITLE_Y_LAT, GRAY_LEVEL_TEXT)
            oled.text(max_lat_str, max_lat_x, MAX_VAL_Y_LAT, GRAY_LEVEL_MAX_VAL)
            oled.hline(bar_start_x, BAR_Y_LAT, BAR_HALF_WIDTH * 2, GRAY_LEVEL_BAR_BG)
            bar_y_start_lat = BAR_Y_LAT - BAR_THICKNESS // 2
            if lat_x >= center_x: oled.fill_rect(center_x, bar_y_start_lat, lat_x - center_x + 1, BAR_THICKNESS, GRAY_LEVEL_INDICATOR)
            else: oled.fill_rect(lat_x, bar_y_start_lat, center_x - lat_x + 1, BAR_THICKNESS, GRAY_LEVEL_INDICATOR)
            oled.vline(center_x, BAR_Y_LAT - (BAR_THICKNESS//2 + 1), BAR_THICKNESS + 3, GRAY_LEVEL_BAR_BG)
            oled.text(current_lat_str, current_lat_x, CURRENT_VAL_Y_LAT, GRAY_LEVEL_CURRENT_VAL)
            # Accel/Decel G Section
            oled.text(title_acc_str, title_acc_x, TITLE_Y_ACCEL, GRAY_LEVEL_TEXT)
            oled.text(max_acc_str, max_acc_x, MAX_VAL_Y_ACCEL, GRAY_LEVEL_MAX_VAL)
            oled.hline(bar_start_x, BAR_Y_ACCEL, BAR_HALF_WIDTH * 2, GRAY_LEVEL_BAR_BG)
            bar_y_start_accel = BAR_Y_ACCEL - BAR_THICKNESS // 2
            if acc_x >= center_x: oled.fill_rect(center_x, bar_y_start_accel, acc_x - center_x + 1, BAR_THICKNESS, GRAY_LEVEL_INDICATOR)
            else: oled.fill_rect(acc_x, bar_y_start_accel, center_x - acc_x + 1, BAR_THICKNESS, GRAY_LEVEL_INDICATOR)
            oled.vline(center_x, BAR_Y_ACCEL - (BAR_THICKNESS//2 + 1), BAR_THICKNESS + 3, GRAY_LEVEL_BAR_BG)
            oled.text(current_acc_str, current_acc_x, CURRENT_VAL_Y_ACCEL, GRAY_LEVEL_CURRENT_VAL)

            # 8. Update Screen
            oled.show()

        except Exception as e:
            print(f"Error in main loop: {e}")
            if oled: # Try to display error on OLED
                try: oled.fill(0); oled.text("LOOP ERROR", 0, 0, 15); oled.show()
                except Exception: pass
            time.sleep(2)

        # Loop Delay Enforcement (Approximate)
        loop_end_time = time.ticks_ms()
        elapsed_loop = time.ticks_diff(loop_end_time, loop_start_time)
        sleep_time = LOOP_DELAY_MS - elapsed_loop
        if sleep_time > 0:
            time.sleep_ms(sleep_time)
        # else: print(f"Warning: Loop took longer than {LOOP_DELAY_MS}ms ({elapsed_loop}ms)") # Optional debug