import BTAnalysisPyInterface as bta


def main():

    # Get in BT-Analysis selected sources
    sources = bta.get_sources()
    source_count = 0

    # Iterate through all sources
    for source in sources:
        source_count += 1
        # Define the columns of your custom table
        columns = [
            "Cycle number",
            "Q charge",
            "Q discharge",
            "Efficiency",
            "Charge Retention",
            "Q charge normalized",
            "Q discharge normalized",
        ]

        # Define the units of your custom table. If "" the default unit for that column will be used.
        units = ["", "", "", "%", "%", "", ""]

        analysis_table = bta.add_table(
            columns=columns, units=units, title=f"Capacity Analysis {source_count}"
        )
        # Get summary data per cycle
        summary = source.get_summary(bta.SummaryType.PER_CYCLE)

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

        # 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 = []

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

                # Add rows to the custom table
        for i in range(len(cycles)):
            try:
                row = {
                    "Cycle number": cycles[i],
                    "Q charge": q_charge[i],
                    "Q discharge": q_discharge[i],
                    "Efficiency": coulombic_efficiency[i],
                    "Charge Retention": capacity_retention[i],
                    "Q charge normalized": q_charge_normalized[i],
                    "Q discharge normalized": q_discharge_normalized[i],
                }
                analysis_table.add_row(row)
            except IndexError:
                pass