#!/usr/bin/env python
#
# See top-level LICENSE file for Copyright information
#
# -*- coding: utf-8 -*-
"""
This script enables the user to convert spec2D FITS files
from IFU instruments into a 3D cube with a defined WCS.
"""

import argparse

from astropy import units
from astropy.io import fits
from astropy.coordinates import SkyCoord
from astropy.wcs import WCS
import numpy as np
import copy, os

from pypeit import msgs, par, io, spec2dobj
from pypeit.spectrographs.util import load_spectrograph
from pypeit.core import datacube as dc_utils
from pypeit.core.flux_calib import load_extinction_data, extinction_correction
from pypeit.core.flexure import calculate_image_offset
from pypeit.core import parse

from IPython import embed

def parse_args(options=None, return_parser=False):

    parser = argparse.ArgumentParser(description='Read in an array of spec2D files and convert them into a datacube',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('file', type = str, default=None, help='filename.coadd3d file')
    parser.add_argument('--det', default=1, type=int, help="Detector")
    parser.add_argument('-o', '--overwrite', default=False, action='store_true',
                        help='Overwrite any existing files/directories')

    if return_parser:
        return parser

    return parser.parse_args() if options is None else parser.parse_args(options)

def coadd_cube(files, parset, overwrite=False):
    """ Main routine to coadd spec2D files into a 3D datacube

    Args:
        files (list):
            List of all spec2D files
        parset (:class:`pypeit.par.core.PypeItPar`):
            An instance of the parameter set.
        overwrite (bool):
            Overwrite the output file, if it exists?
    """
    # Get the detector number
    det = 1 if parset is None else parset['rdx']['detnum']

    # Load the spectrograph
    spec2DObj = spec2dobj.Spec2DObj.from_file(files[0], det)
    specname = spec2DObj.head0['PYP_SPEC']
    spec = load_spectrograph(specname)

    # Grab the parset, if not provided
    if parset is None: parset = spec.default_pypeit_par()
    cubepar = parset['reduce']['cube']

    # Check the output file
    outfile = cubepar['output_filename'] if ".fits" in cubepar['output_filename'] else cubepar['output_filename']+".fits"
    out_whitelight = outfile.replace(".fits", "_whitelight.fits")
    if os.path.exists(outfile) and not overwrite:
        msgs.error("Output filename already exists:"+msgs.newline()+outfile)
    elif os.path.exists(out_whitelight) and cubepar['save_whitelight'] and not overwrite:
        msgs.error("Output filename already exists:"+msgs.newline()+out_whitelight)
    # Check the reference cube and image exist, if requested
    ref_scale = None  # This will be used to correct relative scaling among the various input frames
    if cubepar['standard_cube'] is not None:
        if not os.path.exists(cubepar['standard_cube']):
            msgs.error("Standard cube does not exist:" + msgs.newline() + cubepar['reference_cube'])
        cube = fits.open(cubepar['standard_cube'])
        ref_scale = cube['REFSCALE'].data
    if cubepar['reference_image'] is not None:
        if not os.path.exists(cubepar['reference_image']):
            msgs.error("Reference cube does not exist:" + msgs.newline() + cubepar['reference_image'])
    if cubepar['flux_calibrate']:
        msgs.error("Flux calibration is not currently implemented" + msgs.newline() +
                   "Please set 'flux_calibrate = False'")

    # prep
    numfiles = len(files)
    combine = cubepar['combine']

    all_ra, all_dec, all_wave = np.array([]), np.array([]), np.array([])
    all_sci, all_ivar, all_idx, all_wghts = np.array([]), np.array([]), np.array([]), np.array([])
    all_wcs = []
    dspat = None if cubepar['spatial_delta'] is None else  cubepar['spatial_delta']/3600.0  # binning size on the sky (/3600 to convert to degrees)
    dwv = cubepar['wave_delta']       # binning size in wavelength direction (in Angstroms)
    wave_ref = None
    whitelight_img = None  # This is the whitelight image based on all input spec2d frames
    weights = np.ones(numfiles)  # Weights to use when combining cubes
    for ff, fil in enumerate(files):
        # Load it up
        spec2DObj = spec2dobj.Spec2DObj.from_file(fil, det)
        detector = spec2DObj.detector

        # Setup for PypeIt imports
        msgs.reset(verbosity=2)

        if ref_scale is None:
            ref_scale = spec2DObj.scaleimg.copy()
        # Extract the information
        sciimg = (spec2DObj.sciimg-spec2DObj.skymodel) * (ref_scale/spec2DObj.scaleimg)  # Subtract sky and apply relative sky
        ivar = spec2DObj.ivarraw / (ref_scale/spec2DObj.scaleimg)**2
        waveimg = spec2DObj.waveimg
        bpmmask = spec2DObj.bpmmask

        # Grab the slit edges
        slits = spec2DObj.slits

        wave0 = waveimg[waveimg != 0.0].min()
        diff = waveimg[1:, :] - waveimg[:-1, :]
        dwv = float(np.median(diff[diff != 0.0]))
        msgs.info("Using wavelength solution: wave0={0:.3f}, dispersion={1:.3f} Angstrom/pixel".format(wave0, dwv))

        msgs.info("Constructing slit image")
        slitid_img_init = slits.slit_img(pad=0, initial=True, flexure=spec2DObj.sci_spat_flexure)
        onslit_gpm = (slitid_img_init > 0) & (bpmmask == 0)

        # Grab the WCS of this frame
        wcs = spec.get_wcs(spec2DObj.head0, slits, detector.platescale, wave0, dwv)
        all_wcs.append(copy.deepcopy(wcs))

        # Find the largest spatial scale of all images being combined
        # TODO :: probably need to put this in the DetectorContainer
        pxscl = detector.platescale * parse.parse_binning(detector.binning)[1] / 3600.0  # This should be degrees/pixel
        slscl = spec.get_meta_value([spec2DObj.head0], 'slitwid')
        if dspat is None:
            dspat = max(pxscl, slscl)
        elif max(pxscl, slscl) > dspat:
            dspat = max(pxscl, slscl)

        # Generate an RA/DEC image
        msgs.info("Generating RA/DEC image")
        raimg, decimg, minmax = slits.get_radec_image(wcs, initial=True, flexure=spec2DObj.sci_spat_flexure)

        # Perform the DAR correction
        if wave_ref is None:
            wave_ref = 0.5*(np.min(waveimg[onslit_gpm]) + np.max(waveimg[onslit_gpm]))
        # Get DAR parameters
        raval = spec.get_meta_value([spec2DObj.head0], 'ra')
        decval = spec.get_meta_value([spec2DObj.head0], 'dec')
        obstime = spec.get_meta_value([spec2DObj.head0], 'obstime')
        pressure = spec.get_meta_value([spec2DObj.head0], 'pressure')
        temperature = spec.get_meta_value([spec2DObj.head0], 'temperature')
        rel_humidity = spec.get_meta_value([spec2DObj.head0], 'humidity')
        coord = SkyCoord(raval, decval, unit=(units.deg, units.deg))
        location = spec.location  # TODO :: spec.location should probably end up in the TelescopePar (spec.telescope.location)
        ra_corr, dec_corr = dc_utils.dar_correction(waveimg[onslit_gpm], coord, obstime, location,
                                                    pressure, temperature, rel_humidity, wave_ref=wave_ref)
        raimg[onslit_gpm] += ra_corr
        decimg[onslit_gpm] += dec_corr

        # Get copies of arrays to be saved
        wave_ext = waveimg[onslit_gpm].copy()
        flux_ext = sciimg[onslit_gpm].copy()
        ivar_ext = ivar[onslit_gpm].copy()

        # Perform extinction correction
        msgs.info("Applying extinction correction")
        longitude = spec.telescope['longitude']
        latitude = spec.telescope['latitude']
        airmass = spec2DObj.head0[spec.meta['airmass']['card']]
        extinct = load_extinction_data(longitude, latitude)
        # extinction_correction requires the wavelength is sorted
        wvsrt = np.argsort(wave_ext)
        ext_corr = extinction_correction(wave_ext[wvsrt] * units.AA, airmass, extinct)
        # Correct for extinction
        flux_sav = flux_ext[wvsrt] * ext_corr
        ivar_sav = ivar_ext[wvsrt] / ext_corr ** 2
        # sort back to the original ordering
        resrt = np.argsort(wvsrt)

        # Calculate the weights relative to the zeroth cube
        if ff != 0:
            weights[ff] = np.median(flux_sav[resrt]*np.sqrt(ivar_sav[resrt]))**2

        # Store the information
        numpix = raimg[onslit_gpm].size
        all_ra = np.append(all_ra, raimg[onslit_gpm].copy())
        all_dec = np.append(all_dec, decimg[onslit_gpm].copy())
        all_wave = np.append(all_wave, wave_ext.copy())
        all_sci = np.append(all_sci, flux_sav[resrt].copy())
        all_ivar = np.append(all_ivar, ivar_sav[resrt].copy())
        all_idx = np.append(all_idx, ff*np.ones(numpix))
        all_wghts = np.append(all_wghts, weights[ff]*np.ones(numpix))

    # Grab cos(dec) for convenience
    cosdec = np.cos(np.mean(all_dec) * np.pi / 180.0)

    # Register spatial offsets between all frames if several frames are being combined
    if combine:

        # Check if a reference whitelight image should be used to register the offsets
        if cubepar["reference_image"] is None:
            # Generate white light images
            whitelight_imgs, _, _ = dc_utils.make_whitelight(all_ra, all_dec, all_wave, all_sci, all_wghts, all_idx,
                                                             dspat)
            # ref_idx will be the index of the cube with the highest S/N
            ref_idx = np.argmax(weights)
            reference_image = whitelight_imgs[:, :, ref_idx].copy()
            msgs.info("Calculating spatial translation of each cube relative to cube #{0:d})".format(ref_idx+1))
        else:
            ref_idx = -1  # Don't use an index
            # Load reference information
            reference_image, whitelight_imgs, wlwcs = \
                dc_utils.make_whitelight_fromref(all_ra, all_dec, all_wave, all_sci, all_wghts, all_idx, dspat,
                                                 cubepar['reference_image'])
            msgs.info("Calculating the spatial translation of each cube relative to user-defined 'reference_image'")
        # Calculate the image offsets - check the reference is a zero shift
        ra_shift_ref, dec_shift_ref = calculate_image_offset(reference_image.copy(), reference_image.copy())
        for ff in range(numfiles):
            # Don't correlate the reference image with itself
            if ff == ref_idx:
                continue
            # Calculate the shift
            ra_shift, dec_shift = calculate_image_offset(whitelight_imgs[:, :, ff], reference_image.copy())
            # Convert to reference
            ra_shift -= ra_shift_ref
            dec_shift -= dec_shift_ref
            # Convert pixel shift to degress shift
            ra_shift *= dspat/cosdec
            dec_shift *= dspat
            msgs.info("Spatial shift of cube #{0:d}: RA, DEC (arcsec) = {1:+0.3f}, {2:+0.3f}".format(ff+1, ra_shift*3600.0, dec_shift*3600.0))
            # Apply the shift
            all_ra[all_idx == ff] += ra_shift
            all_dec[all_idx == ff] += dec_shift

        # Generate a white light image of *all* data
        msgs.info("Generating global white light image")
        if cubepar["reference_image"] is None:
            whitelight_img, _, wlwcs = dc_utils.make_whitelight(all_ra, all_dec, all_wave, all_sci, all_wghts,
                                                                np.zeros(all_ra.size), dspat)
        else:
            _, whitelight_img, wlwcs = \
                dc_utils.make_whitelight_fromref(all_ra, all_dec, all_wave, all_sci, all_wghts, np.zeros(all_ra.size),
                                                 dspat, cubepar['reference_image'])

        # Calculate the relative spectral weights of all pixels
        all_wghts = dc_utils.compute_weights(all_ra, all_dec, all_wave, all_sci, all_ivar, all_idx,
                                             whitelight_img[:, :, 0], dspat, dwv,
                                             relative_weights=cubepar['relative_weights'])
    # Check if a whitelight image should be saved
    if cubepar['save_whitelight']:
        # Check if the white light image still needs to be generated - if so, generate it now
        if whitelight_img is None:
            msgs.info("Generating global white light image")
            if cubepar["reference_image"] is None:
                whitelight_img, _, wlwcs = dc_utils.make_whitelight(all_ra, all_dec, all_wave, all_sci, all_wghts,
                                                                    np.zeros(all_ra.size), dspat)
            else:
                _, whitelight_img, wlwcs = \
                    dc_utils.make_whitelight_fromref(all_ra, all_dec, all_wave, all_sci, all_wghts,
                                                     np.zeros(all_ra.size),
                                                     dspat, cubepar['reference_image'])
        # Prepare and save the fits file
        msgs.info("Saving white light image as: {0:s}".format(out_whitelight))
        img_hdu = fits.PrimaryHDU(whitelight_img.T, header=wlwcs.to_header())
        img_hdu.writeto(out_whitelight, overwrite=overwrite)

    # Setup the cube ranges
    ra_min = cubepar['ra_min'] if cubepar['ra_min'] is not None else np.min(all_ra)
    ra_max = cubepar['ra_max'] if cubepar['ra_max'] is not None else np.max(all_ra)
    dec_min = cubepar['dec_min'] if cubepar['dec_min'] is not None else np.min(all_dec)
    dec_max = cubepar['dec_max'] if cubepar['dec_max'] is not None else np.max(all_dec)
    wav_min = cubepar['wave_min'] if cubepar['wave_min'] is not None else np.min(all_wave)
    wav_max = cubepar['wave_max'] if cubepar['wave_max'] is not None else np.max(all_wave)
    if cubepar['wave_delta'] is not None: dwv = cubepar['wave_delta']
    # Generate a master WCS to register all frames
    coord_min = [ra_min, dec_min, wav_min]
    coord_dlt = [dspat, dspat, dwv]
    masterwcs = dc_utils.generate_masterWCS(coord_min, coord_dlt, name=specname)
    msgs.info(msgs.newline()+"-"*40 +
              msgs.newline() + "Parameters of the WCS:" +
              msgs.newline() + "RA   min, max = {0:f}, {1:f}".format(ra_min, ra_max) +
              msgs.newline() + "DEC  min, max = {0:f}, {1:f}".format(dec_min, dec_max) +
              msgs.newline() + "WAVE min, max = {0:f}, {1:f}".format(wav_min, wav_max) +
              msgs.newline() + "Spaxel size = {0:f}''".format(3600.0*dspat) +
              msgs.newline() + "Wavelength step = {0:f} A".format(dwv) +
              msgs.newline() + "-" * 40)

    # Generate the output binning
    if combine:
        numra = int((ra_max-ra_min) * cosdec / dspat)
        numdec = int((dec_max-dec_min)/dspat)
        numwav = int((wav_max-wav_min)/dwv)
        xbins = np.arange(1+numra)-0.5
        ybins = np.arange(1+numdec)-0.5
        spec_bins = np.arange(1+numwav)-0.5
    else:
        slitlength = int(np.round(np.median(slits.get_slitlengths(initial=True, median=True))))
        numwav = int((np.max(waveimg) - wave0) / dwv)
        xbins, ybins, spec_bins = spec.get_datacube_bins(slitlength, minmax, numwav)

    # Make the cube
    msgs.info("Generating pixel coordinates")
    if combine:
        pix_coord = masterwcs.wcs_world2pix(all_ra, all_dec, all_wave * 1.0E-10, 0)
        hdr = masterwcs.to_header()
    else:
        pix_coord = wcs.wcs_world2pix(np.vstack((all_ra, all_dec, all_wave*1.0E-10)).T, 0)
        hdr = wcs.to_header()

    # Find the NGP coordinates for all input pixels
    msgs.info("Generating data cube")
    bins = (xbins, ybins, spec_bins)
    datacube, edges = np.histogramdd(pix_coord, bins=bins, weights=all_sci*all_wghts)
    norm, edges = np.histogramdd(pix_coord, bins=bins, weights=all_wghts)
    norm_cube = (norm > 0) / (norm + (norm == 0))
    datacube *= norm_cube
    # Create the variance cube, including weights
    msgs.info("Generating variance cube")
    all_var = (all_ivar > 0) / (all_ivar + (all_ivar == 0))
    var_cube, edges = np.histogramdd(pix_coord, bins=bins, weights=all_var * all_wghts**2)
    var_cube *= norm_cube**2

    # Save the datacube
    debug = False
    if debug:
        datacube_resid, edges = np.histogramdd(pix_coord, bins=(xbins, ybins, spec_bins), weights=all_sci*np.sqrt(all_ivar))
        norm, edges = np.histogramdd(pix_coord, bins=(xbins, ybins, spec_bins))
        norm_cube = (norm > 0) / (norm + (norm == 0))
        outfile = "datacube_resid.fits"
        msgs.info("Saving datacube as: {0:s}".format(outfile))
        hdu = fits.PrimaryHDU((datacube_resid*norm_cube).T, header=masterwcs.to_header())
        hdu.writeto(outfile, overwrite=overwrite)

    msgs.info("Saving datacube as: {0:s}".format(outfile))
    final_cube = dc_utils.DataCube(datacube.T, var_cube.T, specname,
                                   refscale=ref_scale, fluxed=cubepar['flux_calibrate'])
    final_cube.to_file(outfile, hdr=hdr, overwrite=overwrite)

def main(args):
    if args.file is None:
        msgs.error('You must input a coadd3d file')
    else:
        spectrograph_name, config_lines, spec2d_files = io.read_spec2d_file(args.file, filetype="coadd3d")
        spectrograph = load_spectrograph(spectrograph_name)

        # Parameters
        spectrograph_def_par = spectrograph.default_pypeit_par()
        parset = par.PypeItPar.from_cfg_lines(cfg_lines=spectrograph_def_par.to_config(),
                                              merge_with=config_lines)
        # If detector was passed as an argument override whatever was in the coadd3d file
        if args.det is not None:
            msgs.info("Restricting to detector={}".format(args.det))
            parset['rdx']['detnum'] = int(args.det)

    # Coadd the files
    coadd_cube(spec2d_files, parset, overwrite=args.overwrite)


# Lint as: python2, python3
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Implementation for distributed Shampoo optimizer."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools
import lingvo.compat as tf
from lingvo.core import matrix_functions
from lingvo.core import ops as x_ops
import numpy as np

# pylint: disable=g-direct-tensorflow-import
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.training import optimizer
# pylint: enable=g-direct-tensorflow-import

class PartitionConfig(object):
  """Config for tensor partitioning."""

  def __init__(self, max_dim_size, partition_size):
    """Initialize the `PartitionConfig`.

    Args:
      max_dim_size: Partitions dimensions with size greater than this value.
      partition_size: Size of each partition
    """
    self.max_dim_size = max_dim_size
    self.partition_size = partition_size

class PartitionMetadata(object):
  """Metadata for partitioning."""

  def __init__(self, split_sizes_per_dim, num_splits_per_dim):
    """Initialize the `PartitionMetadata`.

    Args:
      split_sizes_per_dim: Split sizes per dimemsion.
      num_splits_per_dim: Number of splits per dimension ( inferred from
        split_sizes_per_dim).
    """
    self.split_sizes_per_dim = split_sizes_per_dim
    self.num_splits_per_dim = num_splits_per_dim

class TensorPartitioner(object):
  """Shards Tensor's across its axis.

  In cases of TPUs, these partitions are zero cost, and does not involve data
  movement.
  """

  @classmethod
  def partition_metadata(cls, tensor, partition_info):
    """Returns metadata required for partitioning and reforming tensors.

    Args:
      tensor: Tensor to partition.
      partition_info: Partitioning info.

    Returns:
      split_sizes_per_dim and num_splits_per_dim.
    """
    shape = tensor.get_shape()
    # Split if dim is greater than max_dim.
    axis_to_shard = [s > partition_info.max_dim_size for s in shape]
    split_sizes_per_dim = []
    # Compute the number of splits, and the sizes of the splits for
    # each dimension
    for sharded, dim in zip(axis_to_shard, shape):
      dim = int(dim)
      split_sizes_per_dim.append([dim])
      if sharded:
        split_sizes = []
        num_shards = dim // partition_info.partition_size
        if num_shards > 0:
          split_sizes = [partition_info.partition_size] * num_shards
          last_shard_size = dim % partition_info.partition_size
          if last_shard_size > 0:
            split_sizes.append(last_shard_size)
        else:
          split_sizes.append(dim)
        split_sizes_per_dim[-1] = split_sizes
    num_splits_per_dim = [len(v) for v in split_sizes_per_dim]
    return PartitionMetadata(split_sizes_per_dim, num_splits_per_dim)

  @classmethod
  def partition_tensor(cls, tensor, partition_info):
    """Returns partitioned tensors."""
    metadata = (TensorPartitioner.partition_metadata(tensor, partition_info))
    # Split from last to first axis.
    partitioned_tensors = [tensor]
    rank = len(metadata.num_splits_per_dim)
    for raxis, (num_splits, sizes) in enumerate(
        zip(
            reversed(metadata.num_splits_per_dim),
            reversed(metadata.split_sizes_per_dim))):
      if num_splits > 1:
        tmp_partitioned_tensors = []
        for item in partitioned_tensors:
          tmp_partitioned_tensors += tf.split(
              item, sizes, axis=rank - raxis - 1)
        partitioned_tensors = tmp_partitioned_tensors
    return partitioned_tensors

  @classmethod
  def reform_tensor(cls, partitioned_tensors, num_splits_per_dim):
    """Returns a tensor concatenated from the given partitions."""
    # Concatenates tensors across all dimension. Assumes the `partitions` tensor
    # was created by partition_tensor.
    for axis, num_splits in enumerate(num_splits_per_dim):
      if num_splits > 1:
        tmp_partitioned_tensors = []
        num_concat = len(partitioned_tensors) // num_splits
        for i in range(num_concat):
          tensors_to_concat = (
              partitioned_tensors[i * num_splits:(i + 1) * num_splits])
          tmp_partitioned_tensors.append(
              tf.concat(tensors_to_concat, axis=axis))
        partitioned_tensors = tmp_partitioned_tensors
    return partitioned_tensors[0]

class DistributedShampoo(optimizer.Optimizer):
  """Approximates full-matrix AdaGrad per layer.

  Approximates full-matrix AdaGrad with kronecker-products of two statistics
  matrices based on only the first-order gradients of the layer.

  "Second-order optimization made practical.", 2019
  Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, Yoram Singer.
  """

  def __init__(self,
               learning_rate,
               momentum=0.0,
               initial_accumulator_value=0.0,
               start_preconditioning_steps=1000,
               statistics_computation_frequency=1,
               matrix_epsilon=1e-6,
               synchronous_preconditioning=False,
               second_moment_averaging=1.0,
               fallback_to_diagonal_dim=4096,
               max_any_dim=6656,
               block_size=4096,
               block_partition_threshold_size=1000000,
               global_step=None,
               exponent_multiplier=1.0,
               name="DistributedShampoo"):
    """Construct a DistributedShampoo optimizer.

    Args:
      learning_rate: A `Tensor` or a floating point value.  The learning rate.
      momentum: A `Tensor` or a floating point value. Momentum is not applied to
        sparse updates.
      initial_accumulator_value: A floating point value.
      start_preconditioning_steps: A int32 value which indicates when to start
        preconditioning.
      statistics_computation_frequency: A int32 step value which indicates how
        often to compute statistics for preconditioning.
      matrix_epsilon: An epsilon regularizer to make the matrices positive
        definite.
      synchronous_preconditioning: Whether to run preconditioning synchronously.
      second_moment_averaging: 1.0 means sum of gradients squares, while less
        than 1.0 switches to RMSProp style exponential moving averages of the
        second moments.
      fallback_to_diagonal_dim: Fallback to diagonal version of AFMA if the any
        of the dimension is larger than fallback_to_diagonal_dim.
      max_any_dim: If maximum value for any dimension is greater than this value
        we skip preconditioning and fall back to the diagonal.
      block_size: Dimension of the partitioned tensors.
      block_partition_threshold_size: Partitions diemnsions beyond this size.
      global_step: Global step for training.
      exponent_multiplier: A multiplier 'e` for the exponent for the inverse
        calculation. e * -1/(2*rank). Only applies when calculating inverses
        through svd.
      name: Optional name prefix for the operations created when applying
        gradients.
    """
    super(DistributedShampoo, self).__init__(False, name)
    self._learning_rate = learning_rate
    self._momentum = momentum
    self._initial_accumulator_value = initial_accumulator_value
    self._start_preconditioning_steps = start_preconditioning_steps
    self._matrix_epsilon = matrix_epsilon
    self._synchronous_preconditioning = synchronous_preconditioning
    self._second_moment_averaging = second_moment_averaging
    self._fallback_to_diagonal_dim = fallback_to_diagonal_dim
    self._max_any_dim = max_any_dim
    self._block_size = block_size
    # NOTE: On XLA - int64 is not handled properly.
    if global_step is not None:
      self._global_step = tf.cast(tf.identity(global_step), tf.int32)
    else:
      self._global_step = tf.cast(
          tf.identity(tf.train.get_or_create_global_step()), tf.int32)
    self._run_nondiagonal_update = tf.greater_equal(
        self._global_step, self._start_preconditioning_steps)
    start_steps_f = tf.cast(self._start_preconditioning_steps, tf.float32)
    global_step_f = tf.cast(self._global_step, tf.float32)
    self._run_nondiagonal_update_warmup = tf.minimum(
        1.0, tf.maximum((global_step_f - start_steps_f) / start_steps_f, 0.0))
    # Computes statistics every K steps.
    self._statistics_computation_frequency = statistics_computation_frequency
    self._run_statistics_computation = tf.equal(
        tf.math.floormod(self._global_step,
                         self._statistics_computation_frequency), 0)
    # All vars that are preconditioned.
    self._all_vars_for_preconditioning = []
    self._exponent_multiplier = exponent_multiplier
    self._partition_info = PartitionConfig(block_partition_threshold_size,
                                           block_size)
    self._partitioner_metadata = {}

  def _fallback_to_diagonal_for_shape(self, shape):
    """Returns whether we should fallback to the diagonal update given shape."""
    # We fallback to diagonal for the following usecases:
    #
    # (a) Rank <= 1 tensors
    # (b) if any dim of Tensor is > max_any_dim.
    # (c) if all dims are 1 or are greater than fallback_to_diagonal_dim
    #
    if len(shape) <= 1:
      return True
    if any([d > self._max_any_dim for d in shape]):
      return True
    if all([d == 1 for d in shape]):
      return True
    return False

  def _preconditioner_available_for_dims(self, shape):
    """Returns indicator vector if preconditioner exists for each axis."""
    # If any of the dims < fallback_to_diagonal_dim and not 1, we run a
    # a preconditioner for that particular dimension.
    return [d <= self._fallback_to_diagonal_dim and d != 1 for d in shape]

  def _preconditioner_indices(self, shape):
    """Returns indices of the available preconditioner."""
    preconditioners_available_for_dims = (
        self._preconditioner_available_for_dims(shape))
    indices = []
    index = 0
    for is_avail_for_dim_i in preconditioners_available_for_dims:
      indices.append(index)
      if is_avail_for_dim_i:
        index += 1
    return indices

  def _make_named_slot(self, var, val, slot_name):
    _ = self._get_or_make_slot(var, val, slot_name,
                               self._name + "_" + slot_name)

  def make_named_zeros_slot(self, var, slot_name):
    self._zeros_slot(var, slot_name, self._name + "_" + slot_name)

  def _generalized_inverse_pth_root(self, input_t, exponent, epsilon=1e-12):
    input_t_f64 = tf.cast(input_t, tf.float64)
    s, u, v = tf.linalg.svd(
        input_t_f64 +
        tf.eye(tf.shape(input_t_f64)[0], dtype=tf.float64) * epsilon,
        full_matrices=True)
    inv_s = tf.reshape(
        tf.pow(tf.maximum(s, epsilon), tf.cast(exponent, tf.float64)), [1, -1])
    val = tf.matmul(u * inv_s, v, adjoint_b=True)
    return tf.cast(val, tf.float32), tf.reduce_max(tf.abs(u - v))

  def _specialized_inverse_pth_root(self, input_t, exponent, epsilon=1e-12):
    input_t_f64 = tf.cast(input_t, tf.float64)
    val, error = matrix_functions.inlined_matrix_inverse_pth_root(
        input_t_f64,
        tf.shape(input_t_f64)[0],
        exponent,
        iter_count=40,
        ridge_epsilon=epsilon)
    return tf.cast(val, tf.float32), error

  def _inverse_pth_root_graph(self, epsilon):
    graph = tf.Graph()
    with graph.as_default():
      exponent_t = tf.reshape(
          tf.placeholder(dtype=tf.float32, name="exponent", shape=None), [])
      # Apply exponent multiplier.
      exponent_t = exponent_t * self._exponent_multiplier
      input_t = tf.placeholder(dtype=tf.float32, name="input", shape=None)
      # For p = 2, 4 or 8, we use the iterative Newton-Schur method for
      # computing the inverse-pth root.
      either_p_2_4_8 = tf.math.logical_or(
          tf.math.logical_or(
              tf.equal(-1.0 / exponent_t, 2), tf.equal(-1.0 / exponent_t, 4)),
          tf.equal(-1.0 / exponent_t, 8))
      # 4096 is the larger dimension SVD is tractable for.
      greater_than_4096 = tf.greater(tf.shape(input_t)[0], 4096)
      run_specialized_iterative_method = tf.math.logical_and(
          greater_than_4096, either_p_2_4_8)
      specialized_fn = functools.partial(self._specialized_inverse_pth_root,
                                         input_t, exponent_t, epsilon)
      generalized_fn = functools.partial(self._generalized_inverse_pth_root,
                                         input_t, exponent_t, epsilon)
      output, diff = tf.cond(run_specialized_iterative_method, specialized_fn,
                             generalized_fn)

      tf.identity(output, "output")
      tf.identity(tf.cast(diff, tf.float32), "diff")
    return graph.as_graph_def().SerializeToString()

  def _create_slots(self, var_list):
    self._preconditioner_compute_graphdef = self._inverse_pth_root_graph(
        epsilon=self._matrix_epsilon)
    for v in var_list:
      self._make_named_slot(v,
                            tf.ones_like(v) * self._initial_accumulator_value,
                            "accumulator")

      if self._momentum > 0.0:
        self.make_named_zeros_slot(v, "momentum")
      shape = np.array(v.get_shape())
      self._partitioner_metadata[v] = TensorPartitioner.partition_metadata(
          v, self._partition_info)
      partitioned_v = TensorPartitioner.partition_tensor(
          v, self._partition_info)
      if not self._fallback_to_diagonal_for_shape(shape):
        self._all_vars_for_preconditioning.append(v)
        if self._momentum > 0.0:
          self.make_named_zeros_slot(v, "precond_grad_momentum")
        num_partitions = len(partitioned_v)
        for pt_idx, pt_v in enumerate(partitioned_v):
          pt_v_shape = pt_v.get_shape()
          preconditioner_exists_for_dim = (
              self._preconditioner_available_for_dims(pt_v_shape))
          for i, d in enumerate(pt_v_shape):
            if preconditioner_exists_for_dim[i]:
              mat_stat_init = array_ops.zeros([d, d], dtype=pt_v.dtype)
              self._make_named_slot(
                  v, mat_stat_init,
                  self._statistics_key_for_partition_and_dim(
                      i, pt_idx, num_partitions))
              self._make_named_slot(
                  v, mat_stat_init,
                  self._preconditioner_key_for_partition_and_dim(
                      i, pt_idx, num_partitions))

  def _prepare(self):
    learning_rate = self._call_if_callable(self._learning_rate)
    self._learning_rate_tensor = ops.convert_to_tensor(
        learning_rate, name="learning_rate")
    momentum = self._call_if_callable(self._momentum)
    self._momentum_tensor = ops.convert_to_tensor(momentum, name="momentum")

  def invoke_async_preconditioner_computation(self, global_step_int32):
    """Invokes SVD preconditioner and graph runs on the CPU."""
    keys_stats_and_rank = []
    for var in self._all_vars_for_preconditioning:
      shape = var.get_shape()
      if not self._fallback_to_diagonal_for_shape(shape):
        partitioned_v = TensorPartitioner.partition_tensor(
            var, self._partition_info)
        num_partitions = len(partitioned_v)
        for pt_idx, pt_v in enumerate(partitioned_v):
          pt_v_shape = pt_v.get_shape()
          preconditioner_exists_for_dim = (
              self._preconditioner_available_for_dims(pt_v_shape))
          for i in range(len(pt_v_shape)):
            if preconditioner_exists_for_dim[i]:
              rank = sum(preconditioner_exists_for_dim)
              key = self._key_for_var(var, i, pt_idx)
              stat = self.get_slot(
                  var,
                  self._statistics_key_for_partition_and_dim(
                      i, pt_idx, num_partitions))
              keys_stats_and_rank.append((key, stat, rank))

    if not keys_stats_and_rank:
      return tf.no_op()
    keys, stats, ranks = zip(*keys_stats_and_rank)

    return x_ops.compute_preconditioners(
        stats, [-1.0 / (2.0 * r) for r in ranks],
        global_step_int32,
        keys=keys,
        sync=self._synchronous_preconditioning,
        preconditioner_compute_graphdef=self._preconditioner_compute_graphdef)

  def assign_preconditioner_to_host_vars(self):
    """Assign/Grab latest copy of preconditioners."""
    keys_shapes_and_preconditioner_vars = []
    assign_ops = []
    for var in self._all_vars_for_preconditioning:
      shape = var.get_shape()
      if not self._fallback_to_diagonal_for_shape(shape):
        partitioned_v = TensorPartitioner.partition_tensor(
            var, self._partition_info)
        num_partitions = len(partitioned_v)
        for pt_idx, pt in enumerate(partitioned_v):
          pt_shape = pt.get_shape()
          preconditioner_exists_for_dim = (
              self._preconditioner_available_for_dims(pt_shape))
          var_rank = len(pt_shape)
          for i in range(var_rank):
            if preconditioner_exists_for_dim[i]:
              key = self._key_for_var(var, i, pt_idx)
              preconditioner = self.get_slot(
                  var,
                  self._preconditioner_key_for_partition_and_dim(
                      i, pt_idx, num_partitions))
              keys_shapes_and_preconditioner_vars.append(
                  (key, tf.shape(preconditioner), preconditioner))

      if not keys_shapes_and_preconditioner_vars:
        return tf.no_op()

      keys, shapes, preconditioner_vars = zip(
          *keys_shapes_and_preconditioner_vars)

      preconditioner_vals, successes = x_ops.get_preconditioners(
          shapes,
          keys=keys,
          preconditioner_compute_graphdef=(
              self._preconditioner_compute_graphdef))

      for preconditioner_var, preconditioner_val, success in zip(
          preconditioner_vars, preconditioner_vals, successes):
        success_mult = tf.cast(success, preconditioner.dtype)
        assign_ops.append(
            state_ops.assign(preconditioner_var,
                             (1.0 - success_mult) * preconditioner_var +
                             success_mult * preconditioner_val))
    return tf.group(*assign_ops)

  def _statistics_key_for_partition_and_dim(self, dim_index, partition_index,
                                            num_partitions):
    if num_partitions == 1:
      return "mat_statistics_" + str(dim_index)
    else:
      return str(partition_index) + "_mat_statistics_" + str(dim_index)

  def _preconditioner_key_for_partition_and_dim(self, dim_index,
                                                partition_index,
                                                num_partitions):
    if num_partitions == 1:
      return "mat_preconditioner_" + str(dim_index)
    else:
      return str(partition_index) + "_mat_preconditioner_" + str(dim_index)

  def _key_for_var(self, var, dim_index, partition_index):
    return "P_" + str(partition_index) + "_D_" + str(dim_index) + "_" + var.name

  def _updated_statistics(self, var, partitioned_grads):
    """Returns updated Shampoo statistics L_t, R_t, etc.

    Args:
      var: tf.Variable associated with the gradient.
      partitioned_grads: Partitioned gradient tensor.

    Returns:
      A list of updated statistics matrices.
    """
    precond_statistics_update = []
    num_partitions = len(partitioned_grads)
    mat_stats = []
    mat_grads = []
    mat_dims = []
    for pt_idx, pt_grad in enumerate(partitioned_grads):
      pt_shape = pt_grad.get_shape()
      preconditioner_exists_for_dim = (
          self._preconditioner_available_for_dims(pt_shape))
      rank = len(pt_shape)
      # Calculates the preconditioner statistics for each tensor.
      for i in range(rank):
        if preconditioner_exists_for_dim[i]:
          mat_stats.append(
              self.get_slot(
                  var,
                  self._statistics_key_for_partition_and_dim(
                      i, pt_idx, num_partitions)))
          mat_grads.append(pt_grad)
          mat_dims.append(i)

    # axes is the list of indices to reduce - everything but
    # the current i.
    def _update_statistics(dim, stat_var, grad):
      """Update preconditioner statistics."""
      with tf.variable_scope("GradientStatistics"):
        var_rank = len(grad.get_shape())
        axes = list(range(dim)) + list(range(dim + 1, var_rank))
        new_stat = math_ops.tensordot(grad, grad, axes=(axes, axes))
        if self._second_moment_averaging == 1.0:
          updated_stat = state_ops.assign_add(stat_var, new_stat)
        else:
          updated_stat = state_ops.assign_add(
              stat_var, (self._second_moment_averaging - 1.0) * stat_var +
              (1.0 - self._second_moment_averaging) * new_stat)
        return updated_stat

    if self._statistics_computation_frequency <= 1:
      for mat_stat, mat_grad, dim in zip(mat_stats, mat_grads, mat_dims):
        precond_statistics_update.append(
            _update_statistics(dim, mat_stat, mat_grad))
    else:

      # NOTE: We rewrite tf.cond() as a while loop to avoid certain overheads
      # in XLA from buffer allocation.
      def _loop_body(mat_stats, mat_grads, mat_dims, unused_perform_step):
        precond_statistics_update_ops = []
        for mat_stat, mat_grad, dim in zip(mat_stats, mat_grads, mat_dims):
          precond_statistics_update_ops.append(
              _update_statistics(dim, mat_stat, mat_grad))
        with tf.control_dependencies(precond_statistics_update_ops):
          return tf.constant(False)

      loop_body_fn = functools.partial(_loop_body, mat_stats, mat_grads,
                                       mat_dims)
      precond_statistics_update.append(
          tf.while_loop(lambda perform_step: perform_step, loop_body_fn,
                        [self._run_statistics_computation]))

    return precond_statistics_update

  def _compute_preconditioned_raw_grad(self, var, partitioned_grads):
    """Returns preconditioned gradient.

    Args:
      var: tf.Variable associated with the gradient.
      partitioned_grads: Partitioned gradient tensor.

    Returns:
      A preconditioned gradient tensor.
    """

    partitioned_preconditioned_grads = []
    num_partitions = len(partitioned_grads)
    for pt_idx, pt_grad in enumerate(partitioned_grads):
      pt_shape = pt_grad.get_shape()
      rank = len(pt_shape)
      preconditioner_exists_for_dim = (
          self._preconditioner_available_for_dims(pt_shape))
      preconditioner_indices = self._preconditioner_indices(pt_shape)
      mat_preconditioner_list = []
      for i in range(rank):
        if preconditioner_exists_for_dim[i]:
          mat_preconditioner_list.append(
              self.get_slot(
                  var,
                  self._preconditioner_key_for_partition_and_dim(
                      i, pt_idx, num_partitions)))
      precond_grad = pt_grad
      if rank == 2 and all(preconditioner_exists_for_dim):
        # Fast path for speedup.
        precond_grad = tf.matmul(
            tf.matmul(mat_preconditioner_list[0], precond_grad),
            mat_preconditioner_list[1])
      else:
        for i in range(rank):
          if preconditioner_exists_for_dim[i]:
            precond_grad = tf.tensordot(
                precond_grad,
                mat_preconditioner_list[preconditioner_indices[i]],
                axes=([0], [0]))
          else:
            # if preconditioner is not available we transpose it to
            # permute the axis for the next preconditioner.
            precond_grad = tf.transpose(
                precond_grad, perm=list(range(1, rank)) + [0])
      partitioned_preconditioned_grads.append(precond_grad)
    return TensorPartitioner.reform_tensor(
        partitioned_preconditioned_grads,
        self._partitioner_metadata[var].num_splits_per_dim)

  def _preconditioned_update(self, var, partitioned_grads,
                             diagonal_grad_update):
    """Computes the matrix preconditioned update.

    Args:
      var: Variable for which we are computing the preconditioned gradient.
      partitioned_grads: Partitioned gradients.
      diagonal_grad_update: Update as given by diagonal adagrad.

    Returns:
      scaled preconditioned gradient.
    """

    def _l2_norm(v):
      return tf.sqrt(tf.reduce_sum(tf.square(v)))

    precond_grad = self._compute_preconditioned_raw_grad(var, partitioned_grads)
    if self._momentum > 0.0:
      gbar = self.get_slot(var, "precond_grad_momentum")
      matrix_preconditioned_grad = state_ops.assign(
          gbar, gbar * self._momentum_tensor + precond_grad *
          (1.0 - self._momentum_tensor))
    else:
      matrix_preconditioned_grad = precond_grad

    # We use the direction from Shampoo while using the step size scale from
    # diagonal AdaGrad.
    precond_l2_norm = _l2_norm(matrix_preconditioned_grad)
    diagonal_l2_norm = _l2_norm(diagonal_grad_update)
    multiplier = tf.where(
        tf.greater(precond_l2_norm, 0.0),
        tf.maximum(diagonal_l2_norm, 1e-30) /
        (tf.maximum(precond_l2_norm, 1e-30)), 1.0)
    return matrix_preconditioned_grad * multiplier

  def _apply_dense(self, grad, var):
    # Calculates the preconditioner statistics for each tensor.
    partitioned_grads = TensorPartitioner.partition_tensor(
        grad, self._partition_info)
    shape = var.get_shape()
    fallback_to_diagonal = self._fallback_to_diagonal_for_shape(shape)

    precond_statistics_update = []
    if not fallback_to_diagonal:
      precond_statistics_update = self._updated_statistics(
          var, partitioned_grads)

    accumulator = self.get_slot(var, "accumulator")
    accumulator_updated = state_ops.assign_add(accumulator, grad * grad)
    accumulator_inv_sqrt = math_ops.rsqrt(accumulator_updated + 1e-30)
    if self._momentum > 0.0:
      scaled_g = (1.0 - self._momentum_tensor) * (grad * accumulator_inv_sqrt)
      gbar = self.get_slot(var, "momentum")
      gbar_updated = state_ops.assign_add(
          gbar,
          gbar * (self._momentum_tensor - 1.0) + scaled_g)
    else:
      gbar_updated = (grad * accumulator_inv_sqrt)

    if not fallback_to_diagonal:
      # Update the preconditioner statistics followed by computing the
      # preconditioned gradient.
      with ops.control_dependencies(precond_statistics_update):
        s = tf.cast(self._run_nondiagonal_update, tf.float32)
        preconditioned_grad = self._preconditioned_update(
            var, partitioned_grads, gbar_updated)
        # slowly adapt from diagonal to preconditioned gradient.
        w = self._run_nondiagonal_update_warmup
        warmup_update = s * self._learning_rate_tensor * (
            w * preconditioned_grad + (1.0 - w) * gbar_updated)
        fallback_update = (1 - s) * (self._learning_rate_tensor * gbar_updated)
        return state_ops.assign_sub(var, warmup_update + fallback_update)
    else:
      return state_ops.assign_sub(var,
                                  self._learning_rate_tensor * gbar_updated)

  def _resource_apply_dense(self, grad, var):
    return self._apply_dense(grad, var)

  # Sparse gradients are not handled currently and is part of future work.
  def _resource_apply_sparse(self, grad_values, var, grad_indices):
    return tf.no_op()

  def _apply_sparse(self, grad, var):
    return tf.no_op()


# -*- coding: utf-8 -*-
"""
This provides the basis for torque computation via two additional
classes: TorqueMap which has data for a torque derivation and
GalacticTorque which derives from GalacticDisc and contains the torque
maps.
"""
import numpy as np

__author__ = "Eric Emsellem"
__copyright__ = "Eric Emsellem"
__license__ = "mit"

from .disc import GalacticDisc
from .misc_io import add_suffix, AttrDict, guess_stepx
from .transform import extract_radial_profile_fromXY
from .maps_grammar import remap_suffix, _is_flag_density
from .local_units import km_pc
from . import fit_functions as ff
from . import gravpot_functions as gpot

class TorqueMap(object):
    """Main map class for the torques
    """
    def __init__(self, massmap, mass_dname, comp_dname,
                 velprof, vel_dname, pc_per_xyunit=1.0, PAnodes=0.0,
                 factor_hz=12.0):
        """
        """
        self.massmap = massmap
        self.mass_dname = mass_dname
        self.comp_dname = comp_dname
        self.velprof = velprof
        self.vel_dname = vel_dname
        self.factor_hz = factor_hz

        self.dmass = self.massmap.dmaps[self.mass_dname]
        self.dcomp = self.massmap.dmaps[self.comp_dname]
        self.dvel = self.velprof.dprofiles[self.vel_dname]
        self.pc_per_xyunit = pc_per_xyunit

        # Check the coordinates
        self.Xdep = self.massmap.X_londep
        self.Ydep = self.massmap.Y_londep
        self.Rdep = np.sqrt(self.Xdep**2 + self.Ydep**2)
        self.Xdep_pc = self.Xdep * self.pc_per_xyunit
        self.Ydep_pc = self.Ydep * self.pc_per_xyunit
        self.Rdep_pc = self.Rdep * self.pc_per_xyunit
        self.pc_per_x = guess_stepx(self.Xdep_pc)
        self.pc_per_y = guess_stepx(self.Ydep_pc)

        # We also need a maximised X and Y grid to help with the kernel
        self.Xdep_pcexp, self.Ydep_pcexp = self.get_XYpcexp()
        self.Rdep_pcexp = np.sqrt(self.Xdep_pcexp**2 + self.Ydep_pcexp**2)

        # Using the average between the x and y step in parsec
        self.pc_per_pixel = (self.pc_per_x + self.pc_per_y) / 2.
        self.PAnodes = PAnodes

    @property
    def XY_extent(self):
        """Provide the extent - a la matplotlib - for
        the given map grid
        """
        return [np.min(self.Xdep), np.max(self.Xdep),
                np.min(self.Ydep), np.max(self.Ydep)]

    @property
    def XYpc_extent(self):
        """Provide the extent in parsec - a la matplotlib - for
        the given map grid
        """
        return [np.min(self.Xdep_pc), np.max(self.Xdep_pc),
                np.min(self.Ydep_pc), np.max(self.Ydep_pc)]

    def get_XYpcexp(self):
        """Provided an expanded grid for X and Y
        which will be used later for the kernel convolution
        """
        maxXpc = np.max(np.abs(self.Xdep_pc))
        maxYpc = np.max(np.abs(self.Ydep_pc))
        maxX = np.max(np.abs(self.Xdep_pc))
        stepX = self.Xdep_pc[0,1] - self.Xdep_pc[0,0]
        stepY = self.Ydep_pc[1,0] - self.Ydep_pc[0,0]
        npixX = np.int(maxXpc // stepX) + 1
        npixY = np.int(maxYpc // stepY) + 1
        nmaxXpc = npixX * stepX
        nmaxYpc = npixY * stepY
        xlin = np.linspace(-nmaxXpc, nmaxXpc, npixX * 2 + 1)
        ylin = np.linspace(-nmaxYpc, nmaxYpc, npixY * 2 + 1)
        eX, eY = np.meshgrid(xlin, ylin)
        return eX, eY

    def get_mass_profile(self, wedge_size=0.0, wedge_angle=0):
        """Compute the 1d mass profile using the dmass.data
        and a number of bins. A wedge-angle and size can be used to indicate
        where to extract the data. Uses the extract_radial_profile_fromXY
        function.

        Input
        -----
        wedge_size (float): size of the wedge in degrees [0]
        wedge_angle (float): angle of the central wedge (in degrees) [0]

        Creates
        -------
        self.Rmass1d and self.mass1d following the deprojected grid Xdep, Ydep
        and the dmass.data map.
        """
        self.Rmass1d, self.mass1d = extract_radial_profile_fromXY(self.Xdep, self.Ydep,
                                                                  self.dmass.data,
                                                                  nbins=None, verbose=True,
                                                                  wedge_size=wedge_size,
                                                                  wedge_angle=wedge_angle)

    def fit_mass_profile(self):
        """Fit the 1d mass profile
        """
        self.opt_mass, self.fsphe1d, self.fdisc1d, self.fmass1d = ff.fit_disc_sphe(self.Rmass1d, self.mass1d)
        self.bfit_mass1d = self.fmass1d(self.Rmass1d, self.opt_mass[0])
        self.bfit_mass1d_sphe = self.fsphe1d(self.Rmass1d, self.opt_mass[0])
        self.bfit_mass1d_disc = self.fdisc1d(self.Rmass1d, self.opt_mass[0])
        self.bfit_sphe1d = self.fsphe1d(self.massmap._R, self.opt_mass[0])
        self.bfit_sphe1d_dep = self.fsphe1d(self.Rdep, self.opt_mass[0])
        self.massmap.faceon = self.dmass.data - self.bfit_sphe1d_dep + self.bfit_sphe1d
        self.Rl_disc = self.opt_mass[0][3]

    def get_kernel(self, softening=0.0, function="sech2"):
        """Get the kernel array
        """
        hz_pc = self.Rl_disc * self.pc_per_xyunit / self.factor_hz
        self.kernel = gpot.get_gravpot_kernel(self.Rdep_pcexp, hz_pc,
                                              pc_per_pixel=self.pc_per_pixel,
                                              softening=softening,
                                              function=function)

    def get_gravpot(self):
        """Calculate the gravitational potential
        """
        self.gravpot = gpot.get_potential(self.massmap.faceon
                                          * self.pc_per_x * self.pc_per_y,
                                          self.kernel)

    def get_forces(self):
        """Calculate the forces from the potential
        Units should be in km^2.s^-2.pc^-2
        """
        self.Fgrad, self.Fx, self.Fy, self.Frad, self.Ftan = \
            gpot.get_forces(self.Xdep_pc, self.Ydep_pc,
                            self.gravpot, self.PAnodes)

    def get_vrot_from_forces(self):
        """Calculate the velocities from forces
        """
        self.VcU = gpot.get_vrot_from_force(self.Rdep_pc, self.Frad)

    def get_torque_map(self):
        """Compute the torque map

        Units back are units of X * F (hence pc * F)
        Since F is in (km/s)^2 pc-1 -> torque in (km/s)^2
        """
        self.torque_map = gpot.get_torque(self.Xdep_pc, self.Ydep_pc,
                                          self.Fx, self.Fy)

    def get_weighted_torque_map(self):
        """Compute the torque map
        """
        self.torque_w_map = gpot.get_weighted_torque(self.Xdep_pc, self.Ydep_pc,
                                          self.Fx, self.Fy, self.dcomp.data)

    def get_torque_profiles(self, n_rbins=300):
        """Get the profiles from the torque
        """
        self.r_mean, self.v_mean, self.torque_mean, self.torque_mean_w, \
            self.ang_mom_mean, self.dm, self.dm_sum = \
            gpot.get_torque_profiles(self.Xdep_pc, self.Ydep_pc, self.VcU,
                                     self.Fx, self.Fy, self.dcomp.data,
                                     n_rbins=n_rbins, pc_per_pixel=self.pc_per_pixel)
        # T in s
        self.Trot = 2. * np.pi * self.r_mean * km_pc / self.v_mean
        # dloL adimensional
        self.dloL = self.torque_mean_w * self.Trot / self.ang_mom_mean

    def get_torques(self, n_rbins=300):
        """Calculate the torques from existing forces

        Args:
            n_rbins:
        """
        self.get_torque_map()
        self.get_weighted_torque_map()
        self.get_torque_profiles(n_rbins=n_rbins)

    def run_torques(self, softening=0.0, func_kernel="sech2", n_rbins=200):
        """Running the torque calculation from start to end

        Args:
            softening:
            func_kernel:
            n_rbins:

        """
        # Step 1 - extract the radial profile and fit it with bulge and disc
        self.get_mass_profile()

        # Step 2 - Now doing the fit of the spheroid (and disc)
        self.fit_mass_profile()

        # Step 3 - calculate the kernel
        self.get_kernel(softening=softening, function=func_kernel)

        # Step 4 - calculate the potential
        self.get_gravpot()

        # Step 5 - Calculate the forces
        self.get_forces()

        # Step 6 - Get the rotation velocities
        self.get_vrot_from_forces()

        # Step 7 - Normalise the fields with M/L
        # For the moment = passed

        # Step 8 - Calculate the torques
        self.get_torques(n_rbins=n_rbins)

class GalacticTorque(GalacticDisc):
    """Class for functionalities associated with Torques
    """

    def __init__(self, vcfile_name=None, vcfile_type="ROTCUR",
                 Rfinestep=0, vprof_name="Velocity", **kwargs):
        """

        Args:
            vcfile_name:
            vcfile_type:
            **kwargs:
        """
        self.verbose = kwargs.pop("verbose", False)

        # Using GalacticDisc class attributes
        super().__init__(**kwargs)

        # Now the velocity file
        if vcfile_name is not None:
            print("INFO: Adding the provided Vc file")
            velname = self.add_vprofile(filename=vcfile_name, filetype=vcfile_type,
                              Rfinestep=Rfinestep, vprof_name=vprof_name)
        else:
            velname = kwargs.pop("velname", "vel_vel01")

        # And checking the maps
        compname = kwargs.pop("compname", "comp_comp01")
        massname = kwargs.pop("massname", "mass_mass01")
        self.init_torque_components(velname=velname, compname=compname,
                                    massname=massname)

        # Make sure we start with a clean set of torque maps
        self._reset_torquemaps()

    @property
    def velname(self):
        return self.profiles[self.vel_pname]._fullname(self.vel_dname)

    @property
    def massname(self):
        return self.maps[self.mass_mname]._fullname(self.mass_dname)

    @property
    def compname(self):
        return self.maps[self.comp_mname]._fullname(self.comp_dname)

    def init_torque_components(self, **kwargs):
        """Initialise the torque maps components

        Args:
            **kwargs: velname, compname, massname

        """
        self._decode_torque_names(**kwargs)
        self.check_torque_components()
        if self._check_all:
            self.match_comp_mass()

    def _decode_torque_names(self, velname=None, compname=None, massname=None):
        """

        Args:
            velname (str): composite name for the velocity profile
            compname (str): composite name for the component map
            massname (str): composite name for the mass map

        """
        # Decoding the names
        self.vel_pname, self.vel_dname = self._decode_prof_name(velname)
        self.comp_mname, self.comp_dname = self._decode_map_name(compname)
        self.mass_mname, self.mass_dname = self._decode_map_name(massname)

    def check_torque_components(self):
        """Find the components for the torque calculations using the composite
        names for the velocities, component (e.g., gas flux) and mass.

        """
        # Checking the maps and profiles
        # And making sure they are density maps
        if self._check_mass():
            thismap = self.massmap
            if not _is_flag_density(self.massdmap.flag):
                self.mass_dname = thismap.intmap_to_densitymap(self.mass_dname, self)

        if self._check_comp():
            thismap = self.compmap
            if not _is_flag_density(self.compdmap.flag):
                self.comp_dname = thismap.intmap_to_densitymap(self.comp_dname, self)

    @property
    def veldprof(self):
        return self.profiles[self.vel_pname].dprofiles[self.vel_dname]

    @property
    def velprof(self):
        return self.profiles[self.vel_pname]

    @property
    def compdmap(self):
        return self.maps[self.comp_mname].dmaps[self.comp_dname]

    @property
    def compmap(self):
        return self.maps[self.comp_mname]

    @property
    def massdmap(self):
        return self.maps[self.mass_mname].dmaps[self.mass_dname]

    @property
    def massmap(self):
        return self.maps[self.mass_mname]

    @property
    def _matched(self):
        return (self.mass_mname == self.comp_mname)

    @property
    def _check_all(self):
        return all((self._check_comp(), self._check_mass(), self._check_vel()))

    def _check_vel(self):
        return self._has_profile_data(self.vel_pname, self.vel_dname,
                                      order=1)

    def _check_comp(self):
        return self._has_map_data(self.comp_mname, self.comp_dname,
                                  order=0)

    def _check_mass(self):
        return self._has_map_data(self.mass_mname, self.mass_dname,
                                  order=0)

    def _reset_torquemaps(self):
        """Initalise the torquemap Profiles by setting an empty
        'torquemaps' dictionary
        """
        # set up the torquemaps
        self.tmaps = AttrDict()

    @property
    def ntorquemaps(self):
        """Number of existing torquemap profiles
        """
        if hasattr(self, 'tmaps'):
            return len(self.tmaps)
        else:
            return -1

    def match_comp_mass(self, odname1="dmass", odname2="dcomp", **kwargs):
        """Aligning the gas onto the mass map
        """
        if not self._check_all:
            print("WARNING[match_comp_mass]: cannot proceed with match "
                  "as all maps are not yet set up. Please check")
            return

        if not self._matched:
            match_name = self.match_datamaps(self.mass_mname, self.comp_mname,
                                             self.mass_dname, self.comp_dname,
                                             odname1, odname2)
            print("INFO[match_comp_mass]: new map is {}".format(match_name))
            self.mass_mname = match_name
            self.comp_mname = match_name
            self.mass_dname = odname1
            self.comp_dname = odname2
            self.deproject_nodes(match_name)
        else:
            print("WARNING[match_comp_mass]: nothing to match as the data are "
                  "associated with the same map {}".format(self.mass_mname))

    def run_torques(self, torquemap_name=None, softening=0, func_kernel="sech2",
                    n_rbins=200, **kwargs):
        """Running the torques recipes

        Args:
            gas_name:
            mass_name:
            vel_name:
            torquemap_name:

        Returns:

        """
        if torquemap_name is None:
            torquemap_name = "Torq{0:02d}".format(self.ntorquemaps+1)

        # Step 0 - finding the gas, mass and vel if not defined
        # and match the maps
        if not self._check_all:
            # Try to read the maps
            velname = kwargs.pop("velname", self.velname)
            compname = kwargs.pop("compname", self.compname)
            massname = kwargs.pop("massname", self.massname)
            self.init_torque_components(velname=velname, compname=compname,
                                        massname=massname)

        pc_per_xy = self.pc_per_xyunit(self.massmap.XYunit)

        # Defining the structure in which the torques will be calculated
        # Note that PAnodes is now -90 as we put the nodes along the X axis - horizontal
        # During the matching
        newT = TorqueMap(self.massmap, self.mass_dname, self.comp_dname,
                         self.velprof, self.vel_dname, pc_per_xyunit=pc_per_xy,
                         PAnodes=-90.0)

        # Running the torque calculation
        newT.run_torques(softening=softening, n_rbins=n_rbins,
                         func_kernel=func_kernel)

        # Now allocating it to the torquemaps
        self.tmaps[torquemap_name] = newT


# -*- coding: UTF-8 -*-
# !/usr/bin/python
# @time     :2019/8/31 21:33
# @author   :Mo
# @function :topic model of LDA
# @paper    :Latent Dirichlet Allocation

from nlg_yongzhuo.data_preprocess.text_preprocess import extract_chinese
from nlg_yongzhuo.data_preprocess.text_preprocess import cut_sentence
from nlg_yongzhuo.data_preprocess.text_preprocess import jieba_cut
from nlg_yongzhuo.data.stop_words.stop_words import stop_words
# sklearn
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.decomposition import LatentDirichletAllocation
import numpy as np

class LDASum:
    def __init__(self):
        self.stop_words = stop_words.values()
        self.algorithm = 'lda'

    def summarize(self, text, num=8, topic_min=6, judge_topic=None):
        """

        :param text: str
        :param num: int
        :return: list
        """
        # 切句
        if type(text) == str:
            self.sentences = cut_sentence(text)
        elif type(text) == list:
            self.sentences = text
        else:
            raise RuntimeError("text type must be list or str")
        len_sentences_cut = len(self.sentences)
        # 切词
        sentences_cut = [[word for word in jieba_cut(extract_chinese(sentence))
                          if word.strip()] for sentence in self.sentences]
        # 去除停用词等
        self.sentences_cut = [list(filter(lambda x: x not in self.stop_words, sc)) for sc in sentences_cut]
        self.sentences_cut = [" ".join(sc) for sc in self.sentences_cut]
        # 计算每个句子的tf
        vector_c = CountVectorizer(ngram_range=(1, 2), stop_words=self.stop_words)
        tf_ngram = vector_c.fit_transform(self.sentences_cut)
        # 主题数, 经验判断
        topic_num = min(topic_min, int(len(sentences_cut) / 2))  # 设定最小主题数为3
        lda = LatentDirichletAllocation(n_components=topic_num, max_iter=32,
                                        learning_method='online',
                                        learning_offset=50.,
                                        random_state=2019)
        res_lda_u = lda.fit_transform(tf_ngram.T)
        res_lda_v = lda.components_

        if judge_topic:
            ### 方案一, 获取最大那个主题的k个句子
            ##################################################################################
            topic_t_score = np.sum(res_lda_v, axis=-1)
            # 对每列(一个句子topic_num个主题),得分进行排序,0为最大
            res_nmf_h_soft = res_lda_v.argsort(axis=0)[-topic_num:][::-1]
            # 统计为最大每个主题的句子个数
            exist = (res_nmf_h_soft <= 0) * 1.0
            factor = np.ones(res_nmf_h_soft.shape[1])
            topic_t_count = np.dot(exist, factor)
            # 标准化
            topic_t_count /= np.sum(topic_t_count, axis=-1)
            topic_t_score /= np.sum(topic_t_score, axis=-1)
            # 主题最大个数占比, 与主题总得分占比选择最大的主题
            topic_t_tc = topic_t_count + topic_t_score
            topic_t_tc_argmax = np.argmax(topic_t_tc)
            # 最后得分选择该最大主题的
            res_nmf_h_soft_argmax = res_lda_v[topic_t_tc_argmax].tolist()
            res_combine = {}
            for l in range(len_sentences_cut):
                res_combine[self.sentences[l]] = res_nmf_h_soft_argmax[l]
            score_sen = [(rc[1], rc[0]) for rc in sorted(res_combine.items(), key=lambda d: d[1], reverse=True)]
            #####################################################################################
        else:
            ### 方案二, 获取最大主题概率的句子, 不分主题
            res_combine = {}
            for i in range(len_sentences_cut):
                res_row_i = res_lda_v[:, i]
                res_row_i_argmax = np.argmax(res_row_i)
                res_combine[self.sentences[i]] = res_row_i[res_row_i_argmax]
            score_sen = [(rc[1], rc[0]) for rc in sorted(res_combine.items(), key=lambda d: d[1], reverse=True)]
        num_min = min(num, len(self.sentences))
        return score_sen[0:num_min]

if __name__ == '__main__':
    lda = LDASum()
    doc = "多知网5月26日消息，今日，方直科技发公告，拟用自有资金人民币1.2亿元，" \
          "与深圳嘉道谷投资管理有限公司、深圳嘉道功程股权投资基金（有限合伙）共同发起设立嘉道方直教育产业投资基金（暂定名）。" \
          "该基金认缴出资总规模为人民币3.01亿元。" \
          "基金的出资方式具体如下：出资进度方面，基金合伙人的出资应于基金成立之日起四年内分四期缴足，每期缴付7525万元；" \
          "各基金合伙人每期按其出资比例缴付。合伙期限为11年，投资目标为教育领域初创期或成长期企业。" \
          "截止公告披露日，深圳嘉道谷投资管理有限公司股权结构如下:截止公告披露日，深圳嘉道功程股权投资基金产权结构如下:" \
          "公告还披露，方直科技将探索在中小学教育、在线教育、非学历教育、学前教育、留学咨询等教育行业其他分支领域的投资。" \
          "方直科技2016年营业收入9691万元，营业利润1432万元，归属于普通股股东的净利润1847万元。（多知网 黎珊）}}"

    doc = "PageRank算法简介。" \
           "是上世纪90年代末提出的一种计算网页权重的算法! " \
           "当时，互联网技术突飞猛进，各种网页网站爆炸式增长。 " \
           "业界急需一种相对比较准确的网页重要性计算方法。 " \
           "是人们能够从海量互联网世界中找出自己需要的信息。 " \
           "百度百科如是介绍他的思想:PageRank通过网络浩瀚的超链接关系来确定一个页面的等级。 " \
           "Google把从A页面到B页面的链接解释为A页面给B页面投票。 " \
           "Google根据投票来源甚至来源的来源，即链接到A页面的页面。 " \
           "和投票目标的等级来决定新的等级。简单的说， " \
           "一个高等级的页面可以使其他低等级页面的等级提升。 " \
           "具体说来就是，PageRank有两个基本思想，也可以说是假设。 " \
           "即数量假设：一个网页被越多的其他页面链接，就越重）。 " \
           "质量假设：一个网页越是被高质量的网页链接，就越重要。 " \
           "总的来说就是一句话，从全局角度考虑，获取重要的信。 "

    sum = lda.summarize(doc)
    for i in sum:
        print(i)


import numpy as np
import time
import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split
from sklearn import datasets

from Py_FS.wrapper.nature_inspired._utilities import Solution, Data, initialize, sort_agents, display, compute_fitness, Conv_plot
from Py_FS.wrapper.nature_inspired._transfer_functions import get_trans_function

def MA(num_agents, max_iter, train_data, train_label, obj_function=compute_fitness, trans_function_shape='s',  prob_mut=0.2,  save_conv_graph=False):

    # Mayfly Algorithm
    ############################### Parameters ####################################
    #                                                                             #
    #   num_agents: number of mayflies                                            #
    #   max_iter: maximum number of generations                                   #
    #   train_data: training samples of data                                      #
    #   train_label: class labels for the training samples                        #
    #   obj_function: the function to maximize while doing feature selection      #
    #   prob_mut: probability of mutation                                         #
    #   trans_function_shape: shape of the transfer function used                 #
    #   save_conv_graph: boolean value for saving convergence graph               #
    #                                                                             #
    ###############################################################################

    short_name = 'MA'
    agent_name = 'Mayfly'
    train_data, train_label = np.array(train_data), np.array(train_label)
    num_features = train_data.shape[1]
    trans_function = get_trans_function(trans_function_shape)

    # setting up the objectives
    weight_acc = None
    if(obj_function==compute_fitness):
        weight_acc = float(input('Weight for the classification accuracy [0-1]: '))
    obj = (obj_function, weight_acc)
    compute_accuracy = (compute_fitness, 1) # compute_accuracy is just compute_fitness with accuracy weight as 1

    # control parameters
    a1 = 1
    a2 = 1.5
    d = 0.1
    fl = 0.1
    g = 0.8
    beta = 2
    delta = 0.9

    # initialize position and velocities of male and female mayflies' and Leader (the agent with the max fitness)
    male_pos = initialize(num_agents, num_features)
    female_pos = initialize(num_agents, num_features)
    male_vel = np.random.uniform(low = -1, high = 1, size = (num_agents, num_features))
    female_vel = np.random.uniform(low = -1, high = 1, size = (num_agents, num_features))
    male_fitness = np.zeros((num_agents))
    male_accuracy = np.zeros(num_agents)
    female_fitness = np.zeros((num_agents))
    Leader_agent = np.zeros((num_features))
    Leader_fitness = float("-inf")
    Leader_accuracy = float("-inf")
    male_personal_best = np.zeros((num_agents, num_features))
    male_offspring = np.zeros((num_agents, num_features))
    female_offspring = np.zeros((num_agents, num_features))
    vmax_male = np.zeros((num_features))
    vmax_female = np.zeros((num_features))

    # initialize convergence curves
    convergence_curve = {}
    convergence_curve['fitness'] = np.zeros(max_iter)

    # initialize data class
    data = Data()
    val_size = float(input('Enter the percentage of data wanted for valdiation [0, 100]: '))/100
    data.train_X, data.val_X, data.train_Y, data.val_Y = train_test_split(train_data, train_label, stratify=train_label, test_size=val_size)

    # create a solution object
    solution = Solution()
    solution.num_agents = num_agents
    solution.max_iter = max_iter
    solution.num_features = num_features
    solution.obj_function = obj_function

    # rank initial population
    male_pos, male_fitness = sort_agents(male_pos, obj, data)
    female_pos, female_fitness = sort_agents(female_pos, obj, data)

    # start timer
    start_time = time.time()

    # main loop
    for iter_no in range(max_iter):
        print('\n================================================================================')
        print('                          Iteration - {}'.format(iter_no+1))
        print('================================================================================\n')

        #updating velocity limits
        vmax_male, vmax_female = update_max_velocity(male_pos, female_pos)

        for agent in range(num_agents):

            #updating Leader fitness and personal best fitnesses
            if male_fitness[agent] > Leader_fitness:
                Leader_fitness = male_fitness[agent]
                Leader_agent = male_pos[agent]

            if male_fitness[agent] > obj_function(male_personal_best[agent], data.train_X, data.val_X, data.train_Y, data.val_Y):
                male_personal_best[agent] = male_pos[agent]

            #update velocities of male and female mayflies
            male_vel[agent], female_vel[agent] = update_velocity(male_pos[agent], female_pos[agent], male_vel[agent], female_vel[agent], Leader_agent, male_personal_best[agent], a1, a2, d, fl, g, beta, agent, data, obj_function)

            #check boundary condition of velocities of male and female mayflies
            male_vel[agent], female_vel[agent] = check_velocity_limits(male_vel[agent], female_vel[agent], vmax_male, vmax_female)

            #applying transfer functions to update positions of male and female mayflies
            #the updation is done based on their respective velocity values
            for j in range(num_features):
                trans_value = trans_function(male_vel[agent][j])
                if trans_value > np.random.normal(0,1):
                    male_pos[agent][j]=1
                else:
                    male_pos[agent][j]=0

                trans_value = trans_function(female_vel[agent][j])
                if trans_value > np.random.random():
                    female_pos[agent][j]=1
                else:
                    female_pos[agent][j]=0

        #sorting
        male_pos, male_fitness = sort_agents(male_pos, obj, data)
        female_pos, female_fitness = sort_agents(female_pos, obj, data)

        for agent in range(num_agents):

            #generation of offsprings by crossover and mutation between male and female parent mayflies
            male_offspring[agent], female_offspring[agent] = cross_mut(male_pos[agent], female_pos[agent],prob_mut)

        #comparing parents and offsprings and replacing parents wherever necessary
        male_pos = compare_and_replace(male_pos, male_offspring, male_fitness, data, obj)
        female_pos = compare_and_replace(female_pos, female_offspring, female_fitness, data, obj)

        #updating fitness values
        male_pos, male_fitness = sort_agents(male_pos, obj, data)
        female_pos, female_fitness = sort_agents(female_pos, obj, data)

        #updating values of nuptial dance
        d = d * delta
        fl = fl * delta

        #update final information
        display(male_pos, male_fitness, agent_name)
        if(male_fitness[0] > Leader_fitness):
            Leader_agent = male_pos[0].copy()
            Leader_fitness = male_fitness[0].copy()

        convergence_curve['fitness'][iter_no] = np.mean(male_fitness)

    # compute final accuracy
    Leader_agent, Leader_accuracy = sort_agents(Leader_agent, compute_accuracy, data)
    male_pos, male_accuracy = sort_agents(male_pos, compute_accuracy, data)

    print('\n================================================================================')
    print('                                    Final Result                                  ')
    print('================================================================================\n')
    print('Leader ' + agent_name + ' Dimension : {}'.format(int(np.sum(Leader_agent))))
    print('Leader ' + agent_name + ' Fitness : {}'.format(Leader_fitness))
    print('Leader ' + agent_name + ' Classification Accuracy : {}'.format(Leader_accuracy))
    print('\n================================================================================\n')

    # stop timer
    end_time = time.time()
    exec_time = end_time - start_time

    # plot convergence graph
    fig, axes = Conv_plot(convergence_curve)
    if(save_conv_graph):
        plt.savefig('convergence_graph_'+ short_name + '.jpg')
    plt.show()

    # update attributes of solution
    solution.best_agent = Leader_agent
    solution.best_fitness = Leader_fitness
    solution.best_accuracy = Leader_accuracy
    solution.convergence_curve = convergence_curve
    solution.final_population = male_pos
    solution.final_fitness = male_fitness
    solution.final_accuracy = male_accuracy
    solution.execution_time = exec_time

    return solution

def update_max_velocity(male, female):
    size, length = male.shape
    agent1 = []
    agent2 = []
    r = np.random.normal(0,1 , size=(length))
    for j in range(length):
        r[j] *= 2
        agent1.append((male[0][j]-male[size-1][j])*r[j])
        agent2.append((female[0][j]-female[size-1][j])*r[j])

    return (agent1, agent2)

def update_velocity(m_pos, f_pos, m_vel, f_vel, Leader_agent, pbest, a1, a2, d, fl, g, b, i, data, obj_function):
    tot_features = m_pos.shape[0]
    agent1 = np.zeros((tot_features))
    agent2 = np.zeros((tot_features))
    tot_features = len(m_pos)
    if i==0:
        for j in range(tot_features):
            agent1[j] = m_vel[j]+d*np.random.uniform(-1,1)
    else:
        sum = 0
        for j in range(tot_features):
            sum = sum+(m_pos[j]-Leader_agent[j])*(m_pos[j]-Leader_agent[j])
        rg = np.sqrt(sum)
        sum = 0
        for j in range(tot_features):
            sum = sum+(m_pos[j]-pbest[j])*(m_pos[j]-pbest[j])
        rp = np.sqrt(sum)
        for j in range(tot_features):
            agent1[j] = g*m_vel[j]+a1*np.exp(-b*rp*rp)*(pbest[j]-m_pos[j])+a2*np.exp(-b*rg*rg)*(Leader_agent[j]-m_pos[j])
    if obj_function(m_pos, data.train_X, data.val_X, data.train_Y, data.val_Y) >= obj_function(f_pos, data.train_X, data.val_X, data.train_Y, data.val_Y):
        sum = 0
        for j in range(tot_features):
            sum = sum+(m_pos[j]-f_pos[j])*(m_pos[j]-f_pos[j])
        rmf = np.sqrt(sum)
        agent2[j] = g*f_vel[j]+a2*np.exp(-b*rmf*rmf)*(m_pos[j]-f_pos[j])
    else:
        for j in range(tot_features):
            agent2[j] = g*f_vel[j]+fl*np.random.uniform(-1,1)

    return (agent1, agent2)

def check_velocity_limits(m_vel, f_vel, vmax_m, vmax_f):
    tot_features = len(m_vel)
    for j in range(tot_features):
        m_vel[j] = np.minimum(m_vel[j], vmax_m[j])
        m_vel[j] = np.maximum(m_vel[j], -vmax_m[j])
        f_vel[j] = np.minimum(f_vel[j], vmax_f[j])
        f_vel[j] = np.maximum(f_vel[j], -vmax_f[j])

    return (m_vel, f_vel)

def cross_mut(m_pos, f_pos,prob_mut):
    tot_features = len(m_pos)
    offspring1 = np.zeros((tot_features))
    offspring2 = np.zeros((tot_features))
    # partition defines the midpoint of the crossover
    partition = np.random.randint(tot_features//4, np.floor((3*tot_features//4)+1))

    # starting crossover
    for i in range(partition):
        offspring1[i] = m_pos[i]
        offspring2[i] = f_pos[i]

    for i in  range(partition, tot_features):
        offspring1[i] = f_pos[i]
        offspring2[i] = m_pos[i]
    # crossover ended

    # starting mutation
    if np.random.random() <= prob_mut:
        percent = 0.2
        numChange = int(tot_features*percent)
        pos = np.random.randint(0,tot_features-1,numChange)

        for j in pos:
            offspring1[j] = 1-offspring1[j]
        pos=np.random.randint(0,tot_features-1,numChange)
        for j in pos:
            offspring2[j] = 1-offspring2[j]

    # mutation ended

    if np.random.random() >= 0.5:
        return (offspring1, offspring2)
    else:
        return (offspring2, offspring1)

def compare_and_replace(pos, off, fit, data, obj):
    agents, features = pos.shape
    newfit = np.zeros((agents))
    temp_pos = np.zeros((agents, features))
    pos, fit = sort_agents(pos, obj, data)
    # finding fitnesses of offsprings
    off, newfit = sort_agents(off, obj, data)
    i=0
    j=0
    cnt=0
    # merging offsprings and parents and finding the next generation of mayflies
    while(cnt < agents):
        if fit[i] > newfit[j]:
            temp_pos[cnt] = pos[i].copy()
            i+=1
        else:
            temp_pos[cnt] = off[i].copy()
            j+=1
        cnt+=1
    return temp_pos

if __name__ == '__main__':
    data = datasets.load_digits()
    MA(20, 30, data.data, data.target, save_conv_graph=True)