import matplotlib.pyplot as plt
from numba import cuda, jit, njit, prange
from numba.cuda.random import create_xoroshiro128p_states
import numpy as np
from tqdm import tqdm
import scipy.integrate as integrate
import pickle
import time
import tempfile
import h5py
from pynverse import inversefunc

from . import gpu_henon_core as gpu
from . import cpu_henon_core as cpu

from .cpu_henon_core import recursive_accumulation as cpu_accumulate_and_return

def polar_to_cartesian(radius, alpha, theta1, theta2):
    return cpu.polar_to_cartesian(radius, alpha, theta1, theta2)


def cartesian_to_polar(x, px, y, py):
    return cpu.cartesian_to_polar(x, px, y, py)


def circ_trapz(y, x, addendum):
    shape = list(y.shape)
    shape[-1] += 1
    new_y = np.empty(shape)
    new_y[..., :-1] = y[...]
    new_y[..., -1] = y[..., 0]

    new_x = np.empty(x.size + 1)
    new_x[:-1] = x[...]
    new_x[-1] = addendum

    return integrate.trapz(new_y, new_x)


@njit
def modulation(epsilon, n_elements, first_index=0, reversed=False, kind="sps"):
    """Generates a modulation
    
    Parameters
    ----------
    epsilon : float
        intensity of modulation
    n_elements : float
        number of elements
    first_index : int, optional
        starting point of the modulation, by default 0
    kind : string, optional
        kind of modulation: "sps" or "simple" for now, by default "sps"
    
    Returns
    -------
    tuple of ndarray
        (omega_x, omega_y)
    """    
    if "sps":
        coefficients = np.array([1.000e-4,
                                0.218e-4,
                                0.708e-4,
                                0.254e-4,
                                0.100e-4,
                                0.078e-4,
                                0.218e-4])
        modulations = np.array([1 * (2 * np.pi / 868.12),
                                2 * (2 * np.pi / 868.12),
                                3 * (2 * np.pi / 868.12),
                                6 * (2 * np.pi / 868.12),
                                7 * (2 * np.pi / 868.12),
                                10 * (2 * np.pi / 868.12),
                                12 * (2 * np.pi / 868.12)])

        if not reversed:
            number_list = list(range(first_index, first_index + n_elements))
        else:
            number_list = list(range(first_index - n_elements, first_index))[::-1]

        omega_sum = np.array([
            np.sum(coefficients * np.cos(modulations * k)) for k in number_list
        ])
        omega_x = 0.168 * 2 * np.pi * (1 + epsilon * omega_sum)
        omega_y = 0.201 * 2 * np.pi * (1 + epsilon * omega_sum)
        return omega_x, omega_y
    elif "simple":
        omega = 2 * np.pi * 1.15e-3 * np.sqrt(2)
        if not reversed:
            number_list = list(range(first_index, first_index + n_elements))
        else:
            number_list = list(
                range(first_index - n_elements, first_index))[::-1]
        omega_x = 2 * np.pi * 0.201 * \
            (1 + epsilon * np.cos(omega * number_list))
        omega_y = 2 * np.pi * 0.168 * \
            (1 + epsilon * np.cos(omega * number_list))
        return omega_x, omega_y
    else:
        raise NotImplementedError


class partial_track(object):
    """[Updated version] - Basic partial tracker with no particular internal construct. Just dump your coordinates and do whatever you have to do with it!
    """
    def __init__(self):
        pass

    @staticmethod
    def generate_instance(x0, px0, y0, py0, cuda_device=None):
        """Generate an instance of the engine.
        
        Parameters
        ----------
        radius : ndarray
            array of radiuses to consider
        alpha : ndarray
            array of initial alphas
        theta1 : ndarray
            array of initial theta1
        theta2 : ndarray
            array of initial theta2
        epsilon : float
            modulation intensity
        
        Returns
        -------
        class instance
            optimized class instance
        """        
        if cuda_device == None:
            cuda_device = cuda.is_available()
        if cuda_device:
            return gpu_partial_track(x0, px0, y0, py0)
        else:
            return cpu_partial_track(x0, px0, y0, py0)


class cpu_partial_track(partial_track):
    def __init__(self, x0, px0, y0, py0, limit=100.0):

        # save data as members
        self.x0 = x0.copy()
        self.x  = x0.copy()
        self.px0 = px0.copy()
        self.px  = px0.copy()
        self.y0 = y0.copy()
        self.y  = y0.copy()
        self.py0 = py0.copy()
        self.py  = py0.copy()

        # For the modulation
        self.total_iters = 0
        self.step = np.zeros_like(self.x0)

        self.limit = limit

    def compute(self, n_iterations, epsilon, mu=0.0, modulation_kind="sps", full_track=False):
        """Compute the tracking
        
        Returns
        -------
        tuple of ndarray [n_elements]
            (radius, alpha, theta1, theta2, steps)
        """
        omega_x, omega_y = modulation(
            epsilon, n_iterations, self.total_iters, kind=modulation_kind)

        omega_x_cos = np.cos(omega_x)
        omega_x_sin = np.sin(omega_x)
        omega_y_cos = np.cos(omega_y)
        omega_y_sin = np.sin(omega_y)

        if not full_track:
            # Execution
            if mu == 0.0:
                self.x, self.px, self.y, self.py, self.step = cpu.henon_partial_track(
                    self.x, self.px, self.y, self.py, self.step, self.limit, n_iterations, omega_x_sin, omega_x_cos, omega_y_sin, omega_y_cos
                )
            else:
                self.x, self.px, self.y, self.py, self.step = cpu.octo_henon_partial_track(
                    self.x, self.px, self.y, self.py, self.step, self.limit, n_iterations, omega_x_sin, omega_x_cos, omega_y_sin, omega_y_cos, mu
                )

            self.total_iters += n_iterations
            return self.x, self.px, self.y, self.py, self.step
        else:
            data_x = np.ones((n_iterations, self.x.size))
            data_px = np.ones((n_iterations, self.x.size))
            data_y = np.ones((n_iterations, self.x.size))
            data_py = np.ones((n_iterations, self.x.size))
            for i in range(n_iterations):
                # Execution
                if mu == 0.0:
                    self.x, self.px, self.y, self.py, self.step = cpu.henon_partial_track(
                        self.x, self.px, self.y, self.py, self.step, self.limit, 1, omega_x_sin[i:i+1], omega_x_cos[i:i+1], omega_y_sin[i:i+1], omega_y_cos[i:i+1]
                    )
                else:
                    self.x, self.px, self.y, self.py, self.step = cpu.octo_henon_partial_track(
                        self.x, self.px, self.y, self.py, self.step, self.limit, 1, omega_x_sin[i:i+1], omega_x_cos[i:i+1], omega_y_sin[i:i+1], omega_y_cos[i:i+1], mu
                    )
                data_x[i] = self.x
                data_px[i] = self.px
                data_y[i] = self.y
                data_py[i] = self.py
            self.total_iters += n_iterations
            return data_x, data_px, data_y, data_py

    def compute_with_kick(self, n_iterations, epsilon, mu=0.0, kick_module=1e-14, kick_sigma=1e-15, modulation_kind="sps", full_track=False):
        omega_x, omega_y = modulation(
            epsilon, n_iterations, self.total_iters, kind=modulation_kind)

        omega_x_cos = np.cos(omega_x)
        omega_x_sin = np.sin(omega_x)
        omega_y_cos = np.cos(omega_y)
        omega_y_sin = np.sin(omega_y)

        if not full_track:
            # Execution
            if mu == 0.0:
                self.x, self.px, self.y, self.py, self.step = cpu.henon_partial_track_with_kick(
                    self.x, self.px, self.y, self.py, self.step, self.limit, n_iterations, omega_x_sin, omega_x_cos, omega_y_sin, omega_y_cos, kick_module, kick_sigma
                )
            else:
                self.x, self.px, self.y, self.py, self.step = cpu.octo_henon_partial_track_with_kick(
                    self.x, self.px, self.y, self.py, self.step, self.limit, n_iterations, omega_x_sin, omega_x_cos, omega_y_sin, omega_y_cos, mu, kick_module, kick_sigma
                )

            self.total_iters += n_iterations
            return self.x, self.px, self.y, self.py, self.step
        else:
            data_x = np.ones((n_iterations, self.x.size))
            data_px = np.ones((n_iterations, self.x.size))
            data_y = np.ones((n_iterations, self.x.size))
            data_py = np.ones((n_iterations, self.x.size))
            for i in range(n_iterations):
                # Execution
                if mu == 0.0:
                    self.x, self.px, self.y, self.py, self.step = cpu.henon_partial_track_with_kick(
                        self.x, self.px, self.y, self.py, self.step, self.limit, 1, omega_x_sin[
                            i:i+1], omega_x_cos[i:i+1], omega_y_sin[i:i+1], omega_y_cos[i:i+1], kick_module, kick_sigma
                    )
                else:
                    self.x, self.px, self.y, self.py, self.step = cpu.octo_henon_partial_track_with_kick(
                        self.x, self.px, self.y, self.py, self.step, self.limit, 1, omega_x_sin[
                            i:i+1], omega_x_cos[i:i+1], omega_y_sin[i:i+1], omega_y_cos[i:i+1], mu, kick_module, kick_sigma
                    )
                data_x[i] = self.x
                data_px[i] = self.px
                data_y[i] = self.y
                data_py[i] = self.py
            self.total_iters += n_iterations
            return data_x, data_px, data_y, data_py

    def inverse_compute(self, n_iterations, epsilon, modulation_kind="sps"):
        omega_x, omega_y = modulation(
            epsilon, n_iterations, self.total_iters, reversed=True, kind=modulation_kind)

        omega_x_cos = np.cos(omega_x)
        omega_x_sin = np.sin(omega_x)
        omega_y_cos = np.cos(omega_y)
        omega_y_sin = np.sin(omega_y)

        # Execution
        self.x, self.px, self.y, self.py, self.step = cpu.henon_inverse_partial_track(
            self.x, self.px, self.y, self.py, self.step, self.limit, n_iterations, omega_x_sin, omega_x_cos, omega_y_sin, omega_y_cos
        )

        self.total_iters -= n_iterations
        return self.x, self.px, self.y, self.py, self.step

    def inverse_compute_with_kick(self, n_iterations, epsilon, kick_module=1e-14, kick_sigma=1e-15, modulation_kind="sps"):
        omega_x, omega_y = modulation(
            epsilon, n_iterations, self.total_iters, reversed=True, kind=modulation_kind)

        omega_x_cos = np.cos(omega_x)
        omega_x_sin = np.sin(omega_x)
        omega_y_cos = np.cos(omega_y)
        omega_y_sin = np.sin(omega_y)

        # Execution
        self.x, self.px, self.y, self.py, self.step = cpu.henon_inverse_partial_track_with_kick(
            self.x, self.px, self.y, self.py, self.step, self.limit, n_iterations, omega_x_sin, omega_x_cos, omega_y_sin, omega_y_cos, kick_module, kick_sigma
        )

        self.total_iters -= n_iterations
        return self.x, self.px, self.y, self.py, self.step

    def reset(self):
        """Resets the engine
        """        
        self.x = self.x0.copy()
        self.px = self.px0.copy()
        self.y = self.y0.copy()
        self.py = self.py0.copy()
        self.total_iters = 0
        self.step = np.zeros_like(self.x0)

    def get_data(self):
        return self.x, self.px, self.y, self.py, self.step
    
    def get_zero_data(self):
        return self.x0, self.px0, self.y0, self.py0
    
    def add_kick(self, x=None, px=None, y=None, py=None):
        # Add kick
        if x is not None:
            self.x += x
        if px is not None:
            self.px += px
        if y is not None:
            self.y += y
        if py is not None:
            self.py += py


class gpu_partial_track(partial_track):
    def __init__(self, x0, px0, y0, py0, limit=100.0):

        # save data as members
        self.x0 = x0.copy()
        self.x = x0.copy()
        self.d_x = cuda.to_device(x0)

        self.px0 = px0.copy()
        self.px = px0.copy()
        self.d_px = cuda.to_device(px0)
        
        self.y0 = y0.copy()
        self.y = y0.copy()
        self.d_y = cuda.to_device(y0)

        self.py0 = py0.copy()
        self.py = py0.copy()
        self.d_py = cuda.to_device(py0)

        # For the modulation
        self.total_iters = 0
        self.step = np.zeros_like(self.x0)
        self.d_step = cuda.to_device(self.step)

        self.limit = limit

    def compute(self, n_iterations, epsilon, mu=0.0, modulation_kind="sps", full_track=False):
        """Compute the tracking
        
        Returns
        -------
        tuple of ndarray [n_elements]
            (radius, alpha, theta1, theta2, steps)
        """
        threads_per_block = 512
        blocks_per_grid = self.x0.size // 512 + 1

        omega_x, omega_y = modulation(
            epsilon, n_iterations, self.total_iters, kind=modulation_kind)
        
        d_omega_x_sin = cuda.to_device(np.sin(omega_x))
        d_omega_x_cos = cuda.to_device(np.cos(omega_x))
        d_omega_y_sin = cuda.to_device(np.sin(omega_y))
        d_omega_y_cos = cuda.to_device(np.cos(omega_y))

        # Execution
        if not full_track:
            if mu == 0.0:
                gpu.henon_partial_track[blocks_per_grid, threads_per_block](
                    self.d_x, self.d_px, self.d_y, self.d_py, self.d_step, self.limit,
                    n_iterations, d_omega_x_sin, d_omega_x_cos, d_omega_y_sin, d_omega_y_cos
                )
            else:
                gpu.octo_henon_partial_track[blocks_per_grid, threads_per_block](
                    self.d_x, self.d_px, self.d_y, self.d_py, self.d_step, self.limit,
                    n_iterations, d_omega_x_sin, d_omega_x_cos, d_omega_y_sin, d_omega_y_cos, mu
                )
            self.total_iters += n_iterations

            self.d_x.copy_to_host(self.x)
            self.d_y.copy_to_host(self.y)
            self.d_px.copy_to_host(self.px)
            self.d_py.copy_to_host(self.py)
            self.d_step.copy_to_host(self.step)
            
            return self.x, self.px, self.y, self.py, self.step
        else:
            data_x = np.ones((n_iterations, self.x.size))
            data_px = np.ones((n_iterations, self.x.size))
            data_y = np.ones((n_iterations, self.x.size))
            data_py = np.ones((n_iterations, self.x.size))
            for i in range(n_iterations):
                if mu == 0.0:
                    gpu.henon_partial_track[blocks_per_grid, threads_per_block](
                        self.d_x, self.d_px, self.d_y, self.d_py, self.d_step, self.limit,
                        1, d_omega_x_sin[i:i+1], d_omega_x_cos[i:i+1], d_omega_y_sin[i:i+1], d_omega_y_cos[i:i+1]
                    )
                else:
                    gpu.octo_henon_partial_track[blocks_per_grid, threads_per_block](
                        self.d_x, self.d_px, self.d_y, self.d_py, self.d_step, self.limit,
                        1, d_omega_x_sin[i:i+1], d_omega_x_cos[i:i+1], d_omega_y_sin[i:i+1], d_omega_y_cos[i:i+1], mu
                    )
                self.d_x.copy_to_host(data_x[i])
                self.d_y.copy_to_host(data_y[i])
                self.d_px.copy_to_host(data_px[i])
                self.d_py.copy_to_host(data_py[i])
            self.total_iters += n_iterations
            return data_x, data_px, data_y, data_py
    
    def compute_with_kick(self, n_iterations, epsilon, mu=0.0, kick_module=1e-14, kick_sigma=1e-15, modulation_kind="sps", full_track=False):
        """Compute the tracking
        
        Returns
        -------
        tuple of ndarray [n_elements]
            (radius, alpha, theta1, theta2, steps)
        """
        threads_per_block = 512
        blocks_per_grid = self.x0.size // 512 + 1
        rng_states = create_xoroshiro128p_states(threads_per_block * blocks_per_grid, seed=np.random.randint(0, 100000))

        omega_x, omega_y = modulation(
            epsilon, n_iterations, self.total_iters, kind=modulation_kind)
        
        d_omega_x_sin = cuda.to_device(np.sin(omega_x))
        d_omega_x_cos = cuda.to_device(np.cos(omega_x))
        d_omega_y_sin = cuda.to_device(np.sin(omega_y))
        d_omega_y_cos = cuda.to_device(np.cos(omega_y))

        # Execution
        if not full_track:
            if mu == 0.0:
                gpu.henon_partial_track_with_kick[blocks_per_grid, threads_per_block](
                    self.d_x, self.d_px, self.d_y, self.d_py, self.d_step, self.limit,
                    n_iterations, d_omega_x_sin, d_omega_x_cos, d_omega_y_sin, d_omega_y_cos, rng_states, kick_module, kick_sigma
                )
            else:
                gpu.octo_henon_partial_track_with_kick[blocks_per_grid, threads_per_block](
                    self.d_x, self.d_px, self.d_y, self.d_py, self.d_step, self.limit,
                    n_iterations, d_omega_x_sin, d_omega_x_cos, d_omega_y_sin, d_omega_y_cos, mu, rng_states, kick_module, kick_sigma
                )
            self.total_iters += n_iterations

            self.d_x.copy_to_host(self.x)
            self.d_y.copy_to_host(self.y)
            self.d_px.copy_to_host(self.px)
            self.d_py.copy_to_host(self.py)
            self.d_step.copy_to_host(self.step)
            
            return self.x, self.px, self.y, self.py, self.step
        else:
            data_x = np.ones((n_iterations, self.x.size))
            data_px = np.ones((n_iterations, self.x.size))
            data_y = np.ones((n_iterations, self.x.size))
            data_py = np.ones((n_iterations, self.x.size))
            for i in range(n_iterations):
                if mu == 0.0:
                    gpu.henon_partial_track_with_kick[blocks_per_grid, threads_per_block](
                        self.d_x, self.d_px, self.d_y, self.d_py, self.d_step, self.limit,
                        1, d_omega_x_sin[i:i+1], d_omega_x_cos[i:i+1], d_omega_y_sin[i:i+1], d_omega_y_cos[i:i+1], rng_states, kick_module, kick_sigma
                    )
                else:
                    gpu.octo_henon_partial_track_with_kick[blocks_per_grid, threads_per_block](
                        self.d_x, self.d_px, self.d_y, self.d_py, self.d_step, self.limit,
                        1, d_omega_x_sin[i:i+1], d_omega_x_cos[i:i+1], d_omega_y_sin[i:i+1], d_omega_y_cos[i:i+1], mu, rng_states, kick_module, kick_sigma
                    )
                self.d_x.copy_to_host(data_x[i])
                self.d_y.copy_to_host(data_y[i])
                self.d_px.copy_to_host(data_px[i])
                self.d_py.copy_to_host(data_py[i])
            self.total_iters += n_iterations
            return data_x, data_px, data_y, data_py

    def inverse_compute(self, n_iterations, epsilon, modulation_kind="sps"):
        threads_per_block = 512
        blocks_per_grid = self.x0.size // 512 + 1

        omega_x, omega_y = modulation(
            epsilon, n_iterations, self.total_iters, reversed=True, kind=modulation_kind)

        d_omega_x_cos = cuda.to_device(np.cos(omega_x))
        d_omega_x_sin = cuda.to_device(np.sin(omega_x))
        d_omega_y_cos = cuda.to_device(np.cos(omega_y))
        d_omega_y_sin = cuda.to_device(np.sin(omega_y))

        gpu.henon_inverse_partial_track[blocks_per_grid, threads_per_block](
            self.d_x, self.d_px, self.d_y, self.d_py, self.d_step, self.limit,
            n_iterations, d_omega_x_sin, d_omega_x_cos, d_omega_y_sin, d_omega_y_cos
        )
        self.total_iters -= n_iterations

        self.d_x.copy_to_host(self.x)
        self.d_y.copy_to_host(self.y)
        self.d_px.copy_to_host(self.px)
        self.d_py.copy_to_host(self.py)
        self.d_step.copy_to_host(self.step)

        return self.x, self.px, self.y, self.py, self.step

    def inverse_compute_with_kick(self, n_iterations, epsilon, kick_module=1e-14, kick_sigma=1e-15, modulation_kind="sps"):
        threads_per_block = 512
        blocks_per_grid = self.x0.size // 512 + 1
        rng_states = create_xoroshiro128p_states(
            threads_per_block * blocks_per_grid, seed=np.random.randint(0, 100000))

        omega_x, omega_y = modulation(
            epsilon, n_iterations, self.total_iters, reversed=True, kind=modulation_kind)

        d_omega_x_cos = cuda.to_device(np.cos(omega_x))
        d_omega_x_sin = cuda.to_device(np.sin(omega_x))
        d_omega_y_cos = cuda.to_device(np.cos(omega_y))
        d_omega_y_sin = cuda.to_device(np.sin(omega_y))

        gpu.henon_inverse_partial_track_with_kick[blocks_per_grid, threads_per_block](
            self.d_x, self.d_px, self.d_y, self.d_py, self.d_step, self.limit,
            n_iterations, d_omega_x_sin, d_omega_x_cos, d_omega_y_sin, d_omega_y_cos, rng_states, kick_module, kick_sigma
        )
        self.total_iters -= n_iterations

        self.d_x.copy_to_host(self.x)
        self.d_y.copy_to_host(self.y)
        self.d_px.copy_to_host(self.px)
        self.d_py.copy_to_host(self.py)
        self.d_step.copy_to_host(self.step)

        return self.x, self.px, self.y, self.py, self.step

    def reset(self):
        """Resets the engine
        """        
        self.x = self.x0.copy()
        self.px = self.px0.copy()
        self.y = self.y0.copy()
        self.py = self.py0.copy()

        self.d_x = cuda.to_device(self.x0)
        self.d_px = cuda.to_device(self.px0)
        self.d_y = cuda.to_device(self.y0)
        self.d_py = cuda.to_device(self.py0)

        self.total_iters = 0
        self.step = np.zeros_like(self.x0)
        self.d_step = cuda.to_device(self.step)

    def get_data(self):
        return self.x, self.px, self.y, self.py, self.step

    def get_zero_data(self):
        return self.x0, self.px0, self.y0, self.py0

    def add_kick(self, x=None, px=None, y=None, py=None):
        # Add kick
        if x is not None:
            self.x += x
        if px is not None:
            self.px += px
        if y is not None:
            self.y += y
        if py is not None:
            self.py += py
        # Update GPU
        self.d_x = cuda.to_device(self.x)
        self.d_px = cuda.to_device(self.px)
        self.d_y = cuda.to_device(self.y)
        self.d_py = cuda.to_device(self.py)


class uniform_scan(object):
    """With this class we can easly scan a uniform 4D cube of the Hénon map"""
    def __init__(self):
        pass

    def scan(self):
        pass

    def save_values(self, f, label="Henon map"):
        self.db.flush()
        dest = h5py.File(f, mode="w")
        
        dest.attrs["label"] = label
        dest.attrs["epsilon"] = self.db.attrs["epsilon"]
        dest.attrs["top"] = self.db.attrs["top"]
        dest.attrs["steps"] = self.db.attrs["steps"]
        dest.attrs["starting_radius"] = self.db.attrs["starting_radius"]
        dest.attrs["coordinates"] = self.db.attrs["coordinates"]
        dest.attrs["samples"] = self.db.attrs["samples"]
        dest.attrs["max_turns"] = self.db.attrs["max_turns"]

        g = dest.create_group("data")
        self.db.copy("data/times", g)

        dest.create_dataset(
            "/data/weights", (self.samples, self.samples, self.samples, self.samples), dtype=np.float, compression="lzf")
        dest.close()

    @staticmethod
    def generate_instance(epsilon, top, steps, starting_radius=0.0001, cuda_device=None, tempdir=None):
        """Create an uniform scan object

        Parameters
        ----------
        epsilon : float
            modulation intensity
        top : float
            maximum radius
        steps : int
            steps from zero to top (becomes steps * 2 + 1)
        starting_radius : float, optional
            from which position we have to start with the actual computation, by default 0.0001
        cuda_device : bool, optional
            do we have a CUDA capable device (make it manual), by default None

        Returns
        -------
        object
            uniform_scan object
        """        
        if cuda_device == None:
            cuda_device = cuda.is_available()
        if cuda_device:
            return gpu_uniform_scan(epsilon, top, steps, starting_radius, tempdir)
        else:
            return cpu_uniform_scan(epsilon, top, steps, starting_radius, tempdir)


class cpu_uniform_scan(uniform_scan):
    def __init__(self, epsilon, top, steps, starting_radius=0.0001, tempdir=None):
        self.tf = tempfile.TemporaryFile(dir=tempdir)
        self.db = h5py.File(self.tf, mode="w")

        self.samples = steps * 2 + 1
        self.coords = np.linspace(-top, top, self.samples)
        
        self.db.attrs["epsilon"] = epsilon
        self.db.attrs["top"] = top
        self.db.attrs["steps"] = steps
        self.db.attrs["starting_radius"] = starting_radius
        self.db.attrs["samples"] = self.samples
        self.db.attrs["coordinates"] = self.coords

        self.bool_mask = self.db.create_dataset(
            "/data/bool_mask", (self.samples, self.samples, self.samples, self.samples), dtype=np.bool, compression="lzf")

        self.coords2 = np.power(self.coords, 2)
        for i in tqdm(range(len(self.coords)), desc="Make the boolean mask"):
            px, y, py = np.meshgrid(self.coords2, self.coords2, self.coords2)
            self.bool_mask[i] = (
                self.coords[i] ** 2 
                + px
                + y
                + py
                >= starting_radius ** 2
            )

        self.times = self.db.create_dataset(
            "/data/times", (self.samples, self.samples, self.samples, self.samples), dtype=np.int32, compression="lzf")
        
    def scan(self, max_turns, modulation_kind="sps"):
        """Execute a scanning of everything

        Parameters
        ----------
        max_turns : int
            turn limit

        Returns
        -------
        ndarray
            4d array with stable iterations inside
        """
        self.db.attrs["max_turns"] = max_turns

        omega_x, omega_y = modulation(self.db.attrs["epsilon"], max_turns, kind=modulation_kind)

        for i in tqdm(range(len(self.times))):
            px, y, py = np.meshgrid(self.coords, self.coords, self.coords)
            x = np.ones_like(px) * self.coords[i]
            self.times[i] = cpu.henon_map_to_the_end(
                x, px, y, py, 10.0, max_turns, omega_x, omega_y, self.bool_mask[i]
            )

    def scan_octo(self, max_turns, mu, modulation_kind="sps"):
        """Execute a scanning of everything

        Parameters
        ----------
        max_turns : int
            turn limit

        mu : float
            mu parameter

        Returns
        -------
        ndarray
            4d array with stable iterations inside
        """
        self.db.attrs["max_turns"] = max_turns
        self.db.attrs["mu"] = mu

        omega_x, omega_y = modulation(self.db.attrs["epsilon"], max_turns, kind=modulation_kind)

        for i in tqdm(range(len(self.times))):
            px, y, py = np.meshgrid(self.coords, self.coords, self.coords)
            x = np.ones_like(px) * self.coords[i]
            self.times[i] = cpu.octo_henon_map_to_the_end(
                x, px, y, py, 10.0, max_turns, omega_x, omega_y, mu, self.bool_mask[i]
            )


class gpu_uniform_scan(uniform_scan):
    def __init__(self, epsilon, top, steps, starting_radius=0.0001, tempdir=None):
        self.tf = tempfile.TemporaryFile(dir=tempdir)
        self.db = h5py.File(self.tf, mode="w")

        self.samples = steps * 2 + 1
        self.coords = np.linspace(-top, top, self.samples)

        self.db.attrs["epsilon"] = epsilon
        self.db.attrs["top"] = top
        self.db.attrs["steps"] = steps
        self.db.attrs["starting_radius"] = starting_radius
        self.db.attrs["samples"] = self.samples
        self.db.attrs["coordinates"] = self.coords

        self.bool_mask = self.db.create_dataset(
            "/data/bool_mask", (self.samples, self.samples, self.samples, self.samples), dtype=np.bool, compression="lzf")

        self.coords2 = np.power(self.coords, 2)
        for i in tqdm(range(len(self.coords)), desc="Make the boolean mask"):
            px, y, py = np.meshgrid(self.coords2, self.coords2, self.coords2)
            self.bool_mask[i] = (
                self.coords[i] ** 2
                + px
                + y
                + py
                >= starting_radius ** 2
            )

        self.times = self.db.create_dataset(
            "/data/times", (self.samples, self.samples, self.samples, self.samples), dtype=np.int32, compression="lzf")

    def scan(self, max_turns, modulation_kind="sps"):
        """Execute a scanning of everything

        Parameters
        ----------
        max_turns : int
            turn limit

        Returns
        -------
        ndarray
            4d array with stable iterations inside
        """        
        threads_per_block = 512
        blocks_per_grid = 10

        self.db.attrs["max_turns"] = max_turns

        omega_x, omega_y = modulation(self.db.attrs["epsilon"], max_turns, kind=modulation_kind)
        d_omega_x = cuda.to_device(omega_x)
        d_omega_y = cuda.to_device(omega_y)

        t_f = np.empty(shape=(self.samples, self.samples, self.samples), dtype=np.int32).flatten()

        for i in tqdm(range(len(self.times)), smoothing=1.0):
            px, y, py = np.meshgrid(self.coords, self.coords, self.coords)
            x = np.ones_like(px) * self.coords[i]

            d_x = cuda.to_device(x.flatten())
            d_px = cuda.to_device(px.flatten())
            d_y = cuda.to_device(y.flatten())
            d_py = cuda.to_device(py.flatten())
            d_times = cuda.to_device(np.zeros(x.size, dtype=np.int32))
            d_bool_mask = cuda.to_device(np.asarray(self.bool_mask[i]).flatten())

            gpu.henon_map_to_the_end[blocks_per_grid, threads_per_block](
                d_x, d_px, d_y, d_py, d_times, 10.0, max_turns, d_omega_x, d_omega_y, d_bool_mask
            )

            d_times.copy_to_host(t_f)
            self.times[i] = t_f.reshape(x.shape)
    
    def scan_octo(self, max_turns, mu, modulation_kind="sps"):
        """Execute a scanning of everything

        Parameters
        ----------
        max_turns : int
            turn limit

        mu : float
            param

        Returns
        -------
        ndarray
            4d array with stable iterations inside
        """
        threads_per_block = 512
        blocks_per_grid = 10

        self.db.attrs["max_turns"] = max_turns
        self.db.attrs["mu"] = mu

        omega_x, omega_y = modulation(self.db.attrs["epsilon"], max_turns, kind=modulation_kind)
        omega_x_sin = np.sin(omega_x)
        omega_x_cos = np.cos(omega_x)
        omega_y_sin = np.sin(omega_y)
        omega_y_cos = np.cos(omega_y)

        d_omega_x_sin = cuda.to_device(
            np.asarray(omega_x_sin, dtype=np.float64))
        d_omega_x_cos = cuda.to_device(
            np.asarray(omega_x_cos, dtype=np.float64))
        d_omega_y_sin = cuda.to_device(
            np.asarray(omega_y_sin, dtype=np.float64))
        d_omega_y_cos = cuda.to_device(
            np.asarray(omega_y_cos, dtype=np.float64))

        t_f = np.empty(shape=(self.samples, self.samples,
                              self.samples), dtype=np.int32).flatten()

        for i in tqdm(range(len(self.times)), smoothing=1.0):
            px, y, py = np.meshgrid(self.coords, self.coords, self.coords)
            x = np.ones_like(px) * self.coords[i]

            d_x = cuda.to_device(np.asarray(x, dtype=np.float64).flatten())
            d_px = cuda.to_device(np.asarray(px, dtype=np.float64).flatten())
            d_y = cuda.to_device(np.asarray(y, dtype=np.float64).flatten())
            d_py = cuda.to_device(np.asarray(py, dtype=np.float64).flatten())
            d_times = cuda.to_device(np.zeros(x.size, dtype=np.int32))
            d_bool_mask = cuda.to_device(
                np.asarray(self.bool_mask[i]).flatten())

            gpu.octo_henon_map_to_the_end[blocks_per_grid, threads_per_block](
                d_x, d_px, d_y, d_py, d_times, 10.0, max_turns, d_omega_x_sin, d_omega_x_cos, d_omega_y_sin, d_omega_y_cos, np.float64(mu), d_bool_mask
            )

            d_times.copy_to_host(t_f)
            self.times[i] = t_f.reshape(x.shape)


class radial_block(object):
    """More advanced and complete radial_block analysis, replaces the deprecated radial_scan class"""
    def __init__(self):
        pass

    def scan(self):
        pass

    def save_values(self, f, label="Henon map"):
        self.db.flush()
        dest = h5py.File(f, mode="w")

        dest.attrs["label"] = label
        dest.attrs["epsilon"] = self.db.attrs["epsilon"]
        dest.attrs["radial_samples"] = self.db.attrs["radial_samples"]
        dest.attrs["max_radius"] = self.db.attrs["max_radius"]
        dest.attrs["alpha"] = self.db.attrs["alpha"]
        dest.attrs["theta1"] = self.db.attrs["theta1"]
        dest.attrs["theta2"] = self.db.attrs["theta2"]
        dest.attrs["dr"] = self.db.attrs["dr"]
        dest.attrs["max_turns"] = self.db.attrs["max_turns"]
        
        g = dest.create_group("data")
        self.db.copy("data/times", g)

        dest.create_dataset(
            "/data/weights", (self.db.attrs["radial_samples"], len(self.db.attrs["alpha"]), len(self.db.attrs["theta1"]), len(self.db.attrs["theta2"])), dtype=np.float, compression="lzf")
        dest.close()
    
    @staticmethod
    def generate_instance(radial_samples, alpha, theta1, theta2, epsilon, max_radius=1.0, starting_radius=0.0, cuda_device=None, tempdir=None):
        if cuda_device == None:
            cuda_device = cuda.is_available()
        if cuda_device:
            return gpu_radial_block(radial_samples, alpha, theta1, theta2, epsilon, max_radius, starting_radius, tempdir)
        else:
            return cpu_radial_block(radial_samples, alpha, theta1, theta2, epsilon, max_radius, starting_radius, tempdir)

        
class cpu_radial_block(radial_block):
    def __init__(self, radial_samples, alpha, theta1, theta2, epsilon, max_radius=1.0, starting_radius=0.0, tempdir=None):
        assert alpha.size == theta1.size
        assert alpha.size == theta2.size
        assert max_radius > 0.0
        
        self.tf = tempfile.TemporaryFile(dir=tempdir)
        self.db = h5py.File(self.tf, mode="w")

        self.db.attrs["epsilon"] = epsilon
        self.db.attrs["radial_samples"] = radial_samples
        self.db.attrs["starting_radius"] = starting_radius
        self.starting_radius = starting_radius
        self.db.attrs["max_radius"] = max_radius
        self.alpha = alpha
        self.theta1 = theta1
        self.theta2 = theta2
        self.db.attrs["alpha"] = alpha
        self.db.attrs["theta1"] = theta1
        self.db.attrs["theta2"] = theta2

        self.r_list, self.dr = np.linspace(0, max_radius, radial_samples, retstep=True)
        self.db.attrs["dr"] = self.dr

        self.bool_mask = self.db.create_dataset(
            "/data/bool_mask", (radial_samples, len(alpha), len(theta1), len(theta2)), dtype=np.bool, compression="lzf")
        
        for i in tqdm(range(radial_samples)):
            self.bool_mask[i] = self.r_list[i] >= starting_radius

        self.times = self.db.create_dataset(
            "/data/times", (radial_samples, len(alpha), len(theta1), len(theta2)), dtype=np.int32, compression="lzf")
    
    def scan(self, max_turns, modulation_kind="sps"):
        self.db.attrs["max_turns"] = max_turns

        omega_x, omega_y = modulation(self.db.attrs["epsilon"], max_turns, kind=modulation_kind)

        aa, th1, th2 = np.meshgrid(
            self.alpha, self.theta1, self.theta2, indexing='ij'
        )

        for i in tqdm(range(len(self.times)), smoothing=1.0):
            if self.r_list[i] < self.starting_radius:
                self.times[i] = max_turns
            else:
                x, px, y, py = polar_to_cartesian(self.r_list[i], aa, th1, th2)

                self.times[i] = cpu.henon_map_to_the_end(
                    x, px, y, py, 10.0, max_turns, omega_x, omega_y, self.bool_mask[i]
                )
    
    def scan_octo(self, max_turns, mu, modulation_kind="sps"):
        self.db.attrs["max_turns"] = max_turns
        self.db.attrs["mu"] = mu

        omega_x, omega_y = modulation(self.db.attrs["epsilon"], max_turns, kind=modulation_kind)

        aa, th1, th2 = np.meshgrid(
            self.alpha, self.theta1, self.theta2, indexing='ij'
        )

        for i in tqdm(range(len(self.times)), smoothing=1.0):
            if self.r_list[i] < self.starting_radius:
                self.times[i] = max_turns
            else:
                x, px, y, py = polar_to_cartesian(self.r_list[i], aa, th1, th2)

                self.times[i] = cpu.octo_henon_map_to_the_end(
                    x, px, y, py, 10.0, max_turns, omega_x, omega_y, mu, self.bool_mask[i]
                )
        

class gpu_radial_block(radial_block):
    def __init__(self, radial_samples, alpha, theta1, theta2, epsilon, max_radius=1.0, starting_radius=0.0, tempdir=None):
        assert alpha.size == theta1.size
        assert alpha.size == theta2.size
        assert max_radius > 0.0

        self.tf = tempfile.TemporaryFile(dir=tempdir)
        self.db = h5py.File(self.tf, mode="w")

        self.db.attrs["epsilon"] = epsilon
        self.db.attrs["radial_samples"] = radial_samples
        self.db.attrs["starting_radius"] = starting_radius
        self.starting_radius = starting_radius
        self.db.attrs["max_radius"] = max_radius
        self.alpha = alpha
        self.theta1 = theta1
        self.theta2 = theta2
        self.db.attrs["alpha"] = alpha
        self.db.attrs["theta1"] = theta1
        self.db.attrs["theta2"] = theta2

        self.r_list, self.dr = np.linspace(
            0, max_radius, radial_samples + 1, retstep=True)
        self.db.attrs["dr"] = self.dr

        self.bool_mask = self.db.create_dataset(
            "/data/bool_mask", (radial_samples, len(alpha), len(theta1), len(theta2)), dtype=np.bool, compression="lzf")

        for i in tqdm(range(radial_samples)):
            self.bool_mask[i] = self.r_list[i] >= starting_radius

        self.times = self.db.create_dataset(
            "/data/times", (radial_samples, len(alpha), len(theta1), len(theta2)), dtype=np.int32, compression="lzf")

    def scan(self, max_turns, modulation_kind="sps"):
        threads_per_block = 512
        blocks_per_grid = 10

        self.db.attrs["max_turns"] = max_turns

        omega_x, omega_y = modulation(self.db.attrs["epsilon"], max_turns, kind=modulation_kind)
        d_omega_x = cuda.to_device(omega_x)
        d_omega_y = cuda.to_device(omega_y)

        t_f = np.empty(shape=(len(self.alpha), len(
            self.theta1), len(self.theta2)), dtype=np.int32).flatten()
        aa, th1, th2 = np.meshgrid(
            self.alpha, self.theta1, self.theta2, indexing='ij'
        )

        for i in tqdm(range(len(self.times)), smoothing=1.0):
            if self.r_list[i] < self.starting_radius:
                self.times[i] = max_turns
            else:
                x, px, y, py = polar_to_cartesian(self.r_list[i], aa, th1, th2)
                d_x = cuda.to_device(x.flatten())
                d_px = cuda.to_device(px.flatten())
                d_y = cuda.to_device(y.flatten())
                d_py = cuda.to_device(py.flatten())
                d_bool_mask = cuda.to_device(
                    np.asarray(self.bool_mask[i]).flatten())
                d_times = cuda.to_device(np.zeros(x.size, dtype=np.int32))

                gpu.henon_map_to_the_end[blocks_per_grid, threads_per_block](
                    d_x, d_px, d_y, d_py, d_times, 10.0, max_turns, d_omega_x, d_omega_y, d_bool_mask
                )

                d_times.copy_to_host(t_f)
                self.times[i] = t_f.reshape(x.shape)
    
    def scan_octo(self, max_turns, mu, modulation_kind="sps"):
        threads_per_block = 512
        blocks_per_grid = 10

        self.db.attrs["max_turns"] = max_turns
        self.db.attrs["mu"] = mu

        omega_x, omega_y = modulation(self.db.attrs["epsilon"], max_turns, kind=modulation_kind)
        omega_x_sin = np.sin(omega_x)
        omega_x_cos = np.cos(omega_x)
        omega_y_sin = np.sin(omega_y)
        omega_y_cos = np.cos(omega_y)

        d_omega_x_sin = cuda.to_device(
            np.asarray(omega_x_sin, dtype=np.float64))
        d_omega_x_cos = cuda.to_device(
            np.asarray(omega_x_cos, dtype=np.float64))
        d_omega_y_sin = cuda.to_device(
            np.asarray(omega_y_sin, dtype=np.float64))
        d_omega_y_cos = cuda.to_device(
            np.asarray(omega_y_cos, dtype=np.float64))

        t_f = np.empty(shape=(len(self.alpha), len(
            self.theta1), len(self.theta2)), dtype=np.int32).flatten()
        aa, th1, th2 = np.meshgrid(
            self.alpha, self.theta1, self.theta2, indexing='ij'
        )

        for i in tqdm(range(len(self.times)), smoothing=1.0):
            if self.r_list[i] < self.starting_radius:
                self.times[i] = max_turns
            else:
                x, px, y, py = polar_to_cartesian(self.r_list[i], aa, th1, th2)
                d_x = cuda.to_device(np.asarray(x, dtype=np.float64).flatten())
                d_px = cuda.to_device(np.asarray(px, dtype=np.float64).flatten())
                d_y = cuda.to_device(np.asarray(y, dtype=np.float64).flatten())
                d_py = cuda.to_device(np.asarray(py, dtype=np.float64).flatten())
                d_bool_mask = cuda.to_device(np.asarray(self.bool_mask[i]).flatten())
                d_times = cuda.to_device(np.zeros(x.size, dtype=np.int32))

                gpu.octo_henon_map_to_the_end[blocks_per_grid, threads_per_block](
                    d_x, d_px, d_y, d_py, d_times, 10.0, max_turns, d_omega_x_sin, d_omega_x_cos, d_omega_y_sin, d_omega_y_cos, np.float64(mu), d_bool_mask
                )

                d_times.copy_to_host(t_f)
                self.times[i] = t_f.reshape(x.shape)


class uniform_analyzer(object):
    def __init__(self, hdf5_dir):
        self.db = h5py.File(hdf5_dir, mode="r+")
        self.samples = self.db.attrs["samples"]
        self.coords = self.db.attrs["coordinates"]
        self.coords2 = np.power(self.coords, 2)
        self.times = self.db["/data/times"]
        self.weights = self.db["/data/weights"]

    def assign_weights(self, f=lambda x, px, y, py: 1.0, radial_cut=-1.0):
        if radial_cut == -1.0:
            radial_cut = self.db.attrs["top"]

        for i in tqdm(range(self.samples), desc="assigning weights"):
            px, y, py = np.meshgrid(self.coords, self.coords, self.coords)
            self.weights[i] = f(self.coords[i], px, y, py) * (
                np.power(self.coords[i], 2) + np.power(px, 2) + np.power(y, 2) + np.power(py, 2) <= np.power(radial_cut, 2)
            )

    def compute_loss(self, sample_list, normalization=True):
        prelim_values = np.empty(self.samples)        
        for i in tqdm(range(self.samples), desc="baseline integration"):
            prelim_values[i] = integrate.trapz(
                integrate.trapz(
                    integrate.trapz(
                        self.weights[i],
                        x=self.coords
                    ),
                    x=self.coords
                ),
                x=self.coords
            )
        baseline = integrate.trapz(prelim_values, x=self.coords)

        values = np.empty(len(sample_list))
        for j, sample in tqdm(enumerate(sample_list), desc="other integrals...", total=len(sample_list)):
            prelim_values = np.empty(self.samples)
            for i in range(self.samples):
                prelim_values[i] = integrate.trapz(
                    integrate.trapz(
                        integrate.trapz(
                            self.weights[i] * (self.times[i] >= sample),
                            x=self.coords
                        ),
                        x=self.coords
                    ),
                    x=self.coords
                )
            values[j] = integrate.trapz(prelim_values, x=self.coords)

        if normalization:
            values /= baseline

        return values

    def compute_loss_cut(self, cut):
        prelim_values = np.empty(self.samples)
        for i in tqdm(range(self.samples), desc="baseline integration"):
            px, y, py = np.meshgrid(self.coords2, self.coords2, self.coords2)
            prelim_values[i] = integrate.trapz(
                integrate.trapz(
                    integrate.trapz(
                        self.weights[i] * (np.power(self.coords[i], 2) + px + y + py <= np.power(cut, 2)),
                        x=self.coords
                    ),
                    x=self.coords
                ),
                x=self.coords
            )
        return integrate.trapz(prelim_values, self.coords)


class uniform_radial_scanner(object):
    """This class is for analyzing the loss values of a (somewhat) angular uniform scan"""

    def __init__(self, hdf5_dir):
        self.db = h5py.File(hdf5_dir, mode="r+")
        self.max_radius = self.db.attrs["max_radius"]
        self.samples = self.db.attrs["radial_samples"]
        self.r_list, self.dr = np.linspace(
            0, self.max_radius, self.samples, retstep=True)
        self.alpha = self.db.attrs["alpha"]
        self.theta1 = self.db.attrs["theta1"]
        self.theta2 = self.db.attrs["theta2"]
        self.dr = self.db.attrs["dr"]

        self.times = self.db["/data/times"]
        self.weights = self.db["/data/weights"]

    @staticmethod
    def static_extract_radiuses(n_alpha, n_theta1, n_theta2, samples, times, dr, radius_db):
        for index, i1 in enumerate(range(n_alpha)):
            #print(index, "/", n_alpha, flush=True)
            #for i2 in range(n_theta1):
                #for i3 in range(n_theta2):
            temp = times[:, i1, :, :]
            values = np.empty((len(samples), n_theta1, n_theta2))
            for i, sample in enumerate(samples):
                values[i] = np.argmin(temp >= sample, axis=0) - 1
            values[values < 0] = len(values)
            values = (values + 1) * dr
            radius_db[:, i1, :, :] = values

    def compute_DA_standard(self, sample_list, get_radiuses=False):
        self.sample_list = sample_list
        """
        try:
            self.db_sample_list = self.db.require_dataset(
                "/data/DA_samples", shape=(len(sample_list),), dtype=np.int32, exact=True)
        except TypeError:
            del self.db["/data/DA_samples"]
            self.db_sample_list = self.db.require_dataset(
                "/data/DA_samples", shape=(len(sample_list),), dtype=np.int32, exact=True)
        
        try:
            self.radiuses = self.db.require_dataset(
                "/data/DA_radiuses", shape=(len(sample_list), len(self.alpha), len(self.theta1), len(self.theta2)), dtype=np.float, exact=True)
        except TypeError:
            del self.db["/data/DA_radiuses"]
            self.radiuses = self.db.require_dataset(
                "/data/DA_radiuses", shape=(len(sample_list), len(self.alpha), len(self.theta1), len(self.theta2)), dtype=np.float, exact=True)
        """
        radiuses = np.empty((len(sample_list), len(self.alpha), len(self.theta1,), len(self.theta1)))

        # I HAVE A BUG THAT HAS TO BE FOUND
        self.times[0] = np.max(self.times)

        self.static_extract_radiuses(
            len(self.alpha), len(self.theta1), len(self.theta2),
            sample_list, self.times, self.dr, radiuses)

        mod_radiuses = np.power(radiuses, 4)
        mod_radiuses = integrate.trapz(mod_radiuses, x=self.theta2)
        mod_radiuses = integrate.trapz(mod_radiuses, x=self.theta1)
        mod_radiuses = integrate.trapz(mod_radiuses, x=self.alpha)

        self.DA = np.power(
            mod_radiuses / (2 * self.theta1[-1] * self.theta2[-1]), 1/4)
        
        e_alpha = np.mean(
            np.absolute(radiuses[:, 1:] - radiuses[:, :-1]),
            axis=(1, 2, 3)) ** 2
        e_theta1 = np.mean(
            np.absolute(radiuses[:, :, 1:] - radiuses[:, :, :-1]),
            axis=(1, 2, 3)) ** 2
        e_theta2 = np.mean(
            np.absolute(radiuses[:, :, :, 1:] - radiuses[:, :, :, :-1]),
            axis=(1, 2, 3)) ** 2
        e_radius = self.dr ** 2
        self.error = np.sqrt(
            (e_radius + e_alpha + e_theta1 + e_theta2) / 4)
        
        if not get_radiuses:
            return self.DA, self.error
        else:
            return self.DA, self.error, radiuses

    def create_weights_in_dataset(self, file_destination, f=lambda r, a, th1, th2: np.ones_like(a)):
        dest = h5py.File(file_destination, mode="w")
        weights = dest.create_dataset(
            "/data/weights", (self.db.attrs["radial_samples"], len(self.db.attrs["alpha"]), len(self.db.attrs["theta1"]), len(self.db.attrs["theta2"])), dtype=np.float, compression="lzf")
        
        for i in tqdm(range(0, len(self.r_list), 10), desc="assigning weights"):
            rr, aa, th1, th2 = np.meshgrid(
                self.r_list[i: min(i + 10, len(self.r_list))
                       ], self.alpha, self.theta1, self.theta2, indexing='ij'
            )
            weights[i: min(i + 10, len(self.r_list))
                         ] = f(rr, aa, th1, th2)
        
        dest.close()

    def assign_weights_from_file(self, file):
        self.weight_db = h5py.File(file, mode="r")
        self.weights = self.weight_db["/data/weights"]

    def assign_weights(self, f=lambda r, a, th1, th2: np.ones_like(a)):
        """Assign weights to the various radial samples computed (not-so-intuitive to setup, beware...).

        Parameters
        ----------
        f : lambda, optional
            the lambda to assign the weights with, by default returns r
            this lambda has to take as arguments
            r : float
                the radius
            a : float
                the alpha angle
            th1 : float
                the theta1 angle
            th2 : float
                the theta2 angle
        """
        aa, th1, th2 = np.meshgrid(
            self.alpha, self.theta1, self.theta2, indexing='ij'
        )
        for i, r in tqdm(enumerate(self.r_list), desc="assigning weights", total=len(self.r_list)):
            self.weights[i] = f(r, aa, th1, th2) 

    def compute_loss(self, sample_list, cutting_point=-1.0, normalization=True):
        """Compute the loss based on a boolean masking of the various timing values.

        Parameters
        ----------
        sample_list : ndarray
            list of times to use as samples
        cutting_point : float, optional
            radius to set-up as cutting point for normalization purposes, by default -1.0
        normalization : boolean, optional
            execute normalization? By default True

        Returns
        -------
        ndarray
            the values list measured (last element is the cutting point value 1.0 used for renormalization of the other results.)
        """
        if cutting_point == -1.0:
            cutting_point = self.db.attrs["max_radius"]

        # COMPUTING BASELINE
        prelim_values = np.zeros(self.r_list.size)
        top_i = np.argmin(self.r_list <= cutting_point)
        if top_i == 0:
            top_i = len(self.weights)
        prelim_values[: top_i] = integrate.trapz(
            circ_trapz(
                circ_trapz(
                    self.weights[: top_i],
                    self.theta2,
                    np.pi * 2
                ),
                self.theta1,
                np.pi * 2
            ) * np.sin(self.alpha) * np.cos(self.alpha),
            self.alpha
        )
        baseline = integrate.trapz(prelim_values * np.power(self.r_list, 3), self.r_list)

        # COMPUTING THE REST
        values = np.empty(len(sample_list))
        for j, sample in enumerate(sample_list):
            prelim_values = np.zeros(self.r_list.size)
            ## THIS MAKES THE SAME CUTTING FOR THE WHOLE COMPUTATION ####
            top_i = np.argmin(self.r_list <= cutting_point)
            #############################################################
            if top_i == 0:
                top_i = len(self.weights)
            prelim_values[: top_i] = integrate.trapz(
                circ_trapz(
                    circ_trapz(
                        self.weights[: top_i] * (self.times[: top_i] >= sample),
                        self.theta2,
                        np.pi * 2
                    ),
                    self.theta1,
                    np.pi * 2
                ) * np.sin(self.alpha) * np.cos(self.alpha),
                self.alpha
            )
            values[j] = integrate.trapz(prelim_values * np.power(self.r_list, 3), self.r_list)

        if normalization:
            values /= baseline
        return np.abs(values)

    def compute_loss_cut(self, cutting_point=-1.0):
        """Compute the loss based on a simple DA cut.

        Parameters
        ----------
        cutting_point : float
            radius to set-up as cutting point

        Returns
        -------
        float
            the (not-normalized) value
        """
        prelim_values = np.zeros(self.r_list.size)

        for i, r in tqdm(enumerate(self.r_list), desc="integration...", total=len(self.r_list)):
            if r > cutting_point:
                prelim_values[i] = 0.0
            else:
                prelim_values[i] = integrate.trapz(
                    circ_trapz(
                        circ_trapz(
                            self.weights[i],
                            self.theta2,
                            np.pi * 2
                        ),
                        self.theta1,
                        np.pi * 2
                    ) * np.sin(self.alpha) * np.cos(self.alpha),
                    self.alpha
                )
        return integrate.trapz(prelim_values * np.power(self.r_list, 3), self.r_list)

    def compute_inversion_error(self, n_steps):
        try:
            if self.db.attrs["mu"] != 0.0:
                raise NotImplementedError
        except:
            pass
        x, px, y, py = polar_to_cartesian(*np.meshgrid(
            self.r_list, self.alpha, self.theta1, self.theta2, indexing='ij'))
        x = x.flatten()
        px = px.flatten()
        y = y.flatten()
        py = py.flatten()
        engine = partial_track.generate_instance(x, px, y, py)
        engine.compute(n_steps, self.db.attrs["epsilon"])
        n_x, n_px, n_y, n_py, _ = engine.inverse_compute(
            n_steps, self.db.attrs["epsilon"])
        return (np.sqrt(np.power(x - n_x, 2) + np.power(px - n_px, 2) 
            + np.power(y - n_y, 2) + np.power(py - n_py, 2))).reshape((self.r_list.size, self.alpha.size, self.theta1.size, self.theta2.size))

    def compute_basic_lyapunov(self, n_steps, kick=1e-13):
        x, px, y, py = polar_to_cartesian(*np.meshgrid(
            self.r_list, self.alpha, self.theta1, self.theta2, indexing='ij'))
        x = x.flatten()
        px = px.flatten()
        y = y.flatten()
        py = py.flatten()
        alpha = np.random.rand(x.size) * np.pi / 2
        th1 = np.random.rand(x.size) * np.pi * 2
        th2 = np.random.rand(x.size) * np.pi * 2
        engine = partial_track.generate_instance(x, px, y, py)
        x_f, px_f, y_f, py_f, _ = engine.compute(n_steps, self.db.attrs["epsilon"])
        x += kick * np.cos(alpha) * np.cos(th1)
        px += kick * np.cos(alpha) * np.sin(th1)
        y += kick * np.sin(alpha) * np.cos(th2)
        py += kick * np.sin(alpha) * np.sin(th2)
        engine = partial_track.generate_instance(x, px, y, py)
        x_fk, px_fk, y_fk, py_fk, _ = engine.compute(
            n_steps, self.db.attrs["epsilon"])
        return (np.sqrt(np.power(x_f - x_fk, 2) + np.power(px_f - px_fk, 2) + np.power(y_f - y_fk, 2) + np.power(py_f - py_fk, 2)) / kick).reshape((self.r_list.size, self.alpha.size, self.theta1.size, self.theta2.size))

    def compute_multiple_lyapunov(self, n_steps, kick=1e-13):
        x, px, y, py = polar_to_cartesian(*np.meshgrid(
            self.r_list, self.alpha, self.theta1, self.theta2, indexing='ij'))
        x = x.flatten()
        px = px.flatten()
        y = y.flatten()
        py = py.flatten()
        engine = partial_track.generate_instance(x, px, y, py)
        x_f, px_f, y_f, py_f, _ = engine.compute(n_steps, self.db.attrs["epsilon"])
        # l_x
        engine = partial_track.generate_instance(x + kick, px, y, py)
        x_fk, px_fk, y_fk, py_fk, _ = engine.compute(
            n_steps, self.db.attrs["epsilon"])
        l_x = np.sqrt(np.power(x_f - x_fk, 2) + np.power(px_f - px_fk, 2) + np.power(y_f - y_fk, 2) + np.power(py_f - py_fk, 2)) / kick
        # l_px
        engine = partial_track.generate_instance(x, px + kick, y, py)
        x_fk, px_fk, y_fk, py_fk, _ = engine.compute(
            n_steps, self.db.attrs["epsilon"])
        l_px = np.sqrt(np.power(x_f - x_fk, 2) + np.power(px_f - px_fk, 2) + np.power(y_f - y_fk, 2) + np.power(py_f - py_fk, 2)) / kick
        # l_y
        engine = partial_track.generate_instance(x, px, y + kick, py)
        x_fk, px_fk, y_fk, py_fk, _ = engine.compute(
            n_steps, self.db.attrs["epsilon"])
        l_y = np.sqrt(np.power(x_f - x_fk, 2) + np.power(px_f - px_fk, 2) + np.power(y_f - y_fk, 2) + np.power(py_f - py_fk, 2)) / kick
        # l_py
        engine = partial_track.generate_instance(x, px, y, py + kick)
        x_fk, px_fk, y_fk, py_fk, _ = engine.compute(
            n_steps, self.db.attrs["epsilon"])
        l_py = np.sqrt(np.power(x_f - x_fk, 2) + np.power(px_f - px_fk, 2) + np.power(y_f - y_fk, 2) + np.power(py_f - py_fk, 2)) / kick
        return ((l_x + l_px + l_y + l_py) / 4).reshape((self.r_list.size, self.alpha.size, self.theta1.size, self.theta2.size))




def assign_symmetric_gaussian(sigma=1.0, polar=True):
    if polar:
        def f(r, a, th1, th2):
            return (
                np.exp(- 0.5 * np.power(r / sigma, 2))
            )
    else:
        def f(x, px, y, py):
            return(
                np.exp(-0.5 * (np.power(x / sigma, 2.0) + np.power(y / sigma,
                                                                   2.0) + np.power(py / sigma, 2.0) + np.power(px / sigma, 2.0)))
            )
    return f


def assign_uniform_distribution(polar=True):
    if polar:
        def f(r, a, th1, th2):
            return (
                np.ones_like(r)
            )
    else:
        def f(x, px, y, py):
            return (
                np.ones_like(x)
            )
    return f


def assign_generic_gaussian(sigma_x, sigma_px, sigma_y, sigma_py, polar=True):
    if polar:
        def f(r, a, th1, th2):
            x, px, y, py = polar_to_cartesian(r, a, th1, th2)
            x /= sigma_x
            px /= sigma_px
            y /= sigma_y
            py /= sigma_py
            r, a, th1, th2 = cartesian_to_polar(x, px, y, py)
            return (
                np.exp(- np.power(r, 2) * 0.5) / (np.power(2 * np.pi, 2))
            )
    else:
        assert False  # Needs to be implemented lol
    return f


def from_DA_to_loss_sym_gauss(sigma, DA_list, cut_point=None):
    loss_list = np.empty_like(DA_list)
    for i, DA in enumerate(DA_list):
        loss_list[i] = integrate.quad(
            lambda x: np.power(x, 3) * np.power(np.pi, 2) *
            2 * np.exp(- 0.5 * np.power(x / sigma, 2)),
            0.0, DA
        )[0]
    if not(cut_point is None):
        baseline = integrate.quad(
            lambda x: np.power(x, 3) * np.power(np.pi, 2) *
            2 * np.exp(- 0.5 * np.power(x / sigma, 2)),
            0.0, cut_point
        )[0]
        return loss_list / baseline
    else:
        return loss_list


def from_loss_to_DA_sym_gauss(sigma, loss_list, cut_point):
    baseline = integrate.quad(
        lambda x: np.power(x, 3) * np.power(np.pi, 2) *
        2 * np.exp(- 0.5 * np.power(x / sigma, 2)),
        0.0, cut_point
    )[0]
    def f(p):
        return integrate.quad(
            lambda x: np.power(x, 3) * np.power(np.pi, 2) *
            2 * np.exp(- 0.5 * np.power(x / sigma, 2)),
            0.0, p
        )[0] / baseline
    i_f = inversefunc(f, domain=((0.0, cut_point)))
    result = np.empty_like(loss_list)
    for i in range(len(loss_list)):
        result[i] = i_f(loss_list[i])
    return result
