#PH4100 Major Poject - Photometry - Star field simulation program (9)
#Aaron Andrews
#Last edited: 15/11/2016

#('*' indicates newest addition)

#This program can be used to generate a mostly accurate representation of a star field as seen via a telescope/camera system.

###Inputs:
#*Can either retrieve list of known stars (coordinates and magnitudes), or generate a set number of stars with random variables (positions, sigmas, fluxes, correlation)
#Alternatively, can generate stars with specific parameters.
#Can set image size, exposure time, camera gain, plate scale.

###Simulation:
#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
#Background can be generated independently, then added to star field
#*Saturation taken into account

###Ouputs:
#The output is a 2d array representing an image of a star field - each cell stores a grey value (ADU) calculated using the cumulative gaussian profiles for all stars generated, plus any additional layers.
#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


##### Gaussian profile calculation

#2d gaussian function, returns grey value contribution from a single star for any point in the star field
def Gaussian(x,y,xStar,yStar,xSig,ySig,corr,flux): 
    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)))

#returns integrated grey value contribution from a star for any pixel
def IntGaussian(x1,y1,xStar,yStar,xSig,ySig,corr,flux): 
    #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: #can use cumulative distributions
        #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
        Integral=0.25*(math.erf((x1-0.5-xStar)/(math.sqrt(2.0)*xSig))-math.erf((x1+0.5-xStar)/(math.sqrt(2.0)*xSig)))*(math.erf((y1-0.5-yStar)/(math.sqrt(2.0)*ySig))-math.erf((y1+0.5-yStar)/(math.sqrt(2.0)*ySig)))
    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) ))
        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
        Integral=Integral*flux*(1/(2*math.pi*xSig*ySig*np.sqrt(1-corr**2)))
                
    #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 Integral


##### Star field generation and processing - use to construct star field. Best to use functions in sequence

#Generate unfilled array for star field
def MakeField(nxPixel,nyPixel): 
    return np.zeros((nyPixel,nxPixel))

#adds a star to an existing image                        
def AddStar(posX,posY,sigX,sigY,corr,flux,image): 
    Image1=image
    #Loop over all pixels in image
    pxlX=0
    while pxlX<len(image[0]):
        pxlY=0
        while pxlY<len(image):
            #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>1: #prevents integration for negligable gaussian values at pixel coordinates
                Image1[pxlY,pxlX]=Image1[pxlY,pxlX]+IntGaussian(pxlX,pxlY,posX,posY,sigX,sigY,corr,flux)      
            else:
                Image1[pxlY,pxlX]=Image1[pxlY,pxlX]+gaussEst
            
            pxlY=pxlY+1
        pxlX=pxlX+1
    return Image1

#Use calculated pixel values in image to generate random poisson-distributed pixel values
def MakePoisson(imageMeans): 
    Image1=MakeField(len(imageMeans[0]),len(imageMeans))
    pxlX=0
    while pxlX<len(imageMeans[0]):
        pxlY=0
        while pxlY<len(imageMeans):
            Image1[pxlY,pxlX]=np.random.poisson(imageMeans[pxlY,pxlX])
            pxlY=pxlY+1
        pxlX=pxlX+1
    return Image1

#add background (noise) to an image - values are discrete        
def AddBackground(image,backMean,backErr): 
    Image1=MakeField(len(image[0]),len(image))
    pxlX=0
    while pxlX<len(image[0]):
        pxlY=0
        while pxlY<len(image):
            Image1[pxlY,pxlX]=image[pxlY,pxlX]+int(rm.gauss(backMean,backErr)) #error=FWHM*2.35
            pxlY=pxlY+1
        pxlX=pxlX+1
    return Image1

#Simulate pixel saturation (pixels have max value = 2^16-1 = 65535)       
def Saturate(Image): 
    pxlX=0 
    while pxlX<len(Image[0]):
        pxlY=0
        while pxlY<len(Image):
            if Image[pxlY,pxlX]>65535:
                Image[pxlY,pxlX]=65535
            pxlY=pxlY+1
        pxlX=pxlX+1
    return Image
    
    
##### File utilities and operations    

#Set up simulation output file for generating from real star parameters
def setupOutputs_Real(inFile,outFile,filt,nxPixel,nyPixel,telRA,telDEC,pScale,maxFlux,maxPeak,expT,backMean,backError,gain,minSigma,nFrame,nStars):
    imageData=[]
    imageData.append('Readout for '+outFile+'.fits')    
    imageData.append("To recreate, use: SimStarField_Real('"+inFile+"','"+outFile+"','"+filt+"',"+str(nxPixel)+","+str(nyPixel)+",'"+str(telRA)+"','"+str(telDEC)+"',"+str(pScale)+","+str(maxPeak)+","+str(expT)+","+str(gain)+","+str(minSigma)+")")
#SimStarField_Real(inFile,outFile,filt,nxPixel,nyPixel,telRA,telDEC,pScale,maxPeak,expT,gain,minSigma):    
    imageData.append('')
    imageData.append('Input data:\t'+inFile)
    imageData.append('Filter:\t'+filt)
    imageData.append('Frame width:\t'+str(nxPixel))
    imageData.append('Frame height:\t'+str(nyPixel))
    imageData.append('RA (hms):\t'+str(telRA))
    imageData.append('DEC (dms):\t'+str(telDEC))
    imageData.append('Plate scale:\t'+str(pScale))
    imageData.append('Gain:\t'+str(gain))
    imageData.append('Max. flux value:\t'+str(maxFlux))
    imageData.append('Max. peak value:\t'+str(maxPeak))
    imageData.append('Star sigma value:\t'+str(minSigma))
    imageData.append('Exposure time:\t'+str(expT))
    
    imageData=setupOutputs2(imageData,backMean,backError,nFrame,nStars)
    
    return imageData
    
#Setup output for randomly generated star field ###need to finish
def setupOutputs_Rand(outFile,nxPixel,nyPixel,minPeak,maxPeak,minSigma,maxSigma,maxCorr,backMean,backError,gain,randSeed,nFrame,nStars):
    imageData=[]
    imageData.append('Readout for '+outFile+'.fits')    
    imageData.append("To recreate, use: SimStarField_Rand('"+outFile+"',"+str(nxPixel)+","+str(nyPixel)+","+str(minPeak)+","+str(maxPeak)+","+str(minSigma)+","+str(maxSigma)+","+str(maxCorr)+","+str(backMean)+","+str(backError)+","+str(gain)+","+str(randSeed)+")")
#SimStarField_Rand(outFile,nxPixel,nyPixel,nStars,minPeak,maxPeak,minSigma,maxSigma,maxCorr,backMean,backErr,gain,randSeed):
    imageData.append('')
    imageData.append('Frame width:\t'+str(nxPixel))
    imageData.append('Frame height:\t'+str(nyPixel))
    imageData.append('Gain:\t'+str(gain))    
    imageData.append('Peak range: ['+str(minPeak)+','+str(maxPeak)+']')
    imageData.append('Sigma range: ['+str(minSigma)+','+str(maxSigma)+']')
    imageData.append('Correllation range: [-'+str(maxCorr)+','+str(maxCorr)+']') 
    
    imageData=setupOutputs2(imageData,backMean,backError,nFrame,nStars)    
    return imageData
    
#general output data for both real and randomly generated star fields    
def setupOutputs2(imageData,backMean,backError,nFrame,nStars):
    
    imageData.append('Background mean:\t'+str(backMean))
    imageData.append('Background S.D.:\t'+str(backError))
    imageData.append('Stars in frame:\t'+str(nFrame))
    imageData.append('Total stars generated:\t'+str(nStars))
    imageData.append('')
    imageData.append('Star parameters:')
    imageData.append("No.\tX position\tY position\tTotal flux\tPeak value\tX sigma value\tY sigma value\tCorrelation")
    return imageData

#save a star field array as a FITS file     
def saveOutputs(fileName,image): 
    hdu = fits.PrimaryHDU(image)
    hdu.writeto((fileName+'.fits'),clobber=True)
    
#save simulation information+star parameters as text file      
def saveParams(fileName,params):
    file1=open((fileName+'.txt'),'w')
    for n in params: 
        file1.write(n+"\n")
    file1.close()
   
     
##### Utility functions

#convert hour angle coordinates to arcseconds. Input RA as "hh mm ss.xx" (hours, minutes, seconds), DEC as "dd mm ss.xx" (degrees, minutes, seconds)
def ConvertCoords(RA1,DEC1):
        
    DEC2=DEC1.split(" ")
    if len(DEC2)==3:
        arcsecDEC=3600.0*(float(DEC2[0])+(float(DEC2[1])/60.0)+(float(DEC2[2])/3600.0))
    elif len(DEC2)==2:
        arcsecDEC=3600.0*(float(DEC2[0])+(float(DEC2[1])/60.0))
    elif len(DEC2)==1:
        arcsecDEC=3600.0*(float(DEC2[0]))
        
    RA2=RA1.split(" ")
    if len(RA2)==3:
        arcsecRA=(3600.0*360.0/24.0)*(float(RA2[0])+(float(RA2[1])/60.0)+(float(RA2[2])/3600.0))
    elif len(RA2)==2:
        arcsecRA=(3600.0*360.0/24.0)*(float(RA2[0])+(float(RA2[1])/60.0))
    elif len(RA2)==1:
        arcsecRA=(3600.0*360.0/24.0)*(float(RA2[0]))
    #arcsecRA=math.cos((arcsecDEC/3600.0)*math.pi/180.0)*arcsecRA #correction for using curved surface projected onto flat field
    return arcsecRA,arcsecDEC

#Calculate no. of stars within divisions of whole field
#def starFieldDivision(nxPixel,nyPixel,xDiv,yDiv,Params):
    #nSubField=np.zeros((nyPixel,nxPixel))
    #xBounds=np.linspace(0,nxPixel,xDiv+1)
    #yBounds=np.linspace(0,nyPixel,yDiv+1)
    #n=0
    #while n<len(Params):
        #
    
    #return 

##### Get star parameters

#generate uniform random positions and flux values for stars     
def GenRandStars(nxPixel,nyPixel,minPeak,maxPeak,minSigma,maxSigma,maxCorr,gain,nStars,randSeed): 
    
    rm.seed(randSeed)
    extend=15 #additional pixel distance from edges to generate out of frame stars
    
    posX=[0 for x in range(nStars)] #x coordinates
    posY=[0 for x in range(nStars)] #y coordinates  
    peak=[0 for x in range(nStars)] #peak value
    flux=[0 for x in range(nStars)] #total flux
    sigX=[0 for x in range(nStars)] #sigma in x axis
    sigY=[0 for x in range(nStars)] #sigma in y axiz
    corr=[0 for x in range(nStars)] #correllation
 
    n=0 #no. of stars generated within frame
    nOut=0 #no. of stars out of frame
    nT=0 #total no. of stars generated (in and out of frame)
    maxFlux=0 #for finding maximum peak/flux
    while n<nStars: #loop until nStars generated in frame
        
        #Random value generation
        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
        peak[n]=rm.random()*(maxPeak-minPeak)+minPeak
        sigX[n]=rm.random()*(maxSigma-minSigma)+minSigma
        sigY[n]=rm.random()*(maxSigma-minSigma)+minSigma
        corr[n]=rm.random()*maxCorr-maxCorr

        flux[n]=peak[n]*(2*math.pi*sigX[n]*sigY[n]*np.sqrt(1-corr[n]**2))*gain

        if posX[n]>=0 and posX[n]<=nxPixel-1 and posY[n]>=0 and posY[n]<=nyPixel-1: #test if star is in frame.
            if flux[n]>maxFlux:
                maxFlux=flux[n]
            n=n+1
            
        else:
            posX.append(posX[n]) #add out of frame stars to end of output list
            posY.append(posY[n])
            peak.append(peak[n])
            sigX.append(sigX[n])
            sigY.append(sigY[n])
            corr.append(corr[n])
            flux.append(flux[n])
            nOut=nOut+1
        nT=nT+1   
    n=0 
    Params=np.zeros((nT,6))
    while n<nT: #fill output array
        Params[n][0]=posX[n]
        Params[n][1]=posY[n]
        Params[n][2]=flux[n]/maxFlux #relative peak to maxPeak, equal to flux/maxFlux
        Params[n][3]=sigX[n]
        Params[n][4]=sigY[n]
        Params[n][5]=corr[n]
        
        n=n+1
    print "Total stars generated: ",nT
    print "No. inside frame: ",(nT-nOut)
    print "No. outside frame: ",nOut
    
    return Params,(nT-nOut),nT,maxFlux
    
#Retrieve star parameters from file returns array with elements: [xPos,yPos,flux]
def GetStarParams(filename,filt,nxPixel,nyPixel,telRAh,telDECh,pScale): #filt (filter) ='b','v','r' ; telRAh/DECh = image centre position, in hour angles (e.g RA='11 22 33.44', DEC='+66 22 33.44')
    
    telRA,telDEC=ConvertCoords(telRAh,telDECh)
    
    extend=15
    
    posX=[] #x coordinates
    posY=[] #y coordinates  
    flux=[] #total flux        
    
    file1=open((filename+'.txt'),'rb') #open and read data file
    starData1=file1.readlines() #Headers: RA,DEC, BMAG, VMAG, RMAG
    file1.close()
    
    if filt is 'b': #set magnitude value to use from parameter table, and set flux for filter at m=0 (from: http://www.astro.utoronto.ca/~patton/astro/mags.html)
        filtI=2 #filter index
        fluxM0=4260#*1.51e7*0.22/300
    elif filt is 'v': 
        filtI=3
        fluxM0=3640#*1.51e7*0.16/300
    elif filt is 'r': 
        filtI=4
        fluxM0=3080#*1.51e7*0.23/300
    
    nStars=0 #for counting stars to be generated (frame+extended volume)
    nFrame=0 #for counting stars only in frame
    maxFlux=0 #use to determine brightest star in image
    n=1 #first line is headers
    while n<len(starData1):
        starData2=starData1[n].split("\t")
        #retrieve J2000 star coordinates, convert to arc seconds, use to gain positions relative to frame center
        starRA,starDEC=ConvertCoords(starData2[0],starData2[1])       
        #positions in terms of pixel scales, relative to frame center 
        relPosX=(float(starRA-telRA)/pScale)*math.cos((starDEC/3600.0)*math.pi/180.0) #*correction - relative RA seperation dependant on declination
        relPosY=float(starDEC-telDEC)/pScale
       
        #relRA2=relRA1*math.cos(rotation)+relDEC1*math.sin(rotation) #should be aligned, therefore rotation=0. Keep here for reference
        #relDEC2=-relRA1*math.sin(rotation)+relDEC1*math.cos(rotation)
        
        if (relPosX > -(0.5*nxPixel)-extend) and (relPosX < (0.5*nxPixel)+extend) and (relPosY > -(0.5*nyPixel)-extend) and (relPosY < (0.5*nyPixel)+extend) and (starData2[filtI] != ("" or " ")): #check if star is within relevent position range, magnitude has value
            posX.append(relPosX+0.5*nxPixel)
            posY.append(relPosY+0.5*nyPixel)
            flux.append(fluxM0*(10.0**(-0.4*float(starData2[filtI])))) #total flux, photons s^-1 m^-2
            
            nStars=nStars+1
            if (relPosX > -0.5*nxPixel) and (relPosX < 0.5*nxPixel) and (relPosY > -0.5*nyPixel) and (relPosY < 0.5*nyPixel): #test if star is in frame (not extra volume)
                nFrame=nFrame+1
                if (fluxM0*(10.0**(-0.4*float(starData2[filtI]))))>maxFlux: #for finding max. flux in frame only
                    maxFlux=(fluxM0*(10.0**(-0.4*float(starData2[filtI]))))
        n=n+1
     
       
    Params=np.zeros((nStars,3)) #defining output array for star data. each column: [xPos,yPos,flux]
    n=0
    while n<nStars: #fill output array
        Params[n][0]=posX[n]
        Params[n][1]=posY[n]
        Params[n][2]=flux[n]/maxFlux #relative flux in relation to maximum
        n=n+1
    
    print "Stars in frame: ", nFrame
    print "Total stars to generate: ", nStars
    return Params,nFrame,nStars
    
    
##### Composite functions for simulating star fields

#Simulate known stars - get parameters from Aladdin       
def SimStarField_Real(inFile,outFile,filt,nxPixel,nyPixel,telRA,telDEC,pScale,maxPeak,expT,gain,minSigma):
      
    Image1=MakeField(nxPixel,nyPixel)
    
    Params,nFrame,nStars=GetStarParams(inFile,filt,nxPixel,nyPixel,telRA,telDEC,pScale)
    
    sigX=minSigma #set for circular symmetry for now - define in relation to skewing and rotation later (more complicated and not exactly necessary)
    sigY=minSigma
    corr=0.0001 #issues with using value of 0 - fix later
    
    maxFlux=maxPeak*(2*math.pi*sigX*sigY*np.sqrt(1-corr**2))*gain
    
    backMean=(917.858645977)+(0.0759400473796)*expT #from dark/bias frame analysis for backgrounds
    backErr=math.sqrt(((6.44315650804)**2)+(0.0364088264173)*expT)
    
    starData=setupOutputs_Real(inFile,outFile,filt,nxPixel,nyPixel,telRA,telDEC,pScale,maxFlux,maxPeak,expT,backMean,backErr,gain,minSigma,nFrame,nStars)
    
    n=0
    while n<len(Params):
        if (Params[n][2]>0): #check flux is nonzero
            Image1=AddStar(Params[n][0],Params[n][1],sigX,sigY,corr,(Params[n][2]*maxFlux),Image1)
        peak=Params[n][2]*maxFlux/((2*math.pi*sigX*sigY*np.sqrt(1-corr**2))*gain) #calculate from relative 
        starData.append(str(n+1)+"\t"+str(Params[n][0])+"\t"+str(Params[n][1])+"\t"+str(Params[n][2]*maxFlux)+"\t"+str(peak)+"\t"+str(sigX)+"\t"+str(sigY)+"\t"+str(corr)) #fill output file data    
        n=n+1
        print "added star: ", n
        
    Image1=Image1/gain #convert photon flux to ADU
    
    Image1=MakePoisson(Image1)
    
    Image1=AddBackground(Image1,backMean,backErr)
    
    Image1=Saturate(Image1)
    
    saveOutputs(outFile,Image1)
    saveParams(outFile,starData)
     
    return Image1, starData

#Randomised star field - parameters are uniformly distributed *still need to modify field parameter output     
def SimStarField_Rand(outFile,nxPixel,nyPixel,nStars,minPeak,maxPeak,minSigma,maxSigma,maxCorr,backMean,backError,gain,randSeed):
    
    rm.seed(randSeed)
    
    extend=15 
    Image1=MakeField(nxPixel,nyPixel) #Create empty arrays for image data
        
    Params,nFrame,nStars,maxFlux=GenRandStars(nxPixel,nyPixel,minPeak,maxPeak,minSigma,maxSigma,maxCorr,gain,nStars,randSeed)
    
    starData=setupOutputs_Rand(outFile,nxPixel,nyPixel,minPeak,maxPeak,minSigma,maxSigma,maxCorr,backMean,backError,gain,randSeed,nFrame,nStars)

    #Loop over all stars
    n=0 #no. of stars generated within frame
    while n<len(Params):
        if (Params[n][2]>0): #check flux is nonzero
            Image1=AddStar(Params[n][0],Params[n][1],Params[n][3],Params[n][4],Params[n][5],(Params[n][2]*maxFlux),Image1)
        peak=Params[n][2]*maxFlux/((2*math.pi*Params[n][3]*Params[n][4]*np.sqrt(1-Params[n][5]**2))*gain) #calculate from relative 
        starData.append(str(n+1)+"\t"+str(Params[n][0])+"\t"+str(Params[n][1])+"\t"+str(Params[n][2]*maxFlux)+"\t"+str(peak)+"\t"+str(Params[n][3])+"\t"+str(Params[n][4])+"\t"+str(Params[n][5])) #fill output file data    
        n=n+1
        print "added star: ", n
        
    Image1=Image1/gain
                
    Image1=MakePoisson(Image1)
    
    Image1=AddBackground(Image1,backMean,backError)
    
    Image1=Saturate(Image1)
    
    return Image1, starData