import json
import os
import random
import numpy as np


def plant_patterns(n, patterns, prob):
    data = list()
    for _ in range(n):
        t = set()
        trial = 5
        while len(t) == 0 and trial > 0:
            trial = trial - 1
            for j in range(len(patterns)):
                if random.uniform(0, 1.0) <= prob[j]:
                    for x in patterns[j]:
                        t.add(x)
        data.append(sorted(list(t)))
    return data


def generate_one_pattern(first, last):
    length = random.randint(2, 5)
    return sorted(list(set([random.randint(first, last) for _ in range(length)])))


def generate_k_pattern(first, last, k):
    return [generate_one_pattern(first, last) for _ in range(k)]


def generate_k_pattern_without_redundancy(first, last, k):
    patterns = list()

    while len(patterns) != k:
        x = generate_one_pattern(first, last)
        while len([y for y in patterns if len(set(x).intersection(set(y))) > 1]) > 0:
            x = generate_one_pattern(first, last)
        patterns.append(x)

    return patterns


def get_assignment_matrix(m, k, p=0.3):
    return np.random.choice(a=[False, True], size=(m, k), p=[1.0 - p, p])


def make_classification_data(n_k, dim, k, s, lb=0.15, ub=0.3, additive_noise_level=0.05, destructive_noise_level=0.05, p_assignment=0.3):
    S = generate_k_pattern_without_redundancy(0, dim, s)
    A = get_assignment_matrix(k, s, p_assignment)

    X, Y = [], []
    singletons = [[i] for i in range(dim)]
    singletons_fr = [additive_noise_level] * dim
    for k in range(k):
        S_k = [S[i] for i in range(len(A[k])) if A[k][i] == True]
        P_k = [random.uniform(lb, ub) for _ in range(len(S_k))]
        X_k = plant_patterns(n_k, singletons + S_k, singletons_fr + P_k)
        X += X_k
        Y += [k] * n_k

    XY = list(zip(X, Y))
    random.shuffle(XY)
    X, Y = zip(*XY)
    X = destructive_noise(X, destructive_noise_level)
    return X, Y, S, A


def to_sparse_binary_matrix_format(xs):
    return [[i for i in range(len(x)) if x[i] == True] for x in xs]


def make_noise(n, d, k=2, p=0.5):
    y = np.random.randint(0, k, size=n)
    A = to_sparse_binary_matrix_format(np.random.rand(n, d) <= p)
    return A, y.tolist()


def plant_upto_k_patterns(n, patterns, prob, lb, ub):
    prob = np.array(prob)
    prob /= prob.sum()

    data = list()
    elems = list(range(len(patterns)))
    for _ in range(n):
        t = set()
        K = random.randint(lb, ub)
        for e in list(np.random.choice(elems, K, p=prob)):
            for y in patterns[e]:
                t.add(y)
        data.append(sorted(list(t)))
    return data


def destructive_noise(X, level):
    return [[i for i in x if random.uniform(0, 1) > level] for x in X]


def make_classification_data_discrete_sampler(n_k, dim, k, s, lb=0.15, ub=0.3, additive_noise_level=0.05, destructive_noise_level=0.05, p_assignment=0.3, num_lb=2, num_ub=5):
    S = generate_k_pattern_without_redundancy(0, dim, s)
    A = get_assignment_matrix(k, s, p_assignment)

    X, Y = [], []
    singletons = [[i] for i in range(dim)]
    singletons_fr = [additive_noise_level] * dim
    for k in range(k):
        S_k = [S[i] for i in range(len(A[k])) if A[k][i] == True]
        P_k = [random.uniform(lb, ub) for _ in range(len(S_k))]
        X_k = plant_patterns(n_k, singletons, singletons_fr)
        if len(S_k) > 0:
            XX_k = plant_upto_k_patterns(n_k, S_k, P_k, num_lb, num_ub)
            X_k = [list(set(x+y)) for x, y in zip(X_k, XX_k)]
        X += X_k
        Y += [k] * n_k

    XY = list(zip(X, Y))
    random.shuffle(XY)
    X, Y = zip(*XY)
    X = destructive_noise(X, destructive_noise_level)

    return X, Y, S, A


class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)


n_classes = 2
n_samples = 5000  # per class
p_plant = [0.2, 0.35]
trials = 20
n_dim = 100
n_patterns = 50
additive_noise_level = 0.05
destructive_noise_level = 0.0

try:
    os.mkdir("/synth")
except:
    pass

for i in range(3):
    for p in range(0, 110, 10):
        x, y, s, a = make_classification_data(
            n_samples,
            n_dim,
            n_classes,
            n_patterns,
            p_plant[0], p_plant[1],
            additive_noise_level,
            destructive_noise_level,
            p / 100.)
        print(json.dumps({"data": x, "labels": y, "patterns": s, "assignment": a}, cls=NumpyEncoder),
        file=open("/synth/synth-{}-{}.json".format(p, i), "w+"))
