#!/usr/bin/env python3
"""
Test script to verify FastFrame packed transfer hypothesis.

Tests if all FastFrame records come in a single CURVe? transfer,
similar to WFMv3 file format.
"""

import time
from tektronix_base import TektronixOscilloscopeBase


def main():
    """Test FastFrame packed transfer"""

    scope_ip = "192.168.10.105"
    num_frames = 1000
    target_record_length = 5000

    print(f"Connecting to oscilloscope at {scope_ip}...")

    scope = TektronixOscilloscopeBase(resource_name=scope_ip, timeout=20.0)

    try:
        scope.connect()
        print(f"Connected to {scope.identify()}\n")

        # Configure CH1 for 1MHz square wave ±250mV
        print("Configuring CH1...")
        scope.set_channel_scale(1, 0.1)  # 100mV/div
        scope.set_channel_offset(1, 0.0)
        scope.set_channel_coupling(1, 'DC')
        scope.set_channel_termination(1, 50)

        # Configure horizontal for 5000 point record length
        # At 6.25 GS/s: 5000 points / 6.25e9 = 800ns window = 80ns/div
        print("Configuring horizontal for 5000 point record length...")
        scope.set_record_length(target_record_length)
        scope.set_time_scale(80e-9)  # 80ns/div = 800ns window at 6.25 GS/s

        # Check actual settings
        actual_record_length = scope.get_record_length()
        actual_time_scale = scope.get_time_scale()
        actual_sample_rate = scope.get_sample_rate()
        print(f"  Record length: {actual_record_length}")
        print(f"  Time scale: {actual_time_scale*1e9:.0f} ns/div")
        print(f"  Sample rate: {actual_sample_rate/1e6:.1f} MS/s")

        # Configure trigger
        print("Configuring trigger...")
        scope.set_trigger_source('CH1')
        scope.set_trigger_level(1, 0.1)
        scope.set_trigger_slope('RISE')
        scope.set_trigger_mode('NORMAL')

        # Configure FastFrame
        print(f"Configuring FastFrame for {num_frames} frames...")
        scope.set_fastframe_state(True)
        scope.set_fastframe_count(num_frames)

        # Check record length AFTER enabling FastFrame
        actual_record_length_ff = scope.get_record_length()
        actual_sample_rate_ff = scope.get_sample_rate()
        print(f"  Record length after FastFrame: {actual_record_length_ff}")
        print(f"  Sample rate after FastFrame: {actual_sample_rate_ff/1e6:.1f} MS/s")

        ff_state = scope.get_fastframe_state()
        ff_count = scope.get_fastframe_count()
        print(f"  FastFrame: {'enabled' if ff_state else 'disabled'}, {ff_count} frames\n")

        # Configure waveform transfer
        print("Configuring waveform transfer...")
        scope.set_data_encoding('RIBinary')
        scope.set_wfmoutpre_encoding('BINary')
        scope.set_wfmoutpre_byte_count(1)
        scope.set_wfmoutpre_byte_order('MSB')
        scope.set_data_source('CH1')

        # Acquire data
        print("\nAcquiring FastFrame data...")
        time.sleep(0.5)
        scope.write("ACQuire:STATE RUN")

        # Wait for triggers (no polling - just wait)
        print("  Waiting for triggers...")
        time.sleep(3)

        scope.write("ACQuire:STATE STOP")
        time.sleep(0.5)
        print("  Acquisition stopped\n")

        # Transfer all frames individually
        print(f"\n=== Transferring {num_frames} Frames Individually ===\n")

        all_waveforms = []
        start_time = time.time()

        for frame_num in range(1, num_frames + 1):
            # Progress every 100 frames
            if frame_num == 1 or frame_num % 100 == 0:
                elapsed = time.time() - start_time
                if frame_num > 1:
                    rate = frame_num / elapsed
                    eta = (num_frames - frame_num) / rate
                    print(f"  Frame {frame_num}/{num_frames} - {elapsed:.1f}s elapsed - {rate:.1f} fps - ETA {eta:.1f}s")
                else:
                    print(f"  Frame {frame_num}/{num_frames}...")

            # Select frame and wait for it to take effect
            scope.set_fastframe_selected(frame_num)
            time.sleep(0.01)  # Small delay for command processing

            # Transfer curve data
            curve_bytes = scope.transfer_curve()
            waveform = scope.parse_curve_data(curve_bytes, byte_count=1, signed=True, byte_order='MSB')
            all_waveforms.append(waveform)

        end_time = time.time()
        total_time = end_time - start_time

        print(f"\n=== Transfer Complete ===")
        print(f"Frames transferred: {len(all_waveforms)}")
        print(f"Points per frame: {len(all_waveforms[0])}")
        print(f"Total time: {total_time:.2f}s")
        print(f"Rate: {num_frames / total_time:.1f} frames/s")
        print(f"Total data: {sum(len(wf) for wf in all_waveforms)} samples")
        print(f"Throughput: {sum(len(wf) for wf in all_waveforms) / total_time / 1e6:.2f} MS/s")

        # Show statistics for selected frames
        print(f"\n=== Frame Statistics ===")
        for frame_idx in [0, num_frames//2, num_frames-1]:
            wf = all_waveforms[frame_idx]
            print(f"\nFrame {frame_idx + 1}:")
            print(f"  Samples: {len(wf)}")
            print(f"  Min: {min(wf)}, Max: {max(wf)}, Avg: {sum(wf)/len(wf):.1f}")
            print(f"  First 10: {wf[:10]}")

        print("\n=== Test Complete ===")

    except Exception as e:
        print(f"\nError: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()
    finally:
        if scope.is_connected:
            try:
                scope.set_fastframe_state(False)
            except:
                pass
            scope.disconnect()
            print("\nDisconnected")


if __name__ == "__main__":
    main()