#!/usr/bin/env python3
import pyvisa
import yaml
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime
import os
import argparse  # For handling command-line arguments


def detach_kernel_tmc(vid, pid):
    """Detach kernel driver for USBTMC devices."""
    import usb.core, usb.util
    dev = usb.core.find(idVendor=vid, idProduct=pid)
    if not dev:
        return
    try:
        for cfg in dev:
            for intf in cfg:
                if intf.bInterfaceClass == 0xFE:  # USBTMC
                    if dev.is_kernel_driver_active(intf.bInterfaceNumber):
                        dev.detach_kernel_driver(intf.bInterfaceNumber)
        try:
            dev.set_configuration()
        except Exception:
            pass
    except Exception:
        pass


def load_settings(file_path):
    """Load settings from a YAML file."""
    with open(file_path, 'r') as file:
        return yaml.safe_load(file)


def configure_oscilloscope(osc, settings):
    """Configure the oscilloscope with the given settings."""
    osc.write(f":TIMebase:SCALe {settings['timebase_scale']}")
    for channel, config in settings['channels'].items():
        osc.write(f":CHANnel{channel}:OFFSet {config['offset']}")
        osc.write(f":CHANnel{channel}:SCALe {config['scale']}")
    osc.write(":SINGle")  # Set to single trigger mode


def fetch_waveform(osc, channel, timebase_scale):
    """Fetch waveform data from a specific channel."""
    osc.write(f":WAVeform:SOURce CHANnel{channel}")
    osc.write(":WAVeform:FORMat ASCii")
    data = osc.query(":WAVeform:DATA?")
    
    if data.startswith('#'):
        header_length = int(data[1])  # The second character indicates the header length
        data = data[2 + header_length:]  # Skip the header
    
    voltage = np.array([float(x) for x in data.split(',')])
    time = np.linspace(0, timebase_scale * 10, len(voltage))  # 10 divisions
    return time, voltage


def save_waveform_plot(time, voltage, channel, output_dir):
    """Save waveform as a JPG image."""
    plt.figure()
    plt.plot(time, voltage)
    plt.title(f"Channel {channel} Waveform")
    plt.xlabel("Time (s)")
    plt.ylabel("Voltage (V)")
    plt.grid()
    file_path = f"{output_dir}/channel{channel}_waveform.jpg"
    plt.savefig(file_path)
    plt.close()
    print(f"Waveform saved as {file_path}")


def save_raw_timeseries(time, voltage, channel, output_dir):
    """Save raw time series data as a CSV file."""
    file_path = f"{output_dir}/channel{channel}_timeseries.csv"
    data = np.column_stack((time, voltage))
    np.savetxt(file_path, data, delimiter=",", header="Time (s),Voltage (V)", comments="")
    print(f"Raw time series saved as {file_path}")


def calculate_ripple(data):
    """Calculate the ripple (peak-to-peak voltage)."""
    return float(np.ptp(data))  # Convert to plain float


def calculate_average(data):
    """Calculate the average voltage."""
    return float(np.mean(data))  # Convert to plain float


def main():
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description="Automated oscilloscope measurement script.")
    parser.add_argument("SampleNumber", type=int, help="The sample number for the test.")
    parser.add_argument("Supply", type=str, choices=["PoE", "PoELong", "external"], help="The supply type (PoE, PoELong, or external).")
    parser.add_argument("TestCase", type=str, choices=["baseline", "CPU", "EthPrim", "EthSec"], help="The test case (baseline, CPU, EthPrim, or EthSec).")
    args = parser.parse_args()

    # Detach kernel driver if necessary
    VID, PID = 0x0957, 0x1798  # Keysight DSOX2014A
    detach_kernel_tmc(VID, PID)

    # Load settings
    settings_file = "settings.yaml"
    settings = load_settings(settings_file)

    # Ensure timebase_scale is a float
    timebase_scale = float(settings['timebase_scale'])

    # Create a timestamped subfolder for results
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = os.path.join(settings.get('output_directory', '.'), f"measurement_results_{timestamp}")
    os.makedirs(output_dir, exist_ok=True)

    # Connect to the oscilloscope
    rm = pyvisa.ResourceManager("@py")  # Use pyvisa-py backend
    resource_string = f"USB0::0x{VID:04x}::0x{PID:04x}::MY59124583::INSTR"
    oscilloscope = rm.open_resource(resource_string)
    print("Connected to:", oscilloscope.query("*IDN?"))

    # Configure the oscilloscope
    configure_oscilloscope(oscilloscope, settings)

    # Trigger a measurement
    oscilloscope.write(":DIGitize")

    # Fetch and process data for each channel
    results = {
        "SampleNumber": args.SampleNumber,
        "Supply": args.Supply,
        "TestCase": args.TestCase,
        "Channels": {}
    }
    for channel in settings['channels']:
        time, voltage = fetch_waveform(oscilloscope, channel, timebase_scale)
        save_waveform_plot(time, voltage, channel, output_dir)
        save_raw_timeseries(time, voltage, channel, output_dir)
        ripple = calculate_ripple(voltage)
        average = calculate_average(voltage)
        results["Channels"][channel] = {'ripple': ripple, 'average': average}
        print(f"Channel {channel}: Ripple = {ripple:.3f} V, Average = {average:.3f} V")

    # Save results to a YAML file
    results_file = f"{output_dir}/results.yaml"
    with open(results_file, 'w') as file:
        yaml.dump(results, file)
    print(f"Results saved to {results_file}")

    # Close the connection
    oscilloscope.close()
    print("Measurement complete.")


if __name__ == "__main__":
    main()