'''
@brief Reads in MCMC "movie" sample of image files; writes out sorted (pixels in intensity order) in fits form.
@author A. Connors
'''
# @file sort_imagemovie.py
#
# This script is explicitly for:
# reading MC "movies" of images in fits format and sorting pixels in intensity order;
# Finding an extra simple version of confidence bands (i.e. [.10, .90] or [.9545, .6723, .5, .3287, .0455]);
#      and writing out 'images' for each confidence band
# and writing them out in fits form compatible with
# CHASC imaging software.
#

from numpy import *
import pyfits

class SortMovieMatrix:
    def __init__(self,InFitsMovieFileName,NumBurnIn):
#
# 0. Initialize:
        print 'Reading ',InFitsMovieFileName,' with burn-in number:', NumBurnIn
#
# 0.1   Read in fits movie matrix:
        image_inHDULst = pyfits.open(InFitsMovieFileName)
        print 'Header: \n', image_inHDULst[0].header
        naxis3,naxis2,naxis1 = image_inHDULst[0].data.shape
        print naxis3,naxis2,naxis1
        InImage = zeros((naxis3-NumBurnIn,naxis2,naxis1))

#
# 1. Try sorting each pixel by intensity:
        InImage[:][:][:] = image_inHDULst[0].data[NumBurnIn:][:][:]
        SortImage = sort(InImage,axis=0)

        self.data = SortImage
        self.old_header = image_inHDULst[0].header
#        print SortImage
#        print self.data
        image_inHDULst.close()
#etc

class OverlySimpleImageConfBands:
    def __init__(self,SortImageHDU,OutFitsFileName,IntensConfidences=[.10 , .50, .90]):
        numconfbands = len(IntensConfidences)
        naxis3,naxis2,naxis1 = SortImageHDU.data.shape
        print 'Shape of input Sorted MC Movie Images: ',SortImageHDU.data.shape
        self.data    = zeros((numconfbands,naxis2,naxis1))
        TotImages = float(naxis3)

        ll = -1
        howmany = []
        for confvalue in IntensConfidences:
            ll = ll + 1
            print' For ll=',ll,' and confvalue=',confvalue,' Totimages*confvalue is:',TotImages*confvalue
            howmany.append(0.)
            for thisindex in range(naxis3):
                if( abs( confvalue - 0.5 ) < 1.e-6 ):
                    print 'At average piece:'
                    try:
                        self.data[ll] = SortImageHDU.data.mean(axis=0)
                    except:
                        print ' ImageMovie Exception! at: ', ll, UnsortedMovieImageHDU.data
                    print ' \n'
                elif( confvalue < 0.5 ):
                    if( thisindex <= TotImages*confvalue ):
                        self.data[ll] += SortImageHDU.data[thisindex]
                        howmany[ll]   += 1.
                elif( confvalue > 0.5 ):
                    if( thisindex >= TotImages*confvalue ):
                        self.data[ll] += SortImageHDU.data[thisindex]
                        howmany[ll]   += 1.
                # end if
        for lll in range( len(IntensConfidences) ):
            if( howmany[lll] > 0. ):
                self.data[lll] = self.data[lll]/howmany[lll]
        # end for

        image_HDU = pyfits.PrimaryHDU(self.data)
        image_HDU.header.update('NAXIS3',numconfbands,comment='Confidences: '+str(IntensConfidences))

        image_HDU.writeto(OutFitsFileName)
#
#   End.

class SecondSimpleImageConfBands:
    def __init__(self,UnsortedMovieImageHDU,OutFitsFileName,IntensConfidences=[.10, .50 , .90]):
#
# 1/ Initialize handy variables, including ExpectedMSTotalCnts:
#
        numconfbands = len(IntensConfidences)		# Output is confidence bands, plus mean in the middle if conf=0.5
        naxis3,naxis2,naxis1 = UnsortedMovieImageHDU.data.shape
        print 'Shape of input Sorted MC Movie Images: ',UnsortedMovieImageHDU.data.shape
        self.data    = zeros((numconfbands,naxis2,naxis1))
        TotImages = float(naxis3)
        ExpectedMSTotalCnts = []
        for imag in range(naxis3):
            sumit = 0.
            for ll in range(naxis2):
                sumit += sum( UnsortedMovieImageHDU.data[imag][ll][:] )
            print 'imag,sumit: ', imag, ' , ', sumit, '\n'
            ExpectedMSTotalCnts.append(sumit)

#       Next check what ExpectedMSCnts values correspond to each conf band:
        SortedExpectedMSCnts = sort(ExpectedMSTotalCnts)
        CountsConf = []
        for ll in range(numconfbands):
            lin = ll
            lout = ll
            thisindex = int(IntensConfidences[lin]*TotImages)
            CountsConf.append(SortedExpectedMSCnts[thisindex])
            confvalue = IntensConfidences[lin]
            print ' For ll,lin=',ll,lin,' and confvalue=',confvalue,' Totimages*confvalue is:',TotImages*confvalue
            print ' Corresponds to counts: ', CountsConf[lin]
#       end for
        for ll in range(numconfbands):
            lin = ll
            lout = ll
            confvalue = IntensConfidences[lin]
            if (abs( confvalue - 0.5 ) < 1.e-6):
                print 'At average piece:, lin,lout is ', lin, lout
                try:
                    self.data[lout] = UnsortedMovieImageHDU.data.mean(axis=0)
                except:
                    print ' ImageMovie Exception! at: ', lout, UnsortedMovieImageHDU.data
                print ' \n'
            elif( confvalue < 0.5 ):
                print 'At lower limit:, lin,lout is ', lin, lout
                howmany = 0.
                for thisindex in range(naxis3):
                    if ( SortedExpectedMSCnts[thisindex] <= CountsConf[lin] ):
                        self.data[lout] += UnsortedMovieImageHDU.data[thisindex]
                        howmany += 1.
                # Renormalize:
                self.data[lout] = self.data[lout]/howmany
            elif( confvalue > 0.5 ):
                print 'At upper limit:, lin,lout is ', lin, lout
                howmany = 0.
                for thisindex in range(naxis3):
                    if ( SortedExpectedMSCnts[thisindex] >= CountsConf[lin] ):
                        self.data[lout] += UnsortedMovieImageHDU.data[thisindex]
                        howmany += 1.
                # Renormalize:
                self.data[lout] = self.data[lout]/howmany
            print ' \n \n'
    # Now renormalize:

        image_HDU = pyfits.PrimaryHDU(self.data)
        image_HDU.header.update('NAXIS3',numconfbands,comment='Confidences: '+str(IntensConfidences))

        image_HDU.writeto(OutFitsFileName)
        self.header = image_HDU.header
#
#   End.