# -*- coding: utf-8 -*-
# homework3.py
"""
Author:   Raghuveer Parthasarathy
Created on Thu Oct 14 22:05:51 2021
Last modified on Oct 21, 2024

Description
-----------
Scripts for image analysis homework 4
"""

import numpy as np
import matplotlib.pyplot as plt
import os
import scipy.special as scipy_special  # scipy is too big to import all of it.
from skimage import io # input output sub-package 
from skimage import filters as skfilters # filtering sub-package 
import scipy.ndimage as ndi # for convolution

#%% parentDir

parentDir = r'C:\Users\Raghu\Documents\Teaching\Image Analysis Course\Image Analysis Fall 2024\Homework\Images for Homework'


#%% Problem 1: Frequency modulation

print('See Code in frequency_modulation_hw.py')

    
#%% Problem 2: Quantized aggregates. 
    
# See assignment for details
# Based on my MATLAB version: oligomers_Ex4p1_RPSolution.m

# General method:
# High-pass filter and subtract background. Background is defined as the 
#   mean out-of-neighborhood intensity, so be sure the neighborhood is
#   large enough to encapsulate the signals!


# Parameters
# filename =  'emitters_33px_1000ph_noNoise.png' # image filename
fileName = r'emitters_33px_100ph.png'
gridPosition = 33 # positions (x, y) are at multiples of this value
#     (+/- 0.25 px)
neighborhood = 2 # consider around each grid position a region of +/-
#      neighborhood pixels; 

im = io.imread(os.path.join(parentDir, fileName))

# number of particles. It's 10x10, but I'll be a bit more general:
# number in "x" and "y", so [10 10]
# force integer to use in loops later
nParticles_1D = np.floor(np.array(im.shape)/gridPosition).astype(int) 

plt.figure()
plt.imshow(im, cmap = 'gray')
plt.title('Original')

# Let's look at a histogram
plt.figure()
plt.hist(im.ravel(), 50);
# ax = axis; axis([0 ax(2) 0 ax(4)])
plt.xlabel('Intensity')
plt.title('Original')

# Process the image -- high-pass filter

# Gaussian filter
Gstd = 21; # Gaussian width
im_Gaussfilt = skfilters.gaussian(im, Gstd, mode = 'nearest')
print(f'Max of original: {np.max(im)}')
print(f'Max of LP filtered: {np.max(im_Gaussfilt):.3f}')
# We *do* need to rescale; let's check to confirm...
if np.max(im_Gaussfilt) <= 1.0:
    im_Gaussfilt = 255.0*im_Gaussfilt

# high-pass filter
im_HP = im - im_Gaussfilt;
plt.figure()
plt.imshow(im_HP, cmap = 'gray')
plt.title('High-pass filter')

plt.imsave('emitters_33px_100ph_HP21.png', im_HP, cmap='gray')

# Let's look at the new histogram
plt.figure()
plt.hist(im_HP.ravel(), 20);
# ax = axis; axis([0 ax(2) 0 ax(4)])
plt.xlabel('Intensity')
plt.title('High-pass filter')

# It's dominated by the zero-intensity pixels. We could change our bins and 
# look at nonzero values only, or look at neighborhoods around the (known) 
# pixel locations
# Let's do the latter

# Get total intensity in each neighborhood
intensities = np.zeros(nParticles_1D)
# Note that each bright dot is at index gridPosition-1, 2*gridPosition -1 , ...
for j in range(nParticles_1D[0]):
    for k in range(nParticles_1D[1]):
        neighborhoodIdx = np.arange(-neighborhood,neighborhood+1)
        im_thisRegion = im_HP[(j+1)*gridPosition-1 + neighborhoodIdx, 
                              (k+1)*gridPosition-1 + neighborhoodIdx]
        # simple sum of intensities in each neighborhood:
        intensities[j,k] = np.sum(im_thisRegion) 

plt.figure()
plt.hist(intensities.ravel(), 15);
plt.xlabel('Intensity')
plt.title('Region Intensities')
# ax = axis; axis([0 ax(2) 0 ax(4)])
plt.xlabel('Region intensity, bkg sub')
plt.ylabel('Number of Regions')

# Check the background (as reported by median); should be small compared to 
# region intensities:
print('Median intensity: ', np.median(im_HP))
print('Median intensity * neighborhood size: ', 
      np.median(im_HP)*(2*neighborhood+1)**2)

# Look at histogram; clearly shows peaks!
# I'll lazily just divide the range in thirds
intensity_range = np.max(intensities) - np.min(intensities)

intensity_edges = np.min(intensities) + np.array([0.33, 0.67, 1.0])*intensity_range
print('Intensity bins: ', intensity_edges)
n_mer = np.zeros((3,1))
n_mer[0] = np.sum(intensities.ravel()<intensity_edges[0]) # monomers
n_mer[1] = np.sum(np.logical_and(intensities.ravel()>=intensity_edges[0], 
                                intensities.ravel()<intensity_edges[1]))
n_mer[2] = np.sum(intensities.ravel()>=intensity_edges[1])
# Oddly, can't use formatting strings...
print(f'{n_mer[0]} monomers, {n_mer[1]} dimers, {n_mer[2]} trimers')
# ... but I can here!
print(f'total number of dots: {np.sum(n_mer):.0f}')



#%% Problem 3: A high-resolution PSF. 

# Code copied from psf2D.m and slightly modified, to be stand-alone here and 
# to use the same definition of N

def psf2D(N, scale = 0.1, NA = 0.9, lam = 0.53):
    """
    Function to calculate the Point Spread Function (PSF)
    Returns a square 2D array, NxN, normalized to sum==1.
    Inputs:
       N : PSF array size (pixels)
       scale: microns / pixel
       NA : numerical aperture
       lam : free-space wavelength, *microns*
   """

    # force N to be integer:
    N = int(N)
    x, y = np.meshgrid(np.arange(-(N-1)/2,(N+1)/2), np.arange(-(N-1)/2,(N+1)/2))
    d = np.sqrt(x*x + y*y);  # matrix of distances from the center, in pixels
    rpositions = d*scale;  # matrix of distances to the center, microns
    v = (2*np.pi/lam)*NA*rpositions;  # array of 'v' values
    
    psf = 4*(scipy_special.j1(v)/v)**2
    psf[int((N-1)/2), int((N-1)/2)]=1.0 # Correct the undefined central value
    psf = psf/np.sum(psf) # Normalize
    return psf


N = 101
scale = 1/N
lam = 0.5
NA = 0.9

psf = psf2D(N, scale=scale, NA=NA, lam=lam)
plt.figure()
plt.imshow(psf, cmap='gray')

#%% Problem 4: A worse worm image

fileName = r'fetter_Celegans_cellfig10.jpg'
im = io.imread(os.path.join(parentDir, fileName))

plt.figure()
plt.imshow(im, cmap='gray')

# Examining the scale bar, it's about 110 px wide, for 1 um in the focal
# plane, so the scale is 1/110 um/px, or 0.0091 um/px.
scale = 1/110 # um/px
lambda_microns = 0.53 # wavelength
NA = 0.7;
Npx = round(3*lambda_microns/NA/scale)
print('Npx = ', Npx);
psf = psf2D(Npx, scale, NA, lambda_microns)
plt.figure()
plt.imshow(psf, cmap='gray')

# Convolution: very slow! (Minutes!)
im_optical = ndi.convolve(im, psf, mode='nearest') # As in HW2
plt.figure()
plt.imshow(im_optical, cmap='gray')
plt.imsave('fetter_Celegans_Optical.png', im_optical, cmap='gray')


#%% Problem 6-8: Simulated point sources
    
def simPointSource(N, scale=0.1, fineScale = 0.01, NA = 0.9, 
                   lam = 0.53, Nphoton = 1000, xc = 0, yc = 0, 
                   bkgPoissMean=0, makeFigures=False):
    """
    Function to create a simulated point source image: PSF, pixelation, noise
    Returns a square 2D array, NxN, normalized to sum==Nphoton.
    Inputs:
       N : final camera image array size (N x N pixels)
       scale : final camera scale, um/px
       fineScale: microns / pixel for high-resolution PSF
       NA : numerical aperture
       lam : free-space wavelength of light, *microns*
       Nphoton : number of photons, total; rescale output to sum to this
       xc, yc : point source center location, *microns*
       bkgPoissMean : mean per-pixel background intensity
       makeFigures : if true, show fine PSF
     Output: 
       simImage : NxN integer array, simulated image, photons at each px.
   """
    scaleMultiplier = scale / fineScale
    # Check that it's an integer multiple
    if np.abs(scaleMultiplier % 1) > 0.01:
        print('Warning! camera scale is not an integer multiple of the fine scale')
        print(f'Original scaleMultiplier: {scaleMultiplier:.3e}')
        fineScale = scale / np.round(scaleMultiplier)
        print(f'Changing fineScale to be {fineScale:.3e}')
    else:
        scaleMultiplier = np.round(scaleMultiplier)
    
    # High-resolution PSF
    # Could call my PSF function, but I'll make this self-contained.
    N_fine = int(N*scaleMultiplier) # pixels for the fine array
    # High-resolution PSF
    x1D = fineScale*np.arange(-(N_fine-1)/2,(N_fine+1)/2) - xc
    y1D = fineScale*np.arange(-(N_fine-1)/2,(N_fine+1)/2) - yc
    x, y = np.meshgrid(x1D, y1D)  # grid positions, microns
    rpositions = np.sqrt(x*x + y*y);  # matrix of distances from the center, in microns
    v = (2*np.pi/lam)*NA*rpositions;  # array of 'v' values    
    psf = 4*(scipy_special.j1(v)/v)**2
    psf[int((N_fine-1)/2), int((N_fine-1)/2)]=1 # Correct the undefined central value
    if makeFigures:
        plt.figure()
        plt.imshow(psf, cmap='gray')
        plt.title('Fine resolution')
    
    # Pixelate
    simImage = np.zeros((N, N))
    for j in range(N):
        for k in range(N):
            j_range = np.arange(j*scaleMultiplier, (j+1)*scaleMultiplier).astype(int)
            k_range = np.arange(k*scaleMultiplier, (k+1)*scaleMultiplier).astype(int)
            simImage[j,k] = np.sum(psf[j_range, k_range])    
    
    # Rescale so that the total intensity is Nphoton
    simImage = Nphoton*simImage / np.sum(simImage)
    
    # Photon noise
    # Draw the intensity values from a Poisson distribution for each pixel
    # Note: integer values
    simImage = np.random.poisson(simImage, simImage.shape); 
    
    # Add Poisson background
    if bkgPoissMean > 0:
        im_bkg = np.random.poisson(bkgPoissMean, simImage.shape) # Background
        simImage += im_bkg

    return simImage


# Let's show some images!
N = 15
scale = 0.1 # microns/px
lam = 0.5
NA = 0.9
Nphoton = 500

xc = 0.0 # x-offset, microns
yc = 0.0 # y-offset, microns

bkgPoissMean = 0

im = simPointSource(N, scale, fineScale = scale/10, 
                    NA=NA, lam=lam, Nphoton=Nphoton, xc=xc, yc=yc, 
                    bkgPoissMean=bkgPoissMean, makeFigures=True)
plt.figure()
plt.imshow(im, cmap='gray')
# plt.imshow(im)
plt.title(f'Simulated image, Nphoton = {Nphoton}')

plt.savefig(f'SimImage_Nphoton_{Nphoton}.png', bbox_inches='tight')

#%% Problem 8

xc = 0.03 # x-offset, microns
yc = 0.03 # y-offset, microns
Nphoton = 50

bkgPoissMean = 2

im = simPointSource(N, scale, fineScale = scale/10, 
                    NA=NA, lam=lam, Nphoton=Nphoton, xc=xc, yc=yc, 
                    bkgPoissMean=bkgPoissMean, makeFigures=True)
plt.figure()
plt.imshow(im, cmap='gray')
# plt.imshow(im)
plt.title(f'Simulated image, Nphoton = {Nphoton}')
plt.savefig(f'SimImage_Nphoton_{Nphoton}_Offset.png', bbox_inches='tight')
#%% Problem #8

A = 100 # constant value
M = 1000 # number of trials for each sigma
sigma_array = np.logspace(-1,2,100) # from 0.1 to 100

# (a)
RMSE_mean = np.zeros((sigma_array.size,1))
for j in np.arange(0, sigma_array.size):
    x = A + sigma_array[j]*np.random.randn(M,1)
    RMSE_mean[j] = np.sqrt(np.mean((x-A)**2))

plt.figure()
plt.loglog(sigma_array, RMSE_mean)
plt.xlabel('sigma')
plt.ylabel('RMSE_mean')

# (b)
sigma = 20
N_array = np.round(np.logspace(0,3,1000)).astype(int) # from 1 to 1000
RMSE_mean = np.zeros((N_array.size,1))
RMSE_median = np.zeros((N_array.size,1))
for j in np.arange(0, N_array.size):
    x = A + sigma*np.random.randn(M,N_array[j])
    Ahat_mean = np.mean(x, axis=1) # mean along the Ns
    Ahat_median = np.median(x, axis=1) # mean along the Ns
    RMSE_mean[j] = np.sqrt(np.mean((Ahat_mean-A)**2))
    RMSE_median[j] = np.sqrt(np.mean((Ahat_median-A)**2))

plt.figure()
plt.loglog(N_array, RMSE_mean)
plt.loglog(N_array, RMSE_median)
plt.xlabel('N')
plt.ylabel('RMSE')
plt.legend(['Alice (mean)', 'Bob (median)'])

#%%  Misc

# Let's showsome images!
N = 7 # Size of the simulated image, NxN Pixels
scale = 0.1 # camera scale, microns/px
lam = 0.5 # wavelength of light, microns
NA = 0.9 # NA
Nphoton = 2000 # total number of photons

xc = 0.04 # x-offset, microns
yc = 0.0 # y-offset, microns

bkgPoissMean = 2 # Mean of additive Poisson background

N_images = 100
im =  np.zeros((N, N, N_images), dtype=float)
for j in range(N_images):
    im[:,:,j] = simPointSource(N, scale, fineScale = scale/10, 
                    NA=NA, lam=lam, Nphoton=Nphoton, xc=xc, yc=yc, 
                    bkgPoissMean=bkgPoissMean)

plt.figure()
plt.imshow(im[:,:,4])
plt.title('Image 5')

print('Ratio #1')
print(f'   mean: {np.mean(im[3,3,:]/im[3,4,:]):.2f}')
print(f'   std: {np.std(im[3,3,:]/im[3,4,:]):.2f}')
print('Ratio #2')
print(f'   mean: {np.mean(im[3,3,:]/im[3,2,:]):.2f}')
print(f'   std: {np.std(im[3,3,:]/im[3,2,:]):.2f}')
print('Ratio #3')
print(f'   mean: {np.mean(im[3,3,:]/im[2,3,:]):.2f}')
print(f'   std: {np.std(im[3,3,:]/im[2,3,:]):.2f}')
print('Values manually noted.')

def checkImage(simImage):
    # function to do a few numerical checks of PSFs
    # Basing criteria on quick examination of simulated arrays
    # with 0.5 micron light, NA = 0.9, xc= 40 nm, yc = 0
    N = simImage.shape[0] # won't check that array is square
    center_px = int((N-1)/2)
    ratio_center_right1px = simImage[center_px,center_px] / \
        simImage[center_px,center_px+1]
    print(f'Ratio #1: {ratio_center_right1px:.2f}')
    if np.abs(ratio_center_right1px-1.06)<0.18:
        print('Test 1: good')
    else:
        print('Test 1: failed!')
    ratio_center_left1px = simImage[center_px,center_px] / \
        simImage[center_px,center_px-1]
    print(f'Ratio #2: {ratio_center_left1px:.2f}')
    if np.abs(ratio_center_left1px-1.81)<0.4:
        print('Test 2: good')
    else:
        print('Test 2: failed!')
    ratio_center_up1px = simImage[center_px,center_px] / \
        simImage[center_px-1,center_px]
    print(f'Ratio #3: {ratio_center_up1px:.2f}')
    if np.abs(ratio_center_up1px-1.36)<0.26:
        print('Test 3: good')
    else:
        print('Test 3: failed!')
    

checkImage(im[:,:,2])
