use burn::data::dataset::Dataset;

use crate::dqn::prelude::TankItem;
use crate::ffi::prelude::*;
use rand::Rng;

// fn random_action() -> Action {
//     let mut rng = rand::thread_rng();
//     match rng.gen_range(0..2) {
//         0 => Action::AimLeft,
//         1 => Action::Forward,
//         _ => unreachable!(),
//     }
// }

// fn random_item() -> TankItem {
//     let mut previous_info=Info::default();
//     TankItem {
//         previous_state: todo!(),
//         new_state: todo!(),
//         action: todo!(),
//         reward: todo!(),
//     }
// }

pub struct FitDataset;

impl FitDataset {
    /// Get closer to the power station
    fn close_power_station() -> TankItem {
        let mut power_stations = Station::default();

        let mut previous_info = Info::default();
        let mut new_info = Info::default();
        let mut rng = rand::thread_rng();
        previous_info.player.power = rng.gen_range(0..2);
        new_info.player.power = previous_info.player.power;
        previous_info.player.angle = rng.gen_range(0..360);
        new_info.player.angle = previous_info.player.angle;

        TankItem {
            previous_state: todo!(),
            new_state: todo!(),
            action: Action::Forward,
            reward: todo!(),
        }
    }
    /// Flee from power station if power is high
    fn flee_power_station() -> TankItem {
        let mut previous_info = Info::default();
        TankItem {
            previous_state: todo!(),
            new_state: todo!(),
            action: Action::Backward,
            reward: todo!(),
        }
    }
}

impl Dataset<TankItem> for FitDataset {
    fn get(&self, _: usize) -> Option<TankItem> {
        let previous_state = todo!();
        let new_state = todo!();
        let action = Action::AimLeft;
        let reward = 0.0;
        Some(TankItem {
            previous_state,
            new_state,
            action,
            reward,
        })
    }

    fn len(&self) -> usize {
        1
    }
}
