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

#('*' indicates newest addition)
#This program generates a set number of stars with random variables (positions, sigmas, fluxes, correlation)
#*Background can be generated independently, then added to star field
#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 *(only if value of gaussian at pixel centre is not negligable)
#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 image can be saved as FITS file
#*randomly generated paramaters of stars, as returned by 'SimStarField' function, can be output to text file
#*a seed can be set when simulating star fields - use to produce star fields with identical randomised parameters for stars

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,flux): #2d gaussian function, returns grey value contribution from a single star for any point in the star field
    return flux*(1/(2*math.pi*xSig*ySig*np.sqrt(1-corr**2)))*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,flux): #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,epsabs=10,epsrel=10)#2d numerical integration
    return flux*(1/(2*math.pi*xSig*ySig*np.sqrt(1-corr**2)))*Integral


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

def MakeField(nxPixel,nyPixel): #*Generate unfilled array for star field
    return np.zeros((nxPixel,nyPixel))
            
def AddStar(posX,posY,sigX,sigY,corr,flux,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
            gaussEst=Gaussian(pxlX,pxlY,posX,posY,sigX,sigY,corr,flux) #calculating gaussian value at pixel centre
            
            if gaussEst>0.1: #prevents integration for negligable gaussian values at pixel coordinates
                Image1[pxlX,pxlY]=Image1[pxlX,pxlY]+IntGaussian(pxlX,pxlY,posX,posY,sigX,sigY,corr,flux)    
            
            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 AddBackground(image,backMean,backErr): #add background (noise) to an image - best to generate stars first. values are discrete
    Image1=MakeField(len(image),len(image[0]))
    pxlX=0
    while pxlX<len(image):
        pxlY=0
        while pxlY<len(image[0]):
            Image1[pxlX,pxlY]=image[pxlX,pxlY]+int(rm.gauss(backMean,backErr)) #error=FWHM*2.35
            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)
    
def saveParams(filename,params): #save star parameters, as returned by 'SimStarField' function
    file1=open(filename,'w')
    for n in params: 
        file1.write(n+"\n")
    file1.close()
    
#The following functions can be used to create more complex star fields using all the functions above.

def SimStarField(nxPixel,nyPixel,maxPeak,nStars,backMean,backErr,randSeed): #Randomised star field - positions and max intensities are uniformly distributed
    
    rm.seed(randSeed)
    
    extend=15 #additional pixel distance from edges to generate out of frame stars
    
    
    PosX=np.zeros(nStars) #x coordinates
    PosY=np.zeros(nStars) #y coordinates
    
    flux=np.zeros(nStars) #total flux
    peak=np.zeros(nStars) #peak grey value for each star
    sigX=np.zeros(nStars) #sigma values
    sigY=np.zeros(nStars)
    corr=np.zeros(nStars) #correlation between x and y
    
     #calculate max. measurable flux based on max pixel value
    
    #Create empty arrays for image data
    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
    
    #1d list for writing star parameters to output file
    Params=["" for x in range(nStars+1)] #first set of stars in list are in frame - stars out of frame are added to end of list
    
    #Define column headers for output file
    Params[0]="No.\tX position\tY position\tTotal flux\tPeak value\tX sigma value\tY sigma value\tCorrelation"

    #Loop over all stars
    n=0 #no. of stars generated within frame
    nOut=0 #no. of stars out of frame
    nT=0 #total no. of stars generated (ina nd out of frame
    while n<nStars:
        
        #Random value generation (eventually will be replaced with retrieval from database of stars)
        PosX[n]=rm.random()*(nxPixel+2*extend)-extend #generates random value between 0 and 1, multiplies by relevent range
        PosY[n]=rm.random()*(nyPixel+2*extend)-extend #range is extended over fiducial volume
        
        sigX[n]=rm.random()*3+3 #range [3,6]
        sigY[n]=rm.random()*3+3 #range [3,6]
        corr[n]=rm.random()-0.5 #range [-0.5,0.5]
        
        peak[n]=rm.random()*maxPeak
        flux[n]=peak[n]*(2*math.pi*sigX[n]*sigY[n]*np.sqrt(1-corr[n]**2)) #calculate flux
        
        #calc peak from flux value:
        #peak[n]=flux[n]/(2*math.pi*sigX[n]*sigY[n]*np.sqrt(1-corr[n]**2))
        
        Image1Means=AddStar(PosX[n],PosY[n],sigX[n],sigY[n],corr[n],flux[n],Image1Means)
        if PosX[n]>=0 and PosX[n]<=nxPixel-1 and PosY[n]>=0 and PosY[n]<=nyPixel-1: #test if star is in frame. stops once set number of stars appear in frame.
            Params[n+1]=str(n+1)+"\t"+str(PosX[n])+"\t"+str(PosY[n])+"\t"+str(flux[n])+"\t"+str(peak[n])+"\t"+str(sigX[n])+"\t"+str(sigY[n])+"\t"+str(corr[n])
            n=n+1
            print "added star (in frame): ", n #for keeping track of progress
            
        else:
            Params.append(str(nStars+nOut+1)+"\t"+str(PosX[n])+"\t"+str(PosY[n])+"\t"+str(flux[n])+"\t"+str(peak[n])+"\t"+str(sigX[n])+"\t"+str(sigY[n])+"\t"+str(corr[n]))
            nOut=nOut+1
            print "added star (out frame): ", nOut
        nT=nT+1    
        
    Image1=MakePoisson(Image1Means)
    
    Image1=AddBackground(Image1,backMean,backErr)
    
    
    
    return Image1, Params
    
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)