{
 "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 import CustomEnv\n",
    "from stable_baselines3 import PPO\n",
    "import pycuda.driver as cuda\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",
    "from CelerisLab import FlowField\n",
    "from CelerisLab import utils\n",
    "\n",
    "cuda.init()\n",
    "\n",
    "context2 = cuda.Device(2).make_context()\n",
    "context3 = cuda.Device(3).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",
    "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "context2.push()\n",
    "vec_env = CustomEnv(device_id=2)\n",
    "model = PPO.load(os.path.join(parent_dir, \"models\", \"d1a3o12_a3.zip\"), env=vec_env, device=\"cuda:2\")\n",
    "context2.pop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "context3.push()\n",
    "\n",
    "FIFO_LEN = 120\n",
    "SAMPLE_INTERVAL = 800\n",
    "DATA_TYPE = np.float32\n",
    "\n",
    "fifo_states = deque(maxlen=FIFO_LEN)\n",
    "target_states = np.empty((0, 6), dtype=DATA_TYPE)\n",
    "refer_states = np.empty((0, 6), dtype=DATA_TYPE)\n",
    "force_norm_fact = 1.0\n",
    "sens_norm_fact = np.ones(6, dtype=DATA_TYPE)\n",
    "sens_deviation = np.zeros(6, dtype=DATA_TYPE)\n",
    "\n",
    "flow_field = FlowField(config_field, config_cuda, device_id=3)\n",
    "\n",
    "L0 = 20\n",
    "U0 = config_field.velocity\n",
    "NX = flow_field.FIELD_SHAPE[0]\n",
    "NY = flow_field.FIELD_SHAPE[1]\n",
    "center: Tuple[float, float, float] = (10 * L0, (NY - 1) / 2, 0)\n",
    "flow_field.add_cylinder(center, L0 / 2)\n",
    "center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 + 2 * L0, 0)\n",
    "flow_field.add_sensor(center, L0 / 4)\n",
    "center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2, 0)\n",
    "flow_field.add_sensor(center, L0 / 4)\n",
    "center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 - 2 * L0, 0)\n",
    "flow_field.add_sensor(center, L0 / 4)\n",
    "flow_field.get_ddf()\n",
    "flow_field.save_ddf()\n",
    "flow_field.run(int(4*NX/U0), np.zeros(4, dtype=DATA_TYPE))\n",
    "\n",
    "for i in range(FIFO_LEN):\n",
    "    flow_field.run(SAMPLE_INTERVAL, np.zeros(len(flow_field.objects), dtype=DATA_TYPE))\n",
    "    target_states = np.vstack((target_states, flow_field.obs.copy()[2:8]))\n",
    "\n",
    "# previous_lift_test = None\n",
    "# while True:\n",
    "#     flow_field.run(100, np.zeros(len(flow_field.objects), dtype=DATA_TYPE))\n",
    "#     lift_test = flow_field.obs.copy()[1]\n",
    "#     if previous_lift_test is not None and previous_lift_test > 0 and lift_test < 0:\n",
    "#         break\n",
    "#     previous_lift_test = lift_test\n",
    "    \n",
    "context3.pop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "context3.push()\n",
    "\n",
    "flow_field.restore_ddf()\n",
    "flow_field.apply_ddf()\n",
    "center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)\n",
    "flow_field.add_cylinder(center, L0 / 2)\n",
    "center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 + 0.75 * L0, 0)\n",
    "flow_field.add_cylinder(center, L0 / 2)\n",
    "center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 - 0.75 * L0, 0)\n",
    "flow_field.add_cylinder(center, L0 / 2)\n",
    "flow_field.run(int(4*NX/U0), np.zeros(len(flow_field.objects), dtype=DATA_TYPE))\n",
    "flow_field.get_ddf()\n",
    "flow_field.save_ddf()\n",
    "\n",
    "for i in range(FIFO_LEN):\n",
    "    flow_field.run(SAMPLE_INTERVAL, np.zeros(len(flow_field.objects), dtype=DATA_TYPE))\n",
    "    fifo_states.append(flow_field.obs.copy()[2:14])\n",
    "    refer_states = np.vstack((refer_states, flow_field.obs.copy()[2:8]))\n",
    "\n",
    "# previous_lift_test = None\n",
    "# while True:\n",
    "#     flow_field.run(100, np.zeros(len(flow_field.objects), dtype=DATA_TYPE))\n",
    "#     lift_test = flow_field.obs.copy()[1]\n",
    "#     if previous_lift_test is not None and previous_lift_test > 0 and lift_test < 0:\n",
    "#         break\n",
    "#     previous_lift_test = lift_test\n",
    "\n",
    "temp_states = np.array(fifo_states)\n",
    "fifo_states.clear()\n",
    "force_norm_fact = 6 * np.max(np.abs(temp_states[:, 6:12]))\n",
    "\n",
    "for i in range(6):\n",
    "    sens_deviation[i] = np.mean(temp_states[:, i])\n",
    "    sens_norm_fact[i] = 5 * np.max(np.abs(temp_states[:, i] - sens_deviation[i]))\n",
    "\n",
    "context3.pop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "states_list = []\n",
    "force_list = []\n",
    "force_model_list = []\n",
    "sensor_list = []\n",
    "sensor_model_list = []\n",
    "drag_list = []\n",
    "lift_list = []\n",
    "action_list = []\n",
    "\n",
    "context2.push()\n",
    "model_env = model.get_env()\n",
    "obs = model_env.reset()\n",
    "context2.pop()\n",
    "\n",
    "context3.push()\n",
    "flow_field.restore_ddf()\n",
    "flow_field.apply_ddf()\n",
    "context3.pop()\n",
    "\n",
    "previous_lift_test = None\n",
    "i = 0\n",
    "while True:\n",
    "    context2.push()\n",
    "    action, _states = model.predict(observation=obs, deterministic=True)\n",
    "    obs, rewards, dones, info = model_env.step(action)\n",
    "    force_model_list.append(np.array(obs[0][0:6]*force_norm_fact, dtype=DATA_TYPE))\n",
    "    sensor_model_list.append(np.array(obs[0][6:12]*sens_norm_fact+sens_deviation, dtype=DATA_TYPE))\n",
    "    action_list.append(action)\n",
    "    context2.pop()\n",
    "\n",
    "    context3.push()\n",
    "    temp = np.zeros(7, dtype=DATA_TYPE)\n",
    "    temp[4:7] = np.array((action*8+[0,-4,4])*U0, dtype=DATA_TYPE)\n",
    "    flow_field.run(SAMPLE_INTERVAL, temp)\n",
    "    states = flow_field.obs.copy()[2:14]\n",
    "    fifo_states.append(states)\n",
    "    forces = np.array(states[6:12], dtype=DATA_TYPE) / force_norm_fact\n",
    "    sensors = np.array(normalize_sensor(states[:6], sens_deviation, sens_norm_fact), dtype=DATA_TYPE)\n",
    "    # obs = np.hstack((sensors, forces)).reshape(1, 12)\n",
    "    states_list.append(np.array(states))\n",
    "    force_list.append(forces)\n",
    "    sensor_list.append(sensors)\n",
    "\n",
    "    lift_test = flow_field.obs.copy()[1]\n",
    "    if previous_lift_test is not None and previous_lift_test > 0 and lift_test < 0 and i > 320:\n",
    "        break\n",
    "    previous_lift_test = lift_test\n",
    "    i += 1\n",
    "\n",
    "    context3.pop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Actions')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 325324 characters]",
      "text/plain": [
       "<Figure size 1200x1200 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "offset = np.array([0, -4, +4], dtype=np.float32)\n",
    "list_action_offset = [action[0]*8+offset for action in action_list]\n",
    "\n",
    "plt.figure(figsize=(12, 12))\n",
    "plt.subplot(3, 2, 1)\n",
    "plt.plot(target_states)\n",
    "plt.title(\"Target\")\n",
    "plt.subplot(3, 2, 2)\n",
    "plt.plot(refer_states)\n",
    "plt.title(\"No control\")\n",
    "plt.subplot(3, 2, 3)\n",
    "plt.plot(np.array(sensor_model_list)[-132:-12])\n",
    "# plt.plot(np.array(sensor_list)[-120:]*sens_norm_fact+sens_deviation)\n",
    "plt.title(\"Controled\")\n",
    "plt.subplot(3, 2, 4)\n",
    "plt.plot(np.array(list_action_offset)[-132:-12])\n",
    "plt.title(\"Actions\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "filename = \"result1.csv\"\n",
    "data = {}\n",
    "for i in range(6):\n",
    "    data[\"target\" + str(i+1)] = target_states[:, i]\n",
    "for i in range(6):\n",
    "    data[\"reference\" + str(i+1)] = refer_states[:, i]\n",
    "for i in range(6):\n",
    "    data[\"result\" + str(i+1)] = np.array(sensor_model_list)[-132:-12, i]\n",
    "for i in range(3):\n",
    "    data[\"action\" + str(i+1)] = np.array(list_action_offset)[-132:-12, i]\n",
    "\n",
    "df = pd.DataFrame(data)\n",
    "df.to_csv(os.path.join(parent_dir, \"output\", filename), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 106500 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "context3.push()\n",
    "\n",
    "flow_field.get_ddf()\n",
    "\n",
    "# 将一维数组重新整形成三维数组\n",
    "NX = flow_field.FIELD_SHAPE[0]\n",
    "NY = flow_field.FIELD_SHAPE[1]\n",
    "ddf_plot = flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)\n",
    "flag_plot = flow_field.flag.copy().reshape((NY, NX)).transpose(1, 0)\n",
    "\n",
    "context3.pop()\n",
    "\n",
    "# 计算速度场\n",
    "ux = (ddf_plot[:, :, 1] + ddf_plot[:, :, 5] + ddf_plot[:, :, 8] - ddf_plot[:, :, 3] - ddf_plot[:, :, 6] - ddf_plot[:, :, 7]) / U0\n",
    "uy = (ddf_plot[:, :, 2] + ddf_plot[:, :, 5] + ddf_plot[:, :, 6] - ddf_plot[:, :, 4] - ddf_plot[:, :, 7] - ddf_plot[:, :, 8]) / U0\n",
    "\n",
    "# 计算速度大小\n",
    "speed = np.sqrt(ux**2 + uy**2)\n",
    "# speed = np.clip(speed, 0, 0.02)\n",
    "\n",
    "# 绘制标量速度场\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.subplot(2, 1, 1)\n",
    "plt.imshow(speed.T, origin='lower', cmap='viridis', extent=[0, NX, 0, NY])\n",
    "plt.colorbar(label='Speed')\n",
    "plt.title('Scalar Velocity Field')\n",
    "plt.xlabel('X')\n",
    "plt.ylabel('Y')\n",
    "\n",
    "# 绘制flag\n",
    "# plt.subplot(2, 1, 1)\n",
    "# plt.imshow(flag_plot.T, origin='lower', cmap='viridis', extent=[0, NX, 0, NY])\n",
    "# # plt.colorbar(label='Flag')\n",
    "# plt.title('Flag Field')\n",
    "# plt.xlabel('X')\n",
    "# plt.ylabel('Y')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = \"re50.dat\"\n",
    "with open(os.path.join(parent_dir, \"output\", filename), \"w\") as f:\n",
    "    f.write(\"Title= \\\"LBM 2D\\\"\\r\\n\")\n",
    "    f.write(\"VARIABLES= \\\"X\\\",\\\"Y\\\",\\\"flag\\\",\\\"U\\\",\\\"V\\\",\\r\\n\")\n",
    "    f.write(f\"ZONE T= \\\"BOX\\\",I= {NX},J= {NY},F=POINT\\r\\n\")\n",
    "    for j in range(NY):\n",
    "        for i in range(NX):\n",
    "            f.write(f\"{i},{j},{flag_plot[i, j]},{ux[i, j]},{uy[i, j]}\\r\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 示例序列，包含两个周期\n",
    "sequence1 = 0.2*np.sin(np.linspace(0, 12.5 * np.pi, 120))+0.4\n",
    "sequence2 = 0.2*np.sin(np.linspace(0, 12.5 * np.pi, 120) + np.pi / 3) +0.4  # 相位差为 π/4\n",
    "\n",
    "seq1_mean = np.mean(sequence1)\n",
    "seq2_mean = np.mean(sequence2)\n",
    "seq1_std = np.std(sequence1)\n",
    "\n",
    "correlation = np.correlate(sequence1-seq1_mean, sequence2-seq2_mean, mode='full')\n",
    "lags = np.arange(-len(sequence1) + 1, len(sequence1))\n",
    "\n",
    "# 找到最大相关系数及其对应的滞后\n",
    "max_corr = np.max(correlation)\n",
    "max_lag = lags[np.argmax(correlation)]\n",
    "\n",
    "print(f\"对应的滞后: {max_lag}\")\n",
    "\n",
    "# 根据滞后调整序列2\n",
    "\n",
    "aligned_seq2 = np.roll(sequence2, max_lag)\n",
    "period_length = len(sequence1) // 2\n",
    "if max_lag < 0:\n",
    "    seq1_one_period = sequence1[:period_length]-seq1_mean\n",
    "    aligned_seq2_one_period = aligned_seq2[:period_length]-seq2_mean\n",
    "else:\n",
    "    seq1_one_period = sequence1[-period_length:]-seq1_mean\n",
    "    aligned_seq2_one_period = aligned_seq2[-period_length:]-seq2_mean\n",
    "\n",
    "# 计算一个周期的相似度\n",
    "seq_diff = seq1_one_period - aligned_seq2_one_period\n",
    "sim1 = 10*(np.corrcoef(seq1_one_period, aligned_seq2_one_period)[0, 1]-1)\n",
    "sim2 = -np.abs((seq1_mean-seq2_mean)/seq1_std*0.75)\n",
    "sim3 = -np.abs(np.std(seq_diff)/seq1_std*2)\n",
    "similarity = np.exp((sim1 + sim2 + sim3)/3)\n",
    "\n",
    "reward = np.exp((similarity - 1))\n",
    "\n",
    "print(f\"sim1: {sim1}, sim2: {sim2}, sim3: {sim3}\")\n",
    "print(f\"相似度: {similarity}\")\n",
    "print(f\"reward: {reward}\")\n",
    "\n",
    "# 绘制归一化交叉相关图\n",
    "plt.figure(figsize=(12, 6))\n",
    "plt.subplot(2, 1, 1)\n",
    "plt.plot(lags, correlation)\n",
    "plt.title('Normalized Cross-Correlation')\n",
    "plt.xlabel('Lag')\n",
    "plt.ylabel('Correlation')\n",
    "plt.subplot(2, 1, 2)\n",
    "plt.plot(sequence1, label='Sequence 1')\n",
    "plt.plot(aligned_seq2, label='Aligned Sequence 2')\n",
    "plt.show()"
   ]
  }
 ],
 "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
}
