arc_agi/
data.rs

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(rename = "file")]
21    #[serde(flatten)]
22    file: BTreeMap<String, Data>,
23}
24
25#[derive(Deserialize, Serialize, Debug)]
26struct Solution {
27    //#[serde(rename = "file")]
28    #[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    //results.insert(serde_json::to_string_pretty(file).unwrap(), output_data);
55}
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    //results.insert(serde_json::to_string_pretty(file).unwrap(), output_data);
62}
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            //let data = format!("{data}\n");
70
71            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    /*
85    let prefix = if std::path::Path::new("input/arc-prize-2024").is_dir() {
86        "input/arc-prize-2024/arc-agi"
87    } else if std::path::Path::new("kaggle/input").is_dir() {
88        "kaggle/input/arc-prize-2025/arc-agi"
89    } else {
90        "/kaggle/input/arc-prize-2025/arc-agi"
91    */
92    /*
93    let prefix = if std::path::Path::new("input/arc-prize-2025").is_dir() {
94        "input/arc-prize-2025/arc-agi"
95    } else if std::path::Path::new("kaggle/working").is_dir() {
96        "kaggle/working/arc-agi"
97    } else {
98        "/kaggle/working/arc-agi"
99    };
100    */
101    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    };
108//eprintln!("#### {prefix}");
109    let 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    //println!("{}", cfdata.len());
161
162    // If we have data splice in test answers if any
163    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}