import BTAnalysisPyInterface as bta


def main():

    sources = bta.get_sources()
    source_count = 0

    # iterate through all sources selected in BT-Analysis
    for source in sources:
        source_count += 1

        # Define the columns of your custom table
        columns = [
            "Cycle number",
            "Q charge",
            "Q discharge",
            "Efficiency",
            "Energy",
            "Energy charge",
            "Energy discharge",
            "Duration",
            "Q charge normalized",
            "Q discharge normalized",
            "Q normalized",
        ]

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

        # Get DUT mass
        dut_mass = source.get_dut_type_mass()

        # Get the data of the summary per Cycle
        summary = source.get_summary(bta.SummaryType.PER_CYCLE)

        # Get all data needed for the custom table
        cycles = summary.get_column("Cycle number").get_values()
        q_charge = summary.get_column("Q charge").get_values()
        q_discharge = summary.get_column("Q discharge").get_values()
        q = summary.get_column("Q").get_values()
        efficiency = summary.get_column("Efficiency").get_values()
        energy = summary.get_column("Energy").get_values()
        energy_charge = summary.get_column("Energy charge").get_values()
        energy_discharge = summary.get_column("Energy discharge").get_values()
        duration = summary.get_column("Duration").get_values()

        # Include user variables if set
        try:
            var1 = summary.get_column("@Var1").get_values()
            columns.append("Var1")
            units.append("")
        except AttributeError:
            var1 = None
        try:
            var2 = summary.get_column("@Var2").get_values()
            columns.append("Var2")
            units.append("")
        except AttributeError:
            var2 = None
        try:
            var3 = summary.get_column("@Var3").get_values()
            columns.append("Var3")
            units.append("")
        except AttributeError:
            var3 = None
        try:
            var4 = summary.get_column("@Var4").get_values()
            columns.append("Var4")
            units.append("")
        except AttributeError:
            var4 = None
        try:
            var5 = summary.get_column("@Var5").get_values()
            columns.append("Var5")
            units.append("")
        except AttributeError:
            var5 = None
        try:
            var6 = summary.get_column("@Var6").get_values()
            columns.append("Var6")
            units.append("")
        except AttributeError:
            var6 = None
        try:
            var7 = summary.get_column("@Var7").get_values()
            columns.append("Var7")
            units.append("")
        except AttributeError:
            var7 = None
        try:
            var8 = summary.get_column("@Var8").get_values()
            columns.append("Var8")
            units.append("")
        except AttributeError:
            var8 = None
        try:
            var9 = summary.get_column("@Var9").get_values()
            columns.append("Var9")
            units.append("")
        except AttributeError:
            var9 = None
        try:
            var10 = summary.get_column("@Var10").get_values()
            columns.append("Var10")
            units.append("")
        except AttributeError:
            var10 = None

        # Add the custom table to BT-Analysis
        analysis_table = bta.add_table(
            columns=columns, units=units, title=f"Capacity Analysis {source_count}"
        )

        q_charge_normalized = []
        q_discharge_normalized = []
        q_normalized = []

        # Calculate normalized capacity
        for i in range(len(q_charge)):
            q_charge_normalized.append((q_charge[i] / 3.6) / dut_mass)
            q_discharge_normalized.append((q_discharge[i] / 3.6) / dut_mass)
            q_normalized.append((q[i] / 3.6) / dut_mass)

        # Add rows to the custom table
        for i in range(len(cycles)):
            row = {
                "Cycle number": cycles[i],
                "Q charge": q_charge[i],
                "Q discharge": q_discharge[i],
                "Efficiency": efficiency[i] * 100,
                "Energy": energy[i],
                "Energy charge": energy_charge[i],
                "Energy discharge": energy_discharge[i],
                "Duration": duration[i],
                "Q charge normalized": q_charge_normalized[i],
                "Q discharge normalized": q_discharge_normalized[i],
                "Q normalized": q_normalized[i],
            }
            if "Var1" in columns:
                row["Var1"] = var1[i]
            if "Var2" in columns:
                row["Var2"] = var2[i]
            if "Var3" in columns:
                row["Var3"] = var3[i]
            if "Var4" in columns:
                row["Var4"] = var4[i]
            if "Var5" in columns:
                row["Var5"] = var5[i]
            if "Var6" in columns:
                row["Var6"] = var6[i]
            if "Var7" in columns:
                row["Var7"] = var7[i]
            if "Var8" in columns:
                row["Var8"] = var8[i]
            if "Var9" in columns:
                row["Var9"] = var9[i]
            if "Var10" in columns:
                row["Var10"] = var10[i]

            analysis_table.add_row(row)
