import math
from abc import abstractmethod
from pathlib import Path
from typing import List, Literal, Optional, Union
import os

import numpy as np

from . import comm
from . import profiler
from .gpt import GptInitModelParameters

from nebullvm.optional_modules.torch import torch

PathLike = Union[str, Path]

def to_numpy_dtype(maybe_str_dtype: Union[str, np.dtype]):
    assert isinstance(maybe_str_dtype, (str, np.dtype))
    if isinstance(maybe_str_dtype, str):
        try:
            dtype = {
                "fp16": np.float16,
                "float16": np.float16,
                "fp32": np.float32,
                "float32": np.float32,
            }[maybe_str_dtype]
        except KeyError:
            raise ValueError(
                f"Cannot convert to numpy data type, got {maybe_str_dtype}"
            )
    else:
        dtype = maybe_str_dtype
    return dtype

def to_torch_dtype(maybe_str_dtype: Union[str, torch.dtype]):

    if isinstance(maybe_str_dtype, torch.dtype):
        dtype = maybe_str_dtype
    else:
        try:
            dtype = {
                "bf16": torch.bfloat16,
                "fp16": torch.float16,
                "fp32": torch.float32,
                "bfloat16": torch.bfloat16,
                "float16": torch.float16,
                "float32": torch.float32,
            }[maybe_str_dtype]
        except KeyError:
            raise ValueError(
                f"Cannot convert to torch data type, got {maybe_str_dtype}"
            )
    return dtype

def load_weight_from_bin(
    checkpoint_path: PathLike,
    shape: List[int],
    weight_dtype: Union[str, np.dtype],
):
    """Load a weight from a bin file.

    # Args.
        checkpoint_path: str or Path,
            a checkpoint file path of an FT's layer weight.
        shape: list of int, the shape of weight tensor.
        weight_dtype: str or np.dtype, the data type of the stored weight.
    """
    weight_dtype = to_numpy_dtype(weight_dtype)
    return torch.from_numpy(np.fromfile(checkpoint_path, dtype=weight_dtype))

LayernormType = Literal["pre_layernorm", "post_layernorm"]

class GptLayerWeights:
    def __init__(
        self,
        num_heads: int,
        size_per_head: int,
        inter_size: int,
        num_layers: int,
        tensor_para_size: int = 1,
        pipeline_para_size: int = 1,
        has_adapters: bool = False,
        adapter_inter_size: int = 0,
        int8_mode: int = 0,
    ):

        assert num_heads % tensor_para_size == 0, (
            f"num_heads ({num_heads}) is not multiple of "
            "tensor para size ({tensor_para_size})"
        )

        self.num_heads = num_heads
        self.size_per_head = size_per_head
        self.hidden_units = num_heads * size_per_head
        self.num_layers = num_layers

        self.tensor_para_size = tensor_para_size
        self.tensor_para_rank = comm.get_tensor_para_rank()
        self.pipeline_para_size = pipeline_para_size
        self.pipeline_para_rank = comm.get_pipeline_para_rank()

        self.has_adapters = has_adapters
        self.adapter_inter_size = adapter_inter_size

        self.local_num_layers = num_layers // pipeline_para_size
        self.local_num_heads = num_heads // tensor_para_size
        self.local_hidden_units = self.local_num_heads * size_per_head
        self.local_inter_size = inter_size // tensor_para_size
        self.local_adapter_inter_size = (
            self.adapter_inter_size // tensor_para_size
        )

        self.weight_transpose_calibrate_quantize = None
        assert int8_mode in [0, 1], "Invalid int8 mode for GPT. Must be 0 or 1"
        self.int8_mode = int8_mode
        if self.int8_mode == 1:
            quant = (
                torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix  # noqa: E501
            )
            self.weight_transpose_calibrate_quantize = lambda x: quant(
                x, torch.int8
            )

        self.weights = None
        self.int8_weights = None
        self.int8_scales = None

        self.expected_weight_shapes = list()

        # pylint:disable=line-too-long
        # Transformer blocks
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # input layernorm weight
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # input layernorm bias
        self.expected_weight_shapes.extend(
            [(self.hidden_units, self.local_hidden_units * 3)]
            * self.local_num_layers
        )  # attention qkv weight
        self.expected_weight_shapes.extend(
            [(self.local_hidden_units * 3,)] * self.local_num_layers
        )  # attention qkv bias
        self.expected_weight_shapes.extend(
            [(self.local_hidden_units, self.hidden_units)]
            * self.local_num_layers
        )  # attention dense weight
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # attention dense bias
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # post attention layernorm weight
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # post attention layernorm bias
        self.expected_weight_shapes.extend(
            [(self.hidden_units, self.local_inter_size)]
            * self.local_num_layers
        )  # ffn_kernel1
        self.expected_weight_shapes.extend(
            [(self.local_inter_size,)] * self.local_num_layers
        )  # ffn_bias1
        self.expected_weight_shapes.extend(
            [(self.local_inter_size, self.hidden_units)]
            * self.local_num_layers
        )  # ffn_kernel2
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # ffn_bias2

        # Adapters
        if self.has_adapters:
            self.expected_weight_shapes.extend(
                [(self.hidden_units, self.local_adapter_inter_size)]
                * self.local_num_layers
            )  # adaptor1_kernel1
            self.expected_weight_shapes.extend(
                [(self.local_adapter_inter_size,)] * self.local_num_layers
            )  # adaptor1_bias1
            self.expected_weight_shapes.extend(
                [(self.local_adapter_inter_size, self.hidden_units)]
                * self.local_num_layers
            )  # adaptor1_kernel2
            self.expected_weight_shapes.extend(
                [(self.hidden_units,)] * self.local_num_layers
            )  # adaptor1_bias2
            self.expected_weight_shapes.extend(
                [(self.hidden_units, self.local_adapter_inter_size)]
                * self.local_num_layers
            )  # adaptor2_kernel1
            self.expected_weight_shapes.extend(
                [(self.local_adapter_inter_size,)] * self.local_num_layers
            )  # adaptor2_bias1
            self.expected_weight_shapes.extend(
                [(self.local_adapter_inter_size, self.hidden_units)]
                * self.local_num_layers
            )  # adaptor2_kernel2
            self.expected_weight_shapes.extend(
                [(self.hidden_units,)] * self.local_num_layers
            )  # adaptor2_bias2
        # pylint:enable=line-too-long

    @classmethod
    def from_config(cls, config: GptInitModelParameters):
        return cls(
            num_heads=config.head_num,
            size_per_head=config.size_per_head,
            inter_size=4 * config.head_num * config.size_per_head,
            num_layers=config.layer_num,
            tensor_para_size=config.tensor_para_size,
            pipeline_para_size=config.pipeline_para_size,
            has_adapters=config.has_adapters,
            adapter_inter_size=config.adapter_inter_size,
            int8_mode=config.int8_mode,
        )

    @property
    def dtype(self):
        return self.weights[0].dtype

    @property
    def device(self):
        return self.weights[0].device

    def _map(self, func):
        for i in range(len(self.weights)):
            if isinstance(self.weights[i], list):
                for j in range(len(self.weights[i])):
                    self.weights[i][j] = func(self.weights[i][j])
            else:
                self.weights[i] = func(self.weights[i])

    def _map_int8(self, func):
        for i in range(len(self.int8_weights)):
            if isinstance(self.int8_weights[i], list):
                for j in range(len(self.int8_weights[i])):
                    self.int8_weights[i][j] = func(self.int8_weights[i][j])

            else:
                self.int8_weights[i] = func(self.int8_weights[i])
        for i in range(len(self.int8_scales)):
            if isinstance(self.int8_scales[i], list):
                for j in range(len(self.int8_scales[i])):
                    self.int8_scales[i][j] = func(self.int8_scales[i][j])
            else:
                self.int8_scales[i] = func(self.int8_scales[i])

    def float(self):
        if self.dtype == torch.float32:
            return
        self._map(lambda x: x.float())

    def half(self):
        if self.dtype == torch.float16:
            return
        self._map(lambda x: x.half())
        if self.int8_mode == 1:
            self._map_int8(lambda w: w.half())

    def bfloat16(self):
        if self.dtype == torch.bfloat16:
            return
        self._map(lambda x: x.bfloat16())
        if self.int8_mode == 1:
            self._map_int8(lambda w: w.bfloat16())

    def cuda(self, device=None):
        self._map(lambda x: x.cuda(device))
        if self.int8_mode == 1:
            self._map_int8(lambda x: x.cuda(device))

    def to(self, device=None):
        self._map(lambda x: x.to(device))
        if self.int8_mode == 1:
            self._map_int8(lambda x: x.to(device))

    def is_valid_pp_group(self, layer, pp_rank):
        return layer // self.layers_per_device == pp_rank

    def load(
        self,
        checkpoint_path: PathLike,
        compute_dtype: torch.dtype,
        weight_dtype: Optional[Union[str, np.dtype]] = None,
        device: Optional[Union[int, str, torch.device]] = None,
    ):
        """Load checkpoint weights.

        # Args.
            checkpoint_path: str or Path,
                a checkpoint directory where FT checkpoint files locate.
            weight_dtype: str or np.dtype, the data type of stored weights.
        """

        checkpoint_path = Path(checkpoint_path)
        if not checkpoint_path.exists():
            raise FileNotFoundError(
                f"Could not find checkpoint {str(checkpoint_path)}"
            )

        weight_dtype = to_numpy_dtype(weight_dtype)
        print(
            f"Load weights from {str(checkpoint_path)} (data type: {weight_dtype}"  # noqa: E501
        )

        self.weights = list()
        self.int8_weights = list()
        self.int8_scales = list()
        torch.cuda.empty_cache()

        def _load_from_file(fname):
            quant_sub_names = [
                "attention.query_key_value.weight",
                "attention.dense.weight",
                "dense_h_to_4h.weight",
                "dense_4h_to_h.weight",
            ]
            _weight = torch.from_numpy(
                np.fromfile(checkpoint_path / fname, dtype=weight_dtype)
            )
            _weight = _weight.to(compute_dtype)
            weight_index = len(self.weights)
            expected_shape = self.expected_weight_shapes[weight_index]

            try:
                if _weight.nelement() > 0:
                    _weight = _weight.reshape(expected_shape)
            except:  # noqa: E722
                raise ValueError(
                    f"num_heads, size_per_head, vocab_size, and max_seq_len must be the same "  # noqa: E501
                    f"as the ones during training (weight: {fname} expected shape: {expected_shape}, "  # noqa: E501
                    f"got shape: {_weight.shape})."
                )

            should_quantize = any(
                sub_name in fname for sub_name in quant_sub_names
            )
            if self.int8_mode != 0 and should_quantize:
                calibrate = self.weight_transpose_calibrate_quantize
                int8_weight, int8_scales = calibrate(_weight)

                # int8 weights should appear in same order as FP weights.
                # Move to device and add to the int8 list.
                dummy_weight = torch.empty(0, dtype=compute_dtype)
                if device is not None:
                    int8_weight = int8_weight.to(device)
                    int8_scales = int8_scales.to(device)
                    dummy_weight = dummy_weight.to(device)

                self.int8_weights.append(int8_weight)
                self.int8_scales.append(int8_scales)
                self.weights.append(dummy_weight)
            else:
                if device is not None:
                    _weight = _weight.to(device)
                self.weights.append(_weight)

        # Load
        # pylint:disable=line-too-long
        layer_offset = self.local_num_layers * self.pipeline_para_rank
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.input_layernorm.weight.bin"
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.input_layernorm.bias.bin"
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.attention.query_key_value.weight.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.attention.query_key_value.bias.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.attention.dense.weight.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.attention.dense.bias.bin"
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.post_attention_layernorm.weight.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.post_attention_layernorm.bias.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.mlp.dense_h_to_4h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.mlp.dense_h_to_4h.bias.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.mlp.dense_4h_to_h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.mlp.dense_4h_to_h.bias.bin"
            )
            for i in range(self.local_num_layers)
        ]

        if self.has_adapters:
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_attention_adapter.dense_h_to_4h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_attention_adapter.dense_h_to_4h.bias.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_attention_adapter.dense_4h_to_h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_attention_adapter.dense_4h_to_h.bias.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_h_to_4h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_h_to_4h.bias.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_4h_to_h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_4h_to_h.bias.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]

        assert len(self.weights) == len(
            self.expected_weight_shapes
        ), "Incorrect number of weights loaded"

class FtModuleBase:
    def __init__(self):
        self.weight = None

    @classmethod
    @abstractmethod
    def from_config(cls, config: GptInitModelParameters, **kwargs):
        raise NotImplementedError

    @abstractmethod
    def _initialize_model(self, force_init=False):
        raise NotImplementedError

    @abstractmethod
    def forward(self, *args, **kwargs):
        raise NotImplementedError

    def set_weight(self, weight: GptLayerWeights):
        old_weight_dtype = (
            self.weight.dtype if self.weight is not None else None
        )
        self.weight = weight
        if old_weight_dtype is None or old_weight_dtype != self.weight.dtype:
            self._initialize_model(force_init=True)

    @property
    def dtype(self):
        assert self.weight is not None
        return self.weight.dtype

    @property
    def device(self):
        assert self.weight is not None
        return self.weight.device

    def cuda(self, device=None):
        assert torch.cuda.is_available()
        self.weight.cuda(device)
        return self

    def to(self, device=None):
        self.weight.to(device)
        return self

    def float(self):
        self.weight.float()
        self._initialize_model(force_init=True)
        return self

    def half(self):
        self.weight.half()
        self._initialize_model(force_init=True)
        return self

    def bfloat16(self):
        self.weight.bfloat16()
        self._initialize_model(force_init=True)
        return self

class GptContextDecoder(FtModuleBase):
    def __init__(
        self,
        num_heads: int,
        size_per_head: int,
        inter_size: int,
        num_layers: int,
        tensor_para_size: int = 1,
        pipeline_para_size: int = 1,
        remove_padding: bool = True,
        shared_contexts_ratio: float = 1.0,
        layernorm_eps: float = 1e-6,
        layernorm_type: LayernormType = "pre_layernorm",
        activation_type: str = "gelu",
        has_adapters: bool = False,
        adapter_inter_size: int = 0,
        int8_mode: int = 0,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.size_per_head = size_per_head
        self.hidden_size = self.num_heads * self.size_per_head
        self.inter_size = inter_size
        self.num_layers = num_layers

        self.tensor_para_size = tensor_para_size
        self.pipeline_para_size = pipeline_para_size

        self.remove_padding = remove_padding
        self.shared_contexts_ratio = shared_contexts_ratio

        self.layernorm_eps = layernorm_eps
        self.layernorm_type = layernorm_type
        self.activation_type = activation_type
        self.has_adapters = has_adapters
        self.adapter_inter_size = adapter_inter_size

        assert int8_mode in [0, 1]
        self.int8_mode = int8_mode

        self.ft_op = None
        self.weight = None

    def __repr__(self):
        args_dict = dict(
            num_heads=self.num_heads,
            size_per_head=self.size_per_head,
            hidden_size=self.hidden_size,
            inter_size=self.inter_size,
            num_layers=self.num_layers,
            tensor_para_size=self.tensor_para_size,
            pipeline_para_size=self.pipeline_para_size,
            remove_padding=self.remove_padding,
            shared_contexts_ratio=self.shared_contexts_ratio,
            layernorm_eps=self.layernorm_eps,
            layernorm_type=self.layernorm_type,
            activation_type=self.activation_type,
            has_adapters=self.has_adapters,
            adapter_inter_size=self.adapter_inter_size,
            int8_mode=self.int8_mode,
        )
        args_str = ",\n    ".join([f"{k}: {v}" for k, v in args_dict.items()])
        return f"{self.__class__.__name__}[\n{    args_str}\n]"

    @classmethod
    def from_config(cls, config: GptInitModelParameters, **kwargs):
        return cls(
            num_heads=config.head_num,
            size_per_head=config.size_per_head,
            inter_size=4 * config.head_num * config.size_per_head,
            num_layers=config.layer_num,
            tensor_para_size=config.tensor_para_size,
            pipeline_para_size=config.pipeline_para_size,
            remove_padding=kwargs.get("remove_padding", True),
            shared_contexts_ratio=kwargs.get("shared_contexts_ratio", 1.0),
            layernorm_eps=config.layernorm_eps,
            layernorm_type=config.layernorm_type,
            activation_type=config.activation_type,
            has_adapters=config.has_adapters,
            adapter_inter_size=config.adapter_inter_size,
            int8_mode=config.int8_mode,
        )

    def _initialize_model(self, force_init=False):
        if self.weight is None:
            self.weight = GptLayerWeights(
                num_heads=self.num_heads,
                size_per_head=self.size_per_head,
                inter_size=self.inter_size,
                num_layers=self.num_layers,
                tensor_para_size=self.tensor_para_size,
                pipeline_para_size=self.pipeline_para_size,
                has_adapters=self.has_adapters,
                adapter_inter_size=self.adapter_inter_size,
                int8_mode=self.int8_mode,
            )
        if not force_init and self.ft_op is not None:
            return
        if self.ft_op is not None:
            del self.ft_op

        self.ft_op = (
            torch.classes.FasterTransformer.ParallelGptContextDecoderOp(
                self.num_heads,
                self.size_per_head,
                self.inter_size,
                self.num_layers,
                self.tensor_para_size,
                self.pipeline_para_size,
                self.layernorm_eps,
                self.layernorm_type,
                self.activation_type,
                self.has_adapters,
                self.adapter_inter_size,
                self.int8_mode,
                self.weight.weights,
                self.weight.int8_weights,
                self.weight.int8_scales,
                self.remove_padding,
            )
        )

    def forward(
        self,
        input_embeds: torch.Tensor,
        attention_mask: torch.Tensor,
        input_lengths: torch.IntTensor,
        memory_length: Optional[int] = None,
        compact_index: Optional[torch.IntTensor] = None,
        batch_to_compact_index: Optional[torch.IntTensor] = None,
        linear_bias_slopes: Optional[torch.Tensor] = None,
    ):
        """

        # Args.
            input_embeds: Tensor, (batch * beam, max_input_length, hidden_dim),
                input hidden states.
            attention_mask: Tensor, (batch * beam, max_input_length, max_input_length),
                input attention mask.
            input_lengths: (batch * beam,), input sequence lengths.
            memory_length: int, the length of memory to keep key/cache values.
            compact_index: IntTensor, (compact_batch_size,)
                The index of input sequences of a compact batch. If None, the FT op
                doesn't apply the shared context feature and as result the inference
                time may increase.
            batch_to_compact_index: IntTensor, (batch * beam,)
                The index map from the original input batch to the compact batch.
                This must be provided if compact_index is not None.
            linear_bias_slopes: (num_heads,)
                The slope per head of linear attention bias - ALiBi. If None, a base
                self attention will be performed.
        # Returns
            hidden_states: Tensor, (batch * beam, max_input_length, hidden_dim),
                decoder outputs.
            key_cache: Tensor, (num_layers, batch * beam, local_num_heads, size_per_head / x, memory_length, x), # noqa: E501
                key cache of attention of inputs.
                x = 16 / sizeof(T), memory_length = max_input_length or max_input_length + gen_length # noqa: E501
            value_cache: Tensor, (num_layers, batch * beam, local_num_heads, memory_length, hidden_dim) # noqa: E501
                value cache of attention
            last_token_hidden_states: Tensor, (batch * beam, hidden_dim)
                hidden states of the last input token.
        """
        self._initialize_model()
        # outputs: output hidden states
        (
            decoder_ouptut,
            key_cache,
            value_cache,
            last_token_hidden_states,
        ) = self.ft_op.forward(
            input_embeds,
            attention_mask,
            input_lengths,
            memory_length,
            compact_index,
            batch_to_compact_index,
            linear_bias_slopes,
        )
        return decoder_ouptut, key_cache, value_cache, last_token_hidden_states

class GptDecoder(FtModuleBase):
    def __init__(
        self,
        num_heads: int,
        size_per_head: int,
        inter_size: int,
        num_layers: int,
        tensor_para_size: int = 1,
        pipeline_para_size: int = 1,
        layernorm_eps: float = 1e-6,
        layernorm_type: LayernormType = "pre_layernorm",
        activation_type: str = "gelu",
        has_adapters: bool = False,
        adapter_inter_size: int = 0,
        int8_mode: int = 0,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.size_per_head = size_per_head
        self.hidden_size = self.num_heads * self.size_per_head
        self.inter_size = inter_size
        self.num_layers = num_layers

        self.tensor_para_size = tensor_para_size
        self.pipeline_para_size = pipeline_para_size

        self.layernorm_eps = layernorm_eps
        self.layernorm_type = layernorm_type
        self.activation_type = activation_type
        self.has_adapters = has_adapters
        self.adapter_inter_size = adapter_inter_size

        self.int8_mode = int8_mode

        self.ft_op = None
        self.weight = None

    def __repr__(self):
        args_dict = dict(
            num_heads=self.num_heads,
            size_per_head=self.size_per_head,
            hidden_size=self.hidden_size,
            inter_size=self.inter_size,
            num_layers=self.num_layers,
            tensor_para_size=self.tensor_para_size,
            pipeline_para_size=self.pipeline_para_size,
            layernorm_eps=self.layernorm_eps,
            layernorm_type=self.layernorm_type,
            activation_type=self.activation_type,
            has_adapters=self.has_adapters,
            adapter_inter_size=self.adapter_inter_size,
            int8_mode=self.int8_mode,
        )
        args_str = ",\n    ".join(
            [f"{k}: {v}" for k, v in args_dict.items()]
        )  # noqa: E501
        return f"{self.__class__.__name__}[\n    {args_str}\n]"

    @classmethod
    def from_config(cls, config: GptInitModelParameters, **kwargs):
        hidden_dim = config.head_num * config.size_per_head
        return cls(
            num_heads=config.head_num,
            size_per_head=config.size_per_head,
            inter_size=4 * hidden_dim,
            num_layers=config.layer_num,
            tensor_para_size=config.tensor_para_size,
            pipeline_para_size=config.pipeline_para_size,
            layernorm_eps=config.layernorm_eps,
            layernorm_type=config.layernorm_type,
            activation_type=config.activation_type,
            has_adapters=config.has_adapters,
            adapter_inter_size=config.adapter_inter_size,
            int8_mode=config.int8_mode,
        )

    def _initialize_model(self, force_init=False):
        if self.weight is None:
            self.weight = GptLayerWeights(
                num_heads=self.num_heads,
                size_per_head=self.size_per_head,
                inter_size=self.inter_size,
                num_layers=self.num_layers,
                tensor_para_size=self.tensor_para_size,
                pipeline_para_size=self.pipeline_para_size,
                has_adapters=self.has_adapters,
                adapter_inter_size=self.adapter_inter_size,
                int8_mode=self.int8_mode,
            )
        if not force_init and self.ft_op is not None:
            return
        if self.ft_op is not None:
            del self.ft_op
        self.ft_op = torch.classes.FasterTransformer.ParallelGptDecoderOp(
            self.num_heads,
            self.size_per_head,
            self.inter_size,
            self.num_layers,
            self.tensor_para_size,
            self.pipeline_para_size,
            self.layernorm_eps,
            self.layernorm_type,
            self.activation_type,
            self.has_adapters,
            self.adapter_inter_size,
            self.weight.int8_mode,
            self.weight.weights,
            self.weight.int8_weights,
            self.weight.int8_scales,
        )

    def forward(
        self,
        max_input_length: int,
        step: int,
        ite: int,
        input_embeds: torch.Tensor,
        sequence_lengths: torch.IntTensor,
        key_cache: torch.Tensor,
        value_cache: torch.Tensor,
        finished: torch.BoolTensor,
        total_padding_tokens: torch.IntTensor,
        masked_tokens: torch.BoolTensor,
        cache_indirection: Optional[torch.IntTensor] = None,
        linear_bias_slopes: Optional[torch.Tensor] = None,
    ):
        """

        # Args.
            max_input_length: int, maximum input context length.
            step: int, the current step index.
            ite: int, local batch iteration.
            input_embeds: Tensor, (local_batch * beam, hidden_dim),
                input hidden state to decoder.
            sequence_lengths: IntTensor, (local_batch * beam,),
                the current sequence lengths.
            key_cache: Tensor, key cache buffer.
            value_cache: Tensor, value cache buffer.
            finished: BoolTensor, (local_batch * beam,),
                whether to finish sentence generation.
            total_padding_tokens IntTensor, (local_batch * beam,),
                the number of padded tokens.
            masked_tokens: BoolTensor, (local_batch * beam, memory_length),
                a mask tensor that indicates padded tokens.
            cache_indirection: IntTensor, (local_batch * beam,),
                cache of beam positions if needed if beam > 1.
            linear_bias_slopes Tensor, (num_heads,)
                slopes head of linear position bias (ALiBi) (optional).
        # Returns
            IntTensor, (batch * beam,) output token ids.
        """

        self._initialize_model()

        outputs = self.ft_op.forward(
            max_input_length,
            step,
            ite,
            input_embeds,
            sequence_lengths,
            finished,
            total_padding_tokens,
            masked_tokens,
            key_cache,
            value_cache,
            cache_indirection,
            linear_bias_slopes,
        )
        return outputs[0]

class Gpt:
    def __init__(
        self,
        num_heads: int,
        size_per_head: int,
        num_layers: int,
        vocab_size: int,
        start_id: int,
        end_id: int,
        lib_path: PathLike,
        tensor_para_size: int = 1,
        pipeline_para_size: int = 1,
        remove_padding: bool = True,
        shared_contexts_ratio: float = 1.0,
        layernorm_eps: float = 1e-6,
        layernorm_type: LayernormType = "pre_layernorm",
        activation_type: str = "gelu",
        has_positional_encoding: bool = True,
        max_seq_len: int = 0,
        has_pre_decoder_layernorm: bool = False,
        has_post_decoder_layernorm: bool = True,
        has_adapters: bool = False,
        adapter_inter_size: int = 0,
        int8_mode: int = 0,
        inference_data_type: Optional[str] = None,
        weights_data_type: str = "fp32",
        use_fp32_to_compute_logit: bool = False,
        **kwargs,
    ):
        super().__init__()

        inference_data_type = inference_data_type or weights_data_type

        self.config = GptInitModelParameters(
            head_num=num_heads,
            size_per_head=size_per_head,
            layer_num=num_layers,
            max_seq_len=max_seq_len,
            tensor_para_size=tensor_para_size,
            vocab_size=vocab_size,
            start_id=start_id,
            end_id=end_id,
            pipeline_para_size=pipeline_para_size,
            data_type=inference_data_type,
            weights_data_type=weights_data_type,
            layernorm_eps=layernorm_eps,
            layernorm_type=layernorm_type,
            activation_type=activation_type,
            has_positional_encoding=has_positional_encoding,
            has_pre_decoder_layernorm=has_pre_decoder_layernorm,
            has_post_decoder_layernorm=has_post_decoder_layernorm,
            has_adapters=has_adapters,
            adapter_inter_size=adapter_inter_size,
            int8_mode=int8_mode,
            sparse=kwargs.get("sparse", False),
        )
        self.use_fp32_to_compute_logit = use_fp32_to_compute_logit

        self.weight = None
        self.shared_contexts_ratio = shared_contexts_ratio

        torch.classes.load_library(os.path.abspath(lib_path))

        # Embeddings to encode or decode tokens.
        hidden_dim = num_heads * size_per_head

        # Pad vocab size for FT.
        local_vocab_size = math.ceil(
            self.config.vocab_size / self.config.tensor_para_size
        )
        if self.config.data_type == "fp16":
            local_vocab_size = math.ceil(local_vocab_size / 8) * 8
        self.vocab_size_padded = (
            local_vocab_size * self.config.tensor_para_size
        )
        self.vocab_size = self.config.vocab_size

        self.decode_op = torch.classes.FasterTransformer.DynamicDecodeOp(
            self.vocab_size,
            self.vocab_size_padded,
            self.config.tensor_para_size,
            self.config.pipeline_para_size,
            torch.float,
        )

        self._parameters = {}

        def register_param(name, p):
            self._parameters[name] = p
            setattr(self, name, p)

        register_param(
            "context_decoder",
            GptContextDecoder.from_config(
                self.config,
                remove_padding=remove_padding,
                shared_contexts_ratio=shared_contexts_ratio,
                **kwargs,
            ),
        )
        register_param(
            "decoder", GptDecoder.from_config(self.config, **kwargs)
        )

        compute_dtype = to_torch_dtype(inference_data_type)

        if comm.is_pipeline_group_first():
            register_param(
                "word_embedding",
                torch.nn.Embedding(
                    self.vocab_size_padded, hidden_dim, dtype=compute_dtype
                ),
            )
            self._mask_padded_vocab_weights(self.word_embedding.weight)
            if self.config.has_positional_encoding:
                register_param(
                    "position_encoding",
                    torch.nn.Embedding(
                        self.config.max_seq_len,
                        hidden_dim,
                        dtype=compute_dtype,
                    ),
                )
            else:
                self.position_encoding = None
            if self.config.has_pre_decoder_layernorm:
                register_param(
                    "pre_decoder_layernorm",
                    torch.nn.LayerNorm(
                        hidden_dim, eps=layernorm_eps, dtype=compute_dtype
                    ),
                )
            else:
                self.pre_decoder_layernorm = None

        if comm.is_pipeline_group_last():
            if has_post_decoder_layernorm:
                register_param(
                    "post_decoder_layernorm",
                    torch.nn.LayerNorm(
                        hidden_dim, eps=layernorm_eps, dtype=compute_dtype
                    ),
                )
            else:
                self.post_decoder_layernorm = None
            self.lm_head_ctype = (
                compute_dtype
                if not self.use_fp32_to_compute_logit
                else torch.float32
            )
            register_param(
                "lm_head",
                torch.nn.Linear(
                    hidden_dim,
                    self.vocab_size_padded,
                    bias=False,
                    dtype=self.lm_head_ctype,
                ),
            )
            self._mask_padded_vocab_weights(self.lm_head.weight)

    @classmethod
    def from_config(cls, config: GptInitModelParameters, **kwargs):
        return cls(
            num_heads=config.head_num,
            size_per_head=config.size_per_head,
            num_layers=config.layer_num,
            max_seq_len=config.max_seq_len,
            tensor_para_size=config.tensor_para_size,
            vocab_size=config.vocab_size,
            start_id=config.start_id,
            end_id=config.end_id,
            pipeline_para_size=config.pipeline_para_size,
            inference_data_type=config.data_type,
            weights_data_type=config.weights_data_type,
            layernorm_eps=config.layernorm_eps,
            layernorm_type=config.layernorm_type,
            activation_type=config.activation_type,
            has_positional_encoding=config.has_positional_encoding,
            has_pre_decoder_layernorm=config.has_pre_decoder_layernorm,
            has_post_decoder_layernorm=config.has_post_decoder_layernorm,
            has_adapters=config.has_adapters,
            adapter_inter_size=config.adapter_inter_size,
            int8_mode=config.int8_mode,
            **kwargs,
        )

    def load(
        self,
        checkpoint_path: PathLike,
        inference_data_type: Optional[Union[str, torch.dtype]] = None,
        config: Optional[GptInitModelParameters] = None,
        device: Optional[Union[str, int, torch.device]] = None,
    ):

        checkpoint_path = Path(checkpoint_path)
        device = device or comm.get_device()
        config = config or self.config

        compute_dtype = to_torch_dtype(inference_data_type or self.dtype)

        self.weight = GptLayerWeights.from_config(config)
        self.weight.load(
            checkpoint_path, compute_dtype, config.weights_data_type, device
        )

        self.context_decoder.set_weight(self.weight)
        self.decoder.set_weight(self.weight)

        weight_dtype = to_numpy_dtype(config.weights_data_type)

        def _safe_load_from_bin(param: torch.nn.Parameter, fname):
            if (checkpoint_path / fname).exists():
                # np_w is 1-D array since a bin file doesn't have shape info.
                w_ = np.fromfile(checkpoint_path / fname, dtype=weight_dtype)
                param.data = (
                    torch.from_numpy(w_)
                    .reshape(param.data.shape)
                    .to(compute_dtype)
                )
            else:
                raise FileNotFoundError(f"Faile to load {fname}")

        def _safe_load_lm_head_from_bin(param, fname, ctype):
            if (checkpoint_path / fname).exists():
                shape = (
                    self.vocab_size,
                    self.config.head_num * self.config.size_per_head,
                )
                # np_w is 1-D array since a bin file doesn't have shape info.
                w_ = np.fromfile(checkpoint_path / fname, dtype=weight_dtype)
                param.data = param.data.to(ctype)
                param.data[: self.vocab_size, :] = (
                    torch.from_numpy(w_).reshape(shape).to(ctype)
                )
            else:
                print(f"Faile to load {fname}")
                torch.nn.init.normal_(param).to(compute_dtype)
            self._mask_padded_vocab_weights(param)

        # pylint:disable=line-too-long
        if comm.is_pipeline_group_first():
            _safe_load_lm_head_from_bin(
                self.word_embedding.weight, "model.wte.bin", compute_dtype
            )
            self._mask_padded_vocab_weights(self.word_embedding.weight)
            if self.position_encoding is not None:
                _safe_load_from_bin(
                    self.position_encoding.weight, "model.wpe.bin"
                )
            if self.pre_decoder_layernorm is not None:
                _safe_load_from_bin(
                    self.pre_decoder_layernorm.weight,
                    "model.pre_decoder_layernorm.weight.bin",
                )
                _safe_load_from_bin(
                    self.pre_decoder_layernorm.bias,
                    "model.pre_decoder_layernorm.bias.bin",
                )
        if comm.is_pipeline_group_last():
            if self.post_decoder_layernorm is not None:
                _safe_load_from_bin(
                    self.post_decoder_layernorm.weight,
                    "model.final_layernorm.weight.bin",
                )
                _safe_load_from_bin(
                    self.post_decoder_layernorm.bias,
                    "model.final_layernorm.bias.bin",
                )
            if (checkpoint_path / "model.lm_head.weight.bin").exists():
                _safe_load_lm_head_from_bin(
                    self.lm_head.weight,
                    "model.lm_head.weight.bin",
                    self.lm_head_ctype,
                )
            else:
                if self.use_fp32_to_compute_logit:
                    _safe_load_lm_head_from_bin(
                        self.lm_head.weight, "model.wte.bin", torch.float32
                    )
                else:
                    # In this branch we can share the pre and post
                    # decoder embeddings, but ONLY pipeline size is 1.
                    # When pipeline size > 1, these two weights will end up on
                    # different GPUs, so we must load the
                    # post decoder weight again (else case).
                    if comm.get_pipeline_para_size() == 1:
                        self.lm_head.weight = self.word_embedding.weight
                    else:
                        _safe_load_lm_head_from_bin(
                            self.lm_head.weight, "model.wte.bin", compute_dtype
                        )

        self.to(device)

    @property
    def dtype(self):
        assert self.weight is not None
        return self.weight.dtype

    @property
    def device(self):
        assert self.weight is not None
        return self.weight.device

    def cuda(self, device=None):
        assert torch.cuda.is_available()
        for name, param in self._parameters.items():
            setattr(self, name, param.cuda(device))
        return self

    def to(self, device=None):
        for name, param in self._parameters.items():
            setattr(self, name, param.to(device))
        return self

    def float(self):
        for name, param in self._parameters.items():
            setattr(self, name, param.float())
        return self

    def half(self):
        for name, param in self._parameters.items():
            setattr(self, name, param.half())
        return self

    def bfloat16(self):
        for name, param in self._parameters.items():
            setattr(self, name, param.bfloat16())
        return self

    def _mask_padded_vocab_weights(self, weight: torch.Tensor):
        assert self.vocab_size_padded >= self.vocab_size
        if self.vocab_size_padded > self.vocab_size:
            weight.data[self.vocab_size :, ...] = 0  # noqa: E203

    def generate_pad_mask(self, input_lengths, memory_length, init_step=0):
        """Generate a pad mask tensor.

        # Args.
            input_lengths: (batch_size * beam_width,), input lengths
            memory_length: the length of key/value cache memory.
            init_step: int, initial step.
        # Return
            masked_tokens: BoolTensor,
                (batch_size * beam_width, memory_length),
                True if init_step + input_length[i] <= j <
                    init_step + max_input_length,
                where i is a batch-beam index and j is a time step
                modulo by memory_length.
        """
        max_input_length = input_lengths.max()
        input_lengths = input_lengths.unsqueeze(1)
        shift = init_step % memory_length
        step_indices = torch.arange(
            init_step, init_step + memory_length, device=input_lengths.device
        )
        step_indices = (
            step_indices.roll(shift)
            .unsqueeze(0)
            .tile(input_lengths.shape[0], 1)
        )
        masked_tokens = torch.logical_and(
            step_indices >= input_lengths,
            step_indices < init_step + max_input_length,
        )
        return masked_tokens

    def get_local_batch_size(self, batch_size):
        """Get a local batch size by the same way that FT Gpt does."""
        local_batch_size = batch_size
        pp_size = self.decoder.pipeline_para_size
        if pp_size > 1:
            if local_batch_size % pp_size == 0:
                local_batch_size //= pp_size
            while local_batch_size > 1024 and local_batch_size % 2 == 0:
                local_batch_size //= 2
        return local_batch_size

    @torch.no_grad()
    def generate(
        self,
        input_token_ids: torch.IntTensor,
        input_lengths: torch.IntTensor,
        gen_length: int,
        eos_token_id: Optional[int] = None,
        local_batch_size: Optional[int] = None,
        beam_width: int = 1,
        top_k: Optional[torch.IntTensor] = None,
        top_p: Optional[torch.FloatTensor] = None,
        top_p_decay: Optional[torch.FloatTensor] = None,
        top_p_min: Optional[torch.FloatTensor] = None,
        top_p_reset_ids: Optional[torch.IntTensor] = None,
        temperature: Optional[torch.FloatTensor] = None,
        repetition_penalty: Optional[torch.FloatTensor] = None,
        presence_penalty: Optional[torch.FloatTensor] = None,
        min_length: Optional[torch.IntTensor] = None,
        len_penalty: Optional[torch.FloatTensor] = None,
        beam_search_diversity_rate: Optional[torch.FloatTensor] = None,
        stop_words_list: Optional[torch.IntTensor] = None,
        bad_words_list: Optional[torch.IntTensor] = None,
        sequence_limit_lengths: Optional[torch.IntTensor] = None,
        random_seed: Optional[torch.LongTensor] = None,
        memory_length: Optional[int] = None,
        return_output_length: bool = False,
        return_log_probs: bool = False,
    ):
        """

        # Args.
            input_token_ids: IntTensor, (batch_size, max_input_length),
                input hidden state to decoder.
            input_lengths: IntTensor, (batch_size),
                the lengths of input context sequences.
            gen_length: int, the number of tokens to generate.
            local_batch_size: int, optional, a batch size of
                local iteration. (disabled)
            eos_token_id: int, eos token id.
            beam_width: int, number of beams for beam search.
                If 1, sampling decode will be used.
            top_k: IntTensor, (batch_size,) top-k sampling.
                The number of most probable tokens to keep
                for sampling per sentence in a batcch.
            top_p: FloatTensor, (batch_size,), top-p sampling.
                The cumulative probability
                of to filter the set of most probable tokens.
            top_p_decay: FloatTensor, (batch_size,)
                The decay of top-p value for top_p sampling.
            top_p_min: FloatTensor, (batch_size,)
                The minimum top p values in top-p decaying.
            top_p_reset_ids: IntTensor, (batch_size,)
                reset ids for resetting top_p values for top p sampling
            temperature: FloatTensor, (batch_size,),
                The temperature value for smoothing the logit distribution.
            repetition_penalty: FloatTensor, (batch_size,),
                The repetition penalty.
            presence_penalty: FloatTensor, (batch_size,),
                The presence penalty, which is exclusive with
                repetition_penalty.
                Only one of repetition and presence penalties is allowed.
            min_length: IntTensor, (batch_size,),
                Minimum length for each sentences. EOS is masked if length is
                below min.
            len_penalty: FloatTensor, (batch_size,)
                The exponent of the length penalty of beam scores.
            beam_search_diversity_rate: FloatTensor, (batch_size,),
                The diversity rate of beam search.
            stop_words_list: IntTensor, (batch_size, 2, stop_words_length)
                When FT generates words in this list, it will stop the
                generation. An extension of stop id.
            bad_words_list IntTensor, (batch_size, 2, bad_words_length)
                The words in the list will never be sampled.
            sequence_limit_lengths: IntTensor, (batch_size,), The maximum
                length of a generated sequence.
            memory_length: int, the length of cache memory. If None, it will
                be max_input_length + gen_length.
        # Returns
            IntTensor, (batch_size, beam_width, max_seq_length) output
            token ids.
        """
        assert (
            self.weight is not None
        ), "Please call load() first to initialize weights."

        input_token_ids = input_token_ids.type(torch.int32).to(self.device)
        input_lengths = input_lengths.type(torch.int32).to(self.device)

        batch_size = len(input_token_ids)
        max_input_length = input_token_ids.shape[-1]
        max_seq_length = max_input_length + gen_length
        memory_length = memory_length or max_seq_length

        # TODO: Enable local batch later. We currently disable local batching due to # noqa: E501
        #   an input mismatch issue of FT's decode_op: FT's decode_op requires logits # noqa: E501
        #   of shape (batch_size, ...) but we have logits of shape (local_batch_size, ...) # noqa: E501
        #   After fixing FT's side, we will enable local batch.
        # local_batch_size = local_batch_size or self.get_local_batch_size(batch_size) # noqa: E501
        # num_local_batches, last_chunk = divmod(batch_size, local_batch_size)
        # if last_chunk > 0:
        #     num_local_batches += 1
        assert local_batch_size is None or local_batch_size == batch_size
        local_batch_size = batch_size
        num_local_batches = 1

        device = self.device

        eos_token_id = (
            eos_token_id if eos_token_id is not None else self.config.end_id
        )
        assert (
            eos_token_id is not None
        ), "eos_token-id must be specified in generation."
        eos_token_ids = eos_token_id * torch.ones(
            batch_size, dtype=torch.int32, device=device
        )
        assert repetition_penalty is None or presence_penalty is None, (
            "Found ambiguous parameters repetition_penalty and "
            "presence_penalty which are mutually exclusive. "
            "Please provide one of repetition_penalty and presence_penalty."
        )
        # Setup decoder_op prior to calling the forward function.
        self.decode_op.setup(
            batch_size,
            beam_width,
            top_k,
            top_p,
            temperature,
            repetition_penalty,
            presence_penalty,
            min_length,
            len_penalty,
            beam_search_diversity_rate,
            random_seed,
            top_p_decay,
            top_p_min,
            top_p_reset_ids,
        )

        # Prepare input and output arguments.
        if beam_width > 1:
            # Tiling for beam search.
            input_token_ids = input_token_ids.repeat(1, beam_width).view(
                batch_size * beam_width, -1
            )
            input_lengths = (
                input_lengths.view(-1, 1).repeat(1, beam_width).view(-1)
            )
            if sequence_limit_lengths is not None:
                sequence_limit_lengths = (
                    sequence_limit_lengths.view(-1, 1)
                    .repeat(1, beam_width)
                    .view(-1)
                )
            # src/tgt cache indirections.
            cache_indirection = torch.zeros(
                (2, batch_size, beam_width, memory_length),
                dtype=torch.int32,
                device=device,
            )
            parent_ids = torch.zeros(
                max_seq_length,
                batch_size * beam_width,
                dtype=torch.int32,
                device=device,
            )
        else:
            cache_indirection = None
            src_cache_indirection = None
            tgt_cache_indirection = None
            parent_ids = None

        pad_lengths = max_input_length - input_lengths
        # Since tril() doesn't support bf16 dtype,
        # we create of bool type and then cast it to dtype.
        attention_mask = (
            torch.ones(
                (max_input_length, max_input_length),
                dtype=torch.bool,
                device=device,
            )
            .tril()
            .unsqueeze(0)
            .tile(input_token_ids.shape[0], 1, 1)
            .to(self.dtype)
        )
        for b, input_length in enumerate(input_lengths):
            attention_mask[b, input_length:, ...] = 0
        masked_tokens = self.generate_pad_mask(input_lengths, memory_length)
        finished = torch.zeros_like(input_lengths).bool()
        sequence_lengths = (max_input_length - 1) * torch.ones_like(
            input_lengths
        )

        if return_log_probs or beam_width > 1:
            cum_log_probs = torch.zeros(batch_size * beam_width, device=device)
            output_log_probs = torch.zeros(
                (gen_length, batch_size * beam_width), device=device
            )
        else:
            cum_log_probs = None
            output_log_probs = None

        # Contiguous buffer for each decode_op step,
        # it will be transposed tensor for the final output.
        output_token_ids = torch.zeros(
            (max_seq_length, batch_size * beam_width),
            dtype=torch.int32,
            device=device,
        )
        output_token_ids[:max_input_length, ...] = input_token_ids.T

        if comm.is_pipeline_group_first():
            # Prepare input tensors of decoder.
            input_embeds = self.word_embedding(input_token_ids)
            if self.position_encoding is not None:
                position_ids = torch.arange(
                    0, max_input_length, dtype=torch.int, device=device
                )
                position_ids = position_ids.unsqueeze(0).view(
                    -1, max_input_length
                )
                input_embeds += self.position_encoding(position_ids)
            if self.pre_decoder_layernorm is not None:
                input_embeds = self.pre_decoder_layernorm(input_embeds)
        else:
            # Dummy input_embeds
            input_embeds = torch.empty(
                size=(
                    batch_size * beam_width,
                    max_input_length,
                    self.context_decoder.hidden_size,
                ),
                dtype=self.context_decoder.dtype,
                device=device,
            )

        use_shared_contexts = (
            (self.shared_contexts_ratio > 0.0)
            and (max_input_length >= 1)
            and (batch_size > 1)
        )
        batch_to_compact, compact_to_batch = None, None
        if use_shared_contexts:
            find_context_duplications = (
                torch.ops.fastertransformer.find_context_duplications
            )
            batch_to_compact, compact_to_batch = find_context_duplications(
                input_token_ids
            )
            use_shared_contexts = (
                compact_to_batch.shape[0]
                <= self.shared_contexts_ratio * batch_size
            )

            if not use_shared_contexts:
                batch_to_compact, compact_to_batch = None, None

        profiler.start("ft-context-decoder")
        (
            _,
            k_cache,
            v_cache,
            last_token_hidden_states,
        ) = self.context_decoder.forward(
            input_embeds=input_embeds,
            attention_mask=attention_mask,
            input_lengths=input_lengths,
            memory_length=memory_length,
            batch_to_compact_index=batch_to_compact,
            compact_index=compact_to_batch,
        )
        profiler.stop("ft-context-decoder")

        for step in range(max_input_length, max_seq_length):
            src_indir_idx = (step - max_input_length) % 2
            tgt_indir_idx = 1 - src_indir_idx

            is_generation_done = torch.tensor(
                [True], dtype=torch.bool, device=device
            )
            for ite in range(num_local_batches):
                # The indices of the current local batch-beam.
                bbidx = range(
                    ite * local_batch_size * beam_width,
                    min(
                        (ite + 1) * local_batch_size * beam_width,
                        batch_size * beam_width,
                    ),
                )
                if cache_indirection is not None:
                    bidx = range(
                        ite * local_batch_size,
                        min((ite + 1) * local_batch_size, batch_size),
                    )
                    src_cache_indirection = cache_indirection[
                        src_indir_idx, bidx, ...
                    ]
                    tgt_cache_indirection = cache_indirection[
                        tgt_indir_idx, bidx, ...
                    ]

                if step == max_input_length:
                    hidden_states = last_token_hidden_states[bbidx, ...]
                else:
                    if comm.is_pipeline_group_first():
                        input_embeds = self.word_embedding(
                            output_token_ids[step - 1, bbidx]
                        )
                        if self.position_encoding is not None:
                            position_ids = (step - 1) * torch.ones_like(
                                pad_lengths[bbidx]
                            )
                            input_embeds += self.position_encoding(
                                position_ids
                            )
                        if self.pre_decoder_layernorm is not None:
                            input_embeds = self.pre_decoder_layernorm(
                                input_embeds
                            )
                    else:
                        # Dummy input_imbeds
                        input_embeds = torch.empty(
                            size=(len(bbidx), self.decoder.hidden_size),
                            dtype=self.decoder.dtype,
                            device=device,
                        )

                    profiler.start("ft-decoder")
                    hidden_states = self.decoder.forward(
                        max_input_length=max_input_length,
                        step=step,
                        ite=ite,
                        input_embeds=input_embeds,
                        sequence_lengths=sequence_lengths[bbidx],
                        key_cache=k_cache,
                        value_cache=v_cache,
                        finished=finished[bbidx],
                        total_padding_tokens=pad_lengths[bbidx],
                        cache_indirection=src_cache_indirection,
                        masked_tokens=masked_tokens[bbidx, ...],
                    )
                    profiler.stop("ft-decoder")

                if comm.is_pipeline_group_last():
                    if self.post_decoder_layernorm is not None:
                        hidden_states = self.post_decoder_layernorm(
                            hidden_states
                        )

                    # We use logits of fp32 type to avoid overflow issue.
                    if self.use_fp32_to_compute_logit:
                        # The FT GPT op internally uses FP32 compute type
                        # for matrix multiplication.
                        # This will produce the same result with the
                        # end-to-end FT's GPT op.
                        logits = torch.nn.functional.linear(
                            hidden_states.float(), self.lm_head.weight
                        )
                    else:
                        logits = self.lm_head(hidden_states).float()

                    profiler.start("ft-decode")
                    should_stop = self.decode_op.forward(
                        logits.view(batch_size, beam_width, -1),
                        step,
                        max_input_length,
                        ite,
                        local_batch_size,
                        eos_token_ids,
                        top_k,
                        top_p,
                        temperature,
                        repetition_penalty,
                        presence_penalty,
                        min_length,
                        len_penalty,
                        beam_search_diversity_rate,
                        top_p_decay,
                        top_p_min,
                        top_p_reset_ids,
                        None,
                        input_lengths,
                        sequence_limit_lengths,
                        stop_words_list,
                        bad_words_list,
                        src_cache_indirection,
                        output_token_ids.view(-1, batch_size, beam_width),
                        finished,
                        sequence_lengths,
                        cum_log_probs,
                        output_log_probs,
                        parent_ids,
                        tgt_cache_indirection,
                    )
                    profiler.stop("ft-decode")
                    is_generation_done &= should_stop

            # Broadcast from the last pipeline node if needed.
            profiler.start("ft-bcast")
            tensors_to_bcast = [
                output_token_ids[step, ...],
                finished,
                sequence_lengths,
                is_generation_done,
            ]
            if beam_width > 1:
                tensors_to_bcast.append(tgt_cache_indirection)
            self.decode_op.broadcast_from_last_pipeline(tensors_to_bcast)
            profiler.stop("ft-bcast")

            if is_generation_done or finished.all():
                break

        # Transpose (L, batch, beam) -> (batch, beam, L)
        output_token_ids = output_token_ids.view(
            -1, batch_size, beam_width
        ).permute(1, 2, 0)

        # Increase sequence_length by 1 because the sequence length of time step t is t - 1. # noqa: E501
        sequence_lengths += 1

        # Outputs
        output_dict = dict(output_token_ids=output_token_ids)
        if return_output_length:
            output_dict["output_lengths"] = sequence_lengths
        if return_log_probs:
            output_dict["cum_log_probs"] = cum_log_probs
            output_dict["output_log_probs"] = output_log_probs
        return output_dict


import os
import sys
from tqdm import tqdm
from pymongo import MongoClient

# Connect to the mongo db client
client = MongoClient("mongodb+srv://zach:G8GqPsUgP6b9VUvc"
        "@cluster0-svcn3.gcp.mongodb.net/test?retryWrites=true")

# Define the ocr_db to write to in Mongo
db = client[str(sys.argv[1])]

def mongo_writer(file):
	with open (file, "r") as f:
		lines = f.readlines()

	for line in tqdm(lines):
		obj = eval(line)
		collection_name = 'reviews'
		collection = db[collection_name]
		collection.insert_one(obj)

if __name__ == '__main__':

	files = ["ForkPoolWorker-1.txt", "ForkPoolWorker-2.txt", "ForkPoolWorker-3.txt", "ForkPoolWorker-4.txt"]

	for file in files:
		mongo_writer(file)


# -*- encoding:utf-8 -*-
# @brief:  ......
# @date:   2023.05.10 14:40:50

import os
from utils.utils_cmn import CmnUtils
from utils.utils_file import FileUtils
from utils.utils_logger import LoggerUtils
from utils.utils_properties import PropertiesUtils

# -------------------------------------------------------------
class WingEnv:
    g_space_path = None
    g_space_name = None
    g_space_branch = None
    g_space_manifest = None
    g_wing_path = None
    g_remote_host = None
    g_remote_manifest = None
    g_env_path = None

    @classmethod
    def init(cls, spacePath, envPath):
        cls.g_space_path = spacePath
        cls.g_env_path = envPath
        cls.doLoad()

    @classmethod
    def getSpacePath(cls):
        return cls.g_space_path

    @classmethod
    def getEnvPath(cls):
        return cls.g_env_path

    @classmethod
    def getWingPath(cls):
        if cls.g_wing_path is None:
            cls.g_wing_path = os.path.expanduser("~") + os.sep + '.wing/wing'  # such as: /Users/${username}/.wing/wing
        return cls.g_wing_path

    @classmethod
    def getSpaceName(cls):
        return cls.g_space_name

    @classmethod
    def setSpaceName(cls, v):
        cls.g_space_name = v
        cls.doSave()

    @classmethod
    def isLocalMode(cls):
        return CmnUtils.isEmpty(WingEnv.getSpaceRemoteManifestGit())

    @classmethod
    def getSpaceManifestFile(cls):
        return cls.g_space_manifest

    @classmethod
    def setSpaceManifestFile(cls, v):
        cls.g_space_manifest = v
        cls.doSave()

    @classmethod
    def getSpaceBranch(cls):
        return cls.g_space_branch

    @classmethod
    def setSpaceBranch(cls, v):
        cls.g_space_branch = v
        cls.doSave()

    @classmethod
    def getSpaceConfigFile(cls):
        return cls.g_space_path + os.sep + '.wing' + os.sep + 'space.json'

    @classmethod
    def doSave(cls):
        cfg = {}
        if cls.g_space_name is not None: cfg['space'] = cls.g_space_name
        if cls.g_space_branch is not None: cfg['branch'] = cls.g_space_branch
        if cls.g_space_manifest is not None: cfg['manifest'] = cls.g_space_manifest
        # add more here ...
        FileUtils.saveJsonToFile(cls.getSpaceConfigFile(), cfg)

    @classmethod
    def doLoad(cls):
        cfg = FileUtils.loadJsonByFile(cls.getSpaceConfigFile())
        if None == cfg: return
        if 'space' in cfg: cls.g_space_name = cfg['space']
        if 'branch' in cfg: cls.g_space_branch = cfg['branch']
        if 'manifest' in cfg: cls.g_space_manifest = cfg['manifest']
        # add more here ...

    @classmethod
    def getSpaceRemoteManifestGit(cls):
        if cls.g_remote_manifest is None: cls.__do_load_config__()
        # assert cls.g_remote_manifest is not None, 'Error: Invalid space ' + cls.getSpaceName()
        return cls.g_remote_manifest

    @classmethod
    def getSpaceRemoteHost(cls):
        if cls.g_remote_host is None: cls.__do_load_config__()
        # assert cls.g_remote_host is not None, 'Error: Invalid space ' + cls.getSpaceName()
        return cls.g_remote_host

    @classmethod
    def __do_load_config__(cls):
        spaces = PropertiesUtils.getAll(os.path.dirname(cls.getWingPath()) + os.sep + 'space.properties')
        if CmnUtils.isEmpty(spaces):
            LoggerUtils.w('No space found, do "wing -space add" first')
            assert 0
            return
        if cls.getSpaceName() not in spaces:
            LoggerUtils.w('Space "' + cls.getSpaceName() + '" not found, do "wing -space add" first')
            assert 0
            return
        items = spaces[cls.getSpaceName()].split(',')
        cls.g_remote_host = items[0] if 0 < len(items) else ''
        cls.g_remote_manifest = items[1] if 1 < len(items) else ''

    @staticmethod
    def isOsLinux():
        return CmnUtils.isOsLinux()

    @staticmethod
    def isOsWindows():
        return CmnUtils.isOsWindows()

    @staticmethod
    def isOsMac():
        return CmnUtils.isOsMac()


import wave
import numpy as np
import scipy.signal as signal
from real_time_classifier import Yamnet, Predictor

yamnet = Yamnet()
predictor = Predictor()

def eval_wav(file_name):
    with wave.open(file_name, "rb") as f:
        bytes = f.readframes(f.getnframes())
        rate = f.getframerate()
        int_samples = np.frombuffer(bytes, dtype=np.int16)
        s = int_samples.astype(np.float32)
        s = s / (2 ** 15 - 1)
        s = signal.resample(s, 32000)
        embeddings = yamnet.predict(s)
        return predictor.predict(embeddings)

print(eval_wav("training-data/recognized/tmp-1621909207.982633.wav"))


import os

from ckeditor.fields import RichTextField
from django.db import models
from django.template import defaultfilters
from embed_video.fields import EmbedVideoField
from imagekit.models import ImageSpecField
from imagekit.processors import Adjust
from solo.models import SingletonModel

def upload_image_to(instance, filename):
    from django.utils.timezone import now
    filename_base, filename_ext = os.path.splitext(filename)
    return 'images/%s%s%s' % (
        filename_base,
        now().strftime("%Y%m%d%H%M%S"),
        filename_ext.lower(),
    )

class Proyect(models.Model):
    name_es = models.CharField(
        'Name spanish', max_length=140, null=False, blank=False)
    name_en = models.CharField(
        'Name english', max_length=140, null=True, blank=True)
    location = models.CharField(
        'Location', max_length=140, null=True, blank=True)
    date = models.CharField('Fechas', max_length=140, null=True, blank=True)
    description_es = RichTextField()
    description_en = RichTextField()
    socialText_es = models.CharField(
        'Social spanish', max_length=100, null=True, blank=True)
    socialText_en = models.CharField(
        'Social english', max_length=100, null=True, blank=True)
    pub_date = models.DateTimeField(
        'Created', editable=False, auto_now_add=True)
    mainImage = models.ImageField(
        'Main Image', upload_to=upload_image_to, blank=False)
    mainImageBW = ImageSpecField(
        source='mainImage',
        processors=[Adjust(color=0.0)],
        format='JPEG',
        options={
            'quality': 100
        })
    secondImage = models.ImageField(
        'Second Image', upload_to=upload_image_to, blank=True, null=True)
    imageOrientationOpts = (('left', 'left'), ('rigt', 'right'), ('up', 'up'),
                            ('down', 'down'), ('cntr', 'center'))
    imageOrientation = models.CharField(
        'Alignment',
        max_length=4,
        choices=imageOrientationOpts,
        default='cntr')
    home = models.BooleanField('Home', default=False)
    proyects = models.BooleanField('Proyects', default=False)
    slug = models.SlugField('Slug Name', max_length=100)
    order = models.PositiveSmallIntegerField(
        'Order', blank=False, null=False, default=1)

    def save(self, *args, **kwargs):
        self.slug = defaultfilters.slugify(self.name_es)
        super(Proyect, self).save(*args, **kwargs)

    def second_image(self):
        if self.secondImage:
            img = self.secondImage.url
        else:
            img = self.mainImageBW.url
        return '' % img

    second_image.allow_tags = True

    def admin_description(self):
        return '
%s' % self.description_es

    admin_description.allow_tags = True

    def admin_image(self):
        return u"" % self.mainImage.url

    admin_image.allow_tags = True

    def __str__(self):  # Python 3: def __str__(self):
        return self.name_en

    class Meta:
        verbose_name = 'Proyect'
        verbose_name_plural = 'Proyects'

    def getGallery(self):
        s = Image.objects.filter(proyect=self).order_by('order')
        return s

    def getVideos(self):
        s = Video.objects.filter(proyect=self).order_by('order')
        return s

class Image(models.Model):
    idImage = models.AutoField(primary_key=True)
    name = models.CharField('Name', max_length=100, blank=True)
    image = models.ImageField(
        'Image', upload_to=upload_image_to, blank=False, null=False)
    imageOrientationOpts = (('left', 'left'), ('rigt', 'right'),
                            ('cntr', 'center'), ('covr', 'cover'))
    imageOrientation = models.CharField(
        'Tipo', max_length=4, choices=imageOrientationOpts, default='covr')
    imageEffectsOptions = (('no', 'None'), ('bw', 'Black & white'),
                           ('one', 'Effect 1'), ('two', 'Effect 2'),
                           ('tre', 'Effect 3'))
    imageEffect = models.CharField(
        'Effect', max_length=4, choices=imageEffectsOptions, default='no')
    order = models.PositiveSmallIntegerField(
        'Order', blank=False, null=False, default=1)
    blackWhite = ImageSpecField(
        source='image',
        processors=[Adjust(color=0.0)],
        format='JPEG',
        options={
            'quality': 100
        })
    effectOne = ImageSpecField(
        source='image',
        processors=[Adjust(brightness=1.45, contrast=1.10, color=0.45)],
        format='JPEG',
        options={
            'quality': 100
        })
    effectTwo = ImageSpecField(
        source='image',
        processors=[Adjust(brightness=0.85, contrast=1.30, color=0.45)],
        format='JPEG',
        options={
            'quality': 100
        })
    effectThree = ImageSpecField(
        source='image',
        processors=[Adjust(brightness=0.85, contrast=1.25, color=1.20)],
        format='JPEG',
        options={
            'quality': 100
        })
    proyect = models.ForeignKey('Proyect', on_delete=models.CASCADE)

    def imageWithEffect(self):
        tmpImg = ""
        if self.imageEffect == 'no':
            tmpImg = self.image.url
        elif self.imageEffect == 'bw':
            tmpImg = self.blackWhite.url
        elif self.imageEffect == 'one':
            tmpImg = self.effectOne.url
        elif self.imageEffect == 'two':
            tmpImg = self.effectTwo.url
        elif self.imageEffect == 'tre':
            tmpImg = self.effectThree.url
        else:
            tmpImg = self.image.url
        return tmpImg

    def admin_image(self):
        if self.imageEffect == 'no':
            tmpImg = self.image.url
        elif self.imageEffect == 'bw':
            tmpImg = self.blackWhite.url
        elif self.imageEffect == 'one':
            tmpImg = self.effectOne.url
        elif self.imageEffect == 'two':
            tmpImg = self.effectTwo.url
        elif self.imageEffect == 'tre':
            tmpImg = self.effectThree.url
        else:
            tmpImg = self.image.url
        return '' % tmpImg

    admin_image.allow_tags = True

    class Meta:
        verbose_name = 'Imagen'
        verbose_name_plural = 'Imagenes'

    def __str__(self):
        return self.name

class Video(models.Model):
    video = EmbedVideoField()
    name = models.CharField('Name', max_length=100, blank=True)
    order = models.PositiveSmallIntegerField(
        'Order', blank=False, null=False, default=1)
    proyect = models.ForeignKey('Proyect', on_delete=models.CASCADE)

    class Meta:
        verbose_name = 'Video'
        verbose_name_plural = 'Videos'

    def __str__(self):
        return self.name

class SingleInformation(SingletonModel):
    aboutImage = models.ImageField(
        'Imagen about', upload_to=upload_image_to, null=False, blank=False)
    aboutText_es = RichTextField()
    aboutText_en = RichTextField()
    line1 = models.CharField('linea 1', max_length=140, null=True, blank=True)
    line2 = models.CharField('linea 2', max_length=140, null=True, blank=True)
    line3 = models.CharField('linea 3', max_length=140, null=True, blank=True)
    line4 = models.CharField('linea 4', max_length=140, null=True, blank=True)
    published = models.BooleanField('Published', default=False)

    def save(self, *args, **kwargs):
        if self.published:
            try:
                temp = SingleInformation.objects.get(published=True)
                if self != temp:
                    temp.published = False
                    temp.save()
            except SingleInformation.DoesNotExist:
                pass
        super(SingleInformation, self).save(*args, **kwargs)

    def admin_Aboutimage(self):
        return '' % self.aboutImage.url

    admin_Aboutimage.allow_tags = True

    class Meta:
        verbose_name = "About"
        verbose_name_plural = "About"

class Press(models.Model):
    name_es = models.CharField(
        'Name spanish', max_length=140, null=False, blank=False)
    name_en = models.CharField(
        'Name english', max_length=140, null=False, blank=False)
    image = models.ImageField(
        'Main image', upload_to=upload_image_to, null=False, blank=False)
    url = models.URLField('URL')
    description_es = RichTextField()
    description_en = RichTextField()
    pub_date = models.DateTimeField(
        'Created', editable=False, auto_now_add=True)

    def admin_description(self):
        return '
%s
' % self.description_es

    admin_description.allow_tags = True

    class Meta:
        verbose_name = "Noticia"
        verbose_name_plural = "Noticias"

class HomeSliderImages(models.Model):
    url = models.URLField('URL')
    image = models.ImageField('Imagen', upload_to=upload_image_to)

    class Meta:
        verbose_name = "Home Image"
        verbose_name_plural = "Home Images"

    def __str__(self):
        return self.url

class SocialLinks(SingletonModel):
    fb = models.URLField('Facebook')
    tw = models.URLField('Twitter')
    it = models.URLField('Instagram')


import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import sklearn
import joblib
import nibabel as nib
import itertools
from sklearn.linear_model import LogisticRegression
from IPython.display import clear_output
import sys
from subprocess import call
import pickle

def save_obj(obj, name):
    with open(name + '.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

def load_obj(name):
    with open(name + '.pkl', 'rb') as f:
        return pickle.load(f)

def normalize(X):
    X = X - X.mean(0)
    return X

def jitter(size,const=0):
    jit = np.random.normal(0+const, 0.05, size)
    X = np.zeros((size))
    X = X + jit
    return X

def other(target):
    other_objs = [i for i in ['bed', 'bench', 'chair', 'table'] if i not in target]
    return other_objs

def red_vox(n_vox, prop=0.1):
    return int(np.ceil(n_vox * prop))

def get_inds(X, Y, pair, testRun=None):

    inds = {}

    # return relative indices
    if testRun:
        trainIX = Y.index[(Y['label'].isin(pair)) & (Y['run_num'] != int(testRun))]
    else:
        trainIX = Y.index[(Y['label'].isin(pair))]

    # pull training and test data
    trainX = X[trainIX]
    trainY = Y.iloc[trainIX].label

    # Main classifier on 5 runs, testing on 6th
    clf = LogisticRegression(penalty='l2',C=1, solver='lbfgs', max_iter=1000,
                             multi_class='multinomial').fit(trainX, trainY)
    B = clf.coef_[0]  # pull betas

    # retrieve only the first object, then only the second object
    if testRun:
        obj1IX = Y.index[(Y['label'] == pair[0]) & (Y['run_num'] != int(testRun))]
        obj2IX = Y.index[(Y['label'] == pair[1]) & (Y['run_num'] != int(testRun))]
    else:
        obj1IX = Y.index[(Y['label'] == pair[0])]
        obj2IX = Y.index[(Y['label'] == pair[1])]
    # Get the average of the first object, then the second object
    obj1X = np.mean(X[obj1IX], 0)
    obj2X = np.mean(X[obj2IX], 0)

    # Build the importance map
    mult1X = obj1X * B
    mult2X = obj2X * B

    # Sort these so that they are from least to most important for a given category.
    sortmult1X = mult1X.argsort()[::-1]
    sortmult2X = mult2X.argsort()

    # add to a dictionary for later use
    inds[clf.classes_[0]] = sortmult1X
    inds[clf.classes_[1]] = sortmult2X

    return inds

def getEvidence(sub,testEvidence,METADICT=None,FEATDICT=None,filterType=None,roi="V1",include=1):
    # each testRun, each subject, each target axis, each target obj would generate one.
    META = METADICT[sub]
    print('META.shape=',META.shape)
    FEAT = FEATDICT[sub]
    # Using the trained model, get the evidence
    testRun=6
    objects=['bed', 'bench', 'chair', 'table']
    allpairs = itertools.combinations(objects,2)
    for pair in allpairs: #pair=('bed', 'bench')
        # Find the control (remaining) objects for this pair
        altpair = other(pair) #altpair=('chair', 'table')
        for obj in pair: #obj='bed'
            # in the current target axis pair=('bed', 'bench') altpair=('chair', 'table'), display image obj='bed'
            # find the evidence for bed from the (bed chair) and (bed table) classifier

            # get the test data and seperate the test data into category obj and category other
            objIX = META.index[(META['label'].isin([obj])) & (META['run_num'] == int(testRun))]
            otherObjIX = META.index[(META['label'].isin(other(obj))) & (META['run_num'] == int(testRun))]

            obj_X=FEAT[objIX]
            obj_Y=META.iloc[objIX].label
            otherObj_X=FEAT[otherObjIX]
            otherObj_Y=META.iloc[otherObjIX].label

            model_folder = f'/gpfs/milgram/project/turk-browne/jukebox/ntb/projects/sketchloop02/clf/{include}/{roi}/{filterType}/'
            print(f'loading {model_folder}{sub}_{pair[0]}{pair[1]}_{obj}{altpair[0]}.joblib')
            print(f'loading {model_folder}{sub}_{pair[0]}{pair[1]}_{obj}{altpair[1]}.joblib')
            clf1 = joblib.load(f'{model_folder}{sub}_{pair[0]}{pair[1]}_{obj}{altpair[0]}.joblib')
            clf2 = joblib.load(f'{model_folder}{sub}_{pair[0]}{pair[1]}_{obj}{altpair[1]}.joblib')

            if include < 1:
                selectedFeatures=load_obj(f'{model_folder}{sub}_{pair[0]}{pair[1]}_{obj}{altpair[0]}.selectedFeatures')
                obj_X=obj_X[:,selectedFeatures]
                otherObj_X=otherObj_X[:,selectedFeatures]

            s1 = clf1.score(obj_X, obj_Y)
            s2 = clf2.score(obj_X, obj_Y)
            obj_evidence = np.mean([s1, s2])
            print('obj_evidence=',obj_evidence)

            s1 = clf1.score(otherObj_X, otherObj_Y)
            s2 = clf2.score(otherObj_X, otherObj_Y)
            otherObj_evidence = np.mean([s1, s2])
            print('otherObj_evidence=',otherObj_evidence)

            testEvidence = testEvidence.append({
                'sub':sub,
                'testRun':testRun,
                'targetAxis':pair,
                'obj':obj,
                'obj_evidence':obj_evidence,
                'otherObj_evidence':otherObj_evidence,
                'filterType':filterType,
                'include':include,
                'roi':roi
                },
                ignore_index=True)

    return testEvidence

def minimalClass(filterType = 'noFilter',testRun = 6, roi="V1",include = 1): #include is the proportion of features selected

    accuracyContainer = pd.DataFrame(columns=['sub','testRun','targetAxis','obj','altobj','acc','filterType','roi'])
    testEvidence = pd.DataFrame(columns=['sub','testRun','targetAxis','obj','obj_evidence','otherObj_evidence','filterType','roi'])

    working_dir='/gpfs/milgram/project/turk-browne/projects/rtcloud_kp/FilterTesting/neurosketch_realtime_preprocess/'
    os.chdir(working_dir)

    data_dir=f'/gpfs/milgram/project/turk-browne/jukebox/ntb/projects/sketchloop02/features/{filterType}/recognition/'
    files = os.listdir(data_dir)
    feats = [i for i in files if 'metadata' not in i]
    subjects = np.unique([i.split('_')[0] for i in feats if i.split('_')[0] not in ['1121161','0112174']]) # 1121161 has a grid spacing issue and 0112174 lacks one of regressor file
    # If you want to reduce the number of subjects used for testing purposes
    subs=len(subjects)
    # subs=1
    subjects = subjects[:subs]
    print(subjects)

    objects = ['bed', 'bench', 'chair', 'table']
    phases = ['12', '34', '56']

    # THIS CELL READS IN ALL OF THE PARTICIPANTS' DATA and fills into dictionary
    FEATDICT = {}
    METADICT = {}
    subjects_new=[]
    for si, sub in enumerate(subjects[:]):
        print('{}/{}'.format(si+1, len(subjects)))
        diffs = []
        scores = []
        subcount = 0
        for phase in phases:
            _feat = np.load(data_dir+'/{}_{}_{}_featurematrix.npy'.format(sub, roi, phase))
            _feat = normalize(_feat)
            _meta = pd.read_csv(data_dir+'/metadata_{}_{}_{}.csv'.format(sub, roi, phase))
            if phase!='12':
                assert _feat.shape[1]==FEAT.shape[1], 'feat shape not matched'
            FEAT = _feat if phase == "12" else np.vstack((FEAT, _feat))
            META = _meta if phase == "12" else pd.concat((META, _meta))
        META = META.reset_index(drop=True)

        assert FEAT.shape[0] == META.shape[0]

        METADICT[sub] = META
        FEATDICT[sub] = FEAT
        clear_output(wait=True)
        subjects_new.append(sub)

    # Which run to use as test data (leave as None to not have test data)
    subjects=subjects_new

    # Decide on the proportion of crescent data to use for classification
    for si,sub in enumerate(subjects):
        print('{}/{}'.format(si+1, len(subjects)))
        print(sub)
        META = METADICT[sub]
        FEAT = FEATDICT[sub]

        allpairs = itertools.combinations(objects,2)

        # Iterate over all the possible target pairs of objects
        for pair in allpairs:
            # Find the control (remaining) objects for this pair
            altpair = other(pair)

            # pull sorted indices for each of the critical objects, in order of importance (low to high)
            inds = get_inds(FEAT, META, pair, testRun=testRun)

            # Find the number of voxels that will be left given your inclusion parameter above
            nvox = red_vox(FEAT.shape[1], include)

            for obj in pair:
                # foil = [i for i in pair if i != obj][0]
                for altobj in altpair:

                    # establish a naming convention where it is $TARGET_$CLASSIFICATION
                    # Target is the NF pair (e.g. bed/bench)
                    # Classificationis is btw one of the targets, and a control (e.g. bed/chair, or bed/table, NOT bed/bench)
                    naming = '{}{}_{}{}'.format(pair[0], pair[1], obj, altobj)

                    # Pull the relevant inds from your previously established dictionary
                    obj_inds = inds[obj]

                    # If you're using testdata, this function will split it up. Otherwise it leaves out run as a parameter
                    if testRun:
                        trainIX = META.index[(META['label'].isin([obj, altobj])) & (META['run_num'] != int(testRun))]
                        testIX = META.index[(META['label'].isin([obj, altobj])) & (META['run_num'] == int(testRun))]
                    else:
                        trainIX = META.index[(META['label'].isin([obj, altobj]))]
                        testIX = META.index[(META['label'].isin([obj, altobj]))]

                    # pull training and test data
                    trainX = FEAT[trainIX]
                    testX = FEAT[testIX]
                    trainY = META.iloc[trainIX].label
                    testY = META.iloc[testIX].label

                    # If you're selecting high-importance features, this bit handles that
                    if include < 1:
                        trainX = trainX[:, obj_inds[-nvox:]]
                        testX = testX[:, obj_inds[-nvox:]]

                    # Train your classifier
                    clf = LogisticRegression(penalty='l2',C=1, solver='lbfgs', max_iter=1000,
                                            multi_class='multinomial').fit(trainX, trainY)

                    joblib.dump(clf, model_folder + '{}_{}.joblib'.format(sub, naming))
                    save_obj(obj_inds[-nvox:],f'{model_folder}{sub}_{naming}.selectedFeatures')

                    # Monitor progress by printing accuracy (only useful if you're running a test set)
                    acc = clf.score(testX, testY)
                    if (si+1)%10==0:
                        print(naming, acc)
                    accuracyContainer = accuracyContainer.append({
                        'sub':sub,
                        'testRun':testRun,
                        'targetAxis':pair,
                        'obj':obj,
                        'altobj':altobj,
                        'acc':acc,
                        'filterType':filterType,
                        'include':include,
                        'roi':roi
                        },
                        ignore_index=True)

    for sub in subjects:
        testEvidence=getEvidence(sub,testEvidence,
        METADICT=METADICT,FEATDICT=FEATDICT,filterType=filterType,roi=roi,include=include)
    print(accuracyContainer)
    print(testEvidence)
    accuracyContainer.to_csv(f"{model_folder}accuracy.csv")
    testEvidence.to_csv(f'{model_folder}testEvidence.csv')

include=np.float(sys.argv[1])
roi=sys.argv[2]
filterType=sys.argv[3]
model_folder = f'/gpfs/milgram/project/turk-browne/jukebox/ntb/projects/sketchloop02/clf/{include}/{roi}/{filterType}/'
print('model_folder=',model_folder)
call(f"mkdir -p {model_folder}",shell=True)
minimalClass(include = include, roi=roi, filterType = filterType, testRun = 6)

# ## - to run in parallel
# # bash to submit jobs, in the folder of testMiniclass

# # testMiniclass_child.sh
# #!/bin/bash
# #SBATCH --partition=short,scavenge,interactive,long,verylong
# #SBATCH --nodes=1
# #SBATCH --ntasks-per-node=1
# #SBATCH --cpus-per-task=1
# #SBATCH --time=3:00:00
# #SBATCH --job-name=miniClass
# #SBATCH --output=logs/miniClass-%j.out
# #SBATCH --mem-per-cpu=10G
# #SBATCH --mail-type=FAIL
# mkdir -p logs/
# module load miniconda
# source activate /gpfs/milgram/project/turk-browne/kp578/conda_envs/rtAtten
# include=$1
# roi=$2
# filterType=$3
# /usr/bin/time python -u /gpfs/milgram/project/turk-browne/projects/rtSynth_rt/FilterTesting/testMiniclass/testMiniclass.py $include $roi $filterType

# # testMiniclass_parent.py
# from glob import glob
# import os
# from subprocess import call
# for include in [0.1,0.3,0.6,0.9,1]:
#     for roi in ['V1', 'fusiform', 'IT', 'LOC', 'occitemp', 'parahippo']:
#         for filterType in ['noFilter','highPassRealTime','highPassBetweenRuns']:
#             command=f'sbatch testMiniclass_child.sh {include} {roi} {filterType}'
#             print(command)
#             # call(command, shell=True)

# # load and plot data
# testEvidence_noFilter=pd.read_csv(f'/gpfs/milgram/project/turk-browne/jukebox/ntb/projects/sketchloop02/clf/testEvidence_noFilter.csv')
# testEvidence_highPassRealTime=pd.read_csv(f'/gpfs/milgram/project/turk-browne/jukebox/ntb/projects/sketchloop02/clf/testEvidence_highPassRealTime.csv')
# testEvidence_highPassBetweenRuns=pd.read_csv(f'/gpfs/milgram/project/turk-browne/jukebox/ntb/projects/sketchloop02/clf/testEvidence_highPassBetweenRuns.csv')
# def bar(L,labels=None,title=None):
#     import matplotlib.pyplot as plt
#     CTEs = [np.mean(i) for i in L]
#     error = [np.std(i) for i in L]
#     x_pos = np.arange(len(labels))
#     fig, ax = plt.subplots(figsize=(10,10))
#     ax.bar(x_pos, CTEs, yerr=error, align='center', alpha=0.5, ecolor='black', capsize=10)
#     ax.set_ylabel('object evidence')
#     ax.set_xticks(x_pos)
#     ax.set_xticklabels(labels)
#     ax.set_title(title)
#     ax.yaxis.grid(True)
#     plt.tight_layout()
#     plt.xticks(rotation=30)
#     plt.show()
# bar([np.asarray(testEvidence_noFilter['obj_evidence']),
#     np.asarray(testEvidence_noFilter['otherObj_evidence']),
#     [],
#     np.asarray(testEvidence_highPassRealTime['obj_evidence']),
#     np.asarray(testEvidence_highPassRealTime['otherObj_evidence']),
#     [],
#     np.asarray(testEvidence_highPassBetweenRuns['obj_evidence']),
#     np.asarray(testEvidence_highPassBetweenRuns['otherObj_evidence'])],
#     labels=[
#     'noFilter_obj',
#     'noFilter_otherObj',
#     '',
#     'highPassRealTime_obj',
#     'highPassRealTime_otherObj',
#     '',
#     'highpassOffline_obj',
#     'highpassOffline_otherObj'],title=None)


import json

from django.http import HttpResponse, HttpResponseRedirect

# Create your views here.
from django.views.generic.base import View

from apps.operation.forms import UserAskForm
from apps.operation.models import UserFavorite
from libs.fav_help import reduce_fav_num, add_fav_num
from libs.login import LoginRequiredMixinView

class SearchView(View):
    def get(self, request):
        """
        :param type: 是 js 传过来的参数，名字在 js 里面定义的
        :param keywords: 是 js 传过来的参数，名字在 js 里面定义的
        """
        type = request.GET.get('type', '')
        keywords = request.GET.get('keywords', '')
        url = ''
        if type == 'course':
            url = '/course/'
        if type == 'teacher':
            url = '/teacher/'
        if type == 'org':
            url = '/org/'
        return HttpResponseRedirect("{}?type={}&keywords={}".format(url, type, keywords))

class AddUserAskView(View):
    """
    添加用户咨询
    """

    def post(self, request):
        user_ask_form = UserAskForm(request.POST)

        if user_ask_form.is_valid():
            user_ask = user_ask_form.save(commit=True)
            status = 'success'
            msg = '提交成功'
        else:
            status = 'fail'
            msg = '提交失败'

        r = dict(status=status, msg=msg)
        r = json.dumps(r)
        return HttpResponse(r, content_type='application/json')

class AddFavView(LoginRequiredMixinView):
    """
    用户收藏、取消收藏
    """

    def post(self, request):
        fav_id = int(request.POST.get('fav_id', 0))
        fav_type = int(request.POST.get('fav_type', 0))

        if not request.user.is_authenticated:  # 判断用户登录状态
            status = 'fail'
            msg = '用户未登录，收藏失败'
        else:
            exist_record = UserFavorite.objects.filter(user=request.user, fav_id=fav_id, fav_type=fav_type)
            if exist_record:  # 如果记录已存在表示用户取消收藏
                exist_record.delete()
                reduce_fav_num(fav_type, fav_id)
                status = 'success'
                msg = '收藏'
            elif fav_id > 0 and fav_type > 0:
                user_fav = UserFavorite(user_id=request.user.id, fav_type=fav_type, fav_id=fav_id)
                user_fav.save()
                add_fav_num(fav_type, fav_id)
                status = 'success'
                msg = '已收藏'
            else:
                status = 'fail'
                msg = '收藏出错'

        r = dict(status=status, msg=msg)
        r = json.dumps(r)
        return HttpResponse(r, content_type='application/json')