import numpy as np
import pandas as pd

from openlog.datamodel.assay.generic_assay import GenericAssay


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

        if elt == 0 or (elt == 1 and previous == 0):
            label += 1
            res.append(label)
        elif elt == 1 and previous == 1 and values[i] == previous_value:
            res.append(label)
        elif elt == 1 and previous == 1 and values[i] != previous_value:
            label += 1
            res.append(label)

        previous = elt
        previous_value = values[i]

    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").reset_index()
    # 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"], df["y"])

    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:
        """
        Merge continuous data.
        Return a tuple of 2d ndarray : ([from,to], [y])
        """
        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)

    def get_altitude(self, columns=None, planned: bool = False) -> tuple:
        """
        Merge data but return resulting altitudes.
        Return a tuple of 1d ndarray : (from, to)
        """

        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_altitude(planned)

        if x is None:
            return None

        x = pd.DataFrame({"to": x[0].astype(float), "from": x[1].astype(float)})

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

        new_x, _ = merge_contiguous_items(x, y)
        res = (new_x[:, 1], new_x[:, 0])
        # sort in descending order
        res = [np.sort(x)[::-1] for x in res]

        return res
