#PH4100 Major Poject - Photometry - Star field simulation program (3)
#Aaron Andrews
#Last edited: 011/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.
#*Mean value for each pixel (per star) calculated via integration over pixel area
#*
#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.

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 textwrap import wrap

def Gaussian(x,y,xStar,yStar,xSig,ySig,peak): #2d gaussian function, returns grey value contribution from a single star for any point in the star field
    return peak*math.exp(-((((x-xStar)**2)/(2*xSig**2))+(((y-yStar)**2)/(2*ySig**2))))
    
def IntGaussian(x,y,xStar,yStar,xSig,ySig,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)))
                
    #Take pixel coordinates as centre of pixel - integrate between range of [-0.5,+0.5] pixels
    xIntegral,xIntError=si.quad(GaussFuncX,x-0.5,x+0.5)#as numerical integration, returns error
    yIntegral,yIntError=si.quad(GaussFuncY,y-0.5,y+0.5)
    return (peak*xIntegral*yIntegral)

def SimStarField(nxPixel,nyPixel,MaxPixelValue,nStars):
    #Create empty arrays
    Image1=np.zeros((nxPixel,nyPixel)) #array to be output once filled - each cell contains grey values (no. of photons detected)
    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

    #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
        #Loop over all pixels in image
        pxlX=0
        while pxlX<nxPixel:
            pxlY=0
            while pxlY<nyPixel:
                #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[n],PosY[n],sigX,sigY,Peak[n])
                
                pxlY=pxlY+1
            pxlX=pxlX+1
        
        n=n+1
    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)
    GaussCalc=np.zeros(nxPixel)
    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,peak)
        n=n+1
    
    pxlX=0
    while pxlX<nxPixel:
        GaussInt[pxlX]=IntGaussian(pxlX,0,xStar,0,Sigma,Sigma,peak)
        GaussCalc[pxlX]=Gaussian(pxlX,0,xStar,0,Sigma,Sigma,peak)
        
        
        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,GaussCalc,'ro',label='Pixel midpoint')
    plt.plot(GaussCurveX,GaussCurveY,'g',label='Actual Gaussian profile')
    plt.legend(fontsize='small')
    plt.title("\n".join(wrap("Pixel value calculation from integration method compared to calculations using midpoints",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

#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(8,3.75,1,100)