use camino::Utf8Path;
use outfit::constants::ObjectNumber;
use outfit::error_models::ErrorModel;
use outfit::outfit::Outfit;
use outfit::trajectories::trajectory_file::TrajectoryFile;
use outfit::TrajectorySet;

#[test]
fn test_load_traj_from_parquet() {
    let mut env_state = Outfit::new("horizon:DE440", ErrorModel::FCCT14).unwrap();
    let path_file = Utf8Path::new("tests/data/trajectories.parquet");

    let ztf_observer = env_state.get_observer_from_mpc_code(&"I41".into());

    let traj_set =
        TrajectorySet::new_from_parquet(&mut env_state, path_file, ztf_observer, 0.5, 0.5, None)
            .unwrap();
    assert_eq!(traj_set.len(), 4);
    assert_eq!(traj_set.get(&ObjectNumber::Int(1)).unwrap().len(), 3);
}

#[test]
fn test_large_parquet() {
    let mut env_state = Outfit::new("horizon:DE440", ErrorModel::FCCT14).unwrap();
    let path_file = Utf8Path::new("tests/data/test_from_fink.parquet");

    let ztf_observer = env_state.get_observer_from_mpc_code(&"I41".into());

    let mut traj_set =
        TrajectorySet::new_from_parquet(&mut env_state, path_file, ztf_observer, 0.5, 0.5, None)
            .unwrap();

    assert_eq!(traj_set.len(), 2082);
    assert_eq!(traj_set.get(&ObjectNumber::Int(1)).unwrap().len(), 6);
    assert_eq!(traj_set.get(&ObjectNumber::Int(2)).unwrap().len(), 7);
    assert_eq!(traj_set.get(&ObjectNumber::Int(3)).unwrap().len(), 7);
    assert_eq!(traj_set.get(&ObjectNumber::Int(4)).unwrap().len(), 7);

    let path_file = Utf8Path::new("tests/data/trajectories.parquet");
    let rubin_observer = env_state.get_observer_from_mpc_code(&"X05".into());
    traj_set
        .add_from_parquet(&mut env_state, path_file, rubin_observer, 0.5, 0.5, None)
        .unwrap();

    assert_eq!(traj_set.len(), 2082);
    let traj = traj_set.get(&ObjectNumber::Int(1)).unwrap();
    assert_eq!(traj.len(), 9);
    assert_eq!(traj_set.get(&ObjectNumber::Int(2)).unwrap().len(), 10);
    assert_eq!(traj_set.get(&ObjectNumber::Int(3)).unwrap().len(), 10);
    assert_eq!(traj_set.get(&ObjectNumber::Int(4)).unwrap().len(), 10);

    let first_obs = traj.first().unwrap();
    assert_eq!(first_obs.ra, 0.5833713125827732);
    assert_eq!(first_obs.dec, 0.4110543900466954);
    assert_eq!(first_obs.time, 58789.138125000056);
    assert_eq!(
        first_obs.get_observer(&env_state).name,
        Some("Palomar Mountain--ZTF".into())
    );

    let second_obs = traj.get(6).unwrap();
    assert_eq!(second_obs.ra, 3.353185900126074);
    assert_eq!(second_obs.dec, -0.09911550289150597);
    assert_eq!(second_obs.time, 59396.0);
    assert_eq!(
        second_obs.get_observer(&env_state).name,
        Some("Simonyi Survey Telescope, Rubin Observatory".into())
    );
}
