# Bojan Nikolic <bn204@mrao.cam.ac.uk>, <bojan@bnikolic.co.uk>
#
# Initial version October 2007
#
# Generate kolmogorov screens and save to FITS files for later use

import os
import math

from setup import *

import pyfits
import numarray
import iofits4
import pybnlib


import kolmogorovutils

class ScreenGenerator:

    ""

    """

    Keeping everything in pixel coordinates so as to avoid attaching
    physical scale before that is necessary.
    
    """

    def __init__(self, Nx,Ny,Nz):

        self.g= kolmogorovutils.GenerateKolmogorov3D( Nx,Ny,Nz)

    def GenSimpleScreen(self,
                        dz,
                        height=None,
                        offang=0):

        """
        dz is the thickness

        height: height of the layer midplane in pixel units

        returns numarray with the screen
        """

        zstart= self.g[1][2]/2 - dz/2
        zend  = zstart + dz

        # Note reversal of dimensions to get indexing correct
        res=numarray.zeros( (self.g[1][1],  self.g[1][0]) ,
                            numarray.Float64)

        rp = pybnlib.doubleCvt( int(res.__array_data__[0], 16) )

        ext=kolmogorovutils.MkExtent(self.g)

        if height is None or offang is 0:
            cy=dy=0
        else:
            dy=math.tan(offang)
            cy=(height - self.g[1][2]/2 + zstart)*math.tan(offang)

        pybnlib.SkewFlatten(self.g[0],
                            ext,
                            zstart,
                            zend,
                            0,
                            cy,
                            0,
                            dy,
                            rp)
        res.transpose()
        return res

    def ThicknessPHDU(self, dz,
                      height=None,
                      offang=0):

        "Generate PHDU for screens of varying thickness"

        s1=self.GenSimpleScreen(dz,
                                height=height,
                                offang=offang)
        
        phdu=pyfits.PrimaryHDU( )
        phdu.data=s1

        phdu.header.update("Nz", self.g[1][2])
        phdu.header.add_comment("Nz is the vertical thickness of original volume")

        phdu.header.update("dz", dz)
        phdu.header.add_comment("dz is number of slices co-added to make screen")

        if height is None : height = -1
        phdu.header.update("height", height)
        phdu.header.add_comment("Height is the height of the layer midplane. Negative means none ")

        phdu.header.update("Angle", offang)
        phdu.header.add_comment("Angle is the offset angle from vertical")
        

        iofits4.AddBzrVersionInfo(phdu)

        return phdu


def Thickness_test():

    sg=ScreenGenerator(513, 129, 129)
    
    for x in [ 1, 9, 33 ]:

        p=sg.ThicknessPHDU(x, 30)
        iofits4.Write([p],
                      "temp/test_screen_%g.fits" % x,
                      overwrite=1)

        p=sg.ThicknessPHDU(x, 30, math.pi / 180)
        iofits4.Write([p],
                      "temp/test_screen_%g_offset.fits" % x,
                      overwrite=1)                
        
def ReGenThickness():

    sg=ScreenGenerator(4097, 513, 513)
    
    for x in [ 1, 9, 33, 129, 513 ]:

        p=sg.ThicknessPHDU(x)
        iofits4.Write([p],
                      "temp/ts_%g.fits" % x,
                      overwrite=1)

        p=sg.ThicknessPHDU(x,
                           height=40,
                           offang=(1.5 * math.pi/180)
                           )
        iofits4.Write([p],
                      "temp/ts_%g_offset.fits" % x,
                      overwrite=1)        


def TransposeScreens():

    for x in [ 1, 9, 33, 129, 513 ]:

        f=pyfits.open("temp/ts_%g.fits" % x)
        f[0].data.transpose()

        iofits4.Write(f,
                      "temp/ts_%g_t.fits" % x,
                      overwrite=1)        
        

        f=pyfits.open("temp/ts_%g_offset.fits" % x)    
        f[0].data.transpose()

        iofits4.Write(f,
                      "temp/ts_%g_offset_t.fits"% x,
                      overwrite=1)