"""!
Adopted from https://github.com/ujscjj/DPTNet
Modified by Yi Luo {yl3364@columbia.edu}
"""

import numpy as np
import os
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.nn.modules.module import Module
from torch.nn.modules.activation import MultiheadAttention
from torch.nn.modules.container import ModuleList
from torch.nn.init import xavier_uniform_
from torch.nn.modules.dropout import Dropout
from torch.nn.modules.linear import Linear
from torch.nn.modules.rnn import LSTM
from torch.nn.modules.normalization import LayerNorm

class TransformerOptimizer(object):
    """A simple wrapper class for learning rate scheduling"""

    def __init__(self, optimizer, k, d_model, warmup_steps=4000):
        self.optimizer = optimizer
        self.k = k
        self.init_lr = d_model ** (-0.5)
        self.warmup_steps = warmup_steps
        self.step_num = 0
        self.epoch = 0
        self.visdom_lr = None

    def zero_grad(self):
        self.optimizer.zero_grad()

    def step(self, epoch):
        self._update_lr(epoch)
        # self._visdom()
        self.optimizer.step()

    def _update_lr(self, epoch):
        self.step_num += 1
        if self.step_num <= self.warmup_steps:
            lr = self.k * self.init_lr * min(self.step_num ** (-0.5),
                                             self.step_num * (self.warmup_steps ** (-1.5)))
        else:
            lr = 0.0004 * (0.98 ** ((epoch-1)//2))

        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr

    def load_state_dict(self, state_dict):
        self.optimizer.load_state_dict(state_dict)

    def state_dict(self):
        return self.optimizer.state_dict()

    def set_k(self, k):
        self.k = k

    def set_visdom(self, visdom_lr, vis):
        self.visdom_lr = visdom_lr  # Turn on/off visdom of learning rate
        self.vis = vis  # visdom enviroment
        self.vis_opts = dict(title='Learning Rate',
                             ylabel='Leanring Rate', xlabel='step')
        self.vis_window = None
        self.x_axis = torch.LongTensor()
        self.y_axis = torch.FloatTensor()

    def _visdom(self):
        if self.visdom_lr is not None:
            self.x_axis = torch.cat(
                [self.x_axis, torch.LongTensor([self.step_num])])
            self.y_axis = torch.cat(
                [self.y_axis, torch.FloatTensor([self.optimizer.param_groups[0]['lr']])])
            if self.vis_window is None:
                self.vis_window = self.vis.line(X=self.x_axis, Y=self.y_axis,
                                                opts=self.vis_opts)
            else:
                self.vis.line(X=self.x_axis, Y=self.y_axis, win=self.vis_window,
                              update='replace')

class TransformerEncoderLayer(nn.Module):
    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
    This standard encoder layer is based on the paper "Attention Is All You Need".
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
    in a different way during application.
    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of intermediate layer, relu or gelu (default=relu).
    Examples::
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
        >>> src = torch.rand(10, 32, 512)
        >>> out = encoder_layer(src)
    """

    def __init__(self, d_model, nhead, dim_feedforward=256, dropout=0, activation="relu"):
        super(TransformerEncoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        # self.linear1 = Linear(d_model, dim_feedforward)
        self.linear1 = LSTM(d_model, d_model*2, 1, bidirectional=True)
        self.dropout = Dropout(dropout)
        # self.linear2 = Linear(dim_feedforward, d_model)
        self.linear2 = Linear(d_model*2*2, d_model)

        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super(TransformerEncoderLayer, self).__setstate__(state)

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor
        r"""Pass the input through the encoder layer.
        Args:
            src: the sequnce to the encoder layer (required).
            src_mask: the mask for the src sequence (optional).
            src_key_padding_mask: the mask for the src keys per batch (optional).
        Shape:
            see the docs in Transformer class.
        """
        src2 = self.self_attn(src, src, src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        # src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src)[0])))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src

def _get_clones(module, N):
    return ModuleList([copy.deepcopy(module) for i in range(N)])

def _get_activation_fn(activation):
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu

    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))

class SingleTransformer(nn.Module):

    def __init__(self, input_size, dim_feedforward):
        super(SingleTransformer, self).__init__()

        self.transformer = TransformerEncoderLayer(d_model=input_size, nhead=4, dim_feedforward=dim_feedforward, dropout=0)

    def forward(self, input):
        # input shape: batch, seq, dim
        output = input
        output = self.transformer(output.permute(1, 0, 2).contiguous()).permute(1, 0, 2).contiguous()
        return output

# dual-path Transformer
class DPTNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size,
                 dropout=0, num_layers=1):
        super(DPTNet, self).__init__()

        self.input_size = input_size
        self.output_size = output_size
        self.hidden_size = hidden_size

        # dual-path Transformer
        self.row_xfmr = nn.ModuleList([])
        self.col_xfmr = nn.ModuleList([])
        for i in range(num_layers):
            self.row_xfmr.append(SingleTransformer(input_size=input_size, dim_feedforward=hidden_size))
            self.col_xfmr.append(SingleTransformer(input_size=input_size, dim_feedforward=hidden_size))

        self.output = nn.Conv2d(input_size, output_size, 1)

    def forward(self, input):
        # input shape: batch, N, dim1, dim2
        # apply RNN on dim1 first and then dim2

        batch_size, _, dim1, dim2 = input.shape
        output = input
        for i in range(len(self.row_xfmr)):
            row_input = output.permute(0,3,2,1).contiguous().view(batch_size*dim2, dim1, -1)  # B*dim2, dim1, N
            row_output = self.row_xfmr[i](row_input)  # B*dim2, dim1, H
            row_output = row_output.view(batch_size, dim2, dim1, -1).permute(0,3,2,1).contiguous()  # B, N, dim1, dim2
            output = output + row_output

            col_input = output.permute(0,2,3,1).contiguous().view(batch_size*dim1, dim2, -1)  # B*dim1, dim2, N
            col_output = self.col_xfmr[i](col_input)  # B*dim1, dim2, H
            col_output = col_output.view(batch_size, dim1, dim2, -1).permute(0,3,1,2).contiguous()  # B, N, dim1, dim2
            output = output + col_output

        output = self.output(output)

        return output


'''
Some tools to conveniently transform lmfit.Parameters into GalSim.Image's.
'''

import copy

import numpy as np
from scipy.optimize import newton
import galsim
import lmfit

import chroma

class GalTool(object):
    ''' Some generic utilities for drawing ringtest images using GalSim and measuring second moment
    radii.
    '''
    def __init__(self):
        # Subclasses of GalTool must initialize the following:
        #
        #   attributes
        #   ----------
        #   stamp_size - Integer number of pixels in which to draw images
        #   pixel_scale - arcsec / pixel
        #   PSF - either a ChromaticObject or an effective PSF as a GSObject.
        #   offset - tuple defining subpixel offset of image origin from center
        #   gsparams - galsim.GSParams instance defining parameters for GalSim.
        #
        #   methods
        #   -------
        #   _gparam_to_galsim - turn lmfit.Parameters into a galsim.GSObject or
        #                       galsim.ChromaticObject
        raise NotImplementedError("ABC GalTool must be instatiated through a subclass.")

    def get_image(self, gparam, ring_beta=None, ring_shear=None, oversample=1):
        ''' Draw a galaxy image using GalSim.  Potentially rotate and shear the galaxy as part of a
        ring test.  Optionally draw a high-resolution image.

        @param gparam      An lmfit.Parameters object that will be used to initialize a GalSim
                           object.
        @param ring_beta   Angle around ellipticity ring in ring test.
        @param ring_shear  galsim.Shear to apply after rotation as part of ring test.
        @param oversample  Integer factor by which to scale output image resolution and size.
        @returns  galsim.Image
        '''
        # Setup image
        stamp_size = self.stamp_size * oversample
        pixel_scale = self.pixel_scale / float(oversample)
        im = galsim.ImageD(stamp_size, stamp_size, scale=pixel_scale)
        # Get GalSim parameterization of the galaxy.
        gal = self._gparam_to_galsim(gparam)
        # Rotate and shear the galaxy as per the ring test specifications.
        if ring_beta is not None:
            gal = gal.rotate(ring_beta / 2.0 * galsim.radians)
        if ring_shear is not None:
            gal = gal.shear(ring_shear)
        # Convolve with the (chromatic, or effective) PSF.  This will need to be different for the
        # double Sersic case.
        final = galsim.Convolve(gal, self.PSF)
        # And draw the image.
        if isinstance(final, galsim.ChromaticObject):
            final.drawImage(self.bandpass, image=im, offset=self.offset)
        elif isinstance(final, galsim.GSObject):
            final.drawImage(image=im, offset=self.offset)
        else:
            raise ValueError("Don't recognize galaxy object type in GalTool.")
        return im

    def get_PSF_image(self, oversample=1, method='auto'):
        ''' Draw an image of the effective PSF.  Note that the returned PSF includes convolution
        by the pixel response function by default, though this can be overriden using the method
        keywords (e.g., method='fft').

        @param oversample  Integer factor by which to scale output image resolution and size.
        @param method      Method string to pass to galsim drawImage command.
        @returns  galsim.Image
        '''
        stamp_size = self.stamp_size * oversample
        pixel_scale = self.pixel_scale / float(oversample)
        im = galsim.ImageD(stamp_size, stamp_size, scale=pixel_scale)
        if isinstance(self.PSF, galsim.ChromaticObject):
            star = galsim.Gaussian(fwhm=1.e-8) * self.SED
            final = galsim.Convolve(star, self.PSF)
            final.drawImage(self.bandpass, image=im, method=method)
        elif isinstance(self.PSF, galsim.GSObject):
            self.PSF.drawImage(image=im, method=method)
        else:
            raise ValueError("Don't recognize galaxy object type.")
        return im

    def get_r2(self, gparam, oversample=1):
        ''' Compute object second moment radius sqrt(r^2) directly from image.  This may be biased
        if the object wings are significant or the postage stamp size is too small.

        @param gparam   An lmfit.Parameters object that will be used to initialize a GalSim object.
        @returns        Second moment radius (in arcsec)
        '''
        im = self.get_image(gparam, oversample=oversample)
        mx, my, mxx, myy, mxy = chroma.moments(im)
        return np.sqrt(mxx + myy)

    def get_uncvl_image(self, gparam, ring_beta=None, ring_shear=None, oversample=1, center=False):
        ''' Draw a galaxy image, not convolved with a PSF, using GalSim.  Potentially rotate and
        shear the galaxy as part of a ring test.  Optionally draw a high-resolution image.

        @param gparam      An lmfit.Paramters object that will be used to initialize a GalSim object.
        @param ring_beta   Angle around ellipticity ring in ring test.
        @param ring_shear  Shear to apply after rotation as part of ring test. (type=?)
        @param oversample  Integer factor by which to scale output image resolution and size.
        @param center      Force center of profile to (0,0).
        @returns  galsim.Image
        '''
        stamp_size = self.stamp_size * oversample
        pixel_scale = self.pixel_scale / float(oversample)
        im = galsim.ImageD(stamp_size, stamp_size, scale=pixel_scale)
        gal = self._gparam_to_galsim(gparam)
        if center:
            centroid = gal.centroid(self.bandpass)
            gal = gal.shift(-centroid)
        if ring_beta is not None:
            gal = gal.rotate(ring_beta / 2.0 * galsim.radians)
        if ring_shear is not None:
            gal = gal.shear(ring_shear)
        if isinstance(gal, galsim.ChromaticObject):
            gal.drawImage(self.bandpass, image=im, offset=self.offset)
        elif isinstance(gal, galsim.GSObject):
            gal.drawImage(image=im, offset=self.offset)
        else:
            raise ValueError("Don't recognize galaxy object type in GalTool.")
        return im

    def get_uncvl_r2(self, gparam, oversample=1):
        ''' Compute object second moment radius directly from image.  This may be biased if the
        object wings are significant or the postage stamp size is too small.

        @param gparam   An lmfit.Parameters object that will be used to initialize a GalSim object.
        @returns        Second moment radius (in arcsec)
        '''
        im = self.get_uncvl_image(gparam, oversample=oversample)
        mx, my, mxx, myy, mxy = chroma.moments(im)
        return np.sqrt(mxx + myy)

    def compute_AHM(self, gparam, oversample=4):
        ''' Compute the area above half maximum of the convolved image.
        '''
        original_offset = self.offset
        original_scale = self.pixel_scale
        ahms = []
        for i in range(10):
            itry = 0
            while itry < 10:
                xdither = np.random.uniform(-0.5, 0.5, 1)[0]
                ydither = np.random.uniform(-0.5, 0.5, 1)[0]
                rescale = np.random.uniform(0.9, 1.1, 1)[0]
                self.offset = (xdither, ydither)
                self.pixel_scale = original_scale * rescale
                try:
                    im = self.get_image(gparam, oversample=oversample)
                except RuntimeError:
                    itry += 1
                else:
                    break
            if itry >= 10:
                raise RuntimeError("Unable to create image to estimate AHM")
            mx = im.array.max()
            ahms.append(self.pixel_scale**2 * (im.array > mx/2.0).sum() / oversample**2)
        self.offset = original_offset
        self.pixel_scale = original_scale
        return np.mean(ahms), np.std(ahms)/np.sqrt(len(ahms))

    def compute_FWHM(self, gparam, oversample=4):
        ''' Compute FWHM of the convolved galaxy image.
        '''
        ahm, err = self.compute_AHM(gparam, oversample=oversample)
        fwhm = np.sqrt(4.0/np.pi * ahm)
        return fwhm, fwhm * err/ahm * 0.5

    def compute_PSF_AHM(self, oversample=4):
        ''' Compute the area above half maximum of the PSF.
        '''
        original_offset = self.offset
        original_scale = self.pixel_scale
        ahms = []
        for i in range(10):
            itry = 0
            while itry < 10:
                xdither = np.random.uniform(-0.5, 0.5, 1)[0]
                ydither = np.random.uniform(-0.5, 0.5, 1)[0]
                rescale = np.random.uniform(0.9, 1.1, 1)[0]
                self.offset = (xdither, ydither)
                self.pixel_scale = original_scale * rescale
                try:
                    im = self.get_PSF_image(oversample=oversample)
                except RuntimeError:
                    itry += 1
                else:
                    break
            if itry >= 10:
                raise RuntimeError("Unable to create image to estimate AHM")
            if isinstance(im, list):
                im = sum(im)
            mx = im.array.max()
            ahms.append(self.pixel_scale**2 * (im.array > mx/2.0).sum() / oversample**2)
        self.offset = original_offset
        self.pixel_scale = original_scale
        return np.mean(ahms), np.std(ahms)/np.sqrt(len(ahms))

    def compute_PSF_FWHM(self, oversample=4):
        ''' Compute FWHM of the PSF.
        '''
        ahm, err = self.compute_PSF_AHM(oversample=oversample)
        fwhm = np.sqrt(4.0/np.pi * ahm)
        return fwhm, fwhm * err/ahm * 0.5

    def compute_HLA(self, gparam, oversample=4, flux=None):
        ''' Compute the half-light-area of the PSF-convolved galaxy image.
        I.e., the area of the contour containing half the image light.
        '''
        im = self.get_image(gparam, oversample=oversample)
        if flux is None:
            flux = im.array.sum()
        pixel_values = im.array.ravel()
        pixel_values.sort()
        cumulative_sum = np.cumsum(pixel_values[::-1])
        npix = np.interp(0.5, cumulative_sum, np.arange(len(cumulative_sum)))
        return npix * self.pixel_scale**2 / oversample**2

    def compute_HLR(self, gparam, oversample=4, flux=None):
        ''' Compute the half-light-radius of the PSF-convolved galaxy image.
        '''
        return np.sqrt(1.0/np.pi * self.compute_HLA(gparam, oversample, flux))

    def compute_uncvl_HLA(self, gparam, oversample=4, flux=None):
        ''' Compute the half-light-area of the unconvolved galaxy image.
        I.e., the area of the contour containing half the image light.
        '''
        im = self.get_uncvl_image(gparam, oversample=oversample)
        if flux is None:
            flux = im.array.sum()
        pixel_values = im.array.ravel()
        pixel_values.sort()
        cumulative_sum = np.cumsum(pixel_values[::-1])
        npix = np.interp(0.5, cumulative_sum, np.arange(len(cumulative_sum)))
        return npix * self.pixel_scale**2 / oversample**2

    def compute_uncvl_HLR(self, gparam, oversample=4, flux=None):
        ''' Compute the half-light-radius of the unconvolved galaxy image.
        '''
        return np.sqrt(1.0/np.pi * self.compute_uncvl_HLA(gparam, oversample, flux))

class SersicTool(GalTool):
    def __init__(self, PSF, stamp_size, pixel_scale, offset=(0,0),
                 SED=1.0, bandpass=None, gsparams=None):
        self.PSF = PSF
        self.stamp_size = stamp_size
        self.pixel_scale = pixel_scale
        self.offset = offset
        self.gsparams = gsparams
        self.SED = SED
        self.bandpass = bandpass

    def _gparam_to_galsim(self, gparam):
        # Turn lmfit.Parameters into a galsim.ChromaticObject
        gal = galsim.Sersic(n=gparam['n'].value,
                            half_light_radius=gparam['hlr'].value,
                            gsparams=self.gsparams)
        gal = gal.shear(g=gparam['g'].value, beta=gparam['phi'].value * galsim.radians)
        gal = gal.shift(gparam['x0'].value, gparam['y0'].value)
        gal = gal.withFlux(gparam['flux'].value) * self.SED
        return gal

    def set_FWHM(self, gparam, FWHM, oversample=4):
        ''' Set the galaxy PSF-convolved FWHM.
        '''
        def FWHM_resid(scale):
            g1 = copy.deepcopy(gparam)
            g1['hlr'].value *= scale
            current_FWHM = self.compute_FWHM(g1, oversample=oversample)
            return current_FWHM[0] - FWHM
        scale = newton(FWHM_resid, 1.0, tol=0.001)
        gparam['hlr'].value *= scale
        return gparam

    def set_r2(self, gparam, r2, oversample=4):
        ''' Set the second moment radius sqrt(r^2).

        @param gparam      lmfit.Parameters object describing galaxy.
        @param r2          Target second moment radius sqrt(r^2)
        @param oversample  Factor by which to oversample drawn image for computation.
        @returns           New lmfit.Parameters object.
        '''
        def r2_resid(scale):
            g1 = copy.deepcopy(gparam)
            g1['hlr'].value *= scale
            current_r2 = self.get_r2(g1, oversample=oversample)
            return current_r2 - r2
        scale = newton(r2_resid, 1.0)
        gparam['hlr'].value *= scale
        return gparam

    def get_uncvl_r2(self, gparam):
        ''' Get second moment radius sqrt(r^2) of pre-PSF-convolved profile using polynomial
        approximation.
        @gparam   lmfit.Parameters
        '''
        return gparam['hlr'].value * chroma.Sersic_r2_over_hlr(gparam['n'].value)

    def set_uncvl_r2(self, gparam, r2):
        ''' Set the second moment radius sqrt(r^2) of the pre-PSF-convolved profile using
        polynomial approximation.

        @param gparam      lmfit.Parameters object describing galaxy.
        @param r2          Target second moment square radius
        @param oversample  Factor by which to oversample drawn image for r2 computation.
        @returns           New lmfit.Parameters object.
        '''
        gparam1 = copy.deepcopy(gparam)
        r2_now = self.get_uncvl_r2(gparam)
        scale = r2 / r2_now
        gparam1['hlr'].value = gparam['hlr'].value * scale
        return gparam1

    def get_ring_params(self, gparam, ring_beta, ring_shear):
        ''' Compute initial guess parameters for given angle around ellipticity ring during a ring
        test.

        @param gparam      lmfit.Parameters object describing galaxy.
        @param ring_beta   Angle around ellipticity ring in ring test.
        @param ring_shear  Shear to apply after rotation as part of ring test. (type=?)
        @returns           New lmfit.Parameters object.
        '''
        gparam1 = copy.deepcopy(gparam)
        rot_phi = gparam['phi'].value + ring_beta/2.0
        # complex ellipticity
        c_ellip = gparam['g'].value * complex(np.cos(2.0 * rot_phi), np.sin(2.0 * rot_phi))
        c_gamma = ring_shear.g1 + 1j * ring_shear.g2
        # sheared complex ellipticity
        s_c_ellip = chroma.apply_shear(c_ellip, c_gamma)
        s_g = abs(s_c_ellip)
        s_phi = np.angle(s_c_ellip) / 2.0

        gparam1['x0'].value \
          = gparam['x0'].value * np.cos(ring_beta / 2.0) \
          - gparam['y0'].value * np.sin(ring_beta / 2.0)
        gparam1['y0'].value \
          = gparam['x0'].value * np.sin(ring_beta / 2.0) \
          + gparam['y0'].value * np.cos(ring_beta / 2.0)
        gparam1['g'].value = s_g
        gparam1['phi'].value = s_phi
        return gparam1

    @staticmethod
    def default_galaxy():
        '''Setup lmfit.Parameters to represent a single Sersic galaxy.  Pick some default
        parameter values.  Parameters defining the single Sersic galaxy are:
        x0   - the x-coordinate of the galaxy center
        y0   - the y-coordinate of the galaxy center
        n    - the Sersic index.  0.5 gives a Gaussian profile, 1.0 gives an exponential profile,
               4.0 gives a de Vaucouleurs profile.
        hlr  - the galaxy half-light-radius.  This is strictly speaking the half light radius of
               a circularly symmetric profile of the given Sersic index `n`.
        g    - the magnitude of the galaxy ellipticity given in `g` units as used by GalSim.  In
               this convention, the major/minor axis ratio is given by: b/a = (1 - g) / (1 + g)
        phi  - the position angle of the galaxy major axis in radians.  0 indicates that the major
               axis is along the x-axis.
        '''
        gparam = lmfit.Parameters()
        gparam.add('x0', value=0.0)
        gparam.add('y0', value=0.0)
        gparam.add('n', value=4.0, vary=False)
        gparam.add('hlr', value=0.27)
        gparam.add('flux', value=1.0, vary=False)
        gparam.add('g', value=0.2, min=0.0, max=1.0)
        gparam.add('phi', value=0.0)
        return gparam

    def use_effective_PSF(self):
        ''' Integrate a chromatic PSF over wavelength to yield an effective PSF.  Galsim is doing
        this internally anyway, but we do this explicitly here so that it happens only once, instead
        of every time a .drawImage command is executed.
        '''
        star = galsim.Gaussian(fwhm=1.e-8) * self.SED
        prof = galsim.Convolve(star, self.PSF)
        prof0 = prof.evaluateAtWavelength(self.bandpass.effective_wavelength)
        scale = prof0.nyquistScale()
        N = prof0.SBProfile.getGoodImageSize(scale, 1.0)
        im = galsim.ImageD(N, N, scale=scale)
        # careful, don't want to convolve by pixel twice!
        prof.drawImage(self.bandpass, image=im, method='no_pixel')
        self.PSF = galsim.InterpolatedImage(im) # remember the effective PSF

    def apply_perturbative_correction(self, r2byr2=1.0, Vstar=1.e-8, Vgal=1.e-8,
                                      parang=0.0*galsim.degrees):
        if isinstance(self.PSF, galsim.ChromaticObject):
            star = galsim.Gaussian(fwhm=1.e-8) * self.SED
            prof = galsim.Convolve(star, self.PSF)
        elif isinstance(self.PSF, galsim.GSObject):
            prof = self.PSF
        else:
            raise ValueError("Don't recognize galaxy object type.")

        #-----------------------
        # Stellar DCR correction

        # `q` is the axis ratio of a 2D Gaussian representing the 1D DCR kernel. In principle, this
        # should be 0.0, but we need to set it to some small value for computability.
        q = 1.e-4
        sigma = (q * Vstar)**0.5
        kernel = galsim.Gaussian(sigma=sigma)
        kernel = kernel.shear(g1=-(1-q)/(1+q))
        kernel = kernel.rotate(parang)
        prof = galsim.Convolve(galsim.Deconvolve(kernel), prof)

        #----------------------------
        # Chromatic Seeing correction
        prof = prof.dilate(np.sqrt(r2byr2))

        #------------------------
        # Galactic DCR correction
        sigma = (q * Vgal)**0.5
        kernel = galsim.Gaussian(sigma=sigma)
        kernel = kernel.shear(g1=-(1-q)/(1+q))
        kernel = kernel.rotate(parang)
        prof = galsim.Convolve(kernel, prof)

        # and draw into an InterpolatedImage
        prof0 = prof.evaluateAtWavelength(self.bandpass.effective_wavelength)
        scale = prof0.nyquistScale()
        N = prof0.SBProfile.getGoodImageSize(scale, 1.0)
        im = galsim.ImageD(N*9, N*9, scale=scale*0.3)
        if isinstance(prof, galsim.ChromaticObject):
            prof.drawImage(self.bandpass, image=im, method='no_pixel')
        else:
            prof.drawImage(image=im, method='no_pixel')
        self.PSF = galsim.InterpolatedImage(im)

# Note that DoubleSersicTool and FastDoubleSersicTool are both currently untested.
class DoubleSersicTool(GalTool):
    ''' A GalTool to represent a sum of two chroma Sersic profiles.
    '''
    def __init__(self, PSF, stamp_size, pixel_scale, offset=(0,0), SED1=1.0, SED2=1.0,
                 bandpass=None, gsparams=None):
        ''' Initialize a single Sersic profile chromatic galaxy.

        @param PSF          galsim.ChromaticObject representing chromatic PSF
        @param stamp_size   Draw images this many pixels square
        @param pixel_scale  Pixels are this wide in arcsec.
        @param offset       Offset the center of the profile by this amount
        @param SED1         Optional galsim.SED galaxy spectrum for first component
        @param SED2         Optional galsim.SED galaxy spectrum for second component
        @param bandpass     Optional galsim.Bandpass to represent filter being imaged through.
        '''
        self.SED1 = SED1
        self.SED2 = SED2
        self.bandpass = bandpass
        self.PSF = PSF
        self.stamp_size = stamp_size
        self.pixel_scale = pixel_scale
        self.offset = offset
        self.gsparams = gsparams
        self.sersictools = [SersicTool(PSF, stamp_size, pixel_scale, offset=offset,
                                       SED=SED1, bandpass=bandpass, gsparams=gsparams),
                            SersicTool(PSF, stamp_size, pixel_scale, offset=offset,
                                       SED=SED2, bandpass=bandpass, gsparams=gsparams)]

    def indiv_gparam(self, gparam, icomp):
        out = lmfit.Parameters()
        search_string = "_{}".format(icomp+1)
        for p in gparam:
            if p.endswith(search_string):
                out.add(p.replace(search_string, ''), value=gparam[p].value)
        return out

    def get_image(self, gparam, ring_beta=None, ring_shear=None, oversample=1, icomp=None):
        ''' Draw a galaxy image using GalSim.  Potentially rotate and shear the galaxy as part of a
        ring test.  Optionally draw a high-resolution image.

        @param gparam      An lmfit.Parameters object that will be used to initialize a GalSim
                           object.
        @param ring_beta   Angle around ellipticity ring in ring test.
        @param ring_shear  galsim.Shear to apply after rotation as part of ring test.
        @param oversample  Integer factor by which to scale output image resolution and size.
        @param icomp       Which component to return the image of.
        @returns           galsim.Image
        '''
        if icomp is not None:
            return self.sersictools[icomp].get_image(
                self.indiv_gparam(gparam, icomp), ring_beta, ring_shear, oversample)
        else:
            im1 = self.sersictools[0].get_image(self.indiv_gparam(gparam, 0), ring_beta, ring_shear,
                                                oversample)
            im2 = self.sersictools[1].get_image(self.indiv_gparam(gparam, 1), ring_beta, ring_shear,
                                                oversample)
            return im1+im2

    def use_effective_PSF(self):
        ''' Integrate chromatic PSF over wavelength to yield effective PSFs.  GalSim is doing
        this internally anyway, but we do this explicitly here so that it happens only once, instead
        of every time a .drawImage() command is executed.
        '''
        for s in self.sersictools:
            s.use_effective_PSF()
        self.PSF = [s.PSF for s in self.sersictools]

    def get_PSF_image(self, oversample=1, method='auto', icomp=None):
        ''' Draw an image of both effective PSFs.  Note that the returned PSFs include convolution
        by the pixel response function, by default, though this can be overriden using the method
        keyword (e.g., method='fft').

        @param oversample  Integer factor by which to scale output image resolution and size.
        @returns (im1, im2)  Both effective PSFs corresponding to both component SEDs.
        '''
        if icomp is not None:
            return self.sersictools[icomp].get_PSF_image(oversample, method)
        else:
            return [s.get_PSF_image(oversample, method) for s in self.sersictools]

    def set_FWHM(self, gparam, FWHM, oversample=4, icomp=None):
        ''' Set the galaxy PSF-convolved FWHM.
        '''
        if icomp is not None:
            igparam = self.sersictools[icomp].set_FWHM(indiv_gparam(gparam, icomp), FWHM,
                                                       oversample)
            gparam = self.set_indiv_gparam('hlr', icomp, igparam['hlr'].value)
            return gparam
        else:
            def FWHM_resid(scale):
                g1 = copy.deepcopy(gparam)
                g1['hlr_1'].value *= scale
                g1['hlr_2'].value *= scale
                current_FWHM = self.compute_FWHM(g1, oversample=oversample)
                return current_FWHM[0] - FWHM
            scale = newton(FWHM_resid, 1.0, tol=0.01)
            gparam['hlr_1'].value *= scale
            gparam['hlr_2'].value *= scale
            return gparam

    def set_r2(self, gparam, r2, oversample=4):
        ''' Set the second moment radius sqrt(r^2).

        @param gparam      lmfit.Parameters object describing galaxy.
        @param r2          Target second moment radius sqrt(r^2)
        @param oversample  Factor by which to oversample drawn image for r2 computation.
        @returns           New lmfit.Parameters object.
        '''
        def r2_resid(scale):
            g1 = copy.deepcopy(gparam)
            g1['hlr_1'].value *= scale
            g1['hlr_2'].value *= scale
            current_r2 = self.get_r2(g1, oversample=oversample)
            return current_r2 - r2
        scale = newton(r2_resid, 1.0)
        gparam['hlr_1'].value *= scale
        gparam['hlr_2'].value *= scale
        return gparam

    def get_uncvl_r2(self, gparam):
        ''' Get second moment radius of pre-PSF-convolved profile using polynomial approximation.
        @gparam   lmfit.Parameters
        '''
        return chroma.component_Sersic_r2([gparam['n_1'].value, gparam['n_2'].value],
                                          [gparam['flux_1'].value, gparam['flux_2'].value],
                                          [gparam['hlr_1'].value, gparam['hlr_2'].value])

    def set_uncvl_r2(self, gparam, r2):
        ''' Set the second moment square radius of the pre-PSF-convolved profile using polynomial
        approximation.

        @param gparam      lmfit.Parameters object describing galaxy.
        @param r2          Target second moment radius sqrt(r^2)
        @param oversample  Factor by which to oversample drawn image for computation.
        @returns           New lmfit.Parameters object.
        '''
        gparam1 = copy.deepcopy(gparam)
        r2_now = self.get_uncvl_r2(gparam)
        scale = r2 / r2_now
        gparam1['hlr_1'].value *= scale
        gparam1['hlr_2'].value *= scale
        return gparam1

    def get_ring_params(self, gparam, ring_beta, ring_shear):
        ''' Compute initial guess parameters for given angle around ellipticity ring during a ring
        test.

        @param gparam      lmfit.Parameters object describing galaxy.
        @param ring_beta   Angle around ellipticity ring in ring test.
        @param ring_shear  Shear to apply after rotation as part of ring test. (type=?)
        @returns           New lmfit.Parameters object.
        '''
        gparam1 = copy.deepcopy(gparam)

        rot_phi1 = gparam['phi_1'].value + ring_beta/2.0
        # complex ellipticity
        c_ellip1 = gparam['g_1'].value * complex(np.cos(2.0 * rot_phi1), np.sin(2.0 * rot_phi1))
        c_gamma1 = ring_shear.g1 + 1j * ring_shear.g2
        # sheared complex ellipticity
        s_c_ellip1 = chroma.apply_shear(c_ellip1, c_gamma1)
        s_g1 = abs(s_c_ellip1)
        s_phi1 = np.angle(s_c_ellip1) / 2.0

        gparam1['x0_1'].value \
          = gparam['x0_1'].value * np.cos(ring_beta / 2.0) \
          - gparam['y0_1'].value * np.sin(ring_beta / 2.0)
        gparam1['y0_1'].value \
          = gparam['x0_1'].value * np.sin(ring_beta / 2.0) \
          + gparam['y0_1'].value * np.cos(ring_beta / 2.0)
        gparam1['g_1'].value = s_g1
        gparam1['phi_1'].value = s_phi1

        rot_phi2 = gparam['phi_2'].value + ring_beta/2.0
        # complex ellipticity
        c_ellip2 = gparam['g_2'].value * \
          complex(np.cos(2.0 * rot_phi2), np.sin(2.0 * rot_phi2))
        c_gamma2 = ring_shear.g2 + 1j * ring_shear.g2
        # sheared complex ellipticity
        s_c_ellip2 = chroma.apply_shear(c_ellip2, c_gamma2)
        s_g2 = abs(s_c_ellip2)
        s_phi2 = np.angle(s_c_ellip2) / 2.0

        gparam1['x0_2'].value \
          = gparam['x0_2'].value * np.cos(ring_beta / 2.0) \
          - gparam['y0_2'].value * np.sin(ring_beta / 2.0)
        gparam1['y0_2'].value \
          = gparam['x0_2'].value * np.sin(ring_beta / 2.0) \
          + gparam['y0_2'].value * np.cos(ring_beta / 2.0)
        gparam1['g_2'].value = s_g2
        gparam1['phi_2'].value = s_phi2

        return gparam1

    @staticmethod
    def default_galaxy():
        '''Setup lmfit.Parameters to represent a double Sersic galaxy.  Pick some default
        parameter values.  Parameters for each Sersic galaxy component are:
        x0   - the x-coordinate of the galaxy center
        y0   - the y-coordinate of the galaxy center
        n    - the Sersic index.  0.5 gives a Gaussian profile, 1.0 gives an exponential profile,
               4.0 gives a de Vaucouleurs profile.
        hlr  - the galaxy half-light-radius.  This is strictly speaking the half light radius of
               a circularly symmetric profile of the given Sersic index `n`.
        g    - the magnitude of the galaxy ellipticity given in `g` units as used by GalSim.  In
               this convention, the major/minor axis ratio is given by: b/a = (1 - g) / (1 + g)
        phi  - the position angle of the galaxy major axis in radians.  0 indicates that the major
               axis is along the x-axis.
        Parameters for the first component will have suffix _1, parameters for the second component
        will have suffix _2.  By default, the two components are constrained to by concentric and
        coelliptical.  Similar to the single Sersic case, the total flux and both Sersic indices are
        fixed by default.
        '''
        gparam = lmfit.Parameters()
        # bulge first
        gparam.add('x0_1', value=0.0)
        gparam.add('y0_1', value=0.0)
        gparam.add('n_1', value=4.0, vary=False)
        gparam.add('hlr_1', value=0.3, min=0.0)
        gparam.add('flux_1', value=0.25, min=0.0, max=1.0)
        gparam.add('g_1', value=0.2, min=0.0, max=1.0)
        gparam.add('phi_1', value=0.0)
        # then disk
        gparam.add('x0_2', expr='x0_1')
        gparam.add('y0_2', expr='y0_1')
        gparam.add('n_2', value=1.0, vary=False)
        gparam.add('hlr_2', value=0.4, min=0.0)
        gparam.add('flux_2', expr='1.0 - flux_1')
        gparam.add('g_2', expr='g_1')
        gparam.add('phi_2', expr='phi_1')
        #initialize constrained variables
        dummyfit = lmfit.Minimizer(lambda x: 0, gparam)
        dummyfit.prepare_fit()
        return gparam


'''
Soft Actor-Critic version 2
using target Q instead of V net: 2 Q net, 2 target Q net, 1 policy net
add alpha loss compared with version 1
paper: https://arxiv.org/pdf/1812.05905.pdf
'''

import math
import random

# import gym
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Normal
from torch.utils.tensorboard import SummaryWriter

# from IPython.display import clear_output
# import matplotlib.pyplot as plt
# from matplotlib import animation
# from IPython.display import display
# from reacher import Reacher

import argparse
import time

import simulation as sim
from Environment import Environment

import actuator_array as act
from init import action_dim,state_dim

GPU = True
device_idx = 2
if GPU:
    device = torch.device("cuda:" + str(device_idx) if torch.cuda.is_available() else "cpu")
else:
    device = torch.device("cpu")
print(device)

parser = argparse.ArgumentParser(description='Train or test neural net motor controller.')
parser.add_argument('--train', dest='train', action='store_true', default=False)
parser.add_argument('--test', dest='test', action='store_true', default=False)

args = parser.parse_args()

class ReplayBuffer:
    def __init__(self, capacity):
        self.capacity = capacity
        self.buffer = []
        self.position = 0

    def push(self, state, action, reward, next_state, done):
        if len(self.buffer) < self.capacity:
            self.buffer.append(None)
        self.buffer[self.position] = (state, action, reward, next_state, done)
        self.position = int((self.position + 1) % self.capacity)  # as a ring buffer

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        state, action, reward, next_state, done = map(np.stack, zip(*batch)) # stack for each element
        '''
        the * serves as unpack: sum(a,b) <=> batch=(a,b), sum(*batch) ;
        zip: a=[1,2], b=[2,3], zip(a,b) => [(1, 2), (2, 3)] ;
        the map serves as mapping the function on each list element: map(square, [2,3]) => [4,9] ;
        np.stack((1,2)) => array([1, 2])
        '''
        return state, action, reward, next_state, done

    def __len__(self):
        return len(self.buffer)

class SoftQNetwork(nn.Module):
    def __init__(self, num_inputs, num_actions, hidden_size, init_w=3e-3):
        super(SoftQNetwork, self).__init__()

        self.linear1 = nn.Linear(num_inputs + num_actions, hidden_size)
        self.linear2 = nn.Linear(hidden_size, hidden_size)
        self.linear3 = nn.Linear(hidden_size, hidden_size)
        self.linear4 = nn.Linear(hidden_size, 1)

        self.linear4.weight.data.uniform_(-init_w, init_w)
        self.linear4.bias.data.uniform_(-init_w, init_w)

    def forward(self, state, action):
        x = torch.cat([state, action], 1) # the dim 0 is number of samples
        x = F.relu(self.linear1(x))
        x = F.relu(self.linear2(x))
        x = F.relu(self.linear3(x))
        x = self.linear4(x)
        return x

class PolicyNetwork(nn.Module):
    def __init__(self, num_inputs, num_actions, hidden_size, action_range=1., init_w=3e-3, log_std_min=-20, log_std_max=2):
        super(PolicyNetwork, self).__init__()

        self.log_std_min = log_std_min
        self.log_std_max = log_std_max

        self.linear1 = nn.Linear(num_inputs, hidden_size)
        self.linear2 = nn.Linear(hidden_size, hidden_size)
        self.linear3 = nn.Linear(hidden_size, hidden_size)
        self.linear4 = nn.Linear(hidden_size, hidden_size)

        self.mean_linear = nn.Linear(hidden_size, num_actions)
        self.mean_linear.weight.data.uniform_(-init_w, init_w)
        self.mean_linear.bias.data.uniform_(-init_w, init_w)

        self.log_std_linear = nn.Linear(hidden_size, num_actions)
        self.log_std_linear.weight.data.uniform_(-init_w, init_w)
        self.log_std_linear.bias.data.uniform_(-init_w, init_w)

        self.action_range = action_range
        self.num_actions = num_actions

    def forward(self, state):
        x = F.relu(self.linear1(state))
        x = F.relu(self.linear2(x))
        x = F.relu(self.linear3(x))
        x = F.relu(self.linear4(x))

        mean    = (self.mean_linear(x))
        # mean    = F.leaky_relu(self.mean_linear(x))
        log_std = self.log_std_linear(x)
        log_std = torch.clamp(log_std, self.log_std_min, self.log_std_max)

        return mean, log_std

    def evaluate(self, state, epsilon=1e-6):
        '''
        generate sampled action with state as input wrt the policy network;
        '''
        mean, log_std = self.forward(state)
        std = log_std.exp() # no clip in evaluation, clip affects gradients flow

        normal = Normal(0, 1)
        z      = normal.sample(mean.shape)
        action_0 = torch.tanh(mean + std*z.to(device)) # TanhNormal distribution as actions; reparameterization trick
        action = self.action_range*action_0 + 30                ### 加偏置
        # The log-likelihood here is for the TanhNorm distribution instead of only Gaussian distribution. \
        # The TanhNorm forces the Gaussian with infinite action range to be finite. \
        # For the three terms in this log-likelihood estimation: \
        # (1). the first term is the log probability of action as in common \
        # stochastic Gaussian action policy (without Tanh); \
        # (2). the second term is the caused by the Tanh(), \
        # as shown in appendix C. Enforcing Action Bounds of https://arxiv.org/pdf/1801.01290.pdf, \
        # the epsilon is for preventing the negative cases in log; \
        # (3). the third term is caused by the action range I used in this code is not (-1, 1) but with \
        # an arbitrary action range, which is slightly different from original paper.
        log_prob = Normal(mean, std).log_prob(mean+ std*z.to(device)) - torch.log((1. - action_0.pow(2))/2 + epsilon)
        # both dims of normal.log_prob and -log(1-a**2) are (N,dim_of_action);
        # the Normal.log_prob outputs the same dim of input features instead of 1 dim probability,
        # needs sum up across the features dim to get 1 dim prob; or else use Multivariate Normal.
        log_prob = log_prob.sum(dim=1, keepdim=True)
        return action, log_prob, z, mean, log_std

    def get_action(self, state, deterministic):
        state = torch.FloatTensor(state).unsqueeze(0).to(device)
        mean, log_std = self.forward(state)
        std = log_std.exp()

        normal = Normal(0, 1)
        z      = normal.sample(mean.shape).to(device)
        action = self.action_range* torch.tanh(mean + std*z)

        action = self.action_range* torch.tanh(mean).detach().cpu().numpy()[0] if deterministic else action.detach().cpu().numpy()[0]
        action = action + 30
        return action

    def sample_action(self,):
        a=torch.FloatTensor(self.num_actions).uniform_(-1, 1)
        return self.action_range*a.numpy()

class SAC_Trainer():
    def __init__(self, replay_buffer, hidden_dim, action_range):
        self.replay_buffer = replay_buffer

        self.soft_q_net1 = SoftQNetwork(state_dim, action_dim, hidden_dim).to(device)
        self.soft_q_net2 = SoftQNetwork(state_dim, action_dim, hidden_dim).to(device)
        self.target_soft_q_net1 = SoftQNetwork(state_dim, action_dim, hidden_dim).to(device)
        self.target_soft_q_net2 = SoftQNetwork(state_dim, action_dim, hidden_dim).to(device)
        self.policy_net = PolicyNetwork(state_dim, action_dim, hidden_dim, action_range).to(device)
        self.log_alpha = torch.zeros(1, dtype=torch.float32, requires_grad=True, device=device)
        print('Soft Q Network (1,2): ', self.soft_q_net1)
        print('Policy Network: ', self.policy_net)

        for target_param, param in zip(self.target_soft_q_net1.parameters(), self.soft_q_net1.parameters()):
            target_param.data.copy_(param.data)
        for target_param, param in zip(self.target_soft_q_net2.parameters(), self.soft_q_net2.parameters()):
            target_param.data.copy_(param.data)

        self.soft_q_criterion1 = nn.MSELoss()
        self.soft_q_criterion2 = nn.MSELoss()

        soft_q_lr = 3e-4
        policy_lr = 3e-4
        alpha_lr  = 3e-4

        self.soft_q_optimizer1 = optim.Adam(self.soft_q_net1.parameters(), lr=soft_q_lr)
        self.soft_q_optimizer2 = optim.Adam(self.soft_q_net2.parameters(), lr=soft_q_lr)
        self.policy_optimizer = optim.Adam(self.policy_net.parameters(), lr=policy_lr)
        self.alpha_optimizer = optim.Adam([self.log_alpha], lr=alpha_lr)

    def update(self, batch_size, reward_scale=10., auto_entropy=True, target_entropy=-2, gamma=0.99,soft_tau=1e-2):
        state, action, reward, next_state, done = self.replay_buffer.sample(batch_size)
        # print('sample:', state, action,  reward, done)

        state      = torch.FloatTensor(state).to(device)
        next_state = torch.FloatTensor(next_state).to(device)
        action     = torch.FloatTensor(action).to(device)
        reward     = torch.FloatTensor(reward).unsqueeze(1).to(device)  # reward is single value, unsqueeze() to add one dim to be [reward] at the sample dim;
        done       = torch.FloatTensor(np.float32(done)).unsqueeze(1).to(device)

        action = action / 60
        reward = reward / 20

        predicted_q_value1 = self.soft_q_net1(state, action)
        predicted_q_value2 = self.soft_q_net2(state, action)
        new_action, log_prob, z, mean, log_std = self.policy_net.evaluate(state)
        new_next_action, next_log_prob, _, _, _ = self.policy_net.evaluate(next_state)

        new_action = new_action / 60
        new_next_action = new_next_action / 60

        # reward = reward_scale * (reward - reward.mean(dim=0)) / (reward.std(dim=0) + 1e-6) # normalize with batch mean and std; plus a small number to prevent numerical problem
    # Updating alpha wrt entropy
        # alpha = 0.0  # trade-off between exploration (max entropy) and exploitation (max Q)
        if auto_entropy is True:
            alpha_loss = -(self.log_alpha * (log_prob + target_entropy).detach()).mean()
            # print('alpha loss: ',alpha_loss)
            self.alpha_optimizer.zero_grad()
            alpha_loss.backward()
            self.alpha_optimizer.step()
            self.alpha = self.log_alpha.exp()
        else:
            self.alpha = 1.
            alpha_loss = 0

    # Training Q Function
        target_q_min = torch.min(self.target_soft_q_net1(next_state, new_next_action),self.target_soft_q_net2(next_state, new_next_action)) - self.alpha * next_log_prob
        target_q_value = reward + (1 - done) * gamma * target_q_min # if done==1, only reward
        q_value_loss1 = self.soft_q_criterion1(predicted_q_value1, target_q_value.detach())  # detach: no gradients for the variable
        q_value_loss2 = self.soft_q_criterion2(predicted_q_value2, target_q_value.detach())

        self.soft_q_optimizer1.zero_grad()
        q_value_loss1.backward()
        self.soft_q_optimizer1.step()

        self.soft_q_optimizer2.zero_grad()
        q_value_loss2.backward()
        self.soft_q_optimizer2.step()

    # Training Policy Function
        predicted_new_q_value = torch.min(self.soft_q_net1(state, new_action),self.soft_q_net2(state, new_action))
        policy_loss = (self.alpha * log_prob - predicted_new_q_value).mean()

        self.policy_optimizer.zero_grad()
        policy_loss.backward()
        self.policy_optimizer.step()

        # print('q loss: ', q_value_loss1, q_value_loss2)
        # print('policy loss: ', policy_loss )

    # Soft update the target value net
        for target_param, param in zip(self.target_soft_q_net1.parameters(), self.soft_q_net1.parameters()):
            target_param.data.copy_(  # copy data value into target parameters
                target_param.data * (1.0 - soft_tau) + param.data * soft_tau
            )
        for target_param, param in zip(self.target_soft_q_net2.parameters(), self.soft_q_net2.parameters()):
            target_param.data.copy_(  # copy data value into target parameters
                target_param.data * (1.0 - soft_tau) + param.data * soft_tau
            )
        return predicted_new_q_value.mean().item(),q_value_loss1.item(),q_value_loss2.item(),policy_loss.item(),self.alpha.item()

    def save_model(self, path):
        torch.save(self.soft_q_net1.state_dict(), path+'_q1')
        torch.save(self.soft_q_net2.state_dict(), path+'_q2')
        torch.save(self.policy_net.state_dict(), path+'_policy')

    def load_model(self, path):
        self.soft_q_net1.load_state_dict(torch.load(path+'_q1'))
        self.soft_q_net2.load_state_dict(torch.load(path+'_q2'))
        self.policy_net.load_state_dict(torch.load(path+'_policy'))

        self.soft_q_net1.eval()
        self.soft_q_net2.eval()
        self.policy_net.eval()

replay_buffer_size = 1e6
replay_buffer = ReplayBuffer(replay_buffer_size)

# choose env
ENV = ['Reacher', 'Pendulum-v0', 'HalfCheetah-v2'][2]

act_array = act.Actuator_array()
env = Environment()
action_space = [action_dim]
state_space = [state_dim]
action_range = 30

# hyper-parameters for RL training
max_episodes  = 5e6
max_steps   = 1e3   # Pendulum needs 150 steps per episode to learn well, cannot handle 20
frame_idx   = 0
batch_size  = 64
explore_steps = -1  # for random action sampling in the beginning of training
update_itr = 1
AUTO_ENTROPY=True
DETERMINISTIC=False
hidden_dim = 512
rewards     = []
model_path = './model/SAC_1_28_3'

q_value_episode = np.array([])
q_value_loss1_episode = np.array([])
q_value_loss2_episode = np.array([])
policy_loss_episode = np.array([])
alpha_log_episode = np.array([])
writer = SummaryWriter('./logs/SAC_1_28_3')
tmp_avg_reward = 0

sac_trainer=SAC_Trainer(replay_buffer, hidden_dim=hidden_dim, action_range=action_range  )

if __name__ == '__main__':
    if 1:
        # training loop
        for eps in range(int(max_episodes)):
            if ENV == 'Reacher':
                state = env.reset(SCREEN_SHOT)
            else:
                state =  env.reset()
            episode_reward = 0

            for step in range(int(max_steps)):
                if frame_idx > explore_steps:
                    action = sac_trainer.policy_net.get_action(state, deterministic = DETERMINISTIC)
                else:
                    action = sac_trainer.policy_net.sample_action()
                if ENV ==  'Reacher':
                    next_state, reward, done, _ = env.step(action, SPARSE_REWARD, SCREEN_SHOT)
                else:
                    next_state, reward, done = env.step(action)
                    # env.render()

                done = 0
                replay_buffer.push(state, action, reward, next_state, done)

                state = next_state
                episode_reward += reward
                frame_idx += 1

                if len(replay_buffer) > batch_size:
                    for i in range(update_itr):
                        q_value,q_value_loss1,q_value_loss2,policy_loss,alpha_log = sac_trainer.update(batch_size, reward_scale=10., auto_entropy=AUTO_ENTROPY, target_entropy=-1.*action_dim)
                        q_value_episode = np.append(q_value_episode,q_value)
                        q_value_loss1_episode = np.append(q_value_loss1_episode,q_value_loss1)
                        q_value_loss2_episode = np.append(q_value_loss2_episode,q_value_loss2)
                        policy_loss_episode = np.append(policy_loss_episode,policy_loss)
                        alpha_log_episode = np.append(alpha_log_episode,alpha_log)
                # if done:
                #     break

            writer.add_scalar('train/reward', episode_reward, eps)
            writer.add_scalar('train/q_value', q_value_episode.mean(), eps)
            writer.add_scalar('loss/critic_1', q_value_loss1_episode.mean(), eps)
            writer.add_scalar('loss/critic_2', q_value_loss2_episode.mean(), eps)
            writer.add_scalar('loss/policy', policy_loss_episode.mean(), eps)
            writer.add_scalar('train/alpha', alpha_log_episode.mean(), eps)

            print('Episode: ', eps, '| Episode Reward: ', episode_reward)
            # rewards.append(episode_reward)

            if eps % 100 == 0 and eps>0: # 100个episode测试一次如果大于上一次的测试结果，储存下来
                avg_reward = 0.
                episodes = 3
                cal_dist_set = []
                state = env.reset()

                for _  in range(episodes):
                    episode_reward = 0
                    done = False
                    for i in range(1000):
                        action = sac_trainer.policy_net.get_action(state, deterministic = DETERMINISTIC)

                        next_state, reward, done = env.step(action)
                        episode_reward += reward

                        state = next_state
                        if done:                         ## 实际中不可能reset包裹
                            # state = env.reset()
                            pass

                        if env.evl_flag == 1:       ## 有包裹过线
                            cal_dist_set.append(env.cal_dist)
                            env.evl_flag = 0

                    avg_reward += episode_reward
                episode_reward = 0
                avg_reward /= episodes
                state = env.reset()             ### 每个episode前都要reset,reset后需要加mask        之前忘记在测试完成后加reset了
                cal_dist_set = np.array(cal_dist_set)
                dist_success = cal_dist_set[cal_dist_set>=150]
                dist_fail = cal_dist_set[cal_dist_set<150]

                writer.add_scalar('avg_reward/test', avg_reward, eps)
                try:
                    writer.add_histogram('delta_dist_success', dist_success, eps, bins = 500)         # 看delta T是否收敛到设定值
                except:
                    pass
                try:
                    writer.add_histogram('delta_dist_fail', dist_fail, eps, bins = 500)         # 看delta T是否收敛到设定值
                except:
                    pass
                print("----------------------------------------")
                print("Test Episodes: {}, Avg. Reward: {}, Avg. delta_dist: {}".format(eps, round(avg_reward, 2),np.mean(cal_dist_set)))
                print("----------------------------------------")
                if avg_reward > tmp_avg_reward:     # 变大后保存模型
                    tmp_avg_reward = avg_reward
                    sac_trainer.save_model(model_path)          ## 每次test完成后修改mode，及时测试避免错误
                    print("saving model and highest average reward: {}".format(tmp_avg_reward))

    # if args.test:
    #     sac_trainer.load_model(model_path)
    #     for eps in range(10):
    #         if ENV == 'Reacher':
    #             state = env.reset(SCREEN_SHOT)
    #         else:
    #             state =  env.reset()
    #         episode_reward = 0

    #         for step in range(max_steps):
    #             action = sac_trainer.policy_net.get_action(state, deterministic = DETERMINISTIC)
    #             if ENV ==  'Reacher':
    #                 next_state, reward, done, _ = env.step(action, SPARSE_REWARD, SCREEN_SHOT)
    #             else:
    #                 next_state, reward, done, _ = env.step(action)
    #                 env.render()

    #             episode_reward += reward
    #             state=next_state

    #         print('Episode: ', eps, '| Episode Reward: ', episode_reward)


#!/usr/bin/env python
import warnings
warnings.filterwarnings("ignore")
import numpy as np
from astropy.io import fits, ascii
import matplotlib.pyplot as plt
from matplotlib.ticker import AutoMinorLocator
import smart
#import splat
#import splat.model as spmd

#def _constructModelName(teff, logg, metal, en, order, path=None):
#    """
#    Return the full name of the BT-Settl model.
#    """
#    if path is None:
#        path  = '/Users/dinohsu/projects/Models/models/btsettl08/' + \
#        'NIRSPEC-O' + str(order) + '-RAW/'
#    else:
#        path  = path + '/NIRSPEC-O' + str(order) + '-RAW/'
#    full_name = path + 'btsettl08_t'+ str(teff) + '_g' + \
#    '{0:.2f}'.format(float(logg)) + '_z-' + '{0:.2f}'.format(float(metal)) + \
#    '_en' + '{0:.2f}'.format(float(en)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'
#
#    return full_name

class Model():
    """
    The Model class reads in the BT-SETTL or PHOENIXACES models.
    The unit of wavelength is in Angstrom and the unit of the model flux is in erg/s/cm^2/Angstrom.
    (The models in the libraries have the unit of micron, which differed by 10^4 in flux)

    Parameters
    ----------
    1. Read in a BT-Settl model or PHOENIXACES model.
    teff : float
          The effective temperature in Kelvins.
    logg : float
          The log(gravity), given in two decimal digits.
          Ex: logg=4.50
    metal  : float
           The metalicity, given in two decimal digits.
           Ex. metal=0.00
    en   : float
           alpha enhancement. given in two decimal digits.
           Ex. en=0.00

    modelset: str
            available models are
            NIRSPEC: 'btsettl08', 'SONORA_2018'
            APOGEE: 'btsettl08', 'marcs-apogee-dr15', 'phoenix-aces-agss-cond-2011', 'phoenix-btsettl-cifist2011-2015'

    order: int
           This is only for the Keck/NIRSPEC. The order of the model, given from 29 to 80

    path : str
           The path to the model

    2. Creat a model instance with given wavelengths and fluxes
    flux : astropy.table.column.Column
           The input flux.
    wave : astropy.table.column.Column
           The input wavelength.

    Returns
    -------
    flux : astropy.table.column.Column
           The flux retrieved from the model. Our default unit is erg/s/cm^2/Angstrom.
    wave : astropy.table.column.Column
           The wavelength retrieved from the model. Our default unit is Angstrom.

    Examples
    --------
    >>> import smart
    >>> model = smart.Model(teff=2300, logg=5.5, order=33, path='/path/to/models')
    >>> model.plot()
    """
    def __init__(self, **kwargs):
        self.path  = kwargs.get('path')
        self.order = kwargs.get('order')
        self.instrument = kwargs.get('instrument','nirspec')

        if self.order != None and self.instrument == 'nirspec':
            self.teff     = kwargs.get('teff', 2500)
            self.logg     = kwargs.get('logg', 5.00)
            self.metal    = kwargs.get('metal', 0.00)
            self.en       = kwargs.get('en', 0.00)
            self.modelset = kwargs.get('modelset', 'btsettl08')

            wave, flux = smart.forward_model.InterpolateModel.InterpModel(self.teff, self.logg, self.metal, self.en,
                                                                          modelset=self.modelset, order=self.order, instrument=self.instrument)
            #elif self.metal != 0.0:
            #    wave, flux = smart.forward_model.InterpolateModel.InterpModel(Teff=self.teff, Logg=self.logg, Metal=self.metal,
            #    modelset=self.modelset, order=self.order, instrument=self.instrument)

            if self.modelset == 'btsettl08':
                self.wave = wave * 10000 #convert to Angstrom
                self.flux = flux / 10000 #convert from erg/s/cm^2/micron to erg/s/cm^2/Angstrom
            else:
                self.wave = wave
                self.flux = flux

        elif self.instrument == 'apogee':
            self.teff     = kwargs.get('teff', 2500)
            self.logg     = kwargs.get('logg', 5.00)
            self.metal    = kwargs.get('metal', 0.00)
            self.en       = kwargs.get('en', 0.00)
            self.modelset = kwargs.get('modelset', 'btsettl08')

            #wave, flux = smart.forward_model.InterpolateModel.InterpModel(self.teff, self.logg,
            #    modelset=self.modelset, order=self.order, instrument=self.instrument)

            wave, flux = smart.forward_model.InterpolateModel.InterpModel(self.teff, self.logg, self.metal, self.en,
                                                                          modelset=self.modelset, order=self.order, instrument=self.instrument)

            if self.modelset == 'btsettl08':
                self.wave = wave * 10000 #convert to Angstrom
                self.flux = flux / 10000 #convert from erg/s/cm^2/micron to erg/s/cm^2/Angstrom

            else:
                self.wave = wave # Angstrom
                self.flux = flux # erg/s/cm^2/Angstrom

        else:
            self.wave   = kwargs.get('wave', [])
            self.flux   = kwargs.get('flux', [])

    def plot(self, **kwargs):
        """
        Plot the model spectrum.
        """
        if self.order != None:
            name = str(_constructModelName(self.teff, self.logg,
                self.metal, self.en, self.order, self.path))
            output = kwargs.get('output', str(name) + '.pdf')
            ylim = kwargs.get('yrange', [min(self.flux)-.2, max(self.flux)+.2])
            title  = kwargs.get('title')
            save   = kwargs.get('save', False)

            plt.figure(figsize=(16,6))
            plt.plot(self.wave, self.flux, color='k',
                alpha=.8, linewidth=1, label=name)
            plt.legend(loc='upper right', fontsize=12)
            plt.ylim(ylim)

            minor_locator = AutoMinorLocator(5)
            #ax.xaxis.set_minor_locator(minor_locator)
            # plt.grid(which='minor')

            plt.xlabel(r'$\lambda$ [$\mathring{A}$]', fontsize=18)
            plt.ylabel(r'$Flux$', fontsize=18)
            #plt.ylabel(r'$F_{\lambda}$ [$erg/s \cdot cm^{2}$]', fontsize=18)
            if title != None:
                plt.title(title, fontsize=20)
            plt.tight_layout()

            if save == True:
                plt.savefig(output)
            plt.show()
            plt.close()

        else:
            output = kwargs.get('output'+ '.pdf')
            ylim   = kwargs.get('yrange', [min(self.flux)-.2, max(self.flux)+.2])
            title  = kwargs.get('title')
            save   = kwargs.get('save', False)

            plt.figure(figsize=(16,6))
            plt.plot(self.wave, self.flux, color='k', alpha=.8, linewidth=1)
            plt.legend(loc='upper right', fontsize=12)
            plt.ylim(ylim)

            minor_locator = AutoMinorLocator(5)
            #ax.xaxis.set_minor_locator(minor_locator)
            # plt.grid(which='minor')

            plt.xlabel(r'$\lambda$ [$\mathring{A}$]', fontsize=18)
            plt.ylabel(r'$Flux$', fontsize=18)
            #plt.ylabel(r'$F_{\lambda}$ [$erg/s \cdot cm^{2}$]', fontsize=18)
            if title != None:
                plt.title(title, fontsize=20)
            plt.tight_layout()

            if save == True:
                plt.savefig(output)
            plt.show()
            plt.close()


"""
The crux of ``coronagraph`` noise modeling is to determine the photon count rate
incident upon the detector due to both the target planet and an assortment of
different telescope, instrumental, and astrophysical noise sources.
The following classes and functions serve as your interface to the photon count
rate calculations. The core function for these calculations is
:func:`count_rates`, but it may be accessed using the :class:`CoronagraphNoise`
object.
"""

from __future__ import (division as _, print_function as _,
                absolute_import as _, unicode_literals as _)

# Import dependent modules
import numpy as np
import sys, os
import matplotlib.pyplot as plt

from .degrade_spec import downbin_spec
from .convolve_spec import convolve_spec
from .noise_routines import Fstar, Fplan, FpFs, cplan, czodi, cezodi, cspeck, \
    cdark, cread, ctherm, ccic, f_airy, ctherm_earth, construct_lam, \
    set_quantum_efficiency, set_read_noise, set_dark_current, set_lenslet, \
    set_throughput, set_atmos_throughput, \
    exptime_element, get_sky_flux
from .teleplanstar import Telescope, Planet, Star

__all__ = ['count_rates', 'CoronagraphNoise']

class CoronagraphNoise(object):
    """
    The primary interface for ``coronagraph`` noise modeling. This object wraps
    around the functionality of :func:`count_rates`. Simply instantiate a
    `CoronagraphNoise` object by passing it :class:`telescope`, :class:`planet`,
    and :class:`star` objects, and then call
    :func:`CoronagraphNoise.run_count_rates` to perform the photon count rate
    calculation.

    Parameters
    ----------
    telescope : Telescope
        Initialized object containing ``Telescope`` parameters
    planet : Planet
        Initialized object containing ``Planet`` parameters
    star : Star
        Initialized object containing ``Star`` parameters
    texp : float
        Exposure time for which to generate synthetic data [hours]
    wantsnr : float, optional
        Desired signal-to-noise ratio in each pixel
    FIX_OWA : bool, optional
        Set to fix OWA at ``OWA*lammin/D``, as would occur if lenslet array is
        limiting the OWA
    COMPUTE_LAM : bool, optional
        Set to compute lo-res wavelength grid, otherwise the grid input as
        variable ``lam`` is used
    SILENT : bool, optional
        Set to suppress print statements
    NIR : bool, optional
        Re-adjusts pixel size in NIR, as would occur if a second instrument
        was designed to handle the NIR
    THERMAL : bool, optional
        Set to compute thermal photon counts due to telescope temperature
    GROUND : bool, optional
        Set to simulate ground-based observations through atmosphere
    vod : bool, optional
        "Valley of Death" red QE parameterization from Robinson et al. (2016)
    set_fpa : float, optional
        Specify the fraction of planetary signal in Airy pattern, default will
        calculate it from the photometric aperture size `X`
    roll_maneuver : bool, optional
        This assumes an extra factor of 2 hit to the background noise due to a
        telescope roll maneuver needed to subtract out the background. See
        Brown (2005) for more details.

    Note
    ----
    The results of the coronagraph noise calculation will become available as
    attributes of the :class:`CoronagraphNoise` object after
    :func:`CoronagraphNoise.run_count_rates` is called.
    """
    def __init__(self, telescope = Telescope(), planet = Planet(),
                 star = Star(), texp = 10.0, wantsnr=10.0, FIX_OWA = False,
                 COMPUTE_LAM = False, SILENT = False, NIR = False,
                 THERMAL = True, GROUND = False, vod=False, set_fpa=None,
                 roll_maneuver = True):
        """
        """
        self.telescope = telescope
        self.planet = planet
        self.star = star
        self.texp = texp
        self.wantsnr = wantsnr
        self.FIX_OWA = FIX_OWA
        self.COMOUTE_LAM = COMPUTE_LAM
        self.SILENT = SILENT
        self.NIR = NIR
        self.THERMAL = THERMAL
        self.GROUND = GROUND
        self.vod = vod
        self.set_fpa = set_fpa
        self.roll_maneuver = roll_maneuver

        self._computed = False

        return

    def run_count_rates(self, Ahr, lamhr, solhr):
        """
        Calculate the photon count rates and signal to noise on a
        coronagraph observation given a wavelength-dependent planetary
        geometric albedo and stellar flux density.

        Parameters
        ----------
        Ahr : array
            High-res, wavelength-dependent planetary geometric albedo
        lamhr : array
            High-res wavelength grid  [um]
        solhr : array
            High-res TOA solar spectrum [W/m**2/um]

        Calling ``run_count_rates()`` creates the following attributes for
        the ``CoronagraphNoise`` instance:

        Attributes
        ----------
        Ahr : array
            High-res, wavelength-dependent planetary geometric albedo
        lamhr : array
            High-res wavelength grid  [um]
        solhr : array
            High-res TOA solar spectrum [W/m**2/um]
        lam : array
            Observed wavelength grid [$\mu$m]
        dlam : array
            Observed wavelength grid widths [$\mu$m]
        A : array
            Planetary geometric albedo at observed resolution
        Cratio : array
            Planet-to-star flux contrast ratio
        cp : array
            Planetary photon count rate [photons/s]
        csp : array
            Speckle count rate [photons/s]
        cz : array
            Zodi photon count rate [photons/s]
        cez : array
            Exo-zodi photon count rate [photons/s]
        cth : array
            Thermal photon count rate [photons/s]
        cD : array
            Dark current photon count rate [photons/s]
        cR : array
            Read noise photon count rate [photons/s]
        cc : array
            Clock induced charge photon count rate [photons/s]
        cb : array
            Total background photon noise count rate [photons/s]
        DtSNR : array
            Integration time to ``wantsnr`` [hours]
        SNRt : array
            S/N in a ``texp`` hour exposure
        Aobs : array
            Observed albedo with noise
        Asig : array
            Observed uncertainties on albedo
        Cobs : array
            Observed Fp/Fs with noise
        Csig : array
            Observed uncertainties on Fp/Fs
        """

        # Save input arrays
        self.Ahr = Ahr
        self.lamhr = lamhr
        self.solhr = solhr

        # Aperture logic
        accepted_circular = ["circular", "circ", "c"]
        accepted_square = ["square", "s"]
        if self.telescope.aperture.lower() in accepted_circular:
            CIRC = True
        elif self.telescope.aperture.lower() in accepted_square:
            CIRC = False
        else:
            assert False, "telescope.aperture is invalid"

        # Call count_rates
        lam, dlam, A, q, Cratio, cp, csp, cz, cez, cD, cR, cth, cc, DtSNR = \
            count_rates(Ahr, lamhr, solhr,
                        alpha = self.planet.alpha,
                        Phi = self.planet.Phi,
                        Rp = self.planet.Rp,
                        Teff = self.star.Teff,
                        Rs = self.star.Rs,
                        r = self.planet.a,
                        d = self.planet.distance,
                        Nez = self.planet.Nez,
                        mode = self.telescope.mode,
                        filter_wheel = self.telescope.filter_wheel,
                        lammin = self.telescope.lammin,
                        lammax = self.telescope.lammax,
                        Res    = self.telescope.resolution,
                        diam   = self.telescope.diameter,
                        Tput   = self.telescope.throughput,
                        C      = self.telescope.contrast,
                        IWA    = self.telescope.IWA,
                        OWA    = self.telescope.OWA,
                        Tsys   = self.telescope.Tsys,
                        Tdet   = self.telescope.Tdet,
                        emis   = self.telescope.emissivity,
                        De     = self.telescope.darkcurrent,
                        DNHpix = self.telescope.DNHpix,
                        Re     = self.telescope.readnoise,
                        Rc     = self.telescope.Rc,
                        Dtmax  = self.telescope.Dtmax,
                        X      = self.telescope.X,
                        qe     = self.telescope.qe,
                        MzV    = self.planet.MzV,
                        MezV   = self.planet.MezV,
                        A_collect = self.telescope.A_collect,
                        diam_circumscribed = self.telescope.diam_circumscribed,
                        diam_inscribed = self.telescope.diam_inscribed,
                        lam    = self.telescope.lam,
                        dlam   = self.telescope.dlam,
                        Tput_lam = self.telescope.Tput_lam,
                        qe_lam = self.telescope.qe_lam,
                        lammin_lenslet = self.telescope.lammin_lenslet,
                        NIR    = self.NIR,
                        GROUND = self.GROUND,
                        THERMAL = self.THERMAL,
                        CIRC = CIRC,
                        roll_maneuver = self.roll_maneuver,
                        SILENT = self.SILENT,
                        wantsnr = self.wantsnr
                    )

        # Save output arrays
        self.lam     = lam
        self.dlam    = dlam
        self.A       = A
        self.Cratio  = Cratio
        self.cp      = cp
        self.csp     = csp
        self.cz      = cz
        self.cez     = cez
        self.cD      = cD
        self.cR      = cR
        self.cth     = cth
        self.cc      = cc
        self.cb      = cz + cez + csp + cD + cR + cth + cc
        self.DtSNR   = DtSNR

        # Flip the switch
        self._computed = True

        # Make an initial set of fake data
        self.make_fake_data()

        return

    def make_fake_data(self, texp = None):
        """
        Make a fake/synthetic dataset by sampling from a Gaussian.

        Parameters
        ----------
        texp : float, optional
            Exposure time [hours]. If not provided, the ``CoronagraphNoise.texp``
            will be used by default.

        Calling ``make_fake_data()`` creates the following attributes for
        the ``CoronagraphNoise`` instance:

        Attributes
        ----------
        SNRt : array
            S/N in a ``texp`` hour exposure
        Aobs : array
            Observed albedo with noise
        Asig : array
            Observed uncertainties on albedo
        Cobs : array
            Observed Fp/Fs with noise
        Csig : array
            Observed uncertainties on Fp/Fs
        """

        # Ensure that simulation has been run
        assert self._computed

        # Allow new exposure time
        if texp is not None:
            self.texp = texp

        # Convert exposure time to seconds
        Dt = 3600. * self.texp

        # Use telescope roll maneuver
        if self.roll_maneuver:
            # assuming background subtraction (the "2")
            roll_factor = 2.0
        else:
            # standard background noise
            roll_factor = 1.0

        # Calculate signal-to-noise
        SNRt  = self.cp * Dt / np.sqrt((self.cp + roll_factor*self.cb) * Dt)

        # Calculate 1-sigma errors on contrast ratio and albedo
        Csig = self.Cratio/SNRt
        Asig = self.A/SNRt

        # Calculate Gaussian noise
        gaus = np.random.randn(len(self.Cratio))

        # Add gaussian noise to observed data
        Cobs = self.Cratio + Csig * gaus
        Aobs = self.A + Asig * gaus

        # Save attributes
        self.SNRt = SNRt
        self.Asig = Asig
        self.Aobs = Aobs
        self.Csig = Csig
        self.Cobs = Cobs

        return

    def plot_spectrum(self, SNR_threshold = 1.0, Nsig = 6.0, ax0 = None,
                      err_kws = {"fmt" : ".", "c" : "k", "alpha" : 1},
                      plot_kws = {"lw" : 1.0, "c" : "C4", "alpha" : 0.5},
                      draw_box = True):
        """
        Plot noised direct-imaging spectrum.

        Parameters
        ----------
        SNR_threshold : float
            Threshold SNR below which do not plot
        Nsig : float
            Number of standard deviations about median observed points to set
            yaxis limits
        ax0 : `matplotlib.axes`
            Optional axis to provide
        err_kws : dic
            Keyword arguments for `errorbar`
        plot_kws : dic
            Keyword arguments for `plot`
        draw_box : bool
            Draw important quantities in a box?

        Returns
        -------
        fig : `matplotlib.figure.Figure`
            Returns a figure if `ax0` is `None`
        ax : `matplotlib.axes`
            Returns an axis if `ax0` is `None`

        Note
        ----
        Only returns `fig` and `ax` is ``ax0 is None``
        """

        m = [self.SNRt > SNR_threshold]

        scale = 1

        if ax0 is None:
            # Create Plot
            fig, ax = plt.subplots(figsize = (10,8))
            ax.set_xlabel(r"Wavelength [$\mu$m]")
            ax.set_ylabel("Geometic Albedo")
        else:
            ax = ax0

        #ax.plot(lam, scale*RpRs2, alpha = 1.0, ls = "steps-mid")
        ax.errorbar(self.lam[m], scale*self.Aobs[m], yerr=scale*self.Asig[m], zorder = 100, **err_kws)
        #ax.set_yscale("log")

        # Set ylim
        mederr = scale*np.median(self.Asig)
        medy = scale*np.median(self.Aobs)
        ax.set_ylim([medy - Nsig*mederr, medy + Nsig*mederr])

        ylims = ax.get_ylim()
        xlims = ax.get_xlim()

        ax.plot(self.lamhr, scale*self.Ahr, **plot_kws)

        ax.set_ylim(ylims)
        ax.set_xlim(xlims)

        if draw_box:
            # Set string for plot text
            if self.texp > 2.0:
                timestr = "{:.0f}".format(self.texp)+' hours'
            else:
                timestr = "{:.0f}".format(self.texp*60)+' mins'
            plot_text = r'Distance = '+"{:.1f}".format(self.planet.distance)+' pc'+\
            '\n Integration time = '+timestr
            ax.text(0.02, 0.975, plot_text, transform=ax.transAxes, ha = "left", va = "top",
                    bbox=dict(boxstyle="square", fc="w", ec="k", alpha=0.9), zorder=101)

        #ax.legend()

        if ax0 is None:
            return fig, ax
        else:
            return

    def plot_SNR(self, ax0 = None, plot_kws = {"ls" : "steps-mid"}):
        """
        Plot the S/N on the planet as a function of wavelength.

        Parameters
        ----------
        ax0 : `matplotlib.axes`
            Optional axis to provide
        plot_kws : dic
            Keyword arguments for `plot`

        Returns
        -------
        fig : `matplotlib.figure.Figure`
            Returns a figure if `ax0` is `None`
        ax : `matplotlib.axes`
            Returns an axis if `ax0` is `None`

        Note
        ----
        Only returns `fig` and `ax` is ``ax0 is None``
        """
        if ax0 is None:
            # Create Plot
            fig, ax = plt.subplots(figsize = (10,8))
        else:
            ax = ax0

        ax.plot(self.lam, self.SNRt, **plot_kws)
        #ax.set_yscale("log")
        ax.set_xlabel(r"Wavelength [$\mu$m]")
        ax.set_ylabel("S/N on Planet Spectrum in %.1f hrs" %self.texp)
        #ax.legend()

        if ax0 is None:
            return fig, ax
        else:
            return

    def plot_time_to_wantsnr(self, ax0 = None, plot_kws = {"ls" : "steps-mid", "alpha" : 1.0}):
        """
        Plot the exposure time to get a SNR on the planet spectrum.

        Parameters
        ----------
        ax0 : `matplotlib.axes`
            Optional axis to provide
        plot_kws : dic
            Keyword arguments for `plot`

        Returns
        -------
        fig : `matplotlib.figure.Figure`
            Returns a figure if `ax0` is `None`
        ax : `matplotlib.axes`
            Returns an axis if `ax0` is `None`

        Note
        ----
        Only returns `fig` and `ax` is ``ax0 is None``
        """

        if ax0 is None:
            # Create Plot
            fig, ax = plt.subplots(figsize = (10,8))
            ax.set_xlabel(r"Wavelength [$\mu$m]")
            ax.set_ylabel("Hours to S/N = %i on Planet Spectrum" %self.wantsnr)
            ax.set_yscale("log")
        else:
            ax = ax0

        ax.plot(self.lam, self.DtSNR, **plot_kws)

        if ax0 is None:
            return fig, ax
        else:
            return

def count_rates(Ahr, lamhr, solhr,
                alpha, Phi, Rp, Teff, Rs, r, d, Nez,
                mode   = "IFS",
                filter_wheel = None,
                lammin = 0.4,
                lammax = 2.5,
                Res    = 70.0,
                diam   = 10.0,
                Tput   = 0.20,
                C      = 1e-10,
                IWA    = 3.0,
                OWA    = 20.0,
                Tsys   = 260.0,
                Tdet   = 50.0,
                emis   = 0.9,
                De     = 1e-4,
                DNHpix = 3.0,
                Re     = 0.1,
                Rc     = 0.0,
                Dtmax  = 1.0,
                X      = 1.5,
                qe     = 0.9,
                MzV    = 23.0,
                MezV   = 22.0,
                A_collect = None,
                diam_circumscribed = None,
                diam_inscribed = None,
                lam    = None,
                dlam   = None,
                Tput_lam = None,
                qe_lam = None,
                lammin_lenslet = None,
                wantsnr=10.0, FIX_OWA = False, COMPUTE_LAM = False,
                SILENT = False, NIR = False, THERMAL = False, GROUND = False,
                vod=False, set_fpa=None, CIRC = True, roll_maneuver = True):
    """
    Runs coronagraph model (Robinson et al., 2016) to calculate planet and noise
    photon count rates for specified telescope and system parameters.

    Parameters
    ----------
    Ahr : array
        High-res, wavelength-dependent planetary geometric albedo
    lamhr : array
        High-res wavelength grid  [um]
    solhr : array
        High-res TOA solar spectrum [W/m**2/um]
    alpha : float
        Planet phase angle [deg]
    Phi : float
        Planet phase function
    Rp : float
        Planet radius [R_earth]
    Teff : float
        Stellar effective temperature [K]
    Rs : float
        Stellar radius [R_sun]
    r : float
        Planet semi-major axis [AU]
    d : float
        Distance to observed star-planet system [pc]
    Nez : float
        Number of exozodis in exoplanetary disk
    mode : str, optional
        Telescope observing mode: "IFS" or "Imaging"
    filter_wheel : Wheel, optional
        Wheel object containing imaging filters
    lammin : float, optional
        Minimum wavelength [um]
    lammax : float, optional
        Maximum wavelength [um]
    Res : float, optional
        Instrument spectral resolution (``lam / dlam``)
    diam : float, optional
        Telescope diameter [m]
    Tput : float, optional
        Telescope and instrument throughput
    C : float, optional
        Coronagraph design contrast
    IWA : float, optional
        Coronagraph Inner Working Angle (``lam / diam``)
    OWA : float, optional
        Coronagraph Outer Working Angle (``lam / diam``)
    Tsys  : float, optional
        Telescope mirror temperature [K]
    Tdet  : float, optional
        Telescope detector temperature [K]
    emis : float, optional
        Effective emissivity for the observing system (of order unity)
    De : float, optional
        Dark current [counts/s]
    DNHpix : float, optional
        Number of horizontal/spatial pixels for dispersed spectrum
    Re : float, optional
        Read noise counts per pixel
    Rc : float, optional
        Clock induced charge [counts/pixel/photon]
    Dtmax : float, optional
        Detector maximum exposure time [hours]
    X : float, optional
        Width of photometric aperture (``lam / diam``)
    qe : float, optional
        Detector quantum efficiency
    MzV : float, optional
        V-band zodiacal light surface brightness [mag/arcsec**2]
    MezV : float, optional
        V-band exozodiacal light surface brightness [mag/arcsec**2]
    A_collect : float, optional
        Mirror collecting area (m**2) (uses :math:`\pi(D/2)^2` by default)
    diam_circumscribed : float, optional
        Circumscribed telescope diameter [m] used for IWA and OWA (uses `diam`
        if `None` provided)
    diam_inscribed : float, optional
        Inscribed telescope diameter [m] used for lenslet calculations
        (uses `diam` if `None` provided)
    lam : array-like, optional
        Wavelength grid for spectrograph [microns] (uses ``lammin``, ``lammax``,
        and ``resolution`` to determine if ``None`` provided)
    dlam : array-like, optional
        Wavelength grid `widths` for spectrograph [microns] (uses ``lammin``, ``lammax``,
        and ``resolution`` to determine if ``None`` provided)
    Tput_lam : tuple of arrays
        Wavelength-dependent throughput e.g. ``(wls, tputs)``
    qe_lam : tuple of arrays
        Wavelength-dependent throughput e.g. ``(wls, qe)``
    lammin_lenslet : float, optional
        Minimum wavelength to use for lenslet calculation (default is ``lammin``)
    wantsnr : float, optional
        Desired signal-to-noise ratio in each pixel
    FIX_OWA : bool, optional
        Set to fix OWA at ``OWA*lammin/D``, as would occur if lenslet array is
        limiting the OWA
    COMPUTE_LAM : bool, optional
        Set to compute lo-res wavelength grid, otherwise the grid input as
        variable ``lam`` is used
    SILENT : bool, optional
        Set to suppress print statements
    NIR : bool, optional
        Re-adjusts pixel size in NIR, as would occur if a second instrument
        was designed to handle the NIR
    THERMAL : bool, optional
        Set to compute thermal photon counts due to telescope temperature
    GROUND : bool, optional
        Set to simulate ground-based observations through atmosphere
    vod : bool, optional
        "Valley of Death" red QE parameterization from Robinson et al. (2016)
    set_fpa : float, optional
        Specify the fraction of planetary signal in Airy pattern, default will
        calculate it from the photometric aperture size `X`
    CIRC : bool, optional
        Set to use a circular aperture
    roll_maneuver : bool, optional
        This assumes an extra factor of 2 hit to the background noise due to a
        telescope roll maneuver needed to subtract out the background. See
        Brown (2005) for more details.

    Returns
    -------
    lam : ndarray
        Observational wavelength grid [um]
    dlam : ndarray
        Observational spectral element width [um]
    A : ndarray
        Planetary geometric albedo on lam grid
    q : ndarray
        Quantum efficiency grid
    Cratio : ndarray
        Planet-star contrast ratio
    cp : ndarray
        Planetary photon count rate on detector [1/s]
    csp : ndarray
        Speckle photon count rate on detector [1/s]
    cz : ndarray
        Zodiacal photon count rate on detector [1/s]
    cez : ndarray
        Exozodiacal photon count rate on detector [1/s]
    cD : ndarray
        Dark current photon count rate on detector [1/s]
    cR : ndarray
        Read noise photon count rate on detector [1/s]
    cth : ndarray
        Instrument thermal photon count rate on detector [1/s]
    cc : ndarray
        Clock induced charge photon count rate [1/s]
    DtSNR : ndarray
        Exposure time required to get desired S/N (wantsnr) [hours]
    """

    convolution_function = downbin_spec
    #convolution_function = degrade_spec

    # Define a diameter for IWA (circumscribed),
    # collecting area, and lenslet (inscribed)
    if diam_inscribed is None:
        # Defaults to diam
        diam_inscribed = diam
    if A_collect is None:
        # Defaults to diam
        diam_collect = diam
    else:
        # Calculated from provided collecting area
        diam_collect = 2. * (A_collect / np.pi)**0.5
    if diam_circumscribed is None:
        # Defaults to diam
        diam_circumscribed = diam

    # Configure for different telescope observing modes
    if mode == 'Imaging':
        filters = filter_wheel
        IMAGE = True
        COMPUTE_LAM = False
        # sorted filter dict by bandcenters
        tdict = sorted(filters.__dict__.items(), key=lambda x: x[1].bandcenter)
        # Construct array of wavelengths
        lam = np.array([x[1].bandcenter for x in tdict])
        # Construct array of wavelength bin widths (FWHM)
        dlam = np.array([x[1].FWHM for x in tdict])
        Nlam = len(lam)
    elif mode == 'IFS':
        IMAGE = False
        COMPUTE_LAM = True
    else:
        print("Invalid telescope observing mode. Select 'IFS', or 'Imaging'.")
        sys.exit()

    # fraction of planetary signal in Airy pattern
    if set_fpa is None:
        fpa = f_airy(X)
    else:
        fpa = set_fpa * f_airy(X)

    # Set wavelength grid
    # GENERALIZE THIS:
    if COMPUTE_LAM:
        if (lam is None) or (dlam is None):
            lam, dlam = construct_lam(lammin, lammax, Res)
    elif IMAGE:
        pass
    else:
        # Throw error
        print("Error in make_noise: Not computing wavelength grid or providing filters!")
        return None

    # Set Quantum Efficiency
    q = set_quantum_efficiency(lam, qe, NIR=NIR, vod=vod)

    # Set Dark current and Read noise
    De = set_dark_current(lam, De, lammax, Tdet, NIR=NIR)
    Re = set_read_noise(lam, Re, NIR=NIR)

    # Set Angular size of lenslet
    if lammin_lenslet is None: lammin_lenslet = lammin
    theta = set_lenslet(lam, lammin_lenslet, diam_inscribed, X, NIR=True)

    # Set throughput (for inner and outer working angle cutoffs)
    sep  = r/d*np.sin(alpha*np.pi/180.)*np.pi/180./3600. # separation in radians
    T = set_throughput(lam, Tput, diam_circumscribed, sep, IWA, OWA, lammin, FIX_OWA=FIX_OWA, SILENT=SILENT)

    # Apply wavelength-dependent throuput, if needed
    if Tput_lam is not None:
        # Bin input throughput curve to native res
        Tlam = np.interp(lam, Tput_lam[0], Tput_lam[1])
        # Multiply into regular throughput
        T = T * Tlam

    # Apply wavelength-dependent quantum efficiency, if needed
    if qe_lam is not None:
        # Bin input QE curve to native res
        qlam = np.interp(lam, qe_lam[0], qe_lam[1])
        # Multiply into regular QE
        q = q * qlam

    # Modify throughput by atmospheric transmission if GROUND-based
    if GROUND:
        #if GROUND == "ESO":
            # Use ESO SKYCALC
        #    pass
        #else:
        # Use SMART calc
        Tatmos = set_atmos_throughput(lam, dlam, convolution_function)
        # Multiply telescope throughput by atmospheric throughput
        T = T * Tatmos

    # Degrade albedo and stellar spectrum
    if COMPUTE_LAM:
        A = convolution_function(Ahr, lamhr, lam, dlam=dlam)
        Fs = convolution_function(solhr, lamhr, lam, dlam=dlam)
    elif IMAGE:
        # Convolve with filter response
        A = convolve_spec(Ahr, lamhr, filters)
        Fs = convolve_spec(solhr, lamhr, filters)
    else:
        A = Ahr
        Fs = solhr

    # Compute fluxes
    #Fs = Fstar(lam, Teff, Rs, r, AU=True) # stellar flux on planet
    Fp = Fplan(A, Phi, Fs, Rp, d)         # planet flux at telescope
    Cratio = FpFs(A, Phi, Rp, r)

    ##### Compute count rates #####
    cp     =  cplan(q, fpa, T, lam, dlam, Fp, diam_collect)                          # planet count rate
    cz     =  czodi(q, X, T, lam, dlam, diam_collect, MzV)                           # solar system zodi count rate
    cez    =  cezodi(q, X, T, lam, dlam, diam_collect, r, \
        Fstar(lam, Teff, Rs,1. , AU=True), Nez, MezV)                                    # exo-zodi count rate
    csp    =  cspeck(q, T, C, lam, dlam, Fstar(lam,Teff,Rs,d), diam_collect)         # speckle count rate
    cD     =  cdark(De, X, lam, diam_collect, theta, DNHpix, IMAGE=IMAGE)            # dark current count rate
    cR     =  cread(Re, X, lam, diam_collect, theta, DNHpix, Dtmax, IMAGE=IMAGE)     # readnoise count rate
    if THERMAL:
        cth    =  ctherm(q, X, T, lam, dlam, diam_collect, Tsys, emis)               # internal thermal count rate
    else:
        cth = np.zeros_like(cp)

    # Add earth thermal photons if GROUND
    if GROUND:
        # Use ESO SKCALC
        wl_sky, Isky = get_sky_flux()
        # Convolve to instrument resolution
        Itherm = convolution_function(Isky, wl_sky, lam, dlam=dlam)
        # Compute Earth thermal photon count rate
        cthe = ctherm_earth(q, X, T, lam, dlam, diam_collect, Itherm)
        # Add earth thermal photon counts to telescope thermal counts
        cth = cth + cthe
        '''
        if True:
            import matplotlib.pyplot as plt;
            fig2, ax1 = plt.subplots(figsize=(8,6))
            ax1.plot(lam, cthe, c="blue", ls="steps-mid", label="Earth Thermal")
            ax1.plot(lam, cth, c="red", ls="steps-mid", label="Telescope Thermal")
            ax1.plot(lam, cp, c="k", ls="steps-mid", label="Planet")
            ax1.set_ylabel("Photon Count Rate [1/s]")
            ax1.set_xlabel("Wavelength [um]")
            ax1.legend()
            plt.show()
        '''

    # Clock induced charge photon count rate
    # Calculate photon count rate in the scene (everything except read noise)
    cscene = cp + cz + cez + csp + cD + cth
    # Calculate the clock induced charge photon count rate
    cc = ccic(Rc, cscene, X, lam, diam_collect, theta, DNHpix, Dtmax,
              IMAGE=IMAGE, CIRC=CIRC)

    # Calculate total background counts
    cb = (cz + cez + csp + cD + cR + cth + cc)

    # Use telescope roll maneuver
    if roll_maneuver:
        # assuming background subtraction (the "2")
        roll_factor = 2.0
    else:
        # standard background noise
        roll_factor = 1.0

    # Calculate total noise
    cnoise =  cp + roll_factor*cb

    # Calculate total counts
    ctot = cp + cz + cez + csp + cD + cR + cth + cc

    '''
    Giada: where does the factor of 2 come from [above]?

    Ty (Via email): That's due to "background subtraction".
    If you were to take a single exposure, and had the ability
    to post-process the image to the Poisson noise limit, you
    wouldn't have the factor of two.  However, it's not yet
    clear that we'll be able to reach the Poisson, single observation limit.
    Instead, the current idea is that you take two observations (with
    exposure time Delta t/2), with the telescope rotated by a
    small amount between exposures, and then subtract the two images.
    So, for a fixed exoplanet count (cp), the roll technique would
    give you 2x as many noise counts due to background sources as
    would the single-observation technique.
    See also the top of page 4 of Brown (2005).
    '''

    # Exposure time to SNR
    DtSNR = exptime_element(lam, cp, cnoise, wantsnr)

    return lam, dlam, A, q, Cratio, cp, csp, cz, cez, cD, cR, cth, cc, DtSNR


"""
Detector class

----

.. include license and copyright
.. include:: ../include/copy.rst

----

.. include common links, assuming primary doc root is up one directory
.. include:: ../include/links.rst

"""

import os
import numpy

from .efficiency import Efficiency

class Detector(Efficiency):
    """
    Define the detector statistics.

    Args:
        shape (:obj:`tuple`, optional):
            Dimensions of the detector in number of pixels along the
            spectral axis and number of pixels along the spatial
            axis. Can be None, but limits use if it is.
        pixelsize (:obj:`float`, optional):
            The size of the (square) detector pixels in *micron*.
        rn (:obj:`float`, optional):
            Read-noise in electrons.
        dark (:obj:`float`, optional):
            Dark current in electrons per second.
        gain (:obj:`float`, optional):
            Gain of detector amplifier in e- per ADU.
        fullwell (:obj:`float`, optional):
            The full well of the pixels in e-.
        nonlinear (:obj:`float`, optional):
            The fraction of the fullwell above which the detector
            response is nonlinear.
        qe (:obj:`float`, :class:`Efficiency`, optional):
            Detector quantum efficiency.
    """
    # TODO: Allow for multiple amplifiers per detector? Would also need
    # to define amplifier section.
    # TODO: Define overscan and data sections
    def __init__(self, shape=None, pixelsize=15., rn=1., dark=0., gain=1., fullwell=1e4,
                 nonlinear=1., qe=0.9):
        if shape is not None and len(shape) != 2:
            raise ValueError('Shape must contain two integers.')
        self.shape = shape
        self.pixelsize = pixelsize
        self.rn = rn
        self.dark = dark
        self.gain = gain
        self.fullwell = fullwell
        self.nonlinear = nonlinear
        if not isinstance(qe, (Efficiency, float)):
            raise TypeError('Provided quantum efficiency must be type `float` or `Efficiency`.')
        if isinstance(qe, float):
            super(Detector, self).__init__(qe)
        else:
            super(Detector, self).__init__(qe.eta, wave=qe.wave)