#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Author: Tomas Cassanelli
import os
import yaml
import numpy as np
from astropy.io import fits
from astropy.time import Time
from astropy.table import Table, QTable
from astropy import units as apu
from astropy.constants import c as light_speed
from .math_functions import rms
from .aperture import e_rs, phase

__all__ = [
    'extract_data_pyoof', 'extract_data_effelsberg', 'str2LaTeX',
    'store_data_csv', 'uv_ratio', 'store_data_ascii', 'table_pyoof_out'
    ]

def extract_data_pyoof(pathfits):
    """
    Extracts data from the `~pyoof` default FITS file OOF holography
    observations, ready to use for the least squares minimization (see
    `~pyoof.fit_zpoly`). The FITS file has to have the following keys on its
    PrimaryHDU header: ``'FREQ'``, ``'WAVEL'``, ``'MEANEL'``, ``'OBJECT'`` and
    ``'DATE_OBS'``. Besides this three BinTableHDU are required for the data
    itself; ``MINUS OOF``, ``ZERO OOF`` and ``PLUS OOF``. The BinTableHDU
    header has to have the ``'DZ'`` key which includes the radial offset,
    :math:`d_z`. Finally the BinTableHDU has the data files ``'U'``, ``'V'``
    and ``'BEAM'``, which is the :math:`x`- and :math:`y`-axis position in
    radians and the ``'BEAM'`` in a flat array, in mJy.

    Parameters
    ----------
    pathfits : `str`
        Path to the FITS file that contains the three beam maps pre-calibrated,
        using the correct PrimaryHDU and the three BinTableHDU (``MINUS OOF``,
        ``ZERO OOF`` and ``PLUS OOF``).

    Returns
    -------
    data_info : `list`
        It contains all extra data besides the beam map. The output
        corresponds to a list,
        ``[name, pthto, obs_object, obs_date, freq, wavel, d_z, meanel]``.
        These are, name of the FITS file, paht of the FITS file, observed
        object, observation date, frequency, wavelength, radial offset and
        mean elevation, respectively.
    data_obs : `list`
        It contains beam maps and :math:`x`-, and :math:`y`-axis
        (:math:`uv`-plane in Fourier space) data for the least squares
        minimization (see `~pyoof.fit_zpoly`). The list has the following order
        ``[beam_data, u_data, v_data]``. ``beam_data`` is the three beam
        observations, minus, zero and plus out-of-focus, in a flat array.
        ``u_data`` and ``v_data`` are the beam axes in a flat array.

    Raises
    ------
    `ValueError`
        If the input file ``pathfits`` is not a FITS file, or if one or more
        of the header keys are missing in the FITS file:
        ``['FREQ', 'WAVEL', 'MEANEL', 'OBJECT', 'DATE_OBS']``.
    """

    if os.path.splitext(pathfits)[1] != '.fits':
        raise ValueError('File must be a FITS file.')

    hdulist = fits.open(pathfits)  # open FITS file, pyoof format
    # path or directory where the FITS file is located
    pthto = os.path.split(pathfits)[0]
    # name of the fit file to fit
    name = os.path.split(pathfits)[1][:-5]

    if not all(
            k in hdulist[0].header
            for k in ['FREQ', 'WAVEL', 'MEANEL', 'OBJECT', 'DATE_OBS']
            ):
        raise ValueError('Not all needed keys found in FITS header.')

    freq = hdulist[0].header['FREQ'] * apu.Hz
    wavel = hdulist[0].header['WAVEL'] * apu.m
    meanel = hdulist[0].header['MEANEL'] * apu.deg
    obs_object = hdulist[0].header['OBJECT']
    obs_date = hdulist[0].header['DATE_OBS']

    beam_data = np.array([hdulist[i].data['BEAM'] for i in range(1, 4)])
    u_data = np.array([hdulist[i].data['U'] for i in range(1, 4)]) * apu.rad
    v_data = np.array([hdulist[i].data['V'] for i in range(1, 4)]) * apu.rad
    d_z = np.array([hdulist[i].header['DZ'] for i in range(1, 4)]) * apu.m

    data_file = [name, pthto]
    data_info = data_file + [obs_object, obs_date, freq, wavel, d_z, meanel]
    data_obs = [beam_data, u_data, v_data]

    return data_info, data_obs

def extract_data_effelsberg(pathfits):
    """
    Extracts data from the Effelsberg OOF holography observations, ready to
    use for the least squares minimization. This function will only work for
    the Effelsberg telescope beam maps.

    Parameters
    ----------
    pathfits : `str`
        Path to the FITS file that contains the three beam maps pre-calibrated,
        from the Effelsberg telescope.

    Returns
    -------
    data_info : `list`
        It contains all extra data besides the beam map. The output
        corresponds to a list,
        ``[name, pthto, obs_object, obs_date, freq, wavel, d_z, meanel]``.
        These are, name of the FITS file, paht of the FITS file, observed
        object, observation date, frequency, wavelength, radial offset and
        mean elevation, respectively.
    data_obs : `list`
        It contains beam maps and :math:`x`-, and :math:`y`-axis
        (:math:`uv`-plane in Fourier space) data for the least squares
        minimization (see `~pyoof.fit_zpoly`). The list has the following order
        ``[beam_data, u_data, v_data]``. ``beam_data`` is the three beam
        observations, minus, zero and plus out-of-focus, in a flat array.
        ``u_data`` and ``v_data`` are the beam axes in a flat array.

    Raises
    ------
    `ValueError`
        If the input file ``pathfits`` is not a FITS file.
    """

    if os.path.splitext(pathfits)[1] != '.fits':
        raise ValueError('File must be a FITS file.')

    pos = [3, 1, 2]  # Positions for OOF holography observations at Effelsberg
    hdulist = fits.open(pathfits)  # main FITS file OOF holography format

    # Observation frequency
    freq = hdulist[0].header['FREQ'] * apu.Hz
    wavel = light_speed / freq

    # Mean elevation
    meanel = hdulist[0].header['MEANEL'] * apu.deg
    obs_object = hdulist[0].header['OBJECT']        # observed object
    obs_date = hdulist[0].header['DATE_OBS']        # observation date
    d_z = np.array([hdulist[i].header['DZ'] for i in pos]) * apu.m

    beam_data = np.array([hdulist[i].data['fnu'] for i in pos])
    u_data = np.array([hdulist[i].data['DX'] for i in pos]) * apu.rad
    v_data = np.array([hdulist[i].data['DY'] for i in pos]) * apu.rad

    # path or directory where the FITS file is located
    pthto = os.path.split(pathfits)[0]
    # name of the fit file to fit
    name = os.path.split(pathfits)[1][:-5]

    data_info = [name, pthto, obs_object, obs_date, freq, wavel, d_z, meanel]
    data_obs = [beam_data, u_data, v_data]

    return data_info, data_obs

def str2LaTeX(python_string):
    """
    Function that solves the underscore problem in a python string to
    :math:`\\LaTeX` string.

    Parameters
    ----------
    python_string : `str`
        String that needs to be changed.

    Returns
    -------
    LaTeX_string : `str`
        String with the new underscore symbol.
    """

    string_list = list(python_string)
    for idx, string in enumerate(string_list):
        if string_list[idx] == '_':
            string_list[idx] = '\\_'

    LaTeX_string = ''.join(string_list)

    return LaTeX_string

def store_data_csv(name, name_dir, order, save_to_csv):
    """
    Stores all important information in a CSV file after the least squares
    minimization has finished, `~pyoof.fit_zpoly`. All data will be stored in
    the ``pyoof_out/name-number`` directory, with ``name`` the name of the
    FITS file, and ``number`` the `~pyoof.fit_zpoly` number of code execution,
    i.e. the output data is never overwritten.

    Parameters
    ----------
    name : `str`
        File name of the FITS file to be optimized.
    name_dir : `str`
        Path to store all the CSV files. The files will depend on the order of
        the Zernike circle polynomial.
    order : `int`
        Order used for the Zernike circle polynomial, :math:`n`.
    save_to_csv : `list`
        It contains all data that will be stored. The list must have the
        following order, ``[beam_data, u_data, v_data, res_optim, jac_optim,
        grad_optim, phase, cov_ptrue, corr_ptrue]``.
    """

    headers = [
        'Normalized beam', 'u vector radians', 'v vector radians', 'Residual',
        'Jacobian', 'Gradient', 'Phase-error radians',
        'Variance-Covariance matrix (first row fitted parameters idx)',
        'Correlation matrix (first row fitted parameters idx)'
        ]

    fnames = [
        f'beam_data.csv', f'u_data.csv', f'v_data.csv',
        f'res_n{order}.csv', f'jac_n{order}.csv',
        f'grad_n{order}.csv', f'phase_n{order}.csv',
        f'cov_n{order}.csv', f'corr_n{order}.csv'
        ]

    if order != 1:
        headers = headers[3:]
        fnames = fnames[3:]
        save_to_csv = save_to_csv[3:]

    for fname, header, file in zip(fnames, headers, save_to_csv):
        np.savetxt(
            fname=os.path.join(name_dir, fname),
            X=file,
            header=' '.join((header, name))
            )

def store_data_ascii(name, name_dir, order, params_solution, params_init):
    """
    Stores in an ``~astropy.table.table.Table`` format the parameters found by
    the least squares minimization (see `~pyoof.fit_zpoly`).

    Parameters
    ----------
    name : `str`
        File name of the FITS file to be optimized.
    name_dir : `str`
        Path to store all the csv files. The files will depend on the order of
        the Zernike circle polynomial (radial) order.
    order : `int`
        Order used for the Zernike circle polynomial, :math:`n`.
    params_solution : `~numpy.ndarray`
        Contains the best fitted parameters, the illumination function
        coefficients, ``I_coeff`` and the Zernike circle polynomial
        coefficients, ``K_coeff`` in one array.
    params_init : `~numpt.ndarray`
        Contains the initial parameters used in the least squares minimization
        to start finding the best fitted combination of them.
    """

    n = order
    N_K_coeff = (n + 1) * (n + 2) // 2

    # Making nice table :)
    ln = [(j, i) for i in range(0, n + 1) for j in range(-i, i + 1, 2)]
    L = np.array(ln)[:, 0]
    N = np.array(ln)[:, 1]

    params_names = ['i_amp', 'c_dB', 'q', 'x_0', 'y_0']
    for i in range(N_K_coeff):
        params_names.append(f'K({N[i]}, {L[i]})')

    # To store fit information and found parameters in ascii file
    tab = Table(
        data=[params_names, params_solution, params_init],
        names=['parname', 'parfit', 'parinit'],
        )

    tab.write(
        os.path.join(name_dir, f'fitpar_n{n}.csv'),
        overwrite=True
        )

def uv_ratio(u, v):
    """
    Calculates the aspect ratio for the 3 power pattern plots, plus some
    corrections for the text on it. Used in the `function` `~pyoof.plot_beam`
    and `~pyoof.plot_data`

    Parameters
    ----------
    u : `~astropy.units.quantity.Quantity`
        Spatial frequencies from the power pattern in angle units.
    v : `~astropy.units.quantity.Quantity`
        Spatial frequencies from the power pattern in angle units.

    Returns
    -------
    width : `float`
        Width for the power pattern figure.
    height : `float`
        Height for the power pattern figure.
    """

    ratio = (u.max() - u.min()) / (v.max() - v.min())

    if type(ratio) == apu.Quantity:
        ratio = ratio.decompose().value

    height = 5
    width = (ratio * 2.25 * height)

    return width, height

def table_pyoof_out(path_pyoof_out, order):
    """
    Auxiliary function to tabulate all data from a series of observations
    gathered in a common ``pyoof_out/`` directory.

    Note: Piston and tilt are not used in error calculations, the phase
    calculations will only be included if these are manually changed in
    ``core.py``.

    Parameters
    ----------
    path_pyoof_out : `list`
        set of paths to the directory ``pyoof_out/`` or where the output from
        the `~pyoof` package is located.
    order : `int`
        Order used for the Zernike circle polynomial, :math:`n`.

    Returns
    -------
    qt : `~astropy.table.table.QTable`
        `~astropy.table.table.QTable` with units of the most important
        quantities from the `~pyoof` package.
    """

    qt = QTable(
        names=[
            'name', 'tel_name', 'obs-object', 'obs-date', 'meanel',
            'i_amp', 'c_dB', 'q', 'phase-rms', 'e_rs',
            'beam-snr-out-l', 'beam-snr-in', 'beam-snr-out-r'
            ],
        dtype=[np.string_] * 4 + [np.float] * 9
        )

    for p, pyoof_out in enumerate(path_pyoof_out):
        with open(os.path.join(pyoof_out, 'pyoof_info.yml'), 'r') as inputfile:
            pyoof_info = yaml.load(inputfile, Loader=yaml.Loader)

        _phase = np.genfromtxt(
            os.path.join(pyoof_out, f'phase_n{order}.csv')
            ) * apu.rad
        phase_rms = rms(_phase, circ=True)
        phase_e_rs = e_rs(_phase, circ=True)

        # random-surface-error efficiency error
        # cov = np.genfromtxt(os.path.join(pyoof_out, f'cov_n{order}.csv'))
        # idx = np.argwhere(cov[0, :].astype(int) > 5)

        params = Table.read(
            os.path.join(pyoof_out, f'fitpar_n{order}.csv'),
            format='ascii'
            )
        I_coeff = params['parfit'][:5]

        qt.add_row([
            pyoof_info['name'], pyoof_info['tel_name'],
            pyoof_info['obs_object'], pyoof_info['obs_date'],
            pyoof_info['meanel'], I_coeff[0], I_coeff[1], I_coeff[2],
            phase_rms, phase_e_rs
            ] + pyoof_info['snr'])

    # updating units
    qt['phase-rms'] *= apu.rad
    qt['meanel'] *= apu.deg
    qt['obs-date'] = Time(qt['obs-date'], format='isot', scale='utc')
    qt['c_dB'] *= apu.dB

    qt.meta = {'order': order}

    return qt


from __future__ import division
import torch

# general libs
import numpy as np
import torch.nn.functional as F

# my libs
# from Algorithm_Heo.networks.network_CerBerusnet_v6 import CBnet
from .networks.network import NET_GAmap
from libs import utils_custom

# davis

class model():
    def __init__(self):
        self.net = NET_GAmap()
        self.net.cuda()
        self.net.load_state_dict(torch.load('checkpoints/GIS-ckpt_standard.pth'))
        self.net.eval()
        for param in self.net.parameters(): param.requires_grad = False

        self.mean, self.var = torch.Tensor([0.485, 0.456, 0.406]), torch.Tensor([0.229, 0.224, 0.225])
        self.mean, self.var = self.mean.view(1,1,3,1,1).cuda(),  self.var.view(1,1,3,1,1).cuda()

    def init_with_new_video(self, frames, n_obj=1):
        self.n_objects = n_obj
        self.frames = frames.copy() # f h w 3
        self.num_frames, self.height, self.width = self.frames.shape[:3]
        self.scores_nf = np.zeros([self.num_frames])
        self.annotated_frames=[]

        pad_info = utils_custom.apply_pad(self.frames[0])[1]
        self.hpad1, self.wpad1 = pad_info[0][0], pad_info[1][0]
        self.hpad2, self.wpad2 = pad_info[0][1], pad_info[1][1]
        self.padding = pad_info[1] + pad_info[0]
        self.prob_map_of_frames = torch.zeros((self.num_frames, self.n_objects+1, self.height + sum(pad_info[0]), self.width + sum(pad_info[1]))).requires_grad_(False).cuda() # f,1,p_h,p_w {cudatensor}

        self.all_F = torch.unsqueeze(torch.nn.ReflectionPad2d(self.padding)(
            torch.from_numpy(np.transpose(frames, (0,3,1,2))).float() / 255.), dim=1).requires_grad_(False).cuda()  # fr,1,3,p_h,p_w {cudatensor}
        self.all_F = (self.all_F -self.mean) / self.var
        self.current_round_masks = np.zeros([self.num_frames, self.height, self.width])  # f,h,w {numpy}
        self.anno_6chEnc_r4_list = []
        self.anno_3chEnc_r4_list = []
        self.current_round = 1

        self.one_hot_outputs_anno = None
        self.current_round = 1

    def Run_propagation(self, annotated_now, ):

        print('[Propagation running...]')
        self.current_round +=1
        prop_list = utils_custom.get_prop_list(self.annotated_frames, annotated_now, self.num_frames, proportion=0.99)
        annotated_frames_np = np.array(self.annotated_frames)
        prop_fore = sorted(prop_list)[0]
        prop_rear = sorted(prop_list)[-1]

        flag = 0  # 1: propagating backward, 2: propagating forward
        for operating_frame in prop_list:
            if operating_frame == annotated_now:
                if flag == 0:
                    flag += 1
                    adjacent_to_anno = True
                    continue
                elif flag == 1:
                    flag += 1
                    adjacent_to_anno = True
                    continue
                else:
                    raise NotImplementedError
            else:
                print('operating in : {:03d}'.format(operating_frame))
                if adjacent_to_anno:
                    r4_neighbor = self.r4_anno
                    neighbor_pred_onehot = self.anno_onehot_prob
                    adjacent_to_anno = False
                else:
                    r4_neighbor = r4_que
                    neighbor_pred_onehot = targ_onehot_prob

                output_logit, r4_que, score = self.net.forward_prop(
                    self.anno_3chEnc_r4_list, self.all_F[operating_frame].repeat(self.n_objects,1,1,1), self.anno_6chEnc_r4_list,
                    r4_neighbor, neighbor_pred_onehot,
                    anno_fr_list= annotated_frames_np, que_fr= operating_frame)  # [nobj, 1, P_H, P_W]

                output_prob_tmp = F.softmax(output_logit, dim=1) # [nobj, 2, P_H, P_W]
                output_prob_tmp = output_prob_tmp[:, 1] # [nobj, P_H, P_W]
                one_hot_outputs_t = F.softmax(self.soft_aggregation(output_prob_tmp), dim=0) # [nobj+1, P_H, P_W]

                smallest_alpha = 0.5
                if flag==1:
                    sorted_frames = annotated_frames_np[annotated_frames_np < annotated_now]
                    if len(sorted_frames) ==0:
                        alpha = 1
                    else:
                        closest_addianno_frame = np.max(sorted_frames)
                        alpha = smallest_alpha+(1-smallest_alpha)*((operating_frame-closest_addianno_frame)/(annotated_now - closest_addianno_frame))
                else:
                    sorted_frames = annotated_frames_np[annotated_frames_np > annotated_now]
                    if len(sorted_frames) == 0:
                        alpha = 1
                    else:
                        closest_addianno_frame = np.min(sorted_frames)
                        alpha = smallest_alpha+(1-smallest_alpha)*((closest_addianno_frame - operating_frame) / (closest_addianno_frame - annotated_now))

                one_hot_outputs_t = (alpha * one_hot_outputs_t) + ((1 - alpha) * self.prob_map_of_frames[operating_frame])
                self.prob_map_of_frames[operating_frame] = one_hot_outputs_t
                targ_onehot_prob = one_hot_outputs_t.clone()[1:].unsqueeze(1) # [nobj, 1, P_H, P_W]

                self.scores_nf[operating_frame] = score

        self.current_round_masks = torch.argmax(self.prob_map_of_frames,dim=1).cpu().numpy().astype(np.uint8)[:,self.hpad1:-self.hpad2, self.wpad1:-self.wpad2]

        print('[Propagation process is done.]')

    def Run_interaction(self, scribbles):

        print('[Interaction running...]')
        annotated_now = scribbles['annotated_frame']
        scribbles_list = scribbles['scribbles']

        pm_ps_ns_3ch_t=[] # n_obj,3,h,w
        if self.current_round == 1:
            for obj_id in range(1, self.n_objects + 1):
                pos_scrimg, neg_scrimg = utils_custom.scribble_to_image(scribbles_list, annotated_now, obj_id,
                                                                        prev_mask=self.current_round_masks[annotated_now], blur=True,
                                                                        singleimg=False, seperate_pos_neg=True)
                pm_ps_ns_3ch_t.append(np.stack([np.ones_like(pos_scrimg)/2, pos_scrimg, neg_scrimg], axis=0))
            pm_ps_ns_3ch_t = np.stack(pm_ps_ns_3ch_t, axis=0) # n_obj,3,h,w

        else:
            for obj_id in range(1, self.n_objects + 1):
                prev_round_input = (self.current_round_masks[annotated_now] == obj_id).astype(np.float32)  # H,W
                pos_scrimg, neg_scrimg = utils_custom.scribble_to_image(scribbles_list, annotated_now, obj_id,
                                                                        prev_mask=self.current_round_masks[annotated_now], blur=True,
                                                                        singleimg=False, seperate_pos_neg=True)
                pm_ps_ns_3ch_t.append(np.stack([prev_round_input, pos_scrimg, neg_scrimg], axis=0))
            pm_ps_ns_3ch_t = np.stack(pm_ps_ns_3ch_t, axis=0)  # n_obj,3,h,w

        batched_F = self.all_F[annotated_now].repeat(self.n_objects,1,1,1)

        pm_ps_ns_3ch_t = torch.from_numpy(pm_ps_ns_3ch_t).cuda()
        pm_ps_ns_3ch_t = torch.nn.ReflectionPad2d(self.padding)(pm_ps_ns_3ch_t)
        inputs = torch.cat([batched_F, pm_ps_ns_3ch_t], dim=1)

        anno_3chEnc_r4, r2_prev_fromanno = self.net.encoder_3ch.forward(batched_F)
        neighbor_pred_onehot_sal, anno_6chEnc_r4 = self.net.forward_obj_feature_extractor(inputs)  # [nobj, 1, P_H, P_W], # [n_obj,2048,h/16,w/16]

        output_logit, self.r4_anno, score = self.net.forward_prop(
            [anno_3chEnc_r4], batched_F, [anno_6chEnc_r4],
            anno_3chEnc_r4, torch.sigmoid(neighbor_pred_onehot_sal))  # [nobj, 1, P_H, P_W]

        output_prob_tmp = F.softmax(output_logit, dim=1) # [nobj, 2, P_H, P_W]
        output_prob_tmp = output_prob_tmp[:, 1] # [nobj, P_H, P_W]
        one_hot_outputs_t = F.softmax(self.soft_aggregation(output_prob_tmp), dim=0) # [nobj+1, P_H, P_W]

        self.anno_onehot_prob = one_hot_outputs_t.clone()[1:].unsqueeze(1) # [nobj, 1, P_H, P_W]
        self.prob_map_of_frames[annotated_now] = one_hot_outputs_t
        self.current_round_masks[annotated_now] = \
            torch.argmax(self.prob_map_of_frames[annotated_now],dim=0).cpu().numpy().astype(np.uint8)[self.hpad1:-self.hpad2, self.wpad1:-self.wpad2]
        self.scores_nf[annotated_now] = score

        if len(self.anno_6chEnc_r4_list) < self.current_round:
            self.anno_6chEnc_r4_list.append(anno_6chEnc_r4)
            self.anno_3chEnc_r4_list.append(anno_3chEnc_r4)
            self.annotated_frames.append(annotated_now)
        elif len(self.anno_6chEnc_r4_list) == self.current_round:
            self.anno_6chEnc_r4_list[self.current_round-1] = anno_6chEnc_r4
            self.anno_3chEnc_r4_list[self.current_round-1] = anno_3chEnc_r4
        else:
            raise NotImplementedError

        print('[Interaction process is done.]')

    def Get_mask(self):
        return self.current_round_masks

    def Get_mask_index(self, index):
        return self.current_round_masks[index]

    def soft_aggregation(self, ps):
        num_objects, H, W = ps.shape
        em = torch.zeros(num_objects +1, H, W).cuda()
        em[0] =  torch.prod(1-ps, dim=0) # bg prob
        em[1:num_objects+1] = ps # obj prob
        em = torch.clamp(em, 1e-7, 1-1e-7)
        logit = torch.log((em /(1-em)))
        return logit


import os
import re
import sys
import bz2
import bct
import math
import random
import pickle
import numpy as np
import scipy as sp
import pandas as pd
import dionysus as d
import networkx as nx
import cpnet as cpa
import gensim.utils as gu
import gensim.models as gm
import gensim.matutils as gmat
import gensim.parsing.preprocessing as gpp
import mwparserfromhell as mph
import xml.etree.ElementTree as ET
import sklearn.metrics.pairwise as smp

class Dump:
    """``Dump`` loads and parses dumps from wikipedia from
    ``path_xml`` with index ``path_idx``.

    Attributes
    ----------
    idx: dictionary
        ``{'page_name': (byte offset, page id, block size)}``
        Cached. Lazy.
    links: list of strings
        All links.
    article_links: list of strings
        Article links (not files, categories, etc.)
    years: list of int
        Years in the History section of a wikipedia page
        BC denoted as negative values
    page: mwparserfromhell.wikicode
        Current loaded wiki page
    path_xml: string
        Path to the zipped XML dump file.
    path_idx: string
        Path to the zipped index file.
    offset_max: int
        Maximum offset. Set as the size of the zipped dump.
    cache: xml.etree.ElementTree.Node
        Cache of the XML tree in current block
    """
    MAX_YEAR = 2020

    def __init__(self, path_xml, path_idx):
        self._idx = {}
        self._links = []
        self._article_links = []
        self._years = []
        self._page = None
        self.path_xml = path_xml
        self.path_idx = path_idx
        self.offset_max = 0
        self.cache = (0, None) # offset, cache

    @property
    def idx(self):
        if self._idx:
            return self._idx
        else:
            print('Dump: Loading index...')
            with bz2.BZ2File(self.path_idx, 'rb') as file:
                lines = [line for line in file]
            block_end = os.path.getsize(self.path_xml)
            offset_prev = block_end
            for line in reversed(lines):
                offset, pid, name = line.strip().split(b':', 2)
                offset, pid, name = (int(offset), int(pid), name.decode('utf8'))
                block_end = offset_prev if offset < offset_prev else block_end
                self._idx[name] = (offset, pid, block_end-offset)
                offset_prev = offset
            self.offset_max = max([x[0] for x in self._idx.values()])
            print('Dump: Loaded.')
            return self._idx

    @property
    def links(self):
        if self._links:
            return self._links
        elif self.page:
            self._links = [str(x.title) for x in self.page.filter_wikilinks()]
            self._links = [link.split('#')[0] for link in self._links]
            self._links = [link.split(' ') for link in self._links]
            self._links = [[words[0].capitalize()] + words[1:] for words in self._links]
            self._links = [' '.join(words) for words in self._links]
            return self._links
        else:
            return self._links

    @property
    def article_links(self):
        if self._article_links:
            return self._article_links
        elif self.links:
            self._article_links = [x for x in self.links if ':' not in x]
            return self._article_links
        else:
            return self._article_links

    @property
    def years(self):
        if self._years:
            return self._years
        elif self.page:
            history = Dump.get_history(self.page)
            top = self.page.get_sections()[0].strip_code()
            self._years = Dump.filter_years(top + history)
            return self._years
        else:
            return self._years

    @property
    def page(self):
        return self._page

    @page.setter
    def page(self, page):
        self._page = page
        self._links = []
        self._article_links = []
        self._years = []

    def load_page(self, page_name, filter_top=False):
        """Loads & returs page (``mwparserfromhell.wikicode``)
        named ``page_name`` from dump file. Returns only the
        top section if ``filter_top``.
        """
        if page_name not in self.idx.keys():
            self.page = None
            return
        offset, pid, block_size = self.idx[page_name]
        if offset == self.cache[0]:
            root = self.cache[1]
        else:
            xml = Dump.fetch_block(self.path_xml, offset, block_size)
            xml = b'<mediawiki>' + xml + b'</mediawiki>'*(offset != self.offset_max)
            root = ET.fromstring(xml)
            self.cache = (offset, root)
        text = Dump.search_id(root, pid)
        text = Dump.filter_top_section(text) if filter_top else text
        self.page = mph.parse(text, skip_style_tags = True)
        if self.page and 'REDIRECT' in self.page.strip_code():
            redirect = self.page.filter_wikilinks()[0].title
            return self.load_page(str(redirect))
        else:
            return self.page

    @staticmethod
    def fetch_block(path, offset, block_size):
        """ Fetches block of ``block_size`` (``int``) bytes
        at ``offset`` (``int``) in the zipped dump at
        ``path`` (``string``) and returns the uncompressed
        text (``string``).
        """
        with open(path, 'rb') as file:
            file.seek(offset)
            return bz2.decompress(file.read(block_size))

    @staticmethod
    def search_id(root, pid):
        """Returns the text of the page with id ``pid``"""
        for page in root.iter('page'):
            if pid == int(page.find('id').text):
                return page.find('revision').find('text').text

    @staticmethod
    def filter_top_section(text):
        """Returns the top section of text,
        where the first header has the form ``==Heading==``
        """
        head = re.search(r'==.*?==', text)
        idx = head.span(0)[0] if head else len(text)
        return text[:idx] #(text[:idx], text[idx:])

    @staticmethod
    def get_history(page):
        """Returns the text of the history section.
        Returns ``""`` if not found.
        """
        headings = page.filter_headings()
        idx = [i for i, head in enumerate(headings)
                       if 'History' in head or 'history' in head]
        if not idx:
            return ""
        sections = page.get_sections(include_headings=True)
        history = str(sections[idx[0]+1].strip_code())
        return history

    @staticmethod
    def filter_years(text):
        """Filters the years from text."""
        months = ['january', 'february', 'march', 'april', 'may', 'june',
                  'july', 'august', 'september', 'october', 'november', 'december']
        prepositions = ['around', 'after', 'at', 'as',
                        'approximately', 'before', 'between', 'by',
                        'during', 'from', 'in', 'near', 'past',
                        'since', 'until', 'within'] # removed: about, on
        conjugations = ['and']
        articles = ['the']
        times = ['early', 'mid', 'late']
        patterns = months + prepositions + conjugations + articles + times
        re_string = r'\b(' + '|'.join(patterns) + r')\b(\s|-)\b([0-9]{3,4})s?\b(?i)(?!\sMYA)\s?(BCE|BC)?'
        years = [int(match.group(3)) * (-2*bool(match.group(4))+1)
                for match in re.finditer(re_string, text, re.IGNORECASE)]
        re_string = r'([0-9]{1,2})(st|nd|rd|th) century\s?(BCE|BC)?'
        centuries = [(int(match.group(1)) * 100 - 100) * (-2*bool(match.group(2))+1)
                     for match in re.finditer(re_string, text, re.IGNORECASE)]
        years += centuries
        years = [y for y in years if y<Dump.MAX_YEAR]
        return sorted(years + centuries)

class Corpus:
    """``Corpus`` is an ``iterable`` & an ``iterator``
    that uses ``Dump`` to iterate through articles.

    Parameters
    ----------
    dump: wiki.Dump
    output: string
        'doc' for array of documents
        'tag' for TaggedDocument(doc, [self.i])
        'bow' for bag of words [(int, int)]
    dct: gensim.corpus.Dictionary
        used to create BoW representation
    """
    def __init__(self, dump, output='doc', dct=None):
        self.dump = dump
        self.names = list(self.dump.idx.keys())
        self.output = output
        self.dct = dct

    def __iter__(self):
        self.i = 0
        return self

    def __next__(self):
        if self.i < len(self.names):
            sys.stdout.write("\rCorpus index: " + str(self.i+1) +
                             '/' + str(len(self.names)))
            sys.stdout.flush()
            if self.output == 'doc':
                doc = self[self.i]
            elif self.output == 'tag':
                doc = gm.doc2vec.TaggedDocument(self[self.i], [self.i])
            elif self.output == 'bow':
                doc = self.dct.doc2bow(self[self.i])
            self.i += 1
            return doc
        else:
            raise StopIteration

    def __getitem__(self, index):
        doc = self.dump.load_page(self.names[index])
        return gu.simple_preprocess(doc.strip_code())

class GraphContainer():
    """

    Attributes
    ----------
    graph: networkx.DiGraph
        node name is name of wikipedia page
        ``year`` attribute indicates year
    numbered: networkx.DiGraph
        node name is an index (see nodes),
        ``year`` is an index (see years), lazy
    nodes: list
        List of node names,
        indexed by node in ``numbered``, lazy
    years: list
        List of years,
        indexed by ``year`` in ``numbered``, lazy
    nodes_for_year: dict
        ``{int year: [int node_index]}``, lazy
    """
    def __init__(self):
        self.graph = nx.DiGraph()
        self._numbered = None
        self._nodes = []
        self._years = []
        self._nodes_for_year = {}

    @property
    def numbered(self):
        if self._numbered:
            return self._numbered
        else:
            self._numbered = nx.DiGraph()
            for node in self.graph.nodes:
                self._numbered.add_node(self.nodes.index(node),
                                        year = self.years.index(self.graph.nodes[node]['year']))
                self._numbered.add_edges_from([(self.nodes.index(node), self.nodes.index(succ))
                                               for succ in self.graph.successors(node)])
            return self._numbered

    @property
    def nodes(self):
        if self._nodes:
            return self._nodes
        else:
            self._nodes = list(self.graph.nodes)
            return self._nodes

    @property
    def years(self):
        if self._years:
            return self._years
        else:
            self._years = [self.graph.nodes[n]['year']
                           for n in self.graph.nodes()]
            self._years = sorted(list(set(self._years)))
            return self._years

    @property
    def nodes_for_year(self):
        if self._nodes_for_year:
            return self._nodes_for_year
        else:
            self._nodes_for_year = {year: [list(self.graph.nodes).index(n)
                                           for n in self.graph.nodes
                                           if self.graph.nodes[n]['year']==year]
                                    for year in self.years}
            return self._nodes_for_year

class PersistentHomology(GraphContainer):
    """

    Attributes
    ----------
    cliques: list of lists
        lazy
    filtration: dionysus.filtration
        lazy
    persistence: dionysus.reduced_matrix
        lazy
    barcodes: pandas.DataFrame
        lazy
    """

    def __init__(self):
        GraphContainer.__init__(self)
        self._cliques = None
        self._filtration = None
        self._persistence = None
        self._barcodes = None

    @property
    def cliques(self):
        if self._cliques:
            return self._cliques
        else:
            self._cliques = list(nx.algorithms.clique.\
                                 enumerate_all_cliques(nx.Graph(self.numbered)))
            return self._cliques

    @property
    def filtration(self):
        if self._filtration != None:
            return self._filtration
        else:
            self._filtration = d.Filtration()
            nodes_so_far = []
            for year in self.years:
                nodes_now = self.nodes_for_year[year]
                nodes_so_far.extend(nodes_now)
                for clique in self.cliques:
                    if all([n in nodes_so_far for n in clique]):
                        self._filtration.append(d.Simplex(clique, year))
            self._filtration.sort()
            return self._filtration

    @property
    def persistence(self):
        if self._persistence:
            return self._persistence
        else:
            self._persistence = d.homology_persistence(self.filtration)
            return self._persistence

    @property
    def barcodes(self):
        if isinstance(self._barcodes, pd.DataFrame)\
            and len(self._barcodes.index) != 0:
            return self._barcodes
        else:
            self._barcodes = PersistentHomology.compute_barcodes(
                self.filtration, self.persistence, self.graph, self.nodes)
            return self._barcodes

    @staticmethod
    def compute_barcodes(f, m, graph, names):
        """Uses dionysus filtration & persistence
        (in reduced matrix form) to compute barcodes.

        Parameters
        ----------
        f: dionysus.Filtration
        m: dionysus.ReducedMatrix
            (see homology_persistence)
        names: list of strings
            names of node indices

        Returns
        -------
        barcodes: pandas.DataFrame
        """
        print('wiki.Net: computing barcodes... (skip negatives)')
        node_list = list(graph.nodes)
        barcodes = []
        for i, c in enumerate(m):
            if m.pair(i) < i: continue
            sys.stdout.write("\rwiki.Net: barcode {}/{}".\
                             format(i+1,len(m)))
            sys.stdout.flush()
            dim = f[i].dimension()
            birth_year = int(f[i].data)
            birth_simplex = [names[s] for s in f[i]]
            birth_nodes = [n for n in birth_simplex
                           if graph.nodes[n]['year']==birth_year]
            if m.pair(i) != m.unpaired:
                death_year = int(f[m.pair(i)].data)
                death_simplex = [names[s] for s in f[m.pair(i)]]
                death_nodes = [n for n in death_simplex
                               if graph.nodes[n]['year']==death_year]
            else:
                death_year = np.inf
                death_simplex = []
                death_nodes = []
            pair = m.pair(i) if m.pair(i) != m.unpaired else np.inf
            chain = m[pair] if pair != np.inf else m[i]
            simp_comp = [f[entry.index] for entry in chain]
            nodes = [node_list[idx] for simplex in simp_comp for idx in simplex]
            barcodes.append([dim, birth_year, death_year, death_year-birth_year,
                             birth_simplex, death_simplex,
                             birth_nodes, death_nodes, list(set(nodes))])
        print('')
        barcodes.sort(key=lambda x: x[0])
        bar_data = pd.DataFrame(data=barcodes,
                                columns=['dim', 'birth', 'death', 'lifetime',
                                         'birth simplex', 'death simplex',
                                         'birth nodes', 'death nodes',
                                         'homology nodes'])
        return bar_data

class Net(PersistentHomology):
    """``Net`` is a wrapper for ``networkx.DiGraph``.
    Uses ``dionysus`` for persistence homology.

    Attributes
    ----------
    tfidf: scipy.sparse.csc.csc_matrix
        sparse column matrix of tfidfs,
        ordered by nodes, also stored in
        ```self.graph.graph['tfidf']```, lazy
    MAX_YEAR: int
        ``year = MAX_YEAR (2020)`` for nodes with parents
        without years
    YEAR_FILLED_DELTA: int
        ``year = year of parents + YEAR_FILLED_DELTA (1)``
    """
    MAX_YEAR = 2020
    YEAR_FILLED_DELTA = 1

    def __init__(self, path_graph='', path_barcodes=''):
        PersistentHomology.__init__(self)
        self._tfidf = None
        if path_graph:
            self.load_graph(path_graph)
        if path_barcodes:
            self.load_barcodes(path_barcodes)

    @property
    def tfidf(self):
        if self._tfidf:
            return self._tfidf
        elif 'tfidf' in self.graph.graph.keys():
            self._tfidf = self.graph.graph['tfidf']
            return self._tfidf

    def build_graph(self, name='', dump=None, nodes=None, depth_goal=1,
                    filter_top=True, remove_isolates=True, add_years=True,
                    fill_empty_years=True, model=None, dct=None,
                    compute_core_periphery=True, compute_communities=True,
                    compute_community_cores=True):
        """ Builds ``self.graph`` (``networkx.Graph``) from nodes (``list``
        of ``string``). Set ``model`` (from ``gensim``) and ``dct``
        (``gensim.corpora.Dictionary``) for weighted edges.
        """
        self.graph = nx.DiGraph()
        self.graph.name = name
        if not dump:
            raise AttributeError('wiki.Net: Provide wiki.Dump object.')
        print('wiki.Net: traversing Wikipedia...')
        Net.bft(self.graph, dump, nodes, depth_goal=depth_goal,
                nodes=nodes, filter_top=filter_top)
        if remove_isolates:
            print('wiki.Net: removing isolates...')
            self.graph.remove_nodes_from(nx.isolates(self.graph))
        if add_years:
            print('wiki.Net: adding years...')
            for node in self.graph.nodes:
                dump.load_page(node)
                self.graph.nodes[node]['year'] = dump.years[0] if len(dump.years)>0 else []
            self.graph.graph['num_years'] = sum(
                [bool(y) for y in nx.get_node_attributes(self.graph, 'year').values()]
            )
        if fill_empty_years:
            print('wiki.Net: filling empty years...')
            nodes_filled = True
            while nodes_filled:
                nodes_filled = Net.fill_empty_nodes(self.graph, full_parents=True)
            nodes_filled = True
            while nodes_filled:
                nodes_filled = Net.fill_empty_nodes(self.graph, full_parents=False)
            for node in self.graph.nodes:
                if not self.graph.nodes[node]['year']:
                    self.graph.nodes[node]['year'] = Net.MAX_YEAR
        if model and dct:
            print('wiki.Net: calculating weights...')
            self.graph.graph['tfidf'] = Net.compute_tfidf(self.nodes, dump, model, dct)
            Net.set_weights(self.graph)
        if compute_core_periphery:
            print('wiki.Net: computing core-periphery...')
            Net.assign_core_periphery(self.graph)
        if compute_communities:
            print('wiki.Net: computing communities...')
            Net.assign_communities(self.graph)
        if compute_community_cores:
            print('wiki.Net: computing cores within communities...')
            Net.assign_cores_to_communities(self.graph)

    def load_graph(self, path):
        """Loads ``graph`` from ``path``.
        If ``filename.gexf`` then read as ``gexf``.
        Else, use ``pickle``."""
        if path.split('.')[-1]=='gexf':
            self.graph = nx.read_gexf(path)
        else:
            self.graph = nx.read_gpickle(path)

    def save_graph(self, path):
        """Saves ``graph`` at ``path``.
        If ``filename.gexf`` then save as ``gexf``.
        Else, use ``pickle``."""
        if path.split('.')[-1]=='gexf':
            nx.write_gexf(self.graph, path)
        else:
            nx.write_gpickle(self.graph, path)

    def load_barcodes(self, path):
        """Loads ``barcodes`` from ``pickle``."""
        self._barcodes = pickle.load(open(path, 'rb'))

    def save_barcodes(self, path):
        """Saves ``barcodes`` as ``pickle``."""
        pickle.dump(self.barcodes, open(path, 'wb'))

    def randomize(self, null_type,
                  compute_core_periphery=True, compute_communities=True,
                  compute_community_cores=True):
        """Returns a new ``wiki.Net`` with a randomized
        copy of ``graph``. Set ``null_type`` as one of
        ``'year'``, ``'target'``.
        """
        network = Net()
        network.graph = self.graph.copy()
        if null_type == 'year':
            years = list(nx.get_node_attributes(network.graph, 'year').values())
            random.shuffle(years)
            for node in network.graph.nodes:
                network.graph.nodes[node]['year'] = years.pop()
        elif null_type == 'target':
            nodes = list(network.graph.nodes)
            for s, t in self.graph.edges:
                network.graph.remove_edge(s, t)
                nodes.remove(t)
                network.graph.add_edge(s, random.choice(nodes),
                                       weight=self.graph[s][t]['weight'])
                nodes.append(t)
        elif null_type == 'source':
            nodes = list(network.graph.nodes)
            for s, t in self.graph.edges:
                network.graph.remove_edge(s, t)
                nodes.remove(s)
                network.graph.add_edge(random.choice(nodes), t,
                                       weight=self.graph[s][t]['weight'])
                nodes.append(s)
        if compute_core_periphery:
            print('wiki.Net: computing core-periphery...')
            Net.assign_core_periphery(network.graph)
        if compute_communities:
            print('wiki.Net: computing communities...')
            Net.assign_communities(network.graph)
        if compute_community_cores:
            print('wiki.Net: computing cores within communities...')
            Net.assign_cores_to_communities(self.graph)
        return network

    @staticmethod
    def fill_empty_nodes(graph, full_parents=True):
        """ Fills nodes without ``year`` with the ``year`` of parents

        Parameters
        ----------
        graph: networkx.DiGraph
        full_parents: bool
            whether to fill empty nodes that have
            all parents with non-empty 'year'
        Returns
        -------
        bool
            whether at least 1 empty node was filled
        """
        empty_nodes = [n for n in graph.nodes if not graph.nodes[n]['year']]
        for node in empty_nodes:
            years = [graph.nodes[p]['year'] for p in graph.predecessors(node)]
            if not years:
                continue
            if full_parents:
                if [] not in years:
                    graph.nodes[node]['year'] = max(years) \
                                                + Net.YEAR_FILLED_DELTA
                    return True
            else:
                years_filtered = [y for y in years if y]
                if years_filtered:
                    graph.nodes[node]['year'] = max(years_filtered) \
                                                + Net.YEAR_FILLED_DELTA
                    return True
        return False

    @staticmethod
    def bft(graph, dump, queue, depth_goal=1, nodes=None, filter_top=True):
        """Breadth-first traversal of hyperlink graph.

        Parameters
        ----------
        graph: networkx.DiGraph
        dump: wiki.Dump
        queue: list of **strings**
            names of Wikipedia pages
        depth_goal: int
        nodes: list of **strings**
            names of Wikipedia pages
        filter_top: bool
        """
        queue = queue.copy()
        depth = 0
        depth_num_items = len(queue)
        depth_inc_pending = False
        print('wiki.Net: depth = ' + str(depth))
        while queue:
            name = queue.pop(0)
            sys.stdout.write("\rwiki.Net: len(queue) = " + str(len(queue)))
            sys.stdout.flush()
            depth_num_items -= 1
            if depth_num_items == 0:
                depth += 1
                print('\nwiki.Net: depth = ' + str(depth))
                depth_inc_pending = True
            if depth == depth_goal: break
            page = dump.load_page(name, filter_top=filter_top)
            if not page: continue
            links = [l for l in dump.article_links
                     if Net.filter(name, l, graph, nodes)]
            for link in links:
                graph.add_edge(link, name, weight=1)
                if link not in queue:
                    queue.append(link)
            if depth_inc_pending:
                depth_num_items = len(queue)
                depth_inc_pending = False

    @staticmethod
    def filter(page, link, graph, nodes=None):
        """Filter out links"""
        if nodes and link not in nodes:
            return False
        if (page, link) in graph.edges:
            return False
        return True

    @staticmethod
    def compute_tfidf(nodes, dump, model, dct):
        """Compute tf-idf of pages with titles in ``nodes``.

        Parameters
        ----------
        nodes: list of nodes
        dump: wiki.Dump
        model: gensim.modes.tfidfmodel.TfidfModel
        dct: gensim.corpora.Dictionary

        Returns
        -------
        vecs: scipy.sparse.csc.csc_matrix
        """
        pages = [dump.load_page(page) for page in nodes]
        bows = [model[dct.doc2bow(gu.simple_preprocess(page.strip_code()))]
                if page else []
                for page in pages]
        return gmat.corpus2csc(bows)

    @staticmethod
    def set_weights(graph):
        """Set the weights of graph (``networkx.DiGraph``) as
        the cosine similarity between ``graph.graph['tf-idf']]``
        vectors of nodes."""
        vecs = graph.graph['tfidf']
        for n1, n2 in graph.edges:
            v1 = vecs[:,list(graph.nodes).index(n1)].transpose()
            v2 = vecs[:,list(graph.nodes).index(n2)].transpose()
            graph[n1][n2]['weight'] = smp.cosine_similarity(X=v1, Y=v2)[0,0]

    @staticmethod
    def assign_core_periphery(graph):
        """ Compute core-periphery of ``graph`` (``nx.DiGraph``;
        converted to symmetric ``nx.Graph``).
        Assign ``core`` as ``1`` or ``0`` to each node.
        Assign ``coreness`` to ``graph``.
        See ``core_periphery_dir()`` in ``bctpy``.
        """
        # borgatti-everett
        be = bct.core_periphery_dir(nx.convert_matrix.to_numpy_array(graph))
        for i, node in enumerate(graph.nodes):
            graph.nodes[node]['core_be'] = be[0][i]
        graph.graph['coreness_be'] = be[1]
        # rombach
        rb = cpa.Rombach()
        rb.detect(graph)
        if rb.get_coreness() != 0:
            for node, coreness in rb.get_coreness().items():
                graph.nodes[node]['core_rb'] = coreness
            graph.graph['coreness_rb'] = rb.score()[0]
        else:
            for node in graph.nodes:
                graph.nodes[node]['core_rb'] = 0
            graph.graph['coreness_rb'] = 0

    @staticmethod
    def assign_communities(graph):
        """ Compute modular communities of ``graph`` (``nx.DiGraph``).
        Assign community number ``community`` to each node.
        Assign ``modularity`` to ``graph``.
        See ``greedy_modularity_communities`` in ``networkx``.
        """
        communities = nx.algorithms.community\
                        .greedy_modularity_communities(nx.Graph(graph))
        for node in graph.nodes:
            graph.nodes[node]['community'] = [i for i,c in enumerate(communities)
                                              if node in c][0]
        graph.graph['modularity'] = nx.algorithms.community.quality\
                                      .modularity(nx.Graph(graph),
                                                  communities)

    @staticmethod
    def assign_cores_to_communities(graph):
        """"""
        num_comm = max([graph.nodes[n]['community'] for n in graph.nodes])
        community_coreness_be = {i: 0 for i in range(num_comm)}
        community_coreness_rb = {i: 0 for i in range(num_comm)}
        for i in range(num_comm+1):
            community = [n for n in graph.nodes if graph.nodes[n]['community']==i]
            subgraph = graph.subgraph(community).copy()
            matrix = nx.convert_matrix.to_numpy_array(subgraph)
            if (matrix.size>1) & (np.sum(matrix)>0):
                # borgatti-everett
                be = bct.core_periphery_dir(matrix)
                # rombach
                rb = cpa.Rombach()
                rb.detect(subgraph)
                # assign
                community_coreness_be[i] = be[1]
                community_coreness_rb[i] = rb.score()[0]
                cp_rb = rb.get_coreness()
                for j, node in enumerate(subgraph.nodes):
                    graph.nodes[node]['community_core_be'] = be[0][j]
                    graph.nodes[node]['community_core_rb'] = cp_rb[node]
            else:
                community_coreness_be[i] = 0
                community_coreness_rb[i] = 0
                for j, node in enumerate(subgraph.nodes):
                    graph.nodes[node]['community_core_be'] = 1
                    graph.nodes[node]['community_core_rb'] = 1
        graph.graph['community_coreness_be'] = community_coreness_be
        graph.graph['community_coreness_rb'] = community_coreness_rb

class Model(PersistentHomology):
    """

    Attributes
    ----------
    graph: networkx.DiGraph
    graph_parent: networkx.DiGraph
    vectors: scipy.sparse.csc_matrix
    vectors_parent: scipy.sparse.csc_matrix
    seeds: {node string: [scipy.sparse.csc_matrix]}
    thresholds: {node string: [float]}
    year: int
    record: pandas.DataFrame
        record of evolution
    year_start: int
    n_seeds: int
        number of seeds per node
    point, insert, delete: tuple
        See ``mutate()``.
    rvs: lambda n->float
        random values for point mutations & insertions
    dct: gensim.corpora.dictionary
    create: lambda n-> float
        thresholds of cosine similarity with parent
        for node creation
    crossover: float
        threshold of cosine similarity with parent
        for crossing over nodes
    """

    def __init__(self, graph_parent, vectors_parent, year_start, start_nodes,
                 n_seeds, dct, point, insert, delete, rvs,
                 create, crossover=None):
        """

        Parameters
        ----------
        start_nodes: lambda wiki.Model -> list(networkx.Nodes)
        """
        PersistentHomology.__init__(self)
        self.graph_parent = graph_parent
        self.vectors_parent = vectors_parent
        self.year_start = year_start
        self.year = year_start
        self.seeds = {}
        self.thresholds = {}
        self.record = pd.DataFrame()
        nodes = list(graph_parent.nodes)
        self.start_nodes = start_nodes(self)
        self.graph = graph_parent.subgraph(self.start_nodes).copy()
        self.vectors = sp.sparse.hstack([vectors_parent[:,nodes.index(n)]
                                         for n in self.start_nodes])
        self.n_seeds = n_seeds
        self.dct = dct
        self.point = point
        self.insert = insert
        self.delete = delete
        self.rvs = rvs
        self.create = create
        self.crossover = crossover

    def __str__(self):
        return f"Model\tparent: '{self.graph_parent.name}'\n" +\
               f"\tyear_start: {self.year_start}\n" +\
               f"\tstart_nodes: {self.start_nodes}\n" +\
               f"\tn_seeds: {self.n_seeds}\n" +\
               f"\tpoint: ({self.point[0]:.4f}, {self.point[1]:.4f})\n" +\
               f"\tinsert: ({self.insert[0]}, {self.insert[1]:.4f}, {type(self.insert[2])})\n" +\
               f"\tdelete: ({self.delete[0]}, {self.delete[1]:.4f})"

    def __repr__(self):
        return self.__str__()

    def evolve(self, until, record=False):
        """ Evolves a graph based on vector representations
        until `until (lambda wiki.Model) == True`
        """
        year_start = self.year
        while not until(self):
            sys.stdout.write(f"\r{year_start} > {self.year} "+\
                             f"n={self.graph.number_of_nodes()}    ")
            sys.stdout.flush()
            self.initialize_seeds()
            self.mutate_seeds()
            self.create_nodes()
            if record:
                self.record = pd.concat([self.record] + \
                                        [pd.DataFrame({'Year': self.year,
                                                       'Parent': seed,
                                                       'Seed number': i,
                                                       'Seed vectors': seed_vec}, index=[0])
                                         for seed, seed_vecs in self.seeds.items()
                                         for i, seed_vec in enumerate(seed_vecs)],
                                        ignore_index=True,
                                        sort=False)
            self.year += 1
        print('')

    def initialize_seeds(self):
        nodes = list(self.graph.nodes)
        for i, node in enumerate(nodes):
            if node not in self.seeds.keys():
                self.seeds[node] = []
            if node not in self.thresholds.keys():
                self.thresholds[node] = []
            while len(self.seeds[node]) < self.n_seeds:
                self.seeds[node] += [self.vectors[:,i].copy()]
            while len(self.thresholds[node]) < self.n_seeds:
                self.thresholds[node] += [self.create(1)[0]]

    def mutate_seeds(self):
        for node, vecs in self.seeds.items():
            self.seeds[node] = [Model.mutate(vec, self.rvs, self.point,
                                             self.insert, self.delete)
                                for vec in vecs]

    def crossover_seeds(self):
        nodes = list(self.graph.nodes)
        for i in range(len(nodes)):
            seeds_i = sp.sparse.hstack(self.seeds[nodes[i]])
            for j in range(i+1,len(nodes)):
                seeds_j = sp.sparse.hstack(self.seeds[nodes[j]])
                similarity = smp.cosine_similarity(seeds_i.transpose(),
                                                   seeds_j.transpose())
                for k,l in np.argwhere(similarity>self.threshold):
                    cross = Model.crossover(seeds_i[:,k], seeds_j[:,l])
                    choice = np.random.choice(2)
                    self.seeds[nodes[i]][k] = cross if choice else self.vectors[:,i]
                    self.seeds[nodes[j]][l] = cross if not choice else self.vectors[:,j]

    def create_nodes(self):
        nodes = list(self.graph.nodes)
        for i, node in enumerate(nodes):
            parent = self.vectors[:,i]
            seeds = sp.sparse.hstack(self.seeds[node])
            sim_to_parent = smp.cosine_similarity(parent.transpose(), seeds.transpose())
            for j, seed_vec in enumerate(self.seeds[node]):
                if sim_to_parent[0,j] < self.thresholds[node][j]:
                    Model.connect(seed_vec, self.graph, self.vectors, self.dct)
                    self.vectors = sp.sparse.hstack([self.vectors, seed_vec])
                    self.seeds[node].pop(j)
                    self.thresholds[node].pop(j)
        for node in self.graph.nodes:
            if 'year' not in self.graph.nodes[node].keys():
                self.graph.nodes[node]['year'] = self.year

    @staticmethod
    def mutate(x, rvs, point=(0,0), insert=(0,0,None), delete=(0,0)):
        """ Mutates vector ``x`` with point mutations,
        insertions, and deletions. Insertions and point
        mutations draw from a random process ``rvs``.

        Parameters
        ----------
        x: spipy.sparse.csc_matrix
        rvs: lambda (n)-> float
            returns ``n`` random weights in [0,1]
        point: tuple (int n, float p)
            n = number of elements to insert
            p = probability of insertion for each trial
        insert: tuple (n, p, iterable s)
            s = set of elements from which to select
                if None, select from all zero elements
        delete: tuple (n, p)
        max_weight: float
        """
        data = x.data
        idx = x.indices
        if idx.size==0:
            return x
        n_point = np.random.binomial(point[0], point[1])
        i_point = np.random.choice(x.size, size=n_point, replace=False)
        data[i_point] = rvs(n_point)
        # insertion
        n_insert = np.random.binomial(insert[0], insert[1])
        for _ in range(n_insert):
            while True:
                insert_idx = np.random.choice(insert[2]) if insert[2]\
                    else np.random.choice(x.shape[0])
                if insert_idx not in idx: break
            idx = np.append(idx, insert_idx)
            data = np.append(data, rvs(1))
        # deletion
        n_delete = np.random.binomial(delete[0], delete[1])
        i_delete = np.random.choice(idx.size, size=n_delete, replace=False)
        idx = np.delete(idx, i_delete)
        data = np.delete(data, i_delete)
        y = sp.sparse.csc_matrix((data, (idx, np.zeros(idx.shape, dtype=int))),
                                 shape=x.shape)
        return y

    @staticmethod
    def connect(seed_vector, graph, vectors, dct, top_words=10, match_n=6):
        """

        Parameters
        ----------
        seed_vector: scipy.sparse.csc_matrix
        graph: networkx.DiGraph (not optional)
        vectors: scipy.sparse.csc_matrix (not optional)
        dct: gensim.corpora.dictionary (not optional)
        top_words: int (default=5)
        match_n: int
            how many words should be matched by...
        """
        seed_top_words, seed_top_idx = Model.find_top_words(seed_vector, dct)
        seed_name = ' '.join(seed_top_words)
        nodes = list(graph.nodes)
        graph.add_node(seed_name)
        for i, node in enumerate(nodes):
            node_vector = vectors[:,i]
            node_top_words, node_top_idx = Model.find_top_words(node_vector, dct)
            if len(set(seed_top_idx).intersection(set(node_vector.indices))) >= match_n or\
               len(set(node_top_idx).intersection(set(seed_vector.indices))) >= match_n:
                graph.add_edge(node, seed_name)

    @staticmethod
    def find_top_words(x, dct, top_n=10):
        """

        Parameters
        ----------
        x: scipy.sparse.csc_matrix
        dct: gensim.corpora.dictionary
        top_n: int

        Returns
        -------
        words:
        idx_vector:
        """
        top_idx = np.argsort(x.data)[-top_n:]
        idx = [x.indices[i] for i in top_idx]
        words = [dct[i] for i in idx]
        words_nostop = gpp.remove_stopwords(' '.join(words)).split(' ')
        idx_keep = list(map(lambda x: words.index(x), set(words).intersection(words_nostop)))
        idx_nostop = list(map(idx.__getitem__, idx_keep))
        return words_nostop, idx_nostop

    @staticmethod
    def crossover(v1, v2):
        """ Crosses two vectors by combining half of one
        and half of the other.

        Parameters
        ----------
        v1, v2: scipy.sparse.matrix

        Returns
        -------
        v3: scipy.sparse.matrix
        """
        idx1 = np.random.choice(v1.size, size=int(v1.size/2))
        idx2 = np.random.choice(v2.size, size=int(v2.size/2))
        data = np.array([v1.data[i] for i in idx1] +
                        [v2.data[i] for i in idx2])
        idx = np.array([v1.indices[i] for i in idx1] +
                       [v2.indices[i] for i in idx2])
        v3 = sp.sparse.csc_matrix((data, (idx, np.zeros(idx.shape, dtype=int))),
                                  shape=v1.shape)
        return v3

import numpy,import scipy,import networkx


import logging
import numpy as np
import pandas as pd
import os
import video
import settings

class ChunkLoader():
    def __init__(self, set_name, repo_dir, datum_dtype=np.uint8,
                 nclasses=2, augment=False, test_mode=False):
        # assert test_mode is False, 'Test mode not implemented yet'
        np.random.seed(0)
        self.set_name = set_name
        # self.bsz = self.be.bsz
        self.augment = augment
        self.repo_dir = repo_dir
        self.is_training = (set_name == 'train')
        self.chunk_size = settings.chunk_size
        self.chunk_shape = (self.chunk_size, self.chunk_size, self.chunk_size)
        self.chunk_volume = np.prod(self.chunk_shape)
        self.metadata = pd.read_csv(os.path.join(self.repo_dir, set_name + '-metadata.csv'))
        self.data_size = self.metadata.shape[0]
        self.pos_users = self.metadata[self.metadata['flag']==1]['uid']
        self.neg_users = self.metadata[self.metadata['flag']==0]['uid']
        self.nvids = self.metadata.shape[0]
        self.chunks_filled = 0
        self.video_idx = 0
        if not test_mode:
            self.labels = pd.read_csv(os.path.join(self.repo_dir, 'labels.csv'))
            self.nega_labels = pd.read_csv(os.path.join(self.repo_dir, 'candidates.csv'))
        else:
            self.labels = None
        self.test_mode = test_mode
        self.chunks,self.starts,self.targets = [],[],[]
        ##positive points in lables.csv
        self.pos_labels = self.labels[self.labels['uid'].isin(self.pos_users)].shape[0]
        self.pos_neg_ratio = 6.0
        self.chunk_from_neg_users = int(self.pos_labels*self.pos_neg_ratio/len(self.neg_users))
        self.current_uid = self.current_flag = self.current_meta = None

    def reset(self):
        self.chunks,self.starts,self.targets = [],[],[]
    def next_video(self,video_idx):
        self.reset()
        self.current_meta = self.metadata.iloc[video_idx]
        uid = self.current_meta['uid']
        self.current_uid = self.current_meta['uid']
        self.current_flag = int(self.current_meta['flag'])
        data_filename = os.path.join(self.repo_dir, uid + '.' + settings.file_ext)
        vid_shape = (int(self.current_meta['z_len']),
                     int(self.current_meta['y_len']),
                     int(self.current_meta['x_len']))
        vid_data = video.read_blp(data_filename, vid_shape)
        self.video_idx += 1
        self.extract_chunks(vid_data)

        return self.chunks,self.starts,self.targets

    def slice_chunk(self, start, data):
        return data[start[0]:start[0] + self.chunk_size,
               start[1]:start[1] + self.chunk_size,
               start[2]:start[2] + self.chunk_size]#.ravel()

    def extract_one(self, data, data_shape, uid_data,idx):
        # assert uid_data.shape[0] != 0
        if not self.test_mode:
                center = np.array((uid_data['z'].iloc[idx],
                                   uid_data['y'].iloc[idx],
                                   uid_data['x'].iloc[idx]), dtype=np.int32)
                # radius
                rad = 0.5 * uid_data['diam'].iloc[idx]
                if rad == 0:
                    # Assign an arbitrary radius to candidate nodules
                    rad = 20 / settings.resolution

                #comment by lc: low may <0
                low = np.int32(center + rad - self.chunk_size)
                high = np.int32(center - rad)
                for j in range(3):
                    low[j] = max(0, low[j])
                    high[j] = max(low[j] + 1, high[j])
                    high[j] = min(data_shape[j] - self.chunk_size, high[j])
                    low[j] = min(low[j], high[j] - 1)
                start = [np.random.randint(low=low[i], high=high[i]) for i in range(3)]
        else:
            start = self.generate_chunk_start(chunk_idx, data_shape)

        chunk = self.slice_chunk(start, data)

        return chunk,start

    def generate_chunk_start(self, chunk_idx, data_shape):
        chunk_spacing = np.int32((np.array(data_shape) - self.chunk_size) / settings.chunks_per_dim)
        z_chunk_idx = chunk_idx / settings.chunks_per_dim ** 2
        y_chunk_idx = (chunk_idx - z_chunk_idx * settings.chunks_per_dim ** 2) / settings.chunks_per_dim
        x_chunk_idx = chunk_idx - z_chunk_idx * settings.chunks_per_dim ** 2 \
                      - y_chunk_idx * settings.chunks_per_dim

        start = [z_chunk_idx * chunk_spacing[0],
                 y_chunk_idx * chunk_spacing[1],
                 x_chunk_idx * chunk_spacing[2]]
        return start

    def extract_chunks(self, data):
        data_shape = np.array(data.shape, dtype=np.int32)
        if self.current_flag:
            uid_data = self.labels[self.labels['uid'] == self.current_uid]
            for idx in range(uid_data.shape[0]):
                chunk,start = self.extract_one(data, data_shape, uid_data, idx)
                if chunk is None:
                    continue
                self.chunks.append(chunk)
                self.starts.append(start)
                self.targets.append(1)
        else:
            uid_data = self.labels[self.labels['uid'] == self.current_uid]
            for idx in range(min(self.chunk_from_neg_users,uid_data.shape[0])):
                chunk,start = self.extract_one(data, data_shape, uid_data, idx)
                if chunk is None:
                    continue
                self.chunks.append(chunk)
                self.starts.append(start)
                self.targets.append(0)
            ## not enough negative from labels, then to candidates
            if uid_data.shape[0]<self.chunk_from_neg_users:
                left_chunk = self.chunk_from_neg_users-uid_data.shape[0]
                uid_data = self.nega_labels[self.nega_labels['uid'] == self.current_uid]
                for i in range(min(left_chunk,uid_data.shape[0])):
                    idx = np.random.randint(uid_data.shape[0])
                    chunk,start = self.extract_one(data, data_shape, uid_data, idx)
                    if chunk is None:
                        continue
                    self.chunks.append(chunk)
                    self.starts.append(start)
                    self.targets.append(0)


# -*- coding: utf-8 -*-
""" main """

import numpy as np

import settings
from spfs.core.feature_extractors import SpatialPyramidFeatures
from spfs.utils.datasets.handlers import InMemoryDBHandler, LazyDBHandler
from spfs.utils.utils import FeaturesEvaluator

np.random.seed(settings.NUMPY_RANDOM_SEED)

def main():
    pass

if __name__ == '__main__':
    main()

    # spf = SpatialPyramidFeatures(InMemoryDBHandler)
    spf = SpatialPyramidFeatures(LazyDBHandler)
    # func:create_codebook processed in 5602.9401 seconds
    # func:create_codebook processed in 5549.5434 seconds
    # func:create_codebook processed in 17794.9059 seconds
    spf.create_codebook()
    # old time: func:create_spatial_pyramid_features processed in 93547.0923 seconds
    # new time: func:create_spatial_pyramid_features processed in  7969.3614 seconds
    spf.create_spatial_pyramid_features()

    ###########################################################################
    #                        Evaluation with linear SVC                       #
    ###########################################################################
    FeaturesEvaluator.apply_linear_svc()


#!/usr/bin/env python
# Copyright 2021 Calico LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =========================================================================

from __future__ import print_function
from optparse import OptionParser
import json
import os
import pdb
import sys
import time

import h5py
import numpy as np
import pandas as pd
import tensorflow as tf

from basenji import bed
from basenji import dataset
from basenji import seqnn
from basenji import trainer

"""
basenji_predict_tfr.py

Make predictions on TFRecord sequences.
"""

################################################################################
# main
################################################################################
def main():
  usage = 'usage: %prog [options] <params_file> <model_file> <data_dir>'
  parser = OptionParser(usage)
  parser.add_option('--bi', dest='bedgraph_indexes',
      help='Comma-separated list of target indexes to write predictions and targets as bedgraph [Default: %default]')
  parser.add_option('--head', dest='head_i',
      default=0, type='int',
      help='Parameters head [Default: %default]')
  parser.add_option('-o', dest='out_dir',
      default='preds_out',
      help='Output directory for predictions [Default: %default]')
  parser.add_option('--rc', dest='rc',
      default=False, action='store_true',
      help='Average the fwd and rc predictions [Default: %default]')
  parser.add_option('--shifts', dest='shifts',
      default='0',
      help='Ensemble prediction shifts [Default: %default]')
  parser.add_option('-t', dest='targets_file',
      default=None, type='str',
      help='File specifying target indexes and labels in table format')
  parser.add_option('--split', dest='split_label',
      default='test',
      help='Dataset split label for eg TFR pattern [Default: %default]')
  parser.add_option('--tfr', dest='tfr_pattern',
      default=None,
      help='TFR pattern string appended to data_dir/tfrecords for subsetting [Default: %default]')
  (options, args) = parser.parse_args()

  if len(args) != 3:
    parser.error('Must provide parameters, model, and test data directory')
  else:
    params_file = args[0]
    model_file = args[1]
    data_dir = args[2]

  if not os.path.isdir(options.out_dir):
    os.mkdir(options.out_dir)

  # parse shifts to integers
  options.shifts = [int(shift) for shift in options.shifts.split(',')]

  #######################################################
  # inputs

  # read targets
  if options.targets_file is None:
    options.targets_file = '%s/targets.txt' % data_dir
  targets_df = pd.read_csv(options.targets_file, index_col=0, sep='\t')

  # read model parameters
  with open(params_file) as params_open:
    params = json.load(params_open)
  params_model = params['model']
  params_train = params['train']

  # construct eval data
  eval_data = dataset.SeqDataset(data_dir,
    split_label=options.split_label,
    batch_size=params_train['batch_size'],
    mode='eval',
    tfr_pattern=options.tfr_pattern)

  # initialize model
  seqnn_model = seqnn.SeqNN(params_model)
  seqnn_model.restore(model_file, options.head_i)
  seqnn_model.build_ensemble(options.rc, options.shifts)

  #######################################################
  # predict

  # compute predictions
  test_preds = seqnn_model.predict(eval_data).astype('float16')

  # read targets
  test_targets = eval_data.numpy(return_inputs=False)

  # save HDF5
  preds_h5 = h5py.File('%s/preds.h5' % options.out_dir, 'w')
  preds_h5.create_dataset('preds', data=test_preds)
  preds_h5.close()
  targets_h5 = h5py.File('%s/targets.h5' % options.out_dir, 'w')
  targets_h5.create_dataset('targets', data=test_targets)
  targets_h5.close()

  if options.bedgraph_indexes is not None:
    bedgraph_indexes = [int(ti) for ti in options.bedgraph_indexes.split(',')]
    bed.write_bedgraph(test_preds, test_targets, data_dir,
      options.out_dir, options.split_label, bedgraph_indexes)

################################################################################
# __main__
################################################################################
if __name__ == '__main__':
  main()


import math
import os
import random
from collections import deque, defaultdict
import pickle as pkl

import gym
import pathlib
import numpy as np
import re
from omegaconf import OmegaConf
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import distributions as pyd

class eval_mode(object):
    def __init__(self, *models):
        self.models = models

    def __enter__(self):
        self.prev_states = []
        for model in self.models:
            self.prev_states.append(model.training)
            model.train(False)

    def __exit__(self, *args):
        for model, state in zip(self.models, self.prev_states):
            model.train(state)
        return False

def soft_update_params(net, target_net, tau):
    for param, target_param in zip(net.parameters(), target_net.parameters()):
        target_param.data.copy_(tau * param.data +
                                (1 - tau) * target_param.data)

def set_seed_everywhere(seed):
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)

def make_dir(*path_parts):
    dir_path = os.path.join(*path_parts)
    try:
        os.mkdir(dir_path)
    except OSError:
        pass
    return dir_path

def tie_weights(src, trg):
    assert type(src) == type(trg)
    trg.weight = src.weight
    trg.bias = src.bias

def chain(*iterables):
    for it in iterables:
        yield from it

def save(obj, file_path):
    with open(file_path, 'wb') as f:
        torch.save(obj, f)

def load(file_path):
    if isinstance(file_path, str):
        file_path = pathlib.Path(file_path).expanduser()
    with file_path.open('rb') as f:
        return torch.load(f)

def weight_init(m):
    """Custom weight init for Conv2D and Linear layers."""
    if isinstance(m, nn.Linear):
        nn.init.orthogonal_(m.weight.data)
        if hasattr(m.bias, 'data'):
            m.bias.data.fill_(0.0)
    elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
        gain = nn.init.calculate_gain('relu')
        nn.init.orthogonal_(m.weight.data, gain)
        if hasattr(m.bias, 'data'):
            m.bias.data.fill_(0.0)

def mlp(input_dim,
        hidden_dim,
        output_dim,
        hidden_depth,
        output_mod=None,
        use_ln=False):
    if hidden_depth == 0:
        mods = [nn.Linear(input_dim, output_dim)]
    else:
        mods = [nn.Linear(input_dim, hidden_dim)]
        if use_ln:
            mods += [nn.LayerNorm(hidden_dim), nn.Tanh()]
        else:
            mods += [nn.ReLU(inplace=True)]
        for i in range(hidden_depth - 1):
            mods += [nn.Linear(hidden_dim, hidden_dim), nn.ReLU(inplace=True)]
        mods.append(nn.Linear(hidden_dim, output_dim))
    if output_mod is not None:
        mods.append(output_mod)
    trunk = nn.Sequential(*mods)
    return trunk

def to_np(t):
    if t is None:
        return None
    elif t.nelement() == 0:
        return np.array([])
    else:
        return t.cpu().detach().numpy()

def parse_run_overrides(exp_dirs):
    exp_dirs = exp_dirs.split(':')

    def parse_cfg(cfg_list):
        cfg = {}
        for item in cfg_list:
            parts = item.split('=')
            cfg[parts[0]] = parts[1]
        return cfg

    runs = {}
    for exp_dir in exp_dirs:
        exp_dir = pathlib.Path(exp_dir).expanduser()
        for override in exp_dir.glob('**/overrides.yaml'):
            with override.open('rb') as f:
                cfg = parse_cfg(OmegaConf.load(f))
            path = override.parents[1]
            runs[path] = cfg
    return runs

def find_available_seeds(runs, env):
    avail_seeds = {}
    for path, cfg in runs.items():
        if cfg['env'] == env:
            snapshots = {}
            model_dir = path / 'model'
            for snap in model_dir.glob('expl_agent_*.pt'):
                snap_id = int(
                    re.match(r'expl_agent_(\d+).pt', snap.name).group(1))
                snapshots[snap_id] = snap
            avail_seeds[int(cfg['seed'])] = snapshots
    return avail_seeds

def find_pretrained_agent(exp_dirs, env, seed, step):
    runs = parse_run_overrides(exp_dirs)
    avail_seeds = find_available_seeds(runs, env)
    if len(avail_seeds) == 0:
        raise f'cannot find a pretrained agent for {env} {seed}'

    if seed in avail_seeds and step in avail_seeds[seed]:
        return avail_seeds[seed][step]

    for snapshots in avail_seeds.values():
        if step in snapshots:
            return snapshots[step]

    raise f'cannot find a pretrained agent for {env} {seed}'
    return None

class TanhTransform(pyd.transforms.Transform):
    domain = pyd.constraints.real
    codomain = pyd.constraints.interval(-1.0, 1.0)
    bijective = True
    sign = +1

    def __init__(self, cache_size=1):
        super().__init__(cache_size=cache_size)

    @staticmethod
    def atanh(x):
        return 0.5 * (x.log1p() - (-x).log1p())

    def __eq__(self, other):
        return isinstance(other, TanhTransform)

    def _call(self, x):
        return x.tanh()

    def _inverse(self, y):
        # We do not clamp to the boundary here as it may degrade the performance of certain algorithms.
        # one should use `cache_size=1` instead
        return self.atanh(y)

    def log_abs_det_jacobian(self, x, y):
        # We use a formula that is more numerically stable, see details in the following link
        # https://github.com/tensorflow/probability/commit/ef6bb176e0ebd1cf6e25c6b5cecdd2428c22963f#diff-e120f70e92e6741bca649f04fcd907b7
        return 2. * (math.log(2.) - x - F.softplus(-2. * x))

class SquashedNormal(pyd.transformed_distribution.TransformedDistribution):
    def __init__(self, loc, scale):
        self.loc = loc
        self.scale = scale

        self.base_dist = pyd.Normal(loc, scale)
        transforms = [TanhTransform()]
        super().__init__(self.base_dist, transforms)

    @property
    def mean(self):
        mu = self.loc
        for tr in self.transforms:
            mu = tr(mu)
        return mu

class ClippedNormal(pyd.Normal):
    def __init__(self, loc, scale):
        super().__init__(loc, scale)

    def sample(self, sample_shape=torch.Size()):
        x = super().sample(sample_shape)
        return torch.clamp(x, -1.0, 1.0)

    def rsample(self, sample_shape=torch.Size()):
        x = super().rsample(sample_shape)
        return torch.clamp(x, -1.0, 1.0)


import pandas as pd
import numpy
from sklearn.svm import SVR
from sklearn import svm
from sklearn.externals import joblib
import time
import sys
import threading

def predictions():
    print("work")
    # millis = float(round(time.time() * 1000))
    # print(millis)
    clf = joblib.load('Final.pkl')
    # print("work")
    preictioMade = clf.predict([[16., 12., 22., 12., 24.]])
    print("work")
    # add the time for the prediction model so the script will run

    print(str(preictioMade))
    # threading.Timer(1.0, predictions).start()
# 1.513444888000000000e+09
# def retrainers():

if __name__ == '__main__': predictions()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Retinotopic mapping: polar angle

Author: Michele Svanera
Nov. 2017
"""

################################################################################################################
## Imports

from __future__ import division

from psychopy import visual, event, core, logging
#from psychopy import gui  #fetch default gui handler (qt if available)
from psychopy import prefs as pyschopy_prefs

#from itertools import cycle  # import cycle if you want to flicker stimuli
#from psychopy.misc import pol2cart
import copy
import numpy as np
import os
#import matplotlib.pyplot as plt
from datetime import datetime as dt
from time import gmtime, strftime
import threading

################################################################################################################
## Paths and Constants

EPSILON = 1e-5
serif = ['Times', 'Times New Roman']
sans = ['Gill Sans MT', 'Arial', 'Helvetica', 'Verdana']

# Button box
Initial_code = 16777200
Button_coding = [-1,-2,-3,-4,-5]            #red,yellow,green,blue,white(i.e. trigger)
IsMRI = 0                                   #MRI inverses the bit polarities

# Experiment details
Fullscreen = True
Dir_save = 'D:/Mucklis_lab/Retinotopic_mapping/out/'
Log_name = 'LogFile.log'
Frames_durations_name = 'frames_durations.npy'
Fps_update_rate = 1                             # sec
scanner_message = "Waiting for the scanner..."

# Polare angle, i.e. moving beam
Flash_period = 0.2                              # seconds for one B-W cycle (ie 1/Hz)
Cycles_number = 12                               # How many cicles
Rotation_rate = 1./64.                          # revs per sec
#Rotation_rate = 1./54.4                         #changed for Yee-Lee's experiment - change back!
Wedge_width = 22.5                              # degrees
Initial_wedge_pos = 90. - Wedge_width/4.

# Fixation
Rotating_cross = False
Color_change_cross = False
Color_change_rate = 2                           # sec
Rotation_cross_rate = 0.05                      # revs per sec

# Spyder network
Spyder_grid = True
Spyder_rings = 4
Web_size = (1.,1.)

# External cover
External_ring_size = 2.5

Pre_post_stimuli_fixation_time = 12.            # sec.
#Pre_post_stimuli_fixation_time = 13.6            # #changed for Yee-Lee's experiment - change back!
Total_time = Cycles_number * (1/Rotation_rate)

################################################################################################################
## Function

def screenCorrection(mywin,x):
    resX = mywin.size[0]
    resY = mywin.size[1]
    aspect = float(resX) / float(resY)

    if(Fullscreen == True):
        new = x / aspect
    else:
        new = x
    return(new)

def createOutFolder(path_out):
    if not os.path.exists(path_out):
        try:
            os.makedirs(path_out)
        except Exception as e:
            print('Problem with creating an *out* folder, check permissions: '+e)
    else:
        n_folder_name = 1
        path_out_new = path_out + "_" + str(n_folder_name)
        while(os.path.exists(path_out_new) is True):
            n_folder_name += 1
            path_out_new = path_out + "_" + str(n_folder_name)
        try:
            os.makedirs(path_out_new)
        except Exception as e:
            print('Problem with creating an *out* folder, check permissions: '+e)
        path_out = path_out_new
    path_out += '/'

    return path_out

class buttonBoxThread(threading.Thread):
    def __init__(self, thread_id, name):
        threading.Thread.__init__(self)

        self.stimulus_onset_time = 0
        self.local_status = 0
        self.mask = 0x0000 if IsMRI == 0  else 0x001f

        # Open comunication with Vpixx
        DPxOpen()

        #Select the correct device
        DPxSelectDevice('PROPixxCtrl')
        DPxSetMarker()
        self.stimulus_onset_time = DPxGetMarker()
        DPxEnableDinDebounce()              # Filter out button bounce
        self.local_status = DPxSetDinLog()        # Configure logging with default values
        DPxStartDinLog()
        DPxUpdateRegCache()

        # Thread infos
        self.thread_id = thread_id
        self.name = name

        # Every button has: its value {0,1} and time when it changed (as datetime)
        self.button_state = {'time': np.array([dt.now()]*5), 'state': np.zeros((5,),dtype=np.int8)}
        self.button_state['state'][-1] = 1
        self._stop_event = threading.Event()

    def run(self):
        logging.data("Starting " + self.name)

        initial_values = DPxGetDinValue()
        logging.data('Initial button box digital input states = ' + "{0:016b}".format(initial_values) + " (int=%d)".format(initial_values))

        while(True):
            DPxUpdateRegCache()
            DPxGetDinStatus(self.local_status)

            if self.local_status['newLogFrames'] > 0 :           #Something happened
                data_list = DPxReadDinLog(self.local_status)
                self.button_state = self.updateStateButton(self.button_state,data_list)

            if(self.stopped()):
                break

        logging.data("Exiting " + self.name)

    def stop(self):
        self._stop_event.set()

    def stopped(self):
        return self._stop_event.is_set()

    # Take in input a state and check what is changed
    def updateStateButton(self,button_state,data_list):

        for j in range(len(data_list)):
            j_value = bin(data_list[j][1] ^ self.mask)

            # Check every button
            for i_button in range(len(button_state['state'])):
                i_button_detected_value = int(j_value[Button_coding[i_button]])
                if(i_button_detected_value != button_state['state'][i_button]):
                    button_state['state'][i_button] = i_button_detected_value
                    button_state['time'][i_button] = dt.now()

        return button_state

################################################################################################################
## Main

def main(win,globalClock):

    all_changes = []

    ################################ Stimuli prepation ################################
    globalClock = core.Clock()

    # Make two wedges (in opposite contrast) and alternate them for flashing
    grating_texture = np.tile([[1,-1],[-1,1]], (8,8))
    wedge1 = visual.RadialStim(win, tex=grating_texture, color=1, units='pix', size=win.size[1]*1.3,
                               visibleWedge=[Initial_wedge_pos, Initial_wedge_pos+Wedge_width], interpolate=False, autoLog=False,
                               radialCycles=1, angularCycles=4)
    wedge2 = copy.copy(wedge1)
    wedge2.color = -1

    # fixation cross
    fixation = visual.ShapeStim(win,vertices=((0,-1),(0,1),(0,0),(-1,0),(1,0)),lineWidth=4, units="pix",
        size=(20,20),closeShape=False,lineColor='red',autoDraw=True)

    # Spyder network
    web_circle = visual.Circle(win=win,radius=1,edges=200,units='norm',pos=[0, 0],lineWidth=1,opacity=1,interpolate=True,
                            lineColor=[1.0, 1.0, 1.0],lineColorSpace='rgb',fillColor=None,fillColorSpace='rgb')
    web_dimension = (screenCorrection(win,Web_size[0]),Web_size[1])
    web_line = visual.Line(win,name='Line',start=(-1.4, 0),end=(1.4, 0),pos=[0, 0],lineWidth=1,
                           lineColor=[1.0, 1.0, 1.0],lineColorSpace='rgb',opacity=1,interpolate=True)

    # DEBUG stimuli
    if DEBUG_MODE:
        fps_text = visual.TextStim(win, units='norm', height=0.05,pos=(-0.98, +0.93), text='starting...',
                                  font=sans, alignHoriz='left', alignVert='bottom', color='yellow')
        fps_text.autoDraw = True

        orientation_details_string = visual.TextStim(win, text = u"eccentricity..", units='norm', height=0.05,
                                             pos=(0.95, +0.93), alignHoriz='right', alignVert='bottom',
                                             font=sans, color='yellow')
        orientation_details_string.autoDraw = True

    # External Aperture (black ring)
    external_aperture_size = tuple([screenCorrection(win,External_ring_size),External_ring_size])
    external_aperture = visual.Aperture(win, size=external_aperture_size, shape='circle')
    external_aperture.enabled = False

    ################################ Definitions/Functions ################################

    ## handle Rkey presses each frame
    def escapeCondition(which_key):
        for key in event.getKeys():
            if key in ['escape', 'q', which_key]:
                return False
        return True

    ################################ Animation starts ################################
    # display instructions and wait
    message1 = visual.TextStim(win, pos=[0,0.5],text='Hit a key when ready.')
    message1.draw()
    fixation.draw()
    win.flip()
    event.waitKeys()                #pause until there's a keypress

    # Scanner trigger wait
    message3 = visual.TextStim(win,pos=[0,0.25],text=scanner_message,font=serif,alignVert='center',
                               wrapWidth=1.5)
    message3.size = .5
    message3.draw()
    win.flip()

    if BUTTON_BOX:
        button_state = button_thread.button_state
        while 1:
            if(button_state['state'][-1]==0):
                break
            if escapeCondition('f') == False: break
    else:
        event.waitKeys()    #pause until there's a keypress

    # Wait Pre_post_stimuli_fixation_time before stimuli
    # Spyder network
    if Spyder_grid:
        for i_dim in range(Spyder_rings):
            web_circle.setSize(tuple([x*(i_dim+1) * 1./Spyder_rings for x in web_dimension]))
            web_circle.draw()
        for i_dim in range(2):
            web_line.setOri(i_dim * 90)
            web_line.draw()
    win.flip()
    core.wait(Pre_post_stimuli_fixation_time)

    t = i_cycle = last_fps_update = 0
    break_flag = True
    globalClock.reset()
    inizio = globalClock.getTime()
    logging.data('First cycle. Number %d/%d at %f (sec.)' % (i_cycle+1,Cycles_number,inizio))

    while (globalClock.getTime() < Total_time and break_flag==True):
        t = globalClock.getTime()

        # Spyder network
        if Spyder_grid:
            for i_dim in range(Spyder_rings):
                web_circle.setSize(tuple([x*(i_dim+1) * 1./Spyder_rings for x in web_dimension]))
                web_circle.draw()
            for i_dim in range(2):
                web_line.setOri(i_dim * 90)
                web_line.draw()

        # External ring
        external_aperture.enabled = True

        # Setup stimulus
        if t % Flash_period < Flash_period / 2.0:  # more accurate to count frames
            stim = wedge1
        else:
            stim = wedge2
        stim.ori = -t * Rotation_rate * 360.0  # set new rotation
        stim.draw()

        # Fixation
        if Rotating_cross:
            fixation.ori = t * Rotation_cross_rate * 360.0  # set new rotation
        if Color_change_cross:
            if t % Color_change_rate < Color_change_rate / 2.0:  # more accurate to count frames
                fixation.lineColor = 'red'
            else:
                fixation.lineColor = 'green'

        if (t >= ((i_cycle+1)*1/Rotation_rate)):
            logging.data('Change orientation. Number %d/%d at %f (sec.)' % (i_cycle+1,Cycles_number,t))
            new_record = globalClock.getTime() - sum(all_changes)
            all_changes.append(new_record)
            i_cycle += 1

        external_aperture.enabled = False

        if DEBUG_MODE:
            if t - last_fps_update > Fps_update_rate:         # update the fps text every second
                fps_text.text = "%.2f fps" % win.fps()
                last_fps_update += 1
            orientation_details_string.text = 'Pass: %d/%d at %.3f (sec.)' % (i_cycle+1,Cycles_number,np.sum(all_changes))

        win.flip()
        break_flag = escapeCondition('f')
        if break_flag == False: break

    logging.data('Total time spent: %.6f' % (globalClock.getTime() - inizio))
    logging.data('Every frame duration saved in %s' % (path_out+Frames_durations_name))
    logging.data('All durations: ' + str(all_changes))
    logging.data('Mean: ' + str(sum(all_changes)/(len(all_changes)+EPSILON)))

    if DEBUG_MODE:
        orientation_details_string.text = 'Ended at %.3f (sec.)' % (globalClock.getTime())

    win.flip()
    # Wait Pre_post_stimuli_fixation_time after stimuli
    core.wait(Pre_post_stimuli_fixation_time)
    return

if __name__ == "__main__":

    # Experiment variables
    today_date = dt.today().strftime('%Y-%m-%d')        # Date (mm/dd/yy)
    operator = 'MS'                                     # Operator
    DEBUG_MODE = False                                   # Debug mode
    BUTTON_BOX = True                                   # Button box monitoring
    subject_code = 'NIA14'                              # Subject-code
    subject_age = 18                                    # Age
    subject_gender = 'female'                           # Gender

    # Prepare out folder
    path_out = Dir_save + today_date + '_' + subject_code + '_polAng'
    path_out = createOutFolder(path_out)

    if not os.path.exists(Dir_save):
        os.makedirs(Dir_save)
    globalClock = core.Clock()

    # Set the log module to report warnings to the standard output window
    logging.setDefaultClock(globalClock)
    logging.console.setLevel(logging.WARNING)
    lastLog=logging.LogFile(path_out+Log_name,level=logging.DATA,filemode='w',encoding='utf8')
    logging.data("------------- " + strftime("%Y-%m-%d %H:%M:%S", gmtime()) + " -------------")
    logging.data(pyschopy_prefs)
    logging.data("Saving in folder: " + path_out)
    logging.data("Operator: " + operator + "\n")
    logging.data("Subject. Code: " + subject_code + " - Age: " + str(subject_age) + " - Gender: " + subject_gender)
    logging.data('***** Starting *****')

    # Create and start buttonBox thread
    if BUTTON_BOX:
        from pypixxlib._libdpx import DPxOpen, DPxSelectDevice, DPxSetMarker, \
            DPxUpdateRegCache, DPxGetMarker, DPxEnableDinDebounce, DPxGetDinValue, \
            DPxSetDinLog, DPxStartDinLog, DPxGetDinStatus, DPxReadDinLog

        button_thread = buttonBoxThread(1, "button box check")
        button_thread.start()

    # Start window
    win = visual.Window([500,500], monitor="mon", screen=1, units="norm", fullscr=Fullscreen,
                        allowStencil=True) # norm
    resX,resY = win.size
    win.recordFrameIntervals = True

    # Main stimulation
    try:
        main(win, globalClock)
    except Exception as e:
        logging.log(e,level=logging.ERROR)

    # Stop buttonBox thread
    if BUTTON_BOX:
        button_thread.stop()

    logging.data('Overall, %i frames were dropped.' % win.nDroppedFrames)
    np.save(path_out+Frames_durations_name,win.frameIntervals[1:])

    logging.data('***** End *****')

    win.close()
    core.quit()