# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00_GenerateSyntheticImages.ipynb (unless otherwise specified).

__all__ = ['radial_grid', 'IMAGE_SIZE', 'particle_bessel_response', 'background', 'add_noise', 'Particle',
           'TrainingImage']

# Cell
import numpy as np
from scipy.special import jv as bessel

# Cell
def radial_grid(image_size, particle_cx=0, particle_cy=0, particle_radius=1.5):
    """
    Returns array of size (image_size x image_size) where each value in the array
    is the radial distance  from that pixel to the particle, in units of the particle radius.

    image_size: tuple containing width x height in pixels
    particle_cx: horiz center of particle within image (range -1 (left edge) to 1 (right edge))
    particle_cy: vert center of particle within image (range -1 (top edge) to 1 (bottom edge))
    particle_radius: radius of particle (UNITS??)
    """

    # This is in the source code but not in the paper
    MYSTERY_CONST = 0.001

    h=image_size//2

    # particle position in pixels relative to centre of image
    icx = int(h + particle_cx * h)
    icy = int(h + particle_cy * h)

    # calculate radial distance for each pixel to particle centre
    imxs, imys = np.meshgrid( np.arange( 0, image_size), np.arange( 0, image_size))
    return (np.sqrt((imxs - icx)**2 + (imys - icy)**2 + MYSTERY_CONST**2)) / particle_radius


# Cell
IMAGE_SIZE = 51



# Cell
def particle_bessel_response(image_size, particle_cx, particle_cy, particle_radius, particle_intensity, order):
    """Reference to source here"""
    r=radial_grid(image_size, particle_cx, particle_cy, particle_radius)
    return (particle_intensity * 4 * order**2.5 * (bessel(order,r) / r)**2)

# Cell
def background(image_size, background_level = 0.5, gradi = 1, gradd=np.pi/4):
    h=image_size//2
    ix,iy = np.meshgrid(np.arange(-h, h + 1), np.arange(-h, h + 1))
    s=gradi/(np.sqrt(2)*image_size)
    g = s*((ix * np.sin(gradd))+(iy * np.cos(gradd)))
    b = np.ones((image_size,image_size)) * background_level
    return b+g

# Cell
def add_noise(img_without_noise, snr):
    flr=img_without_noise[:].min()
    img_without_noise-=flr
    img_with_noise = np.random.poisson(img_without_noise * snr**2) / snr**2
    return img_with_noise

# Cell
class Particle():
    def __init__(self, cx=0, cy=0, radius=1, bessel_order=1, intensity=1):
        self.cx = cx
        self.cy = cy
        self.radius = radius
        self.bessel_order = bessel_order
        self.intensity = intensity

    def __repr__(self):
        return f"cx={self.cx:.1f}, cy={self.cy:.1f}, r={self.radius:.1f}, order={self.bessel_order}, intensity={self.intensity:.1f}"


# Cell
from PIL import Image

class TrainingImage():
    def __init__(self, image_size, particles = [], background_level=1, gradient_i=1, gradient_d=0, snr=40):
        self.particles = particles
        self.background_level = background_level
        self.gradient_i = gradient_i
        self.gradient_d = gradient_d
        self.snr = snr
        self.image_size = image_size
        self.data = self._generate()

    def add_particle(self, particle, refresh=False):
        self.particles += particle
        if refresh:
            self.data = self._generate()

    def save(self,pth):
        img = (self.data-self.data.min())/self.data.max()
        if pth:
            Image.fromarray(np.uint8(img * 255) , 'L').save(pth)
        return img


    def get_particle_px(self):
        s=self.image_size//2
        return [(s+p.cy*s, s+p.cx*s) for p in self.particles]

    def _generate(self):
        img = background(self.image_size, self.background_level, self.gradient_i, self.gradient_d)
        for p in self.particles:
            img += particle_bessel_response(self.image_size, p.cx, p.cy, p.radius, p.intensity, p.bessel_order)
        img = add_noise(img,snr=self.snr)
        return img

    def show(self, fs=16):
        plt.imshow(self.data,interpolation='none',cmap='Greys_r')
        ax=plt.gca();
        ax.grid(False);
        plt.axis('off')

        # text
        ax.text(1.1,0.9, f"image_size = {image_size}x{image_size}",transform=ax.transAxes, fontsize=fs)
        ax.text(1.1,0.8, f"background intensity = {self.background_level:.1f}",transform=ax.transAxes, fontsize=fs)
        ax.text(1.1,0.7, f"gradient = ({self.gradient_i:.1f},{self.gradient_d:.1f})",transform=ax.transAxes, fontsize=fs)
        ax.text(1.1,0.6, f"SNR = {snr:.1f}",transform=ax.transAxes, fontsize=fs)
        ax.text(1.1,0.5, f"{len(self.particles)} particle(s):",transform=ax.transAxes, fontsize=fs)
        for i,p in enumerate(self.particles):
            ax.text(1.2,0.4-i*0.07, f"{i+1}: {p}",transform=ax.transAxes, fontsize=fs)


        # annotations
        h=self.image_size//2
        for i,p in enumerate(self.particles):
            x=plt.Circle((h+p.cx*h,h+p.cy*h),p.radius,lw=1,edgecolor='r',facecolor='none')
            ax.add_artist(x)


        return None



