# Licensed under a 3-clause BSD style license - see LICENSE.rst
# This old version has been retrieved to use the Spectrum class to do photometry
# Functionalities have been removed

import abc
import math
from copy import deepcopy
from collections import OrderedDict

import numpy as np

from astropy.utils import lazyproperty
from astropy.io import ascii
import astropy.units as u
import astropy.constants as const
from astropy import cosmology

# from ._registry import Registry

__all__ = [
    "get_bandpass",
    "get_magsystem",
    "read_bandpass",
    "Bandpass",
    "Spectrum",
    "MagSystem",
    "SpectralMagSystem",
    "ABMagSystem",
    "CompositeMagSystem",
]

HC_ERG_AA = const.h.cgs.value * const.c.to(u.AA / u.s).value

# _BANDPASSES = Registry()
# _MAGSYSTEMS = Registry()

def get_bandpass(name):
    """Get a Bandpass from the registry by name."""
    if isinstance(name, Bandpass):
        return name
    return _BANDPASSES.retrieve(name)

def get_magsystem(name):
    """Get a MagSystem from the registry by name."""
    if isinstance(name, MagSystem):
        return name
    return _MAGSYSTEMS.retrieve(name)

def read_bandpass(
    fname, fmt="ascii", wave_unit=u.AA, trans_unit=u.dimensionless_unscaled, name=None
):
    """Read bandpass from two-column ASCII file containing wavelength and
    transmission in each line.

    Parameters
    ----------
    fname : str
        File name.
    fmt : {'ascii'}
        File format of file. Currently only ASCII file supported.
    wave_unit : `~astropy.units.Unit` or str, optional
        Wavelength unit. Default is Angstroms.
    trans_unit : `~astropy.units.Unit`, optional
        Transmission unit. Can be `~astropy.units.dimensionless_unscaled`,
        indicating a ratio of transmitted to incident photons, or units
        proportional to inverse energy, indicating a ratio of transmitted
        photons to incident energy. Default is ratio of transmitted to
        incident photons.
    name : str, optional
        Identifier. Default is `None`.

    Returns
    -------
    band : `~sncosmo.Bandpass`
    """

    if fmt != "ascii":
        raise ValueError(
            "format {0} not supported. Supported formats: 'ascii'".format(fmt)
        )
    t = ascii.read(fname, names=["wave", "trans"])
    return Bandpass(
        t["wave"], t["trans"], wave_unit=wave_unit, trans_unit=trans_unit, name=name
    )

class Bandpass(object):
    """Transmission as a function of spectral wavelength.

    Parameters
    ----------
    wave : list_like
        Wavelength. Monotonically increasing values.
    trans : list_like
        Transmission fraction.
    wave_unit : `~astropy.units.Unit` or str, optional
        Wavelength unit. Default is Angstroms.
    trans_unit : `~astropy.units.Unit`, optional
        Transmission unit. Can be `~astropy.units.dimensionless_unscaled`,
        indicating a ratio of transmitted to incident photons, or units
        proportional to inverse energy, indicating a ratio of transmitted
        photons to incident energy. Default is ratio of transmitted to
        incident photons.
    name : str, optional
        Identifier. Default is `None`.

    Examples
    --------
    >>> b = Bandpass([4000., 4200., 4400.], [0.5, 1.0, 0.5])
    >>> b.wave
    array([ 4000.,  4200.,  4400.])
    >>> b.trans
    array([ 0.5,  1. ,  0.5])
    >>> b.dwave
    array([ 200.,  200.,  200.])
    >>> b.wave_eff
    4200.0

    """

    def __init__(
        self,
        wave,
        trans,
        wave_unit=u.AA,
        trans_unit=u.dimensionless_unscaled,
        name=None,
    ):
        wave = np.asarray(wave, dtype=np.float64)
        trans = np.asarray(trans, dtype=np.float64)
        if wave.shape != trans.shape:
            raise ValueError("shape of wave and trans must match")
        if wave.ndim != 1:
            raise ValueError("only 1-d arrays supported")

        # Ensure that units are actually units and not quantities, so that
        # `to` method returns a float and not a Quantity.
        wave_unit = u.Unit(wave_unit)
        trans_unit = u.Unit(trans_unit)

        if wave_unit != u.AA:
            wave = wave_unit.to(u.AA, wave, u.spectral())

        # If transmission is in units of inverse energy, convert to
        # unitless transmission:
        #
        # (transmitted photons / incident photons) =
        #      (photon energy) * (transmitted photons / incident energy)
        #
        # where photon energy = h * c / lambda
        if trans_unit != u.dimensionless_unscaled:
            trans = (HC_ERG_AA / wave) * trans_unit.to(u.erg ** -1, trans)

        # Check that values are monotonically increasing.
        # We could sort them, but if this happens, it is more likely a user
        # error or faulty bandpass definition. So we leave it to the user to
        # sort them.
        if not np.all(np.ediff1d(wave) > 0.0):
            raise ValueError(
                "bandpass wavelength values must be monotonically"
                " increasing when supplied in wavelength or "
                "decreasing when supplied in energy/frequency."
            )
        self.wave = wave
        self._dwave = np.gradient(wave)
        self.trans = trans
        self.name = name

    @property
    def dwave(self):
        """Gradient of wavelengths, numpy.gradient(wave)."""
        return self._dwave

    @lazyproperty
    def wave_eff(self):
        """Effective wavelength of bandpass in Angstroms."""
        weights = self.trans * np.gradient(self.wave)
        return np.sum(self.wave * weights) / np.sum(weights)

    def to_unit(self, unit):
        """Return wavelength and transmission in new wavelength units.

        If the requested units are the same as the current units, self is
        returned.

        Parameters
        ----------
        unit : `~astropy.units.Unit` or str
            Target wavelength unit.

        Returns
        -------
        wave : `~numpy.ndarray`
        trans : `~numpy.ndarray`
        """

        if unit is u.AA:
            return self.wave, self.trans

        d = u.AA.to(unit, self.wave, u.spectral())
        t = self.trans
        if d[0] > d[-1]:
            d = np.flipud(d)
            t = np.flipud(t)
        return d, t

    def __repr__(self):
        name = ""
        if self.name is not None:
            name = " {0!r:s}".format(self.name)
        return "<Bandpass{0:s} at 0x{1:x}>".format(name, id(self))

class Spectrum(object):
    """A spectrum, representing wavelength and spectral density values.

    Parameters
    ----------
    wave : list_like
        Wavelength values.
    flux : list_like
        Spectral flux density values.
    error : list_like, optional
        1 standard deviation uncertainty on flux density values.
    wave_unit : `~astropy.units.Unit`
        Units.
    unit : `~astropy.units.BaseUnit`
        For now, only units with flux density in energy (not photon counts).
    z : float, optional
        Redshift of spectrum (default is `None`)
    dist : float, optional
        Luminosity distance in Mpc, used to adjust flux upon redshifting.
        The default is ``None``.
    meta : OrderedDict, optional
        Metadata.
    """

    def __init__(
        self,
        wave,
        flux,
        error=None,
        unit=(u.erg / u.s / u.cm ** 2 / u.AA),
        wave_unit=u.AA,
        z=None,
        dist=None,
        meta=None,
    ):
        self._wave = np.asarray(wave)
        self._flux = np.asarray(flux)
        self._wunit = wave_unit
        self._unit = unit
        self._z = z
        self._dist = dist

        if error is not None:
            self._error = np.asarray(error)
            if self._wave.shape != self._error.shape:
                raise ValueError("shape of wavelength and variance must match")
        else:
            self._error = None

        if meta is None:
            self.meta = OrderedDict()
        else:
            self.meta = deepcopy(meta)

        if self._wave.shape != self._flux.shape:
            raise ValueError("shape of wavelength and flux must match")
        if self._wave.ndim != 1:
            raise ValueError("only 1-d arrays supported")

    @property
    def wave(self):
        """Wavelength values."""
        return self._wave

    @property
    def flux(self):
        """Spectral flux density values"""
        return self._flux

    @property
    def error(self):
        """Uncertainty on flux density."""
        return self._error

    @property
    def wave_unit(self):
        """Units of wavelength."""
        return self._wunit

    @property
    def unit(self):
        """Units of flux density."""
        return self._unit

    @property
    def z(self):
        """Redshift of spectrum."""
        return self._z

    @z.setter
    def z(self, value):
        self._z = value

    @property
    def dist(self):
        """Distance to object in Mpc."""
        return self._dist

    @dist.setter
    def dist(self, value):
        self._dist = value

    def bandflux(self, band):
        """Perform synthentic photometry in a given bandpass.

        The bandpass transmission is interpolated onto the wavelength grid
        of the spectrum. The result is a weighted sum of the spectral flux
        density values (weighted by transmission values).

        Parameters
        ----------
        band : Bandpass object or name of registered bandpass.

        Returns
        -------
        bandflux : float
            Total flux in ph/s/cm^2. If part of bandpass falls
            outside the spectrum, `None` is returned instead.
        bandfluxerr : float
            Error on flux. Only returned if the `error` attribute is not
            `None`.
        """

        band = get_bandpass(band)
        bwave, btrans = band.to_unit(self._wunit)

        if bwave[0] < self._wave[0] or bwave[-1] > self._wave[-1]:
            return None

        mask = (self._wave > bwave[0]) & (self._wave < bwave[-1])
        d = self._wave[mask]
        f = self._flux[mask]

        # First convert to ergs/s/cm^2/(wavelength unit)...
        target_unit = u.erg / u.s / u.cm ** 2 / self._wunit
        if self._unit != target_unit:
            f = self._unit.to(target_unit, f, u.spectral_density(self._wunit, d))

        # Then convert ergs to photons: photons = Energy / (h * nu).
        f = f / const.h.cgs.value / self._wunit.to(u.Hz, d, u.spectral())

        trans = np.interp(d, bwave, btrans)
        binw = np.gradient(d)
        ftot = np.sum(f * trans * binw)

        if self._error is None:
            return ftot

        else:
            e = self._error[mask]

            # Do the same conversion as above
            if self._unit != target_unit:
                e = self._unit.to(target_unit, e, u.spectral_density(self._wunit, d))
            e = e / const.h.cgs.value / self._wunit.to(u.Hz, d, u.spectral())
            etot = np.sqrt(np.sum((e * binw) ** 2 * trans))
            return ftot, etot

    def redshifted_to(self, z, adjust_flux=False, dist=None, cosmo=None):
        """Return a new Spectrum object at a new redshift.

        The current redshift must be defined (self.z cannot be `None`).
        A factor of (1 + z) / (1 + self.z) is applied to the wavelength.
        The inverse factor is applied to the flux so that the bolometric
        flux (e.g., erg/s/cm^2) remains constant.

        .. note:: Currently this only works for units in ergs.

        Parameters
        ----------
        z : float
            Target redshift.
        adjust_flux : bool, optional
            If True, the bolometric flux is adjusted by
            ``F_out = F_in * (D_in / D_out) ** 2``, where ``D_in`` and
            ``D_out`` are current and target luminosity distances,
            respectively. ``D_in`` is self.dist. If self.dist is ``None``,
            the distance is calculated from the current redshift and
            given cosmology.
        dist : float, optional
            Output distance in Mpc. Used to adjust bolometric flux if
            ``adjust_flux`` is ``True``. Default is ``None`` which means
            that the distance is calculated from the redshift and the
            cosmology.
        cosmo : `~astropy.cosmology.Cosmology` instance, optional
            The cosmology used to estimate distances if dist is not given.
            Default is ``None``, which results in using the default
            cosmology.

        Returns
        -------
        spec : Spectrum object
            A new spectrum object at redshift z.
        """

        if self._z is None:
            raise ValueError(
                "Must set current redshift in order to redshift" " spectrum"
            )

        if self._wunit.physical_type == u.m.physical_type:
            factor = (1.0 + z) / (1.0 + self._z)
        elif self._wunit.physical_type == u.Hz.physical_type:
            factor = (1.0 + self._z) / (1.0 + z)
        else:
            raise ValueError("wavelength must be in wavelength or frequency")

        d = self._wave * factor
        f = self._flux / factor
        if self._error is not None:
            e = self._error / factor
        else:
            e = None

        if adjust_flux:
            if self._dist is None and self._z == 0.0:
                raise ValueError(
                    "When current redshift is 0 and adjust_flux "
                    "is requested, current distance must be "
                    "defined"
                )
            if dist is None and z == 0.0:
                raise ValueError(
                    "When redshift is 0 and adjust_flux "
                    "is requested, dist must be defined"
                )
            if cosmo is None:
                cosmo = cosmology.get_current()

            if self._dist is None:
                dist_in = cosmo.luminosity_distance(self._z)
            else:
                dist_in = self._dist

            if dist is None:
                dist = cosmo.luminosity_distance(z)

            if dist_in <= 0.0 or dist <= 0.0:
                raise ValueError("Distances must be greater than 0.")

            # Adjust the flux
            factor = (dist_in / dist) ** 2
            f *= factor
            if e is not None:
                e *= factor

        return Spectrum(
            d,
            f,
            error=e,
            z=z,
            dist=dist,
            meta=self.meta,
            unit=self._unit,
            wave_unit=self._wunit,
        )

class MagSystem(object):
    """An abstract base class for magnitude systems."""

    __metaclass__ = abc.ABCMeta

    def __init__(self, name=None):
        self._zpbandflux = {}
        self._name = name

    @abc.abstractmethod
    def _refspectrum_bandflux(self, band):
        """Flux of the fundamental spectrophotometric standard."""
        pass

    @property
    def name(self):
        """Name of magnitude system."""
        return self._name

    @name.setter
    def name(self, value):
        self._name = value

    def zpbandflux(self, band):
        """Flux of an object with magnitude zero in the given bandpass.

        Parameters
        ----------
        bandpass : `~sncosmo.spectral.Bandpass` or str

        Returns
        -------
        bandflux : float
            Flux in photons / s / cm^2.
        """

        band = get_bandpass(band)
        try:
            return self._zpbandflux[band]
        except KeyError:
            bandflux = self._refspectrum_bandflux(band)
            self._zpbandflux[band] = bandflux

        return bandflux

    def band_flux_to_mag(self, flux, band):
        """Convert flux (photons / s / cm^2) to magnitude."""
        return -2.5 * math.log10(flux / self.zpbandflux(band))

    def band_mag_to_flux(self, mag, band):
        """Convert magnitude to flux in photons / s / cm^2"""
        return self.zpbandflux(band) * 10.0 ** (-0.4 * mag)

class CompositeMagSystem(MagSystem):
    """A magnitude system defined in a specific set of bands.

    In each band, there is a fundamental standard with a known
    (generally non-zero) magnitude.

    Parameters
    ----------
    bands: iterable of `~sncosmo.Bandpass` or str
        The filters in the magnitude system.
    standards: iterable of `~sncosmo.MagSystem` or str,
        The spectrophotmetric flux standards for each band, in the
        same order as `bands`.
    offsets: list_like
        The magnitude of standard in the given band.
    """

    def __init__(self, bands, standards, offsets, name=None):
        super(CompositeMagSystem, self).__init__(name=name)

        if not len(bands) == len(offsets) == len(standards):
            raise ValueError("Lengths of bands, standards, and offsets " "must match.")

        self._bands = [get_bandpass(band) for band in bands]
        self._standards = [get_magsystem(s) for s in standards]
        self._offsets = offsets

    @property
    def bands(self):
        return self._bands

    @property
    def standards(self):
        return self._standards

    @property
    def offsets(self):
        return self._offsets

    def _refspectrum_bandflux(self, band):
        if band not in self._bands:
            raise ValueError("band not in local magnitude system")
        i = self._bands.index(band)
        standard = self._standards[i]
        offset = self._offsets[i]

        return 10.0 ** (0.4 * offset) * standard.zpbandflux(band)

    def __str__(self):
        s = "CompositeMagSystem {!r}:\n".format(self.name)

        for i in range(len(self._bands)):
            s += "  {!r}: system={!r}  offset={}\n".format(
                self._bands[i].name, self._standards[i].name, self._offsets[i]
            )
        return s

class SpectralMagSystem(MagSystem):
    """A magnitude system defined by a fundamental spectrophotometric
    standard.

    Parameters
    ----------
    refspectrum : `sncosmo.Spectrum`
        The spectrum of the fundamental spectrophotometric standard.
    """

    def __init__(self, refspectrum, name=None):
        super(SpectralMagSystem, self).__init__(name)
        self._refspectrum = refspectrum

    def _refspectrum_bandflux(self, band):
        return self._refspectrum.bandflux(band)

class ABMagSystem(MagSystem):
    """Magnitude system where a source with F_nu = 3631 Jansky at all
    frequencies has magnitude 0 in all bands."""

    def _refspectrum_bandflux(self, band):
        bwave, btrans = band.to_unit(u.Hz)

        # AB spectrum is 3631 x 10^{-23} erg/s/cm^2/Hz
        # Get spectral values in photons/cm^2/s/Hz at bandpass wavelengths
        # by dividing by (h \nu).
        f = 3631.0e-23 / const.h.cgs.value / bwave

        binw = np.gradient(bwave)
        return np.sum(f * btrans * binw)

import numpy,import astropy,from astropy


# Copyright 2021 Bo Chen.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np

import torch
from torch._C import dtype
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist

from mmcv.runner import force_fp32
from mmdet.core import (MlvlPointGenerator, multi_apply, build_assigner, build_sampler,
                    distance2bbox, bbox2distance, images_to_levels, reduce_mean)
from mmdet.core.utils import filter_scores_and_topk, select_single_mlvl

from mmdet.core.bbox.iou_calculators.iou2d_calculator import bbox_overlaps
from ..module.conv import ConvModule, DepthwiseConvModule
from ..module.init_weights import normal_init
from .gfl_head import Integral
from ..builder import HEADS, build_loss

from .gfl_head import GFLHead, Integral

@HEADS.register_module()
class PicoDetATSSHead(GFLHead):
    """
    Modified from GFL, use same loss functions but much lightweight convolution heads
    """

    def __init__(
        self,
        num_classes,
        in_channels,
        feat_channels=96,
        stacked_convs=2,
        kernel_size=5,
        share_cls_reg=False,
        conv_type="DWConv",
        conv_cfg=None,
        norm_cfg=None,
        use_vfl=False,
        loss_cls=dict(
            type='QualityFocalLoss',
            use_sigmoid=True,
            beta=2.0,
            loss_weight=1.0),
        loss_dfl=dict(type='DistributionFocalLoss', loss_weight=0.25),
        reg_max=7,
        strides=[8, 16, 32],
        activation="LeakyReLU",
        sync_num_pos=True,
        init_cfg=dict(
            type='Normal',
            layer='Conv2d',
            std=0.01,
            override=dict(
                type='Normal',
                name='gfl_cls',
                std=0.01,
                bias_prob=0.01)),
        **kwargs):
        self.kernel_size = kernel_size
        self.conv_cfg = conv_cfg
        self.reg_max = reg_max
        self.share_cls_reg = share_cls_reg
        self.activation = activation
        self.ConvModule = ConvModule if conv_type == "Conv" else DepthwiseConvModule
        super(PicoDetATSSHead, self).__init__(
            num_classes, in_channels, stacked_convs=stacked_convs, strides=strides,
            feat_channels=feat_channels, loss_cls=loss_cls, norm_cfg=norm_cfg, init_cfg=init_cfg, **kwargs)
        self.sampling = False
        if self.train_cfg:
            self.assigner = build_assigner(self.train_cfg.assigner)
            # SSD sampling=False so use PseudoSampler
            sampler_cfg = dict(type='PseudoSampler')
            self.sampler = build_sampler(sampler_cfg, context=self)
        self.prior_generator = MlvlPointGenerator(strides, offset=0.5)
        self.sync_num_pos = sync_num_pos
        self.integral = Integral(self.reg_max)
        self.loss_dfl = build_loss(loss_dfl)
        self.use_vfl = use_vfl
        if use_vfl:
            assert loss_cls['type'] == 'VarifocalLoss', 'when set use_vfl=True, loss_cls type must be VarifocalLoss'
            self.loss_cls = build_loss(loss_cls)

    def _init_layers(self):
        self.cls_convs = nn.ModuleList()
        self.reg_convs = nn.ModuleList()
        for _ in self.strides: # 每个head都有一个
            cls_convs, reg_convs = self._buid_not_shared_head()
            self.cls_convs.append(cls_convs)
            self.reg_convs.append(reg_convs)

        self.gfl_cls = nn.ModuleList(
            [
                nn.Conv2d(
                    self.feat_channels,
                    self.cls_out_channels + 4 * (self.reg_max + 1)
                    if self.share_cls_reg
                    else self.cls_out_channels,
                    1,
                    padding=0,
                )
                for _ in self.strides
            ]
        )
        # TODO: if
        # self.gfl_reg = None
        # if not self.share_cls_reg:
        self.gfl_reg = nn.ModuleList(
            [
                nn.Conv2d(self.feat_channels, 4 * (self.reg_max + 1), 1, padding=0)
                for _ in self.strides
            ]
        )

    def _buid_not_shared_head(self):
        cls_convs = nn.ModuleList()
        reg_convs = nn.ModuleList()
        for i in range(self.stacked_convs):
            chn = self.in_channels if i == 0 else self.feat_channels
            cls_convs.append(
                self.ConvModule(
                    chn,
                    self.feat_channels,
                    self.kernel_size,
                    stride=1,
                    padding=(self.kernel_size - 1) // 2,
                    norm_cfg=self.norm_cfg,
                    bias=self.norm_cfg is None,
                    activation=self.activation,
                )
            )
            if not self.share_cls_reg:
                reg_convs.append(
                    self.ConvModule(
                        chn,
                        self.feat_channels,
                        self.kernel_size,
                        stride=1,
                        padding=(self.kernel_size - 1) // 2,
                        norm_cfg=self.norm_cfg,
                        bias=self.norm_cfg is None,
                        activation=self.activation,
                    )
                )

        return cls_convs, reg_convs

    def init_weights(self):
        for m in self.cls_convs.modules():
            if isinstance(m, nn.Conv2d):
                normal_init(m, std=0.01)
        for m in self.reg_convs.modules():
            if isinstance(m, nn.Conv2d):
                normal_init(m, std=0.01)
        # init cls head with confidence = 0.01
        bias_cls = -4.595
        for i in range(len(self.strides)):
            normal_init(self.gfl_cls[i], std=0.01, bias=bias_cls)
            normal_init(self.gfl_reg[i], std=0.01)
        print("Finish initialize PicoDet Head.")

    def forward(self, feats):
        return multi_apply(
            self.forward_single,
            feats,
            self.cls_convs,
            self.reg_convs,
            self.gfl_cls,
            self.gfl_reg,
        )

    def forward_single(self, x, cls_convs, reg_convs, gfl_cls, gfl_reg):
        cls_feat = x
        reg_feat = x
        for cls_conv in cls_convs:
            cls_feat = cls_conv(cls_feat)
        for reg_conv in reg_convs:
            reg_feat = reg_conv(reg_feat)
        if self.share_cls_reg:
            feat = gfl_cls(cls_feat)
            cls_score, bbox_pred = torch.split(
                feat, [self.cls_out_channels, 4 * (self.reg_max + 1)], dim=1
            )
        else:
            cls_score = gfl_cls(cls_feat)
            bbox_pred = gfl_reg(reg_feat)

        if torch.onnx.is_in_onnx_export():
            cls_score = (
                torch.sigmoid(cls_score)
                .reshape(1, self.num_classes, -1)
                .permute(0, 2, 1)
            )
            bbox_pred = bbox_pred.reshape(1, (self.reg_max + 1) * 4, -1).permute(
                0, 2, 1
            )
        return cls_score, bbox_pred

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
                              missing_keys, unexpected_keys, error_msgs):
        """Hack some keys of the model state dict so that can load checkpoints
        of previous version."""
        version = local_metadata.get('version', None)
        version = 'ppdet2mmdet'
        local_metadata['version'] = 'ppdet2mmdet'
        if version is None:
            # the key is different in early versions
            # for example, 'fcos_cls' become 'conv_cls' now
            bbox_head_keys = [
                k for k in state_dict.keys() if k.startswith(prefix)
            ]
            ori_predictor_keys = []
            new_predictor_keys = []
            # e.g. 'fcos_cls' or 'fcos_reg'
            for key in bbox_head_keys:
                ori_predictor_keys.append(key)
                key = key.split('.')
                conv_name = None
                if key[1].endswith('cls'):
                    conv_name = 'conv_cls'
                elif key[1].endswith('reg'):
                    conv_name = 'conv_reg'
                elif key[1].endswith('centerness'):
                    conv_name = 'conv_centerness'
                else:
                    assert NotImplementedError
                if conv_name is not None:
                    key[1] = conv_name
                    new_predictor_keys.append('.'.join(key))
                else:
                    ori_predictor_keys.pop(-1)
            for i in range(len(new_predictor_keys)):
                state_dict[new_predictor_keys[i]] = state_dict.pop(
                    ori_predictor_keys[i])
        super()._load_from_state_dict(state_dict, prefix, local_metadata,
                                      strict, missing_keys, unexpected_keys,
                                      error_msgs)


import os
import argparse
import csv
import pickle as pkl
import pandas as pd
from sklearn.metrics import accuracy_score, precision_score, classification_report, confusion_matrix
import sklearn
from sklearn import metrics
from sklearn.base import BaseEstimator, TransformerMixin
import re
import glob
import json
import numpy as np
import subprocess
import sys
subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'torch==1.4.0'])
subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'torchvision'])
subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'simpletransformers'])
import torch
import torch.distributed as dist
import torch.utils.data
import torch.utils.data.distributed

import simpletransformers
from simpletransformers.classification import ClassificationModel

def load_dataset(path, sep, header):
    data = pd.concat([pd.read_csv(f, sep=sep, header=header) for f in glob.glob('{}/*.csv'.format(path))], ignore_index = True)

    labels = data.iloc[:,0]
    features = data.drop(data.columns[0], axis=1)

    if header==None:
        # Adjust the column names after dropped the 0th column above
        # New column names are 0 (inclusive) to len(features.columns) (exclusive)
        new_column_names = list(range(0, len(features.columns)))
        features.columns = new_column_names

    return features, labels

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model-type', type=str, default='bert')
    parser.add_argument('--model-name', type=str, default='bert-base-cased')
    parser.add_argument('--backend', type=str, default='gloo')
    parser.add_argument('--train-data', type=str, default=os.environ['SM_CHANNEL_TRAIN'])
    parser.add_argument('--validation-data', type=str, default=os.environ['SM_CHANNEL_VALIDATION'])
    parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])
    parser.add_argument('--hosts', type=list, default=json.loads(os.environ['SM_HOSTS']))
    parser.add_argument('--current-host', type=str, default=os.environ['SM_CURRENT_HOST'])
    parser.add_argument('--num-gpus', type=int, default=os.environ['SM_NUM_GPUS'])

    args, _ = parser.parse_known_args()
    model_type = args.model_type
    model_name = args.model_name
    backend = args.backend
    train_data = args.train_data
    validation_data = args.validation_data
    model_dir = args.model_dir
    hosts = args.hosts
    current_host = args.current_host
    num_gpus = args.num_gpus

    # TODO:  Convert to distributed data loader
    #        https://pytorch.org/tutorials/beginner/aws_distributed_training_tutorial.html
    #        https://github.com/aws/sagemaker-python-sdk/issues/1110
    is_distributed = len(args.hosts) > 1 and args.backend is not None
    print('Distributed training - {}'.format(is_distributed))
    use_cuda = args.num_gpus > 0
    print('Number of gpus available - {}'.format(args.num_gpus))
    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
    device = torch.device("cuda" if use_cuda else "cpu")

    if is_distributed:
        # Initialize the distributed environment.
        world_size = len(args.hosts)
        os.environ['WORLD_SIZE'] = str(world_size)
        host_rank = args.hosts.index(args.current_host)
        os.environ['RANK'] = str(host_rank)
        dist.init_process_group(backend=args.backend, rank=host_rank, world_size=world_size)
        print('Initialized the distributed environment: \'{}\' backend on {} nodes. '.format(
            args.backend, dist.get_world_size()) + 'Current host rank is {}. Number of gpus: {}'.format(
            dist.get_rank(), args.num_gpus))

    # TODO:  Change this to use SM_CHANNEL_TRAIN and DistributedDataLoader, etc
    # X_train, y_train = load_dataset(train_data, ',', header=0)
    # X_validation, y_validation = load_dataset(validation_data, ',', header=0)

    df1 = pd.read_csv('./data/amazon_reviews_us_Digital_Software_v1_00.tsv.gz',
                 delimiter='\t',
                 quoting=csv.QUOTE_NONE,
                 compression='gzip',
                 header=0)[:100]
    print(df1.shape)

    df2 = pd.read_csv('./data/amazon_reviews_us_Video_Games_v1_00.tsv.gz',
                 delimiter='\t',
                 quoting=csv.QUOTE_NONE,
                 compression='gzip',
                 header=0)[:100]
    print(df2.shape)

    df = pd.concat([df1, df2])

    print('YES: {}'.format(df.isna().values.any()))
    df = df.dropna()
    df = df.reset_index(drop=True)

    # Enrich the data
    df['is_positive_sentiment'] = (df['star_rating'] >= 4).astype(int)

    df_bert = df[['review_body', 'is_positive_sentiment']]
    df_bert.columns = ['text', 'labels']
    df_bert.head(5)

    print(df_bert.shape)

    df_bert = df_bert #[:200100]
    df_bert.shape

    from sklearn.model_selection import train_test_split

    df_bert_train, df_bert_holdout = train_test_split(df_bert, test_size=0.10)
    df_bert_validation, df_bert_test = train_test_split(df_bert_holdout, test_size=0.50)

    print(df_bert_train.shape)
    print(df_bert_validation.shape)
    print(df_bert_test.shape)

    # TODO:  change output_dir to SM_model_dir or output_path
    bert_args = {
       'output_dir': model_dir,
       'cache_dir': 'cache/',
       'fp16': False,
       'max_seq_length': 128,
       'train_batch_size': 8,
       'eval_batch_size': 8,
       'gradient_accumulation_steps': 1,
       'num_train_epochs': 1,
       'weight_decay': 0,
       'learning_rate': 3e-5,
       'adam_epsilon': 1e-8,
       'warmup_ratio': 0.06,
       'warmup_steps': 0,
       'max_grad_norm': 1.0,
       'logging_steps': 50,
       'evaluate_during_training': False,
       'save_steps': 2000,
       'eval_all_checkpoints': True,
       'use_tensorboard': True,
       'tensorboard_dir': 'tensorboard',
       'overwrite_output_dir': True,
       'reprocess_input_data': False,
    }

    bert_model = ClassificationModel(model_type='distilbert', # bert, distilbert, etc, etc.
                                     model_name='distilbert-base-cased',
                                     args=bert_args,
                                     use_cuda=use_cuda)

    bert_model.train_model(train_df=df_bert_train,
                           eval_df=df_bert_validation,
                           show_running_loss=True)

    # TODO:  use the model_dir that is passed in through args
    #        (currently SM_MODEL_DIR)
#    os.makedirs(model_dir, exist_ok=True)
#    model_path = os.path.join(model_dir, 'bert-model')

#    pkl.dump(bert_model, open(model_path, 'wb'))
#    print('Wrote model to {}'.format(model_path))

#    result, model_outputs, wrong_predictions = bert_model.eval_model(eval_df=df_bert_test, acc=sklearn.metrics.accuracy_score)

#    print(result)

    # Show bad predictions
#    print('Number of wrong predictions: {}'.format(len(wrong_predictions)))
#    print('\n')

#    for prediction in wrong_predictions:
#        print(prediction.text_a)
#        print('\n')

#    predictions, raw_outputs = bert_model.predict(["""I really enjoyed this item.  I highly recommend it."""])

#    print('Predictions: {}'.format(predictions))
#    print('Raw outputs: {}'.format(raw_outputs))

#    predictions, raw_outputs = bert_model.predict(["""This item is awful and terrible."""])

#    print('Predictions: {}'.format(predictions))
#    print('Raw outputs: {}'.format(raw_outputs))


import numpy as np
import pytest

from sklearn.compose import make_column_transformer
from sklearn.datasets import load_breast_cancer, make_classification
from sklearn.exceptions import NotFittedError
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import average_precision_score, precision_recall_curve
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC, SVR
from sklearn.utils import shuffle

from sklearn.metrics import PrecisionRecallDisplay, plot_precision_recall_curve

# TODO: Remove when https://github.com/numpy/numpy/issues/14397 is resolved
pytestmark = pytest.mark.filterwarnings(
    "ignore:In future, it will be an error for 'np.bool_':DeprecationWarning:"
    "matplotlib.*"
)

def test_precision_recall_display_validation(pyplot):
    """Check that we raise the proper error when validating parameters."""
    X, y = make_classification(
        n_samples=100, n_informative=5, n_classes=5, random_state=0
    )

    with pytest.raises(NotFittedError):
        PrecisionRecallDisplay.from_estimator(SVC(), X, y)

    regressor = SVR().fit(X, y)
    y_pred_regressor = regressor.predict(X)
    classifier = SVC(probability=True).fit(X, y)
    y_pred_classifier = classifier.predict_proba(X)[:, -1]

    err_msg = "PrecisionRecallDisplay.from_estimator only supports classifiers"
    with pytest.raises(ValueError, match=err_msg):
        PrecisionRecallDisplay.from_estimator(regressor, X, y)

    err_msg = "Expected 'estimator' to be a binary classifier, but got SVC"
    with pytest.raises(ValueError, match=err_msg):
        PrecisionRecallDisplay.from_estimator(classifier, X, y)

    err_msg = "{} format is not supported"
    with pytest.raises(ValueError, match=err_msg.format("continuous")):
        # Force `y_true` to be seen as a regression problem
        PrecisionRecallDisplay.from_predictions(y + 0.5, y_pred_classifier, pos_label=1)
    with pytest.raises(ValueError, match=err_msg.format("multiclass")):
        PrecisionRecallDisplay.from_predictions(y, y_pred_regressor, pos_label=1)

    err_msg = "Found input variables with inconsistent numbers of samples"
    with pytest.raises(ValueError, match=err_msg):
        PrecisionRecallDisplay.from_predictions(y, y_pred_classifier[::2])

    X, y = make_classification(n_classes=2, n_samples=50, random_state=0)
    y += 10
    classifier.fit(X, y)
    y_pred_classifier = classifier.predict_proba(X)[:, -1]
    err_msg = r"y_true takes value in {10, 11} and pos_label is not specified"
    with pytest.raises(ValueError, match=err_msg):
        PrecisionRecallDisplay.from_predictions(y, y_pred_classifier)

# FIXME: Remove in 1.2
def test_plot_precision_recall_curve_deprecation(pyplot):
    """Check that we raise a FutureWarning when calling
    `plot_precision_recall_curve`."""

    X, y = make_classification(random_state=0)
    clf = LogisticRegression().fit(X, y)
    deprecation_warning = "Function plot_precision_recall_curve is deprecated"
    with pytest.warns(FutureWarning, match=deprecation_warning):
        plot_precision_recall_curve(clf, X, y)

@pytest.mark.parametrize("constructor_name", ["from_estimator", "from_predictions"])
@pytest.mark.parametrize("response_method", ["predict_proba", "decision_function"])
def test_precision_recall_display_plotting(pyplot, constructor_name, response_method):
    """Check the overall plotting rendering."""
    X, y = make_classification(n_classes=2, n_samples=50, random_state=0)
    pos_label = 1

    classifier = LogisticRegression().fit(X, y)
    classifier.fit(X, y)

    y_pred = getattr(classifier, response_method)(X)
    y_pred = y_pred if y_pred.ndim == 1 else y_pred[:, pos_label]

    # safe guard for the binary if/else construction
    assert constructor_name in ("from_estimator", "from_predictions")

    if constructor_name == "from_estimator":
        display = PrecisionRecallDisplay.from_estimator(
            classifier, X, y, response_method=response_method
        )
    else:
        display = PrecisionRecallDisplay.from_predictions(
            y, y_pred, pos_label=pos_label
        )

    precision, recall, _ = precision_recall_curve(y, y_pred, pos_label=pos_label)
    average_precision = average_precision_score(y, y_pred, pos_label=pos_label)

    np.testing.assert_allclose(display.precision, precision)
    np.testing.assert_allclose(display.recall, recall)
    assert display.average_precision == pytest.approx(average_precision)

    import matplotlib as mpl

    assert isinstance(display.line_, mpl.lines.Line2D)
    assert isinstance(display.ax_, mpl.axes.Axes)
    assert isinstance(display.figure_, mpl.figure.Figure)

    assert display.ax_.get_xlabel() == "Recall (Positive label: 1)"
    assert display.ax_.get_ylabel() == "Precision (Positive label: 1)"

    # plotting passing some new parameters
    display.plot(alpha=0.8, name="MySpecialEstimator")
    expected_label = f"MySpecialEstimator (AP = {average_precision:0.2f})"
    assert display.line_.get_label() == expected_label
    assert display.line_.get_alpha() == pytest.approx(0.8)

@pytest.mark.parametrize(
    "constructor_name, default_label",
    [
        ("from_estimator", "LogisticRegression (AP = {:.2f})"),
        ("from_predictions", "Classifier (AP = {:.2f})"),
    ],
)
def test_precision_recall_display_name(pyplot, constructor_name, default_label):
    """Check the behaviour of the name parameters"""
    X, y = make_classification(n_classes=2, n_samples=100, random_state=0)
    pos_label = 1

    classifier = LogisticRegression().fit(X, y)
    classifier.fit(X, y)

    y_pred = classifier.predict_proba(X)[:, pos_label]

    # safe guard for the binary if/else construction
    assert constructor_name in ("from_estimator", "from_predictions")

    if constructor_name == "from_estimator":
        display = PrecisionRecallDisplay.from_estimator(classifier, X, y)
    else:
        display = PrecisionRecallDisplay.from_predictions(
            y, y_pred, pos_label=pos_label
        )

    average_precision = average_precision_score(y, y_pred, pos_label=pos_label)

    # check that the default name is used
    assert display.line_.get_label() == default_label.format(average_precision)

    # check that the name can be set
    display.plot(name="MySpecialEstimator")
    assert (
        display.line_.get_label()
        == f"MySpecialEstimator (AP = {average_precision:.2f})"
    )

@pytest.mark.parametrize(
    "clf",
    [
        make_pipeline(StandardScaler(), LogisticRegression()),
        make_pipeline(
            make_column_transformer((StandardScaler(), [0, 1])), LogisticRegression()
        ),
    ],
)
def test_precision_recall_display_pipeline(pyplot, clf):
    X, y = make_classification(n_classes=2, n_samples=50, random_state=0)
    with pytest.raises(NotFittedError):
        PrecisionRecallDisplay.from_estimator(clf, X, y)
    clf.fit(X, y)
    display = PrecisionRecallDisplay.from_estimator(clf, X, y)
    assert display.estimator_name == clf.__class__.__name__

def test_precision_recall_display_string_labels(pyplot):
    # regression test #15738
    cancer = load_breast_cancer()
    X, y = cancer.data, cancer.target_names[cancer.target]

    lr = make_pipeline(StandardScaler(), LogisticRegression())
    lr.fit(X, y)
    for klass in cancer.target_names:
        assert klass in lr.classes_
    display = PrecisionRecallDisplay.from_estimator(lr, X, y)

    y_pred = lr.predict_proba(X)[:, 1]
    avg_prec = average_precision_score(y, y_pred, pos_label=lr.classes_[1])

    assert display.average_precision == pytest.approx(avg_prec)
    assert display.estimator_name == lr.__class__.__name__

    err_msg = r"y_true takes value in {'benign', 'malignant'}"
    with pytest.raises(ValueError, match=err_msg):
        PrecisionRecallDisplay.from_predictions(y, y_pred)

    display = PrecisionRecallDisplay.from_predictions(
        y, y_pred, pos_label=lr.classes_[1]
    )
    assert display.average_precision == pytest.approx(avg_prec)

@pytest.mark.parametrize(
    "average_precision, estimator_name, expected_label",
    [
        (0.9, None, "AP = 0.90"),
        (None, "my_est", "my_est"),
        (0.8, "my_est2", "my_est2 (AP = 0.80)"),
    ],
)
def test_default_labels(pyplot, average_precision, estimator_name, expected_label):
    """Check the default labels used in the display."""
    precision = np.array([1, 0.5, 0])
    recall = np.array([0, 0.5, 1])
    display = PrecisionRecallDisplay(
        precision,
        recall,
        average_precision=average_precision,
        estimator_name=estimator_name,
    )
    display.plot()
    assert display.line_.get_label() == expected_label

@pytest.mark.parametrize("constructor_name", ["from_estimator", "from_predictions"])
@pytest.mark.parametrize("response_method", ["predict_proba", "decision_function"])
def test_plot_precision_recall_pos_label(pyplot, constructor_name, response_method):
    # check that we can provide the positive label and display the proper
    # statistics
    X, y = load_breast_cancer(return_X_y=True)
    # create an highly imbalanced version of the breast cancer dataset
    idx_positive = np.flatnonzero(y == 1)
    idx_negative = np.flatnonzero(y == 0)
    idx_selected = np.hstack([idx_negative, idx_positive[:25]])
    X, y = X[idx_selected], y[idx_selected]
    X, y = shuffle(X, y, random_state=42)
    # only use 2 features to make the problem even harder
    X = X[:, :2]
    y = np.array(["cancer" if c == 1 else "not cancer" for c in y], dtype=object)
    X_train, X_test, y_train, y_test = train_test_split(
        X,
        y,
        stratify=y,
        random_state=0,
    )

    classifier = LogisticRegression()
    classifier.fit(X_train, y_train)

    # sanity check to be sure the positive class is classes_[0] and that we
    # are betrayed by the class imbalance
    assert classifier.classes_.tolist() == ["cancer", "not cancer"]

    y_pred = getattr(classifier, response_method)(X_test)
    # we select the correcponding probability columns or reverse the decision
    #  function otherwise
    y_pred_cancer = -1 * y_pred if y_pred.ndim == 1 else y_pred[:, 0]
    y_pred_not_cancer = y_pred if y_pred.ndim == 1 else y_pred[:, 1]

    if constructor_name == "from_estimator":
        display = PrecisionRecallDisplay.from_estimator(
            classifier,
            X_test,
            y_test,
            pos_label="cancer",
            response_method=response_method,
        )
    else:
        display = PrecisionRecallDisplay.from_predictions(
            y_test,
            y_pred_cancer,
            pos_label="cancer",
        )
    # we should obtain the statistics of the "cancer" class
    avg_prec_limit = 0.65
    assert display.average_precision < avg_prec_limit
    assert -np.trapz(display.precision, display.recall) < avg_prec_limit

    # otherwise we should obtain the statistics of the "not cancer" class
    if constructor_name == "from_estimator":
        display = PrecisionRecallDisplay.from_estimator(
            classifier,
            X_test,
            y_test,
            response_method=response_method,
            pos_label="not cancer",
        )
    else:
        display = PrecisionRecallDisplay.from_predictions(
            y_test,
            y_pred_not_cancer,
            pos_label="not cancer",
        )
    avg_prec_limit = 0.95
    assert display.average_precision > avg_prec_limit
    assert -np.trapz(display.precision, display.recall) > avg_prec_limit


from torch.utils.data import DataLoader, TensorDataset
from model import Zi2ZiModel
import os
import sys
import argparse
import torch
import random
import time
import math
import numpy as np
import pandas as pd
from PIL import Image, ImageDraw, ImageFont
import torchvision.transforms as transforms
from torchvision.utils import save_image, make_grid
import time
from tqdm import tqdm

writer_dict = {
        '智永': 0, ' 隸書-趙之謙': 1, '張即之': 2, '張猛龍碑': 3, '柳公權': 4, '標楷體-手寫': 5, '歐陽詢-九成宮': 6,
        '歐陽詢-皇甫誕': 7, '沈尹默': 8, '美工-崩雲體': 9, '美工-瘦顏體': 10, '虞世南': 11, '行書-傅山': 12, '行書-王壯為': 13,
        '行書-王鐸': 14, '行書-米芾': 15, '行書-趙孟頫': 16, '行書-鄭板橋': 17, '行書-集字聖教序': 18, '褚遂良': 19, '趙之謙': 20,
        '趙孟頫三門記體': 21, '隸書-伊秉綬': 22, '隸書-何紹基': 23, '隸書-鄧石如': 24, '隸書-金農': 25,  '顏真卿-顏勤禮碑': 26,
        '顏真卿多寶塔體': 27, '魏碑': 28
    }

parser = argparse.ArgumentParser(description='Infer')
parser.add_argument('--experiment_dir', required=True,
                    help='experiment directory, data, samples,checkpoints,etc')
parser.add_argument('--gpu_ids', default=[], nargs='+', help="GPUs")
parser.add_argument('--image_size', type=int, default=256,
                    help="size of your input and output image")
parser.add_argument('--L1_penalty', type=int, default=100, help='weight for L1 loss')
parser.add_argument('--Lconst_penalty', type=int, default=15, help='weight for const loss')
# parser.add_argument('--Ltv_penalty', dest='Ltv_penalty', type=float, default=0.0, help='weight for tv loss')
parser.add_argument('--Lcategory_penalty', type=float, default=1.0,
                    help='weight for category loss')
parser.add_argument('--embedding_num', type=int, default=40,
                    help="number for distinct embeddings")
parser.add_argument('--embedding_dim', type=int, default=128, help="dimension for embedding")
parser.add_argument('--batch_size', type=int, default=1, help='number of examples in batch')
parser.add_argument('--lr', type=float, default=0.001, help='initial learning rate for adam')
parser.add_argument('--resume', type=int, default=140000, help='resume from previous training')
parser.add_argument('--input_nc', type=int, default=1)

parser.add_argument('--charset', type=str, choices=['s', 'm', 'l'], default='l')
parser.add_argument('--canvas_size', type=int, default=256)
parser.add_argument('--char_size', type=int, default=256)
parser.add_argument('--src_font', type=str, default='charset/gbk/方正新楷体_GBK(完整).TTF')
parser.add_argument('--label', type=int, default=0)

def draw_single_char(ch, font, canvas_size):
    img = Image.new("RGB", (canvas_size, canvas_size), (255, 255, 255))
    draw = ImageDraw.Draw(img)
    draw.text((0, 0), ch, (0, 0, 0), font=font)
    img = img.convert('L')
    return img

def get_charset(charset_size):
    if charset_size == 's':
        with open('charset/charset_s.txt', 'r', encoding='utf-8') as charset_s_file:
            charset_txt = charset_s_file.readlines()
            charset = [s.strip() for s in charset_txt]
    else:
        charset_csv = pd.read_csv('charset/all_abooks.unigrams_desc.Clean.rate.csv')
        if charset_size == 'm':
            charset = charset_csv[['char']][charset_csv['acc_rate'] <= 0.999].values.squeeze(axis=-1).tolist()
        elif charset_size == 'l':
            charset = charset_csv[['char']][charset_csv['acc_rate'] <= 0.9999].values.squeeze(axis=-1).tolist()
        else:
            raise ValueError('charset_size should be s, m or l.')
        return ''.join(charset)

def main():
    args = parser.parse_args()
    data_dir = os.path.join(args.experiment_dir, "data")
    checkpoint_dir = os.path.join(args.experiment_dir, "checkpoint")
    sample_dir = os.path.join(args.experiment_dir, "sample")
    infer_dir = os.path.join(args.experiment_dir, "infer")

    # train_dataset = DatasetFromObj(os.path.join(data_dir, 'train.obj'), augment=True, bold=True, rotate=True, blur=True)
    # val_dataset = DatasetFromObj(os.path.join(data_dir, 'val.obj'))
    # dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)

    model = Zi2ZiModel(
        input_nc=args.input_nc,
        embedding_num=args.embedding_num,
        embedding_dim=args.embedding_dim,
        Lconst_penalty=args.Lconst_penalty,
        Lcategory_penalty=args.Lcategory_penalty,
        save_dir=checkpoint_dir,
        gpu_ids=args.gpu_ids,
        is_training=False
    )
    model.setup()
    model.print_networks(True)
    model.load_networks(args.resume)
    sys.stdout.flush()

    src = get_charset(args.charset)
    font = ImageFont.truetype(args.src_font, size=args.char_size)
    img_list = [transforms.Normalize(0.5, 0.5)(
        transforms.ToTensor()(
            draw_single_char(ch, font, args.canvas_size)
        )
    ).unsqueeze(dim=0) for ch in src]
    label_list = [args.label for _ in src]

    img_list = torch.cat(img_list, dim=0)
    label_list = torch.tensor(label_list)

    dataset = TensorDataset(label_list, img_list, img_list)
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)

    global writer_dict
    writer_dict_inv = {v: k for k, v in writer_dict.items()}

    for label_idx in range(29):
        dir_path = os.path.join(infer_dir, writer_dict_inv[label_idx])
        if not os.path.isdir(dir_path):
            os.mkdir(dir_path)

    for label_idx in range(29):
        writer = writer_dict_inv[label_idx]
        for i, batch in enumerate(dataloader):
            model.set_input(torch.ones_like(batch[0]) * label_idx, batch[2], batch[1])
            model.forward()
            tensor_to_plot = model.fake_B.detach()
            tensor_to_plot = tensor_to_plot.squeeze(0)
            tensor_to_plot = tensor_to_plot.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).squeeze(-1)
            tensor_to_plot = tensor_to_plot.to('cpu', torch.uint8)
            img_ndarray = tensor_to_plot.numpy()
            im = Image.fromarray(img_ndarray, mode='L')
            im.save(os.path.join(infer_dir, writer, src[i] + '.png'))
        print('writer: ' + writer + ' complete.')

if __name__ == '__main__':
    with torch.no_grad():
        main()


import skimage.io as io
import skimage.transform as skt
import numpy as np
from PIL import Image, ImageOps
from src.models.class_patcher import patcher
from src.utils.imgproc import *

class patcher(patcher):
    def __init__(self, body='./body/body_supica.png', **options):
        super().__init__('スピカ', body=body, pantie_position=[19, 0], **options)
        self.mask = io.imread('./mask/mask_supica.png')
        self.components = io.imread('./material/supica_components.png')[:, :, -1] / 255.0
        self.components_shade = io.imread('./material/supica_components_shade.png')[:, :, -1] / 255.0
        self.components_position = [1221, 17]

    def gen_components(self, image):
        pantie = np.array(image)
        ribbon = pantie[19:58, 5:35, :3] / 255.0
        base_color = np.mean(np.mean(ribbon[5:12, 16:20], axis=0), axis=0)
        shade_color = np.mean(np.mean(ribbon[8:14, 7:15], axis=0), axis=0)
        components = self.components[:, :, None] * base_color
        shade = self.components_shade[:, :, None] * (1 - shade_color)
        components -= shade
        components = np.dstack((components, components[:, :, 0] > 0))
        return Image.fromarray(np.uint8(np.clip(components, 0, 1) * 255))

    def convert(self, image):
        pantie = np.array(image)
        patch = np.copy(pantie[-170:, 546:, :])
        pantie[-150:, 546:, :] = 0
        patch = skt.resize(patch[::-1, ::-1, :], (patch.shape[0], 50), anti_aliasing=True, mode='reflect')
        [pr, pc, d] = patch.shape
        pantie[142:142 + pr, :pc, :] = np.uint8(patch * 255)
        pantie = np.uint8(resize(pantie, [1.38, 1.62]) * 255)[:, 11:]
        pantie = np.bitwise_and(pantie, self.mask)[:, 1:]
        pantie = np.concatenate((pantie[:, ::-1], pantie), axis=1)
        components = self.gen_components(image)
        return Image.fromarray(pantie)

    def patch(self, image, transparent=False):
        pantie = self.convert(image)
        components = self.gen_components(image)
        if transparent:
            patched = Image.new("RGBA", self.body_size)
        else:
            patched = self.body.copy()
        self.paste(patched, pantie, self.pantie_position)
        self.paste(patched, components, self.components_position)
        self.paste(patched, ImageOps.mirror(components), [2048 - self.components_position[0] - components.width, self.components_position[1]])
        return patched


import numpy as np
import ogg.vorbis
import pylab as p
import scikits.openopt
from scikits import audiolab
import cPickle
import os
import sys
sys.path.append('../../')
from channel_config import ChannelConfig
from PyQt4 import QtCore, QtGui
from utils import Utils
from pylab_settings import DispSettings

class BatchSoundAnalyser(object):
    def __init__(self, i_display=False):
        self.sound_analyser = SoundAnalyser()
        self.times = None
        self.display = i_display
        self.utils = Utils()
        self.disp = DispSettings()
        self.disp.setDispSoundAmplitudes()

    def computeEndOfSounds(self, i_alphabet, i_ground_truth=None, o_file_name=None  ):
        self.times = {}
        for letter in i_alphabet:
            self.times [letter] =  self.sound_analyser.getEndOfSound()
        if o_file_name is not None:
            filename =  i_sound_dir + o_file_name
            print "Saving time information to ", filename
            self.disp.savePickle(self.times, filename)

    def displayChannelSounds(self, i_sound_dir, i_channel_config, i_timing_info, i_plot_axis=True):
        #Get the channel colors (assume 5 for now)
        channel_colours = [ QtGui.QColor("red"),
                                                    QtGui.QColor("green"),
                                                    QtGui.QColor("blue"),
                                                    QtGui.QColor("Brown"),
                                                    QtGui.QColor("Purple")]
        colours = [tuple( list( np.array(c.getRgb()[0:-1])/255. )) for c in channel_colours]
        #Other plot parameters
        self.disp.newFigure()
        n_channels =  len(i_channel_config)
        n_letters_group = len(i_channel_config[0])
        n_letters = n_channels*n_letters_group
        (tmax, tmin, ymax,ymin) = (max(i_timing_info),  0, 1.5, -1.5)
        p.axis( [tmin, tmax, ymin, ymax])

        #Initialise all counters
        channel_cnt = -1
        letter_cnt = 0
        time_cnt = 0
        yticks = []
        ytick_labels = []
        for n in range(0, n_letters):
            channel_cnt += 1
            if channel_cnt == n_channels:
                channel_cnt = 0
                letter_cnt += 1
            letter = i_channel_config[channel_cnt][letter_cnt]
            if letter == "*":
                continue;
            print "channel_cnt = ", channel_cnt , " letter_cnt = ", letter_cnt,
            print " letter = ", letter
            sound_file = i_sound_dir + letter + ".wav"
            self.sound_analyser.setSoundFile(sound_file)
            x = i_timing_info[time_cnt]
            y = -(channel_cnt+1) * (ymax - ymin)
            ytick_labels.append(1.0)
            ytick_labels.append(-1.0)
            yticks.append(y+1.0)
            yticks.append(y-1.0)
            time_cnt += 1
            self.sound_analyser.plotSound(  x, y, colours[channel_cnt])
            p.text(  x, y+ymax-0.4, str(letter) )

        #Some final plots parameters
        if i_plot_axis:
            p.axis('on')
        else:
            p.axis('off')
        p.grid('on')
        ax = p.gca()
        ax.set_yticks( tuple(yticks) )
        ax.set_yticklabels(  ytick_labels )
        p.xlabel("Time (seconds)")
        p.ylabel("Normalised Amplitude Range (Per Channel)")
        p.axis('tight')
        self.disp.saveFig("alphabet_sounds", i_save_eps=True)
        p.show()

    def displaySounds(self,  i_sound_dir,  i_alphabet, i_labels=True, i_ground_truth=None, i_plot_axis=False):
        nrows = np.int32(np.round(np.sqrt(len(i_alphabet))))
        ncols = 1
        while (1):
            if nrows*ncols >= len(i_alphabet):
                break
            ncols += 1
        p.figure(facecolor='w')
        total_diff = 0.0
        for n in range(0, len(i_alphabet)):
            letter = i_alphabet[n]
            sound_file = i_sound_dir + letter + ".ogg"
            self.sound_analyser.setSoundFile(sound_file)
            if self.times is not None:
                self.sound_analyser.setEndOfSound(self.times[letter])
                if self.display:
                    diff = np.abs( self.times[letter] -  i_ground_truth[letter] )
                    total_diff += diff
                    print "Letter: %s, time  = %0.2f, grnd_truth = %0.2f, diff=%0.2f" % (letter, self.times[letter],  i_ground_truth[letter], diff)
            p.subplot(nrows, ncols, n+1)
            self.sound_analyser. plotAmplitude( i_ground_truth=i_ground_truth, i_labels= i_labels)
            p.title(letter)
            if i_plot_axis:
                p.axis('on')
            else:
                p.axis('off')
        if self.display:
            print "Average diff=0.2f" % (total_diff / len(i_alphabet))
        p.show()

class SoundAnalyser(object):
    def __init__(self, i_sound_file=None, i_display=False):
        self.rms_thresh = 0.001
        self.display = i_display
        if i_sound_file is not None:
            self.setSoundFile(i_sound_file)

    def setSoundFile(self, i_sound_file):
        self.sound_file = i_sound_file
        (   self.sound, self.n_channels, self.sampling_freq,
            self.file_duration, self.nsamples, self.letter ) = self.soundInfo(i_sound_file)
        self.rms_vals = None
        self.time_end = None#Time when the sound ends

    def soundInfo(self, i_sound_file):
        print "Sound file = ", i_sound_file
        (snd, sampling_freq, n_bits) = audiolab.wavread(i_sound_file)
        if snd.ndim == 1:
            n_channels = 1
            channel1 = np.array(snd)
        else:
            n_channels = snd.shape[1]
            channel1 = np.array(snd[:,0])
        n_samples = snd.shape[0]
        file_duration = float(n_samples) / float( sampling_freq )
        letter = str(i_sound_file).split('/')[-1]
        if letter == "..ogg":
            letter = "."
        else:
            letter = letter.split('.')[0]

        if self.display:
            print "File info: ", i_sound_file
            print "Number of samples:  " , n_samples
            print "Number of channels: ", n_channels
            print "File duration: ", file_duration, " seconds"
            print "Currrent letter: ", letter
        return (channel1, n_channels, sampling_freq, file_duration, n_samples, letter)

    def rms(self, i_samples=None, i_fast=True):
        if i_fast:
            return  np.sqrt(np.mean(i_samples**2))
        (fourier, n_unique_points) = self.fft(i_samples)
        return np.sqrt(np.sum(fourier))

    def fft(self, i_samples=None):
        if i_samples is None:
            fourier = p.fft(self.sound)
        else:
            fourier = p.fft(i_samples)
        n_unique_pts = np.ceil((self.nsamples+1)/2.0)
        fourier= np.abs(fourier[0:n_unique_pts])
        fourier = fourier / float(self.nsamples)# scale by the number of points so that
                                            #the magnitude does not depend on the length
                                            # of the signal or on its sampling frequency
        fourier= fourier**2  # square it to get the power

        # multiply by two (see technical document for details)
        # odd nfft excludes Nyquist point
        if (self.nsamples % 2) > 0: # we've got odd number of points fft
            fourier[1:self.nsamples] = fourier[1:self.nsamples] * 2
        else:
            fourier[1:self.nsamples-1] = fourier[1:self.nsamples- 1] * 2 # we've got even number of points fft
        return (fourier, n_unique_pts)

    def setEndOfSound(self, i_time):
        self.time_end  = float(i_time)

    def getEndOfSound(self):
        if self.time_end  is None:
            if self.rms_vals is None:
                rms_vals = [self.rms(i_samples=self.sound[-n:], i_fast=True) for n in range(1,self.nsamples+1)]
                rms_vals.reverse()
                self.rms_vals = np.array(rms_vals)
            idx = np.nonzero( np.array( self.rms_vals ) >= self.rms_thresh )[0]
            idx = idx[-1]
            t =    np.arange(0.0, float(self.nsamples), 1.0) / float(self.sampling_freq)
            self.time_end = t[idx]
        return self.time_end

####################################### Display functions

    def plotRms(self, i_labels=False):
        t_end =self.getEndOfSound()
        t = np.arange(0.0, self.nsamples, 1.0) / self.sampling_freq
        max_y = max(self.rms_vals)
        min_y = min(self.rms_vals)
        p.plot( t, self.rms_vals, color='k' )
        p.hold('True')
        p.plot([t_end, t_end], [min_y, max_y],'r')
        if i_labels:
            p.ylabel('Accummulating RMS')
            p.xlabel('Time (ms)')

    def plotSound(self, x, y, i_color='k'):
        t =  np.arange(0.0, self.nsamples , 1.0) / self.sampling_freq
        t += x
        max_y = max(self.sound)
        min_y = min(self.sound)
        p.plot(t, self.sound + y, color=i_color)
        p.plot([x+0.15, x+0.15], [-1.2+y, 1.5+y ],'k', linewidth=3)

    def plotAmplitude(self, i_ground_truth=None, i_labels=False, i_time_offset=0.0, i_plot_sound_end=True):
        t_end =self.getEndOfSound() + i_time_offset
        t =  np.arange(0.0, self.nsamples , 1.0) / self.sampling_freq
        t += i_time_offset
        max_y = max(self.sound)
        min_y = min(self.sound)
        p.plot(t, self.sound, color='k')
        p.hold('True')
        if i_plot_sound_end:
            p.plot([t_end, t_end], [min_y, max_y],'r', linewidth=3)
        if i_ground_truth is not None:
            t_end_grnd_truth =  i_ground_truth[self.letter] + i_time_offset
            p.plot([t_end_grnd_truth, t_end_grnd_truth], [min_y, max_y],'b')
        if i_labels:
            p.ylabel('Amplitude')
            p.xlabel('Time (ms)')

    def plotPowerSpectrum(self,i_labels=False ):
        (fourier, n_unique_pts)=  self.fft()
        freq = np.arange(0, n_unique_pts, 1.0) * (self.sampling_freq / self.nsamples)
        print "RMS from amplitude: ",  np.sqrt(np.mean(self.sound**2))
        print "RMS from power spectrum: ",  np.sqrt(np.sum(fourier))
        p.plot(freq , 10.*np.log10(fourier), color='k')
        if i_labels:
            p.xlabel('Frequency (kHz)')
            p.ylabel('Power (dB)')

    def displaySound(self, i_sound_file=None, i_labels=True, i_ground_truth=None):
        if i_sound_file is not None:
            self.setSoundFile(i_sound_file)
        p.figure()
        p.subplot(2,2,1)
        sound_analyser.plotAmplitude( i_ground_truth=i_ground_truth, i_labels=i_labels)
        p.subplot(2,2,2)
        sound_analyser.plotPowerSpectrum(i_labels=i_labels)
        p.subplot(2,2,3)
        sound_analyser.plotRms(i_labels=i_labels)
        p.show()

class BatchExtractor(object):
    def __init__(self, i_alphabet_sequence, i_file_name, out_dir,  i_start=0.1, i_period=0.4):
        self.start = i_start
        self.file_name = i_file_name
        self.period = i_period
        self.seq = np.array(i_alphabet_sequence)
        self.out_dir = out_dir

    def compute(self):
        print "Sound file = ", self.file_name
        #can also  call oggread
        (snd, sampling_freq, n_bits) =  audiolab.wavread(self.file_name)
        n_samples = snd.shape[0]
        file_duration = float(n_samples) / float( sampling_freq )
        print "n_samples = ", n_samples, " file_duration = ", file_duration, " seconds"

        #The number of samples in a file - assuming all files have the same length = self.period
        letter_samples = np.int32(self.period* float( sampling_freq ))
        #The sample number to start from
        n_start = np.int32( self.start * float(sampling_freq))
        print "letter_duration = ", letter_samples, " m_start = ", n_start
        p.figure()
        for n in range(0, len(self.seq)):
            letter = self.seq[n]
            filename =  self.out_dir + letter+ ".ogg"
            amplitude = snd[n_start:n_start+letter_samples]
            audiolab.oggwrite(amplitude, filename, fs=sampling_freq, enc='vorbis')
            n_start += letter_samples
            #Display
            p.subplot(7,4,n+1)
            t =  1000.0 * np.arange(0.0, letter_samples , 1.0) / sampling_freq
            print "len t = ", len(t), " letter_samples = ", letter_samples, " lend(Data) = ", len(amplitude)
            p.plot(t, amplitude, color='k')
            max_y = max(amplitude)
            min_y = min(amplitude)
            p.plot([0.5*t[-1], 0.5*t[-1]], [min_y, max_y],'r', linewidth=3)
            p.ylabel(letter)
            #t += i_time_offset
            #n_end -=letter_samples
            #p.ylabel('Amplitude')
            #p.xlabel('Time (s)')
        p.show()

if __name__ ==  "__main__":
    #print "Run letter_offset_display for some examples, and saving of all end of alphabet sounds"
    c = ChannelConfig(i_nchannels=5,i_sound_overlap=0.65, i_file_length=0.21, i_root_dir='../../')
    timing_info = c.getClickTimes()
    alphabet_loader = c.getAlphabetLoader()
    alphabet = alphabet_loader.getAlphabet(i_with_spaces=True, i_group=True)

    sound_dir  = "../alphabet_fast/channels5/"
    print "sound dir = ", sound_dir
    print "timing_info = ", timing_info
    print "alphabet = ", alphabet

    a = BatchSoundAnalyser()
    a.displayChannelSounds( sound_dir ,  alphabet,  timing_info, i_plot_axis=True)

    #alphabet = ['a','b','c','d','e','f','g','h','i','j','k', 'l', 'm','n','o','p','q','r','s','t','u','v','w','x','y','z','_','.']
    #extractor = BatchExtractor(alphabet, file_name, output_dir,  i_start=0.77, i_period=0.5)
    #extractor.compute()


"""
    This is the main control file, used to generate synthetic data or to model
    real (or synthetic) datasets.
"""

import numpy as np

# locators
output_dir = 'storage/'
template_dir = output_dir+'obs_templates/'
storage_dir = output_dir+'synth_storage/'
reduced_dir = output_dir+'data/'
casalogs_dir = output_dir+'CASA_logs/'
simobs_dir = '/pool/asha0/casa-release-5.7.2-4.el7/data/alma/simmos/'

# naming
basename = 'simple3-default'
in_MS = storage_dir+basename+'/'+basename
dataname = reduced_dir+basename+'/'+basename

# observation settings
template = ['lmm']
config = ['6']
ttotal = ['20min']
tinteg = ['30s']
date = ['2022/07/15']
HA_0 = ['-0.5h']
RMS = [3.7]             # desired naturally-weighted RMS in mJy/beam/channel

# spectral settings
dnu_native = [122.0703125 * 1e3]	# native channel spacing (Hz)
nu_rest = 230.538e9	# rest frequency (Hz)
V_tune  = [4.0e3]	# LSRK velocity tuning for central channel (m/s)
V_span  = [15.0e3]	# +/- velocity width around vtune for simulation (m/s)
nover = 5	# spectral over-sampling factor (for SRF convolution)

# spatial settings
RA   = '16:00:00.00'	# phase center RA
DEC  = '-40:00:00.00'	# phase center DEC
RA_pieces = [np.float(RA.split(':')[i]) for i in np.arange(3)]
RAdeg = 15 * np.sum(np.array(RA_pieces) / [1., 60., 3600.])
DEC_pieces = [np.float(DEC.split(':')[i]) for i in np.arange(3)]
DECdeg = np.sum(np.array(DEC_pieces) / [1., 60., 3600.])

# reduction settings
tavg = ['']
V_bounds = [(5.2 - 10)*1e3, (5.2 + 10)*1e3]
bounds_pad = 3

# Model parameters
incl  = 40.
PA    = 130.
mstar = 2.0
r_l   = 400.
z0    = 1.9
psi   = 1.
T0    = 180.
q     = -0.5
Tmaxb = 20.
sigV0 = 326.
tau0  = 715.
p     = -1.
Vsys  = 0.0e3	#5.2e3
dx    = 0.
dy    = 0.
pars  = np.array([incl, PA, mstar, r_l, z0, psi, T0, q, Tmaxb,
                  sigV0, tau0, p, Vsys, dx, dy])

# Fixed parameters
FOV  = [12.8]		# full field of view (arcsec)
Npix = [512]  		# number of pixels per FOV
dist = 150.		# distance (pc)

cfg_dict = {}


import base64
import io
import logging
import re
import subprocess
import time
import uuid
import tempfile

import librosa
from ml_serving.utils import helpers
import numpy as np
import soundfile as sf

from synthesizer.inference import Synthesizer
from encoder import inference as enc_inference
from vocoder import inference as voc_inference

LOG = logging.getLogger(__name__)
PARAMS = {
    'vocoder': None,
    'encoder': None,
    'synthesizer': None,
}
SR_REGEX = re.compile('([0-9]+) Hz')

class CachedVector:
    def __init__(self, vector, time):
        self.vector = vector
        self.time = time

class AudioGen:
    def __init__(self, encoder, vocoder, synthesizer):
        self.synthesizer = Synthesizer(synthesizer)
        enc_inference.load_model(encoder)
        voc_inference.load_model(vocoder)

        self.speaker_embedding_size = 256
        self.keep_cache_sec = 86400
        self.cache = {}

        self.test()

    def test(self):
        print("\tTesting the encoder...")
        enc_inference.embed_utterance(np.zeros(enc_inference.sampling_rate))

        embed = np.random.rand(self.speaker_embedding_size)
        embed /= np.linalg.norm(embed)
        embeds = [embed, np.zeros(self.speaker_embedding_size)]
        texts = ["test 1", "test 2"]
        print("\tTesting the synthesizer... (loading the model will output a lot of text)")
        mels = self.synthesizer.synthesize_spectrograms(texts, embeds)

        mel = np.concatenate(mels, axis=1)
        no_action = lambda *args: None
        print("\tTesting the vocoder...")
        voc_inference.infer_waveform(mel, target=200, overlap=50, progress_callback=no_action)
        print("All test passed! You can now synthesize speech.\n\n")

    def _save_tmp(self, bytes_):
        fname = tempfile.mktemp(suffix='.wav')
        with open(fname, 'wb') as f:
            f.write(bytes_)

        return fname

    def _get_sample_rate(self, fname):
        pipe = subprocess.Popen(
            ['ffprobe', '-hide_banner', fname],
            stderr=subprocess.PIPE,
        )
        output = pipe.stderr.read().decode()
        pipe.communicate()
        groups = re.findall(SR_REGEX, output)
        if len(groups) < 1:
            raise RuntimeError('Can not estimate sample rate for file. Please ensure that file is correct.')

        return int(groups[0])

    def get_embedding(self, speech_example_bytes):
        fname = self._save_tmp(speech_example_bytes)
        sample_rate = self._get_sample_rate(fname)

        original_wav, sampling_rate = librosa.load(fname, sr=sample_rate)
        preprocessed_wav = enc_inference.preprocess_wav(original_wav, sampling_rate, trim_silence=True)

        embed = enc_inference.embed_utterance(preprocessed_wav)
        return embed

    def synthesize(self, embed, text):
        embeds = [embed]
        texts = [text]
        # If you know what the attention layer alignments are, you can retrieve them here by
        # passing return_alignments=True
        specs = self.synthesizer.synthesize_spectrograms(texts, embeds)
        spec = specs[0]

        # Synthesizing the waveform is fairly straightforward. Remember that the longer the
        # spectrogram, the more time-efficient the vocoder.
        def progress(i, seq_len, b_size, gen_rate):
            pass

        generated_wav = voc_inference.infer_waveform(spec, progress_callback=progress)

        generated_wav = np.pad(generated_wav, (0, self.synthesizer.sample_rate), mode="constant")

        # Trim excess silences to compensate for gaps in spectrograms (issue #53)
        generated_wav = enc_inference.preprocess_wav(generated_wav)

        audio_bytes = io.BytesIO()
        # sf.write('test.wav', generated_wav.astype(np.float32), self.synthesizer.sample_rate)
        sf.write(
            audio_bytes,
            generated_wav.astype(np.float32),
            self.synthesizer.sample_rate,
            format='wav'
        )
        return audio_bytes.getvalue()

    def generate_speech(self, speech_example_bytes, text):
        embed = self.get_embedding(speech_example_bytes)

        return self.synthesize(embed, text)

    def cache_vector(self, vector):
        vector_id = str(uuid.uuid4())
        LOG.info(f'[Cache] New vector ID={vector_id}')
        self.cache[vector_id] = CachedVector(vector, time.time())
        self._invalidate_cache()
        return vector_id

    def get_cached_vector(self, key):
        result = self.cache[key]
        self._invalidate_cache()
        return result.vector

    def _invalidate_cache(self):
        keys = list(self.cache.keys())
        now = time.time()
        for key in keys:
            if now - self.cache[key].time >= self.keep_cache_sec:
                LOG.info(f'[Cache] Expired vector ID={key}')
                del self.cache[key]

def init_hook(ctx, **params):
    PARAMS.update(params)
    audio_gen = AudioGen(
        encoder=PARAMS['encoder'],
        vocoder=PARAMS['vocoder'],
        synthesizer=PARAMS['synthesizer'],
    )
    return audio_gen

def process(inputs, ctx, **kwargs):
    audio_gen: AudioGen = ctx.global_ctx

    speech_id = helpers.get_param(inputs, 'speech_id')
    generate_speech = helpers.boolean_string(
        helpers.get_param(inputs, 'generate_speech', default=True)
    )
    audio_bytes = helpers.get_param(inputs, 'audio', raw_bytes=True)
    if isinstance(audio_bytes, str):
        audio_bytes = base64.decodebytes(audio_bytes.encode())
    text = helpers.get_param(inputs, 'text')

    if audio_bytes is not None and len(audio_bytes) > 0:
        vector = audio_gen.get_embedding(audio_bytes)
        speech_id = audio_gen.cache_vector(vector)
    else:
        vector = audio_gen.get_cached_vector(speech_id)

    result = {'speech_id': speech_id}

    # Generate speech from speech_id
    if generate_speech:
        generated_speech = audio_gen.synthesize(vector, text)
        result['audio'] = generated_speech

    return result


from cvzone.HandTrackingModule import HandDetector
import cv2
import numpy as np
import pyautogui
import cvzone

cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
hand_detector = HandDetector(maxHands=1)
cvw, cvh = 640, 480
sw, sh = pyautogui.size()
frame_cut = 0.2
smooth = 7
cx, cy, px, py = 0, 0, 0, 0
min_area, max_area = 5000, 11000
cap.set(3, cvw)
cap.set(4, cvh)

while True:
    success, img = cap.read()
    img = hand_detector.findHands(img, draw=True)
    lmlist, bboxs = hand_detector.findPosition(img, draw=True)

    if bboxs:
        bbox = bboxs['bbox']
        cv2.rectangle(img, (int(frame_cut * cvw), int(frame_cut * cvh)), (int(cvw - (frame_cut * cvw)), int(cvh - (frame_cut * cvh))), (255, 0, 0), 5)
        fingers = hand_detector.fingersUp()
        total_finger = fingers.count(1)
        index_finger = fingers[1]
        if index_finger == 1:
            x, y = lmlist[8][0], lmlist[8][1]
            x, y = np.interp(x, (0 + frame_cut * cvw, cvw - frame_cut * cvw), (0, sw)), np.interp(y, (0 + frame_cut * cvh, cvh - frame_cut * cvh), (0, sh))
            cx = px + (x - px) / smooth
            cy = py + (y - py) / smooth
            try:
                pyautogui.moveTo(int(cx), int(cy))
            except:
                pass
            distance, img, info = hand_detector.findDistance(8, 12, img, draw=True)
            area = bbox[2] * bbox[3]
            if distance <= 20:
                pyautogui.click(int(cx), int(cy), button='left')
            px, py = cx, cy

    cv2.imshow("Video", img)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break


import tensorflow as tf
import tensorflow_hub as hub
import numpy as np
import pyautogui
import win32api, win32con, win32gui
import cv2
import math
import time

detector = hub.load("https://tfhub.dev/tensorflow/centernet/resnet50v1_fpn_512x512/1")
size_scale = 3

while True:
    # Get rect of Window
    hwnd = win32gui.FindWindow(None, 'Counter-Strike: Global Offensive')
    #hwnd = win32gui.FindWindow("UnrealWindow", None) # Fortnite
    rect = win32gui.GetWindowRect(hwnd)
    region = rect[0], rect[1], rect[2] - rect[0], rect[3] - rect[1]

    # Get image of screen
    ori_img = np.array(pyautogui.screenshot(region=region))
    ori_img = cv2.resize(ori_img, (ori_img.shape[1] // size_scale, ori_img.shape[0] // size_scale))
    image = np.expand_dims(ori_img, 0)
    img_w, img_h = image.shape[2], image.shape[1]

    # Detection
    result = detector(image)
    result = {key:value.numpy() for key,value in result.items()}
    boxes = result['detection_boxes'][0]
    scores = result['detection_scores'][0]
    classes = result['detection_classes'][0]

    # Check every detected object
    detected_boxes = []
    for i, box in enumerate(boxes):
        # Choose only person(class:1)
        if classes[i] == 1 and scores[i] >= 0.5:
            ymin, xmin, ymax, xmax = tuple(box)
            if ymin > 0.5 and ymax > 0.8: # CS:Go
            #if int(xmin * img_w * 3) < 450: # Fortnite
                continue
            left, right, top, bottom = int(xmin * img_w), int(xmax * img_w), int(ymin * img_h), int(ymax * img_h)
            detected_boxes.append((left, right, top, bottom))
            #cv2.rectangle(ori_img, (left, top), (right, bottom), (255, 255, 0), 2)

    print("Detected:", len(detected_boxes))

    # Check Closest
    if len(detected_boxes) >= 1:
        min = 99999
        at = 0
        centers = []
        for i, box in enumerate(detected_boxes):
            x1, x2, y1, y2 = box
            c_x = ((x2 - x1) / 2) + x1
            c_y = ((y2 - y1) / 2) + y1
            centers.append((c_x, c_y))
            dist = math.sqrt(math.pow(img_w/2 - c_x, 2) + math.pow(img_h/2 - c_y, 2))
            if dist < min:
                min = dist
                at = i

        # Pixel difference between crosshair(center) and the closest object
        x = centers[at][0] - img_w/2
        y = centers[at][1] - img_h/2 - (detected_boxes[at][3] - detected_boxes[at][2]) * 0.45

        # Move mouse and shoot
        scale = 1.7 * size_scale
        x = int(x * scale)
        y = int(y * scale)
        win32api.mouse_event(win32con.MOUSEEVENTF_MOVE, x, y, 0, 0)
        time.sleep(0.05)
        win32api.mouse_event(win32con.MOUSEEVENTF_LEFTDOWN, x, y, 0, 0)
        time.sleep(0.1)
        win32api.mouse_event(win32con.MOUSEEVENTF_LEFTUP, x, y, 0, 0)

    #ori_img = cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB)
    #cv2.imshow("ori_img", ori_img)
    #cv2.waitKey(1)

    time.sleep(0.1)


import logging
import numpy as np
import math
from matplotlib import pyplot
import time
from queue import Queue
from datetime import datetime, timedelta

from .input_handler import InputHandler
from .utils import get_config
from .grabscreen import grab_screen
import cv2

OWN_INVENTORY_ORIGIN = (0.6769531, 0.567361)
TRADE_PARTNER_ORIGIN = (0.17578, 0.2118)
INVITATION_TIMEOUT = 30
TRADE_TIMEOUT = 16
TRADE_RETRIES = 2
TRADE_TIMEOUT_AFTER_INITIATED = 30
RETURN_TRADE_TIMEOUT = 15

class Trader:
    def __init__(self, resolution, accept_trades):
        self.resolution = resolution
        self.config = get_config("trader")
        self.input_handler = InputHandler(self.resolution)
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s %(message)s",
            datefmt="[%H:%M:%S %d-%m-%Y]",
        )
        self.log = logging.getLogger("trader")
        if accept_trades:
            self.trade_queue = Queue()
            self.invite_message = self.config["invite_message"]
            self.trade_queue_set = set()
            with open(self.config["log_path"], "r", encoding="utf8") as f:
                for i, l in enumerate(f):
                    pass
                byte_pos = f.tell()
            self.byte_pos = byte_pos

    def stash_items(self):
        stash_loc = self._find_stash()
        self.log.info("Found stash at %s" % stash_loc)
        stash_loc_br = [stash_loc[0] + 0.005, stash_loc[1] + 0.005]
        self.input_handler.click(*stash_loc, *stash_loc_br)
        self.input_handler.rnd_sleep(min=1200, mean=1500)
        slots = self.find_nonempty_inventory_slots(origin=OWN_INVENTORY_ORIGIN)
        item_locations = np.argwhere(slots == 1)
        for loc in item_locations:
            self.input_handler.inventory_click(
                loc[0], loc[1], inventory_base=OWN_INVENTORY_ORIGIN, ctrl_click=True
            )
        self.input_handler.click_keys([0x81])

        items_stashed = len(item_locations)

        if items_stashed > 0:
            self.log.info("Stashed %s items" % items_stashed)
        else:
            self.log.info("No stashing needed, inventory empty")

    def verify_empty_inventory(self):
        self.input_handler.click_hotkey("i")
        slots = self.find_nonempty_inventory_slots(origin=OWN_INVENTORY_ORIGIN)
        self.input_handler.click_hotkey("i")
        non_empty = int(np.sum(slots))
        self.log.info("Non-empty slots found: %d" % non_empty)
        return non_empty == 0

    def wait_for_trade(self):
        self.log.info("Waiting for trading partner")
        trade_partner_in_area = False
        while not trade_partner_in_area:
            while self.trade_queue.empty():
                messages = self.tail(set_seeker=True)
                for message in messages:
                    if (
                        message["text"] == self.invite_message
                        and message["type"] == "whisper"
                    ):
                        self._put_in_trade_queue(message["user"])
                time.sleep(0.1)
            trade_partner = self._get_from_trade_queue()
            trade_partner_in_area = self._invite_and_wait_for_player(trade_partner)
        return trade_partner

    def _invite_and_wait_for_player(self, username):
        self.input_handler.type("/invite %s" % username)
        timeout_at = datetime.now() + timedelta(seconds=INVITATION_TIMEOUT)
        player_in_ho = False
        sought_message = username + " has joined the area."
        while datetime.now() < timeout_at:
            messages = self.tail(max_age=timedelta(seconds=10))
            for message in messages:
                if message["type"] == "info" and message["text"] == sought_message:
                    return True
            time.sleep(0.1)
        return False

    def _put_in_trade_queue(self, user):
        if user not in self.trade_queue_set:
            self.trade_queue_set.add(user)
            self.trade_queue.put(user)

    def _get_from_trade_queue(self):
        user = self.trade_queue.get()
        self.trade_queue_set.remove(user)
        return user

    def tail(self, set_seeker=False, max_age=None):
        file_handle = open(self.config["log_path"], "r", encoding="utf8")
        file_handle.seek(self.byte_pos)
        lines = []
        while True:
            line = file_handle.readline()
            if line:
                lines.append(line[:-1])
            else:
                break
        if set_seeker:
            self.byte_pos = file_handle.tell()
        file_handle.close()
        if max_age is None:
            messages = [self._split_chat_msg(line) for line in lines]
        else:
            messages = [
                self._split_chat_msg(line)
                for line in lines
                if self._split_chat_msg(line)["time"] > datetime.now() - max_age
            ]
        if messages:
            self.log.info("interpreted messages: %s" % messages)
        return messages

    def get_items(self, username):
        self.log.info("Trying to open trade with %s" % username)
        initiated = self._initiate_trade(username)
        if not initiated:
            self.log.info("Failed to open trade with %s" % username)
            return None
        self.log.info("Opened trade with %s" % username)
        trade_successful = self._get_partners_trade_items()
        return trade_successful

    def get_jewel_locations(self):
        self.input_handler.click_hotkey("i")
        slots = self.find_nonempty_inventory_slots(origin=OWN_INVENTORY_ORIGIN)
        item_locations = np.argwhere(slots == 1)
        jewel_descriptions = []
        for item_location in item_locations:
            item = self.input_handler.inventory_copy(
                item_location[0], item_location[1], OWN_INVENTORY_ORIGIN, speed_factor=3
            )
            item_type = item.split("\n")[3].strip()

            if item_type != "Timeless Jewel":
                slots[item_location[0], item_location[1]] = 0
                self.log.info("Received item: %s" % item)
            else:
                item_desc = item.split("\n")[10].strip()
                self.log.info("Received jewel: %s" % item_desc)
                jewel_descriptions.append(item_desc)
        jewel_locations = np.argwhere(slots == 1)
        self.input_handler.click_hotkey("i")
        return jewel_locations, jewel_descriptions

    def return_items(self, username, locations):
        self.log.info("Trying to do a return trade with %s" % username)
        initiated = self._initiate_trade(username)
        if not initiated:
            self.log.info("Failed to do a return trade with %s" % username)
            return None
        self.log.info("Opened return trade with %s" % username)
        trade_successful = self._return_partners_trade_items(locations)
        return trade_successful

    def _return_partners_trade_items(self, locations):
        for location in locations:
            self.input_handler.inventory_click(
                location[0],
                location[1],
                OWN_INVENTORY_ORIGIN,
                ctrl_click=True,
                speed_factor=2,
            )
        self._accept_trade()
        timeout = datetime.now() + timedelta(seconds=RETURN_TRADE_TIMEOUT)
        sought_message = "Trade accepted."
        while datetime.now() < timeout:
            messages = self.tail(max_age=timedelta(seconds=15))
            for message in messages:
                if message["type"] == "info" and message["text"] == sought_message:
                    return True
        self.input_handler.click_keys([0x81])  # Escape trade screen
        return False

    def _initiate_trade(self, username):
        for try_idx in range(TRADE_RETRIES):
            self.input_handler.type("/tradewith %s" % username)
            timeout = datetime.now() + timedelta(seconds=TRADE_TIMEOUT)
            while datetime.now() < timeout:
                trade_initiated = self._trade_window_open()
                if trade_initiated:
                    return True
                time.sleep(0.1)
        return False

    def _trade_window_open(self):
        img_bgr = grab_screen((0, 0, *self.resolution))
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGRA2RGB)
        target_color = (183, 125, 66)
        x = int(0.29765 * self.resolution[0])
        y = int(0.08889 * self.resolution[1])
        actual_color = img_rgb[y, x]
        return self._close_colors(target_color, actual_color, max_distance=25)

    def _get_partners_trade_items(self):
        timeout = datetime.now() + timedelta(seconds=TRADE_TIMEOUT_AFTER_INITIATED)
        while datetime.now() < timeout:
            self._hover_trade_partner()
            self.input_handler.inventory_click(
                -1, -1, inventory_base=TRADE_PARTNER_ORIGIN
            )
            if self._partner_has_accepted():
                self._hover_trade_partner()
                self._accept_trade()
                self.input_handler.rnd_sleep()
                return True
        return False

    def _hover_trade_partner(self):
        slots = self.find_nonempty_inventory_slots(origin=TRADE_PARTNER_ORIGIN)
        item_locations = np.argwhere(slots == 1)
        for loc in item_locations:
            self.input_handler.inventory_click(
                loc[0],
                loc[1],
                inventory_base=TRADE_PARTNER_ORIGIN,
                button=None,
                speed_factor=3,
            )

    def _accept_trade(self):
        self.input_handler.click(0.18086, 0.76041, 0.19148, 0.7607)

    def _partner_has_accepted(self):
        img_bgr = grab_screen((0, 0, *self.resolution))
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGRA2RGB)
        target_color = (49, 62, 42)
        x = int(0.32617 * self.resolution[0])
        y = int(0.18264 * self.resolution[1])
        actual_color = img_rgb[y, x]
        return self._close_colors(target_color, actual_color)

    def find_nonempty_inventory_slots(self, origin):
        slots = np.zeros((12, 5))
        img_bgr = grab_screen((0, 0, *self.resolution))
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGRA2RGB)
        target_color = (0, 0, 0)
        for x_idx in range(12):
            for y_idx in range(5):
                slot_color = self._color_at_slot(img_rgb, x_idx, y_idx, origin)
                empty = self._close_colors(target_color, slot_color)
                if not empty:
                    self.log.info("Slot (%d, %d) was not empty" % (x_idx, y_idx))
                    slots[x_idx, y_idx] = 1
        self.input_handler.rnd_sleep(min=300, mean=500)
        return slots

    def _color_at_slot(self, img, slot_x, slot_y, origin):
        x = int((origin[0] + slot_x * 0.02735) * self.resolution[0])
        y = int((origin[1] + slot_y * 0.04930555) * self.resolution[1])
        return img[y, x]

    def _close_colors(self, target_color, actual_color, max_distance=16):
        self.log.debug(actual_color)
        distance = math.sqrt(
            (actual_color[0] - target_color[0]) ** 2
            + (actual_color[1] - target_color[1]) ** 2
            + (actual_color[2] - target_color[2]) ** 2
        )
        self.log.debug("Measured color distance of: %s" % distance)
        return distance < max_distance

    def _find_stash(self):
        img_rgb = grab_screen((0, 0, *self.resolution))
        img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGBA2GRAY)
        template = cv2.imread("data/images/stash.png", 0)
        res = cv2.matchTemplate(img_gray, template, cv2.TM_CCOEFF_NORMED)
        threshold = 0.8
        stash_corner = np.where(res >= threshold)
        if len(stash_corner[0]) > 0:
            return [
                stash_corner[1][0] / self.resolution[0],
                stash_corner[0][0] / self.resolution[1],
            ]
        return None

    def _split_chat_msg(self, line):
        msg = {}
        if len(line) < 30:
            msg["type"] = "other"
            return msg
        time = line[:20]
        ms = line[26:29] + "000"
        full_time = time + ms
        msg["time"] = datetime.strptime(full_time, "%Y/%m/%d %H:%M:%S %f")
        last_info_idx = line.find("]") + 2
        content = line[last_info_idx:]
        if content[0] == ":":
            msg["type"] = "info"
        elif content[0] == "#":
            msg["type"] = "global"
        elif content[0] == "%":
            msg["type"] = "party"
        elif content[0] == "@":
            msg["type"] = "whisper"
        else:
            msg["type"] = "other"

        msg["user"] = None
        msg["text"] = None
        if msg["type"] in ["global", "party", "whisper"]:
            name_start = last_info_idx + len(content.split(" ")[0])
            name_end = name_start + line[name_start:].find(":")
            name = line[name_start:name_end].lower()
            if ">" in name:
                name = name[name.find(">") + 2 :]
            msg["user"] = name
            msg["text"] = line[name_end + 2 :]
        elif msg["type"] == "info":
            msg["text"] = content[2:]
        return msg


import cv2
import io
import numpy as np
import os
recognizer = cv2.face.LBPHFaceRecognizer_create()
recognizer.read('trainer/trainer.yml')
cascadePath = "haarcascade_frontalface_default.xml"
faceCascade = cv2.CascadeClassifier(cascadePath);
font = cv2.FONT_HERSHEY_SIMPLEX
#iniciate id counter
id = 0
# names related to ids: example ==> Marcelo: id=1,  etc
names = ['None', 'Vishal', 'Muthuraja', 'Kishore', 'Prem', 'Aditya']
# Initialize and start realtime video capture
with picamera.PiCamera() as cam:
    cam.resolution = (640,480) # set video widht
    cam.capture(io.BytesIO, format='jpeg')

buff = numpy.fromstring(stream.getvalue(), dtype=numpy.uint8)
cam = cv2.imdecode(buff,1)

# Define min window size to be recognized as a face
minW = 0.1*cam.get(3)
minH = 0.1*cam.get(4)
while True:
    ret, img =cam.read()
#    img = cv2.flip(img, -1) # Flip vertically
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)

    faces = faceCascade.detectMultiScale(
        gray,
        scaleFactor = 1.2,
        minNeighbors = 5,
        minSize = (int(minW), int(minH)),
       )
    for(x,y,w,h) in faces:
        cv2.rectangle(img, (x,y), (x+w,y+h), (0,255,0), 2)
        id, confidence = recognizer.predict(gray[y:y+h,x:x+w])
        # Check if confidence is less them 100 ==> "0" is perfect match
        if (confidence < 100):
            print(id)
            id = names[id]
            confidence = "  {0}%".format(round(100 - confidence))
        else:
            id = "unknown"
            confidence = "  {0}%".format(round(100 - confidence))

        cv2.putText(img, str(id), (x+5,y-5), font, 1, (255,255,255), 2)
        cv2.putText(img, str(confidence), (x+5,y+h-5), font, 1, (255,255,0), 1)

    cv2.imshow('camera',img)
    k = cv2.waitKey(10) & 0xff # Press 'ESC' for exiting video
    if k == 27:
        break
# Do a bit of cleanup
print("\n [INFO] Exiting Program and cleanup stuff")
cam.release()
cv2.destroyAllWindows()