export { ConcatGather };
export { ConcatGatherPool };

import * as Pool from "../../../util/Pool.js";
import { ConcatGather as ChannelShuffler_ConcatGather } from "../ChannelShuffler_ConcatGather.js";

/**
 *
 */
class ConcatGather extends ChannelShuffler_ConcatGather {

  constructor( concatenatedShape, outputGroupCount ) {
    super( concatenatedShape, outputGroupCount );
  }

//!!! (2022/06/25 Remarked) Inherits from Recyclable.Base instead.
//   /**
//    * After calling this method, this object should be viewed as disposed and should not be operated again.
//    *
//    * Sub-class should override this method for recycling to its pool (and NEVER call super.disposeResources_and_recycleToPool()).
//    */
//   disposeResources_and_recycleToPool() {
//     this.disposeResources();
//     ConcatGatherPool.Singleton.recycle( this );
//   }

  gather_map( concatenatedTensor ) {
    // shuffle and split by gather (one operation achieves two operations).
    return this.shuffledChannelIndicesTensor1dArray.map(
      shuffledChannelIndicesTensor1d =>
        concatenatedTensor.gather( shuffledChannelIndicesTensor1d, this.shuffleInfo.lastAxisId )
    );
  }


  concatGather_dispose_direct_call_loop( tensorArray ) {
    let concatenatedTensor = tf.concat( tensorArray, this.shuffleInfo.lastAxisId );

    let shuffledSplitedTensorArray = this.gather_loop( concatenatedTensor );
    concatenatedTensor.dispose();

    return shuffledSplitedTensorArray;
  }


  concatGather_dispose_finally_call_map( tensorArray ) {
    let concatenatedTensor = tf.concat( tensorArray, this.shuffleInfo.lastAxisId );

    try {
      return this.gather_map( concatenatedTensor );

    } finally {
      concatenatedTensor.dispose();
    }
  }

  concatGather_dispose_direct_call_map( tensorArray ) {
    let concatenatedTensor = tf.concat( tensorArray, this.shuffleInfo.lastAxisId );

    let shuffledSplitedTensorArray = this.gather_map( concatenatedTensor );
    concatenatedTensor.dispose();

    return shuffledSplitedTensorArray;
  }


  concatGather_dispose_finally_loop( tensorArray ) {
    let concatenatedTensor = tf.concat( tensorArray, this.shuffleInfo.lastAxisId );

    try {
      let shuffledSplitedTensorArray = new Array( this.shuffledChannelIndicesTensor1dArray.length );
      for ( let i = 0; i < shuffledSplitedTensorArray.length; ++i ) {
        // shuffle and split by gather (one operation achieves two operations).
        shuffledSplitedTensorArray[ i ] = concatenatedTensor.gather( this.shuffledChannelIndicesTensor1dArray[ i ], this.shuffleInfo.lastAxisId );
      }

      return shuffledSplitedTensorArray;

    } finally {
      concatenatedTensor.dispose();
    }
  }

  concatGather_dispose_direct_loop( tensorArray ) {
    let concatenatedTensor = tf.concat( tensorArray, this.shuffleInfo.lastAxisId );

    let shuffledSplitedTensorArray = new Array( this.shuffledChannelIndicesTensor1dArray.length );
    for ( let i = 0; i < shuffledSplitedTensorArray.length; ++i ) {
      // shuffle and split by gather (one operation achieves two operations).
      shuffledSplitedTensorArray[ i ] = concatenatedTensor.gather( this.shuffledChannelIndicesTensor1dArray[ i ], this.shuffleInfo.lastAxisId );
    }

    concatenatedTensor.dispose();

    return shuffledSplitedTensorArray;
  }


  concatGather_dispose_finally_map( tensorArray ) {
    let concatenatedTensor = tf.concat( tensorArray, this.shuffleInfo.lastAxisId );

    try {
      // shuffle and split by gather (one operation achieves two operations).
      let shuffledSplitedTensorArray = this.shuffledChannelIndicesTensor1dArray.map(
        shuffledChannelIndicesTensor1d =>
          concatenatedTensor.gather( shuffledChannelIndicesTensor1d, this.shuffleInfo.lastAxisId )
      );
      return shuffledSplitedTensorArray;

    } finally {
      concatenatedTensor.dispose();
    }
  }

  concatGather_dispose_direct_map( tensorArray ) {
    let concatenatedTensor = tf.concat( tensorArray, this.shuffleInfo.lastAxisId );

    // shuffle and split by gather (one operation achieves two operations).
    let shuffledSplitedTensorArray = this.shuffledChannelIndicesTensor1dArray.map(
      shuffledChannelIndicesTensor1d =>
        concatenatedTensor.gather( shuffledChannelIndicesTensor1d, this.shuffleInfo.lastAxisId )
    );

    concatenatedTensor.dispose();

    return shuffledSplitedTensorArray;
  }


  concatGather_tidy_map( tensorArray ) {
    return tf.tidy( "ChannelShuffler.ConcatGather.concatGather", () => {
      let concatenatedTensor = tf.concat( tensorArray, this.shuffleInfo.lastAxisId );

      // shuffle and split by gather (one operation achieves two operations).
      let shuffledSplitedTensorArray = this.shuffledChannelIndicesTensor1dArray.map(
        shuffledChannelIndicesTensor1d =>
          concatenatedTensor.gather( shuffledChannelIndicesTensor1d, this.shuffleInfo.lastAxisId )
      );
      return shuffledSplitedTensorArray;
    });
  }

}


/**
 * Providing ChannelShuffler.PerformanceTest.ConcatGather
 *
 */
class ConcatGatherPool extends Pool.Root {

  constructor() {
    super( "ChannelShuffler.PerformanceTest.ConcatGatherPool", ConcatGather, ConcatGather.setAsConstructor );
  }

}


/**
 * Used as default ChannelShuffler.PerformanceTest.ConcatGather provider for conforming to Recyclable interface.
 */
ConcatGather.Pool = new ConcatGatherPool();


@@ -21,37 +21,8 @@
 r %7D;
-%0Aexport %7B ConcatGatherPool %7D;
 %0A%0Aim
@@ -242,582 +242,896 @@
 %0A%0A  
-constructor( concatenatedShape, outputGroupCount ) %7B%0A    super( concatenatedShape, outputGroupCount );%0A  %7D%0A%0A//!!! (2022/06/25 Remarked) Inherits from Recyclable.Base instead.%0A//   /**%0A//    * After calling this method, this object should be viewed as disposed and should not be operated again.%0A//    *%0A//    * Sub-class should override this method for recycling to its pool (and NEVER call super.disposeResources_and_recycleToPool()).%0A//    */%0A//   disposeResources_and_recycleToPool() %7B%0A//     this.disposeResources();%0A//     ConcatGatherPool.Singleton.recycle( this );%0A// 
+/**%0A   * Used as default ChannelShuffler.PerformanceTest.ConcatGather provider for conforming to Recyclable interface.%0A   */%0A  static Pool = new Pool.Root( %22ChannelShuffler.PerformanceTest.ConcatGatherPool%22, ConcatGather, ConcatGather.setAsConstructor );%0A%0A  constructor( concatenatedShape, outputGroupCount ) %7B%0A    super( concatenatedShape, outputGroupCount );%0A    ConcatGather.setAsConstructor_self.call( this, concatenatedShape, outputGroupCount );%0A  %7D%0A%0A  /** @override */%0A  static setAsConstructor_self( concatenatedShape, outputGroupCount ) %7B%0A  %7D%0A%0A  /** @override */%0A  static setAsConstructor( concatenatedShape, outputGroupCount ) %7B%0A    super.setAsConstructor( concatenatedShape, outputGroupCount );%0A    ConcatGather.setAsConstructor_self.call( this, concatenatedShape, outputGroupCount );%0A    return this;%0A  %7D%0A%0A  /** @override */%0A  disposeResources() %7B%0A   super.disposeResources();%0A
   %7D%0A
@@ -5249,421 +5249,4 @@
 %0A%7D%0A%0A
-%0A/**%0A * Providing ChannelShuffler.PerformanceTest.ConcatGather%0A *%0A */%0Aclass ConcatGatherPool extends Pool.Root %7B%0A%0A  constructor() %7B%0A    super( %22ChannelShuffler.PerformanceTest.ConcatGatherPool%22, ConcatGather, ConcatGather.setAsConstructor );%0A  %7D%0A%0A%7D%0A%0A%0A/**%0A * Used as default ChannelShuffler.PerformanceTest.ConcatGather provider for conforming to Recyclable interface.%0A */%0AConcatGather.Pool = new ConcatGatherPool();%0A%0A

Update ChannelShuffler_PerformanceTest_ConcatGather.js