# flake8: noqa T001, E501
# -*- coding: utf-8 -*-
"""
Custom Graphs Example

Demonstrates non-time-series visualizations:
    - I-V characteristic curves
    - Capacity vs voltage plots
    - Nyquist plots (EIS data)
    - Cycle filtering (FilterCriteria vs manual)

Shows how to use any columns for X/Y axes and filter data efficiently.
"""

import BTAnalysisPyInterface as bta


def main():
    bta.print("Setting up custom graphs...")

    for source in bta.get_sources():

        # ========== I-V CURVE ==========
        iv_graph = bta.add_graph(
            title="I-V Characteristic Curve", xAxis="I", y1Axis="U"
        )

        if source.has_category("Potentiostat measure"):
            for point in source.data_points("Potentiostat measure", ["I", "U"]):
                iv_graph.add_point(point["I"], point["U"])
            bta.print("I-V curve complete!")

        # ========== CAPACITY VS VOLTAGE ==========
        capacity_graph = bta.add_graph(
            title="Capacity vs Voltage", xAxis="Q", y1Axis="U"
        )

        if source.has_column("Potentiostat measure", "Q"):
            for point in source.data_points("Potentiostat measure", ["Q", "U"]):
                capacity_graph.add_point(point["Q"], point["U"])
            bta.print("Capacity plot complete!")
        else:
            bta.print("Skipped capacity plot (no Q data)")

        # ========== NYQUIST PLOT (EIS) ==========
        if source.has_column("Potentiostat measure", "Re(Z)") and source.has_column(
            "Potentiostat measure", "-Im(Z)"
        ):
            nyquist_graph = bta.add_graph(
                title="Nyquist Plot (EIS)", xAxis="Re(Z)", y1Axis="-Im(Z)"
            )

            for point in source.data_points(
                "Potentiostat measure", ["Re(Z)", "-Im(Z)"]
            ):
                nyquist_graph.add_point(point["Re(Z)"], point["-Im(Z)"])

            bta.print("Nyquist plot complete!")
        else:
            bta.print("Skipped Nyquist plot (no EIS data)")

        # ========== FILTERED CYCLES - METHOD 1: FilterCriteria (Recommended) ==========
        filter_cycles = bta.FilterCriteria()
        filter_cycles.cycles = [1, 5, 10]

        cycles_graph_filtered = bta.add_graph(
            title="Voltage (Cycles 1, 5, 10) - FilterCriteria", xAxis="Time", y1Axis="U"
        )

        if source.has_column("Potentiostat measure", "Cycle number"):
            for point in source.data_points(
                "Potentiostat measure", ["Time", "U"], filter_cycles
            ):
                cycles_graph_filtered.add_point(point["Time"], point["U"])
            bta.print("Filtered cycles (FilterCriteria) complete!")

        # ========== FILTERED CYCLES - METHOD 2: Manual with NaN ==========
        cycles_graph_manual = bta.add_graph(
            title="Voltage (Cycles 1, 5, 10) - Manual Filter", xAxis="Time", y1Axis="U"
        )

        target_cycles = {1, 5, 10}
        if source.has_column("Potentiostat measure", "Cycle number"):
            i = 0
            for point in source.data_points(
                "Potentiostat measure", ["Time", "U", "Cycle number"]
            ):
                i += 1
                try:  # Manage points that are not in a cycle
                    cycle = int(point["Cycle number"])
                    if cycle in target_cycles:
                        cycles_graph_manual.add_point(point["Time"], point["U"])
                    else:
                        # Add NaN for filtered-out cycles to create gaps in the plot
                        cycles_graph_manual.add_point(point["Time"], float("nan"))
                except ValueError:
                    print(f"No Cycle Number found for point {i}")
            bta.print("Filtered cycles (manual) complete!")
        else:
            bta.print("Skipped cycle filtering (no cycle data)")

    bta.print("Custom graphs ready!")
