from numba import njit
import numpy as np


class Problem:
    def __init__(self, problem):
        self.problem = problem
        self.output = None
        self.B, self.L, self.D, self.S, self.N, self.T, self.M, self.idx = self.read_data()
        self.remaining_libraries = np.array(list(range(self.L)))
        self.current_day = 0

    def read_data(self):
        filename = f'input/{self.problem}.txt'
        lines = [line.strip('\n') for line in open(filename, 'r')]
        B, L, D = map(int, lines[0].split())
        S = np.array(list(map(int, lines[1].split())))

        N = np.zeros(L)
        T = np.zeros(L)
        M = np.zeros(L)
        idx = -np.ones(shape=(L, B), dtype=int)
        for i in range(L):
            n, t, m = map(int, lines[2*i + 2].split())
            N[i] = n
            T[i] = t
            M[i] = m
            books = np.array(list(map(int, lines[2*i + 3].split())))
            book_values = S[books]
            sorted_idx = np.argsort(-book_values)
            idx[i][:books.shape[0]] = books[sorted_idx]

        return B, L, D, S, N, T, M, idx

    def run(self):
        info, books = njit_run(self.current_day, self.remaining_libraries, self.D, self.N, self.T, self.M, self.S, self.idx)
        return info, books

    def save(self, info, books):
        with open(f'output/{self.problem}.out', 'w') as f:
            for i in range(len(info)):
                for arr in [info, books]:
                    out = ' '.join(map(str, arr[i]))
                    f.write(f'{out}\n')


@njit(fastmath=True)
def njit_run(current_day, remaining_libraries, D, N, T, M, S, idx):
    info = []
    books = []
    score = 0
    while current_day < D:
        best_library, best_library_value, best_books = calculate_best_library(remaining_libraries, current_day, D, N, T, M, S, idx)
        if best_library_value <= 0:
            break
        score += best_library_value
        current_day += T[best_library]
        remaining_libraries[best_library] = -1
        remove_books(remaining_libraries, set(best_books), idx)
        if best_books.shape[0] != 0:
            info.append([best_library, best_books.shape[0]])
            books.append(best_books)
        print('Progress:', current_day/D)
    print('Score:', score)
    return info, books


@njit(fastmath=True)
def calculate_best_library(remaining_libraries, current_day, deadline, N, T, M, S, idx):
    best_library_value = -1

    for lib in remaining_libraries:
        if lib == -1:
            continue
        constructed_by = current_day + T[lib]
        remaining_days = max(0, deadline - constructed_by)
        lib_size = np.where(idx[lib] == -1)[0][0]
        books_by_deadline = int(min(lib_size, M[lib]*remaining_days))
        books = idx[lib][:books_by_deadline]

        library_value = 0
        for b in books:
            library_value += S[b]

        if library_value > best_library_value:
            best_library = lib
            best_library_value = library_value
            best_books = books

    return best_library, best_library_value, best_books


@njit(fastmath=True)
def remove_books(remaining_libraries, best_books, idx):
    for lib in remaining_libraries:
        if lib == -1:
            continue
        lib_size = np.where(idx[lib] == -1)[0][0]
        for i in range(lib_size):
            if idx[lib][i] in best_books:
                idx[lib][i] = -1
        new_idx = np.where(idx[lib] != -1)[0]
        idx[lib][:new_idx.shape[0]] = idx[lib][new_idx]
        idx[lib][new_idx.shape[0]:] = -1


if __name__ == '__main__':
    for problem in ['c']:
        print(problem)
        cls = Problem(problem)
        info, books = cls.run()
        cls.save(info, books)