{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple, Union\n",
    "from collections import deque\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from gym_env_erase import CustomEnv\n",
    "from stable_baselines3 import PPO\n",
    "import pycuda.driver as cuda\n",
    "from scipy.optimize import curve_fit\n",
    "from scipy.fft import fft, fftfreq\n",
    "import pickle\n",
    "import sys\n",
    "import os\n",
    "\n",
    "current_dir = os.path.dirname(os.path.abspath(\"__file__\"))\n",
    "parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))\n",
    "sys.path.append(parent_dir)\n",
    "\n",
    "from CelerisLab import FlowField\n",
    "from CelerisLab import utils\n",
    "\n",
    "cuda.init()\n",
    "\n",
    "context2 = cuda.Device(2).make_context()\n",
    "\n",
    "config_cuda = utils.load_cuda_config(os.path.join(parent_dir, \"configs\", \"config_cuda.json\"))\n",
    "config_field = utils.load_flow_field_config(os.path.join(parent_dir, \"configs\", \"config_flowfield.json\"))\n",
    "\n",
    "FIFO_LEN = 120\n",
    "SAMPLE_INTERVAL = 800\n",
    "DATA_TYPE = np.float32\n",
    "CONV_LEN = 60\n",
    "SAMP_RATE = 60\n",
    "\n",
    "def normalize_sensor(states, diviation, norm_factor):\n",
    "    return (states - diviation) / norm_factor\n",
    "\n",
    "def force2cdcl(force, U0, Diam):\n",
    "    return force / (0.5 * U0**2 * Diam)\n",
    "\n",
    "def calc_lag(target, state):\n",
    "    target_mean = np.mean(target)\n",
    "    state_mean = np.mean(state)\n",
    "\n",
    "    correlation = np.correlate(target - target_mean, state - state_mean, \"full\")\n",
    "    lags = np.arange(-len(target) + 1, len(target))\n",
    "    max_lag = lags[np.argmax(correlation)]\n",
    "    return max_lag\n",
    "    \n",
    "def calc_sim(target, state, lag):\n",
    "    target_mean = np.mean(target)\n",
    "    state_mean = np.mean(state)\n",
    "    target_std = np.std(target)\n",
    "\n",
    "    aligned_state = np.roll(state, lag)\n",
    "\n",
    "    if lag >= 0:\n",
    "        seq_target = target[-CONV_LEN:]-target_mean\n",
    "        seq_state = aligned_state[-CONV_LEN:]-state_mean\n",
    "    else:\n",
    "        seq_target = target[:CONV_LEN]-target_mean\n",
    "        seq_state = aligned_state[:CONV_LEN]-state_mean\n",
    "\n",
    "    seq_diff = seq_target - seq_state\n",
    "    sim_cor = 10*(np.corrcoef(seq_target, seq_state)[0, 1] - 1)\n",
    "    sim_div = -np.abs((target_mean - state_mean) / target_std * 0.75)\n",
    "    sim_amp = -np.abs(np.std(seq_diff) / target_std * 2)\n",
    "\n",
    "    return np.exp((sim_cor + sim_div + sim_amp) / 3)\n",
    "\n",
    "def sine_wave(x, amplitude, frequency, phase, offset):\n",
    "    return amplitude * np.sin(2 * np.pi * frequency * x + phase) + offset\n",
    "\n",
    "def estimate_frequency(data):\n",
    "    N = len(data)\n",
    "    yf = fft(data)\n",
    "    xf = fftfreq(N, 1 / SAMP_RATE)\n",
    "\n",
    "    idx = np.argmax(np.abs(yf[1:N//2])) + 1\n",
    "    frequency = np.abs(xf[idx])\n",
    "    return frequency\n",
    "\n",
    "def fit_sine_wave(data):\n",
    "    N = len(data)\n",
    "    x = np.arange(N) / SAMP_RATE\n",
    "\n",
    "    initial_frequency = estimate_frequency(data)\n",
    "\n",
    "    initial_amplitude = np.std(data) * np.sqrt(2)\n",
    "    initial_phase = 0\n",
    "    initial_offset = np.mean(data)\n",
    "\n",
    "    popt, _ = curve_fit(sine_wave, x, data, p0=[initial_amplitude, initial_frequency, initial_phase, initial_offset])\n",
    "\n",
    "    amplitude, frequency, phase, offset = popt\n",
    "    if amplitude < 0:\n",
    "        amplitude = -amplitude\n",
    "        phase += np.pi\n",
    "    return amplitude, frequency, phase, offset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "context2.push()\n",
    "vec_env = CustomEnv(device_id=2)\n",
    "context2.pop()\n",
    "model = PPO.load(os.path.join(parent_dir, \"models\", \"d1a3o12_re100_imit_re50.zip\"), env=vec_env, device=\"cuda:2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_reward = 0\n",
    "best_seed = 0\n",
    "for seed in range(1, 11):\n",
    "    model.set_random_seed(seed=seed)\n",
    "    context2.push()\n",
    "    model_env = model.get_env()\n",
    "    obs = model_env.reset()\n",
    "    context2.pop()\n",
    "    reward_list = deque(maxlen=FIFO_LEN)\n",
    "    for _ in range(240):\n",
    "        context2.push()\n",
    "        action, _states = model.predict(observation=obs, deterministic=True)\n",
    "        obs, rewards, dones, info = model_env.step(action)\n",
    "        states = model_env.envs[0].flow_field.obs.copy()[2:14]\n",
    "        context2.pop()\n",
    "        reward_list.append(rewards)\n",
    "    if np.mean(reward_list) > max_reward:\n",
    "        max_reward = np.mean(np.array(reward_list)[-FIFO_LEN:])\n",
    "        best_seed = seed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.set_random_seed(seed=best_seed)\n",
    "states_list = []\n",
    "reward_list = []\n",
    "cl_list, cd_list, diff_list = [], [], []\n",
    "action_list = deque(maxlen=FIFO_LEN)\n",
    "fifo_states = deque(maxlen=FIFO_LEN)\n",
    "\n",
    "context2.push()\n",
    "model_env = model.get_env()\n",
    "obs = model_env.reset()\n",
    "context2.pop()\n",
    "\n",
    "for _ in range(240):\n",
    "    context2.push()\n",
    "    # action, _states = model.predict(observation=obs, deterministic=True)\n",
    "    action = np.array([0.0, 0.5, -0.5], dtype=np.float32)[np.newaxis, :]\n",
    "    obs, rewards, dones, info = model_env.step(action)\n",
    "    states = model_env.envs[0].flow_field.obs.copy()[2:14]\n",
    "    context2.pop()\n",
    "\n",
    "    fifo_states.append(states)\n",
    "    reward_list.append(rewards)\n",
    "    action_list.append(action[0,:])\n",
    "    states_list.append(np.array(states))\n",
    "    cd_list.append(obs[0, 3])\n",
    "    cl_list.append(obs[0, 4])\n",
    "    diff_list.append(obs[0, 5])\n",
    "\n",
    "ave_reward = np.mean(np.array(reward_list)[-CONV_LEN:])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 105456 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "reward_cd = np.exp(-np.abs(cd_list) * 20)\n",
    "reward_cl = np.exp(-np.abs(cl_list) * 80)\n",
    "reward_sim = np.exp(-np.abs(diff_list) * 20)\n",
    "reward_sum = 0.3*reward_cd + 0.3*reward_cl + 0.4*reward_sim\n",
    "plt.figure()\n",
    "plt.plot(reward_cd, label=\"cd\")\n",
    "plt.plot(reward_cl, label=\"cl\")\n",
    "plt.plot(reward_sim, label=\"sim\")\n",
    "plt.plot(reward_sum, label=\"sum\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Actions')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 161512 characters]",
      "text/plain": [
       "<Figure size 1500x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "offset = np.array([0, -4, +4], dtype=np.float32)\n",
    "list_action_offset = [action*8+offset for action in action_list]\n",
    "target_states = model_env.envs[0].target_states.copy()\n",
    "# for i in range(6):\n",
    "#     norm_factor = model_env.envs[0].sens_norm_fact[i]\n",
    "#     diviation = model_env.envs[0].sens_deviation[i]\n",
    "#     target_states[:, i] = target_states[:, i] * norm_factor + diviation\n",
    "force_list = np.array(states_list)[:, 6:12]*100/3\n",
    "sens_list = np.array(states_list)[:, 0:6]\n",
    "drag = force_list[:, 0] + force_list[:, 2] + force_list[:, 4]\n",
    "lift = force_list[:, 1] + force_list[:, 3] + force_list[:, 5]\n",
    "plt.figure(figsize=(15, 12))\n",
    "plt.subplot(2, 3, 1)\n",
    "plt.plot(np.array(states_list)[-120:, 0:6])\n",
    "plt.title(\"Sensors\")\n",
    "# plt.subplot(2, 3, 2)\n",
    "# plt.plot(np.array(target_states)[:, 0:6])\n",
    "# plt.title(\"Target Sensors\")\n",
    "plt.subplot(2, 3, 3)\n",
    "plt.plot(drag, label=\"Drag\")\n",
    "plt.plot(lift, label=\"Lift\")\n",
    "plt.title(\"Forces\")\n",
    "plt.legend()\n",
    "plt.subplot(2, 3, 4)\n",
    "plt.plot(np.array(list_action_offset))\n",
    "plt.title(\"Actions\")\n",
    "# plt.subplot(2, 3, 5)\n",
    "# plt.plot(reward_list)\n",
    "# plt.title(\"Rewards\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "amp, freq, pha_ref, offset = fit_sine_wave(np.array(states_list)[-CONV_LEN:, 7])\n",
    "\n",
    "amp_action, pha_action, offset_action = [], [], []\n",
    "freq_action = []\n",
    "for i in range(3):\n",
    "    amp, freq, pha, offset = fit_sine_wave(np.array(action_list)[:, i])\n",
    "    amp_action.append(amp)\n",
    "    pha_action.append(pha-pha_ref)\n",
    "    freq_action.append(freq)\n",
    "    offset_action.append(offset)\n",
    "freq_action = np.mean(freq_action)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "variables = {\n",
    "    'amp_action': amp_action,\n",
    "    'pha_action': pha_action,\n",
    "    'freq_action': freq_action,\n",
    "    'offset_action': offset_action,\n",
    "    'best_seed': best_seed\n",
    "}\n",
    "with open(os.path.join(parent_dir, \"output\", \"parm_e0.pkl\"), 'wb') as f:\n",
    "    pickle.dump(variables, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(parent_dir, \"output\", \"parm_e0.pkl\"), 'rb') as f:\n",
    "    variables = pickle.load(f)\n",
    "amp_action = variables['amp_action']\n",
    "pha_action = variables['pha_action']\n",
    "freq_action = variables['freq_action']\n",
    "offset_action = variables['offset_action']\n",
    "best_seed = variables['best_seed']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.set_random_seed(seed=best_seed)\n",
    "states_list_fit = []\n",
    "reward_list_fit = []\n",
    "action_list_fit = deque(maxlen=FIFO_LEN)\n",
    "fifo_states_fit = deque(maxlen=FIFO_LEN)\n",
    "\n",
    "context2.push()\n",
    "model_env = model.get_env()\n",
    "obs = model_env.reset()\n",
    "context2.pop()\n",
    "\n",
    "for _ in range(FIFO_LEN):\n",
    "    context2.push()\n",
    "    action = np.array([0.0, 0.0, 0.0], dtype=np.float32)[np.newaxis, :]\n",
    "    obs, rewards, dones, info = model_env.step(action)\n",
    "    states = model_env.envs[0].flow_field.obs.copy()[2:14]\n",
    "    # states = model_env.envs[0].flow_field.obs.copy()[0:12]\n",
    "    context2.pop()\n",
    "\n",
    "    fifo_states_fit.append(states)\n",
    "\n",
    "for i in range(240):\n",
    "    context2.push()\n",
    "    _, _, pha_state, _ = fit_sine_wave(np.array(fifo_states_fit)[-20:, 7])\n",
    "    temps = []\n",
    "    for j in range(3):\n",
    "        temp = sine_wave(CONV_LEN / SAMP_RATE, amp_action[j], freq_action, pha_state + pha_action[j], offset_action[j])\n",
    "        # temp = sine_wave(i / SAMP_RATE, amp_action[j], freq_action, pha_action[j], offset_action[j])\n",
    "        temps.append(temp)\n",
    "    action = np.array(temps, dtype=np.float32)[np.newaxis, :]\n",
    "    obs, rewards, dones, info = model_env.step(action)\n",
    "    states = model_env.envs[0].flow_field.obs.copy()[2:14]\n",
    "    # states = model_env.envs[0].flow_field.obs.copy()[0:12]\n",
    "    context2.pop()\n",
    "\n",
    "    fifo_states_fit.append(states)\n",
    "    action_list_fit.append(action[0,:])\n",
    "    states_list_fit.append(np.array(states))\n",
    "    reward_list_fit.append(rewards)\n",
    "\n",
    "ave_reward_fit = np.mean(np.array(reward_list_fit)[-CONV_LEN:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 164692 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "context2.push()\n",
    "model_env.envs[0].render()\n",
    "context2.pop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "context2.push()\n",
    "model_env.envs[0].save_field(filename=\"re200_d2_fit.dat\")\n",
    "context2.pop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Fitted Forces')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 424588 characters]",
      "text/plain": [
       "<Figure size 1200x1200 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "offset = np.array([0, -4, +4], dtype=np.float32)\n",
    "# list_action_offset = [action*8+offset for action in action_list]\n",
    "fit_action_offset = [action*8+offset for action in action_list_fit]\n",
    "# force_list = np.array(states_list)[:, 6:12]*100/3\n",
    "# sens_list = np.array(states_list)[:, 0:6]\n",
    "force_list_fit = np.array(states_list_fit)[:, 6:12]*100/3\n",
    "sens_list_fit = np.array(states_list_fit)[:, 0:6]\n",
    "\n",
    "# drag = force_list[:, 0] + force_list[:, 2] + force_list[:, 4]\n",
    "drag_fit = force_list_fit[:, 0] + force_list_fit[:, 2] + force_list_fit[:, 4]\n",
    "# lift = force_list[:, 1] + force_list[:, 3] + force_list[:, 5]\n",
    "lift_fit = force_list_fit[:, 1] + force_list_fit[:, 3] + force_list_fit[:, 5]\n",
    "\n",
    "plt.figure(figsize=(12, 12))\n",
    "plt.subplot(3, 2, 1)\n",
    "plt.plot(np.array(list_action_offset))\n",
    "plt.title(\"DRL Actions\")\n",
    "plt.subplot(3, 2, 2)\n",
    "plt.plot(np.array(fit_action_offset))\n",
    "plt.title(\"Fitted Actions\")\n",
    "plt.subplot(3, 2, 3)\n",
    "plt.plot(sens_list)\n",
    "plt.title(\"DRL Sensors\")\n",
    "plt.subplot(3, 2, 4)\n",
    "plt.plot(sens_list_fit)\n",
    "plt.title(\"Fitted Sensors\")\n",
    "plt.subplot(3, 2, 5)\n",
    "plt.plot(drag, label=\"Drag\")\n",
    "plt.plot(lift, label=\"Lift\")\n",
    "plt.legend()\n",
    "plt.title(\"DRL Forces\")\n",
    "plt.subplot(3, 2, 6)\n",
    "plt.plot(drag_fit, label=\"Drag\")\n",
    "plt.plot(lift_fit, label=\"Lift\")\n",
    "plt.legend()\n",
    "plt.title(\"Fitted Forces\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pycuda_3_10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
