from __future__ import absolute_import

import collections
import math
import numpy as np
import scipy.constants as sc

def black_body(temperature, wavenumber):
    """Function to calculate Planck function. 
       temperature - Kelvin
       wavenumber - frequency in cm-1
    """
    freq = wavenumber * sc.c * 100.0
    jnu = 2.0 * sc.h * pow(freq,3) / (pow(sc.c,2) *
      (math.exp((sc.h * freq) / (sc.k * temperature)) - 1.0))

    return jnu


class SkyGenerator(object):
    """Class to generate a model sky.
    """

    def __init__(self, parameters, previous_results):
        self.parameters = parameters
        self.previous_results = previous_results
        self.result = collections.OrderedDict()

    def run(self):
        print 'SkyGenerator.run'

        fts = self.previous_results['fts']
        frequency_axis = fts['fts_wn']
        nspec = len(frequency_axis)

        beamsgenerator = self.previous_results['beamsgenerator']
        npix = beamsgenerator['npix']
        spatial_axis = beamsgenerator['spatial axis [arcsec]']

        # oversampling factor
#        oversample = 10
        oversample = 1

#        print 'creating sky model with spatial dims [%s,%s] spectral dim %s',\
#          (oversample * npix, oversample * npix, nspec)
#        print 'spatial oversampling: %s' % oversample
        skymodel = np.zeros([npix, npix, nspec], np.float)

        sky = self.parameters['substages']['Sky']
        columns = sky['SourceNum'].keys()

        self.result['sources'] = collections.OrderedDict()

        for column in columns:
            temp = sky['SourceNum'][column]
            sourcenum = int(round(temp))
            if sourcenum not in self.result['sources'].keys():
                self.result['sources'][sourcenum] = {}

            type = sky['Type'][column]

            temp = sky['x pos [asec]'][column]
            xpos = float(temp)

            temp = sky['y pos [asec]'][column]
            ypos = float(temp)

            temp = sky['Temp'][column]
            temperature = float(temp)

            temp = sky['cutoffmin'][column]
            cutoffmin = float(temp)

            temp = sky['cutoffmax'][column]
            cutoffmax = float(temp)

            temp = sky['emissivity'][column]
            emissivity = float(temp)

            print 'generating source:%s type:%s xpos:%s ypos:%s temperature:%s cutoffmin:%s cutoffmax:%s e:%s' % (
              sourcenum, type, xpos, ypos, temperature, cutoffmin, cutoffmax,
              emissivity)

            if type.upper().strip() == 'POINT':
                print 'creating point source'
                source_spectrum = self._create_point_source(xpos, ypos,
                  temperature, 
                  cutoffmin, cutoffmax, emissivity, skymodel, 
                  spatial_axis, frequency_axis)
            else:
                source_spectrum = None
                print "source type '%s' not yet implemented" % type

            self.result['sources'][sourcenum]['spectrum'] = source_spectrum

        print 'sources', self.result['sources']
        self.result['sky model'] = skymodel
        self.result['spatial axis'] = spatial_axis
        self.result['frequency axis'] = frequency_axis 

        return self.result

    def _create_point_source(self, xpos, ypos, temperature, cutoffmin,
      cutoffmax, emissivity, skymodel, spatial_axis, frequency_axis):
        """Create a point source.
        """

        # calculate xpos, ypos in units of pixel - numpy arrays [row,col]
        nx = len(spatial_axis)
        colpos = float(nx-1) * float (xpos - spatial_axis[0]) / (spatial_axis[-1] - spatial_axis[0])
        rowpos = float(nx-1) * float (ypos - spatial_axis[0]) / (spatial_axis[-1] - spatial_axis[0])

        if colpos < 0 or colpos > (nx-1) or rowpos < 0 or rowpos > (nx-1):
            # point source is outside modelled area 
            return

        # calculate fourier phase shift to move point at [0,0] to [rowpos, colpos]
        shiftx = np.zeros([nx], np.complex)
        shiftx[:nx/2] = np.arange(nx/2, dtype=np.complex)
        shiftx[nx/2:] = np.arange(-nx/2, 0, dtype=np.complex)
        shiftx = np.exp((-2.0j * np.pi * colpos * shiftx) / float(nx))

        shifty = np.zeros([nx], np.complex)
        shifty[:nx/2] = np.arange(nx/2, dtype=np.complex)
        shifty[nx/2:] = np.arange(-nx/2, 0, dtype=np.complex)
        shifty = np.exp((-2.0j * np.pi * rowpos * shifty) / float(nx))

        shift = np.ones([nx,nx], np.complex)
        for j in range(nx):
            shift[j,:] *= shiftx
        for i in range(nx):
            shift[:,i] *= shifty

        # calculate black-body spectrum, modified for cutoffs in 
        # sensitivity and source emissivity
        bbmod = np.zeros(np.shape(frequency_axis))
        print 'frequency axis', frequency_axis
        # ignore floating point errors
        old_settings = np.seterr(all='ignore')
        for iwn,wn in enumerate(frequency_axis):
            # simulate real-life 'rounded' cutoffs numerically  
            f1 = 1.0 / (1.0 + pow(cutoffmin/wn, 18) + pow(wn/cutoffmax, 24))
            print pow(cutoffmin/wn,18), pow(wn/cutoffmax,24)
            print temperature, wn, black_body(temperature, wn), f1, emissivity
            bbmod[iwn] = black_body(temperature, wn) * f1 * emissivity
        # restore fp behaviour
        ignore = np.seterr(**old_settings)

        # go through freq planes and add point source to each
        for iwn,wn in enumerate(frequency_axis):
            # create point in frequency space
            temp = np.zeros([nx,nx])
            temp[0,0] = bbmod[iwn]
            # 2d fft
            temp = np.fft.fft2(temp)
            # apply phase shift to move point to required offset
            temp *= shift
            # transform back
            temp = np.fft.ifft2(temp)
            temp = np.real(temp)

            # add to sky model
            skymodel[:,:,iwn] += temp

        return bbmod

    def __repr__(self):
        return 'SkyGenerator'