{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "paraview_dir = os.path.join(parent_dir, \"ParaView\")\n",
    "\n",
    "os.environ['PYTHONPATH'] = f\"{paraview_dir}/lib:{paraview_dir}/lib/python3.9/site-packages\"\n",
    "os.environ['LD_LIBRARY_PATH'] = os.path.join(paraview_dir, \"bin\")\n",
    "\n",
    "sys.path.append(parent_dir)\n",
    "sys.path.append(os.path.join(paraview_dir, \"lib\"))\n",
    "sys.path.append(os.path.join(paraview_dir, \"bin\"))\n",
    "sys.path.append(os.path.join(paraview_dir, \"lib/python3.9/site-packages\"))\n",
    "\n",
    "# from paraview.simple import *\n",
    "import paraview.simple as pv\n",
    "from CelerisLab import FlowField\n",
    "from CelerisLab import utils\n",
    "\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 pickle\n",
    "\n",
    "FIFO_LEN = 120\n",
    "SAMPLE_INTERVAL = 800\n",
    "DATA_TYPE = np.float32\n",
    "CONV_LEN = 60\n",
    "SAMP_RATE = 60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "renderView1 = pv.GetActiveViewOrCreate('RenderView')\n",
    "renderView1.ViewSize = [1200, 480]\n",
    "renderView1.InteractionMode = '2D'\n",
    "renderView1.OrientationAxesVisibility = 0\n",
    "renderView1.CameraPosition = [0.0, 0.0, 3.0]\n",
    "renderView1.UseLight = 0\n",
    "name_files = []\n",
    "for i in range(1, 30):\n",
    "    name_files.append(os.path.join(parent_dir, f\"output/d1a0/field.dat.{i:02}\"))\n",
    "dat = pv.TecplotReader(registrationName='field.dat*', FileNames=name_files)\n",
    "dat.DataArrayStatus = ['flag', 'U', 'V']\n",
    "calculator1 = pv.Calculator(registrationName='Calculator1', Input=dat)\n",
    "calculator1.ResultArrayName = 'W'\n",
    "calculator1.Function = '0'\n",
    "mergeVectorComponents1 = pv.MergeVectorComponents(registrationName='MergeVectorComponents1', Input=calculator1)\n",
    "mergeVectorComponents1.XArray = 'U'\n",
    "mergeVectorComponents1.YArray = 'V'\n",
    "mergeVectorComponents1.ZArray = 'W'\n",
    "gradient1 = pv.Gradient(registrationName='Gradient1', Input=mergeVectorComponents1)\n",
    "gradient1.ScalarArray = ['POINTS', 'Vector']\n",
    "gradient1.ComputeGradient = 0\n",
    "gradient1.ComputeVorticity = 1\n",
    "calculator1 = pv.Calculator(registrationName='Calculator1', Input=dat)\n",
    "calculator1.ResultArrayName = 'Sensor'\n",
    "calculator1.Function = '1/(flag-17)'\n",
    "calculator1.ReplacementValue = 1000.0\n",
    "calculator2 = pv.Calculator(registrationName='Calculator2', Input=dat)\n",
    "calculator2.ResultArrayName = 'Solid'\n",
    "calculator2.Function = '1/(flag-2)/(flag-9)'\n",
    "calculator2.ReplacementValue = 1000.0\n",
    "\n",
    "vorticityTF2D = pv.GetTransferFunction2D('Vorticity')\n",
    "vorticityTF2D.ScalarRangeInitialized = 1\n",
    "vorticityTF2D.Range = [-0.1, 0.1, 0.0, 1.0]\n",
    "vorticityLUT = pv.GetColorTransferFunction('Vorticity')\n",
    "vorticityLUT.AutomaticRescaleRangeMode = 'Never'\n",
    "vorticityLUT.TransferFunction2D = vorticityTF2D\n",
    "vorticityLUT.RGBPoints = [-0.1, 0.23137254902, 0.298039215686, 0.752941176471, -0.0125, 1.0, 1.0, 1.0, 0.0125, 1.0, 1.0, 1.0, 0.1, 0.705882352941, 0.0156862745098, 0.149019607843]\n",
    "vorticityLUT.NumberOfTableValues = 16\n",
    "vorticityLUT.ScalarRangeInitialized = 1.0\n",
    "vorticityLUT.VectorComponent = 2\n",
    "vorticityLUT.VectorMode = 'Component'\n",
    "vorticityPWF = pv.GetOpacityTransferFunction('Vorticity')\n",
    "vorticityPWF.Points = [-0.1, 0.0, 0.5, 0.0, 0.1, 1.0, 0.5, 0.0]\n",
    "vorticityPWF.ScalarRangeInitialized = 1\n",
    "\n",
    "sensorTF2D = pv.GetTransferFunction2D('Sensor')\n",
    "sensorLUT = pv.GetColorTransferFunction('Sensor')\n",
    "sensorLUT.EnableOpacityMapping = 1\n",
    "sensorLUT.TransferFunction2D = sensorTF2D\n",
    "sensorLUT.RGBPoints = [-0.14285714285714285, 1.0, 1.0, 1.0, 1000.0, 0.4, 0.8, 0.4]\n",
    "sensorLUT.ColorSpace = 'RGB'\n",
    "sensorLUT.NanColor = [1.0, 0.0, 0.0]\n",
    "sensorLUT.NumberOfTableValues = 10\n",
    "sensorLUT.ScalarRangeInitialized = 1.0\n",
    "sensorPWF = pv.GetOpacityTransferFunction('Sensor')\n",
    "sensorPWF.Points = [-0.14285714285714285, 0.0, 0.5, 0.0, 499.87855928449983, 0.0, 0.5, 0.0, 1000.0, 1.0, 0.5, 0.0]\n",
    "sensorPWF.ScalarRangeInitialized = 1\n",
    "\n",
    "solidTF2D = pv.GetTransferFunction2D('Solid')\n",
    "solidLUT = pv.GetColorTransferFunction('Solid')\n",
    "solidLUT.EnableOpacityMapping = 1\n",
    "solidLUT.TransferFunction2D = solidTF2D\n",
    "solidLUT.RGBPoints = [-0.14285714285714285, 1.0, 1.0, 1.0, 1000.0, 0.0, 0.0, 0.0]\n",
    "solidLUT.ColorSpace = 'RGB'\n",
    "solidLUT.NanColor = [1.0, 0.0, 0.0]\n",
    "solidLUT.NumberOfTableValues = 10\n",
    "solidLUT.ScalarRangeInitialized = 1.0\n",
    "solidPWF = pv.GetOpacityTransferFunction('Solid')\n",
    "solidPWF.Points = [-0.14285714285714285, 0.0, 0.5, 0.0, 500.0, 0.0, 0.5, 0.0, 1000.0, 1.0, 0.5, 0.0]\n",
    "solidPWF.ScalarRangeInitialized = 1\n",
    "\n",
    "Display = pv.Show(gradient1, renderView1)\n",
    "Display.ColorArrayName = ['POINTS', 'Vorticity']\n",
    "Display.LookupTable = vorticityLUT\n",
    "Display.ScalarOpacityFunction = vorticityPWF\n",
    "calculator1Display = pv.Show(calculator1, renderView1, 'StructuredGridRepresentation')\n",
    "calculator1Display.ColorArrayName = ['POINTS', 'Sensor']\n",
    "calculator1Display.LookupTable = sensorLUT\n",
    "calculator1Display.ScalarOpacityFunction = sensorPWF\n",
    "calculator2Display = pv.Show(calculator2, renderView1, 'StructuredGridRepresentation')\n",
    "calculator2Display.ColorArrayName = ['POINTS', 'Solid']\n",
    "calculator2Display.LookupTable = solidLUT\n",
    "calculator2Display.ScalarOpacityFunction = solidPWF\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pv.Render()\n",
    "renderView1.CameraParallelScale = 240\n",
    "pv.SaveScreenshot(os.path.join(parent_dir, \"output\", \"test.png\"), view=renderView1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pv.SaveAnimation(filename=os.path.join(parent_dir, \"output\", \"d1a0\", \"anim\", \"anim.png\"), viewOrLayout=renderView1, SuffixFormat='.%02d', FrameWindow=[0, 29], FrameStride=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# help(dat)"
   ]
  }
 ],
 "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
}