import BTAnalysisPyInterface as bta


def main():
    
    # Define three necessary graphs
    cycling_graph_normalized = bta.add_graph(
        title="Capacity Normalized vs Cycle",
        xAxis="Cycle number",
        y1Axis="Normalized Capacity (mA.h/g)",
    )
    cycling_graph = bta.add_graph(
        title="Capacity vs Cycle", 
        xAxis="Cycle number", 
        y1Axis="Capacity",
    )
    efficiency_graph = bta.add_graph(
        title="Coulombic Efficiency vs Cycle",
        xAxis="Cycle number",
        y1Axis="Coulombic Efficiency(%)",
    )

    retention_graph = bta.add_graph(
        title="Capacity Retention vs Cycle",
        xAxis="Cycle number",
        y1Axis="Capacity Retention (%)",
    )
    
    # Get in BT-Analysis selected sources
    sources = bta.get_sources()

    # Define colors for the traces in the graphs
    colors_discharge = ["#800080", "#ffff00", "#E1C16E", "#A52A2A"]
    colors_charge = ["#ff0000", "#0000ff", "#000000", "ffc0cb"]
    
    # Iterate through all sources
    for source in sources:
        index = sources.index(source)
        
        # Define traces for each graph
        cycling_graph_normalized.configure_trace(
            f"Discharge normalized Source {index + 1}",
            color=colors_discharge[index],
            line_style=bta.LineStyle.DASHED,
            line_width=2.0,
            marker=bta.MarkerStyle.DIAMOND,
            marker_size=4.0,
        )

        cycling_graph.configure_trace(
            f"Discharge Source {index + 1}",
            color=colors_discharge[index],
            line_style=bta.LineStyle.DASHED,
            line_width=2.0,
            marker=bta.MarkerStyle.DIAMOND,
            marker_size=4.0,
        )

        retention_graph.configure_trace(
            f"Charge Retention Source {index + 1}",
            color=colors_charge[index],
            line_style=bta.LineStyle.DASHED,
            line_width=2.0,
            marker=bta.MarkerStyle.DIAMOND,
            marker_size=4.0,
        )

        efficiency_graph.configure_trace(
            f"Coulombic Efficiency Source {index + 1}",
            color=colors_discharge[index],
            line_style=bta.LineStyle.DASHED,
            line_width=2.0,
            marker=bta.MarkerStyle.DIAMOND,
            marker_size=4.0,
        )

        # Get summary data per cycle
        summary = source.get_summary(bta.SummaryType.PER_CYCLE)

        cycles = summary.get_column("Cycle number")
        q_discharge = summary.get_column("Q discharge")
        q_charge = summary.get_column("Q charge")

        # Get DUT capacity
        dut_capacity = source.get_dut_type_capacity()

        # Define empty lists and empty series for calculation
        q_discharge_normalized = []
        q_charge_normalized = []

        coulombic_efficiency = []
        capacity_retention = []

        q_discharge_normalized_series = bta.Series("q_discharge_normalized")
        q_charge_normalized_series = bta.Series("q_charge_normalized")
        coulombic_efficiency_series = bta.Series("coulombic_efficiency")
        capacity_retention_series = bta.Series("capacity_retention")

        # Calculate normalized capacity, coulombic efficiency and capacity retention
        for i in range(len(q_discharge)):
            q_discharge_normalized.append(q_discharge.get_values()[i] / dut_capacity)
            q_charge_normalized.append(q_charge.get_values()[i] / dut_capacity)
            try:
                coulombic_efficiency.append(
                    (q_discharge.get_values()[i] / q_charge.get_values()[i]) * 100
                )
            except ZeroDivisionError:
                pass
            try:
                capacity_retention.append(
                    (q_discharge.get_values()[i + 1] / q_discharge.get_values()[i])
                    * 100
                )
            except IndexError:
                pass
            except ZeroDivisionError:
                pass

        q_charge_normalized_series.add_values(q_charge_normalized)
        q_discharge_normalized_series.add_values(q_discharge_normalized)
        coulombic_efficiency_series.add_values(coulombic_efficiency)
        capacity_retention_series.add_values(capacity_retention)

        # Add data to the graphs
        cycling_graph.add_series(
            f"Discharge Source {index + 1}", cycles, q_discharge, bta.GraphAxis.Y1
        )

        cycling_graph_normalized.add_series(
            f"Discharge normalized Source {index + 1}",
            cycles,
            q_discharge_normalized_series,
            bta.GraphAxis.Y1,
        )

        efficiency_graph.add_series(
            f"Coulombic Efficiency Source {index + 1}",
            cycles,
            coulombic_efficiency_series,
            bta.GraphAxis.Y1,
        )

        # Change the trace colors for each graph before adding the next data series

        cycling_graph_normalized.configure_trace(
            f"Charge normalized Source {index + 1}",
            color=colors_charge[index],
            line_style=bta.LineStyle.DASHED,
            line_width=2.0,
            marker=bta.MarkerStyle.DIAMOND,
            marker_size=4.0,
        )

        cycling_graph.configure_trace(
            f"Charge Source {index + 1}",
            color=colors_charge[index],
            line_style=bta.LineStyle.DASHED,
            line_width=2.0,
            marker=bta.MarkerStyle.DIAMOND,
            marker_size=4.0,
        )

        # Add data to the graphs
        cycling_graph.add_series(
            f"Charge Source {index + 1}", cycles, q_charge, bta.GraphAxis.Y1
        )

        cycling_graph_normalized.add_series(
            f"Charge normalized Source {index + 1}",
            cycles,
            q_charge_normalized_series,
            bta.GraphAxis.Y1,
        )
        
        retention_graph.add_series(
            f"Charge Retention Source {index + 1}",
            cycles,
            capacity_retention_series,
            bta.GraphAxis.Y1,
        )
