import numpy as np


# Paul Tol's color scheme
RED = (204/255, 102/255, 119/255)   # CC6677
GREEN = (17/255, 119/255, 51/255)   # 117733
BLUE = (136/255, 204/255, 238/255)  # 88CCEE
GRAY = (221/255, 221/255, 221/255)  # DDDDDD


def assignment_to_model(assignment, sample_size):
    """Convert an assignment to a model."""
    return [x + 1 if x in assignment else -(x + 1) for x in range(sample_size)]

def model_to_assignment(model):
    """Convert a model to an assignment."""
    return [x - 1 for x in model if x > 0]


def bitmask_to_list(bitmask, sample_size):
    """Convert bitmask to list of integers."""
    if isinstance(bitmask, str):
        bitmask = int(bitmask, 2)
    return [i for i in range(sample_size) if (bitmask >> i) & 1]

def list_to_bitmask(int_list, sample_size):
    """Convert list of integers to bitmask."""
    bitmask = 0

    for i in int_list:
        bitmask |= 1 << i

    return format(bitmask, f'0{sample_size}b')

def calc_samples_prototype(samples, prototype, axis=0, flatten=True):
    samples = preprocess_samples(samples, flatten)

    if prototype == "mean":
        return np.mean(samples, axis=axis)
    elif prototype == "median":
        return np.median(samples, axis=axis)
    else:
        raise ValueError(f"Unknown prototype: {prototype}. Use 'mean' or 'median'.")

def calc_samples_std(samples, axis=0, flatten=True):
    samples = preprocess_samples(samples, flatten)
    return np.std(samples, axis=axis)

def preprocess_samples(samples, flatten=True):
    samples = np.array(samples)

    if flatten:
        samples = samples.reshape(samples.shape[0], -1)

    return samples

def average_score(elements, score_func, **kwargs):
    score = 0

    for element in elements:
        score += score_func(element, **kwargs)

    return score / len(elements)

def interp_color(c1, c2, t):
    """Interpolate between two RGB tuples (0-1 range) by t (0=first, 1=second)."""
    return tuple((1-t)*a + t*b for a, b in zip(c1, c2))

def group_samples_by_label(samples, labels, num_labels):
    samples_by_label = {n: [] for n in range(num_labels)}

    for sample, label in zip(samples, labels):
        label = np.argmax(label)
        samples_by_label[label].append(sample)

    return samples_by_label
