import pandas as pd

from .chart import Chart


class Table(Chart):
    def map_module_result(self):
        raw_data = self._data_retrieval_task.wait_for_result()
        df = pd.DataFrame(
            raw_data, columns=["segment", "attribute_name", "value", "type"]
        )

        # saving initial sorting order
        sorted_segments = df["segment"].drop_duplicates().tolist()

        column_types = (
            df[["attribute_name", "type"]]
            .drop_duplicates()
            .set_index("attribute_name")
            .to_dict()["type"]  # either "numeric" or "percent"
        )
        df = df.pivot_table(
            values="value", index="segment", columns="attribute_name"
        ).reset_index()  # columns: ["segment", "attribute1", "attribute2", ...]
        columns = []
        for i, column in enumerate(df.columns):
            columns.append(
                {
                    "id": column,
                    "key": f"{i}",
                    "label": self.conf["inputs"]["attribute_labels"].get(
                        column, column.capitalize()
                    ),
                    "type": "number" if i >= 1 else "string",
                    "range": [min(df[column].min(), 0.0), df[column].max()]
                    if i >= 1
                    else None,
                    "valueFormat": f"{{{i}}}%"
                    if column_types.get(column) == "percent"
                    else f"{{{i}}}",
                }
            )
        data = []  # [{'0': 'Most Intense Superfans', '1': 3.14, '2': 42, ...}]
        for _, row in df.iterrows():
            data_row = {}
            for col in columns:
                value = row[col["id"]]
                data_row[col["key"]] = value
            data.append(data_row)

        # coming back to initial sort order by first column
        data.sort(key=lambda row: sorted_segments.index(row["0"]))
        return columns, data

    def fill_json_template(self):
        columns, data = self.map_module_result()
        return {
            "id": self.conf["id"],
            "type": "Table",
            "label": self.conf["label"],
            "options": columns,
            "data": data,
        }