
import math
import numpy as np
from cheshire.Grid import *


class InfiniteWell1DSampler(Grid1D):
    """
    Class that generates parameters for the 1D infinite well.
    """

    def sample_params(self, min_l=5, max_l=23, min_c=-8, max_c=8):
        """
        Random parameter method that generates parameters for infinite
        square wells.

        Args:
            **min_l (float)**: The minimum possible allowed length for the well size. Default is 5.
            **max_l (float)**: The maximum possible allowed length for the well size. Default is 23.
            **min_c (float)**: The minimum possible center point for the well. Default is -8.
            **max_c (float)**: The maximum possible center point for the well. Default is 8.

        Returns:
            A dictionary of randomly generated infinite well parameters.
        """

        assert min_l <= max_l
        assert min_c <= max_c

        # Rescale arguments based on the size of the grid
        min_l = self.rescale_by_size(min_l)
        max_l = self.rescale_by_size(max_l)
        min_c = self.rescale_by_size(min_c)
        max_c = self.rescale_by_size(max_c)

        l_x = np.random.uniform(low=min_l*self.x_max/20, high=max_l*self.x_max/20)
        c_x = np.random.uniform(low=min_c, high=max_c)

        return dict(l_x=l_x, c_x=c_x)


class InfiniteWell2DSampler(Grid2D):
    """
    Class that generates parameters for the 2D infinite well.
    """

    def sample_params(self, min_l=4, max_l=15, min_c=-8, max_c=8):
        """
        Random parameter method that generates parameters for 2D infinite square wells.

        Args:
            **min_l (float)**: The minimum possible allowed length for the well size. Default is 4.
            **max_l (float)**: The maximum possible allowed length for the well size. Default is 15.
            **min_c (float)**: The minimum possible center point for the well. Default is -8.
            **max_c (float)**: The maximum possible center point for the well. Default is 8.

        Returns:
            A dictionary of randomly generated infinite well parameters.
        """

        assert min_l <= max_l
        assert min_c <= max_c

        # Rescale arguments based on the size of the grid
        min_l = self.rescale_by_size(min_l)
        max_l = self.rescale_by_size(max_l)
        min_c = self.rescale_by_size(min_c)
        max_c = self.rescale_by_size(max_c)

        while True:
            l_x, l_y = self.__sample__(min_l=min_l, max_l=max_l)
            if (min_l <= l_x) & \
                    (l_x <= max_l) & \
                    (min_l <= l_y) & \
                    (l_y <= max_l):
                break

        switch = np.random.uniform(low=0, high=1)
        if switch > 0.5:
            temp = l_x
            l_x = l_y
            l_y = temp

        l_x = l_x
        l_y = l_y
        c_x = np.random.uniform(low=min_c, high=max_c)
        c_y = np.random.uniform(low=min_c, high=max_c)

        return dict(l_x=l_x, l_y=l_y, c_x=c_x, c_y=c_y)

    def __sample__(self, min_l, max_l):
        """
        Sample parameters. Sampling from the energy first ensures that the desired
        energy range is obtained
        """

        energy = np.random.uniform(low=0, high=0.4)
        l_x = np.random.uniform(low=min_l*self.x_max/20, high=max_l*self.x_max/20)

        if 2*energy/math.pi**2 <= 1/(l_x**2):
            l_y = np.nan
        elif 2*energy/math.pi**2 <= 1/(l_x**2) > max_l:
            l_y = np.nan
        elif 2*energy/math.pi**2 <= 1/(l_x**2) < min_l:
            l_y = np.nan
        else:
            l_y = 1 / np.sqrt((2*energy)/(math.pi**2) - 1/(l_x**2))
        return l_x, l_y


class SimpleHarmonicOscillator2DSampler(Grid2D):
    """
    Class that generates parameters for the 2D simple harmonic oscillator.
    """

    def sample_params(self, min_kx=0, max_kx=0.16, min_ky=0, max_ky=0.16, min_cx=-8, max_cx=8, min_cy=-8, max_cy=8):
        """
        Random parameter method that generates parameters for simple harmonic oscillator potentials.

        Args:
            **min_kx (float)**: The minimum possible value for the well width along the x axis. Default is 0.
            **max_kx (float)**: The maximum possible value for the well width along the x axis. Default is 0.16.
            **min_ky (float)**: The minimum possible value for the well width along the y axis. Default is 0.
            **max_ky (float)**: The maximum possible value for the well width along the y axis. Default is 0.16.
            **min_cx (float)**: The minimum possible value for the center (in a.u.) of the simple harmonic oscillator
                along the x axis. Default is -8.
            **max_cx (float)**: The maximum possible value for the center (in a.u.) of the simple harmonic oscillator
                along the x axis. Default is 8.
            **min_cy (float)**: The minimum possible value for the center (in a.u.) of the simple harmonic oscillator
                along the y axis. Default is -8.
            **max_cy (float)**: The maximum possible value for the center (in a.u.) of the simple harmonic oscillator
                along the y axis. Default is 8.

        Returns:
            A dictionary of randomly generated simple harmonic oscillator parameters.
        """

        assert min_kx <= max_kx
        assert min_ky <= max_ky
        assert min_cx <= max_cx
        assert min_cy <= max_cy

        # Rescale arguments
        min_kx = self.rescale_by_size(min_kx)
        max_kx = self.rescale_by_size(max_kx)
        min_ky = self.rescale_by_size(min_ky)
        max_ky = self.rescale_by_size(max_ky)
        min_cx = self.rescale_by_size(min_cx)
        max_cx = self.rescale_by_size(max_cx)
        min_cy = self.rescale_by_size(min_cy)
        max_cy = self.rescale_by_size(max_cy)

        # Sample parameters
        k_x = np.random.uniform(low=min_kx, high=max_kx)
        k_y = np.random.uniform(low=min_ky, high=max_ky)
        c_x = np.random.uniform(low=min_cx, high=max_cx)
        c_y = np.random.uniform(low=min_cy, high=max_cy)

        return dict(k_x=k_x, k_y=k_y, c_x=c_x, c_y=c_y)


class DoubleInvertedGaussian2DSampler(Grid2D):
    """
    Class that generates parameters for the 2D double inverted Gaussian.
    """

    def sample_params(self, min_a1=2, max_a1=4, min_a2=2, max_a2=4, min_cx1=-8, max_cx1=8, min_cy1=-8, max_cy1=8,
                      min_cx2=-8, max_cx2=8, min_cy2=-8, max_cy2=8, min_kx1=1.6, max_kx1=8, min_ky1=1.6, max_ky1=8,
                      min_kx2=1.6, max_kx2=8, min_ky2=1.6, max_ky2=8):
        """
        Random parameter method that generates parameters for double inverted Gaussian potentials.

        Args:
            **min_a1 (float)**: The minimum possible value of the amplitude of the first Gaussian. Default is 2. Units
                are in Hartree energy.
            **max_a1 (float)**: The maximum possible value of the amplitude of the first Gaussian. Default is 4. Units
                are in Hartree energy.
            **min_a2 (float)**: The minimum possible value of the amplitude of the second Gaussian. Default is 2. Units
                are in Hartree energy.
            **max_a2 (float)**: The maximum possible value of the amplitude of the second Gaussian. Default is 4. Units
                are in Hartree energy.
            **min_cx1 (float)**: The minimum possible position of the center along the x axis of the first Gaussian.
                Default is -8 a.u.
            **max_cx1 (float)**: The maximum possible position of the center along the x axis of the first Gaussian.
                Default is 8 a.u.
            **min_cy1 (float)**: The minimum possible position of the center along the y axis of the first Gaussian.
                Default is -8 a.u.
            **max_cy1 (float)**: The maximum possible position of the center along the y axis of the first Gaussian.
                Default is 8 a.u.
            **min_cx2 (float)**: The minimum possible position of the center along the x axis of the second Gaussian.
                Default is -8 a.u.
            **max_cx2 (float)**: The maximum possible position of the center along the x axis of the second Gaussian.
                Default is 8 a.u.
            **min_cy2 (float)**: The minimum possible position of the center along the y axis of the second Gaussian.
                Default is -8 a.u.
            **max_cy2 (float)**: The maximum possible position of the center along the y axis of the second Gaussian.
                Default is 8 a.u.
            **min_kx1 (float)**: The minimum possible value for the constant that determines the width along the x axis
                of the first Gaussian. Default is 1.6.
            **max_kx1 (float)**: The maximum possible value for the constant that determines the width along the x axis
                of the first Gaussian. Default is 8.
            **min_ky1 (float)**: The minimum possible value for the constant that determines the width along the y axis
                of the first Gaussian. Default is 1.6.
            **max_ky1 (float)**: The maximum possible value for the constant that determines the width along the y axis
                of the first Gaussian. Default is 8.
            **min_kx2 (float)**: The minimum possible value for the constant that determines the width along the x axis
                of the second Gaussian. Default is 1.6.
            **max_kx2 (float)**: The maximum possible value for the constant that determines the width along the x axis
                of the second Gaussian. Default is 8.
            **min_ky2 (float)**: The minimum possible value for the constant that determines the width along the y axis
                of the second Gaussian. Default is 1.6.
            **max_ky2 (float)**: The maximum possible value for the constant that determines the width along the y axis
                of the second Gaussian. Default is 8.

        Returns:
            A dictionary of randomly generated double inverted Gaussian parameters.
        """

        assert min_a1 > 0
        assert min_a2 > 0
        assert min_a1 <= max_a1
        assert min_a2 <= max_a2
        assert min_cx1 <= max_cx1
        assert min_cy1 <= max_cy1
        assert min_cx2 <= max_cx2
        assert min_cy2 <= max_cy2
        assert min_kx1 <= max_kx1
        assert min_ky1 <= max_ky1
        assert min_kx2 <= max_kx2
        assert min_ky2 <= max_ky2

        # Rescale arguments
        min_cx1 = self.rescale_by_size(min_cx1)
        min_cy1 = self.rescale_by_size(min_cy1)
        min_cx2 = self.rescale_by_size(min_cx2)
        min_cy2 = self.rescale_by_size(min_cy2)
        min_kx1 = self.rescale_by_size(min_kx1)
        min_ky1 = self.rescale_by_size(min_ky1)
        min_kx2 = self.rescale_by_size(min_kx2)
        min_ky2 = self.rescale_by_size(min_ky2)

        # Sample parameters
        a1 = np.random.uniform(low=min_a1, high=max_a1)
        a2 = np.random.uniform(low=min_a2, high=max_a2)
        c_x1 = np.random.uniform(low=min_cx1, high=max_cx1)
        c_y1 = np.random.uniform(low=min_cy1, high=max_cy1)
        c_x2 = np.random.uniform(low=min_cx2, high=max_cx2)
        c_y2 = np.random.uniform(low=min_cy2, high=max_cy2)
        k_x1 = np.random.uniform(low=min_kx1, high=max_kx1)
        k_y1 = np.random.uniform(low=min_ky1, high=max_ky1)
        k_x2 = np.random.uniform(low=min_kx2, high=max_kx2)
        k_y2 = np.random.uniform(low=min_ky2, high=max_ky2)

        return dict(a1=a1, a2=a2, c_x1=c_x1, c_y1=c_y1, c_x2=c_x2, c_y2=c_y2, k_x1=k_x1, k_y1=k_y1, k_x2=k_x2,
                    k_y2=k_y2)


class Random2DSampler(Grid2D):
    """
    Class that generates parameters for a random 2D potential.
    """

    def sample_params(self, min_k=2, max_k=7, min_r=80, max_r=180, min_sig1=6, max_sig1=10, min_sig2=10, max_sig2=16,
                      p_range=[0.5, 1.0, 1.5, 2.0]):
        """
        Random parameter method that generates parameters for random potentials.

        Args:
            **min_k (int)**: Determines the low end of the number of points used to create the convex hull. Default is
                2.
            **max_k (int)**: Determines the high end of the number of points used to create the convex hull. Default is
                7.
            **min_r (float)**: The minimum possible resolution size of the random blob. Default is 80 pixels (scaled
                for 256 x 256).
            **max_r (float)**: The maximum possible resolution size of the random blob. Default is 180 pixels (scaled
                for 256 x 256).
            **min_sig1 (float)**: The minimum possible variance of the first Gaussian blur. Default is 6.
            **max_sig1 (float)**: The maximum possible variance of the first Gaussian blur. Default is 10.
            **min_sig2 (float)**: The minimum possible variance of the second Gaussian blur. Default is 10.
            **max_sig2 (float)**: The maximum possible variance of the second Gaussian blur. Default is 16.
            **p_range (list)**: The range of exponents to sample from when exponentiating the potential for contrasting.
                Default is [1, 2].

        Returns:
            A dictionary of randomly generated random potential parameters.
        """

        assert min_k <= max_k
        assert min_r <= max_r
        assert min_sig1 <= max_sig1
        assert min_sig1 > 0
        assert min_sig2 <= max_sig2
        assert min_sig2 > 0
        assert all([isinstance(i, float) for i in p_range])

        # Rescale arguments
        min_r = self.rescale_by_grid(min_r)
        max_r = self.rescale_by_grid(max_r)
        min_sig1 = self.rescale_by_grid(min_sig1)
        max_sig1 = self.rescale_by_grid(max_sig1)
        min_sig2 = self.rescale_by_grid(min_sig2)
        max_sig2 = self.rescale_by_grid(max_sig2)

        # Sample parameters
        k = np.random.randint(low=min_k, high=max_k+1)
        r = np.random.uniform(low=min_r, high=max_r)
        sig1 = np.random.uniform(low=min_sig1, high=max_sig1)
        sig2 = np.random.uniform(low=min_sig2, high=max_sig2)
        p = np.random.choice(p_range)

        return dict(k=k, r=r, sig1=sig1, sig2=sig2, p=p)


class Coulomb2DSampler(Grid2D):
    """
    Class that generates parameters for the 2D Coulomb potential.
    """

    def sample_params(self, min_z=1, max_z=118, min_cx=-8, max_cx=8, min_cy=-8, max_cy=8, alpha=10**-9):
        """
        Random parameter method that generates parameters for Coulomb potentials.

        Args:
            **min_z (int)**: The minimum number of allowable protons.
            **max_z (int)**: The maximum number of allowable protons.
            **min_cx (float)**: The minimum possible value for the center (in a.u.) of the Coulomb potential along the x
                axis. Default is -8.
            **max_cx (float)**: The maximum possible value for the center (in a.u.) of the Coulomb potential along the x
                axis. Default is 8.
            **min_cy (float)**: The minimum possible value for the center (in a.u.) of the Coulomb potential along the y
                axis. Default is -8.
            **max_cy (float)**: The maximum possible value for the center (in a.u.) of the Coulomb potential along the y
                axis. Default is 8.

        Returns:
            A dictionary of randomly generated Coulomb potential parameters.
        """

        assert min_z <= max_z
        assert min_z >= 1

        # Sample parameters
        z = np.random.randint(low=min_z, high=max_z)
        c_x = np.random.uniform(low=min_cx, high=max_cx)
        c_y = np.random.uniform(low=min_cy, high=max_cy)
        alpha = alpha

        return dict(z=z, c_x=c_x, c_y=c_y, alpha=alpha)


import math
import numpy as np
from cheshire.Grid import *


class InfiniteWell1DSampler(Grid1D):
    """
    Class that generates parameters for the 1D infinite well.
    """

    def sample_params(self, min_l=5, max_l=23, min_c=-8, max_c=8):
        """
        Random parameter method that generates parameters for infinite
        square wells.

        Args:
            **min_l (float)**: The minimum possible allowed length for the well size. Default is 5.
            **max_l (float)**: The maximum possible allowed length for the well size. Default is 23.
            **min_c (float)**: The minimum possible center point for the well. Default is -8.
            **max_c (float)**: The maximum possible center point for the well. Default is 8.

        Returns:
            A dictionary of randomly generated infinite well parameters.
        """

        assert min_l <= max_l
        assert min_c <= max_c

        # Rescale arguments based on the size of the grid
        min_l = self.rescale_by_size(min_l)
        max_l = self.rescale_by_size(max_l)
        min_c = self.rescale_by_size(min_c)
        max_c = self.rescale_by_size(max_c)

        l_x = np.random.uniform(low=min_l*self.x_max/20, high=max_l*self.x_max/20)
        c_x = np.random.uniform(low=min_c, high=max_c)

        return dict(l_x=l_x, c_x=c_x)


class InfiniteWell2DSampler(Grid2D):
    """
    Class that generates parameters for the 2D infinite well.
    """

    def sample_params(self, min_l=4, max_l=15, min_c=-8, max_c=8):
        """
        Random parameter method that generates parameters for 2D infinite square wells.

        Args:
            **min_l (float)**: The minimum possible allowed length for the well size. Default is 4.
            **max_l (float)**: The maximum possible allowed length for the well size. Default is 15.
            **min_c (float)**: The minimum possible center point for the well. Default is -8.
            **max_c (float)**: The maximum possible center point for the well. Default is 8.

        Returns:
            A dictionary of randomly generated infinite well parameters.
        """

        assert min_l <= max_l
        assert min_c <= max_c

        # Rescale arguments based on the size of the grid
        min_l = self.rescale_by_size(min_l)
        max_l = self.rescale_by_size(max_l)
        min_c = self.rescale_by_size(min_c)
        max_c = self.rescale_by_size(max_c)

        while True:
            l_x, l_y = self.__sample__(min_l=min_l, max_l=max_l)
            if (min_l <= l_x) & \
                    (l_x <= max_l) & \
                    (min_l <= l_y) & \
                    (l_y <= max_l):
                break

        switch = np.random.uniform(low=0, high=1)
        if switch > 0.5:
            temp = l_x
            l_x = l_y
            l_y = temp

        l_x = l_x
        l_y = l_y
        c_x = np.random.uniform(low=min_c, high=max_c)
        c_y = np.random.uniform(low=min_c, high=max_c)

        return dict(l_x=l_x, l_y=l_y, c_x=c_x, c_y=c_y)

    def __sample__(self, min_l, max_l):
        """
        Sample parameters. Sampling from the energy first ensures that the desired
        energy range is obtained
        """

        energy = np.random.uniform(low=0, high=0.4)
        l_x = np.random.uniform(low=min_l*self.x_max/20, high=max_l*self.x_max/20)

        if 2*energy/math.pi**2 <= 1/(l_x**2):
            l_y = np.nan
        elif 2*energy/math.pi**2 <= 1/(l_x**2) > max_l:
            l_y = np.nan
        elif 2*energy/math.pi**2 <= 1/(l_x**2) < min_l:
            l_y = np.nan
        else:
            l_y = 1 / np.sqrt((2*energy)/(math.pi**2) - 1/(l_x**2))
        return l_x, l_y


class SimpleHarmonicOscillator2DSampler(Grid2D):
    """
    Class that generates parameters for the 2D simple harmonic oscillator.
    """

    def sample_params(self, min_kx=0, max_kx=0.16, min_ky=0, max_ky=0.16, min_cx=-8, max_cx=8, min_cy=-8, max_cy=8):
        """
        Random parameter method that generates parameters for simple harmonic oscillator potentials.

        Args:
            **min_kx (float)**: The minimum possible value for the well width along the x axis. Default is 0.
            **max_kx (float)**: The maximum possible value for the well width along the x axis. Default is 0.16.
            **min_ky (float)**: The minimum possible value for the well width along the y axis. Default is 0.
            **max_ky (float)**: The maximum possible value for the well width along the y axis. Default is 0.16.
            **min_cx (float)**: The minimum possible value for the center (in a.u.) of the simple harmonic oscillator
                along the x axis. Default is -8.
            **max_cx (float)**: The maximum possible value for the center (in a.u.) of the simple harmonic oscillator
                along the x axis. Default is 8.
            **min_cy (float)**: The minimum possible value for the center (in a.u.) of the simple harmonic oscillator
                along the y axis. Default is -8.
            **max_cy (float)**: The maximum possible value for the center (in a.u.) of the simple harmonic oscillator
                along the y axis. Default is 8.

        Returns:
            A dictionary of randomly generated simple harmonic oscillator parameters.
        """

        assert min_kx <= max_kx
        assert min_ky <= max_ky
        assert min_cx <= max_cx
        assert min_cy <= max_cy

        # Rescale arguments
        min_kx = self.rescale_by_size(min_kx)
        max_kx = self.rescale_by_size(max_kx)
        min_ky = self.rescale_by_size(min_ky)
        max_ky = self.rescale_by_size(max_ky)
        min_cx = self.rescale_by_size(min_cx)
        max_cx = self.rescale_by_size(max_cx)
        min_cy = self.rescale_by_size(min_cy)
        max_cy = self.rescale_by_size(max_cy)

        # Sample parameters
        k_x = np.random.uniform(low=min_kx, high=max_kx)
        k_y = np.random.uniform(low=min_ky, high=max_ky)
        c_x = np.random.uniform(low=min_cx, high=max_cx)
        c_y = np.random.uniform(low=min_cy, high=max_cy)

        return dict(k_x=k_x, k_y=k_y, c_x=c_x, c_y=c_y)


class DoubleInvertedGaussian2DSampler(Grid2D):
    """
    Class that generates parameters for the 2D double inverted Gaussian.
    """

    def sample_params(self, min_a1=2, max_a1=4, min_a2=2, max_a2=4, min_cx1=-8, max_cx1=8, min_cy1=-8, max_cy1=8,
                      min_cx2=-8, max_cx2=8, min_cy2=-8, max_cy2=8, min_kx1=1.6, max_kx1=8, min_ky1=1.6, max_ky1=8,
                      min_kx2=1.6, max_kx2=8, min_ky2=1.6, max_ky2=8):
        """
        Random parameter method that generates parameters for double inverted Gaussian potentials.

        Args:
            **min_a1 (float)**: The minimum possible value of the amplitude of the first Gaussian. Default is 2. Units
                are in Hartree energy.
            **max_a1 (float)**: The maximum possible value of the amplitude of the first Gaussian. Default is 4. Units
                are in Hartree energy.
            **min_a2 (float)**: The minimum possible value of the amplitude of the second Gaussian. Default is 2. Units
                are in Hartree energy.
            **max_a2 (float)**: The maximum possible value of the amplitude of the second Gaussian. Default is 4. Units
                are in Hartree energy.
            **min_cx1 (float)**: The minimum possible position of the center along the x axis of the first Gaussian.
                Default is -8 a.u.
            **max_cx1 (float)**: The maximum possible position of the center along the x axis of the first Gaussian.
                Default is 8 a.u.
            **min_cy1 (float)**: The minimum possible position of the center along the y axis of the first Gaussian.
                Default is -8 a.u.
            **max_cy1 (float)**: The maximum possible position of the center along the y axis of the first Gaussian.
                Default is 8 a.u.
            **min_cx2 (float)**: The minimum possible position of the center along the x axis of the second Gaussian.
                Default is -8 a.u.
            **max_cx2 (float)**: The maximum possible position of the center along the x axis of the second Gaussian.
                Default is 8 a.u.
            **min_cy2 (float)**: The minimum possible position of the center along the y axis of the second Gaussian.
                Default is -8 a.u.
            **max_cy2 (float)**: The maximum possible position of the center along the y axis of the second Gaussian.
                Default is 8 a.u.
            **min_kx1 (float)**: The minimum possible value for the constant that determines the width along the x axis
                of the first Gaussian. Default is 1.6.
            **max_kx1 (float)**: The maximum possible value for the constant that determines the width along the x axis
                of the first Gaussian. Default is 8.
            **min_ky1 (float)**: The minimum possible value for the constant that determines the width along the y axis
                of the first Gaussian. Default is 1.6.
            **max_ky1 (float)**: The maximum possible value for the constant that determines the width along the y axis
                of the first Gaussian. Default is 8.
            **min_kx2 (float)**: The minimum possible value for the constant that determines the width along the x axis
                of the second Gaussian. Default is 1.6.
            **max_kx2 (float)**: The maximum possible value for the constant that determines the width along the x axis
                of the second Gaussian. Default is 8.
            **min_ky2 (float)**: The minimum possible value for the constant that determines the width along the y axis
                of the second Gaussian. Default is 1.6.
            **max_ky2 (float)**: The maximum possible value for the constant that determines the width along the y axis
                of the second Gaussian. Default is 8.

        Returns:
            A dictionary of randomly generated double inverted Gaussian parameters.
        """

        assert min_a1 > 0
        assert min_a2 > 0
        assert min_a1 <= max_a1
        assert min_a2 <= max_a2
        assert min_cx1 <= max_cx1
        assert min_cy1 <= max_cy1
        assert min_cx2 <= max_cx2
        assert min_cy2 <= max_cy2
        assert min_kx1 <= max_kx1
        assert min_ky1 <= max_ky1
        assert min_kx2 <= max_kx2
        assert min_ky2 <= max_ky2

        # Rescale arguments
        min_cx1 = self.rescale_by_size(min_cx1)
        min_cy1 = self.rescale_by_size(min_cy1)
        min_cx2 = self.rescale_by_size(min_cx2)
        min_cy2 = self.rescale_by_size(min_cy2)
        min_kx1 = self.rescale_by_size(min_kx1)
        min_ky1 = self.rescale_by_size(min_ky1)
        min_kx2 = self.rescale_by_size(min_kx2)
        min_ky2 = self.rescale_by_size(min_ky2)

        # Sample parameters
        a1 = np.random.uniform(low=min_a1, high=max_a1)
        a2 = np.random.uniform(low=min_a2, high=max_a2)
        c_x1 = np.random.uniform(low=min_cx1, high=max_cx1)
        c_y1 = np.random.uniform(low=min_cy1, high=max_cy1)
        c_x2 = np.random.uniform(low=min_cx2, high=max_cx2)
        c_y2 = np.random.uniform(low=min_cy2, high=max_cy2)
        k_x1 = np.random.uniform(low=min_kx1, high=max_kx1)
        k_y1 = np.random.uniform(low=min_ky1, high=max_ky1)
        k_x2 = np.random.uniform(low=min_kx2, high=max_kx2)
        k_y2 = np.random.uniform(low=min_ky2, high=max_ky2)

        return dict(a1=a1, a2=a2, c_x1=c_x1, c_y1=c_y1, c_x2=c_x2, c_y2=c_y2, k_x1=k_x1, k_y1=k_y1, k_x2=k_x2,
                    k_y2=k_y2)


class Random2DSampler(Grid2D):
    """
    Class that generates parameters for a random 2D potential.
    """

    def sample_params(self, min_k=2, max_k=7, min_r=80, max_r=180, min_sig1=6, max_sig1=10, min_sig2=10, max_sig2=16,
                      p_range=[0.5, 1.0, 1.5, 2.0]):
        """
        Random parameter method that generates parameters for random potentials.

        Args:
            **min_k (int)**: Determines the low end of the number of points used to create the convex hull. Default is
                2.
            **max_k (int)**: Determines the high end of the number of points used to create the convex hull. Default is
                7.
            **min_r (float)**: The minimum possible resolution size of the random blob. Default is 80 pixels (scaled
                for 256 x 256).
            **max_r (float)**: The maximum possible resolution size of the random blob. Default is 180 pixels (scaled
                for 256 x 256).
            **min_sig1 (float)**: The minimum possible variance of the first Gaussian blur. Default is 6.
            **max_sig1 (float)**: The maximum possible variance of the first Gaussian blur. Default is 10.
            **min_sig2 (float)**: The minimum possible variance of the second Gaussian blur. Default is 10.
            **max_sig2 (float)**: The maximum possible variance of the second Gaussian blur. Default is 16.
            **p_range (list)**: The range of exponents to sample from when exponentiating the potential for contrasting.
                Default is [1, 2].

        Returns:
            A dictionary of randomly generated random potential parameters.
        """

        assert min_k <= max_k
        assert min_r <= max_r
        assert min_sig1 <= max_sig1
        assert min_sig1 > 0
        assert min_sig2 <= max_sig2
        assert min_sig2 > 0
        assert all([isinstance(i, float) for i in p_range])

        # Rescale arguments
        min_r = self.rescale_by_grid(min_r)
        max_r = self.rescale_by_grid(max_r)
        min_sig1 = self.rescale_by_grid(min_sig1)
        max_sig1 = self.rescale_by_grid(max_sig1)
        min_sig2 = self.rescale_by_grid(min_sig2)
        max_sig2 = self.rescale_by_grid(max_sig2)

        # Sample parameters
        k = np.random.randint(low=min_k, high=max_k+1)
        r = np.random.uniform(low=min_r, high=max_r)
        sig1 = np.random.uniform(low=min_sig1, high=max_sig1)
        sig2 = np.random.uniform(low=min_sig2, high=max_sig2)
        p = np.random.choice(p_range)

        return dict(k=k, r=r, sig1=sig1, sig2=sig2, p=p)


class Coulomb2DSampler(Grid2D):
    """
    Class that generates parameters for the 2D Coulomb potential.
    """

    def sample_params(self, min_z=1, max_z=118, min_cx=-8, max_cx=8, min_cy=-8, max_cy=8, alpha=10**-9):
        """
        Random parameter method that generates parameters for Coulomb potentials.

        Args:
            **min_z (int)**: The minimum number of allowable protons.
            **max_z (int)**: The maximum number of allowable protons.
            **min_cx (float)**: The minimum possible value for the center (in a.u.) of the Coulomb potential along the x
                axis. Default is -8.
            **max_cx (float)**: The maximum possible value for the center (in a.u.) of the Coulomb potential along the x
                axis. Default is 8.
            **min_cy (float)**: The minimum possible value for the center (in a.u.) of the Coulomb potential along the y
                axis. Default is -8.
            **max_cy (float)**: The maximum possible value for the center (in a.u.) of the Coulomb potential along the y
                axis. Default is 8.

        Returns:
            A dictionary of randomly generated Coulomb potential parameters.
        """

        assert min_z <= max_z
        assert min_z >= 1

        # Sample parameters
        z = np.random.randint(low=min_z, high=max_z)
        c_x = np.random.uniform(low=min_cx, high=max_cx)
        c_y = np.random.uniform(low=min_cy, high=max_cy)
        alpha = alpha

        return dict(z=z, c_x=c_x, c_y=c_y, alpha=alpha)

Class that generates parameters for the 1D infinite well. Random parameter method that generates parameters for infinite         square wells.          Args:             **min_l (float)**: The minimum possible allowed length for the well size. Default is 5.             **max_l (float)**: The maximum possible allowed length for the well size. Default is 23.             **min_c (float)**: The minimum possible center point for the well. Default is -8.             **max_c (float)**: The maximum possible center point for the well. Default is 8.          Returns:             A dictionary of randomly generated infinite well parameters. # Rescale arguments based on the size of the grid Class that generates parameters for the 2D infinite well. Random parameter method that generates parameters for 2D infinite square wells.          Args:             **min_l (float)**: The minimum possible allowed length for the well size. Default is 4.             **max_l (float)**: The maximum possible allowed length for the well size. Default is 15.             **min_c (float)**: The minimum possible center point for the well. Default is -8.             **max_c (float)**: The maximum possible center point for the well. Default is 8.          Returns:             A dictionary of randomly generated infinite well parameters. # Rescale arguments based on the size of the grid Sample parameters. Sampling from the energy first ensures that the desired         energy range is obtained Class that generates parameters for the 2D simple harmonic oscillator. Random parameter method that generates parameters for simple harmonic oscillator potentials.          Args:             **min_kx (float)**: The minimum possible value for the well width along the x axis. Default is 0.             **max_kx (float)**: The maximum possible value for the well width along the x axis. Default is 0.16.             **min_ky (float)**: The minimum possible value for the well width along the y axis. Default is 0.             **max_ky (float)**: The maximum possible value for the well width along the y axis. Default is 0.16.             **min_cx (float)**: The minimum possible value for the center (in a.u.) of the simple harmonic oscillator                 along the x axis. Default is -8.             **max_cx (float)**: The maximum possible value for the center (in a.u.) of the simple harmonic oscillator                 along the x axis. Default is 8.             **min_cy (float)**: The minimum possible value for the center (in a.u.) of the simple harmonic oscillator                 along the y axis. Default is -8.             **max_cy (float)**: The maximum possible value for the center (in a.u.) of the simple harmonic oscillator                 along the y axis. Default is 8.          Returns:             A dictionary of randomly generated simple harmonic oscillator parameters. # Rescale arguments # Sample parameters Class that generates parameters for the 2D double inverted Gaussian. Random parameter method that generates parameters for double inverted Gaussian potentials.          Args:             **min_a1 (float)**: The minimum possible value of the amplitude of the first Gaussian. Default is 2. Units                 are in Hartree energy.             **max_a1 (float)**: The maximum possible value of the amplitude of the first Gaussian. Default is 4. Units                 are in Hartree energy.             **min_a2 (float)**: The minimum possible value of the amplitude of the second Gaussian. Default is 2. Units                 are in Hartree energy.             **max_a2 (float)**: The maximum possible value of the amplitude of the second Gaussian. Default is 4. Units                 are in Hartree energy.             **min_cx1 (float)**: The minimum possible position of the center along the x axis of the first Gaussian.                 Default is -8 a.u.             **max_cx1 (float)**: The maximum possible position of the center along the x axis of the first Gaussian.                 Default is 8 a.u.             **min_cy1 (float)**: The minimum possible position of the center along the y axis of the first Gaussian.                 Default is -8 a.u.             **max_cy1 (float)**: The maximum possible position of the center along the y axis of the first Gaussian.                 Default is 8 a.u.             **min_cx2 (float)**: The minimum possible position of the center along the x axis of the second Gaussian.                 Default is -8 a.u.             **max_cx2 (float)**: The maximum possible position of the center along the x axis of the second Gaussian.                 Default is 8 a.u.             **min_cy2 (float)**: The minimum possible position of the center along the y axis of the second Gaussian.                 Default is -8 a.u.             **max_cy2 (float)**: The maximum possible position of the center along the y axis of the second Gaussian.                 Default is 8 a.u.             **min_kx1 (float)**: The minimum possible value for the constant that determines the width along the x axis                 of the first Gaussian. Default is 1.6.             **max_kx1 (float)**: The maximum possible value for the constant that determines the width along the x axis                 of the first Gaussian. Default is 8.             **min_ky1 (float)**: The minimum possible value for the constant that determines the width along the y axis                 of the first Gaussian. Default is 1.6.             **max_ky1 (float)**: The maximum possible value for the constant that determines the width along the y axis                 of the first Gaussian. Default is 8.             **min_kx2 (float)**: The minimum possible value for the constant that determines the width along the x axis                 of the second Gaussian. Default is 1.6.             **max_kx2 (float)**: The maximum possible value for the constant that determines the width along the x axis                 of the second Gaussian. Default is 8.             **min_ky2 (float)**: The minimum possible value for the constant that determines the width along the y axis                 of the second Gaussian. Default is 1.6.             **max_ky2 (float)**: The maximum possible value for the constant that determines the width along the y axis                 of the second Gaussian. Default is 8.          Returns:             A dictionary of randomly generated double inverted Gaussian parameters. # Rescale arguments # Sample parameters Class that generates parameters for a random 2D potential. Random parameter method that generates parameters for random potentials.          Args:             **min_k (int)**: Determines the low end of the number of points used to create the convex hull. Default is                 2.             **max_k (int)**: Determines the high end of the number of points used to create the convex hull. Default is                 7.             **min_r (float)**: The minimum possible resolution size of the random blob. Default is 80 pixels (scaled                 for 256 x 256).             **max_r (float)**: The maximum possible resolution size of the random blob. Default is 180 pixels (scaled                 for 256 x 256).             **min_sig1 (float)**: The minimum possible variance of the first Gaussian blur. Default is 6.             **max_sig1 (float)**: The maximum possible variance of the first Gaussian blur. Default is 10.             **min_sig2 (float)**: The minimum possible variance of the second Gaussian blur. Default is 10.             **max_sig2 (float)**: The maximum possible variance of the second Gaussian blur. Default is 16.             **p_range (list)**: The range of exponents to sample from when exponentiating the potential for contrasting.                 Default is [1, 2].          Returns:             A dictionary of randomly generated random potential parameters. # Rescale arguments # Sample parameters Class that generates parameters for the 2D Coulomb potential. Random parameter method that generates parameters for Coulomb potentials.          Args:             **min_z (int)**: The minimum number of allowable protons.             **max_z (int)**: The maximum number of allowable protons.             **min_cx (float)**: The minimum possible value for the center (in a.u.) of the Coulomb potential along the x                 axis. Default is -8.             **max_cx (float)**: The maximum possible value for the center (in a.u.) of the Coulomb potential along the x                 axis. Default is 8.             **min_cy (float)**: The minimum possible value for the center (in a.u.) of the Coulomb potential along the y                 axis. Default is -8.             **max_cy (float)**: The maximum possible value for the center (in a.u.) of the Coulomb potential along the y                 axis. Default is 8.          Returns:             A dictionary of randomly generated Coulomb potential parameters. # Sample parameters