{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Split Dataset into validation / test in 50:50 fashion\n",
    "\n",
    "The digests are stored in a json. The dataset can then be filtered with\n",
    "\n",
    "```python\n",
    "from sec_certs.model.evaluation import get_validation_dgsts\n",
    "from sec_cers.dataset import CCDataset\n",
    "\n",
    "dset = CCDataset.from_json() # or call FIPSDataset.from_json()\n",
    "validation_dgsts = get_validation_dgsts('/path/to/validation_set.json')\n",
    "validation_certs = [x for x in dset if x.dgst in validation_dgsts]\n",
    "y_valid = [(x.heuristics.verified_cpe_matches) for x in validation_certs]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from typing import Union\n",
    "from pathlib import Path\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sec_certs.dataset import CCDataset, FIPSDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def split_dataset(dset_type: str, path: Union[str, Path], validation_outpath: Union[str, Path] = './validation_set.json', test_outpath: Union[str, Path] = './test_set.json'):\n",
    "    if dset_type == 'cc':\n",
    "        dset: CCDataset = CCDataset.from_json(path)\n",
    "    elif dset_type == 'fips':\n",
    "        dset: FIPSDataset = FIPSDataset.from_json(path)\n",
    "    else:\n",
    "        raise ValueError(f'type variable must be cc or fips, {dset_type} was given')\n",
    "\n",
    "    cpe_rich_certs = [x for x in dset if x.heuristics.verified_cpe_matches]\n",
    "    cpe_free_certs = [x for x in dset if not x.heuristics.verified_cpe_matches]\n",
    "\n",
    "    x_valid_cpe_rich, x_test_cpe_rich = train_test_split(cpe_rich_certs, test_size=0.5)\n",
    "    x_valid_cpe_free, x_test_cpe_free = train_test_split(cpe_free_certs, test_size=0.5)\n",
    "\n",
    "    validation_set = [x.dgst for x in x_valid_cpe_rich + x_valid_cpe_free]\n",
    "    test_set = [x.dgst for x in x_test_cpe_rich + x_test_cpe_free]\n",
    "\n",
    "    with Path(validation_outpath).open('w') as handle:\n",
    "        json.dump(validation_set, handle, indent=4)\n",
    "\n",
    "    with Path(test_outpath).open('w') as handle:\n",
    "        json.dump(test_set, handle, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "split_dataset('fips', '/path/to/fips_dataset.json')"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "9c2c998cb2293d5512c6ce922cbfe08310964e398b5d5fc2b51147c001724304"
  },
  "kernelspec": {
   "display_name": "Python 3.8.1 64-bit ('3.8.1': pyenv)",
   "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.8.1"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}