Skip to content

Commit

Permalink
boost tracked, tests pass
Browse files Browse the repository at this point in the history
  • Loading branch information
scottleedavis committed Dec 4, 2024
1 parent 0701ad3 commit 1ca1504
Show file tree
Hide file tree
Showing 7 changed files with 115 additions and 104 deletions.
Empty file removed examples/summary.md
Empty file.
177 changes: 93 additions & 84 deletions src/convert.rs
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,7 @@ pub fn parse_frames(data: &Value, file: &mut dyn Write) -> Result<(), Box<dyn st
let empty_array: Vec<Value> = vec![];
let frames = data.as_array().unwrap_or(&empty_array);
let mut player_map: HashMap<String, String> = HashMap::new();
let mut player_actor_map: HashMap<String, String> = HashMap::new();
let mut team_map: HashMap<String, String> = HashMap::new();
let mut car_map: HashMap<String, String> = HashMap::new();
let mut car_boost_map: HashMap<String, i64> = HashMap::new();
Expand Down Expand Up @@ -244,35 +245,40 @@ pub fn parse_frames(data: &Value, file: &mut dyn Write) -> Result<(), Box<dyn st
}

} else if name == "Engine.Pawn:PlayerReplicationInfo" ||
name == "TAGame.CarComponent_TA:Vehicle" ||
name == "Engine.Pawn:PlayerReplicationInfo" {
if let Some(value_int) = update
.get("value")
.and_then(|value| value.get("flagged_int"))
.and_then(|value| value.get("int"))
.and_then(|int_value| int_value.as_i64())

.and_then(|int_value| int_value.as_i64())
{
if let Some(_name) = player_map.get(&value_int.to_string()) {
car_map.insert(actor_id.clone(),value_int.to_string() );
if value_int > 0 {
player_actor_map.insert(actor_id.clone(), value_int.to_string());
}

}
} else if name == "TAGame.CarComponent_TA:Vehicle" {
if let Some(value_int) = update
.get("value")
.and_then(|value| value.get("flagged_int"))
.and_then(|value| value.get("int"))
.and_then(|int_value| int_value.as_i64())

{
car_map.insert(actor_id.clone(),value_int.to_string() );
}
} else if name == "\"TAGame.CarComponent_Boost_TA:ReplicatedBoost\"" {
if let Some(cname) = car_map.get(&actor_id).map(String::as_str) {
if let Some(value_int) = update
.get("value")
.and_then(|value| value.get("boost"))
.and_then(|value| value.get("boostAmount"))
.and_then(|int_value| int_value.as_i64())

{
car_boost_map.insert(actor_id.clone(), value_int);
} else if name == "TAGame.CarComponent_Boost_TA:ReplicatedBoost" {
// Extract boost value
if let Some(value_int) = update
.get("value")
.and_then(|value| value.get("boost"))
.and_then(|value| value.get("boostAmount"))
.and_then(|boost| boost.as_i64())
{
if let Some(cname) = car_map.get(&actor_id).map(String::as_str) {
car_boost_map.insert(cname.to_string(), value_int);
}
} else {
//where is this boost
}

}
}
}
}
Expand All @@ -290,8 +296,10 @@ pub fn parse_frames(data: &Value, file: &mut dyn Write) -> Result<(), Box<dyn st
let obj_name = spawned.get("object_name").unwrap_or(&Value::Null).to_string();
if let Some(cname) = car_map.get(&actor_id).map(String::as_str) {
if obj_name == "\"Archetypes.Car.Car_Default\""{
let pname = player_map.get(cname).map(String::as_str).unwrap_or("Unknown");
let tname = team_map.get(cname).map(String::as_str).unwrap_or("Unknown");

let paname = player_actor_map.get(cname).map(String::as_str).unwrap_or("Unknown");
let pname = player_map.get(paname).map(String::as_str).unwrap_or("Unknown");
let tname = team_map.get(paname).map(String::as_str).unwrap_or("Unknown");
let boost = car_boost_map.get(cname).copied().unwrap_or(0);

let location_x = spawned.pointer("/initialization/location/x")
Expand Down Expand Up @@ -359,70 +367,65 @@ pub fn parse_frames(data: &Value, file: &mut dyn Write) -> Result<(), Box<dyn st

if name == "\"TAGame.RBActor_TA:ReplicatedRBState\"" {

// let value = serde_json::to_string(update.get("value").unwrap_or(&Value::Null))
// .unwrap_or_else(|_| "{}".to_string())
// .replace("\"", "\\\"");

if let Some(cname) = car_map.get(&actor_id).map(String::as_str) {
if cname != "Unknown" {
let pname = player_map.get(cname).map(String::as_str).unwrap_or("Unknown");
let tname = team_map.get(cname).map(String::as_str).unwrap_or("Unknown");
let boost = car_boost_map.get(cname).copied().unwrap_or(0);

let location_x = update.pointer("/value/rigid_body_state/location/x")
.and_then(Value::as_i64)
.unwrap_or(0);
let location_y = update.pointer("/value/rigid_body_state/location/y")
.and_then(Value::as_i64)
.unwrap_or(0);
let location_z = update.pointer("/value/rigid_body_state/location/z")
.and_then(Value::as_i64)
.unwrap_or(0);

let rotation_x = update.pointer("/value/rigid_body_state/rotation/quaternion/x")
.and_then(Value::as_f64)
.unwrap_or(0.0);
let rotation_y = update.pointer("/value/rigid_body_state/rotation/quaternion/y")
.and_then(Value::as_f64)
.unwrap_or(0.0);
let rotation_z = update.pointer("/value/rigid_body_state/rotation/quaternion/z")
.and_then(Value::as_f64)
.unwrap_or(0.0);
let rotation_w = update.pointer("/value/rigid_body_state/rotation/quaternion/w")
.and_then(Value::as_f64)
.unwrap_or(0.0);

let angular_velocity_x = update.pointer("/value/rigid_body_state/angular_velocity/x")
.and_then(Value::as_i64)
.unwrap_or(0);
let angular_velocity_y = update.pointer("/value/rigid_body_state/angular_velocity/y")
.and_then(Value::as_i64)
.unwrap_or(0);
let angular_velocity_z = update.pointer("/value/rigid_body_state/angular_velocity/z")
.and_then(Value::as_i64)
.unwrap_or(0);

let linear_velocity_x = update.pointer("/value/rigid_body_state/linear_velocity/x")
.and_then(Value::as_f64)
.unwrap_or(0.0);
let linear_velocity_y = update.pointer("/value/rigid_body_state/linear_velocity/y")
.and_then(Value::as_f64)
.unwrap_or(0.0);
let linear_velocity_z = update.pointer("/value/rigid_body_state/linear_velocity/z")
.and_then(Value::as_f64)
.unwrap_or(0.0);

lines.push(format!(
"{},{},{},\"{}\",{},{},{},{},{},{},{},{},{},{},{},{},{},{}",
frame_index, time, tname, pname, boost,
location_x, location_y, location_z,
rotation_x, rotation_y, rotation_z, rotation_w,
angular_velocity_x, angular_velocity_y, angular_velocity_z,
linear_velocity_x, linear_velocity_y, linear_velocity_z
));
}
} else if actor_id == ball_id {
if let Some(paname) = player_actor_map.get(&actor_id).map(String::as_str) {

let pname = player_map.get(paname).map(String::as_str).unwrap_or("Unknown");
let tname = team_map.get(paname).map(String::as_str).unwrap_or("Unknown");
let boost = car_boost_map.get(&actor_id).copied().unwrap_or(0);

let location_x = update.pointer("/value/rigid_body_state/location/x")
.and_then(Value::as_i64)
.unwrap_or(0);
let location_y = update.pointer("/value/rigid_body_state/location/y")
.and_then(Value::as_i64)
.unwrap_or(0);
let location_z = update.pointer("/value/rigid_body_state/location/z")
.and_then(Value::as_i64)
.unwrap_or(0);

let rotation_x = update.pointer("/value/rigid_body_state/rotation/quaternion/x")
.and_then(Value::as_f64)
.unwrap_or(0.0);
let rotation_y = update.pointer("/value/rigid_body_state/rotation/quaternion/y")
.and_then(Value::as_f64)
.unwrap_or(0.0);
let rotation_z = update.pointer("/value/rigid_body_state/rotation/quaternion/z")
.and_then(Value::as_f64)
.unwrap_or(0.0);
let rotation_w = update.pointer("/value/rigid_body_state/rotation/quaternion/w")
.and_then(Value::as_f64)
.unwrap_or(0.0);

let angular_velocity_x = update.pointer("/value/rigid_body_state/angular_velocity/x")
.and_then(Value::as_i64)
.unwrap_or(0);
let angular_velocity_y = update.pointer("/value/rigid_body_state/angular_velocity/y")
.and_then(Value::as_i64)
.unwrap_or(0);
let angular_velocity_z = update.pointer("/value/rigid_body_state/angular_velocity/z")
.and_then(Value::as_i64)
.unwrap_or(0);

let linear_velocity_x = update.pointer("/value/rigid_body_state/linear_velocity/x")
.and_then(Value::as_f64)
.unwrap_or(0.0);
let linear_velocity_y = update.pointer("/value/rigid_body_state/linear_velocity/y")
.and_then(Value::as_f64)
.unwrap_or(0.0);
let linear_velocity_z = update.pointer("/value/rigid_body_state/linear_velocity/z")
.and_then(Value::as_f64)
.unwrap_or(0.0);

lines.push(format!(
"{},{},{},\"{}\",{},{},{},{},{},{},{},{},{},{},{},{},{},{}",
frame_index, time, tname, pname, boost,
location_x, location_y, location_z,
rotation_x, rotation_y, rotation_z, rotation_w,
angular_velocity_x, angular_velocity_y, angular_velocity_z,
linear_velocity_x, linear_velocity_y, linear_velocity_z
));

} else if actor_id == ball_id {

let location_x = update.pointer("/value/rigid_body_state/location/x")
.and_then(Value::as_i64)
Expand Down Expand Up @@ -483,6 +486,12 @@ pub fn parse_frames(data: &Value, file: &mut dyn Write) -> Result<(), Box<dyn st
}
}

// println!("Player Map: {:#?}", player_map);
// println!("Player Actor Map: {:#?}", player_actor_map);
// println!("Team Map: {:?}", team_map);
// println!("Car Map: {:#?}", car_map);
// println!("Boost Map: {:#?}", car_boost_map);

let mut writer = BufWriter::new(file);
for line in &lines {
writeln!(writer, "{}", line)?;
Expand Down
1 change: 1 addition & 0 deletions src/extract.rs
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@ pub fn extract_replay(input: &str) -> io::Result<()> {
} else {
"./rattletrap"
})
.arg("--compact")
.arg("--input")
.arg(input)
.arg("--output")
Expand Down
14 changes: 7 additions & 7 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ fn main() {
let args: Vec<String> = env::args().collect();

if args.len() < 2 {
println!("Usage: rl_replay_ai <command> [options]");
println!("Usage: rattlebrain <command> [options]");
println!("Commands:");
println!(" analysis <path/some.replay> - Analyze replay data. (runs extract, convert, plot and ai in sequence)");
println!(" extract <path/some.replay> - Extract replay data to CSV.");
Expand All @@ -27,7 +27,7 @@ fn main() {
match command.as_str() {
"extract" => {
if args.len() < 3 {
println!("Usage: rocket-league-replay-ai-analysis extract <input>");
println!("Usage: rattlebrain extract <input>");
return;
}
let input = &args[2];
Expand All @@ -39,7 +39,7 @@ fn main() {
}
"convert" => {
if args.len() < 3 {
println!("Usage: rocket-league-replay-ai-analysis convert <input>");
println!("Usage: rattlebrain convert <input>");
return;
}
let input = &args[2];
Expand Down Expand Up @@ -71,7 +71,7 @@ fn main() {
}
"plot" => {
if args.len() < 3 {
println!("Usage: rocket-league-replay-ai-analysis plot <csv>");
println!("Usage: rattlebrain plot <csv>");
return;
}
let csv_file = &args[2];
Expand All @@ -83,7 +83,7 @@ fn main() {
}
"analysis" => {
if args.len() < 3 {
println!("Usage: rocket-league-replay-ai-analysis analysis <file.replay>");
println!("Usage: rattlebrain analysis <file.replay>");
return;
}
let input = &args[2];
Expand Down Expand Up @@ -138,7 +138,7 @@ fn main() {
}
"ai" => {
if args.len() < 3 {
println!("Usage: rocket-league-replay-ai-analysis ai <input>");
println!("Usage: rattlebrain ai <input>");
return;
}
let query = &args[2];
Expand All @@ -150,7 +150,7 @@ fn main() {
}
_ => {
println!("Unknown command: {}", command);
println!("Usage: rocket-league-replay-ai-analysis <command> [options]");
println!("Usage: rattlebrain <command> [options]");
}
}
}
24 changes: 12 additions & 12 deletions src/plot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,22 +8,22 @@ use std::collections::HashMap;

#[derive(Debug, Deserialize)]
pub struct GameState {
time: f64,
// time: f64,
team: Option<u32>,
player_name: String,
location_x: f64,
location_y: f64,
location_z: f64,
rotation_x: f64,
rotation_y: f64,
rotation_z: f64,
rotation_w: f64,
angular_velocity_x: f64,
angular_velocity_y: f64,
angular_velocity_z: f64,
linear_velocity_x: f64,
linear_velocity_y: f64,
linear_velocity_z: f64,
// location_z: f64,
// rotation_x: f64,
// rotation_y: f64,
// rotation_z: f64,
// rotation_w: f64,
// angular_velocity_x: f64,
// angular_velocity_y: f64,
// angular_velocity_z: f64,
// linear_velocity_x: f64,
// linear_velocity_y: f64,
// linear_velocity_z: f64,
}

pub fn plot_csv(file_path: &str) -> Result<Vec<GameState>, Box<dyn Error>> {
Expand Down
1 change: 1 addition & 0 deletions tests/extract_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ async fn test_extract_replay_json_schema_validation() {

// Run `rattletrap` to extract replay data to JSON
let output_status = Command::new(rattletrap_path)
.arg("--compact")
.arg("--input")
.arg(input_replay)
.arg("--output")
Expand Down
2 changes: 1 addition & 1 deletion tests/main_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ mod main_tests {
.expect("Failed to execute process");

let stdout = String::from_utf8(output.stdout).unwrap();
assert!(stdout.contains("Usage: rl_replay_ai <command> [options]"));
assert!(stdout.contains("Usage: rattlebrain <command> [options]"));
}

#[test]
Expand Down

0 comments on commit 1ca1504

Please sign in to comment.