import numpy as np
import pandas as pd

from openlog.datamodel.assay.generic_assay import GenericAssay


def label_objects_1d(binary_sequence):
    res = []
    label = 0
    for i, elt in enumerate(binary_sequence):
        if i == 0:
            res.append(label)
            previous = elt
            continue

        if elt == 0 or (elt == 1 and previous == 0):
            label += 1
            previous = elt
            res.append(label)
        elif elt == 1 and previous == 1:
            previous = elt
            res.append(label)

    return res


def merge_contiguous_items(x: pd.DataFrame, y: pd.DataFrame) -> tuple:
    """
    Merge contiguous categories.
    """

    # create a key based on combination of columns values
    df = x.copy()
    df["y"] = y.apply(lambda x: "_".join([str(elt) for elt in x]), axis=1)
    df = df[["from", "to", "y"]]
    df = df.sort_values("from")
    # identify contiguous blocks
    df["shift_m"] = df["y"].shift(-1)
    df["shift_p"] = df["y"].shift(1)
    df["contiguous"] = np.where(
        (df["y"] == df["shift_m"]) | (df["y"] == df["shift_p"]), 1, 0
    )

    # label blocks
    df["block"] = label_objects_1d(df["contiguous"])

    if df["block"].nunique() == 1:
        res = df.copy()
    else:
        # merge by block
        res = (
            df.groupby(["y", "block"]).agg({"from": np.min, "to": np.max}).reset_index()
        )
    res = res.sort_values("from")
    # dispatch columns from key
    new_x = res[["from", "to"]].values
    new_y = res["y"].str.split("_", expand=True).values
    return (new_x, new_y)


class CategoryMerger:
    def __init__(self, assay: GenericAssay, columns: [str]):
        """
        Class merging contiguous categories within an extended categorical assay.
        Columns are a list of category columns names.
        If multiple columns are provided the combination of values is take into account (concatened to have a unique key).

        """
        self.assay = assay
        self.columns = columns

    def get_values(self, columns=None) -> tuple:
        """
        Return merged data.
        """
        if columns is None:
            columns = self.columns

        y = [
            self.assay.get_all_values(column, remove_none=True)[1] for column in columns
        ]
        y = pd.concat([pd.Series(elt) for elt in y], axis=1)
        x = self.assay.get_all_values(columns[0], remove_none=True)[0]
        x = x.astype(float)

        if x.shape[0] == 0:
            y = y.values
            return x, y

        x = pd.DataFrame(x, columns=["from", "to"])
        new_x, new_y = merge_contiguous_items(x, y)
        return (new_x, new_y)
