1use std::fs;
2use std::collections::BTreeMap;
3use serde_derive::{Deserialize, Serialize};
4use crate::grid::Grid;
5
6#[derive(Deserialize, Serialize, Debug)]
7pub struct Data {
8 pub train: Vec<IO>,
9 pub test: Vec<IO>
10}
11
12#[derive(Deserialize, Serialize, Debug)]
13pub struct IO {
14 pub input: Vec<Vec<usize>>,
15 pub output: Option<Vec<Vec<usize>>>,
16}
17
18#[derive(Deserialize, Serialize, Debug)]
19struct Training {
20 #[serde(flatten)]
22 file: BTreeMap<String, Data>,
23}
24
25#[derive(Deserialize, Serialize, Debug)]
26struct Solution {
27 #[serde(flatten)]
29 file: BTreeMap<String, Vec<Vec<Vec<usize>>>>,
30}
31
32#[derive(Deserialize, Serialize, Debug)]
33struct Input {
34 input: Vec<Vec<usize>>,
35}
36
37#[derive(Deserialize, Serialize, Debug)]
38pub struct Output {
39 #[serde(flatten)]
40 file: BTreeMap<String, Vec<OutputData>>,
41}
42
43#[derive(Deserialize, Serialize, Clone, Debug)]
44pub struct OutputData {
45 attempt_1: Vec<Vec<usize>>,
46 attempt_2: Vec<Vec<usize>>,
47}
48
49pub fn add_dummy_output(file: &str, no_answers: usize, results: &mut BTreeMap<String, Vec<OutputData>>) {
50 let attempt = vec![vec![0, 0], vec![0, 0]];
51 let output_data: Vec<OutputData> = (0 .. no_answers).map(|_| OutputData { attempt_1: attempt.clone(), attempt_2: attempt.clone() }).collect();
52
53 results.insert(file.to_string(), output_data);
54 }
56
57pub fn add_real_output(file: &str, answers: &[Grid], results: &mut BTreeMap<String, Vec<OutputData>>) {
58 let output_data: Vec<OutputData> = answers.iter().map(|ans| OutputData { attempt_1: ans.to_vec(), attempt_2: ans.to_vec() }).collect();
59
60 results.insert(file.to_string(), output_data);
61 }
63
64pub fn create_output(live: bool, file: &BTreeMap<String, Vec<OutputData>>) {
65 let output = Output { file: file.clone() };
66
67 match serde_json::to_string(&output) {
68 Ok(data) => {
69 fs::write(if live { "submission.json" } else { "/kaggle/working/submission.json" }, data).expect("Unable to write file");
72 },
73 Err(e) => eprintln!("Failed to write submission.json: {e}"),
74 }
75}
76
77pub fn dir(dir: &str) -> Vec<String> {
78 let paths = fs::read_dir(dir).unwrap();
79
80 paths.map(|p| p.unwrap().path().display().to_string()).collect()
81}
82
83pub fn load_files(data: &str) -> BTreeMap<String, Data> {
84 let prefix = if std::path::Path::new("input/arc-prize-2025").is_dir() {
102 "input/arc-prize-2025/arc-agi"
103 } else if std::path::Path::new("/kaggle/input/arc-prize-2025").is_dir() {
104 "/kaggle/input/arc-prize-2025/arc-agi"
105 } else {
106 "/kaggle/input/arc-prize-2025/arc-agi"
107 };
108let training_suffix = "challenges.json";
110 let solution_suffix = "solutions.json";
111 let training = format!("{}_{}_{}", prefix, data, training_suffix);
112 let solution = format!("{}_{}_{}", prefix, data, solution_suffix);
113
114 let mut tdata =
115 match std::fs::read_to_string(&training) {
116 Ok(cf) => {
117 let training: Result<Training, _> = serde_json::from_str(&cf);
118
119 match training {
120 Ok(files) => {
121 files.file
122 },
123 Err(e) => {
124 eprintln!("{e:?}");
125
126 BTreeMap::new()
127 }
128 }
129 },
130 Err(e) => {
131 eprintln!("{training}: {e}");
132
133 BTreeMap::new()
134 }
135 };
136
137 let sdata =
138 match std::fs::read_to_string(&solution) {
139 Ok(sf) => {
140 let solutions: Result<Solution, _> = serde_json::from_str(&sf);
141
142 match solutions {
143 Ok(files) => {
144 files.file
145 },
146 Err(e) => {
147 eprintln!("{e:?}");
148
149 BTreeMap::new()
150 }
151 }
152 },
153 Err(e) => {
154 eprintln!("{solution}: {e}");
155
156 BTreeMap::new()
157 }
158 };
159
160 if !tdata.is_empty() {
164 for (file, data) in tdata.iter_mut() {
165 if sdata.contains_key(file) {
166 for (i, tests) in data.test.iter_mut().enumerate() {
167 if tests.output.is_none() {
168 tests.output = Some(sdata[file][i].clone());
169 }
170 }
171 }
172 }
173 }
174
175 tdata
176}