import numpy as np
import galsim

import Grid
import mathutil as util

class InjectionCatalog(object):

    def __init__(self, image):
        self.ramin, self.ramax = image.ramin, image.ramax
        self.decmin, self.decmax = image.decmin, image.decmax
        self.ra_boundary_cross = image.ra_boundary_cross
        self.Npix_x, self.Npix_y = image.Npix_x, image.Npix_y
        self.pixel_scale = image.pixel_scale
        self.wcs = image.wcs
        print(self.wcs)

        # This is also set for a given image, as it may depend on its area
        self.objs_per_real = image.objs_per_real

        # TODO: output truth catalog after injection
        self.truth_outfile = {}
        pass

    def generate_positions(self, config):
        
        bg = Grid.BaseGrid()

        self.nobjects = self.objs_per_real
        ps = config['pos_sampling']
        pstype = ps['type']
        
        if pstype == 'uniform' and (self.objs_per_real is not None):
            ra = util.sample_uniform_ra(self.ramin, self.ramax, self.objs_per_real, 
                                        boundary_cross=self.ra_boundary_cross)
            dec = util.sample_uniform_dec(self.decmin, self.decmax, self.objs_per_real, unit='deg')
            self.pos = np.column_stack((ra, dec))
        
        elif pstype in bg._valid_grid_types:
            grid_kwargs = self._build_grid_kwargs(pstype, ps)
            gtype = pstype
            
            img_grid = Grid._build_grid(gtype, **grid_kwargs)
            self.pos = img_grid.pos
            self.nobjects = np.shape(img_grid.pos)[0]
        
        elif pstype in bg._valid_mixed_types:
            pass
        
        else:
            raise ValueError('Position sampling type {} is not valid!'.format(gtype))

        # Generate object indices (in input catalog)
        # TODO
        # self.index = np.random.choice(xrange(input_nobjects), size=self.nobjects)

        # Generate object rotation angles, if desired
        if config['rotate_objs'] is True:
            rot = util.sample_uniform(0., 360., self.nobjects)
            self.rotate = np.array([str(r) + ' deg' for r in rot])
        else:
            self.rotate = None


    def _build_grid_kwargs(self, pstype, ps):
        if pstype == 'MixedGrid':
            gtype = ps['grid_type']
        else:
            gtype = pstype

        gs = ps['grid_spacing']
        
        try:
            r = ps['rotate']
            if (isinstance(r, str) and (r.lower() == 'random')):
                if gtype == 'RectGrid':
                    self.grid_rot_angle = np.random.uniform(0., np.pi/2.)
                elif gtype == 'HexGrid':
                    self.grid_rot_angle = np.random.uniform(0., np.pi/3.)
            else:
                unit = ps['angle_unit']
                if unit == 'deg':
                    if (r >= 0.) and (r < 360.):
                        self.grid_rot_angle = float(r)
                    else:
                        raise ValueError('Grid rotation of {} '.format(r) + 'deg is not valid!')
                else:
                    if (r >= 0.) and (r < 2*np.pi):
                        self.grid_rot_angle = float(r)
                    else:
                        raise ValueError('Grid rotation of {} '.format(r) + 'rad is not valid!')
        except KeyError:
            self.grid_rot_angle = 0.0

        # Offset grid if desired
        try:
            offset = ps['offset']
            if (isinstance(offset, str)) and (offset.lower() == 'random'):
                self.grid_offset = [np.random.uniform(-gs/2., gs/2.),
                                    np.random.uniform(-gs/2., gs/2.)]
            else:
                if isinstance(offset, list):
                    self.grid_offset = list(offset)
                else:
                    raise ValueError('Grid offset of {} '.format(offset) + 'is not valid!')
        except KeyError:
            self.grid_offset = [0., 0.]

        try:
            self.angle_unit = ps['angle_unit']
        except KeyError:
            self.angle_unit = 'rad'

        grid_kwargs = dict(grid_spacing=gs,
                        wcs=self.wcs,
                        Npix_x=self.Npix_x,
                        Npix_y=self.Npix_y,
                        pixelscale=self.pixel_scale,
                        rot_angle=self.grid_rot_angle,
                        angle_unit=self.angle_unit,
                        pos_offset=self.grid_offset)
        return grid_kwargs
