#include "jngen.h"
#include <bits/stdc++.h>

using namespace std;

const int MIN_N = 2;
const int MAX_N = 100000;

template <typename T> void append(vector<T> &dest, const vector<T> &orig) {
	dest.insert(dest.end(), orig.begin(), orig.end());
}

string output_tc_file(int n, int type = 0) {
    ostringstream oss;
    int m = rnd.next(1, n);
    oss << n << " " << m << "\n";
    
    for (int i = 1; i <= n; i++) {
        oss << rnd.next(0, 1) << (i == n ? "" : " ");
    }
    oss << "\n";
    
    Tree tree;
    if (type == 0) tree = Tree::random(n);
    else if (type == 1) tree = Tree::star(n);
    else if (type == 2) tree = Tree::bamboo(n);
    else if (type == 3) tree = Tree::caterpillar(n, max(1, rnd.next(1, n)));
    else tree = Tree::random(n);
    
    vector<pair<int, int>> edges = tree.edges();
    for (auto e : edges) {
        oss << e.first + 1 << " " << e.second + 1 << "\n";
    }
    return oss.str();
}

string generate_random_tc_file(int min_n, int max_n, int type) {
    int n = rnd.next(min_n, max_n);
    return output_tc_file(n, type);
}

string generate_manual_sample(int n, int m, vector<int> cats, const vector<pair<int, int>>& edges) {
    ostringstream oss;
    oss << n << " " << m << "\n";
    for (int i = 0; i < n; i++) {
        oss << cats[i] << (i == n - 1 ? "" : " ");
    }
    oss << "\n";
    for (auto e : edges) {
        oss << e.first << " " << e.second << "\n";
    }
    return oss.str();
}

vector<string> generate_sample_tests() {
	vector<string> tests;
    tests.push_back(generate_manual_sample(4, 1, {1, 1, 0, 0}, {{1, 2}, {1, 3}, {1, 4}}));
    tests.push_back(generate_manual_sample(7, 1, {1, 0, 1, 1, 0, 0, 0}, {{1, 2}, {1, 3}, {2, 4}, {2, 5}, {3, 6}, {3, 7}}));
	return tests;
}

vector<string> generate_manual_tests() {
	vector<string> tests;
	tests.push_back(output_tc_file(10));
	return tests;
}

vector<string> generate_random_tests() { 
    vector<string> tests;
    
    for (int i = 0; i < 5; i++) {
        tests.push_back(generate_random_tc_file(MIN_N, 10, i % 4));
    }
    
    for (int i = 0; i < 5; i++) {
        tests.push_back(generate_random_tc_file(10, 1000, i % 4));
    }
    
    for (int i = 0; i < 5; i++) {
        tests.push_back(generate_random_tc_file(10000, 20000, i % 4));
    }

    return tests;
}

vector<string> generate_extreme_tests(){
    vector<string> tests;   
    tests.push_back(output_tc_file(MAX_N, 0));
    tests.push_back(output_tc_file(MAX_N, 1));
    tests.push_back(output_tc_file(MAX_N, 2));
    tests.push_back(output_tc_file(MAX_N, 3));
    return tests;
}

int main(int argc, char *argv[]) {
	registerGen(argc, argv, 1);
	vector<string> tests;
	size_t test = 0;
	append(tests, generate_sample_tests());
	append(tests, generate_manual_tests());
	append(tests, generate_random_tests());
	append(tests, generate_extreme_tests());
	for (const auto &t : tests) {
		startTest(++test);
		cout << t;
	}
	return 0;
}