use crate::actors::*;
use clap::{arg_enum, Clap};
use log::Level;
use simple_logger::init_with_level;

#[derive(Clap)]
pub struct Opts {
    /// Add more flags to increase verbosity to log at the debug or trace level.
    #[clap(
        short = "v",
        long = "verbose",
        parse(from_occurrences),
        conflicts_with("quiet")
    )]
    pub verbose: u8,
    /// Add more flags to decrease verbosity to the warn, error, or silent level.
    #[clap(short = "q", long = "quiet", parse(from_occurrences))]
    pub quiet: u8,
    #[clap(subcommand)]
    pub subcmd: SubCommand,
}

#[derive(Clap)]
pub enum SubCommand {
    /// Play the game, without training an actor
    Play(Play),
    /// Train an actor to play the game.
    Train(Train),
}

#[derive(Clap)]
pub struct Play {
    /// Disable gravity. Useful for debugging.
    #[clap(short = "G", long = "no-gravity")]
    pub no_gravity: bool,
}

#[derive(Clap)]
pub struct Train {
    /// Which agent to use for training.
    pub agent: Agent,

    /// The rate at which temporal agents learn at.
    #[clap(short = "a", long = "alpha", default_value = "0.3")]
    pub learning_rate: f64,

    /// The rate at which agents explore new actions. Range is [0, 1].
    #[clap(short = "e", long = "epsilon", default_value = "0.1")]
    pub exploration_prob: f64,

    /// The discount rate for future states.
    #[clap(short = "g", long = "gamma", default_value = "0.7")]
    pub discount_rate: f64,

    /// Stop learning during the evaluation of the agent. This sets the learning
    /// rate to 0 when displaying the results.
    #[clap(short = "L", long = "no-learn")]
    pub no_learn_during_evaluation: bool,

    /// Stop exploring during the evaluation of the agent. This sets the
    /// exploration rate to 0 when displaying the results.
    #[clap(short = "E", long = "no-explore")]
    pub no_explore_during_evaluation: bool,

    /// Number of episodes to train the agent
    #[clap(short = "n", long = "num", default_value = "10")]
    pub episodes: usize,
    // #[clap(long = "use-epsilon-decreasing")]
    // pub epsilon_decreasing: bool,
}

arg_enum! {
    #[derive(Debug)]
    pub enum Agent {
        QLearning,
        ApproximateQLearning,
        HeuristicGenetic
    }
}

pub fn init_verbosity(opts: &Opts) -> Result<(), Box<dyn std::error::Error>> {
    match (opts.quiet, opts.verbose) {
        (0, 0) => init_with_level(Level::Info)?,
        (0, 1) => init_with_level(Level::Debug)?,
        (0, _) => init_with_level(Level::Trace)?,
        (1, 0) => init_with_level(Level::Warn)?,
        (2, 0) => init_with_level(Level::Error)?,
        _ => (),
    };

    Ok(())
}
