# flake8: noqa T001, E501
# -*- coding: utf-8 -*-
"""
NumPy/SciPy Signal Processing Example

Demonstrates advanced data processing:
    - Moving average smoothing
    - Low-pass filtering
    - Statistical analysis
    - Comparison visualization
    - CSV export

Requirements:
    - numpy (installs automatically)
    - scipy (installs automatically)
"""

import BTAnalysisPyInterface as bta

try:
    import numpy as np

    HAS_NUMPY = True
except ImportError:
    HAS_NUMPY = False
    bta.print("NumPy will install automatically")

try:
    from scipy import signal
    from scipy import stats as sp_stats

    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False
    bta.print("SciPy will install automatically")


def moving_average(data, window_size):
    """Apply moving average smoothing"""
    return np.convolve(data, np.ones(window_size) / window_size, mode="same")


def main():
    bta.print("Setting up signal processing examples...")

    for source in bta.get_sources():

        if not HAS_NUMPY:
            bta.print("NumPy not available, skipping processing")
            continue

        # Collect data
        times = []
        voltages = []

        for point in source.data_points("Potentiostat measure", ["Time", "U"]):
            times.append(point["Time"])
            voltages.append(point["U"])

        if not voltages:
            bta.print("No voltage data found")
            continue

        # Convert to NumPy arrays
        times = np.array(times)
        voltages = np.array(voltages)

        # ========== COMPARISON GRAPH (Raw + Smoothed + Filtered) ==========
        comparison_graph = bta.add_graph(
            title="Signal Processing Comparison", xAxis="Time", y1Axis="Voltage"
        )

        # Add raw data
        raw_series = bta.Series("Raw")
        raw_series.add_values(voltages.tolist())

        time_series = bta.Series("Time")
        time_series.add_values(times.tolist())

        comparison_graph.add_series("Raw", time_series, raw_series)

        # Add smoothed data
        smoothed = moving_average(voltages, window_size=10)
        smoothed_series = bta.Series("Smoothed")
        smoothed_series.add_values(smoothed.tolist())
        comparison_graph.add_series("Smoothed", time_series, smoothed_series)

        # Add filtered data (if SciPy available)
        if HAS_SCIPY:
            sos = signal.butter(4, 0.1, output="sos")
            filtered = signal.sosfilt(sos, voltages)
            filtered_series = bta.Series("Filtered")
            filtered_series.add_values(filtered.tolist())
            comparison_graph.add_series("Filtered", time_series, filtered_series)

            # Style traces for clarity
            comparison_graph.configure_trace("Raw", color="#CCCCCC", line_width=1.0)
            comparison_graph.configure_trace(
                "Smoothed", color="#2E86AB", line_width=2.0
            )
            comparison_graph.configure_trace(
                "Filtered", color="#C73E1D", line_width=2.0
            )

            bta.print("Comparison graph created (Raw + Smoothed + Filtered)")
        else:
            # Style without filtered
            comparison_graph.configure_trace("Raw", color="#CCCCCC", line_width=1.0)
            comparison_graph.configure_trace(
                "Smoothed", color="#2E86AB", line_width=2.0
            )
            bta.print("Comparison graph created (Raw + Smoothed)")

        # ========== STATISTICS TABLE ==========
        stats_table = bta.add_table(
            columns=["StatID", "Value"], title="Voltage Statistics"
        )

        stats_list = [
            ("Count", float(len(voltages))),
            ("Mean (V)", float(np.mean(voltages))),
            ("Std Dev (V)", float(np.std(voltages))),
            ("Min (V)", float(np.min(voltages))),
            ("Max (V)", float(np.max(voltages))),
        ]

        # Add SciPy stats if available
        if HAS_SCIPY:
            stats_list.append(("Skewness", float(sp_stats.skew(voltages))))
            stats_list.append(("Kurtosis", float(sp_stats.kurtosis(voltages))))

        # Populate and display
        bta.print("=== Voltage Statistics ===")
        for stat_id, (metric, value) in enumerate(stats_list):
            stats_table.add_row({"StatID": float(stat_id), "Value": value})
            bta.print(f"[{stat_id}] {metric}: {value:.4f}")

        # Export statistics
        stats_table.export_to_csv("voltage_statistics.csv")
        bta.print("Statistics exported to voltage_statistics.csv")

        # ========== EXPORT PROCESSED DATA TABLE ==========
        if HAS_SCIPY:
            results_table = bta.add_table(
                columns=["Time", "Raw", "Smoothed", "Filtered"],
                title="Processing Results",
            )

            for t, raw, smooth, filt in zip(times, voltages, smoothed, filtered):
                results_table.add_row(
                    {
                        "Time": float(t),
                        "Raw": float(raw),
                        "Smoothed": float(smooth),
                        "Filtered": float(filt),
                    }
                )

            # Export with high precision
            csv_options = bta.CsvExportOptions()
            csv_options.decimal_places = 6
            results_table.export_to_csv("processed_data.csv", csv_options)
            bta.print("Processed data exported to processed_data.csv")

        bta.print(f"Processing complete! ({len(voltages)} samples)")

    bta.print("Signal processing examples complete!")
