# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

"""CMSIS-NN integration tests: softmax"""

import pytest
import sys

import tvm
from tvm import relay
from tvm.relay.op.contrib import cmsisnn
import numpy as np

def count_num_calls(mod):
    class CallCounter(relay.ExprVisitor):
        def __init__(self):
            super().__init__()
            self.count = 0

        def visit_call(self, call):
            if isinstance(call.op, tvm.ir.Op):
                self.count += 1

            super().visit_call(call)

    counter = CallCounter()
    for var in mod.get_global_vars():
        counter.visit(mod[var.name_hint])
    return counter.count

def make_module(func):
    func = relay.Function(relay.analysis.free_vars(func), func)
    mod = tvm.IRModule.from_expr(func)
    return relay.transform.InferType()(mod)

def make_model(shape, zero_point, scale, in_dtype, out_dtype):
    a = relay.var("a", shape=shape, dtype=in_dtype)
    dequantize = relay.qnn.op.dequantize(
        a,
        input_scale=relay.const(scale, "float32"),
        input_zero_point=relay.const(zero_point, "int32"),
    )
    softmax = relay.nn.softmax(dequantize)
    model = relay.qnn.op.quantize(
        softmax,
        output_scale=relay.const(scale, "float32"),
        output_zero_point=relay.const(zero_point, "int32"),
        out_dtype=out_dtype,
    )
    return model

def test_softmax_int8():
    model = make_model([1, 16, 16, 3], 64, 0.02, "int8", "int8")
    orig_mod = make_module(model)
    cmsisnn_mod = cmsisnn.partition_for_cmsisnn(orig_mod)

    attrs = [
        cmsisnn_mod[var.name_hint].attrs
        for var in cmsisnn_mod.get_global_vars()
        if cmsisnn_mod[var.name_hint].attrs
    ]
    assert any(attrs), "At least one function with external attributes was expected."

    compilers = [
        key == "Compiler" and value == "cmsisnn" for attr in attrs for key, value in attr.items()
    ]
    assert any(compilers), "Module does not contain function for cmsisnn target."

    assert count_num_calls(orig_mod) == count_num_calls(
        cmsisnn_mod
    ), "Number of calls changed during partitioning"

@pytest.mark.parametrize("in_dtype,out_dtype", [["uint8", "int8"], ["int8", "uint8"]])
def test_softmax_not_int8(in_dtype, out_dtype):
    model = make_model([1, 16, 16, 3], 64, 0.02, in_dtype, out_dtype)
    orig_mod = make_module(model)
    cmsisnn_mod = cmsisnn.partition_for_cmsisnn(orig_mod)

    attrs = [
        cmsisnn_mod[var.name_hint].attrs
        for var in cmsisnn_mod.get_global_vars()
        if cmsisnn_mod[var.name_hint].attrs
    ]
    assert not any(attrs), "No function should have an external attribute."

if __name__ == "__main__":
    sys.exit(pytest.main([__file__] + sys.argv[1:]))


# -*- coding: utf-8 -*-
# Copyright 2019 the HERA Project
# Licensed under the MIT License

"""Module for delay filtering data and related operations."""

import numpy as np

from . import io
from . import version
from .vis_clean import VisClean
from . import vis_clean
import pickle
import random
import glob
import os
import warnings
from pyuvdata import UVCal
from copy import deepcopy

class DelayFilter(VisClean):
    """
    DelayFilter object.

    Used for delay CLEANing and filtering.
    See vis_clean.VisClean for CLEAN functions.
    """
    def run_filter(self, **kwargs):
        '''
        wrapper for run_delay_filter. Backwards compatibility. See run_delay_filter for documentation.
        '''
        self.run_delay_filter(**kwargs)

    def run_delay_filter(self, to_filter=None, weight_dict=None, horizon=1., standoff=0.15, min_dly=0.0, mode='clean',
                         skip_wgt=0.1, tol=1e-9, cache_dir=None, read_cache=False, write_cache=False,
                         skip_flagged_edges=False, **filter_kwargs):
        '''
        Run uvtools.dspec.vis_filter on data.

        Run a delay-filter on (a subset of) the data stored in the object.
        Uses stored flags unless explicitly overridden with weight_dict.

        Arguments:
            to_filter: list of visibilities to filter in the (i,j,pol) format.
                If None (the default), all visibilities are filtered.
            weight_dict: dictionary or DataContainer with all the same keys as self.data.
                Linear multiplicative weights to use for the delay filter. Default, use np.logical_not
                of self.flags. uvtools.dspec.delay_filter will renormalize to compensate.
            horizon: coefficient to bl_len where 1 is the horizon [freq filtering]
            standoff: fixed additional delay beyond the horizon (in nanosec) to filter [freq filtering]
            min_dly: max delay (in nanosec) used for freq filter is never below this.
            mode: string specifying filtering mode. See fourier_filter or uvtools.dspec.fourier_filter for supported modes.
            skip_wgt: skips filtering rows with very low total weight (unflagged fraction ~< skip_wgt).
                Model is left as 0s, residual is left as data, and info is {'skipped': True} for that
                time. Skipped channels are then flagged in self.flags.
                Only works properly when all weights are all between 0 and 1.
            tol : float, optional. To what level are foregrounds subtracted.
            cache_dir: string, optional, path to cache file that contains pre-computed dayenu matrices.
                see uvtools.dspec.dayenu_filter for key formats.
            read_cache: bool, If true, read existing cache files in cache_dir before running.
            write_cache: bool. If true, create new cache file with precomputed matrices
                that were not in previously loaded cache files.
            skip_flagged_edges : bool, if true do not include frequencies at the edge of the band that are fully flagged. Instead
                filter over frequencies bounded by the edge flags.
            filter_kwargs: see fourier_filter for a full list of filter_specific arguments.

        Results are stored in:
            self.clean_resid: DataContainer formatted like self.data with only high-delay components
            self.clean_model: DataContainer formatted like self.data with only low-delay components
            self.clean_info: Dictionary of info from uvtools.dspec.delay_filter with the same keys as self.data
        '''
        # read in cache
        if not mode == 'clean':
            if read_cache:
                filter_cache = io.read_filter_cache_scratch(cache_dir)
            else:
                filter_cache = {}
            keys_before = list(filter_cache.keys())
        else:
            filter_cache = None
        # loop over all baselines in increments of Nbls
        self.vis_clean(keys=to_filter, data=self.data, flags=self.flags, wgts=weight_dict,
                       ax='freq', x=self.freqs, cache=filter_cache, mode=mode,
                       horizon=horizon, standoff=standoff, min_dly=min_dly, tol=tol,
                       skip_wgt=skip_wgt, overwrite=True,
                       skip_flagged_edges=skip_flagged_edges, **filter_kwargs)
        if not mode == 'clean':
            if write_cache:
                filter_cache = io.write_filter_cache_scratch(filter_cache, cache_dir, skip_keys=keys_before)

def load_delay_filter_and_write(datafile_list, baseline_list=None, calfile_list=None,
                                Nbls_per_load=None, spw_range=None, cache_dir=None,
                                read_cache=False, write_cache=False, avg_red_bllens=False,
                                factorize_flags=False, time_thresh=0.05, external_flags=None,
                                res_outfilename=None, CLEAN_outfilename=None, filled_outfilename=None,
                                clobber=False, add_to_history='', polarizations=None,
                                skip_flagged_edges=False, overwrite_flags=False,
                                flag_yaml=None, **filter_kwargs):
    '''
    Uses partial data loading and writing to perform delay filtering.
    While this function reads from multiple files (in datafile_list)
    it always writes to a single file for the resid, filled, and model files.

    Arguments:
        datafile_list: list of data files to perform cross-talk filtering on
        baseline_list: list of antenna-pair 2-tuples to filter and write out from the datafile_list.
                       If None, load all baselines in files. Default is None.
        Nbls_per_load: int, the number of baselines to load at once.
            If None, load all baselines at once. default : None.
        calfile_list: optional list of calibration files to apply to data before xtalk filtering
        spw_range: 2-tuple or 2-list, spw_range of data to filter.
        cache_dir: string, optional, path to cache file that contains pre-computed dayenu matrices.
            see uvtools.dspec.dayenu_filter for key formats.
        read_cache: bool, If true, read existing cache files in cache_dir before running.
        write_cache: bool. If true, create new cache file with precomputed matrices
            that were not in previously loaded cache files.
        avg_red_bllens: bool, if True, round baseline lengths to redundant average. Default is False.
        factorize_flags: bool, optional
            If True, factorize flags before running delay filter. See vis_clean.factorize_flags.
        time_thresh : float
            Fractional threshold of flagged pixels across time needed to flag all times
            per freq channel. It is not recommend to set this greater than 0.5.
            Fully flagged integrations do not count towards triggering time_thresh.
        external_flags : str, optional, path to external flag files to apply
        res_outfilename: path for writing the filtered visibilities with flags
        CLEAN_outfilename: path for writing the CLEAN model visibilities (with the same flags)
        filled_outfilename: path for writing the original data but with flags unflagged and replaced
            with CLEAN models wherever possible
        clobber: if True, overwrites existing file at the outfilename
        add_to_history: string appended to the history of the output file
        polarizations: list of polarizations to include and write.
        skip_flagged_edges: if true, skip flagged edges in filtering.
        flag_yaml: path to manual flagging text file.
        filter_kwargs: additional keyword arguments to be passed to DelayFilter.run_delay_filter()
    '''
    if baseline_list is not None and Nbls_per_load is not None:
        raise NotImplementedError("baseline loading and partial i/o not yet implemented.")
    hd = io.HERAData(datafile_list, filetype='uvh5', axis='blt')
    if baseline_list is not None and len(baseline_list) == 0:
        warnings.warn("Length of baseline list is zero."
                      "This can happen under normal circumstances when there are more files in datafile_list then baselines."
                      "in your dataset. Exiting without writing any output.", RuntimeWarning)
    else:
        if baseline_list is None:
            if len(hd.filepaths) > 1:
                baseline_list = list(hd.bls.values())[0]
            else:
                baseline_list = hd.bls
        if spw_range is None:
            spw_range = [0, hd.Nfreqs]
        freqs = hd.freq_array.flatten()[spw_range[0]:spw_range[1]]
        baseline_antennas = []
        for blpolpair in baseline_list:
            baseline_antennas += list(blpolpair[:2])
        baseline_antennas = np.unique(baseline_antennas).astype(int)
        if calfile_list is not None:
            cals = io.HERACal(calfile_list)
            cals.read(antenna_nums=baseline_antennas, frequencies=freqs)
        else:
            cals = None
        if polarizations is None:
            if len(hd.filepaths) > 1:
                polarizations = list(hd.pols.values())[0]
            else:
                polarizations = hd.pols
        if Nbls_per_load is None:
            Nbls_per_load = len(baseline_list)
        for i in range(0, len(baseline_list), Nbls_per_load):
            df = DelayFilter(hd, input_cal=cals)
            df.read(bls=baseline_list[i:i + Nbls_per_load], frequencies=freqs)
            if avg_red_bllens:
                df.avg_red_baseline_vectors()
            if external_flags is not None:
                df.apply_flags(external_flags, overwrite_flags=overwrite_flags)
            if flag_yaml is not None:
                df.apply_flags(flag_yaml, overwrite_flags=overwrite_flags, filetype='yaml')
            if factorize_flags:
                df.factorize_flags(time_thresh=time_thresh, inplace=True)
            df.run_delay_filter(cache_dir=cache_dir, read_cache=read_cache, write_cache=write_cache,
                                skip_flagged_edges=skip_flagged_edges, **filter_kwargs)
            df.write_filtered_data(res_outfilename=res_outfilename, CLEAN_outfilename=CLEAN_outfilename,
                                   filled_outfilename=filled_outfilename, partial_write=Nbls_per_load < len(baseline_list),
                                   clobber=clobber, add_to_history=add_to_history,
                                   extra_attrs={'Nfreqs': df.Nfreqs, 'freq_array': df.hd.freq_array})
            df.hd.data_array = None  # this forces a reload in the next loop

# ----------------------------------------
# Arg-parser for delay-filtering.
# can be linear or clean.
# ---------------------------------------

def delay_filter_argparser():
    '''
    Arg parser for commandline operation of delay filters.

    Parameters:
        mode, str : optional. Determines sets of arguments to load.
            can be 'clean', 'dayenu', or 'dpss_leastsq'.
        multifile: bool, optional.
            If True, add calfilelist and filelist
            arguments.
    Returns:
        argparser for delay-domain filtering for specified filtering mode
    '''
    a = vis_clean._filter_argparser()
    filt_options = a.add_argument_group(title='Options for the delay filter')
    filt_options.add_argument("--standoff", type=float, default=15.0, help='fixed additional delay beyond the horizon (default 15 ns)')
    filt_options.add_argument("--horizon", type=float, default=1.0, help='proportionality constant for bl_len where 1.0 (default) is the horizon\
                              (full light travel time)')
    filt_options.add_argument("--min_dly", type=float, default=0.0, help="A minimum delay threshold [ns] used for filtering.")
    return a


import argparse, bisect
import collections
import cPickle as pickle
import hashlib, imp
import json, random
import os, shutil, sys
import tempfile
import time

import numpy as np
import tensorflow as tf
import tf_utils
import _jsonnet

label_names = ['trigger_chans', 'trigger_funcs', 'action_chans',
               'action_funcs']

class keydefaultdict(collections.defaultdict):
  def __missing__(self, key):
    if self.default_factory is None:
      raise KeyError(key)
    else:
      ret = self[key] = self.default_factory(key)
      return ret

def without_name(d):
  d = d.copy()
  del d['name']
  return d

def variable_summaries(var):
  name = var.name
  with tf.name_scope("summaries"):
    mean = tf.reduce_mean(var)
    tf.scalar_summary('vars-mean/' + name, mean)
    with tf.name_scope('stddev'):
      stddev = tf.sqrt(tf.reduce_sum(tf.square(var - mean)))
    tf.scalar_summary('vars-sttdev/' + name, stddev)
    tf.scalar_summary('vars-max/' + name, tf.reduce_max(var))
    tf.scalar_summary('vars-min/' + name, tf.reduce_min(var))
    tf.histogram_summary(name, var)

def make_train_op(global_step, loss, optim_config):
  lr = getattr(tf.train, optim_config['lr_scheduler']['name'])(
      global_step=global_step,
      **without_name(optim_config['lr_scheduler']))
  tf.scalar_summary('lr', lr)
  tvars = tf.trainable_variables()
  #  for var in tvars:
  #    variable_summaries(var)
  grads = tf.gradients(loss, tvars)
  if optim_config['max_grad_norm'] > 0:
    grads, grads_norm = tf.clip_by_global_norm(grads,
                                               optim_config['max_grad_norm'])
  else:
    grads_norm = tf.global_norm(grads)
  tf.scalar_summary('grads/norm', grads_norm)
  optimizer = getattr(tf.train, optim_config['optimizer']['name'])(
      learning_rate=lr,
      **without_name(optim_config['optimizer']))
  return optimizer.apply_gradients(zip(grads, tvars), global_step=global_step)

class IFTTTModel(object):
  def __init__(self, optim_config, model_config, arch_config, num_labels,
               label_types, vocab_size, memory_size, is_training):
    batch_size = optim_config['batch_size']
    self.ids = tf.placeholder(tf.int32, name='input',
                              shape=[batch_size, memory_size])
    self.ids_lengths = tf.placeholder(
        tf.int32, name='input_lengths', shape=[batch_size])
    self.labels = tf.placeholder(tf.int32,
                                 name='labels',
                                 shape=[batch_size, len(num_labels)])
    self.expanded_labels = tf.placeholder(
        tf.int32,
        name='expanded_labels',
        # batch size x seq length x label types
        shape=[optim_config['batch_size'], memory_size, len(num_labels)])

    with tf.name_scope("input_length_summaries"):
        max_length = tf.reduce_max(self.ids_lengths)
        tf.scalar_summary("input_length/max", max_length)
        tf.scalar_summary("input_length/min", tf.reduce_min(self.ids_lengths))
        mean_length = tf.reduce_mean(tf.cast(self.ids_lengths, tf.float32))
        tf.scalar_summary("input_length/mean", mean_length)
        tf.scalar_summary("input_length/fill_ratio",
                          mean_length / tf.cast(max_length, tf.float32))

    ids = tf.minimum(vocab_size - 1, self.ids, name='make_unk')
    if arch_config:
      label_groupings = arch_config['label_groupings']

      if arch_config['share_word_embeddings']:
        embedding = tf.get_variable(
            'embedding_for_' + '_'.join(str(i) for i in label_types),
            [vocab_size, model_config['embedding_size']])
        rnn_inputs = [tf.nn.embedding_lookup(embedding, ids)
                      ] * len(label_groupings)
      else:
        rnn_inputs = []
        for grouping in label_groupings:  #enumerate(label_groupings):
          embedding = tf.get_variable(
              'embedding_for_' + '_'.join(str(i) for i in grouping),
              [vocab_size, model_config['embedding_size']])
          rnn_inputs.append(tf.nn.embedding_lookup(embedding, ids))

      outputs = []
      for grouping, rnn_input in zip(label_groupings, rnn_inputs):
        with tf.variable_scope('labels_' + '_'.join(str(i) for i in grouping)):
          if model_config['name'] == 'rnn':
            outputs.append(tf_utils.models.rnn(
              rnn_input,
              self.ids_lengths,
              getattr(tf.nn.rnn_cell, model_config['cell_type']),
              int(model_config['num_layers']),
              int(model_config['num_units']),
              model_config['keep_prob'],
              is_training,
              bidirectional=model_config['bidirectional'],
              debug=False))
          elif model_config['name'] == 'Dict':
            outputs.append((rnn_input,tf.reduce_sum(rnn_input,1)))

    else:
      label_groupings = [label_types]
      embedding = tf.get_variable('embedding',
                                  [vocab_size, model_config['embedding_size']])
      rnn_input = tf.nn.embedding_lookup(embedding, ids)
      if model_config['name'] == 'rnn':
        outputs = [tf_utils.models.rnn(
          rnn_input,
          self.ids_lengths,
          getattr(tf.nn.rnn_cell, model_config['cell_type']),
          int(model_config['num_layers']),
          int(model_config['num_units']),
          model_config['keep_prob'],
          is_training,
          bidirectional=model_config['bidirectional'],
          debug=False)]
      else:
        outputs = [(rnn_input,tf.reduce_sum(rnn_input,1))]

    losses = []
    all_preds = []
    self.all_probs = []
    decoder_type = model_config.get('decoder', 'standard')

    for i, num_classes in enumerate(num_labels):
        with tf.variable_scope('label_{}'.format(label_types[i])):
            for (states, output), group in zip(outputs, label_groupings):
                if label_types[i] in group:
                    break
            else:
                raise ValueError('label_types[i] {} not found in {}'.format(
                    label_types[i], label_groupings))

            if decoder_type in ['LA', 'attention']:
                if model_config['name'].find('rnn') != -1:
                    vec_size = model_config['embedding_size'] * 2
                else:
                    vec_size = model_config['embedding_size']

                PREP = tf.get_variable("PREP", [1, vec_size])
                softmax_w = tf.get_variable(
                    "softmax_w", [vec_size, num_classes],
                    initializer = tf.contrib.layers.xavier_initializer())
                TA = tf.get_variable("TA", [memory_size, vec_size])
                m = states + TA

                if decoder_type == 'LA':
                    B = tf.get_variable("B", [vec_size, memory_size])
                    m_t = tf.reshape(m,
                                     [optim_config['batch_size'] * memory_size,
                                      vec_size])
                    d_t = tf.matmul(m_t, B)
                    d_softmax = tf.nn.softmax(d_t)
                    d = tf.reshape(d_softmax, [optim_config['batch_size'],
                                               memory_size,
                                               memory_size])
                    dotted_prep = tf.reduce_sum(states * PREP, 2)
                else: # 'attention'
                    dotted_prep = tf.reduce_sum(m * PREP, 1)

                probs_prep = tf.nn.softmax(dotted_prep)
                preps = []
                preps.append(probs_prep)
                for _ in xrange(1):
                    probs_prep = preps[-1]
                    if decoder_type == 'LA':
                        probs_prep_temp = tf.expand_dims(probs_prep, -1)
                        probs_temp = tf.batch_matmul(d, probs_prep_temp)
                        probs = tf.squeeze(probs_temp)
                        output_probs = tf.nn.l2_normalize(probs, 1)
                    else: # 'attention'
                        probs_prep_temp = tf.transpose(
                            tf.expand_dims(probs_prep, -1), [0, 2, 1])
                        dotted = tf.reduce_sum(m * probs_prep_temp, 2)
                        output_probs = tf.nn.softmax(dotted)
                    preps.append(output_probs)

                probs_temp = tf.expand_dims(preps[-1], 1)
                c_temp = tf.transpose(m, [0, 2, 1])
                o_k = tf.reduce_sum(c_temp * probs_temp, 2)
                logits = tf.matmul(o_k, softmax_w)
            else:
                w = tf.get_variable('softmax_w', [output.get_shape()[1],
                                              num_classes])
                b = tf.get_variable('softmax_b', [num_classes])
                logits = tf.nn.xw_plus_b(output, w, b)

            preds = tf.cast(tf.argmax(logits, 1), tf.int32)
            all_preds.append(preds)
            probs = tf.nn.softmax(logits)
            self.all_probs.append(probs)
            loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
                logits, self.labels[:, i]))
            losses.append(loss)
            tf.scalar_summary(
                'train_acc/{}'.format(label_names[label_types[i]]),
                tf.reduce_mean(tf.cast(
                    tf.equal(preds, self.labels[:, i]), tf.float32)))

    self.preds = tf.transpose(tf.pack(all_preds))
    self.loss = tf.add_n(losses)

    tf.scalar_summary('loss', self.loss)

    self.global_step = tf.Variable(0, trainable=False, name='global_step')
    if is_training:
      self.train_op = make_train_op(self.global_step, self.loss, optim_config)

    self.summaries = tf.merge_all_summaries()


import numpy as np
import pandas as pd

from .BaseEmbed import BaseEmbed

pd.set_option('display.max_rows', 500)
pd.set_option('display.max_columns', 500)
pd.set_option('display.width', 1000)
pd.options.display.width = 0
from tqdm import tqdm

from ..logging import getLogger
from ..utils import unit_length, clean_text, is_1d_array
from typing import List, Union

Feature = List[Union[List[str], str]]

class FlairGlove100Embed(BaseEmbed):
    def __init__(self, n_dims=100, make_unit_length=True):
        from flair.embeddings import WordEmbeddings, DocumentPoolEmbeddings
        super().__init__(n_dims=n_dims, make_unit_length=make_unit_length)
        embeddings = [WordEmbeddings('glove')]
        self.embeddings = DocumentPoolEmbeddings(embeddings, fine_tune_mode='none')
        self.log = getLogger(type(self).__name__)

    def get_sentence_vector(self, text):
        from flair.embeddings import Sentence
        sentence = Sentence(clean_text(text))
        _ = self.embeddings.embed(sentence)
        a = sentence.get_embedding()
        result = a.detach().cpu().numpy()
        if np.sum(result[0:5]) == 0:
            result = np.random.randn(self.n_dims)
        return result

    def fit(self, feature: Feature, **kwargs):
        super().fit(feature, **kwargs)
        self.log.debug("End Fitting FlairEmbedding")
        return

    def transform(self, feature: Feature, **kwargs) -> np.ndarray:
        self.log.debug("Start Transform TextEmbedding...")
        outputs = np.vstack([np.array([self.get_sentence_vector(i) for i in t]).mean(0) if is_1d_array(
            t) else self.get_sentence_vector(t) for t in tqdm(feature)])
        outputs = unit_length(outputs, axis=1) if self.make_unit_length else outputs
        self.log.debug("End Transform TextEmbedding")
        return self.check_output_dims(outputs, feature)

class FlairGlove100AndBytePairEmbed(FlairGlove100Embed):
    def __init__(self, n_dims=200, make_unit_length=True):
        from flair.embeddings import WordEmbeddings, DocumentPoolEmbeddings, BytePairEmbeddings
        super().__init__(n_dims=n_dims, make_unit_length=make_unit_length)
        embeddings = [WordEmbeddings('glove'), BytePairEmbeddings('en')]
        self.embeddings = DocumentPoolEmbeddings(embeddings, fine_tune_mode='none')
        self.log = getLogger(type(self).__name__)


from curses import raw
from multiprocessing.spawn import old_main_modules
from pyexpat import features
from black import out
import mmcv
from mmcv.runner import force_fp32
from argparse import ArgumentParser
import torch
import numpy as np
from torch import nn as nn
from mmdet3d.models.voxel_encoders.utils import VFELayer, get_paddings_indicator
from mmcv.runner import load_checkpoint
from mmdet3d.models import build_model
from mmdet3d.ops import DynamicScatter
from mmcv.parallel import collate, scatter
from mmdet3d.core import bbox3d2result
from mmdet3d.core.bbox import get_box_type
from copy import deepcopy
from mmdet3d.datasets.pipelines import Compose
from mmdet3d.models.middle_encoders.pillar_scatter import PointPillarsScatter
from mmdet3d.models.backbones import SECOND
from mmdet3d.models.necks import SECONDFPN
import cv2
from mmdet.models.builder import (HEADS)
from mmdet.core import build_bbox_coder
from mmcv.runner import BaseModule
from mmcv.cnn import build_conv_layer, build_norm_layer

class HardVFE(nn.Module):
    def __init__(self,
                 cfg,
                 return_point_feats=False,
                 fusion_layer=None,
                 prune_index = None,
                 test=False):
        super().__init__()
        if prune_index:
            self.feat_channels = prune_index
        else:
            self.feat_channels = cfg.model['pts_voxel_encoder']['feat_channels']
        assert len(self.feat_channels) > 0
        self.in_channels = cfg.model['pts_voxel_encoder']['in_channels']
        self._with_distance = cfg.model['pts_voxel_encoder']['with_distance']
        self._with_cluster_center = cfg.model['pts_voxel_encoder']['with_cluster_center']
        self._with_voxel_center = cfg.model['pts_voxel_encoder']['with_voxel_center']
        self.return_point_feats = return_point_feats
        self.fp16_enabled = False
        if self._with_cluster_center:
            self.in_channels += 3
        if self._with_voxel_center:
            self.in_channels += 3
        if self._with_distance:
            self.in_channels += 1

        # Need pillar (voxel) size and x/y offset to calculate pillar offset
        self.vx = cfg['voxel_size'][0]
        self.vy = cfg['voxel_size'][1]
        self.vz = cfg['voxel_size'][2]
        self.point_cloud_range = cfg.model['pts_voxel_encoder']['point_cloud_range']
        self.x_offset = self.vx / 2 + self.point_cloud_range[0]
        self.y_offset = self.vy / 2 + self.point_cloud_range[1]
        self.z_offset = self.vz / 2 + self.point_cloud_range[2]
        self.norm_cfg = cfg.model['pts_voxel_encoder']['norm_cfg']
        self.scatter = DynamicScatter(cfg['voxel_size'], self.point_cloud_range, True)
        self.fusion_layer = None
        self.test = test

        feat_channels = [self.in_channels] + list(self.feat_channels)
        vfe_layers = []
        for i in range(len(feat_channels) - 1):
            in_filters = feat_channels[i]
            out_filters = feat_channels[i + 1]
            print("in_filter: {:d}, out_filter: {:d}.".format(in_filters, out_filters))
            if i > 0:
                in_filters *= 2
            # TODO: pass norm_cfg to VFE
            # norm_name, norm_layer = build_norm_layer(norm_cfg, out_filters)
            if i == (len(feat_channels) - 2):
                cat_max = False
                max_out = True
                if fusion_layer:
                    max_out = False
            else:
                max_out = True
                cat_max = True
            vfe_layers.append(
                VFELayer(
                    in_filters,
                    out_filters,
                    norm_cfg=self.norm_cfg,
                    max_out=max_out,
                    cat_max=cat_max))
            self.vfe_layers = nn.ModuleList(vfe_layers)
        self.num_vfe = len(vfe_layers)

    @force_fp32(out_fp16=True)
    def forward(self, *args):
        if self.training:
            return self.forward_train(*args)
        elif self.test:
            return self.forward_train(*args)
        else:
            return self.forward_test(*args)

    @force_fp32(out_fp16=True)
    def forward_test(self,voxel_feats):
        for i, vfe in enumerate(self.vfe_layers):
            voxel_feats = vfe(voxel_feats)
        return voxel_feats

    @force_fp32(out_fp16=True)
    def forward_train(self,
                    features,
                    num_points,
                    coors,
                    img_feats=None,
                    img_metas=None):
        """Forward functions.

        Args:
            features (torch.Tensor): Features of voxels, shape is MxNxC.
            num_points (torch.Tensor): Number of points in each voxel.
            coors (torch.Tensor): Coordinates of voxels, shape is Mx(1+NDim).
            img_feats (list[torch.Tensor], optional): Image fetures used for
                multi-modality fusion. Defaults to None.
            img_metas (dict, optional): [description]. Defaults to None.

        Returns:
            tuple: If `return_point_feats` is False, returns voxel features and
                its coordinates. If `return_point_feats` is True, returns
                feature of each points inside voxels.
        """
        features_ls = [features]
        # Find distance of x, y, and z from cluster center
        if self._with_cluster_center:
            points_mean = (
                features[:, :, :3].sum(dim=1, keepdim=True) /
                num_points.type_as(features).view(-1, 1, 1))
            # TODO: maybe also do cluster for reflectivity
            f_cluster = features[:, :, :3] - points_mean
            features_ls.append(f_cluster)

        # Find distance of x, y, and z from pillar center
        if self._with_voxel_center:
            f_center = features.new_zeros(
                size=(features.size(0), features.size(1), 3))
            f_center[:, :, 0] = features[:, :, 0] - (
                coors[:, 3].type_as(features).unsqueeze(1) * self.vx +
                self.x_offset)
            f_center[:, :, 1] = features[:, :, 1] - (
                coors[:, 2].type_as(features).unsqueeze(1) * self.vy +
                self.y_offset)
            f_center[:, :, 2] = features[:, :, 2] - (
                coors[:, 1].type_as(features).unsqueeze(1) * self.vz +
                self.z_offset)
            features_ls.append(f_center)

        if self._with_distance:
            points_dist = torch.norm(features[:, :, :3], 2, 2, keepdim=True)
            features_ls.append(points_dist)

        # Combine together feature decorations
        voxel_feats = torch.cat(features_ls, dim=-1)
        # The feature decorations were calculated without regard to whether
        # pillar was empty.
        # Need to ensure that empty voxels remain set to zeros.
        voxel_count = voxel_feats.shape[1]
        mask = get_paddings_indicator(num_points, voxel_count, axis=0)
        voxel_feats *= mask.unsqueeze(-1).type_as(voxel_feats)
        raw_feats = voxel_feats
        for i, vfe in enumerate(self.vfe_layers):
            voxel_feats = vfe(voxel_feats)

        if (self.fusion_layer is not None and img_feats is not None):
            voxel_feats = self.fusion_with_mask(features, mask, voxel_feats,
                                                coors, img_feats, img_metas)
        return voxel_feats, raw_feats

    def fusion_with_mask(self, features, mask, voxel_feats, coors, img_feats,
                         img_metas):
        """Fuse image and point features with mask.

        Args:
            features (torch.Tensor): Features of voxel, usually it is the
                values of points in voxels.
            mask (torch.Tensor): Mask indicates valid features in each voxel.
            voxel_feats (torch.Tensor): Features of voxels.
            coors (torch.Tensor): Coordinates of each single voxel.
            img_feats (list[torch.Tensor]): Multi-scale feature maps of image.
            img_metas (list(dict)): Meta information of image and points.

        Returns:
            torch.Tensor: Fused features of each voxel.
        """
        # the features is consist of a batch of points
        batch_size = coors[-1, 0] + 1
        points = []
        for i in range(batch_size):
            single_mask = (coors[:, 0] == i)
            points.append(features[single_mask][mask[single_mask]])

        point_feats = voxel_feats[mask]
        point_feats = self.fusion_layer(img_feats, points, point_feats,
                                        img_metas)

        voxel_canvas = voxel_feats.new_zeros(
            size=(voxel_feats.size(0), voxel_feats.size(1),
                  point_feats.size(-1)))
        voxel_canvas[mask] = point_feats
        out = torch.max(voxel_canvas, dim=1)[0]

        return out

class Backbone(nn.Module):
    def __init__(self, cfg, pts_backbone, pts_neck, pts_bbox_head):
        super().__init__()
        self.pts_backbone = pts_backbone
        self.pts_neck = pts_neck
        self.cfg = cfg
        self.pts_bbox_head = pts_bbox_head
        self.bbox_coder = build_bbox_coder(cfg['pts_bbox_head']['bbox_coder'])
        self.box_code_size = self.bbox_coder.code_size
        loss_cls = cfg['pts_bbox_head']['loss_cls']
        self.use_sigmoid_cls = loss_cls.get('use_sigmoid', False)
        self.dir_limit_offset = cfg['pts_bbox_head']['dir_limit_offset']
        self.dir_offset = cfg['pts_bbox_head']['dir_offset']

    def forward(self, input):
        print(input.shape)
        x = input[:22*200*200]
        x = x.reshape(-1, 22, 200, 200)
        anchors = input[22*200*200:]
        anchors = anchors.reshape(-1, 9)

        x = self.pts_backbone(x)
        x = self.pts_neck(x)
        x = self.pts_bbox_head(x)
        bboxes, scores, labels = self.get_boxes(*x, anchors)
        return bboxes, scores, labels

    def decode(self, anchors, deltas):
        """Apply transformation `deltas` (dx, dy, dz, dw, dh, dl, dr, dv*) to
        `boxes`.

        Args:
            anchors (torch.Tensor): Parameters of anchors with shape (N, 7).
            deltas (torch.Tensor): Encoded boxes with shape
                (N, 7+n) [x, y, z, w, l, h, r, velo*].

        Returns:
            torch.Tensor: Decoded boxes.
        """
        cas, cts = [], []
        xa, ya, za, wa, la, ha, ra, *cas = torch.split(anchors, 1, dim=-1)
        xt, yt, zt, wt, lt, ht, rt, *cts = torch.split(deltas, 1, dim=-1)
        za = za + ha / 2
        diagonal = torch.sqrt(la**2 + wa**2)
        xg = xt * diagonal + xa
        yg = yt * diagonal + ya
        zg = zt * ha + za

        lg = torch.exp(lt) * la
        wg = torch.exp(wt) * wa
        hg = torch.exp(ht) * ha
        rg = rt + ra
        zg = zg - hg / 2
        cgs = [t + a for t, a in zip(cts, cas)]
        return torch.cat([xg, yg, zg, wg, lg, hg, rg, *cgs], dim=-1)

    def get_boxes(self, cls_scores, bbox_preds, dir_cls_preds, anchors):
        cls_score = cls_scores[0].detach()
        bbox_pred = bbox_preds[0].detach()
        dir_cls_pred = dir_cls_preds[0].detach()
        cls_score = cls_scores[0].detach()
        bbox_pred = bbox_preds[0].detach()
        dir_cls_pred = dir_cls_preds[0].detach()
        bboxes, scores, labels = self.get_bboxes_single(cls_score, bbox_pred,
                                            dir_cls_pred, anchors)
        return bboxes, scores, labels

    def get_bboxes_single(self, cls_scores, bbox_preds, dir_cls_preds, anchors):
        cls_score = cls_scores[0]
        bbox_pred = bbox_preds[0]
        dir_cls_pred = dir_cls_preds[0]
        dir_cls_pred = dir_cls_pred.permute(1, 2, 0).reshape(-1, 2)
        dir_cls_score = torch.max(dir_cls_pred, dim=-1)[1]
        cls_score = cls_score.permute(1, 2, 0).reshape(-1, self.pts_bbox_head.num_classes)
        if self.use_sigmoid_cls:
            scores = cls_score.sigmoid()
        else:
            scores = cls_score.softmax(-1)
        bbox_pred = bbox_pred.permute(1, 2, 0).reshape(-1, self.box_code_size)

        nms_pre = 1000
        if self.use_sigmoid_cls:
            max_scores, _ = scores.max(dim=1)
        else:
            max_scores, _ = scores[:, :-1].max(dim=1)
        _, topk_inds = max_scores.topk(nms_pre)
        topk_inds = topk_inds.long()
        anchors = anchors[topk_inds, :]
        bbox_pred = bbox_pred[topk_inds, :]
        scores = scores[topk_inds, :]
        dir_cls_score = dir_cls_score[topk_inds]
        bboxes = self.decode(anchors, bbox_pred)
        return bboxes, scores, dir_cls_score

class PRUNED_SECOND(BaseModule):
    """Backbone network for SECOND/PointPillars/PartA2/MVXNet.

    Args:
        in_channels (int): Input channels.
        out_channels (list[int]): Output channels for multi-scale feature maps.
        layer_nums (list[int]): Number of layers in each stage.
        layer_strides (list[int]): Strides of each stage.
        norm_cfg (dict): Config dict of normalization layers.
        conv_cfg (dict): Config dict of convolutional layers.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 layer_nums=[3, 5, 5],
                 layer_strides=[2, 2, 2],
                 norm_cfg=dict(type='BN', eps=1e-3, momentum=0.01),
                 conv_cfg=dict(type='Conv2d', bias=False),
                 init_cfg=None,
                 pretrained=None):
        super().__init__()
        self.layer_nums = layer_nums

        in_filters = [in_channels, *out_channels[:-1]]
        # note that when stride > 1, conv2d with same padding isn't
        # equal to pad-conv2d. we should use pad-conv2d.
        blocks = []
        count = 0
        for i, layer_num in enumerate(layer_nums):
            block = [
                build_conv_layer(
                    conv_cfg,
                    in_filters[count],
                    out_channels[count],
                    3,
                    stride=layer_strides[i],
                    padding=1),
                build_norm_layer(norm_cfg, out_channels[count])[1],
                nn.ReLU(inplace=True),
            ]
            count += 1
            for j in range(layer_num):
                block.append(
                    build_conv_layer(
                        conv_cfg,
                        out_channels[count-1],
                        out_channels[count],
                        3,
                        padding=1))
                block.append(build_norm_layer(norm_cfg, out_channels[count])[1])
                block.append(nn.ReLU(inplace=True))
                count += 1

            block = nn.Sequential(*block)
            blocks.append(block)

        self.blocks = nn.ModuleList(blocks)
        self.init_cfg = dict(type='Kaiming', layer='Conv2d')

    def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor): Input with shape (N, C, H, W).

        Returns:
            tuple[torch.Tensor]: Multi-scale features.
        """
        outs = []
        for i in range(len(self.blocks)):
            x = self.blocks[i](x)
            outs.append(x)
        return tuple(outs)

def parse_model(model):
    for name, parameters in model.named_parameters():
        print(name, ':', parameters.size())

def visualize_feature(feature, i=0, file = 'visualize.jpg'):
    if isinstance(feature, torch.Tensor):
        feature = feature.cpu().detach().numpy()
    feature = feature.squeeze(0)
    img = feature[i,:,:] * 255
    cv2.imwrite(file, img)

def get_thre(model, module = nn.BatchNorm2d, ratio = 0.7):
    total = 0
    for m in model.modules():
        if isinstance(m, module):
            total += m.weight.data.shape[0]
    bn = torch.zeros(total)
    index = 0
    for m in model.modules():
        if isinstance(m, module):
            size = m.weight.data.shape[0]
            bn[index:(index + size)] = m.weight.data.abs().clone()
            index += size
    y, i = torch.sort(bn)
    thre_index = int(total * ratio)
    thre = y[thre_index]
    return thre

def prune(model, thre, module = nn.BatchNorm2d):
    pruned = 0
    cfg_index = []
    cfg_mask = []
    index = 0
    for m in model.modules():
        if isinstance(m, module):
            weight_copy = m.weight.data.abs().clone()
            mask = weight_copy.cpu().ge(thre).float().cuda()
            pruned = pruned + mask.shape[0] - torch.sum(mask)
            m.weight.data.mul_(mask)
            m.bias.data.mul_(mask)
            cfg_index.append(int(torch.sum(mask)))
            cfg_mask.append(mask.clone())
            print('layer index: {:d} \t total channel: {:d} \t remaining channel: {:d}'.format(index, mask.shape[0], int(torch.sum(mask))))
        elif isinstance(m, nn.MaxPool2d):
            cfg_index.append('M')
        index += 1
    return cfg_index, cfg_mask

def convert_SyncBN(config):
    """Convert config's naiveSyncBN to BN.

    Args:
         config (str or :obj:`mmcv.Config`): Config file path or the config
            object.
    """
    if isinstance(config, dict):
        for item in config:
            if item == 'norm_cfg':
                config[item]['type'] = config[item]['type']. \
                                    replace('naiveSyncBN', 'BN')
            else:
                convert_SyncBN(config[item])

def build_and_prune_vfe_model(config, checkpoint=None, device='cuda:0'):
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        f'but got {type(config)}')
    config.model.pretrained = None
    convert_SyncBN(config.model)
    config.model.train_cfg = None
    # original model
    model = build_model(config.model, test_cfg=config.get('test_cfg'))
    if checkpoint is not None:
        checkpoint_load = load_checkpoint(model, checkpoint, map_location='cpu')
        if 'CLASSES' in checkpoint_load['meta']:
            model.CLASSES = checkpoint_load['meta']['CLASSES']
        else:
            model.CLASSES = config.class_names
        if 'PALETTE' in checkpoint_load['meta']:  # 3D Segmentor
            model.PALETTE = checkpoint_load['meta']['PALETTE']
    model.cfg = config  # save the config in the model for convenience
    torch.cuda.set_device(device)
    model.to(device)
    model.eval()

    cfg_index, cfg_mask = prune(model, get_thre(model, module=nn.BatchNorm1d, ratio=0.7), nn.BatchNorm1d)
    # parse_model(model.pts_voxel_encoder)

    # VFE model
    pts_voxel_encoder = HardVFE(config, prune_index=cfg_index)
    pts_voxel_encoder.to(device).eval()
    ori_modules = list(model.pts_voxel_encoder.modules())
    prune_modules = list(pts_voxel_encoder.modules())
    layer_id_in_cfg = 0

    start_mask = torch.ones(pts_voxel_encoder.in_channels)
    end_mask = cfg_mask[layer_id_in_cfg]
    linear_count = 0

    for layer_id in range(len(prune_modules)):
        m0 = ori_modules[layer_id]
        m1 = prune_modules[layer_id]
        if isinstance(m0, nn.BatchNorm1d):
            idx1 = np.squeeze(np.argwhere(np.asarray(end_mask.cpu().numpy())))
            if idx1.size == 1:
                idx1 = np.resize(idx1, (1,))
            m1.weight.data = m0.weight.data[idx1.tolist()].clone()
            m1.bias.data = m0.bias.data[idx1.tolist()].clone()
            m1.running_mean = m0.running_mean[idx1.tolist()].clone()
            m1.running_var = m0.running_var[idx1.tolist()].clone()
        elif isinstance(m0, nn.Linear):
            if isinstance(ori_modules[layer_id-1], nn.BatchNorm1d):
                idx0 = np.squeeze(np.argwhere(np.asarray(start_mask.cpu().numpy())))
                idx1 = np.squeeze(np.argwhere(np.asarray(end_mask.cpu().numpy())))
                # print('In shape: {:d}, Out shape {:d}.'.format(idx0.size, idx1.size))
                if idx0.size == 1:
                    idx0 = np.resize(idx0, (1,))
                if idx1.size == 1:
                    idx1 = np.resize(idx1, (1,))
                if linear_count == 1:   ## VFElayer 的repeat操作
                    repeated = np.array(idx0 + int(m0.weight.data.size()[1] / 2))
                    idx0 = np.concatenate((idx0, repeated))
                w1 = m0.weight.data[:, idx0.tolist()].clone()
                w1 = w1[idx1.tolist(),:].clone()
                m1.weight.data = w1.clone()
                linear_count += 1
                layer_id_in_cfg += 1
                start_mask = end_mask.clone()
                if layer_id_in_cfg < len(cfg_mask):
                    end_mask = cfg_mask[layer_id_in_cfg]

    # print('-----------------------')
    # parse_model(pts_voxel_encoder)
    return model, pts_voxel_encoder, cfg_mask

def build_backbone_model(model, in_channels, cfg_mask, device='cuda:0'):
    pts_backbone = SECOND(in_channels, out_channels=[64, 128, 256])
    pts_backbone.to(device).eval()
    ori_modules = list(model.pts_backbone.modules())
    prune_modules = list(pts_backbone.modules())
    flag = False
    for layer_id in range(len(prune_modules)):
        m0 = ori_modules[layer_id]
        m1 = prune_modules[layer_id]
        if not flag :
            if isinstance(m0, nn.Conv2d):       #only the first conv layer is different
                idx = np.squeeze(np.argwhere(np.asarray(cfg_mask.cpu().numpy())))
                if idx.size == 1:
                    idx = np.resize(idx, (1,))
                w1 = m0.weight.data[:,idx.tolist(),:,:].clone()
                m1.weight.data = w1.clone()
                flag = True
                continue

        if isinstance(m0, nn.Conv2d):
            m1.weight.data = m0.weight.data.clone()
        if isinstance(m0, nn.BatchNorm2d):
            m1.weight.data = m0.weight.data.clone()
            m1.bias.data = m0.bias.data.clone()
            m1.running_mean = m0.running_mean.clone()
            m1.running_var = m0.running_var.clone()

    return pts_backbone

def prune_backbone_model(model, in_channel, ratio=0.1, device='cuda:0'):
    cfg_index, cfg_mask = prune(model, get_thre(model, ratio=ratio))    # bn2d
    pruned_backbone = PRUNED_SECOND(in_channel, cfg_index)
    pruned_backbone.to(device).eval()
    # parse_model(pruned_backbone)
    layer_id_in_cfg = 0

    output_layer_mask = []
    id = -1
    for layer_num in pruned_backbone.layer_nums:
        id += 1 + layer_num
        output_layer_mask.append(cfg_mask[id])

    start_mask = torch.ones(in_channel)
    end_mask = cfg_mask[layer_id_in_cfg]

    ori_modules = list(model.modules())
    prune_modules = list(pruned_backbone.modules())
    for layer_id in range(len(ori_modules)):
        m0 = ori_modules[layer_id]
        m1 = prune_modules[layer_id]
        if isinstance(m0, nn.BatchNorm2d):
            idx1 = np.squeeze(np.argwhere(np.asarray(end_mask.cpu().numpy())))
            # print('shape {:d}.'.format(idx1.size))
            if idx1.size == 1:
                idx1 = np.resize(idx1, (1,))
            m1.weight.data = m0.weight.data[idx1.tolist()].clone()
            m1.bias.data = m0.bias.data[idx1.tolist()].clone()
            m1.running_mean = m0.running_mean[idx1.tolist()].clone()
            m1.running_var = m0.running_var[idx1.tolist()].clone()
            layer_id_in_cfg += 1
            start_mask = end_mask.clone()
            if layer_id_in_cfg < len(cfg_mask):
                end_mask = cfg_mask[layer_id_in_cfg]
        elif isinstance(m0, nn.Conv2d):
            if isinstance(ori_modules[layer_id+1], nn.BatchNorm2d):
                idx0 = np.squeeze(np.argwhere(np.asarray(start_mask.cpu().numpy())))
                idx1 = np.squeeze(np.argwhere(np.asarray(end_mask.cpu().numpy())))
                # print('In shape: {:d}, Out shape {:d}.'.format(idx0.size, idx1.size))
                if idx0.size == 1:
                    idx0 = np.resize(idx0, (1,))
                if idx1.size == 1:
                    idx1 = np.resize(idx1, (1,))
                w1 = m0.weight.data[:, idx0.tolist(), :, :].clone()
                w1 = w1[idx1.tolist(), :, :, :].clone()
                m1.weight.data = w1.clone()
    return pruned_backbone, output_layer_mask

def build_neck_model(model, in_channels, cfg_mask, device='cuda:0'):
    pts_neck = SECONDFPN(in_channels, out_channels=[128, 128, 128])
    pts_neck.to(device).eval()
    ori_modules = list(model.pts_neck.modules())
    prune_modules = list(pts_neck.modules())
    cfg_mask_id = 0
    for layer_id in range(len(prune_modules)):
        m0 = ori_modules[layer_id]
        m1 = prune_modules[layer_id]
        if isinstance(m0, nn.ConvTranspose2d):       # 3 channels , the first convtranspose layer in every channel is different
            idx = np.squeeze(np.argwhere(np.asarray(cfg_mask[cfg_mask_id].cpu().numpy())))
            if idx.size == 1:
                idx = np.resize(idx, (1,))
            w1 = m0.weight.data[idx.tolist(),:,:,:].clone()
            m1.weight.data = w1.clone()
            cfg_mask_id += 1
        if isinstance(m0, nn.BatchNorm2d):
            m1.weight.data = m0.weight.data.clone()
            m1.bias.data = m0.bias.data.clone()
            m1.running_mean = m0.running_mean.clone()
            m1.running_var = m0.running_var.clone()
    return pts_neck

def preprocess(model, pcd):
    """Inference point cloud with the detector.

    Args:
        model (nn.Module): The loaded detector.
        pcd (str): Point cloud files.

    Returns:
        tuple: Predicted results and data from pipeline.
    """
    cfg = model.cfg
    device = next(model.parameters()).device  # model device
    # build the data pipeline
    test_pipeline = deepcopy(cfg.data.test.pipeline)
    test_pipeline = Compose(test_pipeline)
    box_type_3d, box_mode_3d = get_box_type(cfg.data.test.box_type_3d)
    data = dict(
        pts_filename=pcd,
        box_type_3d=box_type_3d,
        box_mode_3d=box_mode_3d,
        # for ScanNet demo we need axis_align_matrix
        ann_info=dict(axis_align_matrix=np.eye(4)),
        sweeps=[],
        # set timestamp = 0
        timestamp=[0],
        img_fields=[],
        bbox3d_fields=[],
        pts_mask_fields=[],
        pts_seg_fields=[],
        bbox_fields=[],
        mask_fields=[],
        seg_fields=[])
    data = test_pipeline(data)
    data = collate([data], samples_per_gpu=1)
    if next(model.parameters()).is_cuda:
        # scatter to specified GPU
        data = scatter(data, [device.index])[0]
    else:
        # this is a workaround to avoid the bug of MMDataParallel
        data['img_metas'] = data['img_metas'][0].data
        data['points'] = data['points'][0].data
    return data

# get api
def get_prune_model(cfg):
    device = "cuda:0"
    # checkpoint = "checkpoints/hv_pointpillars_secfpn_sbn-all_4x8_2x_nus-3d_20200620_230725-0817d270.pth"
    # pcd = "/data/testdata/QT_js_for_nuscenes_10Hz/1606813512902872600.bin"
    checkpoint = "work_dirs_0218/hv_pointpillars_secfpn_sbn-all_4x8_2x_custom-3d_0302/epoch_24.pth"
    # pcd = "data/2022-02-18/testing/velodyne/1645148902202493613.bin"
    pcd = "demo/1645179542771259609.bin"
    model, prune_pts_vfe, cfg_mask = build_and_prune_vfe_model(cfg, checkpoint, device=device)

    data = preprocess(model, pcd)
    pts = data['points'][0]
    voxels, num_points, coors = model.voxelize(pts)
    _, raw_feats = model.pts_voxel_encoder(voxels, num_points, coors)
    v_features = prune_pts_vfe(raw_feats)
    batch_size = coors[-1, 0] + 1
    # pts_middle_encoder = PointPillarsScatter(v_features.shape[1], output_shape=[400, 400])
    pts_middle_encoder = PointPillarsScatter(v_features.shape[1], output_shape=[200, 200])
    feature_map = pts_middle_encoder(v_features, coors, batch_size)
    pts_backbone = build_backbone_model(model, v_features.shape[1], cfg_mask[-1], device=device)
    prune_pts_backbone, backbone_cfg_mask = prune_backbone_model(pts_backbone, v_features.shape[1], ratio=0.1, device=device)
    b_features = prune_pts_backbone(feature_map)
    pts_neck = build_neck_model(model, [b_features[0].shape[1], b_features[1].shape[1], b_features[2].shape[1]], backbone_cfg_mask, device=device)
    return prune_pts_vfe, pts_middle_encoder, prune_pts_backbone, pts_neck

def main():
    parser = ArgumentParser()
    parser.add_argument('pcd', help='Point cloud file')
    parser.add_argument('config', help='Config file')
    parser.add_argument('checkpoint', help='Checkpoint file')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    args = parser.parse_args()

    # build the model from a config file and a checkpoint file
    model, prune_model, cfg_mask = build_and_prune_vfe_model(args.config, args.checkpoint, device=args.device)

    data = preprocess(model, args.pcd)
    img_metas = data['img_metas'][0]
    pts = data['points'][0]
    np_pts = pts[0].cpu().numpy()
    np_pts.tofile('pts.bin')
    voxels, num_points, coors = model.voxelize(pts)
    voxel_features, raw_feats = model.pts_voxel_encoder(voxels, num_points, coors)

    batch_size = coors[-1, 0] + 1
    feature_map = model.pts_middle_encoder(voxel_features, coors, batch_size)

    r_np_raw_feats = np.loadtxt('0_Model_pfe_input_gather_feature.txt',  dtype=np.float32).reshape(-1, 64, 10)
    r_np_raw_feats = r_np_raw_feats[:3687, :, :]

    # np_raw_feats = raw_feats.cpu().numpy()
    # sort_ind = []
    # for i in range(r_np_raw_feats.shape[0]):
    #     sort_ind.append([r_np_raw_feats[i,0,0], r_np_raw_feats[i,0,1]])
    # sort_ind = np.array(sort_ind)
    # modify_r_np_raw_feats = np.empty_like(r_np_raw_feats)
    # np_coors = coors.cpu().numpy()
    # check_coors = np.zeros_like(np_coors)
    # for i in range(np_raw_feats.shape[0]):
    #     value = np.array([np_raw_feats[i,0,0], np_raw_feats[i,0,1]])
    #     idx = (np.sqrt((sort_ind[:,0] - value[0])*(sort_ind[:,0] - value[0]) + (sort_ind[:,1] - value[1])*(sort_ind[:,1] - value[1]))).argmin()
    #     modify_r_np_raw_feats[i,:,:] = r_np_raw_feats[idx,:,:]
    #     check_coors[i,2] = int((value[1] + 25) * 4)
    #     check_coors[i,3] = int((value[0] + 25) * 4)

    # coors_range = np.array([-25, -25, -5, 25, 25, 3])
    # voxel_size = np.array([0.25, 0.25, 8])
    # voxelmap_shape = (coors_range[3:] - coors_range[:3]) / voxel_size
    # voxelmap_shape = tuple(np.round(voxelmap_shape).astype(np.int32).tolist())
    # voxelmap_shape = voxelmap_shape[::-1]
    # print(voxelmap_shape)
    # coor_to_voxelidx = -np.ones(shape=voxelmap_shape, dtype=np.int32)
    # coor = np.zeros(shape=(3, ), dtype=np.int32)
    # ndim = 3
    # ndim_minus_1 = ndim - 1
    # voxel_num = 0
    # for i in range(np_pts.shape[0]):
    #     for j in range(ndim):
    #         c = np.floor((np_pts[i, j] - coors_range[j]) / voxel_size[j])
    #         coor[ndim_minus_1 - j] = c
    #     voxelidx = coor_to_voxelidx[coor[0], coor[1], coor[2]]
    #     if voxelidx == -1:
    #         voxelidx = voxel_num
    #         voxel_num += 1
    #         coor_to_voxelidx[coor[0], coor[1], coor[2]] = voxelidx

    # swap_r_np_raw_feats = np.empty_like(r_np_raw_feats)
    # for i in range(r_np_raw_feats.shape[0]):
    #     for j in range(ndim):
    #         c = np.floor((r_np_raw_feats[i,0,j] - coors_range[j]) / voxel_size[j]).astype(np.int32)
    #         coor[ndim_minus_1 - j] = c
    #     voxelidx = coor_to_voxelidx[coor[0], coor[1], coor[2]]
    #     swap_r_np_raw_feats[voxelidx,:,:] = r_np_raw_feats[i,:,:]

    r_raw_feats = torch.from_numpy(r_np_raw_feats).cuda(args.device)

    # vfe prune
    v_features = prune_model(r_raw_feats)
    # # vfe onnx export
    # if isinstance(args.config, str):
    #     config = mmcv.Config.fromfile(args.config)
    # dummy_input = torch.ones(config.model['pts_voxel_layer']['max_voxels'][1], config.model['pts_voxel_layer']['max_num_points'] , prune_model.in_channels).cuda()
    # export_onnx_file = './pts_pfe_prune.onnx'
    # torch.onnx.export(prune_model,
    #                 dummy_input,
    #                 export_onnx_file,
    #                 opset_version=12,
    #                 verbose=True,
    #                 do_constant_folding=True) # 输出名

    r_np_v_features = np.loadtxt('1_Model_pfe_output_buffers_[1].txt',  dtype=np.float32).reshape(-1, 22)
    r_np_v_features = r_np_v_features[:3687, :]
    r_v_features = torch.from_numpy(r_np_v_features).cuda(args.device)
    v_features = r_v_features

    pts_middle_encoder = PointPillarsScatter(v_features.shape[1], output_shape=[200, 200])
    rfeature_map = pts_middle_encoder(v_features, coors, batch_size)
    print(rfeature_map.shape)
    pts_backbone = build_backbone_model(model, v_features.shape[1], cfg_mask[-1], device=args.device)

    r_np_feature_map = np.loadtxt('2_Model_backbone_input_dev_scattered_feature.txt',  dtype=np.float32).reshape(-1, 200, 200)
    r_feature_map = torch.from_numpy(r_np_feature_map).cuda(args.device).unsqueeze(0)
    print(r_feature_map.shape)
    # debug
    x = model.pts_backbone(feature_map)
    rx = pts_backbone(r_feature_map)
    x = model.pts_neck(rx)
    bbox_list = [dict() for i in range(len(img_metas))]
    outs = model.pts_bbox_head(x)
    bbox_list = model.pts_bbox_head.get_bboxes(
        *outs, img_metas, rescale=False)
    bbox_results = [
        bbox3d2result(bboxes, scores, labels)
        for bboxes, scores, labels in bbox_list
    ]
    boxes = bbox_results[0]['boxes_3d'].tensor.numpy()
    np.savetxt('./boxes.txt', boxes[:,:7], fmt="%f", delimiter=" ")
    exit()
    # # backbone onnx export
    # if isinstance(args.config, str):
    #     config = mmcv.Config.fromfile(args.config)
    # pts_backbone = Backbone(config.model, pts_backbone, model.pts_neck, model.pts_bbox_head)
    # pts_backbone.to(args.device).eval()
    # anchors = np.load('./np_anchors.npy')
    # anchors = torch.from_numpy(anchors).cuda(args.device)
    # rfeature_map = rfeature_map.reshape(-1)
    # ranchors = anchors.reshape(-1)
    # input = torch.cat((rfeature_map, ranchors))
    # # export to onnx
    # export_onnx_file = './pts_backbone_prune.onnx'
    # torch.onnx.export(pts_backbone,
    #                 input,
    #                 export_onnx_file,
    #                 opset_version=12,
    #                 verbose=True,
    #                 do_constant_folding=True) # 输出名

    prune_pts_backbone, backbone_cfg_mask = prune_backbone_model(pts_backbone, v_features.shape[1], ratio=0.1, device=args.device)
    v_features = prune_pts_backbone(rfeature_map)
    pts_neck = build_neck_model(model, [v_features[0].shape[1], v_features[1].shape[1], v_features[2].shape[1]], backbone_cfg_mask, device=args.device)

    # debug
    x = model.pts_backbone(feature_map)
    x = model.pts_neck(x)
    rx = pts_neck(v_features)
    bbox_list = [dict() for i in range(len(img_metas))]
    outs = model.pts_bbox_head(rx)
    bbox_list = model.pts_bbox_head.get_bboxes(
        *outs, img_metas, rescale=False)
    bbox_results = [
        bbox3d2result(bboxes, scores, labels)
        for bboxes, scores, labels in bbox_list
    ]
    boxes = bbox_results[0]['boxes_3d'].tensor.numpy()
    np.savetxt('./boxes.txt', boxes[:,:7], fmt="%f", delimiter=" ")

    # backbone onnx export
    if isinstance(args.config, str):
        config = mmcv.Config.fromfile(args.config)
    pts_backbone_output = Backbone(config.model, prune_pts_backbone, pts_neck, model.pts_bbox_head)
    pts_backbone_output.to(args.device).eval()
    anchors = np.load('./np_anchors.npy')
    anchors = torch.from_numpy(anchors).cuda(args.device)
    rfeature_map = rfeature_map.reshape(-1)
    np_rfeature_map = rfeature_map.detach().cpu().numpy()
    np.save('./np_rfeature_map.npy', np_rfeature_map)
    ranchors = anchors.reshape(-1)
    input = torch.cat((rfeature_map, ranchors))

    # export to onnx
    export_onnx_file = './pts_backbone_prune.onnx'
    torch.onnx.export(pts_backbone_output,
                    input,
                    export_onnx_file,
                    opset_version=12,
                    verbose=True,
                    do_constant_folding=True) # 输出名

import onnxruntime
def test_onnx():
    # test onnx
    input_data = np.loadtxt('2_Model_backbone_input_dev_scattered_feature.txt',  dtype=np.float32).reshape(-1, 22, 200, 200)
    # input_data = np.load('./rfeature_map.npy')
    anchors = np.load('./np_anchors.npy')
    # np.savetxt('generate_anchors.txt', anchors, fmt="%f", delimiter=" ")
    # anchors = np.zeros(anchors.shape, dtype=anchors.dtype)

    input_data = input_data.reshape(-1)
    print(input_data.shape)
    anchors = anchors.reshape(-1)
    print(anchors.shape)
    input_data = np.concatenate([input_data, anchors])
    print(input_data.shape)

    sess = onnxruntime.InferenceSession('./pts_backbone_prune.onnx')
    outputs = sess.run(["271", "222", "223"], {"0":input_data})
    # outputs = sess.run(["134", "193", "194", "195"], {"0":input_data})
    np_preds = outputs[0]
    np_scores = outputs[1]
    np_clss = outputs[2]
    # np_anchor = outputs[0]

    print(np_preds[:1])
    print(np_scores[:1])
    print(np_clss[:1])
    # print(np_anchor[:1])
    print(np_preds.shape)
    print(np_scores.shape)
    print(np_clss.shape)
    # print(np_anchor.shape)

    filter_pred = []
    for i in range(np_scores.shape[0]):
        score = np_scores[i,:][0]
        if score > 0.1:
            filter_pred.append(np_preds[i,:])
    filter_pred = np.array(filter_pred)
    print(filter_pred.shape)
    np.savetxt('./boxes_onnx.txt', filter_pred[:,:7], fmt="%f", delimiter=" ")

if __name__ == '__main__':
    # test_onnx()
    main()


import os, cv2
import numpy as np
import tensorflow as tf

image_folder = r'G:\Data\coco2017\val2017'
imgs = os.listdir(image_folder)
shapes = []
for img in imgs:
    img = cv2.imread(os.path.join(image_folder, img))
    shapes.append(img.shape)

shapes = np.array(shapes)
print(max(shapes[:,0]))


#!/usr/bin/env python

"""

 Background:
 ===========
 CTD_discreet_cal_corrections.py

 Purpose:
 ========
 Apply salinity or oxygen calibration corrections to CTD data

 USAGE:
 ========
 python CTD_discreet_cal_corrections.py {data} --secondary_oxygen {linear_correction_slope linear_correction_offset}
 python CTD_discreet_cal_corrections.py {data} --secondary_salinity {offset_correction}

 History:
 ========
 2019-08-15: Python 3 print statments and f-strings

 Compatibility:
 ==============
 python >=3.8
 python 2.7 - not supported

"""
import argparse
import datetime
import os
import sys

# must be python 3.8 or greater
try:
    assert sys.version_info >= (3, 8)
except AssertionError:
    sys.exit("Must be running python 3.8 or greater")

import numpy as np
from netCDF4 import Dataset

__author__ = "Shaun Bell"
__email__ = "shaun.bell@noaa.gov"
__created__ = datetime.datetime(2014, 10, 30)
__modified__ = datetime.datetime(2014, 10, 30)
__version__ = "0.1.0"
__status__ = "Development"
__keywords__ = "CTD", "SeaWater", "Cruise", "derivations"

"""--------------------------------netcdf Routines---------------------------------------"""

def get_global_atts(nchandle):

    g_atts = {}
    att_names = nchandle.ncattrs()

    for name in att_names:
        g_atts[name] = nchandle.getncattr(name)

    return g_atts

def get_vars(nchandle):
    return nchandle.variables

def ncreadfile_dic(nchandle, params):
    data = {}
    for j, v in enumerate(params):
        if v in nchandle.variables.keys():  # check for nc variable
            data[v] = nchandle.variables[v][:]

        else:  # if parameter doesn't exist fill the array with zeros
            data[v] = None
    return data

def repl_var(nchandle, var_name, val):
    nchandle.variables[var_name][:] = val[:]
    return

"""----------------------------- MAIN -------------------------------------------------"""

parser = argparse.ArgumentParser(
    description="Apply corrections from Discreet oxygen or salinity samples"
)
parser.add_argument(
    "DataPath", metavar="DataPath", type=str, help="full path to directory for cruise"
)
parser.add_argument(
    "--primary_oxygen",
    nargs="+",
    type=float,
    help="apply linear correction to primary oxygen",
)
parser.add_argument(
    "--secondary_oxygen",
    nargs="+",
    type=float,
    help="apply linear correction to secondary oxygen",
)
parser.add_argument(
    "--primary_salinity", type=float, help="apply offset correction to primary salinity"
)
parser.add_argument(
    "--secondary_salinity",
    type=float,
    help="apply offset correction to secondary salinity",
)
parser.add_argument(
    "--fluorometer",
    type=float,
    help="apply offset correction to fluorometer (0.0 - deep noise)",
)

args = parser.parse_args()

# get all .nc files from chosen directory
full_path = [args.DataPath +
             x for x in os.listdir(args.DataPath) if x.endswith(".nc")]

if args.primary_oxygen:  # O_65
    for ifile, ncfile in enumerate(full_path):
        print(f"Working on file {ncfile} \n")
        # nc readin
        nchandle = Dataset(ncfile, "a")
        global_atts = get_global_atts(nchandle)
        vars_dic = get_vars(nchandle)
        data = ncreadfile_dic(nchandle, ["O_65"])
        O2corr = (data["O_65"] * args.primary_oxygen[0]) + \
            args.primary_oxygen[1]
        O2corr[0, data["O_65"][0, :, 0, 0] >= 1e10, 0, 0] = 1e35

        repl_var(nchandle, "O_65", O2corr)

        # Look for existing program and edit comments / scoot down one level and add new
        for i in range(1, 10):
            if ("PROG_CMNT0" + str(i)) in global_atts.keys():
                nchandle.setncattr(
                    "PROG_CMNT0" +
                    str(i + 1), global_atts["PROG_CMNT0" + str(i)]
                )
            else:
                nchandle.setncattr(
                    "PROG_CMNT01", __file__.split("/")[-1] + " v" + __version__
                )
            if ("EDIT_CMNT0" + str(i)) in global_atts.keys():
                nchandle.setncattr(
                    "EDIT_CMNT0" +
                    str(i + 1), global_atts["EDIT_CMNT0" + str(i)]
                )
            else:
                nchandle.setncattr(
                    "EDIT_CMNT01",
                    "Primary Oxygen Cal Factor of: "
                    + str(args.primary_oxygen[0])
                    + "x + "
                    + str(args.primary_oxygen[1])
                    + " applied",
                )
        nchandle.close()

if args.secondary_oxygen:  # CTDOXY_4221
    for ifile, ncfile in enumerate(full_path):
        print(f"Working on file {ncfile} \n")
        # nc readin
        nchandle = Dataset(ncfile, "a")
        global_atts = get_global_atts(nchandle)
        vars_dic = get_vars(nchandle)
        data = ncreadfile_dic(nchandle, ["CTDOXY_4221"])
        O2corr = (
            data["CTDOXY_4221"] * args.secondary_oxygen[0]
        ) + args.secondary_oxygen[1]
        O2corr[0, data["CTDOXY_4221"][0, :, 0, 0] >= 1e10, 0, 0] = 1e35

        repl_var(nchandle, "CTDOXY_4221", O2corr)

        # Look for existing program and edit comments / scoot down one level and add new
        for i in range(1, 10):
            if ("PROG_CMNT0" + str(i)) in global_atts.keys():
                nchandle.setncattr(
                    "PROG_CMNT0" +
                    str(i + 1), global_atts["PROG_CMNT0" + str(i)]
                )
            else:
                nchandle.setncattr(
                    "PROG_CMNT01", __file__.split("/")[-1] + " v" + __version__
                )
            if ("EDIT_CMNT0" + str(i)) in global_atts.keys():
                nchandle.setncattr(
                    "EDIT_CMNT0" +
                    str(i + 1), global_atts["EDIT_CMNT0" + str(i)]
                )
            else:
                nchandle.setncattr(
                    "EDIT_CMNT01",
                    "Secondary Oxygen Cal Factor of: "
                    + str(args.secondary_oxygen[0])
                    + "x + "
                    + str(args.secondary_oxygen[1])
                    + " applied",
                )

        nchandle.close()

if args.primary_salinity:  # S_41
    for ifile, ncfile in enumerate(full_path):
        print(f"Working on file {ncfile} \n")
        # nc readin
        nchandle = Dataset(ncfile, "a")
        global_atts = get_global_atts(nchandle)
        vars_dic = get_vars(nchandle)
        data = ncreadfile_dic(nchandle, ["S_41"])
        repl_var(nchandle, "S_41", data["S_41"] + args.primary_salinity)

        # Look for existing program and edit comments / scoot down one level and add new
        for i in range(1, 10):
            if ("PROG_CMNT0" + str(i)) in global_atts.keys():
                nchandle.setncattr(
                    "PROG_CMNT0" +
                    str(i + 1), global_atts["PROG_CMNT0" + str(i)]
                )
            else:
                nchandle.setncattr(
                    "PROG_CMNT01", __file__.split("/")[-1] + " v" + __version__
                )
            if ("EDIT_CMNT0" + str(i)) in global_atts.keys():
                nchandle.setncattr(
                    "EDIT_CMNT0" +
                    str(i + 1), global_atts["EDIT_CMNT0" + str(i)]
                )
            else:
                nchandle.setncattr(
                    "EDIT_CMNT01",
                    "Primary Salinity Cal Factor of: "
                    + str(args.primary_salinity)
                    + " applied",
                )

        nchandle.close()

if args.secondary_salinity:  # S_42
    for ifile, ncfile in enumerate(full_path):
        print(f"Working on file {ncfile} \n")
        # nc readin
        nchandle = Dataset(ncfile, "a")
        global_atts = get_global_atts(nchandle)
        vars_dic = get_vars(nchandle)
        data = ncreadfile_dic(nchandle, ["S_42"])
        repl_var(nchandle, "S_42", data["S_42"] + args.secondary_salinity)

        # Look for existing program and edit comments / scoot down one level and add new
        for i in range(1, 10):
            if ("PROG_CMNT0" + str(i)) in global_atts.keys():
                nchandle.setncattr(
                    "PROG_CMNT0" +
                    str(i + 1), global_atts["PROG_CMNT0" + str(i)]
                )
            else:
                nchandle.setncattr(
                    "PROG_CMNT01", __file__.split("/")[-1] + " v" + __version__
                )
            if ("EDIT_CMNT0" + str(i)) in global_atts.keys():
                nchandle.setncattr(
                    "EDIT_CMNT0" +
                    str(i + 1), global_atts["EDIT_CMNT0" + str(i)]
                )
            else:
                nchandle.setncattr(
                    "EDIT_CMNT01",
                    "Secondary Salinity Cal Factor of: "
                    + str(args.secondary_salinity)
                    + " applied",
                )

        nchandle.close()

if args.fluorometer:  # fWS_973, or F_903
    for ifile, ncfile in enumerate(full_path):
        print(f"Working on file {ncfile} \n")
        # nc readin
        nchandle = Dataset(ncfile, "a")
        global_atts = get_global_atts(nchandle)
        vars_dic = get_vars(nchandle)
        try:
            data = ncreadfile_dic(nchandle, ["fWS_973"])
            repl_var(nchandle, "fWS_973", data["fWS_973"] + args.fluorometer)
        except:
            try:
                data = ncreadfile_dic(nchandle, ["F_903"])
                repl_var(nchandle, "F_903", data["F_903"] + args.fluorometer)
            except:
                print("No valid Fluorometer key - F_903 or fWS_973 found")

        # Look for existing program and edit comments / scoot down one level and add new
        for i in range(1, 10):
            if ("PROG_CMNT0" + str(i)) in global_atts.keys():
                nchandle.setncattr(
                    "PROG_CMNT0" +
                    str(i + 1), global_atts["PROG_CMNT0" + str(i)]
                )
            else:
                nchandle.setncattr(
                    "PROG_CMNT01", __file__.split("/")[-1] + " v" + __version__
                )
            if ("EDIT_CMNT0" + str(i)) in global_atts.keys():
                nchandle.setncattr(
                    "EDIT_CMNT0" +
                    str(i + 1), global_atts["EDIT_CMNT0" + str(i)]
                )
            else:
                nchandle.setncattr(
                    "EDIT_CMNT01",
                    "Fluorometer offset of: " +
                    str(args.fluorometer) + " applied",
                )

        nchandle.close()


#coding=utf-8
import random
import numpy as np
class CARD():
	def __init__(self,_type,_size):
		assert (_type in range(4))
		assert (_size in range(13))
		self._type=_type
		self._size=_size
	def __str__(self):
		return '[type:%s,size:%s]'%(self._type,self._size)
	def __cmp__(self,s):
		if self._size<s._size:
			return -1
		return 1 if self._size>s._size else 0

class PLAYER():
	def __init__(self,id):
		assert(id in range(6))
		self.id=id
		self.see_card=False
		self.cards=[]
		self.money=0
		self.alive=True
		self.last_action=0
	def seecard(self):
		self.see_card=True
	#游戏结束后进行反省,修改策略参数，以期使自身获取更多的累积奖励
	def rethink(self):
		pass
	def choose_action(self):
		return random.randint(0,14)
	def handle_reward(self,money):
		self.money+=money
	def size(self):
		c=self.cards
		if c[0]._size==c[1]._size and c[1]._size==c[2]._size:
			return 60000+c[0]._size
		if c[0]._type==c[1]._type and c[1]._type==c[2]._type:
			if c[2]._size==c[1]._size+1 and c[1]._size==c[0]._size+1:
				return 50000+c[0]._size
			return 40000+c[0]._size+c[1]._size*13+c[2]._size*13*13
		if c[2]._size==c[1]._size+1 and c[1]._size==c[0]._size+1:
			return 30000+c[0]._size
		if c[0]._size==c[1]._size:
			return 20000+c[1]._size*100+c[2]._size
		if c[1]==c[2]:
			return 20000+c[1]._size*100+c[0]._size
		if c[0]._size==2 and c[1]._size==3 and c[2]._size==5:
			return 235
		return 10000+c[0]._size+c[1]._size*13+c[2]._size*13*13

	def __cmp__(self,s):
		if self.size()<s.size():
			return -1
		return 1 if self.size()>s.size() else 0

	def output(self):
		print('playerid=>',self.id,'cards=>',[str(card) for card in self.cards],'money=>',self.money,'size=>',self.size(),'alive=>',self.alive)

class GAME():
	def __init__(self,nPlayers,max_epoch=10):
		self.nPlayers=nPlayers
		self.players=[PLAYER(id) for id in range(nPlayers)]	#保存未淘汰的用户
		self.cards=[]	     #保存扑克牌
		self.reward_pool=0   #当前奖池累积
		self.min_add_money=1 #最小加注数
		self.max_epoch=max_epoch
		self.cur_pid=0
		self.epoch=0
		self.actions=np.zeros([20,2,6])
		for i in range(4):
			for j in range(13):
				self.cards.append(CARD(i,j))
		random.shuffle(self.cards)
		self.action_history=[] #动作的历史记录
		self.actions=[  #可选动作
		'no_see',
		'see_cards',
		'give_up',
		'compare_0',
		'compare_1',
		'compare_2',
		'compare_3',
		'compare_4',
		'compare_5',
		'add_money_1',
		'add_money_2',
		'add_money_4',
		'add_money_8',
		'add_money_16',
		'add_money_32',
		]

	#发牌操作
	def deal(self):
		for i in range(3):
			for p in self.players:
				card=self.cards.pop()
				p.cards.append(card)
		for p in self.players:
			p.cards.sort(key=lambda s:s._size)
	#淘汰一个玩家
	def out(self,p):
		p.alive=False

	def alive_players(self):
		ap=0
		for p in self.players:
			if p.alive:
				ap+=1
		return ap
	#加注操作,每执行一次加注操作用户就暂时损失一笔钱
	def add_money(self,p,money):
		self.reward(p,-1*money)

	#对player进行奖励
	def reward(self,p,money):
		self.reward_pool-=money
		p.handle_reward(money)
		return p.id,money

	#选择看牌阶段
	def gambling_1(self):
		for p in self.players:
			if not p.alive:
				continue
			a=p.choose_action()
			if a>1: #不允许执行除看牌或无操作以外的操作
				self.reward(p,-10)
				self.out(p)
			elif a==1:
				if p.see_card:#不允许二次看牌
					self.reward(p,-10)
					self.out(p)
				else:
					p.seecard()
			self.action_history.append([p.id,a])
			p.last_action=a
			if self.alive_players()<=1:
				return 'over'

	#选择加注、比牌、弃牌阶段
	def gambling_2(self):
		for p in self.players:
			if not p.alive:
				continue
			a=p.choose_action()
			if a==0 or a==1:#必须要操作，并且不是看牌操作
				#违反规则给与惩罚，并且该玩家出局
				self.reward(p,-10)
				self.out(p)
			if a==2:#选择弃牌
				self.out(p)
			if a in range(3,9):#选择比牌操作，但不能选择比牌对象是自身或淘汰的玩家
				compare_id=a-3
				pCp=None
				for p1 in self.players:
					if p1.id==compare_id:
						pCp=p1
						break
				if p.id==compare_id or not pCp:
					#违反规则给与惩罚，并且该玩家出局
					self.reward(p,-10)
					self.out(p)
				else:
					self.compare(p,pCp)
			if a in range(9,15):#选择加注操作
				money=2**(a-9)
				if self.min_add_money<=money:
					self.min_add_money=money
					self.add_money(p,2**(a-9))
				else:
					self.reward(p,-10)
					self.out(p)
			#加入到游戏的动作历史集中，作为游戏状态的一个属性。
			self.action_history.append([p.id,a])
			p.last_action=a
			if self.alive_players()<=1:
				return 'over'

	#比牌操作
	def compare(self,p1,p2):
		#比牌者必须拿出两倍的最小加注数放入奖池
		self.reward(p1,-1*self.min_add_money*2)
		if p1.size()>p2.size():
			self.out(p2)
		else:
			self.out(p1)

	#摊牌操作，所有玩家都亮出底牌，比较大小后给与奖励后游戏结束
	def showhand(self):
		maxSize=0
		for p in self.players:
			if p.alive and maxSize <p.size():
				maxSize=p.size()
		for p in self.players:
			if p.alive and p.size()<maxSize:
				self.out(p)
		ap=self.alive_players()
		for p in self.players:
			if p.alive and p.size()==maxSize:
				self.reward(p,self.reward_pool/ap)
				p.rethink()

	def play(self):
		#清空奖励池
		self.reward_pool=0
		#洗牌、发牌、押底
		self.deal()
		for p in self.players:
			self.reward(p,-1)
		#玩家开始轮流操作
		for self.epoch in range(self.max_epoch):
			print(game.spectrum())
			#选择加注、比牌、弃牌
			if 'over'==self.gambling_1() or 'over'==self.gambling_2():
				return self.showhand()
		return self.showhand()

	def replay(self):
		print('reward_pool:',self.reward_pool)
		print('players:')
		for p in self.players:
			p.output()
		print('actions:')
		for a in self.action_history:
			print(a)
			pid=a[0]
			aid=a[1]
			print('player:%s %s'%(pid,self.actions[aid]))

	def spectrum(self):
		player_state=np.zeros([6,30])
		alive=player_state[:,0]
		isme=player_state[:,1]
		money=player_state[:,2]
		last_action=player_state[:,3:18]
		cards=np.reshape(player_state[:,18:30],[6,3,4])
		for p in self.players:
			alive[p.id]= 1 if p.alive else 0
			isme[p.id]=1 if p.id==self.cur_pid else 0
			money[p.id]=p.money
			last_action[p.id]=p.last_action
			for i in range(3):
				cards[p.id,i,p.cards[i]._type]=p.cards[i]._size
		game_state=np.zeros(22)
		game_state[0]=self.reward_pool
		game_state[1]=self.min_add_money
		game_state[2+self.epoch]=1
		return game_state,player_state

if __name__=='__main__':
	for i in range(1):
		game=GAME(nPlayers=6)
		game.play()
		print('game %d'%i)
		game.replay()