use crate::burnrl::environment::TrictracEnvironment;
use crate::burnrl::utils::{soft_update_linear, Config};
use burn::backend::{ndarray::NdArrayDevice, NdArray};
use burn::module::Module;
use burn::nn::{Linear, LinearConfig};
use burn::optim::AdamWConfig;
use burn::record::{CompactRecorder, Recorder};
use burn::tensor::activation::{relu, softmax};
use burn::tensor::backend::{AutodiffBackend, Backend};
use burn::tensor::Tensor;
use burn_rl::agent::{SACActor, SACCritic, SACNets, SACOptimizer, SACTrainingConfig, SAC};
use burn_rl::base::{Action, Agent, ElemType, Environment, Memory, Model, State};
use std::time::SystemTime;

#[derive(Module, Debug)]
pub struct Actor<B: Backend> {
    linear_0: Linear<B>,
    linear_1: Linear<B>,
    linear_2: Linear<B>,
}

impl<B: Backend> Actor<B> {
    pub fn new(input_size: usize, dense_size: usize, output_size: usize) -> Self {
        Self {
            linear_0: LinearConfig::new(input_size, dense_size).init(&Default::default()),
            linear_1: LinearConfig::new(dense_size, dense_size).init(&Default::default()),
            linear_2: LinearConfig::new(dense_size, output_size).init(&Default::default()),
        }
    }
}

impl<B: Backend> Model<B, Tensor<B, 2>, Tensor<B, 2>> for Actor<B> {
    fn forward(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
        let layer_0_output = relu(self.linear_0.forward(input));
        let layer_1_output = relu(self.linear_1.forward(layer_0_output));

        softmax(self.linear_2.forward(layer_1_output), 1)
    }

    fn infer(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
        self.forward(input)
    }
}

impl<B: Backend> SACActor<B> for Actor<B> {}

#[derive(Module, Debug)]
pub struct Critic<B: Backend> {
    linear_0: Linear<B>,
    linear_1: Linear<B>,
    linear_2: Linear<B>,
}

impl<B: Backend> Critic<B> {
    pub fn new(input_size: usize, dense_size: usize, output_size: usize) -> Self {
        Self {
            linear_0: LinearConfig::new(input_size, dense_size).init(&Default::default()),
            linear_1: LinearConfig::new(dense_size, dense_size).init(&Default::default()),
            linear_2: LinearConfig::new(dense_size, output_size).init(&Default::default()),
        }
    }

    fn consume(self) -> (Linear<B>, Linear<B>, Linear<B>) {
        (self.linear_0, self.linear_1, self.linear_2)
    }
}

impl<B: Backend> Model<B, Tensor<B, 2>, Tensor<B, 2>> for Critic<B> {
    fn forward(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
        let layer_0_output = relu(self.linear_0.forward(input));
        let layer_1_output = relu(self.linear_1.forward(layer_0_output));

        self.linear_2.forward(layer_1_output)
    }

    fn infer(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
        self.forward(input)
    }
}

impl<B: Backend> SACCritic<B> for Critic<B> {
    fn soft_update(this: Self, that: &Self, tau: ElemType) -> Self {
        let (linear_0, linear_1, linear_2) = this.consume();

        Self {
            linear_0: soft_update_linear(linear_0, &that.linear_0, tau),
            linear_1: soft_update_linear(linear_1, &that.linear_1, tau),
            linear_2: soft_update_linear(linear_2, &that.linear_2, tau),
        }
    }
}

#[allow(unused)]
const MEMORY_SIZE: usize = 4096;

type MyAgent<E, B> = SAC<E, B, Actor<B>>;

#[allow(unused)]
pub fn run<
    E: Environment + AsMut<TrictracEnvironment>,
    B: AutodiffBackend<InnerBackend = NdArray>,
>(
    conf: &Config,
    visualized: bool,
) -> impl Agent<E> {
    let mut env = E::new(visualized);
    env.as_mut().max_steps = conf.max_steps;
    let state_dim = <<E as Environment>::StateType as State>::size();
    let action_dim = <<E as Environment>::ActionType as Action>::size();

    let actor = Actor::<B>::new(state_dim, conf.dense_size, action_dim);
    let critic_1 = Critic::<B>::new(state_dim, conf.dense_size, action_dim);
    let critic_2 = Critic::<B>::new(state_dim, conf.dense_size, action_dim);
    let mut nets = SACNets::<B, Actor<B>, Critic<B>>::new(actor, critic_1, critic_2);

    let mut agent = MyAgent::default();

    let config = SACTrainingConfig {
        gamma: conf.gamma,
        tau: conf.tau,
        learning_rate: conf.learning_rate,
        min_probability: conf.min_probability,
        batch_size: conf.batch_size,
        clip_grad: Some(burn::grad_clipping::GradientClippingConfig::Value(
            conf.clip_grad,
        )),
    };

    let mut memory = Memory::<E, B, MEMORY_SIZE>::default();

    let optimizer_config = AdamWConfig::new().with_grad_clipping(config.clip_grad.clone());

    let mut optimizer = SACOptimizer::new(
        optimizer_config.clone().init(),
        optimizer_config.clone().init(),
        optimizer_config.clone().init(),
        optimizer_config.init(),
    );

    let mut step = 0_usize;

    for episode in 0..conf.num_episodes {
        let mut episode_done = false;
        let mut episode_reward = 0.0;
        let mut episode_duration = 0_usize;
        let mut state = env.state();
        let mut now = SystemTime::now();

        while !episode_done {
            if let Some(action) = MyAgent::<E, _>::react_with_model(&state, &nets.actor) {
                let snapshot = env.step(action);

                episode_reward += <<E as Environment>::RewardType as Into<ElemType>>::into(
                    snapshot.reward().clone(),
                );

                memory.push(
                    state,
                    *snapshot.state(),
                    action,
                    snapshot.reward().clone(),
                    snapshot.done(),
                );

                if config.batch_size < memory.len() {
                    nets = agent.train::<MEMORY_SIZE, _>(nets, &memory, &mut optimizer, &config);
                }

                step += 1;
                episode_duration += 1;

                if snapshot.done() || episode_duration >= conf.max_steps {
                    env.reset();
                    episode_done = true;

                    println!(
                        "{{\"episode\": {episode}, \"reward\": {episode_reward:.4}, \"steps count\": {episode_duration}, \"duration\": {}}}",
                    now.elapsed().unwrap().as_secs()
                    );
                    now = SystemTime::now();
                } else {
                    state = *snapshot.state();
                }
            }
        }
    }

    let valid_agent = agent.valid(nets.actor);
    if let Some(path) = &conf.save_path {
        if let Some(model) = valid_agent.model() {
            save_model(model, path);
        }
    }
    valid_agent
}

pub fn save_model(model: &Actor<NdArray<ElemType>>, path: &String) {
    let recorder = CompactRecorder::new();
    let model_path = format!("{path}.mpk");
    println!("info: Modèle de validation sauvegardé : {model_path}");
    recorder
        .record(model.clone().into_record(), model_path.into())
        .unwrap();
}

pub fn load_model(dense_size: usize, path: &String) -> Option<Actor<NdArray<ElemType>>> {
    let model_path = format!("{path}.mpk");
    // println!("Chargement du modèle depuis : {model_path}");

    CompactRecorder::new()
        .load(model_path.into(), &NdArrayDevice::default())
        .map(|record| {
            Actor::new(
                <TrictracEnvironment as Environment>::StateType::size(),
                dense_size,
                <TrictracEnvironment as Environment>::ActionType::size(),
            )
            .load_record(record)
        })
        .ok()
}