{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import struct\n",
    "\n",
    "import numpy as np\n",
    "from scipy import signal, optimize\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import rocof_test_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "fs = 1000 # Hz\n",
    "ff = 50 # Hz\n",
    "duration = 60 # seconds\n",
    "# test_data = rocof_test_data.sample_waveform(rocof_test_data.test_close_interharmonics_and_flicker(),\n",
    "#                                             duration=20,\n",
    "#                                             sampling_rate=fs,\n",
    "#                                             frequency=ff)[0]\n",
    "# test_data = rocof_test_data.sample_waveform(rocof_test_data.gen_noise(fmin=10, amplitude=1),\n",
    "#                                             duration=20,\n",
    "#                                             sampling_rate=fs,\n",
    "#                                             frequency=ff)[0]\n",
    "\n",
    "\n",
    "#gen = rocof_test_data.gen_noise(fmin=10, amplitude=1)\n",
    "# gen = rocof_test_data.gen_noise(fmin=60, amplitude=0.2)\n",
    "# gen = rocof_test_data.test_harmonics()\n",
    "# gen = rocof_test_data.gen_interharmonic(*rocof_test_data.test_interharmonics)\n",
    "# gen = rocof_test_data.test_amplitude_steps()\n",
    "# gen = rocof_test_data.test_amplitude_and_phase_steps()\n",
    "test_data = []\n",
    "test_labels = [ fun.__name__.replace('test_', '') for fun in rocof_test_data.all_tests ]\n",
    "for gen in rocof_test_data.all_tests:\n",
    "    test_data.append(rocof_test_data.sample_waveform(gen(),\n",
    "                                            duration=duration,\n",
    "                                            sampling_rate=fs,\n",
    "                                            frequency=ff)[0])\n",
    "# d = 10 # seconds\n",
    "# test_data = np.sin(2*np.pi * ff * np.linspace(0, d, int(d*fs)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "spr_fmt = f'{fs}Hz' if fs<1000 else f'{fs/1e3:f}'.rstrip('.0') + 'kHz'\n",
    "for label, data in zip(test_labels, test_data):\n",
    "    with open(f'rocof_test_data/rocof_test_{label}_{spr_fmt}.bin', 'wb') as f:\n",
    "        for sample in data:\n",
    "            f.write(struct.pack('<f', sample))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis_periods = 10\n",
    "window_len = 256 # fs * analysis_periods/ff\n",
    "nfft_factor = 1\n",
    "sigma = window_len/8 # samples\n",
    "quantization_bits = 14\n",
    "\n",
    "ffts = []\n",
    "for item in test_data:\n",
    "    f, t, Zxx = signal.stft((item * (2**(quantization_bits-1) - 1)).round().astype(np.int16).astype(float),\n",
    "                fs = fs,\n",
    "                window=('gaussian', sigma),\n",
    "                nperseg = window_len,\n",
    "                nfft = window_len * nfft_factor)\n",
    "                #boundary = 'zeros')\n",
    "    ffts.append((f, t, Zxx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(129, 470)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Zxx.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.90625"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1000/256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "96e3c06718434807813d909d3a1f4d0e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(len(test_data), figsize=(8, 20), sharex=True)\n",
    "fig.tight_layout(pad=2, h_pad=0.1)\n",
    "\n",
    "for fft, ax, label in zip(test_data, ax.flatten(), test_labels):\n",
    "    ax.plot((item * (2**(quantization_bits-1) - 1)).round())\n",
    "    \n",
    "    ax.set_title(label, pad=-20, color='white', bbox=dict(boxstyle=\"square\", ec=(0,0,0,0), fc=(0,0,0,0.8)))\n",
    "    ax.grid()\n",
    "    ax.set_ylabel('f [Hz]')\n",
    "ax.set_xlabel('simulation time t [s]')\n",
    "ax.set_xlim([5000, 5200])\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af9cb4a349eb42c6b89d7951dabdd71e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(len(test_data), figsize=(8, 20), sharex=True)\n",
    "fig.tight_layout(pad=2, h_pad=0.1)\n",
    "\n",
    "for fft, ax, label in zip(ffts, ax.flatten(), test_labels):\n",
    "    f, t, Zxx = fft\n",
    "    ax.pcolormesh(t[1:], f[:250], np.abs(Zxx[:250,1:]))\n",
    "    ax.set_title(label, pad=-20, color='white')\n",
    "    ax.grid()\n",
    "    ax.set_ylabel('f [Hz]')\n",
    "    ax.set_ylim([30, 75]) # Hz\n",
    "ax.set_xlabel('simulation time t [s]')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([  0.     ,   3.90625,   7.8125 ,  11.71875,  15.625  ,  19.53125,\n",
       "        23.4375 ,  27.34375,  31.25   ,  35.15625,  39.0625 ,  42.96875,\n",
       "        46.875  ,  50.78125,  54.6875 ,  58.59375,  62.5    ,  66.40625,\n",
       "        70.3125 ,  74.21875,  78.125  ,  82.03125,  85.9375 ,  89.84375,\n",
       "        93.75   ,  97.65625, 101.5625 , 105.46875, 109.375  , 113.28125,\n",
       "       117.1875 , 121.09375, 125.     , 128.90625, 132.8125 , 136.71875,\n",
       "       140.625  , 144.53125, 148.4375 , 152.34375, 156.25   , 160.15625,\n",
       "       164.0625 , 167.96875, 171.875  , 175.78125, 179.6875 , 183.59375,\n",
       "       187.5    , 191.40625, 195.3125 , 199.21875, 203.125  , 207.03125,\n",
       "       210.9375 , 214.84375, 218.75   , 222.65625, 226.5625 , 230.46875,\n",
       "       234.375  , 238.28125, 242.1875 , 246.09375, 250.     , 253.90625,\n",
       "       257.8125 , 261.71875, 265.625  , 269.53125, 273.4375 , 277.34375,\n",
       "       281.25   , 285.15625, 289.0625 , 292.96875, 296.875  , 300.78125,\n",
       "       304.6875 , 308.59375, 312.5    , 316.40625, 320.3125 , 324.21875,\n",
       "       328.125  , 332.03125, 335.9375 , 339.84375, 343.75   , 347.65625,\n",
       "       351.5625 , 355.46875, 359.375  , 363.28125, 367.1875 , 371.09375,\n",
       "       375.     , 378.90625, 382.8125 , 386.71875, 390.625  , 394.53125,\n",
       "       398.4375 , 402.34375, 406.25   , 410.15625, 414.0625 , 417.96875,\n",
       "       421.875  , 425.78125, 429.6875 , 433.59375, 437.5    , 441.40625,\n",
       "       445.3125 , 449.21875, 453.125  , 457.03125, 460.9375 , 464.84375,\n",
       "       468.75   , 472.65625, 476.5625 , 480.46875, 484.375  , 488.28125,\n",
       "       492.1875 , 496.09375, 500.     ])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "930f346b46d64489bdc790ec11ef9a69",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "slice(12, 15, None)\n",
      "slice(12, 15, None)\n",
      "slice(12, 15, None)\n",
      "slice(12, 15, None)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/user/safety-reset/lab-windows/env/lib/python3.8/site-packages/scipy/optimize/minpack.py:807: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "slice(12, 15, None)\n",
      "slice(12, 15, None)\n",
      "slice(12, 15, None)\n",
      "slice(12, 15, None)\n",
      "slice(12, 15, None)\n",
      "slice(12, 15, None)\n"
     ]
    }
   ],
   "source": [
    "fig, axs = plt.subplots(len(test_data), figsize=(8, 20), sharex=True)\n",
    "fig.tight_layout(pad=2.2, h_pad=0, w_pad=1)\n",
    "\n",
    "for fft, ax, label in zip(ffts, axs.flatten(), test_labels):\n",
    "    f, f_t, Zxx = fft\n",
    "    \n",
    "    n_f, n_t = Zxx.shape\n",
    "    # start, stop = 180, 220\n",
    "    # start, stop = 90, 110\n",
    "    # start, stop = 15, 35\n",
    "    # bounds_f = slice(start // 4 * nfft_factor, stop // 4 * nfft_factor)\n",
    "    f_min, f_max = 30, 70 # Hz\n",
    "    bounds_f = slice(np.argmax(f > f_min), np.argmin(f < f_max))\n",
    "    \n",
    "\n",
    "    f_mean = np.zeros(Zxx.shape[1])\n",
    "    for t in range(1, Zxx.shape[1] - 1):\n",
    "        frame_f = f[bounds_f]\n",
    "        frame_step = frame_f[1] - frame_f[0]\n",
    "        time_step = f_t[1] - f_t[0]\n",
    "        #if t == 10:\n",
    "        #    axs[-1].plot(frame_f, frame_Z)\n",
    "        frame_Z = np.abs(Zxx[bounds_f, t])\n",
    "        # frame_f = f[180:220]\n",
    "        # frame_Z = np.abs(Zxx[180:220, 40])\n",
    "        # frame_f = f[15:35]\n",
    "        # frame_Z = np.abs(Zxx[15:35, 40])\n",
    "        # plt.plot(frame_f, frame_Z)\n",
    "\n",
    "        # peak_f = frame_f[np.argmax(frame)]\n",
    "        # plt.axvline(peak_f, color='red')\n",
    "\n",
    "#         def gauss(x, *p):\n",
    "#             A, mu, sigma, o = p\n",
    "#             return A*np.exp(-(x-mu)**2/(2.*sigma**2)) + o\n",
    "        \n",
    "        def gauss(x, *p):\n",
    "            A, mu, sigma = p\n",
    "            return A*np.exp(-(x-mu)**2/(2.*sigma**2))\n",
    "        \n",
    "        f_start = frame_f[np.argmax(frame_Z)]\n",
    "        A_start = np.max(frame_Z)\n",
    "        p0 = [A_start, f_start, 1.]\n",
    "        try:\n",
    "            coeff, var = optimize.curve_fit(gauss, frame_f, frame_Z, p0=p0)\n",
    "            # plt.plot(frame_f, gauss(frame_f, *coeff))\n",
    "            #print(coeff)\n",
    "            A, mu, sigma, *_ = coeff\n",
    "            f_mean[t] = mu\n",
    "        except RuntimeError:\n",
    "            f_mean[t] = np.nan\n",
    "    ax.plot(f_t[1:-1], f_mean[1:-1])\n",
    "    \n",
    "#     b, a = signal.butter(3,\n",
    "#                          1/5, # Hz\n",
    "#                          btype='lowpass',\n",
    "#                          fs=1/time_step)\n",
    "#     filtered = signal.lfilter(b, a, f_mean[1:-1], axis=0)\n",
    "#     ax.plot(f_t[1:-1], filtered)\n",
    "    \n",
    "    ax.set_title(label, pad=-20)\n",
    "    ax.set_ylabel('f [Hz]')\n",
    "    ax.grid()\n",
    "    if not label in ['off_frequency', 'sweep_phase_steps']:\n",
    "        ax.set_ylim([49.90, 50.10])\n",
    "        var = np.var(f_mean[1:-1])\n",
    "        ax.text(0.5, 0.1, f'σ²={var * 1e3:.3g} mHz²', transform=ax.transAxes, ha='center')\n",
    "        ax.text(0.5, 0.25, f'σ={np.sqrt(var) * 1e3:.3g} mHz', transform=ax.transAxes, ha='center')\n",
    "#         ax.text(0.5, 0.2, f'filt. σ²={np.var(filtered) * 1e3:.3g} mHz', transform=ax.transAxes, ha='center')\n",
    "    else:\n",
    "        f_min, f_max = min(f_mean[1:-1]), max(f_mean[1:-1])\n",
    "        delta = f_max - f_min\n",
    "        ax.set_ylim(f_min - delta * 0.1, f_max + delta * 0.3)\n",
    "        \n",
    "ax.set_xlabel('simulation time t [s]')\n",
    "None"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "labenv",
   "language": "python",
   "name": "labenv"
  },
  "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}