#PH4100 Major Poject - Photometry - Star field simulation program (5)
#Aaron Andrews
#Last edited: 18/10/2016

#('*' indicates newest addition)
#This program generates a set number of stars with random positions and max. grey values.
#Stars are assumed to have 2d gaussian profiles. *correlation coefficient for x and y is an adjustable variable
#Mean value for each pixel (per star) calculated via integration over pixel area
#Mean value used to generate random poisson value for each pixel
#The output is a 2d array representing an image of a star field - each cell stores a grey value calculated using the cumulative gaussian profiles for all stars generated.
#The star field representation generated by this particular program is not accurate.
#*Output can be saved as FITS file

import math
import random as rm
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import scipy.integrate as si #for integration of gaussian function over pixels
from astropy.io import fits

from textwrap import wrap

#The first two functions are used for calculations of star gaussian profiles

def Gaussian(x,y,xStar,yStar,xSig,ySig,corr,peak): #2d gaussian function, returns grey value contribution from a single star for any point in the star field
    return peak*math.exp(-(1/(2*(1-corr**2)))*( (((x-xStar)**2)/(xSig**2)) + (((y-yStar)**2)/(ySig**2)) - 2*corr*((x-xStar)/xSig)*((y-yStar)/ySig) ))
#normalisation: (1/(2*math.pi*xSig/8ySignp.sqrt(1-corr**2)))

def IntGaussian(x1,y1,xStar,yStar,xSig,ySig,corr,peak): #returns integrated grey value contribution from a star for any pixel
    #integrate x and y components of gaussian seperately - define functions for integration first
    #GaussFuncX=lambda x:math.exp(-(((x-xStar)**2)/(2*xSig**2)))
    #GaussFuncY=lambda y:math.exp(-(((y-yStar)**2)/(2*ySig**2)))
    if corr==0:
        GaussFunc=lambda y,x:math.exp(-(((x-xStar)**2)/(2*xSig**2)))*math.exp(-(((y-yStar)**2)/(2*ySig**2))) #Failsafe for integration - without it, would cause integration to return value of 1
    else:
        GaussFunc=lambda y,x:math.exp(-(1/(2*(1-corr**2)))*( (((x-xStar)**2)/(xSig**2)) + (((y-yStar)**2)/(ySig**2)) - 2*corr*((x-xStar)/xSig)*((y-yStar)/ySig) ))
                
    #Take pixel coordinates as centre of pixel - integrate between range of [-0.5,+0.5] pixels
    #xIntegral,xIntError=si.quad(GaussFuncX,x1-0.5,x1+0.5)#as numerical integration, returns error
    #yIntegral,yIntError=si.quad(GaussFuncY,y1-0.5,y1+0.5)
    Integral,IntError=si.dblquad(GaussFunc,x1-0.5,x1+0.5,lambda y:y1-0.5,lambda y:y1+0.5)#2d numerical integration
    return (peak*Integral)


#The following functions can be used to manually create a star field and generate stars with set parameters.
#(Optional) Generated image can be converted into array of discrete (poisson) pixel values, and saved as a FITS file.

def MakeField(nxPixel,nyPixel): #*Generate unfilled array for star field
    return np.zeros((nxPixel,nyPixel))
            
def AddStar(posX,posY,sigX,sigY,corr,maxN,image): #*adds a star to an existing image
    Image1=image
    #Loop over all pixels in image
    pxlX=0
    while pxlX<len(image):
        pxlY=0
        while pxlY<len(image[0]):
            #For every pixel, a grey value as determined by the gaussian profile of each star is calculated, then added to the existing grey value of the pixel.
            #Image1[pxlX,pxlY]=Image1[pxlX,pxlY]+int(Gaussian(pxlX,pxlY,PosX[n],PosY[n],sigX,sigY,Peak[n])) #no integration
            Image1[pxlX,pxlY]=Image1[pxlX,pxlY]+IntGaussian(pxlX,pxlY,posX,posY,sigX,sigY,corr,maxN)    
            pxlY=pxlY+1
        pxlX=pxlX+1
    return Image1

def MakePoisson(imageMeans): #Use calculated values to generate random poisson-distributed pixel values
    Image1=MakeField(len(imageMeans),len(imageMeans[0]))
    pxlX=0
    while pxlX<len(imageMeans):
        pxlY=0
        while pxlY<len(imageMeans[0]):
            Image1[pxlX,pxlY]=np.random.poisson(imageMeans[pxlX,pxlY])
            pxlY=pxlY+1
        pxlX=pxlX+1
    return Image1
    
def saveFile(fileName,image): #save a star field array as a FITS file
    #nxPixel=len(image) #no. of columns
    #nyPixel=len(image[0]) #no. of rows
    hdu = fits.PrimaryHDU(image)
    hdu.writeto(fileName,clobber=True)
    
#The following functions can be used to create more complex star fields using all the functions above.

def SimStarField(nxPixel,nyPixel,MaxPixelValue,nStars): #Randomised star field - positions and max intensities are uniformly distributed
    #Create empty arrays
    Image1=MakeField(nxPixel,nyPixel) #array to be output once filled - each cell contains grey values (no. of photons detected)
    Image1Means=MakeField(nxPixel,nyPixel) #filled first with mean pixel values to be used for poisson distributed random variables
    PosX=np.zeros(nStars) #x coordinates
    PosY=np.zeros(nStars) #y coordinates
    Peak=np.zeros(nStars) #peak grey value for each star
    sigX=3 #standard deviations, set to arbitrary values (for now)
    sigY=3
    corr=0

    #Loop over all stars
    n=0
    while n<nStars:
        #Random value generation (eventually will be replaced with retrieval from database of stars)
        PosX[n]=rm.random()*nxPixel #generates random value between 0 and 1, multiplies by relevent range
        PosY[n]=rm.random()*nyPixel
        Peak[n]=rm.random()*MaxPixelValue
        #Image1[int(PosX[n]),int(PosY[n])]=1
        Image1Means=AddStar(nxPixel,nyPixel,PosX[n],PosY[n],sigX,sigY,corr,Peak[n],Image1)
        n=n+1
    
    Image1=MakePoisson(Image1Means)
    
    return Image1
    
def IntegrationTest(nxPixel,xStar,Sigma,peak): #demonstration of integrated pixel values vs centre-point calculation
    
    axisX=np.linspace(0,nxPixel-1,nxPixel)
    
    GaussInt=np.zeros(nxPixel) #Pixel values calculated by integration
    GaussCalc=np.zeros(nxPixel) #Pixel values calculated with gaussian by inputting bin midpoint
    GaussPois=np.zeros(nxPixel) #Calculation with integration, then used as mean for random distributed poisson variable
    GaussCurveX=np.linspace(0,nxPixel,1000)
    GaussCurveY=np.zeros(1000)
    
    n=0
    while n<1000:
        GaussCurveY[n]=Gaussian(GaussCurveX[n],0,xStar,0,Sigma,Sigma,0,peak)
        n=n+1
    
    pxlX=0
    while pxlX<nxPixel:
        GaussInt[pxlX]=IntGaussian(pxlX,0,xStar,0,Sigma,Sigma,0,peak)
        GaussCalc[pxlX]=Gaussian(pxlX,0,xStar,0,Sigma,Sigma,0,peak)
        GaussPois[pxlX]=np.random.poisson(GaussInt[pxlX])
        
        
        pxlX=pxlX+1
        
    plt.figure(1)
    #plt.bar(axisX,GaussInt,edgecolor='b',color='b',align='center',fill=False,width=1,label='Integration method')
    #plt.plot(axisX,GaussInt,'bo',label='Integrated pixel value')
    plt.bar(axisX,GaussPois,edgecolor='maroon',color='maroon',align='center',fill=False,width=1,label='Poisson distributed pixel values')
    #plt.plot(axisX,GaussCalc,'ro',label='Pixel midpoint')
    plt.plot(GaussCurveX,GaussCurveY,'g',label='Actual Gaussian profile')
    plt.legend(fontsize='small')
    plt.title("\n".join(wrap("Pixel values generated via poisson distribution, using integrated pixel value as mean",60)))
    plt.xlabel('x axis position (pixels)')
    plt.ylabel('Value')
    plt.xlim(-0.5,nxPixel-0.5)
    plt.ylim(0,peak*1.2)
    plt.show(1)
    
    return 0
    
def MakeTestImage():
    return 0
    


    
    return 0

#Defining example variables
ImageX=640 #pixels, x axis
ImageY=480 #pixels, y axis
MaxGrey=65535

Nstars=50 #number of stars to generate

#SimStarField(ImageX,ImageY,MaxGrey,Nstars)
#IntegrationTest(16,7.75,2,100)