import numpy as np
import os
import copy
import galsim
from galsim import DeVaucouleurs
from galsim import Exponential
import descwl

from .shifts import get_shifts, get_pair_shifts
from .constants import SCALE
from .cache_tools import cached_catalog_read

DEFAULT_FIXED_GAL_CONFIG = {
    "mag": 17.0,
    "hlr": 0.5,
    "morph": "exp",
}

def make_galaxy_catalog(
    *,
    rng,
    gal_type,
    coadd_dim=None,
    buff=0,
    layout=None,
    gal_config=None,
    sep=None,
):
    """
    rng: numpy.random.RandomState
        Numpy random state
    gal_type: string
        'fixed', 'varying' or 'wldeblend'
    coadd_dim: int
        Dimensions of coadd
    buff: int, optional
        Buffer around the edge where no objects are drawn.  Ignored for
        layout 'grid'.  Default 0.
    layout: string, optional
        'grid' or 'random'.  Ignored for gal_type "wldeblend", otherwise
        required.
    gal_config: dict or None
        Can be sent for fixed galaxy catalog.  See DEFAULT_FIXED_GAL_CONFIG
        for defaults mag, hlr and morph
    sep: float, optional
        Separation of pair in arcsec for layout='pair'
    """
    if layout == 'pair':
        if sep is None:
            raise ValueError(
                f'send sep= for gal_type {gal_type} and layout {layout}'
            )
        gal_config = get_fixed_gal_config(config=gal_config)

        if gal_type in ['fixed', 'exp']:  # TODO remove exp
            cls = FixedPairGalaxyCatalog
        else:
            cls = PairGalaxyCatalog

        galaxy_catalog = cls(
            rng=rng,
            mag=gal_config['mag'],
            hlr=gal_config['hlr'],
            morph=gal_config['morph'],
            sep=sep,
        )

    else:
        if coadd_dim is None:
            raise ValueError(
                f'send coadd_dim= for gal_type {gal_type} and layout {layout}'
            )

        if gal_type == 'wldeblend':
            galaxy_catalog = WLDeblendGalaxyCatalog(
                rng=rng,
                coadd_dim=coadd_dim,
                buff=buff,
            )
        elif gal_type in ['fixed', 'varying', 'exp']:  # TODO remove exp
            if layout is None:
                raise ValueError("send layout= for gal_type '%s'" % gal_type)

            gal_config = get_fixed_gal_config(config=gal_config)

            if gal_type == 'fixed':
                cls = FixedGalaxyCatalog
            else:
                cls = GalaxyCatalog

            galaxy_catalog = cls(
                rng=rng,
                coadd_dim=coadd_dim,
                buff=buff,
                layout=layout,
                mag=gal_config['mag'],
                hlr=gal_config['hlr'],
                morph=gal_config['morph'],
            )

        else:
            raise ValueError(f'bad gal_type "{gal_type}"')

    return galaxy_catalog

def get_fixed_gal_config(config=None):
    """
    get the configuration for fixed galaxies, with defaults in place

    Parameters
    ----------
    config: dict, optional
        The input config. Over-rides defaults

    Returns
    -------
    the config dict
    """
    out_config = copy.deepcopy(DEFAULT_FIXED_GAL_CONFIG)

    if config is not None:
        for key in config:
            if key not in out_config:
                raise ValueError("bad key for fixed gals: '%s'" % key)
        out_config.update(config)
    return out_config

class FixedGalaxyCatalog(object):
    """
    Galaxies of fixed galsim type, flux, and size and shape.

    Same for all bands

    Parameters
    ----------
    rng: np.random.RandomState
        The random number generator
    coadd_dim: int
        dimensions of the coadd
    layout: string
        The layout of objects, either 'grid' or 'random'
    mag: float
        Magnitude of all objects. Objects brighter than magntiude 17 (e.g., 14
        since mags are opposite) tend to cause the Rubin Observatory science
        pipeline detection algorithm to misdetect isolted objects in unphysical
        ways. This effect causes the shear response to be non-linear and so
        metadetect will fail. For this reason, you should use the default
        magnitude of 17 or fainter for this kind of galaxy.
    hlr: float
        Half light radius of all objects
    buff: int, optional
        Buffer region with no objects, on all sides of image.  Ingored
        for layout 'grid'.  Default 0.
    morph: str
        Galaxy morphology, 'exp', 'dev' or 'bd', 'bdk'.  Default 'exp'
    """
    def __init__(self, *, rng, coadd_dim, layout, mag, hlr, buff=0, morph='exp'):
        self.gal_type = 'fixed'
        self.morph = morph
        self.mag = mag
        self.hlr = hlr

        self.shifts_array = get_shifts(
            rng=rng,
            coadd_dim=coadd_dim,
            buff=buff,
            layout=layout,
        )

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

    def get_objlist(self, *, survey):
        """
        get a list of galsim objects

        Parameters
        ----------
        band: string
            Get objects for this band.  For the fixed
            catalog, the objects are the same for every band

        Returns
        -------
        [galsim objects], [shifts]
        """

        flux = survey.get_flux(self.mag)

        sarray = self.shifts_array
        objlist = []
        shifts = []
        for i in range(len(self)):
            objlist.append(self._get_galaxy(flux))
            shifts.append(galsim.PositionD(sarray['dx'][i], sarray['dy'][i]))

        return objlist, shifts

    def _get_galaxy(self, flux):
        """
        get a galaxy object

        Parameters
        ----------
        flux: float
            Flux of object

        Returns
        --------
        galsim.GSObject
        """

        if self.morph == 'exp':
            gal = _generate_exp(hlr=self.hlr, flux=flux)
        elif self.morph == 'dev':
            gal = _generate_dev(hlr=self.hlr, flux=flux)
        elif self.morph == 'bd':
            gal = _generate_bd(hlr=self.hlr, flux=flux)
        elif self.morph == 'bdk':
            gal = _generate_bdk(hlr=self.hlr, flux=flux)
        else:
            raise ValueError(f"bad gal type '{self.morph}'")

        return gal

class GalaxyCatalog(FixedGalaxyCatalog):
    """
    Galaxies of fixed galsim type, but varying properties.

    Same for all bands

    Parameters
    ----------
    rng: np.random.RandomState
        The random number generator
    coadd_dim: int
        dimensions of the coadd
    layout: string
        The layout of objects, either 'grid' or 'random'
    mag: float
        Magnitude of all objects. Objects brighter than magntiude 17 (e.g., 14
        since mags are opposite) tend to cause the Rubin Observatory science
        pipeline detection algorithm to misdetect isolted objects in unphysical
        ways. This effect causes the shear response to be non-linear and so
        metadetect will fail. For this reason, you should use the default
        magnitude of 17 or fainter for this kind of galaxy.
    hlr: float
        Half light radius of all objects
    buff: int, optional
        Buffer region with no objects, on all sides of image.  Ingored
        for layout 'grid'.  Default 0.
    morph: str
        Galaxy morphology, 'exp', 'dev' or 'bd', 'bdk'.  Default 'exp'
    """
    def __init__(self, *, rng, coadd_dim, layout, mag, hlr, buff=0, morph='exp'):
        super().__init__(
            rng=rng, coadd_dim=coadd_dim, buff=buff, layout=layout,
            mag=mag, hlr=hlr, morph=morph,
        )
        self.gal_type = 'varying'

        # we use this to ensure the same galaxies are generated in different
        # bands
        self.morph_seed = rng.randint(0, 2**31)
        self.gs_morph_seed = rng.randint(0, 2**31)

    def get_objlist(self, *, survey):
        """
        get a list of galsim objects

        Parameters
        ----------
        band: string
            Get objects for this band.  For the fixed
            catalog, the objects are the same for every band

        Returns
        -------
        [galsim objects], [shifts]
        """

        self._morph_rng = np.random.RandomState(self.morph_seed)
        self._gs_morph_rng = galsim.BaseDeviate(seed=self.gs_morph_seed)
        return super().get_objlist(survey=survey)

    def _get_galaxy(self, flux):
        """
        get a galaxy object

        Parameters
        ----------
        flux: float
            Flux of object

        Returns
        --------
        galsim.GSObject
        """

        if self.morph == 'exp':
            gal = _generate_exp(
                hlr=self.hlr, flux=flux, vary=True, rng=self._morph_rng,
            )
        elif self.morph == 'dev':
            gal = _generate_dev(
                hlr=self.hlr, flux=flux, vary=True, rng=self._morph_rng,
            )
        elif self.morph == 'bd':
            gal = _generate_bd(
                hlr=self.hlr, flux=flux,
                vary=True, rng=self._morph_rng,
            )
        elif self.morph == 'bdk':
            gal = _generate_bdk(
                hlr=self.hlr, flux=flux,
                vary=True,
                rng=self._morph_rng, gsrng=self._gs_morph_rng,
            )
        else:
            raise ValueError(f"bad morph '{self.morph}'")

        return gal

def _generate_exp(hlr, flux, vary=False, rng=None):
    gal = Exponential(half_light_radius=hlr, flux=flux)

    if vary:
        g1, g2 = _generate_g1g2(rng)
        gal = gal.shear(g1=g1, g2=g2)

    return gal

def _generate_dev(hlr, flux, vary=False, rng=None):
    gal = DeVaucouleurs(half_light_radius=hlr, flux=flux)
    if vary:
        g1, g2 = _generate_g1g2(rng)
        gal = gal.shear(g1=g1, g2=g2)

    return gal

def _generate_bd(
    hlr, flux,
    vary=False,
    rng=None,
    max_bulge_shift_frac=0.1,  # fraction of hlr
    max_bulge_rot=np.pi/4,
):

    if vary:
        bulge_frac = _generate_bulge_frac(rng)
    else:
        bulge_frac = 0.5

    disk_frac = (1.0 - bulge_frac)

    bulge = DeVaucouleurs(half_light_radius=hlr, flux=flux * bulge_frac)
    disk = Exponential(half_light_radius=hlr, flux=flux * disk_frac)

    if vary:
        bulge = _shift_bulge(rng, bulge, hlr, max_bulge_shift_frac)

    if vary:
        g1disk, g2disk = _generate_g1g2(rng)

        g1bulge, g2bulge = g1disk, g2disk
        if vary:
            g1bulge, g2bulge = _rotate_bulge(rng, max_bulge_rot, g1bulge, g2bulge)

        bulge = bulge.shear(g1=g1bulge, g2=g2bulge)
        disk = disk.shear(g1=g1disk, g2=g2disk)

    return galsim.Add(bulge, disk)

def _generate_bdk(
    hlr, flux,
    vary=False,
    rng=None,
    gsrng=None,
    knots_hlr_frac=0.25,
    max_knots_disk_frac=0.1,  # fraction of disk light
    max_bulge_shift_frac=0.1,  # fraction of hlr
    max_bulge_rot=np.pi/4,
):

    if vary:
        bulge_frac = _generate_bulge_frac(rng)
    else:
        bulge_frac = 0.5

    all_disk_frac = (1.0 - bulge_frac)

    knots_hlr = knots_hlr_frac * hlr
    if vary:
        knots_sub_frac = _generate_knots_sub_frac(rng, max_knots_disk_frac)
    else:
        knots_sub_frac = max_knots_disk_frac

    disk_frac = (1 - knots_sub_frac) * all_disk_frac
    knots_frac = knots_sub_frac * all_disk_frac

    bulge = DeVaucouleurs(half_light_radius=hlr, flux=flux * bulge_frac)
    disk = Exponential(half_light_radius=hlr, flux=flux * disk_frac)

    if gsrng is None:
        # fixed galaxy, so fix the rng
        gsrng = galsim.BaseDeviate(123)

    knots = galsim.RandomKnots(
        npoints=10,
        half_light_radius=knots_hlr,
        flux=flux * knots_frac,
        rng=gsrng,
    )

    if vary:
        bulge = _shift_bulge(rng, bulge, hlr, max_bulge_shift_frac)

    if vary:
        g1disk, g2disk = _generate_g1g2(rng)

        g1bulge, g2bulge = g1disk, g2disk
        if vary:
            g1bulge, g2bulge = _rotate_bulge(rng, max_bulge_rot, g1bulge, g2bulge)

        bulge = bulge.shear(g1=g1bulge, g2=g2bulge)
        disk = disk.shear(g1=g1disk, g2=g2disk)
        knots = knots.shear(g1=g1disk, g2=g2disk)

    return galsim.Add(bulge, disk, knots)

def _generate_bulge_frac(rng):
    assert rng is not None, 'send rng to generate bulge fraction'
    return rng.uniform(low=0.0, high=1.0)

def _generate_g1g2(rng, std=0.2):
    assert rng is not None, 'send rng to vary shape'
    while True:
        g1, g2 = rng.normal(scale=std, size=2)
        g = np.sqrt(g1**2 + g2**2)
        if abs(g) < 0.9999:
            break

    return g1, g2

def _generate_bulge_shift(rng, hlr, max_bulge_shift_frac):
    bulge_shift = rng.uniform(low=0.0, high=max_bulge_shift_frac*hlr)
    bulge_shift_angle = rng.uniform(low=0, high=2*np.pi)
    bulge_shiftx = bulge_shift * np.cos(bulge_shift_angle)
    bulge_shifty = bulge_shift * np.sin(bulge_shift_angle)

    return bulge_shiftx, bulge_shifty

def _shift_bulge(rng, bulge, hlr, max_bulge_shift_frac):
    bulge_shiftx, bulge_shifty = _generate_bulge_shift(
        rng, hlr, max_bulge_shift_frac,
    )
    return bulge.shift(bulge_shiftx, bulge_shifty)

def _rotate_bulge(rng, max_bulge_rot, g1, g2):
    assert rng is not None, 'send rng to rotate bulge'
    bulge_rot = rng.uniform(low=-max_bulge_rot, high=max_bulge_rot/4)
    return _rotate_shape(g1, g2, bulge_rot)

def _rotate_shape(g1, g2, theta_radians):
    twotheta = 2.0 * theta_radians

    cos2angle = np.cos(twotheta)
    sin2angle = np.sin(twotheta)
    g1rot = g1 * cos2angle + g2 * sin2angle
    g2rot = -g1 * sin2angle + g2 * cos2angle

    return g1rot, g2rot

def _generate_knots_sub_frac(rng, max_knots_disk_frac):
    assert rng is not None, 'send rng to generate knots sub frac'
    return rng.uniform(low=0.0, high=max_knots_disk_frac)

class FixedPairGalaxyCatalog(FixedGalaxyCatalog):
    """
    A pair of galaxies of fixed galsim type, flux, and size

    Same for all bands

    Parameters
    ----------
    rng: np.random.RandomState
        The random number generator
    mag: float
        Magnitude of all objects. Objects brighter than magntiude 17 (e.g., 14
        since mags are opposite) tend to cause the Rubin Observatory science
        pipeline detection algorithm to misdetect isolted objects in unphysical
        ways. This effect causes the shear response to be non-linear and so
        metadetect will fail. For this reason, you should use the default
        magnitude of 17 or fainter for this kind of galaxy.
    hlr: float
        Half light radius of all objects
    sep: float
        Separation of pair in arcsec
    morph: str
        Galaxy morphology, 'exp', 'dev' or 'bd', 'bdk'.  Default 'exp'
    """
    def __init__(self, *, rng, mag, hlr, sep, morph='exp'):
        self.gal_type = 'fixed'
        self.morph = morph
        self.mag = mag
        self.hlr = hlr
        self.rng = rng

        self.shifts_array = get_pair_shifts(
            rng=rng,
            sep=sep,
        )

class PairGalaxyCatalog(GalaxyCatalog):
    """
    A pair of galaxies of fixed galsim type, flux, and size

    Same for all bands

    Parameters
    ----------
    rng: np.random.RandomState
        The random number generator
    mag: float
        Magnitude of all objects. Objects brighter than magntiude 17 (e.g., 14
        since mags are opposite) tend to cause the Rubin Observatory science
        pipeline detection algorithm to misdetect isolted objects in unphysical
        ways. This effect causes the shear response to be non-linear and so
        metadetect will fail. For this reason, you should use the default
        magnitude of 17 or fainter for this kind of galaxy.
    hlr: float
        Half light radius of all objects
    sep: float
        Separation of pair in arcsec
    morph: str
        Galaxy morphology, 'exp', 'dev' or 'bd', 'bdk'.  Default 'exp'
    """
    def __init__(self, *, rng, mag, hlr, sep, morph='exp'):
        self.gal_type = 'varying'
        self.morph = morph
        self.mag = mag
        self.hlr = hlr
        self.rng = rng

        self.morph_seed = rng.randint(0, 2**31)
        self.gs_morph_seed = rng.randint(0, 2**31)

        self.shifts_array = get_pair_shifts(
            rng=rng,
            sep=sep,
        )

class WLDeblendGalaxyCatalog(object):
    """
    Catalog of galaxies from wldeblend

    Parameters
    ----------
    rng: np.random.RandomState
        The random number generator
    coadd_dim: int
        Dimensions of the coadd
    buff: int, optional
        Buffer region with no objects, on all sides of image.  Ingored
        for layout 'grid'.  Default 0.
    """
    def __init__(self, *, rng, coadd_dim, buff=0):
        self.gal_type = 'wldeblend'
        self.rng = rng

        self._wldeblend_cat = read_wldeblend_cat(rng)

        # one square degree catalog, convert to arcmin
        gal_dens = self._wldeblend_cat.size / (60 * 60)
        area = ((coadd_dim - 2*buff)*SCALE/60)**2
        nobj_mean = area * gal_dens
        nobj = rng.poisson(nobj_mean)

        self.shifts_array = get_shifts(
            rng=rng,
            coadd_dim=coadd_dim,
            buff=buff,
            layout="random",
            nobj=nobj,
        )

        num = len(self)
        self.indices = self.rng.randint(
            0,
            self._wldeblend_cat.size,
            size=num,
        )

        self.angles = self.rng.uniform(low=0, high=360, size=num)

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

    def get_objlist(self, *, survey):
        """
        get a list of galsim objects

        Parameters
        ----------
        survey: WLDeblendSurvey
            The survey object

        Returns
        -------
        [galsim objects], [shifts]
        """

        builder = descwl.model.GalaxyBuilder(
            survey=survey.descwl_survey,
            no_disk=False,
            no_bulge=False,
            no_agn=False,
            verbose_model=False,
        )

        band = survey.filter_band

        sarray = self.shifts_array
        objlist = []
        shifts = []
        for i in range(len(self)):
            objlist.append(self._get_galaxy(builder, band, i))
            shifts.append(galsim.PositionD(sarray['dx'][i], sarray['dy'][i]))

        return objlist, shifts

    def _get_galaxy(self, builder, band, i):
        """
        Get a galaxy

        Parameters
        ----------
        builder: descwl.model.GalaxyBuilder
            Builder for this object
        band: string
            Band string, e.g. 'r'
        i: int
            Index of object

        Returns
        -------
        galsim.GSObject
        """
        index = self.indices[i]

        angle = self.angles[i]

        galaxy = builder.from_catalog(
            self._wldeblend_cat[index],
            0,
            0,
            band,
        ).model.rotate(
            angle * galsim.degrees,
        )

        return galaxy

def read_wldeblend_cat(rng):
    """
    Read the catalog from the cache, but update the position angles each time

    Parameters
    ----------
    rng: np.random.RandomState
        The random number generator

    Returns
    -------
    array with fields
    """
    fname = os.path.join(
        os.environ.get('CATSIM_DIR', '.'),
        'OneDegSq.fits',
    )

    # not thread safe
    cat = cached_catalog_read(fname)
    return cat


from typing import List

import numpy

import torch

def intersect(box_a, box_b):
    """ We resize both tensors to [A,B,2] without new malloc:
    [A,2] -> [A,1,2] -> [A,B,2]
    [B,2] -> [1,B,2] -> [A,B,2]
    Then we compute the area of intersect between box_a and box_b.
    Args:
      box_a: (tensor) bounding boxes, Shape: [A,4].
      box_b: (tensor) bounding boxes, Shape: [B,4].
    Return:
      (tensor) intersection area, Shape: [A,B].
    """
    A = box_a.size(0)
    B = box_b.size(0)
    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
                       box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
                       box_b[:, :2].unsqueeze(0).expand(A, B, 2))
    inter = torch.clamp((max_xy - min_xy), min=0)
    return inter[:, :, 0] * inter[:, :, 1]
    # inter[:, :, 0] is the width of intersection and inter[:, :, 1] is height

def jaccard_tensor(box_a: torch.Tensor, box_b: torch.Tensor) -> torch.Tensor:
    """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap
    is simply the intersection over union of two boxes.  Here we operate on
    ground truth boxes and default boxes.
    E.g.:
        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
    Args:
        box_a: (tensor) Ground truth bounding boxes, Shape: [A,4]
        box_b: (tensor) Prior boxes from priorbox layers, Shape: [B,4]
    Return:
        jaccard overlap: (tensor) Shape: [A, B]
    """
    inter = intersect(box_a, box_b)
    area_a = ((box_a[:, 2] - box_a[:, 0]) *
              (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
    area_b = ((box_b[:, 2] - box_b[:, 0]) *
              (box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
    union = area_a + area_b - inter
    return inter / union  # [A,B]

class YoloV3Loss(torch.nn.Module):
    """
    YoloV3 损失函数
    """

    def __init__(self, config: dict) -> None:
        super().__init__()

        self.lambda_xy = 1.0  # 预测框中心误差权重
        self.lambda_wh = 1.0  # 预测框大小误差权重
        self.lambda_noobj = 1.0  # 预测框置信度误差权重
        self.lambda_obj = 1.0  # 预测框置信度误差权重
        self.lambda_class = 1.0  # 预测框类别误差权重
        self.lambda_conf = 1.0  # 预测框类别误差权重

        self.normd_anchors = numpy.asarray(config["anchors"]).astype(numpy.float)
        self.normd_anchors[:, :, 0] /= config["image_width"]
        self.normd_anchors[:, :, 1] /= config["image_height"]
        self.normd_anchors = self.normd_anchors.reshape((9, 2))
        self.normd_anchors_box = torch.cat(
            (
                torch.zeros((self.normd_anchors.shape[0], 2)),
                torch.from_numpy(self.normd_anchors)
            ), 1)

        self.classes = config["classes"]
        self.bbox_attrs = 4 + 1 + self.classes

        self.ignore_threshold = 0.5  # iou 忽略的阈值

        self.cuda = config["cuda"]

    def pyramid_target(self, tensord_target_list: List[torch.Tensor]):
        pyramid_target_list_13 = []
        pyramid_target_list_26 = []
        pyramid_target_list_52 = []

        for tensord_target in tensord_target_list:
            tensord_target_13 = []
            tensord_target_26 = []
            tensord_target_52 = []

            target_box = tensord_target[:, :4].clone().detach()
            target_box[:, 0] = 0
            target_box[:, 1] = 0
            normd_anch_ious = jaccard_tensor(target_box, self.normd_anchors_box)
            max_anch_ious_index = torch.argmax(normd_anch_ious, dim=-1)

            for box_index, anch_index in enumerate(max_anch_ious_index):
                if anch_index in [0, 1, 2]:
                    tensord_target_13.append(tensord_target[box_index].numpy())
                elif anch_index in [3, 4, 5]:
                    tensord_target_26.append(tensord_target[box_index].numpy())
                elif anch_index in [6, 7, 8]:
                    tensord_target_52.append(tensord_target[box_index].numpy())
                else:
                    raise Exception("unexpected error")

            pyramid_target_list_13.append(torch.as_tensor(tensord_target_13))
            pyramid_target_list_26.append(torch.as_tensor(tensord_target_26))
            pyramid_target_list_52.append(torch.as_tensor(tensord_target_52))

        return pyramid_target_list_13, pyramid_target_list_26, pyramid_target_list_52

    def decode_pyramid_target(self,
                              pyramid_target_list: List[torch.Tensor],
                              pyramid_normd_anchors: numpy.ndarray,
                              pyramid_features: int
                              ) -> (
            torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,
            torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,
            torch.Tensor, torch.Tensor
    ):
        assert pyramid_features in [13, 26, 52]

        if pyramid_features == 13:
            pyramid_anch_index_list = [0, 1, 2]
        elif pyramid_features == 26:
            pyramid_anch_index_list = [3, 4, 5]
        elif pyramid_features == 52:
            pyramid_anch_index_list = [6, 7, 8]
        else:
            raise Exception("unexpected error")

        batch_size = len(pyramid_target_list)

        target_x = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        target_y = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        target_w = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        target_h = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)

        target_loss_weight_xw = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        target_loss_weight_yh = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)

        target_obj_conf = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        target_class_conf_list = torch.zeros(batch_size, 3, pyramid_features, pyramid_features,
                                             self.classes)

        target_obj_mask = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        target_noobj_mask = torch.ones(batch_size, 3, pyramid_features, pyramid_features)

        # 遍历这一个批次所有的图片
        for bs_i, pyramid_target in enumerate(pyramid_target_list):
            if pyramid_target.shape[0] == 0:
                continue

            truth_feature_box = pyramid_target[:, 0:4] * pyramid_features

            truth_grid_x = torch.floor(truth_feature_box[:, 0]).int()
            truth_grid_y = torch.floor(truth_feature_box[:, 1]).int()

            truth_x = truth_feature_box[:, 0] - truth_grid_x
            truth_y = truth_feature_box[:, 1] - truth_grid_y

            target_box = pyramid_target[:, :4].clone().detach()
            target_box[:, 0] = 0
            target_box[:, 1] = 0
            normd_anch_ious = jaccard_tensor(target_box, self.normd_anchors_box)
            max_anch_ious_index = torch.argmax(normd_anch_ious, dim=-1)

            for box_i, anch_i in enumerate(max_anch_ious_index):
                if anch_i not in pyramid_anch_index_list:
                    continue
                pyramid_anch_i = anch_i % 3

                truth_w_box = torch.log(pyramid_target[box_i][2] / self.normd_anchors[anch_i][0])
                truth_h_box = torch.log(pyramid_target[box_i][3] / self.normd_anchors[anch_i][1])

                target_x[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = truth_x[box_i]
                target_y[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = truth_y[box_i]

                target_w[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = truth_w_box
                target_h[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = truth_h_box

                target_loss_weight_xw[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = \
                    pyramid_target[box_i][2]
                target_loss_weight_yh[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = \
                    pyramid_target[box_i][3]

                target_obj_conf[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = 1
                target_class_conf_list[
                    bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i], pyramid_target[box_i][4].int()] = 1

                target_obj_mask[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = 1
                target_noobj_mask[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = 0

        # print("loss in cuda") if self.cuda else print("loss not in cuda")

        if self.cuda:
            return target_x.cuda(), \
                   target_y.cuda(), \
                   target_w.cuda(), \
                   target_h.cuda(), \
                   target_loss_weight_xw.cuda(), \
                   target_loss_weight_yh.cuda(), \
                   target_obj_conf.cuda(), \
                   target_class_conf_list.cuda(), \
                   target_obj_mask.cuda(), \
                   target_noobj_mask.cuda()

        return target_x, \
               target_y, \
               target_w, \
               target_h, \
               target_loss_weight_xw, \
               target_loss_weight_yh, \
               target_obj_conf, \
               target_class_conf_list, \
               target_obj_mask, \
               target_noobj_mask

    def compute_loss(self, predict_feature: torch.Tensor, decoded_target) -> (
            torch.Tensor, torch.Tensor):
        (target_x, target_y, target_w, target_h, target_loss_weight_xw, target_loss_weight_yh, target_obj_conf,
         target_class_conf_list, target_obj_mask, target_noobj_mask) = decoded_target

        predict_feature = predict_feature.view(
            predict_feature.shape[0],
            3,
            self.bbox_attrs,
            predict_feature.shape[2],
            predict_feature.shape[3],
        ).permute(0, 1, 3, 4, 2).contiguous()

        predict_x = torch.sigmoid(predict_feature[..., 0])
        predict_y = torch.sigmoid(predict_feature[..., 1])
        predict_w = predict_feature[..., 2]
        predict_h = predict_feature[..., 3]
        predict_obj_conf = torch.sigmoid(predict_feature[..., 4])
        predict_class_conf_list = torch.sigmoid(predict_feature[..., 5:])

        target_loss_scale = 2 - target_loss_weight_xw * target_loss_weight_yh

        loss_x = torch.sum(torch.nn.BCELoss()(predict_x, target_x) * target_loss_scale * target_obj_mask)
        loss_y = torch.sum(torch.nn.BCELoss()(predict_y, target_y) * target_loss_scale * target_obj_mask)

        loss_w = torch.sum(torch.nn.MSELoss()(predict_w, target_w) * 0.5 * target_loss_scale * target_obj_mask)
        loss_h = torch.sum(torch.nn.MSELoss()(predict_h, target_h) * 0.5 * target_loss_scale * target_obj_mask)

        loss_conf = self.lambda_obj * torch.sum(
            torch.nn.BCELoss()(predict_obj_conf, target_obj_mask) * target_obj_mask) + \
                    self.lambda_noobj * torch.sum(
            torch.nn.BCELoss()(predict_obj_conf, target_obj_mask) * target_noobj_mask)

        loss_class = torch.sum(torch.nn.BCELoss()(predict_class_conf_list[target_obj_mask == 1],
                                                  target_class_conf_list[target_obj_mask == 1]))

        print("\n---------------------------------------")
        print(loss_x, loss_y)
        print(loss_w, loss_h)
        print(loss_conf, loss_class)
        print("---------------------------------------\n")

        loss = loss_x * self.lambda_xy + loss_y * self.lambda_xy + \
               loss_w * self.lambda_wh + loss_h * self.lambda_wh + \
               loss_conf * self.lambda_conf + loss_class * self.lambda_class

        return loss, torch.sum(target_obj_mask)

    def forward(self, predict_feature_list,
                tensord_target_list: List[torch.Tensor]) -> torch.Tensor:
        # pyramid_target_list_13, pyramid_target_list_26, pyramid_target_list_52 = \
        #     self.pyramid_target(tensord_target_list)
        # pyramid_normd_anchors_13, pyramid_normd_anchors_26, pyramid_normd_anchors_52 = \
        #     self.normd_anchors[0:3], self.normd_anchors[3, 6], self.normd_anchors[6, 9]

        target_13 = self.decode_pyramid_target(tensord_target_list, None, 13)
        target_26 = self.decode_pyramid_target(tensord_target_list, None, 26)
        target_52 = self.decode_pyramid_target(tensord_target_list, None, 52)

        loss_13, loss_13_num = self.compute_loss(predict_feature_list[0], target_13)
        loss_26, loss_26_num = self.compute_loss(predict_feature_list[1], target_26)
        loss_52, loss_52_num = self.compute_loss(predict_feature_list[2], target_52)

        loss_list = []

        if not torch.isnan(loss_13):
            loss_list.append(loss_13)
        if not torch.isnan(loss_26):
            loss_list.append(loss_26)
        if not torch.isnan(loss_52):
            loss_list.append(loss_52)

        assert len(loss_list) != 0

        loss = sum(loss_list)

        loss_num = loss_13_num + loss_26_num + loss_52_num

        return loss / loss_num


import numpy as np
import scipy.signal as signal
from scipy.interpolate import interp1d
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from astropy.io import fits
import emcee
import corner
import copy
import time
import os
import sys
import smart

def makeModel(teff, logg=5, metal=0, vsini=1, rv=0, tell_alpha=1.0, airmass=1.0, pwv=0.5, wave_offset=0, flux_offset=0,**kwargs):
	"""
	Return a forward model.

	Parameters
	----------
	teff   : effective temperature

	data   : an input science data used for continuum correction

	Optional Parameters
	-------------------

	Returns
	-------
	model: a synthesized model
	"""

	# read in the parameters
	order        = kwargs.get('order', '33')
	modelset     = kwargs.get('modelset', 'btsettl08')
	instrument   = kwargs.get('instrument', 'nirspec')
	veiling      = kwargs.get('veiling', 0)    # flux veiling parameter
	lsf          = kwargs.get('lsf', 4.5)   # instrumental LSF
	include_fringe_model = kwargs.get('include_fringe_model', False)

	if instrument == 'apogee':
		try:
			import apogee_tools as ap
		except ImportError:
			print('Need to install the package "apogee_tools" (https://github.com/jbirky/apogee_tools) \n')
		xlsf       = kwargs.get('xlsf', np.linspace(-7.,7.,43))   # APOGEE instrumental LSF sampling
		wave_off1  = kwargs.get('wave_off1') # wavelength offset for chip a
		wave_off2  = kwargs.get('wave_off2') # wavelength offset for chip b
		wave_off3  = kwargs.get('wave_off3') # wavelength offset for chip c
		c0_1       = kwargs.get('c0_1')      # constant flux offset for chip a
		c0_2       = kwargs.get('c0_2')      # linear flux offset for chip a
		c1_1       = kwargs.get('c1_1')      # constant flux offset for chip b
		c1_2       = kwargs.get('c1_2')      # linear flux offset for chip b
		c2_1       = kwargs.get('c2_1')      # constant flux offset for chip c
		c2_2       = kwargs.get('c2_2')      # linear flux offset for chip c

	tell       = kwargs.get('tell', True) # apply telluric
	#tell_alpha = kwargs.get('tell_alpha', 1.0) # Telluric alpha power
	binary     = kwargs.get('binary', False) # make a binary model

	# assume the secondary has the same metallicity
	if binary:
		teff2       = kwargs.get('teff2')
		logg2       = kwargs.get('logg2')
		rv2         = kwargs.get('rv2')
		vsini2      = kwargs.get('vsini2')
		flux_scale = kwargs.get('flux_scale', 0.8)

	data       = kwargs.get('data', None) # for continuum correction and resampling

	output_stellar_model = kwargs.get('output_stellar_model', False)

	if data is not None and instrument == 'nirspec':
		order = data.order
		# read in a model
		#print('teff ',teff,'logg ',logg, 'z', z, 'order', order, 'modelset', modelset)
		#print('teff ',type(teff),'logg ',type(logg), 'z', type(z), 'order', type(order), 'modelset', type(modelset))
		model    = smart.Model(teff=teff, logg=logg, metal=metal, order=str(order), modelset=modelset, instrument=instrument)

	#elif data is not None and instrument == 'apogee':
	elif instrument == 'apogee':
		model    = smart.Model(teff=teff, logg=logg, metal=metal, modelset=modelset, instrument=instrument)
		# Dirty fix here
		model.wave = model.wave[np.where(model.flux != 0)]
		model.flux = model.flux[np.where(model.flux != 0)]

		# apply vmicro
		vmicro = 2.478 - 0.325*logg
		model.flux = smart.broaden(wave=model.wave, flux=model.flux, vbroad=vmicro, rotate=False, gaussian=True)

	elif data is None and instrument == 'nirspec':
		model    = smart.Model(teff=teff, logg=logg, metal=metal, order=str(order), modelset=modelset, instrument=instrument)

	# wavelength offset
	#model.wave += wave_offset

	# apply vsini
	model.flux = smart.broaden(wave=model.wave, flux=model.flux, vbroad=vsini, rotate=True, gaussian=False)

	# apply rv (including the barycentric correction)
	model.wave = rvShift(model.wave, rv=rv)

	# flux veiling
	model.flux += veiling

	## if binary is True: make a binary model
	if binary:
		model2      = smart.Model(teff=teff2, logg=logg2, metal=metal, order=str(order), modelset=modelset, instrument=instrument)
		# apply vsini
		model2.flux = smart.broaden(wave=model2.wave, flux=model2.flux, vbroad=vsini2, rotate=True, gaussian=False)
		# apply rv (including the barycentric correction)
		model2.wave = rvShift(model2.wave, rv=rv2)
		# linearly interpolate the model2 onto the model1 grid
		fit = interp1d(model2.wave, model2.flux)

		select_wavelength = np.where( (model.wave < model2.wave[-1]) & (model.wave > model2.wave[0]) )
		model.flux = model.flux[select_wavelength]
		model.wave = model.wave[select_wavelength]

		# combine the models together and scale the secondary flux
		model.flux += flux_scale * fit(model.wave)

	if output_stellar_model:
		stellar_model = copy.deepcopy(model)
		if binary:
			model2.flux = flux_scale * fit(model.wave)

	# apply telluric
	if tell is True:
		model = smart.applyTelluric(model=model, tell_alpha=tell_alpha, airmass=airmass, pwv=pwv)

	# fringe
	if include_fringe_model is True:
		#print('adding the fringe model')
		s1, s2, s3, s4, s5 = 0, 150, 400, 600, -1
		piecewise_fringe_model = [s1, s2, s3, s4, s5]
		model.flux *= smart.double_sine_fringe(model, data, piecewise_fringe_model, teff, logg, vsini, rv, airmass, pwv, wave_offset, flux_offset, lsf, modelset)

	# instrumental LSF
	if instrument == 'nirspec':
		model.flux = smart.broaden(wave=model.wave, flux=model.flux, vbroad=lsf, rotate=False, gaussian=True)
	elif instrument == 'apogee':
		model.flux = ap.apogee_hack.spec.lsf.convolve(model.wave, model.flux, lsf=lsf, xlsf=xlsf).flatten()
		model.wave = ap.apogee_hack.spec.lsf.apStarWavegrid()
		# Remove the NANs
		model.wave = model.wave[~np.isnan(model.flux)]
		model.flux = model.flux[~np.isnan(model.flux)]

	if output_stellar_model:
		stellar_model.flux = smart.broaden(wave=stellar_model.wave, flux=stellar_model.flux, vbroad=lsf, rotate=False, gaussian=True)
		if binary:
			model2.flux = smart.broaden(wave=model2.wave, flux=model2.flux, vbroad=lsf, rotate=False, gaussian=True)

	# wavelength offset
	model.wave += wave_offset

	if output_stellar_model:
		stellar_model.wave += wave_offset
		if binary:
			model2.wave = stellar_model.wave

	# integral resampling
	if data is not None:
		if instrument == 'nirspec':
			model.flux = np.array(smart.integralResample(xh=model.wave, yh=model.flux, xl=data.wave))
			model.wave = data.wave

			if output_stellar_model:
				stellar_model.flux = np.array(smart.integralResample(xh=stellar_model.wave, yh=stellar_model.flux, xl=data.wave))
				stellar_model.wave = data.wave
				if binary:
					model2.flux = np.array(smart.integralResample(xh=model2.wave, yh=model2.flux, xl=data.wave))
					model2.wave = data.wave

		# contunuum correction
		if data.instrument == 'nirspec':
			niter = 5 # continuum iteration
			if output_stellar_model:
				model, cont_factor = smart.continuum(data=data, mdl=model, prop=True)
				for i in range(niter):
					model, cont_factor2 = smart.continuum(data=data, mdl=model, prop=True)
					cont_factor *= cont_factor2
				stellar_model.flux *= cont_factor
				if binary:
					model2.flux *= cont_factor
			else:
				model = smart.continuum(data=data, mdl=model)
				for i in range(niter):
					model = smart.continuum(data=data, mdl=model)
		elif data.instrument == 'apogee':
			## set the order in the continuum fit
			deg         = 5
			## because of the APOGEE bands, continuum is corrected from three pieces of the spectra
			data0       = copy.deepcopy(data)
			model0      = copy.deepcopy(model)

			# wavelength offset
			model0.wave += wave_off1

			range0      = np.where((data0.wave >= data.oriWave0[0][-1]) & (data0.wave <= data.oriWave0[0][0]))
			data0.wave  = data0.wave[range0]
			data0.flux  = data0.flux[range0]
			if data0.wave[0] > data0.wave[-1]:
				data0.wave = data0.wave[::-1]
				data0.flux = data0.flux[::-1]
			model0.flux = np.array(smart.integralResample(xh=model0.wave, yh=model0.flux, xl=data0.wave))
			model0.wave = data0.wave
			model0      = smart.continuum(data=data0, mdl=model0, deg=deg)
			# flux corrections
			model0.flux = (model0.flux + c0_1) * np.e**(-c0_2)

			data1       = copy.deepcopy(data)
			model1      = copy.deepcopy(model)

			# wavelength offset
			model1.wave += wave_off2

			range1      = np.where((data1.wave >= data.oriWave0[1][-1]) & (data1.wave <= data.oriWave0[1][0]))
			data1.wave  = data1.wave[range1]
			data1.flux  = data1.flux[range1]
			if data1.wave[0] > data1.wave[-1]:
				data1.wave = data1.wave[::-1]
				data1.flux = data1.flux[::-1]
			model1.flux = np.array(smart.integralResample(xh=model1.wave, yh=model1.flux, xl=data1.wave))
			model1.wave = data1.wave
			model1      = smart.continuum(data=data1, mdl=model1, deg=deg)

			# flux corrections
			model1.flux = (model1.flux + c1_1) * np.e**(-c1_2)

			data2       = copy.deepcopy(data)
			model2      = copy.deepcopy(model)

			# wavelength offset
			model2.wave += wave_off3

			range2      = np.where((data2.wave >= data.oriWave0[2][-1]) & (data2.wave <= data.oriWave0[2][0]))
			data2.wave  = data2.wave[range2]
			data2.flux  = data2.flux[range2]
			if data2.wave[0] > data2.wave[-1]:
				data2.wave = data2.wave[::-1]
				data2.flux = data2.flux[::-1]

			model2.flux = np.array(smart.integralResample(xh=model2.wave, yh=model2.flux, xl=data2.wave))
			model2.wave = data2.wave
			model2      = smart.continuum(data=data2, mdl=model2, deg=deg)
			# flux corrections
			model2.flux = (model2.flux + c2_1) * np.e**(-c2_2)

			## scale the flux to be the same as the data
			#model0.flux *= (np.std(data0.flux)/np.std(model0.flux))
			#model0.flux -= np.median(model0.flux) - np.median(data0.flux)

			#model1.flux *= (np.std(data1.flux)/np.std(model1.flux))
			#model1.flux -= np.median(model1.flux) - np.median(data1.flux)

			#model2.flux *= (np.std(data2.flux)/np.std(model2.flux))
			#model2.flux -= np.median(model2.flux) - np.median(data2.flux)

			model.flux  = np.array( list(model2.flux) + list(model1.flux) + list(model0.flux) )
			model.wave  = np.array( list(model2.wave) + list(model1.wave) + list(model0.wave) )

	if instrument == 'nirspec':
		# flux offset
		model.flux += flux_offset
		if output_stellar_model:
			stellar_model.flux += flux_offset
			if binary:
				model2.flux += flux_offset
	#model.flux **= (1 + flux_exponent_offset)

	if output_stellar_model:
		if not binary:
			return model, stellar_model
		else:
			return model, stellar_model, model2
	else:
		return model

# Fringe Model; testing

def doub_sine(wave, a1, k1, a2, k2):
    # the initial guess is determined from the best frequency
    # wave (i.e. kx) multiplicative effects
    return (1 + a1**2 + 2 * a1*np.sin( k1 * wave )) * ( 1 + a2**2 + 2 * a2 * np.sin( k2*wave ))

def makeModelFringe(teff, logg=5, metal=0, vsini=1, rv=0, tell_alpha=1.0, airmass=1.0, pwv=0.5, wave_offset=0, flux_offset=0,
	a1_1=0.01, k1_1=2.10, a2_1=0.01, k2_1=0.85, a1_2=0.01, k1_2=2.10, a2_2=0.01, k2_2=0.85,
	a1_3=0.01, k1_3=2.10, a2_3=0.01, k2_3=0.85, a1_4=0.01, k1_4=2.10, a2_4=0.01, k2_4=0.85, **kwargs):
	"""
	Return a forward model.

	Parameters
	----------
	teff   : effective temperature

	data   : an input science data used for continuum correction

	Optional Parameters
	-------------------

	Returns
	-------
	model: a synthesized model
	"""

	# read in the parameters
	order        = kwargs.get('order', '33')
	modelset     = kwargs.get('modelset', 'btsettl08')
	instrument   = kwargs.get('instrument', 'nirspec')
	veiling      = kwargs.get('veiling', 0)    # flux veiling parameter
	lsf          = kwargs.get('lsf', 4.5)   # instrumental LSF

	if instrument == 'apogee':
		try:
			import apogee_tools as ap
		except ImportError:
			print('Need to install the package "apogee_tools" (https://github.com/jbirky/apogee_tools) \n')
		xlsf       = kwargs.get('xlsf', np.linspace(-7.,7.,43))   # APOGEE instrumental LSF sampling
		wave_off1  = kwargs.get('wave_off1') # wavelength offset for chip a
		wave_off2  = kwargs.get('wave_off2') # wavelength offset for chip b
		wave_off3  = kwargs.get('wave_off3') # wavelength offset for chip c
		c0_1       = kwargs.get('c0_1')      # constant flux offset for chip a
		c0_2       = kwargs.get('c0_2')      # linear flux offset for chip a
		c1_1       = kwargs.get('c1_1')      # constant flux offset for chip b
		c1_2       = kwargs.get('c1_2')      # linear flux offset for chip b
		c2_1       = kwargs.get('c2_1')      # constant flux offset for chip c
		c2_2       = kwargs.get('c2_2')      # linear flux offset for chip c

	tell       = kwargs.get('tell', True) # apply telluric
	#tell_alpha = kwargs.get('tell_alpha', 1.0) # Telluric alpha power
	binary     = kwargs.get('binary', False) # make a binary model

	# assume the secondary has the same metallicity
	if binary:
		teff2       = kwargs.get('teff2')
		logg2       = kwargs.get('logg2')
		rv2         = kwargs.get('rv2')
		vsini2      = kwargs.get('vsini2')
		flux_scale = kwargs.get('flux_scale', 0.8)

	data       = kwargs.get('data', None) # for continuum correction and resampling

	output_stellar_model = kwargs.get('output_stellar_model', False)

	if data is not None and instrument == 'nirspec':
		order = data.order
		# read in a model
		#print('teff ',teff,'logg ',logg, 'z', z, 'order', order, 'modelset', modelset)
		#print('teff ',type(teff),'logg ',type(logg), 'z', type(z), 'order', type(order), 'modelset', type(modelset))
		model    = smart.Model(teff=teff, logg=logg, metal=metal, order=str(order), modelset=modelset, instrument=instrument)

	#elif data is not None and instrument == 'apogee':
	elif instrument == 'apogee':
		model    = smart.Model(teff=teff, logg=logg, metal=metal, modelset=modelset, instrument=instrument)
		# Dirty fix here
		model.wave = model.wave[np.where(model.flux != 0)]
		model.flux = model.flux[np.where(model.flux != 0)]

		# apply vmicro
		vmicro = 2.478 - 0.325*logg
		model.flux = smart.broaden(wave=model.wave, flux=model.flux, vbroad=vmicro, rotate=False, gaussian=True)

	elif data is None and instrument == 'nirspec':
		model    = smart.Model(teff=teff, logg=logg, metal=metal, order=str(order), modelset=modelset, instrument=instrument)

	# wavelength offset
	#model.wave += wave_offset

	# apply vsini
	model.flux = smart.broaden(wave=model.wave, flux=model.flux, vbroad=vsini, rotate=True, gaussian=False)

	# apply rv (including the barycentric correction)
	model.wave = rvShift(model.wave, rv=rv)

	# flux veiling
	model.flux += veiling

	## if binary is True: make a binary model
	if binary:
		model2      = smart.Model(teff=teff2, logg=logg2, metal=metal, order=str(order), modelset=modelset, instrument=instrument)
		# apply vsini
		model2.flux = smart.broaden(wave=model2.wave, flux=model2.flux, vbroad=vsini2, rotate=True, gaussian=False)
		# apply rv (including the barycentric correction)
		model2.wave = rvShift(model2.wave, rv=rv2)
		# linearly interpolate the model2 onto the model1 grid
		fit = interp1d(model2.wave, model2.flux)

		select_wavelength = np.where( (model.wave < model2.wave[-1]) & (model.wave > model2.wave[0]) )
		model.flux = model.flux[select_wavelength]
		model.wave = model.wave[select_wavelength]

		# combine the models together and scale the secondary flux
		model.flux += flux_scale * fit(model.wave)

	if output_stellar_model:
		stellar_model = copy.deepcopy(model)
		if binary:
			model2.flux = flux_scale * fit(model.wave)

	# apply telluric
	if tell is True:
		model = smart.applyTelluric(model=model, tell_alpha=tell_alpha, airmass=airmass, pwv=pwv)

	# instrumental LSF
	if instrument == 'nirspec':
		model.flux = smart.broaden(wave=model.wave, flux=model.flux, vbroad=lsf, rotate=False, gaussian=True)
	elif instrument == 'apogee':
		model.flux = ap.apogee_hack.spec.lsf.convolve(model.wave, model.flux, lsf=lsf, xlsf=xlsf).flatten()
		model.wave = ap.apogee_hack.spec.lsf.apStarWavegrid()
		# Remove the NANs
		model.wave = model.wave[~np.isnan(model.flux)]
		model.flux = model.flux[~np.isnan(model.flux)]

	if output_stellar_model:
		stellar_model.flux = smart.broaden(wave=stellar_model.wave, flux=stellar_model.flux, vbroad=lsf, rotate=False, gaussian=True)
		if binary:
			model2.flux = smart.broaden(wave=model2.wave, flux=model2.flux, vbroad=lsf, rotate=False, gaussian=True)

	# wavelength offset
	model.wave += wave_offset

	if output_stellar_model:
		stellar_model.wave += wave_offset
		if binary:
			model2.wave = stellar_model.wave

	# integral resampling
	if data is not None:
		if instrument == 'nirspec':
			model.flux = np.array(smart.integralResample(xh=model.wave, yh=model.flux, xl=data.wave))
			model.wave = data.wave

			if output_stellar_model:
				stellar_model.flux = np.array(smart.integralResample(xh=stellar_model.wave, yh=stellar_model.flux, xl=data.wave))
				stellar_model.wave = data.wave
				if binary:
					model2.flux = np.array(smart.integralResample(xh=model2.wave, yh=model2.flux, xl=data.wave))
					model2.wave = data.wave

		# fringe modeling
		if instrument == 'nirspec':
			# Define the four piece-wise fringe model to include the k(wavelength) variation.
			s1, s2, s3, s4, s5 = 0, 200, 400, 600, -1

			model.flux[s1:s2] = model.flux[s1:s2] * doub_sine(wave=model.wave[s1:s2], a1=a1_1, k1=k1_1, a2=a2_1, k2=k2_1)
			model.flux[s2:s3] = model.flux[s2:s3] * doub_sine(wave=model.wave[s2:s3], a1=a1_2, k1=k1_2, a2=a2_2, k2=k2_2)
			model.flux[s3:s4] = model.flux[s3:s4] * doub_sine(wave=model.wave[s3:s4], a1=a1_3, k1=k1_3, a2=a2_3, k2=k2_3)
			model.flux[s4:s5] = model.flux[s4:s5] * doub_sine(wave=model.wave[s4:s5], a1=a1_4, k1=k1_4, a2=a2_4, k2=k2_4)

			if output_stellar_model:
				stellar_model.flux[s1:s2] = stellar_model.flux[s1:s2] * doub_sine(wave=model.wave[s1:s2], a1=a1_1, k1=k1_1, a2=a2_1, k2=k2_1)
				stellar_model.flux[s2:s3] = stellar_model.flux[s2:s3] * doub_sine(wave=model.wave[s2:s3], a1=a1_2, k1=k1_2, a2=a2_2, k2=k2_2)
				stellar_model.flux[s3:s4] = stellar_model.flux[s3:s4] * doub_sine(wave=model.wave[s3:s4], a1=a1_3, k1=k1_3, a2=a2_3, k2=k2_3)
				stellar_model.flux[s4:s5] = stellar_model.flux[s4:s5] * doub_sine(wave=model.wave[s4:s5], a1=a1_4, k1=k1_4, a2=a2_4, k2=k2_4)

		# contunuum correction
		if data.instrument == 'nirspec':
			niter = 5 # continuum iteration
			if output_stellar_model:
				model, cont_factor = smart.continuum(data=data, mdl=model, prop=True)
				for i in range(niter):
					model, cont_factor2 = smart.continuum(data=data, mdl=model, prop=True)
					cont_factor *= cont_factor2
				stellar_model.flux *= cont_factor
				if binary:
					model2.flux *= cont_factor
			else:
				model = smart.continuum(data=data, mdl=model)
				for i in range(niter):
					model = smart.continuum(data=data, mdl=model)

		elif data.instrument == 'apogee':
			## set the order in the continuum fit
			deg         = 5
			## because of the APOGEE bands, continuum is corrected from three pieces of the spectra
			data0       = copy.deepcopy(data)
			model0      = copy.deepcopy(model)

			# wavelength offset
			model0.wave += wave_off1

			range0      = np.where((data0.wave >= data.oriWave0[0][-1]) & (data0.wave <= data.oriWave0[0][0]))
			data0.wave  = data0.wave[range0]
			data0.flux  = data0.flux[range0]
			if data0.wave[0] > data0.wave[-1]:
				data0.wave = data0.wave[::-1]
				data0.flux = data0.flux[::-1]
			model0.flux = np.array(smart.integralResample(xh=model0.wave, yh=model0.flux, xl=data0.wave))
			model0.wave = data0.wave
			model0      = smart.continuum(data=data0, mdl=model0, deg=deg)
			# flux corrections
			model0.flux = (model0.flux + c0_1) * np.e**(-c0_2)

			data1       = copy.deepcopy(data)
			model1      = copy.deepcopy(model)

			# wavelength offset
			model1.wave += wave_off2

			range1      = np.where((data1.wave >= data.oriWave0[1][-1]) & (data1.wave <= data.oriWave0[1][0]))
			data1.wave  = data1.wave[range1]
			data1.flux  = data1.flux[range1]
			if data1.wave[0] > data1.wave[-1]:
				data1.wave = data1.wave[::-1]
				data1.flux = data1.flux[::-1]
			model1.flux = np.array(smart.integralResample(xh=model1.wave, yh=model1.flux, xl=data1.wave))
			model1.wave = data1.wave
			model1      = smart.continuum(data=data1, mdl=model1, deg=deg)

			# flux corrections
			model1.flux = (model1.flux + c1_1) * np.e**(-c1_2)

			data2       = copy.deepcopy(data)
			model2      = copy.deepcopy(model)

			# wavelength offset
			model2.wave += wave_off3

			range2      = np.where((data2.wave >= data.oriWave0[2][-1]) & (data2.wave <= data.oriWave0[2][0]))
			data2.wave  = data2.wave[range2]
			data2.flux  = data2.flux[range2]
			if data2.wave[0] > data2.wave[-1]:
				data2.wave = data2.wave[::-1]
				data2.flux = data2.flux[::-1]

			model2.flux = np.array(smart.integralResample(xh=model2.wave, yh=model2.flux, xl=data2.wave))
			model2.wave = data2.wave
			model2      = smart.continuum(data=data2, mdl=model2, deg=deg)
			# flux corrections
			model2.flux = (model2.flux + c2_1) * np.e**(-c2_2)

			## scale the flux to be the same as the data
			#model0.flux *= (np.std(data0.flux)/np.std(model0.flux))
			#model0.flux -= np.median(model0.flux) - np.median(data0.flux)

			#model1.flux *= (np.std(data1.flux)/np.std(model1.flux))
			#model1.flux -= np.median(model1.flux) - np.median(data1.flux)

			#model2.flux *= (np.std(data2.flux)/np.std(model2.flux))
			#model2.flux -= np.median(model2.flux) - np.median(data2.flux)

			model.flux  = np.array( list(model2.flux) + list(model1.flux) + list(model0.flux) )
			model.wave  = np.array( list(model2.wave) + list(model1.wave) + list(model0.wave) )

	if instrument == 'nirspec':
		# flux offset
		model.flux += flux_offset
		if output_stellar_model:
			stellar_model.flux += flux_offset
			if binary:
				model2.flux += flux_offset
	#model.flux **= (1 + flux_exponent_offset)

	if output_stellar_model:
		if not binary:
			return model, stellar_model
		else:
			return model, stellar_model, model2
	else:
		return model

# Fringe Model; testing

def rvShift(wavelength, rv):
	"""
	Perform the radial velocity correction.

	Parameters
	----------
	wavelength 	: 	numpy array
					model wavelength (in Angstroms)

	rv 			: 	float
					radial velocity shift (in km/s)

	Returns
	-------
	wavelength 	: 	numpy array
					shifted model wavelength (in Angstroms)
	"""
	return wavelength * ( 1 + rv / 299792.458)

def applyTelluric(model, tell_alpha=1.0, airmass=1.5, pwv=0.5):
	"""
	Apply the telluric model on the science model.

	Parameters
	----------
	model 	:	model object
				BT Settl model
	alpha 	: 	float
				telluric scaling factor (the power on the flux)

	Returns
	-------
	model 	: 	model object
				BT Settl model times the corresponding model

	"""
	# read in a telluric model
	wavelow  = model.wave[0] - 10
	wavehigh = model.wave[-1] + 10
	#telluric_model = smart.getTelluric(wavelow=wavelow, wavehigh=wavehigh, alpha=alpha, airmass=airmass)

	telluric_model = smart.Model()
	telluric_model.wave, telluric_model.flux = 	smart.InterpTelluricModel(wavelow=wavelow, wavehigh=wavehigh, airmass=airmass, pwv=pwv)

	# apply the telluric alpha parameter
	telluric_model.flux = telluric_model.flux**(tell_alpha)

	#if len(model.wave) > len(telluric_model.wave):
	#	print("The model has a higher resolution ({}) than the telluric model ({})."\
	#		.format(len(model.wave),len(telluric_model.wave)))
	#	model.flux = np.array(smart.integralResample(xh=model.wave,
	#		yh=model.flux, xl=telluric_model.wave))
	#	model.wave = telluric_model.wave
	#	model.flux *= telluric_model.flux

	#elif len(model.wave) < len(telluric_model.wave):
	## This should be always true
	telluric_model.flux = np.array(smart.integralResample(xh=telluric_model.wave, yh=telluric_model.flux, xl=model.wave))
	telluric_model.wave = model.wave
	model.flux *= telluric_model.flux

	#elif len(model.wave) == len(telluric_model.wave):
	#	model.flux *= telluric_model.flux

	return model

def convolveTelluric(lsf, telluric_data, alpha=1.0, airmass='1.0', pwv='1.5'):
	"""
	Return a convolved telluric transmission model given a telluric data and lsf.
	"""
	# get a telluric standard model
	wavelow               = telluric_data.wave[0]  - 50
	wavehigh              = telluric_data.wave[-1] + 50
	telluric_model        = smart.getTelluric(wavelow=wavelow,wavehigh=wavehigh, airmass=airmass, pwv=pwv)
	telluric_model.flux **= alpha
	# lsf
	telluric_model.flux = smart.broaden(wave=telluric_model.wave, flux=telluric_model.flux,
		vbroad=lsf, rotate=False, gaussian=True)
	# resample
	telluric_model.flux = np.array(smart.integralResample(xh=telluric_model.wave,
		yh=telluric_model.flux, xl=telluric_data.wave))
	telluric_model.wave = telluric_data.wave
	return telluric_model

def getLSF2(telluric_data, continuum=True, test=False, save_path=None):
	"""
	Return a best LSF value from a telluric data.
	"""

	data = copy.deepcopy(telluric_data)

	def bestParams(data, i, alpha, c2, c0):

		data2          = copy.deepcopy(data)
		data2.wave     = data2.wave + c0
		telluric_model = smart.convolveTelluric(i, data2, alpha=alpha)
		model          = smart.continuum(data=data2, mdl=telluric_model)
		#plt.figure(2)
		#plt.plot(model.wave, model.flux+c2, 'r-', alpha=0.5)
		#plt.plot(data.wave*c1+c0, data.flux, 'b-', alpha=0.5)
		#plt.close()
		#plt.show()
		#sys.exit()
		return model.flux + c2

	def bestParams2(theta, data):

		i, alpha, c2, c0, c1 = theta
		data2                = copy.deepcopy(data)
		data2.wave           = data2.wave*c1 + c0
		telluric_model       = smart.convolveTelluric(i, data2, alpha=alpha)
		model                = smart.continuum(data=data2, mdl=telluric_model)
		return np.sum(data.flux - (model.flux + c2))**2

	from scipy.optimize import curve_fit, minimize

	popt, pcov = curve_fit(bestParams, data, data.flux, p0=[4.01, 1.01, 0.01, 1.01], maxfev=1000000, epsfcn=0.1)

	#nll = lambda *args: bestParams2(*args)
	#results = minimize(nll, [3., 1., 0.1, -10., 1.], args=(data))
	#popt = results['x']

	data.wave      = data.wave+popt[3]

	telluric_model = smart.convolveTelluric(popt[0], data, alpha=popt[1])
	model          = smart.continuum(data=data, mdl=telluric_model)

	#model.flux * np.e**(-popt[2]) + popt[3]
	model.flux + popt[2]

	return popt[0]

def getLSF(telluric_data, alpha=1.0, continuum=True,test=False,save_path=None):
	"""
	Return a best LSF value from a telluric data.
	"""
	lsf_list = []
	test_lsf = np.arange(3.0,13.0,0.1)

	data = copy.deepcopy(telluric_data)
	if continuum is True:
		data = smart.continuumTelluric(data=data)

	data.flux **= alpha
	for i in test_lsf:
		telluric_model = smart.convolveTelluric(i,data)
		if telluric_data.order == 59:
			telluric_model.flux **= 3
			# mask hydrogen absorption feature
			data2          = copy.deepcopy(data)
			tell_mdl       = copy.deepcopy(telluric_model)
			mask_pixel     = 450
			data2.wave     = data2.wave[mask_pixel:]
			data2.flux     = data2.flux[mask_pixel:]
			data2.noise    = data2.noise[mask_pixel:]
			tell_mdl.wave  = tell_mdl.wave[mask_pixel:]
			tell_mdl.flux  = tell_mdl.flux[mask_pixel:]

			chisquare = smart.chisquare(data2,tell_mdl)

		else:
			chisquare = smart.chisquare(data,telluric_model)
		lsf_list.append([chisquare,i])

		if test is True:
			plt.plot(telluric_model.wave,telluric_model.flux+(i-3)*10+1,
				'r-',alpha=0.5)

	if test is True:
		plt.plot(data.wave,data.flux,
			'k-',label='telluric data',alpha=0.5)
		plt.title("Test LSF",fontsize=15)
		plt.xlabel("Wavelength ($\AA$)",fontsize=12)
		plt.ylabel("Transmission + Offset",fontsize=12)
		plt.minorticks_on()
		if save_path is not None:
			plt.savefig(save_path+\
				"/{}_O{}_lsf_data_mdl.png"\
				.format(data.name, data.order))
		#plt.show()
		plt.close()

		fig, ax = plt.subplots()
		for i in range(len(lsf_list)):
			ax.plot(lsf_list[i][1],lsf_list[i][0],'k.',alpha=0.5)
		ax.plot(min(lsf_list)[1],min(lsf_list)[0],'r.',
			label="best LSF {} km/s".format(min(lsf_list)[1]))
		ax.set_xlabel("LSF (km/s)",fontsize=12)
		ax.set_ylabel("$\chi^2$",fontsize=11)
		plt.minorticks_on()
		plt.legend(fontsize=10)
		if save_path is not None:
			plt.savefig(save_path+\
				"/{}_O{}_lsf_chi2.png"\
				.format(data.name, data.order))
		#plt.show()
		plt.close()

	lsf = min(lsf_list)[1]

	if telluric_data.order == 61 or telluric_data.order == 62 \
	or telluric_data.order == 63: #or telluric_data.order == 64:
		lsf = 5.5
		print("The LSF is obtained from orders 60 and 65 (5.5 km/s).")

	return lsf

def getAlpha(telluric_data,lsf,continuum=True,test=False,save_path=None):
	"""
	Return a best alpha value from a telluric data.
	"""
	alpha_list = []
	test_alpha = np.arange(0.1,7,0.1)

	data = copy.deepcopy(telluric_data)
	if continuum is True:
		data = smart.continuumTelluric(data=data)

	for i in test_alpha:
		telluric_model = smart.convolveTelluric(lsf,data,
			alpha=i)
		#telluric_model.flux **= i
		if data.order == 59:
			# mask hydrogen absorption feature
			data2          = copy.deepcopy(data)
			tell_mdl       = copy.deepcopy(telluric_model)
			mask_pixel     = 450
			data2.wave     = data2.wave[mask_pixel:]
			data2.flux     = data2.flux[mask_pixel:]
			data2.noise    = data2.noise[mask_pixel:]
			tell_mdl.wave  = tell_mdl.wave[mask_pixel:]
			tell_mdl.flux  = tell_mdl.flux[mask_pixel:]

			chisquare = smart.chisquare(data2,tell_mdl)

		else:
			chisquare = smart.chisquare(data,telluric_model)
		alpha_list.append([chisquare,i])

		if test is True:
			plt.plot(telluric_model.wave,telluric_model.flux+i*10,
				'k-',alpha=0.5)

	if test is True:
		plt.plot(telluric_data.wave,telluric_data.flux,
			'r-',alpha=0.5)
		plt.rc('font', family='sans-serif')
		plt.title("Test Alpha",fontsize=15)
		plt.xlabel("Wavelength ($\AA$)",fontsize=12)
		plt.ylabel("Transmission + Offset",fontsize=12)
		plt.minorticks_on()
		if save_path is not None:
			plt.savefig(save_path+\
				"/{}_O{}_alpha_data_mdl.png"\
				.format(telluric_data.name,
					telluric_data.order))
		plt.show()
		plt.close()

		fig, ax = plt.subplots()
		plt.rc('font', family='sans-serif')
		for i in range(len(alpha_list)):
			ax.plot(alpha_list[i][1],alpha_list[i][0],'k.',alpha=0.5)
		ax.plot(min(alpha_list)[1],min(alpha_list)[0],'r.',
			label="best alpha {}".format(min(alpha_list)[1]))
		ax.set_xlabel(r"$\alpha$",fontsize=12)
		ax.set_ylabel("$\chi^2$",fontsize=12)
		plt.minorticks_on()
		plt.legend(fontsize=10)
		if save_path is not None:
			plt.savefig(save_path+\
				"/{}_O{}_alpha_chi2.png"\
				.format(telluric_data.name,
					telluric_data.order))
		plt.show()
		plt.close()

	alpha = min(alpha_list)[1]

	return alpha

def getFringeFrequecy(tell_data, test=False):
	"""
	Use the Lomb-Scargle Periodogram to identify
	the fringe pattern.
	"""
	tell_sp  = copy.deepcopy(tell_data)

	## continuum correction
	tell_sp  = smart.continuumTelluric(data=tell_sp, order=tell_sp.order)

	## get a telluric model
	lsf      = smart.getLSF(tell_sp)
	alpha    = smart.getAlpha(tell_sp,lsf)
	tell_mdl = smart.convolveTelluric(lsf=lsf,
		telluric_data=tell_sp,alpha=alpha)

	## fit the fringe pattern in the residual
	pgram_x = np.array(tell_sp.wave,float)[10:-10]
	pgram_y = np.array(tell_sp.flux - tell_mdl.flux,float)[10:-10]
	offset  = np.mean(pgram_y)
	pgram_y -= offset
	mask    = np.where(pgram_y - 1.5 * np.absolute(np.std(pgram_y)) > 0)
	pgram_x = np.delete(pgram_x, mask)
	pgram_y = np.delete(pgram_y, mask)
	pgram_x = np.array(pgram_x, float)
	pgram_y = np.array(pgram_y, float)

	#f = np.lismartace(0.01,10,100000)
	f = np.lismartace(1.0,10,100000)

	## Lomb Scargle Periodogram
	pgram = signal.lombscargle(pgram_x, pgram_y, f)

	if test:
		fig, ax = plt.subplots(figsize=(16,6))
		ax.plot(f,pgram, 'k-', label='residual',alpha=0.5)
		ax.set_xlabel('frequency')
		plt.legend()
		plt.show()
		plt.close()

	return f[np.argmax(pgram)]

def initModelFit(sci_data, lsf, modelset='btsettl08'):
	"""
	Conduct simple chisquare fit to obtain the initial parameters
	for the forward modeling MCMC.

	The function would calculate the chisquare for teff, logg, vini, rv, and alpha.

	Parameters
	----------
	data 				:	spectrum object
							input science data

	lsf 				:	float
							line spread function for the NIRSPEC

	Returns
	-------
	best_params_dic 	:	dic
							a dictionary that stores the best parameters for
							teff, logg, vsini, rv, and alpha

	chisquare 			:	int
							minimum chisquare

	"""
	data            = copy.deepcopy(sci_data)

	## set up the parameter grid for chisquare computation
	teff_array      = np.arange(1200,3001,100)
	logg_array      = np.arange(3.5,5.51,0.5)
	vsini_array     = np.arange(10,101,10)
	rv_array        = np.arange(-200,201,50)
	alpha_array     = np.arange(0.5,2.01,0.5)
	chisquare_array = np.empty(len(teff_array)*len(logg_array)*len(vsini_array)*len(rv_array)*len(alpha_array))\
	.reshape(len(teff_array),len(logg_array),len(vsini_array),len(rv_array),len(alpha_array))

	time1 = time.time()
	for i, teff in enumerate(teff_array):
		for j, logg in enumerate(logg_array):
			for k, vsini in enumerate(vsini_array):
				for l, rv in enumerate(rv_array):
					for m, alpha in enumerate(alpha_array):
						model = smart.makeModel(teff, logg, 0.0, vsini, rv, alpha, 0, 0,
							lsf=lsf, order=str(data.order), data=data, modelset=modelset)
						chisquare_array[i,j,k,l,m] = smart.chisquare(data, model)
	time2 = time.time()
	print("total time:",time2-time1)

	ind = np.unravel_index(np.argmin(chisquare_array, axis=None), chisquare_array.shape)
	print("ind ",ind)
	chisquare       = chisquare_array[ind]

	best_params_dic = {'teff':teff_array[ind[0]], 'logg':logg_array[ind[1]],
	'vsini':vsini_array[ind[2]], 'rv':rv_array[ind[3]], 'alpha':alpha_array[ind[4]]}

	print(best_params_dic, chisquare)

	return best_params_dic , chisquare

import numpy,import scipy,from scipy,from astropy


# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
dust --- Models for dust
========================

.. autosummary::
   :toctree: generated/

   Activity
   --------
   acrit

   Dust Models
   -----------
   AfrhoRadiation
   AfrhoScattered
   AfrhoThermal

   Phase functions
   ---------------
   phaseK
   phaseH
   phaseHM

"""

import numpy as np
import astropy.units as u
from astropy.units import Quantity

__all__ = [
    'acrit',

    'AfrhoRadiation',
    'AfrhoScattered',
    'AfrhoThermal',

    'phaseK',
    'phaseH',
    'phaseHM'
]

def acrit(Q, vth, R, rho_g=Quantity(1, u.g / u.cm**3),
          rho_n=Quantity(0.3, u.g / u.cm**3), f_active=1.0):
    """Maximum liftable grain radius from a spherical nucleus.

    See Meech and Svoren 2004, Comets II.

    Parameters
    ----------
    Q : Quantity
      Mass production rate of the driving gas.
    vth : Quantity
      Gas expansion speed.
    R : Quantity
      Nucleus radius.
    rho_g : Quantity, optional
      Grain density.
    rho_n : Quantity, optional
      Nucleus mean density.
    f_active : float, optional
      Active fraction of the nucleus.

    Returns
    -------
    a : Quantity
      Maximum liftable grain radius.

    """

    from numpy import pi
    import astropy.constants as c

    a = 9 * Q * vth / (64 * pi**2 * rho_g * rho_n * R**3 * c.G)
    return a.to(u.um)

class AfrhoRadiation(object):
    """Light from a comet coma parameterized by Afrho, or similar.

    Methods
    -------
    fluxd : Total flux density from the object.

    Notes
    -----
    Afrho should refer to the value at zero phase angle.

    Inheriting classes should override `fluxd`, and `__init__`
    functions.  `__init__` should take a single argument, `Afrho` (or
    equivalent) as a Quantity.

    As much as possible, keyword arguments must have the same meaning
    in all derived models.

    """

    def __init__(self, Afrho, **kwargs):
        pass

    def __call__(self, *args, **kwargs):
        return self.fluxd(*args, **kwargs)

    def fluxd(self, geom, wave, unit=None):
        pass

class AfrhoScattered(AfrhoRadiation):
    """Scattered light from a comet coma parameterized by Afrho.

    If you use this model, please reference A'Hearn et al. (1984, AJ
    89, 579-591) as the source of the Afrho parameter.

    Parameters
    ----------
    Afrho : Quantity
      The product of albedo at zero phase, A, dust filling factor, f,
      and observer's aperture radius, rho.
    phasef : function, optional
      The phase function of the coma.  Set to `None` to use `phaseK`.

    Methods
    -------
    fluxd : Total flux density from the coma.

    """

    def __init__(self, Afrho, phasef=None, **kwargs):
        self.Afrho = Afrho
        if phasef is None:
            self.phasef = phaseK
        else:
            self.phasef = phasef

    def fluxd(self, geom, wave, rap, unit=u.Unit('W / (m2 um)')):
        """Flux density.

        Parameters
        ----------
        geom : dict of Quantities
          A dictionary-like object with the keys 'rh' (heliocentric
          distance), 'delta' (observer-target distance), and 'phase'
          (phase angle).
        wave : Quantity
          The wavelengths at which to compute the emission.
        rap : Quantity
          The aperture radius, angular or projected distance at the
          comet.
        unit : astropy Units, optional
          The return units.  Must be spectral flux density.

        Returns
        -------
        fluxd : Quantityu
          The flux density from the coma.

        Raises
        ------
        ValueError : If `rap` has incorrect units.

        """

        from ..calib import solar_flux

        if rap.unit.is_equivalent(u.cm):
            rho = rap.to(self.Afrho.unit)
        elif rap.unit.is_equivalent(u.arcsec):
            rho = geom['delta'].to(self.Afrho.unit) * rap.to(u.rad).value
        else:
            raise ValueError("rap must have angular or length units.")

        fsun = solar_flux(wave, unit=unit) / geom['rh'].to(u.au).value**2
        fluxd = (self.Afrho
                 * self.phasef(np.abs(geom['phase'].to(u.deg).value))
                 * rho * fsun / 4.0 / geom['delta'].to(self.Afrho.unit)**2)

        return fluxd

class AfrhoThermal(AfrhoRadiation):
    """Thermal emisson from a coma parameterized by efrho.

    If you use this model, please cite and reference Kelley et
    al. (2013, Icarus 225, 475-494).  They define `epsilon-f-rho` as
    the product of IR emissivity (`epsilon`), dust filling factor
    (`f`), and observer's aperture radius (`rho`).

    The default `ef2af` is 3.5, which assumes `epsilion` is
    approximately 0.9, `A` is approximately 0.25, and the scattering
    and emission filling factors are the same.  This value can roughly
    reproduce the spectral shape of 73P-C/Schwassmann-Wachmann in
    Fig. 16 of Sitko et al. (2011, AJ 142, 80) for `Tscale = 1.12`.

    The default long-wavelength slope, `beta = 0.89+/-0.10`, is from
    an analysis of Hyakutake JCMT data by Jewitt and Matthews (1997,
    AJ 113, 1145).  The break-point, `wave0` = 70 um, is based on my
    own analysis, combining the Jewitt and Matthews fluxes with mid-IR
    fluxes from Mason et al. (1998, ApJ 507, 398).

    Parameters
    ----------
    Afrho : Quantity
      The product of albedo at zero phase, A, dust filling factor, f,
      and observer's aperture radius, rho.
    ef2af : float, optional
      The ratio of epsilon-f_therm to A-f_sca, where f_therm and f_sca
      are the effective thermal and scattered light filling factors,
      (they are not necessarily the same).
    Tscale : float, optional
      The isothermal blackbody sphere temperature scale factor that
      characterizes the spectral shape of the thermal emission.
    beta : float, optional
    wave0 : Quantity, optional
      Scale wavelengths longer than `wave0` by `(wave / wave0)**-beta`.

    Methods
    -------
    fluxd : Total flux density from the coma.

    """

    def __init__(self, Afrho, ef2af=3.5, Tscale=1.1, beta=0.89,
                 wave0=70 * u.um, **kwargs):
        assert isinstance(Afrho, u.Quantity)
        self.Afrho = Afrho
        self.ef2af = ef2af
        self.Tscale = Tscale
        self.beta = beta
        self.wave0 = wave0

    def fluxd(self, geom, wave, rap, unit=u.Unit('W / (m2 um)')):
        """Flux density.

        Parameters
        ----------
        geom : dict of Quantities
          A dictionary-like object with the keys 'rh' (heliocentric
          distance), 'delta' (observer-target distance), and 'phase'
          (phase angle).
        wave : Quantity
          The wavelengths at which to compute the emission.
        rap : Quantity
          The aperture radius, angular or projected distance at the
          comet.
        unit : astropy Units, optional
          The return units.  Must be spectral flux density.

        Returns
        -------
        fluxd : Quantity
          The flux density from the coma.

        Raises
        ------
        ValueError : If `rap` has incorrect units.

        """

        from ..util import planck

        if rap.unit.is_equivalent(u.cm):
            rho = rap.to(self.Afrho.unit)
        elif rap.unit.is_equivalent(u.arcsec):
            rho = geom['delta'].to(self.Afrho.unit) * rap.to(u.rad).value
        else:
            raise ValueError("rap must have angular or length units.")

        T = self.Tscale * 278 / np.sqrt(geom['rh'].to(u.au).value)
        B = planck(wave, T, unit=unit / u.sr).value
        efrho = self.Afrho * self.ef2af
        d = geom['delta'].to(self.Afrho.unit).value
        fluxd = efrho.value * np.pi * B * rho.value / d**2

        if any(wave > self.wave0):
            eps = np.ones(len(wave))
            i = wave > self.wave0
            eps[i] *= (wave[i] / self.wave0)**-self.beta
            fluxd *= eps

        return fluxd * unit

def phaseK(phase):
    """Phase function derived from Kolokolova et al. (2004, Comets II).

    The phase function of K04 is scaled to phasef(0) = 1.0.

    Parameters
    ----------
    phase : float or array
        Phase angle. [degrees]

    Returns
    -------
    phi : float or ndarray
        The phase function.

    Notes
    -----
    To estimate the phase function, I fit a polynomial function to the
    solid line of Kolokolova et al. (2004, Comets II):

      a = array([0.27, 0.21, 0.17, 0.15, 0.14, 0.135, 0.135, 0.135, 0.15,
           0.175, 0.225, 0.3, 0.43, 0.62, 0.775])
      b = array([0.0,  10,   20,   30,   40,   60,    70,   80,   100,
           110,   120,   130, 140,  150,  156])
      fit = poly1d(polyfit(b, a / 0.27, 4))
      plot(b, a, 'o')
      plot(b, fit(b) * min(a), 'r-')

    """

    phasef = np.poly1d([  3.14105489e-08,  -7.84714255e-06,   7.34255521e-04,
                         -3.09608957e-02,   1.00920684e+00])
    return phasef(np.abs(phase))

def phaseH(phase):
    """Halley phase function from Schleicher et al. (1998).

    The Halley phase function is from Schleicher et al. (1998, Icarus
    132, 397-417).  The Comet Halley observations were at phases less
    than 70 degrees.

    Parameters
    ----------
    phase : float or array
        Phase angle. [degrees]

    Returns
    -------
    phi : float or ndarray
        The phase function.

    """

    phasef = np.poly1d([0.000177, -0.01807, 1])
    return phasef(np.abs(phase))

def phaseHM(phase):
    """Halley-Marcus phase function from Schleicher et al. (2011).

    The Halley phase function is first published in Schleicher and
    Bair (2011, AJ 141, 117), but only described in detail by
    Schleicher and Marcus (May 2010) at:

      http://asteroid.lowell.edu/comet/dustphase.html

      "To distinguish this curve from others, we designate this as the
      HM phase function, for the sources of the two components: Halley
      and Marcus, where the Halley curve for smaller phase angles
      comes from our previous work (Schleicher et al. 1998) while Joe
      Marcus has fit a Henyey-Greenstein function to a variety of mid-
      and large-phase angle data sets (Marcus 2007); see here for
      details. Note that we do not consider our composite curve to be
      a definitive result, but rather appropriate for performing
      first-order adjustments to dust measurements for changing phase
      angle."

    Parameters
    ----------
    phase : float or array
        Phase angle. [degrees]

    Returns
    -------
    phi : float or ndarray
        The phase function.

    """

    from scipy.interpolate import splrep, splev

    th = np.arange(181)
    ph = np.array(
        [  1.0000e+00,   9.5960e-01,   9.2170e-01,   8.8590e-01,
           8.5220e-01,   8.2050e-01,   7.9060e-01,   7.6240e-01,
           7.3580e-01,   7.1070e-01,   6.8710e-01,   6.6470e-01,
           6.4360e-01,   6.2370e-01,   6.0490e-01,   5.8720e-01,
           5.7040e-01,   5.5460e-01,   5.3960e-01,   5.2550e-01,
           5.1220e-01,   4.9960e-01,   4.8770e-01,   4.7650e-01,
           4.6590e-01,   4.5590e-01,   4.4650e-01,   4.3770e-01,
           4.2930e-01,   4.2150e-01,   4.1420e-01,   4.0730e-01,
           4.0090e-01,   3.9490e-01,   3.8930e-01,   3.8400e-01,
           3.7920e-01,   3.7470e-01,   3.7060e-01,   3.6680e-01,
           3.6340e-01,   3.6030e-01,   3.5750e-01,   3.5400e-01,
           3.5090e-01,   3.4820e-01,   3.4580e-01,   3.4380e-01,
           3.4210e-01,   3.4070e-01,   3.3970e-01,   3.3890e-01,
           3.3850e-01,   3.3830e-01,   3.3850e-01,   3.3890e-01,
           3.3960e-01,   3.4050e-01,   3.4180e-01,   3.4320e-01,
           3.4500e-01,   3.4700e-01,   3.4930e-01,   3.5180e-01,
           3.5460e-01,   3.5760e-01,   3.6090e-01,   3.6450e-01,
           3.6830e-01,   3.7240e-01,   3.7680e-01,   3.8150e-01,
           3.8650e-01,   3.9170e-01,   3.9730e-01,   4.0320e-01,
           4.0940e-01,   4.1590e-01,   4.2280e-01,   4.3000e-01,
           4.3760e-01,   4.4560e-01,   4.5400e-01,   4.6270e-01,
           4.7200e-01,   4.8160e-01,   4.9180e-01,   5.0240e-01,
           5.1360e-01,   5.2530e-01,   5.3750e-01,   5.5040e-01,
           5.6380e-01,   5.7800e-01,   5.9280e-01,   6.0840e-01,
           6.2470e-01,   6.4190e-01,   6.5990e-01,   6.7880e-01,
           6.9870e-01,   7.1960e-01,   7.4160e-01,   7.6480e-01,
           7.8920e-01,   8.1490e-01,   8.4200e-01,   8.7060e-01,
           9.0080e-01,   9.3270e-01,   9.6640e-01,   1.0021e+00,
           1.0399e+00,   1.0799e+00,   1.1223e+00,   1.1673e+00,
           1.2151e+00,   1.2659e+00,   1.3200e+00,   1.3776e+00,
           1.4389e+00,   1.5045e+00,   1.5744e+00,   1.6493e+00,
           1.7294e+00,   1.8153e+00,   1.9075e+00,   2.0066e+00,
           2.1132e+00,   2.2281e+00,   2.3521e+00,   2.4861e+00,
           2.6312e+00,   2.7884e+00,   2.9592e+00,   3.1450e+00,
           3.3474e+00,   3.5685e+00,   3.8104e+00,   4.0755e+00,
           4.3669e+00,   4.6877e+00,   5.0418e+00,   5.4336e+00,
           5.8682e+00,   6.3518e+00,   6.8912e+00,   7.4948e+00,
           8.1724e+00,   8.9355e+00,   9.7981e+00,   1.0777e+01,
           1.1891e+01,   1.3166e+01,   1.4631e+01,   1.6322e+01,
           1.8283e+01,   2.0570e+01,   2.3252e+01,   2.6418e+01,
           3.0177e+01,   3.4672e+01,   4.0086e+01,   4.6659e+01,
           5.4704e+01,   6.4637e+01,   7.7015e+01,   9.2587e+01,
           1.1237e+02,   1.3775e+02,   1.7060e+02,   2.1348e+02,
           2.6973e+02,   3.4359e+02,   4.3989e+02,   5.6292e+02,
           7.1363e+02,   8.8448e+02,   1.0533e+03,   1.1822e+03,
           1.2312e+03])

    C = splrep(th, ph)
    return splev(np.abs(phase), C)

# update module docstring
from ..util import autodoc
autodoc(globals())
del autodoc

import numpy,from numpy,from scipy,import astropy,from astropy


from __future__ import annotations
from typing import Optional, Tuple, Union, List, Dict
import torch
import numpy as np
import scipy.special
from dqc.hamilton.base_hamilton import BaseHamilton
from dqc.hamilton.hcgto_pbc import HamiltonCGTO_PBC
from dqc.system.base_system import BaseSystem
from dqc.grid.base_grid import BaseGrid
from dqc.grid.factory import get_predefined_grid
from dqc.system.mol import _parse_basis, _get_nelecs_spin, \
                           _get_orb_weights, AtomZsType, AtomPosType
from dqc.utils.datastruct import CGTOBasis, AtomCGTOBasis, ZType, BasisInpType, \
                                 SpinParam, DensityFitInfo
from dqc.utils.safeops import safe_cdist
from dqc.utils.periodictable import get_atom_mass
from dqc.api.parser import parse_moldesc
from dqc.hamilton.intor.lattice import Lattice
from dqc.hamilton.intor.pbcintor import PBCIntOption
from dqc.utils.cache import Cache

__all__ = ["Sol"]

class Sol(BaseSystem):
    """
    Describe the system of a solid (i.e. periodic boundary condition system).

    Arguments
    ---------
    * soldesc: str or 2-elements tuple
        Description of the molecule system.
        If string, it can be described like ``"H 1 0 0; H -1 0 0"``.
        If tuple, the first element of the tuple is the Z number of the atoms while
        the second element is the position of the atoms: ``(atomzs, atomposs)``.
    * basis: str, CGTOBasis, list of str, or CGTOBasis
        The string describing the gto basis. If it is a list, then it must have
        the same length as the number of atoms.
    * grid: int
        Describe the grid.
        If it is an integer, then it uses the default grid with specified level
        of accuracy.
    * spin: int, float, torch.Tensor, or None
        The difference between spin-up and spin-down electrons.
        It must be an integer or ``None``.
        If ``None``, then it is ``num_electrons % 2``.
        For floating point atomzs and/or charge, the ``spin`` must be specified.
    * charge: int, float, or torch.Tensor
        The charge of the molecule.
    * orb_weights: SpinParam[torch.Tensor] or None
        Specifiying the orbital occupancy (or weights) directly. If specified,
        ``spin`` and ``charge`` arguments are ignored.
    * dtype: torch.dtype
        The data type of tensors in this class.
    * device: torch.device
        The device on which the tensors in this class are stored.
    """

    def __init__(self,
                 soldesc: Union[str, Tuple[AtomZsType, AtomPosType]],
                 alattice: torch.Tensor,
                 basis: Union[str, List[CGTOBasis], List[str], List[List[CGTOBasis]]],
                 *,
                 grid: Union[int, str] = "sg3",
                 spin: Optional[ZType] = None,
                 lattsum_opt: Optional[Union[PBCIntOption, Dict]] = None,
                 dtype: torch.dtype = torch.float64,
                 device: torch.device = torch.device('cpu'),
                 ):
        self._dtype = dtype
        self._device = device
        self._grid_inp = grid
        self._basis_inp = basis
        self._grid: Optional[BaseGrid] = None
        charge = 0  # we can't have charged solids for now

        # get the AtomCGTOBasis & the hamiltonian
        # atomzs: (natoms,) dtype: torch.int or dtype for floating point
        # atompos: (natoms, ndim)
        atomzs, atompos = parse_moldesc(soldesc, dtype, device)
        allbases = _parse_basis(atomzs, basis)  # list of list of CGTOBasis
        atombases = [AtomCGTOBasis(atomz=atz, bases=bas, pos=atpos)
                     for (atz, bas, atpos) in zip(atomzs, allbases, atompos)]
        self._atombases = atombases
        self._atompos = atompos  # (natoms, ndim)
        self._atomzs = atomzs  # (natoms,) int-type
        nelecs_tot: torch.Tensor = torch.sum(atomzs)

        # get the number of electrons and spin and orbital weights
        nelecs, spin, frac_mode = _get_nelecs_spin(nelecs_tot, spin, charge)
        assert not frac_mode, "Fractional Z mode for pbc is not supported"
        _orb_weights, _orb_weights_u, _orb_weights_d = _get_orb_weights(
            nelecs, spin, frac_mode, dtype, device)

        # initialize cache
        self._cache = Cache()

        # save the system's properties
        self._spin = spin
        self._charge = charge
        self._numel = nelecs
        self._orb_weights = _orb_weights
        self._orb_weights_u = _orb_weights_u
        self._orb_weights_d = _orb_weights_d
        self._alattice_inp = alattice
        self._lattice = Lattice(self._alattice_inp)
        self._lattsum_opt = PBCIntOption.get_default(lattsum_opt)

    def densityfit(self, method: Optional[str] = None,
                   auxbasis: Optional[BasisInpType] = None) -> BaseSystem:
        """
        Indicate that the system's Hamiltonian uses density fit for its integral.

        Arguments
        ---------
        method: Optional[str]
            Density fitting method. Available methods in this class are:

            * ``"gdf"``: Density fit with gdf compensating charge to perform
                the lattice sum. Ref https://doi.org/10.1063/1.4998644 (default)

        auxbasis: Optional[BasisInpType]
            Auxiliary basis for the density fit. If not specified, then it uses
            ``"cc-pvtz-jkfit"``.
        """
        if method is None:
            method = "gdf"
        if auxbasis is None:
            # TODO: choose the auxbasis properly
            auxbasis = "cc-pvtz-jkfit"

        # get the auxiliary basis
        assert auxbasis is not None
        auxbasis_lst = _parse_basis(self._atomzs, auxbasis)
        atomauxbases = [AtomCGTOBasis(atomz=atz, bases=bas, pos=atpos)
                        for (atz, bas, atpos) in zip(self._atomzs, auxbasis_lst, self._atompos)]

        # change the hamiltonian to have density fit
        df = DensityFitInfo(method=method, auxbases=atomauxbases)
        self._hamilton = HamiltonCGTO_PBC(self._atombases, df=df, latt=self._lattice,
                                          lattsum_opt=self._lattsum_opt,
                                          cache=self._cache.add_prefix("hamilton"))
        return self

    def get_hamiltonian(self) -> BaseHamilton:
        """
        Returns the Hamiltonian that corresponds to the system, i.e.
        :class:`~dqc.hamilton.HamiltonCGTO_PBC`
        """
        return self._hamilton

    def set_cache(self, fname: str, paramnames: Optional[List[str]] = None) -> BaseSystem:
        """
        Setup the cache of some parameters specified by `paramnames` to be read/written
        on a file.
        If the file exists, then the parameters will not be recomputed, but just
        loaded from the cache instead.
        Cache is usually used for repeated calculations where the cached parameters
        are not changed (e.g. running multiple systems with slightly different environment.)

        Arguments
        ---------
        fname: str
            The file to store the cache.
        paramnames: list of str or None
            List of parameter names to be read/write from the cache.
        """
        self._cache.set(fname, paramnames)
        return self

    def get_orbweight(self, polarized: bool = False) -> Union[torch.Tensor, SpinParam[torch.Tensor]]:
        if not polarized:
            return self._orb_weights
        else:
            return SpinParam(u=self._orb_weights_u, d=self._orb_weights_d)

    def get_nuclei_energy(self) -> torch.Tensor:
        # self._atomzs: (natoms,)
        # self._atompos: (natoms, ndim)

        # r12: (natoms, natoms)
        r12_inf = safe_cdist(self._atompos, self._atompos, add_diag_eps=True, diag_inf=True)
        r12 = safe_cdist(self._atompos, self._atompos, add_diag_eps=True)
        z12 = self._atomzs.unsqueeze(-2) * self._atomzs.unsqueeze(-1)  # (natoms, natoms)

        precision = self._lattsum_opt.precision
        eta = self._lattice.estimate_ewald_eta(precision) * 2
        vol = self._lattice.volume()
        rcut = scipy.special.erfcinv(float(vol.detach()) * eta * eta / (2 * np.pi) * precision) / eta
        gcut = scipy.special.erfcinv(precision * np.sqrt(np.pi) / 2 / eta) * 2 * eta

        # get the shift vector in real space and in reciprocal space
        ls = self._lattice.get_lattice_ls(rcut=rcut, exclude_zeros=True)  # (nls, ndim)
        # gv: (ngv, ndim), gvweights: (ngv,)
        gv, gvweights = self._lattice.get_gvgrids(gcut=gcut, exclude_zeros=True)
        gv_norm2 = torch.einsum("gd,gd->g", gv, gv)  # (ngv)

        # get the shift in position
        atpos_shift = self._atompos - ls.unsqueeze(-2)  # (nls, natoms, ndim)
        r12_ls = safe_cdist(atpos_shift, self._atompos)  # (nls, natoms, natoms)

        # calculate the short range
        short_range_comp1 = torch.erfc(eta * r12_ls) / r12_ls  # (nls, natoms, natoms)
        short_range_comp2 = torch.erfc(eta * r12) / r12_inf  # (natoms, natoms)
        short_range1 = torch.sum(z12 * short_range_comp1)  # scalar
        short_range2 = torch.sum(z12 * short_range_comp2)  # scalar
        short_range = short_range1 + short_range2

        # calculate the long range sum
        coul_g = 4 * np.pi / gv_norm2 * gvweights  # (ngv,)
        # this part below is quicker, but raises warning from pytorch
        si = torch.exp(1j * torch.matmul(self._atompos, gv.transpose(-2, -1)))  # (natoms, ngv)
        zsi = torch.einsum("a,ag->g", self._atomzs.to(si.dtype), si)  # (ngv,)
        zexpg2 = zsi * torch.exp(-gv_norm2 / (4 * eta * eta))
        long_range = torch.einsum("a,a,a->", zsi.conj(), zexpg2, coul_g.to(zsi.dtype)).real  # (scalar)

        # # alternative way to compute the long-range part
        # r12_pair = self._atompos.unsqueeze(-2) - self._atompos  # (natoms, natoms, ndim)
        # long_range_exp = torch.exp(-gv_norm2 / (4 * eta * eta)) * coul_g  # (ngv,)
        # long_range_cos = torch.cos(torch.einsum("gd,abd->gab", gv, -r12_pair))  # (ngv, natoms, natoms)
        # long_range = torch.sum(long_range_exp[:, None, None] * long_range_cos * z12)  # scalar

        # background interaction
        vbar1 = -torch.sum(self._atomzs ** 2) * (2 * eta / np.sqrt(np.pi))
        vbar2 = -torch.sum(self._atomzs) ** 2 * np.pi / (eta * eta * vol)
        vbar = vbar1 + vbar2  # (scalar)

        eii = short_range + long_range + vbar
        return eii * 0.5

    def setup_grid(self) -> None:
        self._grid = get_predefined_grid(self._grid_inp, self._atomzs, self._atompos,
                                         lattice=self._lattice,
                                         dtype=self._dtype, device=self._device)

    def get_grid(self) -> BaseGrid:
        if self._grid is None:
            raise RuntimeError("Please run mol.setup_grid() first before calling get_grid()")
        return self._grid

    def requires_grid(self) -> bool:
        return False

    def getparamnames(self, methodname: str, prefix: str = "") -> List[str]:
        pass

    def make_copy(self, **kwargs) -> Sol:
        """
        Returns a copy of the system identical to the orginal except for new
        parameters set in the kwargs.

        Arguments
        ---------
        **kwargs
            Must be the same kwargs as Sol.
        """
        # create dictionary of all parameters
        parameters = {
            'soldesc': (self.atomzs, self.atompos),
            'alattice': self._alattice_inp,
            'basis': self._basis_inp,
            'grid': self._grid_inp,
            'spin': self._spin,
            'lattsum_opt': self._lattsum_opt,
            'dtype': self._dtype,
            'device': self._device
        }
        # update dictionary with provided kwargs
        parameters.update(kwargs)
        # create new system
        return Sol(**parameters)

    ################### properties ###################
    @property
    def atompos(self) -> torch.Tensor:
        return self._atompos

    @property
    def atomzs(self) -> torch.Tensor:
        return self._atomzs

    @property
    def atommasses(self) -> torch.Tensor:
        # returns the atomic mass (only for non-isotope for now)
        return torch.tensor([get_atom_mass(int(atomz)) for atomz in self._atomzs],
                            dtype=self._dtype, device=self._device)

    @property
    def spin(self) -> ZType:
        return self._spin

    @property
    def charge(self) -> ZType:
        return self._charge

    @property
    def numel(self) -> ZType:
        return self._numel

    @property
    def efield(self) -> Optional[Tuple[torch.Tensor, ...]]:
        # solid with external efield has not been implemented
        return None