"""
Simple library for Tektronix MSO4000 series oscilloscopes
Provides connection checking and waveform recording to CSV
"""
import pyvisa
import numpy as np
import time
import csv
from datetime import datetime
from typing import Optional, Tuple, List


class MSO4000:
    """Simple interface for Tektronix MSO4000 series oscilloscopes"""
    
    def __init__(self, channel: int = 1):
        """
        Initialize MSO4000 interface
        
        Args:
            channel: Channel number (1-4)
        """
        self.channel = channel
        self.scope = None
        self.rm = None
        self.timebase = None
        self.ch_scale = None
    
    def check_connection(self) -> Tuple[bool, Optional[str]]:
        """
        Check if oscilloscope is connected and accessible
        
        Returns:
            Tuple of (success: bool, device_info: str or None)
        """
        try:
            self.rm = pyvisa.ResourceManager()
            resources = self.rm.list_resources()
            
            # Prefer USB resources
            usb_resources = [r for r in resources if r.startswith('USB')]
            if not usb_resources:
                return False, "No USB oscilloscope found"
            
            resource_name = usb_resources[0]
            self.scope = self.rm.open_resource(resource_name)
            self.scope.timeout = 5000
            
            # Query identification
            idn = self.scope.query("*IDN?")
            device_info = idn.strip()
            
            return True, device_info
            
        except Exception as e:
            return False, f"Connection error: {e}"
    
    def connect(self) -> bool:
        """
        Connect to oscilloscope and configure for waveform transfer
        
        Returns:
            True if successful, False otherwise
        """
        success, info = self.check_connection()
        if not success:
            print(f"❌ {info}")
            return False
        
        try:
            print(f"✓ Connected to: {info}")
            
            # Configure for waveform transfer
            self.scope.write("*CLS")
            self.scope.write(f":DATA:SOURCE CH{self.channel}")
            self.scope.write(":DATA:ENCDG RIBINARY")
            self.scope.write(":WFMOUTPRE:BYT_NR 2")
            self.scope.write(":DATA INIT")
            
            # Get scale info
            self.timebase = float(self.scope.query(":HORIZONTAL:SCALE?"))
            self.ch_scale = float(self.scope.query(f":CH{self.channel}:SCALE?"))
            
            print(f"✓ Configured: Timebase={self.timebase} s/div, CH{self.channel}={self.ch_scale} V/div")
            return True
            
        except Exception as e:
            print(f"❌ Configuration error: {e}")
            return False
    
    def capture_waveform(self) -> Optional[np.ndarray]:
        """
        Capture a single waveform from the oscilloscope
        
        Returns:
            Voltage array (numpy array) or None if failed
        """
        if not self.scope:
            print("❌ Not connected. Call connect() first.")
            return None
        
        try:
            self.scope.timeout = 30000  # 30 seconds for data transfer
            
            # Trigger and get waveform
            self.scope.write(":TRIGGER:FORCE")
            time.sleep(0.2)
            
            # Get data using write + read_raw for binary data
            self.scope.write("CURVE?")
            curve_data_bytes = self.scope.read_raw()
            
            # Parse binary data
            curve_data_str = curve_data_bytes.decode('latin1', errors='ignore')
            
            if curve_data_str.startswith('#'):
                digit_count = int(curve_data_str[1])
                byte_count = int(curve_data_str[2:2+digit_count])
                header_len = 2 + digit_count
                
                # Extract binary data
                binary_data = curve_data_bytes[header_len:header_len+byte_count]
                
                # Convert to array (16-bit signed integers, big-endian)
                voltage_array = np.frombuffer(binary_data, dtype='>i2')
                
                # Scale to voltage
                voltage_scaled = (voltage_array / 32768.0) * self.ch_scale * 5.0
            else:
                # ASCII fallback
                values = [float(x) for x in curve_data_str.split(',')]
                voltage_scaled = np.array(values)
            
            return voltage_scaled
            
        except Exception as e:
            print(f"❌ Error capturing waveform: {e}")
            return None
    
    def record_to_csv(self, filename: Optional[str] = None, 
                      max_records: Optional[int] = None,
                      delay: float = 0.5) -> Optional[str]:
        """
        Record multiple waveforms and save to CSV
        
        Args:
            filename: Output CSV filename (auto-generated if None)
            max_records: Maximum number of waveforms to record (None = until interrupted)
            delay: Delay between captures in seconds
            
        Returns:
            Filename of saved CSV or None if failed
        """
        if not self.scope:
            print("❌ Not connected. Call connect() first.")
            return None
        
        if filename is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"waveform_data_{timestamp}.csv"
        
        all_data = []
        timestamps = []
        record_count = 0
        
        print(f"\n📝 Recording waveforms... (Press Ctrl+C to stop)")
        print(f"   Output: {filename}\n")
        
        try:
            while True:
                if max_records and record_count >= max_records:
                    print(f"\n✓ Record limit reached ({max_records} records)")
                    break
                
                print(f"[{record_count+1}] Capturing... ", end="", flush=True)
                
                voltage_data = self.capture_waveform()
                
                if voltage_data is not None:
                    all_data.append(voltage_data)
                    timestamps.append(time.time())
                    record_count += 1
                    
                    min_v = voltage_data.min()
                    max_v = voltage_data.max()
                    mean_v = voltage_data.mean()
                    print(f"✓ {len(voltage_data):6d} samples | "
                          f"Min: {min_v:8.4f}V | Max: {max_v:8.4f}V | Mean: {mean_v:8.4f}V")
                    
                    time.sleep(delay)
                else:
                    print("✗ FAILED")
                    time.sleep(1)
        
        except KeyboardInterrupt:
            print("\n⏹ Recording interrupted by user")
        
        if not all_data:
            print("❌ No data to export!")
            return None
        
        # Export to CSV
        print(f"\n💾 Exporting to {filename}...")
        
        try:
            with open(filename, 'w', newline='') as csvfile:
                max_samples = max(len(wf) for wf in all_data)
                
                # Create header
                headers = ['Sample_Index', 'Time_sec'] + \
                         [f'Waveform_{i+1}_V' for i in range(record_count)]
                
                writer = csv.writer(csvfile)
                writer.writerow(headers)
                
                # Calculate time per point
                num_points = len(all_data[0]) if all_data else 1
                time_per_point = (self.timebase * 10.0) / num_points
                
                # Write data
                for sample_idx in range(max_samples):
                    row = [sample_idx, sample_idx * time_per_point]
                    
                    for waveform in all_data:
                        if sample_idx < len(waveform):
                            row.append(waveform[sample_idx])
                        else:
                            row.append('')
                    
                    writer.writerow(row)
            
            print(f"✓ Exported {record_count} waveforms, {max_samples} samples each")
            
            import os
            file_size = os.path.getsize(filename)
            print(f"✓ File size: {file_size / 1024 / 1024:.2f} MB")
            
            return filename
            
        except Exception as e:
            print(f"❌ Error exporting: {e}")
            return None
    
    def get_acquisition_info(self) -> Optional[dict]:
        """
        Get acquisition parameters: sampling rate, window size, number of samples
        
        Returns:
            Dictionary with acquisition info or None if not connected
        """
        if not self.scope or not self.timebase:
            print("❌ Not connected. Call connect() first.")
            return None
        
        try:
            # Capture one waveform to get number of samples
            waveform = self.capture_waveform()
            if waveform is None:
                return None
            
            num_samples = len(waveform)
            time_window = self.timebase * 10.0  # 10 divisions
            sampling_rate = num_samples / time_window  # Hz
            time_per_sample = time_window / num_samples  # seconds
            
            info = {
                'num_samples': num_samples,
                'time_window_sec': time_window,
                'sampling_rate_hz': sampling_rate,
                'time_per_sample_sec': time_per_sample,
                'timebase_s_div': self.timebase,
                'channel': self.channel,
                'ch_scale_v_div': self.ch_scale,
            }
            
            # Print nicely formatted info
            print("\n" + "="*60)
            print("ACQUISITION PARAMETERS")
            print("="*60)
            print(f"Channel:              CH{self.channel}")
            print(f"Timebase:             {self.timebase*1e3:.3f} ms/div")
            print(f"Channel Scale:        {self.ch_scale:.3f} V/div")
            print(f"Number of Samples:    {num_samples}")
            print(f"Time Window:          {time_window*1e3:.3f} ms ({time_window:.6f} s)")
            print(f"Sampling Rate:        {sampling_rate/1e6:.2f} MHz ({sampling_rate:.0f} Hz)")
            print(f"Time per Sample:      {time_per_sample*1e9:.2f} ns ({time_per_sample*1e6:.3f} µs)")
            print("="*60 + "\n")
            
            return info
            
        except Exception as e:
            print(f"❌ Error getting acquisition info: {e}")
            return None
    
    def close(self):
        """Close connection to oscilloscope"""
        if self.scope:
            self.scope.close()
            self.scope = None
            print("✓ Connection closed")


# Convenience functions
def check_connection() -> Tuple[bool, Optional[str]]:
    """
    Quick function to check if oscilloscope is connected
    
    Returns:
        Tuple of (success: bool, device_info: str or None)
    """
    scope = MSO4000()
    return scope.check_connection()


def quick_record(channel: int = 1, max_records: Optional[int] = None,
                filename: Optional[str] = None) -> Optional[str]:
    """
    Quick function to connect and record waveforms
    
    Args:
        channel: Channel number (1-4)
        max_records: Maximum number of waveforms (None = until interrupted)
        filename: Output CSV filename (auto-generated if None)
        
    Returns:
        Filename of saved CSV or None if failed
    """
    scope = MSO4000(channel=channel)
    
    if not scope.connect():
        return None
    
    try:
        filename = scope.record_to_csv(filename=filename, max_records=max_records)
        return filename
    finally:
        scope.close()


if __name__ == "__main__":
    # Example usage
    print("=" * 70)
    print("MSO4000 Library - Example Usage")
    print("=" * 70 + "\n")
    
    # Check connection
    print("[1] Checking connection...")
    success, info = check_connection()
    if success:
        print(f"✓ {info}\n")
    else:
        print(f"❌ {info}\n")
        exit(1)
    
    # Record waveforms
    print("[2] Recording waveforms...")
    scope = MSO4000(channel=1)
    
    if scope.connect():
        filename = scope.record_to_csv(max_records=5)  # Record 5 waveforms
        if filename:
            print(f"\n✓ Success! Data saved to: {filename}")
        scope.close()

