{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We will start by fixing the seed for reproducibility\n",
    "seed = 0\n",
    "import torch\n",
    "\n",
    "torch.manual_seed(seed)\n",
    "import numpy as np\n",
    "\n",
    "np.random.seed(seed)\n",
    "import random\n",
    "\n",
    "random.seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import pgmuvi\n",
    "except (ImportError, ModuleNotFoundError):\n",
    "    %pip install git+https://github.com/ICSM/pgmuvi.git\n",
    "    import pgmuvi"

   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now generate some data from a more complex case, maybe drawing from a specific PSD to test if pgmuvi reconstructs it\n",
    "# timestamps_2d = #generate random x data here\n",
    "# wavelengths_2d = #generate random x data here\n",
    "# fluxes_2d = #generate random y data here\n",
    "# flux_err_2d =\n",
    "P0, P1 = 1000, 220\n",
    "n_data = 50\n",
    "n_data1 = n_data\n",
    "n_data2 = n_data // 2\n",
    "jd_min, jd_max = 2450000, 2455000\n",
    "timestamps1_2d = torch.Tensor(np.random.uniform(jd_min, jd_max, size=n_data1))\n",
    "timestamps2_2d = torch.Tensor(np.random.uniform(jd_min, jd_max, size=n_data2))\n",
    "timestamps_2d = torch.concat((timestamps1_2d, timestamps2_2d), dim=0)\n",
    "\n",
    "bands1_2d = torch.Tensor(np.array([0.854] * n_data1))\n",
    "bands2_2d = torch.Tensor(np.array([2.206] * n_data2))\n",
    "bands_2d = torch.concat((bands1_2d, bands2_2d), dim=0)\n",
    "\n",
    "flux1_2d = torch.Tensor(np.sin(timestamps1_2d * (2 * np.pi / P0)))\n",
    "flux1_2d += 0.1 * torch.randn_like(flux1_2d)\n",
    "flux1_err_2d = 0.1 * flux1_2d\n",
    "\n",
    "flux2_2d = torch.Tensor(np.sin(timestamps2_2d * (2 * np.pi / P1) + 0.255 * jd_min))\n",
    "flux2_2d += 0.18 * torch.randn_like(flux2_2d)\n",
    "flux2_err_2d = 0.02 * flux2_2d\n",
    "\n",
    "fluxes_2d = torch.concat((flux1_2d, flux2_2d), dim=0)\n",
    "flux_err_2d = torch.concat((flux1_err_2d, flux2_err_2d), dim=0)\n",
    "\n",
    "print(timestamps_2d.shape, bands_2d.shape, fluxes_2d.shape, flux_err_2d.shape)\n",
    "\n",
    "timestamps_bands_2d = torch.stack((timestamps_2d, bands_2d), 0).T\n",
    "timestamps_bands_2d.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pgmuvi.lightcurve import Lightcurve\n",
    "\n",
    "lightcurve_2d = Lightcurve(\n",
    "    timestamps_bands_2d, fluxes_2d, yerr=flux_err_2d, xtransform=\"minmax\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lightcurve_2d._xdata_transformed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_2d = lightcurve_2d.fit(model=\"2D\", likelihood=\"learn\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_2d[\"covar_module.mixture_means\"][-1][..., 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lightcurve_2d.print_results()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}