from colicoords.synthetic_data import add_readout_noise, draw_poisson
from colicoords import load
import numpy as np
import mahotas as mh
from tqdm import tqdm
import os
import tifffile


def chunk_list(l, sizes):
    prev = 0
    for s in sizes:
        result = l[prev:prev+s]
        prev += s
        yield result


def generate_images(cell_list, num_images, cell_per_img, cell_per_img_std, shape):
    nums = np.round(np.random.normal(cell_per_img, cell_per_img_std, num_images)).astype(int)
    nums = nums[nums > 0]

    assert sum(nums) < len(cell_list), 'Not enough cells'
    chunked = [chunk for chunk in tqdm(chunk_list(cell_list, nums))]
    dicts = [generate_image(cells, shape) for cells in tqdm(chunked)]
    out_dict = {}
    for i, d in enumerate(dicts):
        for k, v in d.items():
            if 'storm' in k:
                v['frame'] = i + 1
                if k in out_dict:
                    out_dict[k] = np.append(out_dict[k], v)
                else:
                    out_dict[k] = v
            else:
                if k in out_dict:
                    out_dict[k][i] = v
                else:
                    out_dict[k] = np.zeros((num_images, *shape))
                    out_dict[k][i] = v

    return out_dict


def generate_image(cells, shape, max_dist=5):
    thetas = 360 * np.random.rand(len(cells))
    data_list = [cell.data.rotate(theta) for cell, theta in zip(cells, thetas)]
    assert all([data.names == data_list[0].names for data in data_list]), 'All cells must have the same data elements'
    out_dict = {name: np.zeros(shape) for name, dclass in zip(data_list[0].names, data_list[0].dclasses) if dclass != 'storm'}
    for i, data in enumerate(data_list):
        valid_position = False
        while not valid_position:
            pos_x = int(np.round(shape[1] * np.random.rand()))
            pos_y = int(np.round(shape[0] * np.random.rand()))

            min1 = pos_y - int(np.floor(data.shape[0]))
            max1 = min1 + data.shape[0]

            min2 = pos_x - int(np.floor(data.shape[1]))
            max2 = min2 + data.shape[1]

            # Crop the data for when the cell is on the border of the image
            d_min1 = np.max([0 - min1, 0])
            d_max1 = np.min([data.shape[0] + (shape[0] - pos_y), data.shape[0]])

            d_min2 = np.max([0 - min2, 0])
            d_max2 = np.min([data.shape[1] + (shape[1] - pos_x), data.shape[1]])

            data_cropped = data[d_min1:d_max1, d_min2:d_max2]

            # Limit image position to the edges of the image
            min1 = np.max([min1, 0])
            max1 = np.min([max1, shape[0]])
            min2 = np.max([min2, 0])
            max2 = np.min([max2, shape[1]])

            temp_binary = np.zeros(shape)
            temp_binary[min1:max1, min2:max2] = data_cropped.binary_img
            out_binary = (out_dict['binary'] > 0).astype(int)
            distance_map = mh.distance(1 - out_binary, metric='euclidean')

            if np.any(distance_map[temp_binary.astype(bool)] < max_dist):
                continue

            valid_position = True

        for name in data.names:
            data_elem = data_cropped.data_dict[name]
            if data_elem.dclass == 'storm':
                data_elem['x'] += min2
                data_elem['y'] += min1

                xmax, ymax = shape[1], shape[0]
                bools = (data_elem['x'] < 0) + (data_elem['x'] > xmax) + (data_elem['y'] < 0) + (data_elem['y'] > ymax)
                data_out = data_elem[~bools].copy()
                if name in out_dict:
                    out_dict[name] = np.append(out_dict[name], data_out)
                else:
                    out_dict[name] = data_out

                continue
            elif data_elem.dclass == 'binary':
                out_dict[name][min1:max1, min2:max2] += ((i+1)*data_elem)
            else:
                out_dict[name][min1:max1, min2:max2] += data_elem

    return out_dict


def gen_image_from_storm(storm_table, shape, sigma=1.54, sigma_std=0.3):
    xmax = shape[1]
    ymax = shape[0]
    step = 1
    xi = np.arange(step / 2, xmax, step)
    yi = np.arange(step / 2, ymax, step)

    x_coords = np.repeat(xi, len(yi)).reshape(len(xi), len(yi)).T
    y_coords = np.repeat(yi, len(xi)).reshape(len(yi), len(xi))
    x, y = storm_table['x'], storm_table['y']

    img = np.zeros_like(x_coords)
    intensities = storm_table['intensity']
    sigma = sigma * np.ones_like(x) if not sigma_std else np.random.normal(sigma, sigma_std, size=len(x))
    for _sigma, _int, _x, _y in zip(sigma, intensities, x, y):
        img += _int * np.exp(-(((_x - x_coords) / _sigma) ** 2 + ((_y - y_coords) / _sigma) ** 2) / 2)

    return img


def gen_im(data_dir):
    """Generate microscopy images from a list of cell objects by placing them randomly oriented in the image."""
    cell_list = load(os.path.join(data_dir, 'cell_obj', 'cells_final_selected.hdf5'))

    out_dict = generate_images(cell_list, 1000, 10, 3, (512, 512))

    if not os.path.exists(os.path.join(data_dir, 'images')):
        os.mkdir(os.path.join(data_dir, 'images'))

    np.save(os.path.join(data_dir, 'images', 'binary.npy'), out_dict['binary'])
    np.save(os.path.join(data_dir, 'images', 'brightfield.npy'), out_dict['brightfield'])
    np.save(os.path.join(data_dir, 'images', 'foci_inner.npy'), out_dict['foci_inner'])
    np.save(os.path.join(data_dir, 'images', 'foci_outer.npy'), out_dict['foci_outer'])
    np.save(os.path.join(data_dir, 'images', 'storm_inner.npy'), out_dict['storm_inner'])
    np.save(os.path.join(data_dir, 'images', 'storm_outer.npy'), out_dict['storm_outer'])

    tifffile.imsave(os.path.join(data_dir, 'images', 'binary.tif'), out_dict['binary'])
    tifffile.imsave(os.path.join(data_dir, 'images', 'brightfield.tif'), out_dict['brightfield'])
    tifffile.imsave(os.path.join(data_dir, 'images', 'foci_inner.tif'), out_dict['foci_inner'])
    tifffile.imsave(os.path.join(data_dir, 'images', 'foci_outer.tif'), out_dict['foci_outer'])
    np.savetxt(os.path.join(data_dir, 'images', 'storm_inner.txt'), out_dict['storm_inner'])
    np.savetxt(os.path.join(data_dir, 'images', 'storm_outer.txt'), out_dict['storm_inner'])


def noise_bf(data_dir):
    """add poissonian and readout noise to brightfield images"""
    noise = 20
    img_stack = np.load(os.path.join(data_dir, 'images', 'brightfield.npy'))
    for photons in [10000, 1000, 500]:
        ratio = 1.0453  # ratio between 'background' (no cells) and cell wall
        img = (photons*(ratio-1))*img_stack + photons
        img = draw_poisson(img)
        img = add_readout_noise(img, noise)
        tifffile.imsave(os.path.join(data_dir, 'images', 'bf_noise_{}_photons.tif'.format(photons)), img)
        np.save(os.path.join(data_dir, 'images', 'bf_noise_{}_photons.npy'.format(photons)), img)


if __name__ == '__main__':
    np.random.seed(42)
    data_dir = r'.'
    if not os.path.exists(os.path.join(data_dir, 'images')):
        os.mkdir(os.path.join(data_dir, 'images'))
    gen_im(data_dir)
    noise_bf(data_dir)

[((113, 9, 113, 40), 'numpy.arange', 'np.arange', ({(113, 19, 113, 27): 'step / 2', (113, 29, 113, 33): 'xmax', (113, 35, 113, 39): 'step'}, {}), '(step / 2, xmax, step)', True, 'import numpy as np\n'), ((114, 9, 114, 40), 'numpy.arange', 'np.arange', ({(114, 19, 114, 27): 'step / 2', (114, 29, 114, 33): 'ymax', (114, 35, 114, 39): 'step'}, {}), '(step / 2, ymax, step)', True, 'import numpy as np\n'), ((120, 10, 120, 33), 'numpy.zeros_like', 'np.zeros_like', ({(120, 24, 120, 32): 'x_coords'}, {}), '(x_coords)', True, 'import numpy as np\n'), ((167, 4, 167, 22), 'numpy.random.seed', 'np.random.seed', ({(167, 19, 167, 21): '(42)'}, {}), '(42)', True, 'import numpy as np\n'), ((48, 22, 48, 37), 'numpy.zeros', 'np.zeros', ({(48, 31, 48, 36): 'shape'}, {}), '(shape)', True, 'import numpy as np\n'), ((131, 21, 131, 84), 'os.path.join', 'os.path.join', ({(131, 34, 131, 42): 'data_dir', (131, 44, 131, 54): '"""cell_obj"""', (131, 56, 131, 83): '"""cells_final_selected.hdf5"""'}, {}), "(data_dir, 'cell_obj', 'cells_final_selected.hdf5')", False, 'import os\n'), ((138, 12, 138, 58), 'os.path.join', 'os.path.join', ({(138, 25, 138, 33): 'data_dir', (138, 35, 138, 43): '"""images"""', (138, 45, 138, 57): '"""binary.npy"""'}, {}), "(data_dir, 'images', 'binary.npy')", False, 'import os\n'), ((139, 12, 139, 63), 'os.path.join', 'os.path.join', ({(139, 25, 139, 33): 'data_dir', (139, 35, 139, 43): '"""images"""', (139, 45, 139, 62): '"""brightfield.npy"""'}, {}), "(data_dir, 'images', 'brightfield.npy')", False, 'import os\n'), ((140, 12, 140, 62), 'os.path.join', 'os.path.join', ({(140, 25, 140, 33): 'data_dir', (140, 35, 140, 43): '"""images"""', (140, 45, 140, 61): '"""foci_inner.npy"""'}, {}), "(data_dir, 'images', 'foci_inner.npy')", False, 'import os\n'), ((141, 12, 141, 62), 'os.path.join', 'os.path.join', ({(141, 25, 141, 33): 'data_dir', (141, 35, 141, 43): '"""images"""', (141, 45, 141, 61): '"""foci_outer.npy"""'}, {}), "(data_dir, 'images', 'foci_outer.npy')", False, 'import os\n'), ((142, 12, 142, 63), 'os.path.join', 'os.path.join', ({(142, 25, 142, 33): 'data_dir', (142, 35, 142, 43): '"""images"""', (142, 45, 142, 62): '"""storm_inner.npy"""'}, {}), "(data_dir, 'images', 'storm_inner.npy')", False, 'import os\n'), ((143, 12, 143, 63), 'os.path.join', 'os.path.join', ({(143, 25, 143, 33): 'data_dir', (143, 35, 143, 43): '"""images"""', (143, 45, 143, 62): '"""storm_outer.npy"""'}, {}), "(data_dir, 'images', 'storm_outer.npy')", False, 'import os\n'), ((145, 20, 145, 66), 'os.path.join', 'os.path.join', ({(145, 33, 145, 41): 'data_dir', (145, 43, 145, 51): '"""images"""', (145, 53, 145, 65): '"""binary.tif"""'}, {}), "(data_dir, 'images', 'binary.tif')", False, 'import os\n'), ((146, 20, 146, 71), 'os.path.join', 'os.path.join', ({(146, 33, 146, 41): 'data_dir', (146, 43, 146, 51): '"""images"""', (146, 53, 146, 70): '"""brightfield.tif"""'}, {}), "(data_dir, 'images', 'brightfield.tif')", False, 'import os\n'), ((147, 20, 147, 70), 'os.path.join', 'os.path.join', ({(147, 33, 147, 41): 'data_dir', (147, 43, 147, 51): '"""images"""', (147, 53, 147, 69): '"""foci_inner.tif"""'}, {}), "(data_dir, 'images', 'foci_inner.tif')", False, 'import os\n'), ((148, 20, 148, 70), 'os.path.join', 'os.path.join', ({(148, 33, 148, 41): 'data_dir', (148, 43, 148, 51): '"""images"""', (148, 53, 148, 69): '"""foci_outer.tif"""'}, {}), "(data_dir, 'images', 'foci_outer.tif')", False, 'import os\n'), ((149, 15, 149, 66), 'os.path.join', 'os.path.join', ({(149, 28, 149, 36): 'data_dir', (149, 38, 149, 46): '"""images"""', (149, 48, 149, 65): '"""storm_inner.txt"""'}, {}), "(data_dir, 'images', 'storm_inner.txt')", False, 'import os\n'), ((150, 15, 150, 66), 'os.path.join', 'os.path.join', ({(150, 28, 150, 36): 'data_dir', (150, 38, 150, 46): '"""images"""', (150, 48, 150, 65): '"""storm_outer.txt"""'}, {}), "(data_dir, 'images', 'storm_outer.txt')", False, 'import os\n'), ((156, 24, 156, 75), 'os.path.join', 'os.path.join', ({(156, 37, 156, 45): 'data_dir', (156, 47, 156, 55): '"""images"""', (156, 57, 156, 74): '"""brightfield.npy"""'}, {}), "(data_dir, 'images', 'brightfield.npy')", False, 'import os\n'), ((160, 14, 160, 31), 'colicoords.synthetic_data.draw_poisson', 'draw_poisson', ({(160, 27, 160, 30): 'img'}, {}), '(img)', False, 'from colicoords.synthetic_data import add_readout_noise, draw_poisson\n'), ((161, 14, 161, 43), 'colicoords.synthetic_data.add_readout_noise', 'add_readout_noise', ({(161, 32, 161, 35): 'img', (161, 37, 161, 42): 'noise'}, {}), '(img, noise)', False, 'from colicoords.synthetic_data import add_readout_noise, draw_poisson\n'), ((24, 55, 24, 68), 'tqdm.tqdm', 'tqdm', ({(24, 60, 24, 67): 'chunked'}, {}), '(chunked)', False, 'from tqdm import tqdm\n'), ((62, 21, 62, 42), 'numpy.max', 'np.max', ({(62, 28, 62, 41): '[0 - min1, 0]'}, {}), '([0 - min1, 0])', True, 'import numpy as np\n'), ((63, 21, 63, 80), 'numpy.min', 'np.min', ({(63, 28, 63, 79): '[data.shape[0] + (shape[0] - pos_y), data.shape[0]]'}, {}), '([data.shape[0] + (shape[0] - pos_y), data.shape[0]])', True, 'import numpy as np\n'), ((65, 21, 65, 42), 'numpy.max', 'np.max', ({(65, 28, 65, 41): '[0 - min2, 0]'}, {}), '([0 - min2, 0])', True, 'import numpy as np\n'), ((66, 21, 66, 80), 'numpy.min', 'np.min', ({(66, 28, 66, 79): '[data.shape[1] + (shape[1] - pos_x), data.shape[1]]'}, {}), '([data.shape[1] + (shape[1] - pos_x), data.shape[1]])', True, 'import numpy as np\n'), ((71, 19, 71, 36), 'numpy.max', 'np.max', ({(71, 26, 71, 35): '[min1, 0]'}, {}), '([min1, 0])', True, 'import numpy as np\n'), ((72, 19, 72, 43), 'numpy.min', 'np.min', ({(72, 26, 72, 42): '[max1, shape[0]]'}, {}), '([max1, shape[0]])', True, 'import numpy as np\n'), ((73, 19, 73, 36), 'numpy.max', 'np.max', ({(73, 26, 73, 35): '[min2, 0]'}, {}), '([min2, 0])', True, 'import numpy as np\n'), ((74, 19, 74, 43), 'numpy.min', 'np.min', ({(74, 26, 74, 42): '[max2, shape[1]]'}, {}), '([max2, shape[1]])', True, 'import numpy as np\n'), ((76, 26, 76, 41), 'numpy.zeros', 'np.zeros', ({(76, 35, 76, 40): 'shape'}, {}), '(shape)', True, 'import numpy as np\n'), ((79, 27, 79, 74), 'mahotas.distance', 'mh.distance', (), '', True, 'import mahotas as mh\n'), ((122, 20, 122, 35), 'numpy.ones_like', 'np.ones_like', ({(122, 33, 122, 34): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((124, 22, 124, 102), 'numpy.exp', 'np.exp', ({(124, 29, 124, 101): '(-(((_x - x_coords) / _sigma) ** 2 + ((_y - y_coords) / _sigma) ** 2) / 2)'}, {}), '(-(((_x - x_coords) / _sigma) ** 2 + ((_y - y_coords) / _sigma) ** 2) / 2\n    )', True, 'import numpy as np\n'), ((135, 26, 135, 58), 'os.path.join', 'os.path.join', ({(135, 39, 135, 47): 'data_dir', (135, 49, 135, 57): '"""images"""'}, {}), "(data_dir, 'images')", False, 'import os\n'), ((136, 17, 136, 49), 'os.path.join', 'os.path.join', ({(136, 30, 136, 38): 'data_dir', (136, 40, 136, 48): '"""images"""'}, {}), "(data_dir, 'images')", False, 'import os\n'), ((169, 26, 169, 58), 'os.path.join', 'os.path.join', ({(169, 39, 169, 47): 'data_dir', (169, 49, 169, 57): '"""images"""'}, {}), "(data_dir, 'images')", False, 'import os\n'), ((170, 17, 170, 49), 'os.path.join', 'os.path.join', ({(170, 30, 170, 38): 'data_dir', (170, 40, 170, 48): '"""images"""'}, {}), "(data_dir, 'images')", False, 'import os\n'), ((19, 20, 19, 80), 'numpy.random.normal', 'np.random.normal', ({(19, 37, 19, 49): 'cell_per_img', (19, 51, 19, 67): 'cell_per_img_std', (19, 69, 19, 79): 'num_images'}, {}), '(cell_per_img, cell_per_img_std, num_images)', True, 'import numpy as np\n'), ((31, 34, 31, 59), 'numpy.append', 'np.append', ({(31, 44, 31, 55): 'out_dict[k]', (31, 57, 31, 58): 'v'}, {}), '(out_dict[k], v)', True, 'import numpy as np\n'), ((38, 34, 38, 64), 'numpy.zeros', 'np.zeros', ({(38, 43, 38, 63): '(num_images, *shape)'}, {}), '((num_images, *shape))', True, 'import numpy as np\n'), ((55, 31, 55, 54), 'numpy.floor', 'np.floor', ({(55, 40, 55, 53): 'data.shape[0]'}, {}), '(data.shape[0])', True, 'import numpy as np\n'), ((58, 31, 58, 54), 'numpy.floor', 'np.floor', ({(58, 40, 58, 53): 'data.shape[1]'}, {}), '(data.shape[1])', True, 'import numpy as np\n'), ((96, 37, 96, 72), 'numpy.append', 'np.append', ({(96, 47, 96, 61): 'out_dict[name]', (96, 63, 96, 71): 'data_out'}, {}), '(out_dict[name], data_out)', True, 'import numpy as np\n'), ((52, 44, 52, 60), 'numpy.random.rand', 'np.random.rand', ({}, {}), '()', True, 'import numpy as np\n'), ((53, 44, 53, 60), 'numpy.random.rand', 'np.random.rand', ({}, {}), '()', True, 'import numpy as np\n')]