"""
Cosmic Rays
===========

This simple class can be used to include cosmic ray events to an image.
By default the cosmic ray events are drawn from distributions describing
the length and energy of the events. Such distributions can be generated
for example using Stardust code (http://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=04636917).
The energy of the cosmic ray events can also be set to constant for
testing purposes. The class can be used to draw a single cosmic ray
event or up to a covering fraction.

:requires: NumPy
:requires: SciPy

:version: 0.2

"""
import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline


class cosmicrays():
    """
    Cosmic ray generation class. Can either draw events from distributions or
    set the energy of the events to a constant.

    :param log: logger instance
    :param image: image to which cosmic rays are added to (a copy is made not to change the original numpy array)
    :param crInfo: column information (cosmic ray file)
    :param information: cosmic ray track information (file containing track length and energy information) and
                        exposure time.
    """
    def __init__(self, log, image, exptime, crInfo=None, information=None):
        """
        Cosmic ray generation class. Can either draw events from distributions or
        set the energy of the events to a constant.

        :param log: logger instance
        :param image: image to which cosmic rays are added to (a copy is made not to change the original numpy array)
        :param crInfo: column information (cosmic ray file)
        :param information: cosmic ray track information (file containing track length and energy information) and
                            exposure time.
        """
               
        self.exptime=exptime
        
        self.log = log

        #image and size
        self.image = image.copy()
        self.ysize, self.xsize = self.image.shape

        if crInfo is not None:
            self.cr = crInfo
        else:
            self._readCosmicrayInformation()


##################################################


    def _cosmicRayIntercepts(self, lum, x0, y0, l, phi):
        """
        Derive cosmic ray streak intercept points.

        :param lum: luminosities of the cosmic ray tracks
        :param x0: central positions of the cosmic ray tracks in x-direction
        :param y0: central positions of the cosmic ray tracks in y-direction
        :param l: lengths of the cosmic ray tracks
        :param phi: orientation angles of the cosmic ray tracks

        :return: cosmic ray map (image)
        :rtype: nd-array
        """
        #create empty array
        crImage = np.zeros((self.ysize, self.xsize), dtype=np.float64)

        #x and y shifts
        dx = l * np.cos(phi) / 2.
        dy = l * np.sin(phi) / 2.
        mskdx = np.abs(dx) < 1e-8
        mskdy = np.abs(dy) < 1e-8
        dx[mskdx] = 0.
        dy[mskdy] = 0.

        #pixels in x-direction
        ilo = np.round(x0.copy() - dx)
        msk = ilo < 0.
        ilo[msk] = 0
        ilo = ilo.astype(int)

        ihi = 1 + np.round(x0.copy() + dx)
        msk = ihi > self.xsize
        ihi[msk] = self.xsize
        ihi = ihi.astype(int)

        #pixels in y-directions
        jlo = np.round(y0.copy() - dy)
        msk = jlo < 0.
        jlo[msk] = 0
        jlo = jlo.astype(int)

        jhi = 1 + np.round(y0.copy() + dy)
        msk = jhi > self.ysize
        jhi[msk] = self.ysize
        jhi = jhi.astype(int)

        #loop over the individual events
        for i, luminosity in enumerate(lum):
            n = 0  # count the intercepts

            u = []
            x = []
            y = []

            #Compute X intercepts on the pixel grid
            if ilo[i] < ihi[i]:
                for xcoord in range(ilo[i], ihi[i]):
                    ok = (xcoord - x0[i]) / dx[i]
                    if np.abs(ok) <= 0.5:
                        n += 1
                        u.append(ok)
                        x.append(xcoord)
                        y.append(y0[i] + ok * dy[i])
            else:
                for xcoord in range(ihi[i], ilo[i]):
                    ok = (xcoord - x0[i]) / dx[i]
                    if np.abs(ok) <= 0.5:
                        n += 1
                        u.append(ok)
                        x.append(xcoord)
                        y.append(y0[i] + ok * dy[i])

            #Compute Y intercepts on the pixel grid
            if jlo[i] < jhi[i]:
                for ycoord in range(jlo[i], jhi[i]):
                    ok = (ycoord - y0[i]) / dy[i]
                    if np.abs(ok) <= 0.5:
                        n += 1
                        u.append(ok)
                        x.append(x0[i] + ok * dx[i])
                        y.append(ycoord)
            else:
                for ycoord in range(jhi[i], jlo[i]):
                    ok = (ycoord - y0[i]) / dy[i]
                    if np.abs(ok) <= 0.5:
                        n += 1
                        u.append(ok)
                        x.append(x0[i] + ok * dx[i])
                        y.append(ycoord)

            #check if no intercepts were found
            if n < 1:
                xc = int(np.floor(x0[i]))
                yc = int(np.floor(y0[i]))
                crImage[yc, xc] += luminosity

            #Find the arguments that sort the intersections along the track
            u = np.asarray(u)
            x = np.asarray(x)
            y = np.asarray(y)

            args = np.argsort(u)

            u = u[args]
            x = x[args]
            y = y[args]

            #Decide which cell each interval traverses, and the path length
            for i in range(1, n - 1):
                w = u[i + 1] - u[i]
                cx = int(1 + np.floor((x[i + 1] + x[i]) / 2.))
                cy = int(1 + np.floor((y[i + 1] + y[i]) / 2.))

                if 0 <= cx < self.xsize and 0 <= cy < self.ysize:
                    crImage[cy, cx] += (w * luminosity)

        return crImage


############################################################################


    def _drawEventsToCoveringFactor(self, coveringFraction=3.0, limit=1000, verbose=False):
        """
        Generate cosmic ray events up to a covering fraction and include it to a cosmic ray map (self.cosmicrayMap).

        :param coveringFraction: covering fraction of cosmic rya events in per cent of total number of pixels
        :type coveringFraction: float
        :param limit: limiting energy for the cosmic ray event [None = draw from distribution]
        :type limit: None or float
        :param verbose: print out information to stdout
        :type verbose: bool


        :return: None
        """
        self.cosmicrayMap = np.zeros((self.ysize, self.xsize))

        #how many events to draw at once, too large number leads to exceeding the covering fraction
        cr_n = int(295 * self.exptime / 565. * coveringFraction / 1.4)

        covering = 0.0

        while covering < coveringFraction:
            #pseudo-random numbers taken from a uniform distribution between 0 and 1
            np.random.seed()
            luck = np.random.rand(cr_n)

            #draw the length of the tracks
            ius = InterpolatedUnivariateSpline(self.cr['cr_cdf'], self.cr['cr_u'])
            self.cr['cr_l'] = ius(luck)

            if limit is None:
                ius = InterpolatedUnivariateSpline(self.cr['cr_cde'], self.cr['cr_v'])
                self.cr['cr_e'] = ius(luck)
            else:
                #set the energy directly to the limit
                self.cr['cr_e'] = np.asarray([limit,])

            #Choose the properties such as positions and an angle from a random Uniform dist
            np.random.seed()
            cr_x = self.xsize * np.random.rand(int(np.floor(cr_n)))
            
            np.random.seed()
            cr_y = self.ysize * np.random.rand(int(np.floor(cr_n)))
            
            np.random.seed()
            cr_phi = np.pi * np.random.rand(int(np.floor(cr_n)))

            #find the intercepts
            self.cosmicrayMap += self._cosmicRayIntercepts(self.cr['cr_e'], cr_x, cr_y, self.cr['cr_l'], cr_phi)

            #count the covering factor
            area_cr = np.count_nonzero(self.cosmicrayMap)
            covering = 100.*area_cr / (self.xsize*self.ysize)

            # text = 'The cosmic ray covering factor is %i pixels i.e. %.3f per cent' % (area_cr, covering)
            # self.log.info(text)

###################################################33

    def addUpToFraction(self, coveringFraction, limit=None, verbose=False):
        """
        Add cosmic ray events up to the covering Fraction.

        :param coveringFraction: covering fraction of cosmic rya events in per cent of total number of pixels
        :type coveringFraction: float
        :param limit: limiting energy for the cosmic ray event [None = draw from distribution]
        :type limit: None or float
        :param verbose: print out information to stdout
        :type verbose: bool

        :return: image with cosmic rays
        :rtype: ndarray
        """
        self._drawEventsToCoveringFactor(coveringFraction, limit=limit, verbose=verbose)

        #paste cosmic rays
        self.image += self.cosmicrayMap

        return self.image




