import numpy as np

# Galaxies with fixed size and flux
from descwl_shear_sims.galaxies import FixedGalaxyCatalog

from descwl_shear_sims.sim import make_sim

# convenience function to make a PSF
from descwl_shear_sims.psfs import make_fixed_psf

seed = 8312
rng = np.random.RandomState(seed)

ntrial = 2
coadd_dim = 351
buff = 50

for trial in range(ntrial):
    print('trial: %d/%d' % (trial+1, ntrial))

    # galaxy catalog; you can make your own
    galaxy_catalog = FixedGalaxyCatalog(
        rng=rng,
        coadd_dim=coadd_dim,
        buff=buff,
        layout='random',
        mag=25,
        hlr=1.0,
    )

    # make a constant gaussian psf
    psf = make_fixed_psf(psf_type='gauss')

    # generate some simulation data, with a particular shear

    sim_data = make_sim(
        rng=rng,
        galaxy_catalog=galaxy_catalog,
        coadd_dim=coadd_dim,
        g1=0.02,
        g2=0.00,
        psf=psf,
    )

    # get the first i band exposure
    sim_data['band_data']['i'][0]

    # full list of outputs in the dictionary
    # sim_data: dict
    #     band_data: a dict keyed by band name, holding a list of Exposure
    #        objects
    #     coadd_wcs: lsst.afw.geom.makeSkyWcs
    #     psf_dims: (int, int)
    #     coadd_dims: (int, int)
    #     coadd_bbox: lsst.geom.Box2I
    #     bright_info: structured array
    #         fields are
    #         ra, dec: sky position of bright stars
    #         radius_pixels: radius of mask in pixels
    #         has_bleed: bool, True if there is a bleed trail
    #     se_wcs: list of WCS

import numpy as np

# use galaxy models from WeakLensingDeblending.  Note you need
# to get the data for this, see below for downloading instructions
from descwl_shear_sims.galaxies import WLDeblendGalaxyCatalog

# The star catalog class
from descwl_shear_sims.stars import StarCatalog
from descwl_shear_sims.sim import make_sim

# for making a power spectrum PSF
from descwl_shear_sims.psfs import make_ps_psf

# convert coadd dims to SE dims, need for this PSF
from descwl_shear_sims.sim import get_se_dim

seed = 8312
rng = np.random.RandomState(seed)

ntrial = 2
coadd_dim = 351
buff = 50
rotate = True
dither = True

# this is the single epoch image sized used by the sim, we need
# it for the power spectrum psf
se_dim = get_se_dim(coadd_dim=coadd_dim, rotate=rotate, dither=dither)

for trial in range(ntrial):
    print('trial: %d/%d' % (trial+1, ntrial))

    # galaxy catalog; you can make your own
    galaxy_catalog = WLDeblendGalaxyCatalog(
        rng=rng,
        coadd_dim=coadd_dim,
        buff=buff,
    )
    # star catalog; you can make one of these too
    star_catalog = StarCatalog(
        rng=rng,
        coadd_dim=coadd_dim,
        buff=buff,
    )
    # make a power-spectrum PSF, again you can make your own PSF
    psf = make_ps_psf(rng=rng, dim=se_dim)

    # generate some simulation data, with a particular shear,
    # and dithering, rotation, cosmic rays, bad columns, star bleeds
    # turned on.  By sending the star catalog we generate stars and
    # some can be saturated and bleed

    sim_data = make_sim(
        rng=rng,
        galaxy_catalog=galaxy_catalog,
        star_catalog=star_catalog,
        coadd_dim=coadd_dim,
        g1=0.02,
        g2=0.00,
        psf=psf,
        dither=dither,
        rotate=rotate,
        bands=['r', 'i', 'z'],
        noise_factor=0.58,
        cosmic_rays=True,
        bad_columns=True,
        star_bleeds=True,
    )