<template>
  <v-navigation-drawer v-model="store.rosBarOpened" temporary location="right">
    <v-container fluid>
      <v-row justify="space-between">
        <v-col cols="6" md="12">
          <v-text-field
            v-model="formState.ipaddress"
            label="IP Address"
            :disabled="connectionState"
          />
          <v-text-field
            v-model="formState.port"
            label="Port"
            :disabled="connectionState"
          />
          <v-switch
            @change="updateConnectionState"
            true-value="connected"
            false-value="not connected"
            :label="`${formState.connect}`"
          />
        </v-col>
      </v-row>
    </v-container>
  </v-navigation-drawer>
</template>

<script setup>
import * as ROS from "roslib";
import { useBWStore } from "@/store/BaseWebsiteStore";
import { usePStore } from "@/store/PolicyStore";
import { useCStore } from "@/store/ControlStore";
import { useRStore } from "@/store/RewardStore";
import { useRLStore } from "@/store/RLStore";

import { reactive, ref, watch } from "vue";

const store = useBWStore();
const pstore = usePStore();
const cstore = useCStore();
const rstore = useRStore();
const rlstore = useRLStore();

const formState = reactive({
  ipaddress: "localhost",
  port: "9090",
  connect: "not connected",
  // connectionState: computed(() => formState.connect !== 'not connected'),
});

const connectionState = ref(false);
const stateCounter = ref([0, 0, 0]);

// Open Connection to ROSBridge and subscribe the topics.
const ros = new ROS.Ros();
const ros_eval = new ROS.Ros();

ros.on("connection", () => {
  stateCounter.value[0] += 1;
});
ros_eval.on("connection", () => {});

ros.on("error", (error) => {
  stateCounter.value[1] += 1;
  console.log(error);
});
ros_eval.on("error", () => {});

ros.on("close", () => {
  stateCounter.value[2] += 1;
});
ros_eval.on("close", () => {});

function updateConnectionState() {
  connectionState.value = !connectionState.value;
  if (formState.connect === "not connected") {
    formState.connect = "connected";
    const rosUrl = `ws://${formState.ipaddress}:${formState.port}`;
    ros.connect(rosUrl);
    ros_eval.connect(rosUrl);
  } else {
    formState.connect = "not connected";
    ros.close();
    ros_eval.close();
  }
}

const lastStateCounter = [0, 0, 0];

// ROS2 State Observer
watch(stateCounter.value, (newValue) => {
  let currentState = 0;
  for (let i = 0; i < lastStateCounter.length; i += 1) {
    if (newValue[i] - lastStateCounter[i] === 1) {
      currentState = i;
    }
    lastStateCounter[i] = newValue[i];
  }
  if (currentState === 0) {
    store.setRosConState("connected");
  } else if (currentState === 1) {
    store.setRosConState("error");
  } else if (currentState === 2) {
    store.setRosConState("not connected");
  }
});

// State Subscriber
const state_subscriber = new ROS.Topic({
  ros: ros,
  name: "/active_bo_state",
  messageType: "active_bo_msgs/msg/ActiveBOState",
});

state_subscriber.subscribe((msg) => {
  rlstore.setCurrentRun(msg.current_run);
  rlstore.setCurrentEpisode(msg.current_episode);
  rlstore.setBestReward(msg.best_reward);
  rlstore.setLastUserReward(msg.last_user_reward);

  if (msg.current_episode === 1) {
    pstore.resetWeights_Fixed();
  }
});

// RL Service + Feedback Subscriber
const rl_feedback_subscriber = new ROS.Topic({
  ros: ros,
  name: "/rl_feedback",
  messageType: "active_bo_msgs/msg/ImageFeedback",
});

rl_feedback_subscriber.subscribe((msg) => {
  rlstore.setDim(msg.height, msg.width);
  rlstore.setRgbArrays(msg.red, msg.green, msg.blue);
  rlstore.setCurrentTime(msg.current_time);
});

const active_rl_eval_sub = new ROS.Topic({
  ros: ros_eval,
  name: "/active_rl_eval_request",
  messageType: "active_bo_msgs/msg/ActiveRLEvalRequest",
});

const pendingRequest = ref(false);
active_rl_eval_sub.subscribe((msg) => {
  const nr_steps = msg.nr_steps;
  const nr_weights = msg.nr_weights;

  const pol_x = msg.policy.slice(0, nr_steps);
  const pol_y = msg.policy.slice(nr_steps);

  const weights_x = msg.weights.slice(0, nr_weights);
  const weights_y = msg.weights.slice(nr_weights);

  pstore.setPolicy(pol_x);
  pstore.setPolicy_y(pol_y);
  pstore.setWeights(weights_x);
  pstore.setWeights_y(weights_y);
  pendingRequest.value = true;
});

const active_rl_eval_pub = new ROS.Topic({
  ros: ros_eval,
  name: "/active_rl_eval_response",
  messageType: "active_bo_msgs/msg/ActiveRLEvalResponse",
});

watch(
  () => cstore.getSendWeights,
  () => {
    if (!pendingRequest.value) {
      return;
    }
    console.log("Button pressed!");
    const policy = pstore.policy.concat(pstore.policy_y);
    const weights = pstore.weights.concat(pstore.weights_y);
    const weights_fixed = pstore.weights_fixed.concat(pstore.weights_fixed_y);
    const active_eval_response = new ROS.Message({
      policy: policy,
      weights: weights,
      weight_preference: weights_fixed,
    });

    console.log(active_eval_response);

    active_rl_eval_pub.publish(active_eval_response);
    console.log("New Policy/ Weights published");
    pendingRequest.value = false;
  }
);

const active_bo_pending = ref(false);

const active_bo_response = new ROS.Topic({
  ros: ros,
  name: "/active_bo_response",
  messageType: "active_bo_msgs/msg/ActiveBOResponse",
});

active_bo_response.subscribe((msg) => {
  const nr_steps = msg.nr_steps;
  const nr_weights = msg.nr_weights;

  const pol_x = msg.best_policy.slice(0, nr_steps);
  const pol_y = msg.best_policy.slice(nr_steps);

  const weights_x = msg.best_weights.slice(0, nr_weights);
  const weights_y = msg.best_weights.slice(nr_weights);

  pstore.setPolicy(pol_x);
  pstore.setPolicy_y(pol_y);
  pstore.setWeights(weights_x);
  pstore.setWeights_y(weights_y);
  rstore.setMean(msg.reward_mean);
  rstore.setStd(msg.reward_std);
  active_bo_pending.value = false;
});

const active_bo_request = new ROS.Topic({
  ros: ros,
  name: "/active_bo_request",
  messageType: "active_bo_msgs/msg/ActiveBORequest",
});

watch(
  () => cstore.getRunner,
  () => {
    const request_msg = new ROS.Message({
      env: cstore.env,
      metric: cstore.metric,
      fixed_seed: cstore.fixed_seed,
      nr_weights: pstore.nr_weights,
      max_steps: pstore.max_steps,
      nr_dims: cstore.env_dim[cstore.env],
      nr_episodes: cstore.nr_episodes,
      nr_runs: cstore.nr_runs,
      acquisition_function: cstore.acq_fun,
      metric_parameter: cstore.metric_parameter,
      metric_parameter_2: cstore.improvement_2,
      save_result: cstore.save_result,
      overwrite: cstore.overwrite,
    });
    active_bo_pending.value = true;
    active_bo_request.publish(request_msg);
  }
);
</script>

<style scoped></style>
