import json
from pathlib import Path
from typing import Optional, List, Union, Tuple, Dict

import requests
import torch
import torch.nn as nn
import random
import tempfile
import os

from huggingface_hub import hf_hub_download
from torch.nn import CrossEntropyLoss

from transformers.utils import (
add_start_docstrings_to_model_forward,
add_end_docstrings,
replace_return_docstrings,
)

from transformers import AutoModelForSeq2SeqLM
from transformers.models.bart.modeling_bart import (
    BartForConditionalGeneration,
    _expand_mask, logger,
    shift_tokens_right,
    BART_INPUTS_DOCSTRING,
    _CONFIG_FOR_DOC,
    BART_GENERATION_EXAMPLE,
    BartModel,
    BartDecoder

)
from .adapter import Adapter
from transformers.modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions,
    Seq2SeqModelOutput,
    BaseModelOutput,
    Seq2SeqLMOutput
)

from huggingface_hub.utils import validate_hf_hub_args,HfFolder
from huggingface_hub.hf_api import HfApi
from huggingface_hub.repository import Repository
from huggingface_hub.constants import CONFIG_NAME, PYTORCH_WEIGHTS_NAME

class KeyBartAdapter(BartForConditionalGeneration):
    def __init__(self,adapter_hid_dim:int) -> None:
        keyBart = AutoModelForSeq2SeqLM.from_pretrained("bloomberg/KeyBART")
        self.__fix_weights__(keyBart)

        super().__init__(keyBart.model.config)

        self.lm_head = keyBart.lm_head
        self.model = BartPlus(keyBart, adapter_hid_dim)
        self.register_buffer("final_logits_bias", torch.zeros((1, self.model.shared.num_embeddings)))
        self.config = keyBart.config

    def __fix_weights__(self,keyBart:BartForConditionalGeneration):
        for i in keyBart.model.parameters():
            i.requires_grad = False
        for i in keyBart.lm_head.parameters():
            i.requires_grad = False

    def save_pretrained(
        self,
        save_directory: Union[str, Path],
        config: Optional[dict] = None,
        push_to_hub: bool = False,
        **kwargs,
    ):
        """
        Save weights in local directory.

        Parameters:
            save_directory (`str` or `Path`):
                Specify directory in which you want to save weights.
            config (`dict`, *optional*):
                Specify config (must be dict) in case you want to save
                it.
            push_to_hub (`bool`, *optional*, defaults to `False`):
                Whether or not to push your model to the Hugging Face Hub after
                saving it. You can specify the repository you want to push to with
                `repo_id` (will default to the name of `save_directory` in your
                namespace).
            kwargs:
                Additional key word arguments passed along to the
                [`~utils.PushToHubMixin.push_to_hub`] method.
        """
        os.makedirs(save_directory, exist_ok=True)

        # saving model weights/files
        self._save_pretrained(save_directory)

        # saving config
        if isinstance(config, dict):
            path = os.path.join(save_directory, CONFIG_NAME)
            with open(path, "w") as f:
                json.dump(config, f)

        if push_to_hub:
            kwargs = kwargs.copy()  # soft-copy to avoid mutating input
            if config is not None:  # kwarg for `push_to_hub`
                kwargs["config"] = config

            if (
                # If a deprecated argument is passed, we have to use the deprecated
                # version of `push_to_hub`.
                # TODO: remove this possibility in v0.12
                kwargs.get("repo_url") is not None
                or kwargs.get("repo_path_or_name") is not None
                or kwargs.get("organization") is not None
                or kwargs.get("git_user") is not None
                or kwargs.get("git_email") is not None
                or kwargs.get("skip_lfs_files") is not None
            ):
                if kwargs.get("repo_path_or_name") is None:
                    # Repo name defaults to `save_directory` name
                    kwargs["repo_path_or_name"] = save_directory
            elif kwargs.get("repo_id") is None:
                # Repo name defaults to `save_directory` name
                kwargs["repo_id"] = Path(save_directory).name

            return self.push_to_hub(**kwargs)

    @classmethod
    @validate_hf_hub_args
    def from_pretrained(
            cls,
            pretrained_model_name_or_path: str,
            force_download: bool = False,
            resume_download: bool = False,
            proxies: Optional[Dict] = None,
            token: Optional[Union[str, bool]] = None,
            cache_dir: Optional[str] = None,
            local_files_only: bool = False,
            revision: Optional[str] = None,
            **model_kwargs,
    ):
        r"""
        Download and instantiate a model from the Hugging Face Hub.

                Parameters:
                    pretrained_model_name_or_path (`str` or `os.PathLike`):
                        Can be either:
                            - A string, the `model id` of a pretrained model
                              hosted inside a model repo on huggingface.co.
                              Valid model ids can be located at the root-level,
                              like `bert-base-uncased`, or namespaced under a
                              user or organization name, like
                              `dbmdz/bert-base-german-cased`.
                            - You can add `revision` by appending `@` at the end
                              of model_id simply like this:
                              `dbmdz/bert-base-german-cased@main` Revision is
                              the specific model version to use. It can be a
                              branch name, a tag name, or a commit id, since we
                              use a git-based system for storing models and
                              other artifacts on huggingface.co, so `revision`
                              can be any identifier allowed by git.
                            - A path to a `directory` containing model weights
                              saved using
                              [`~transformers.PreTrainedModel.save_pretrained`],
                              e.g., `./my_model_directory/`.
                            - `None` if you are both providing the configuration
                              and state dictionary (resp. with keyword arguments
                              `config` and `state_dict`).
                    force_download (`bool`, *optional*, defaults to `False`):
                        Whether to force the (re-)download of the model weights
                        and configuration files, overriding the cached versions
                        if they exist.
                    resume_download (`bool`, *optional*, defaults to `False`):
                        Whether to delete incompletely received files. Will
                        attempt to resume the download if such a file exists.
                    proxies (`Dict[str, str]`, *optional*):
                        A dictionary of proxy servers to use by protocol or
                        endpoint, e.g., `{'http': 'foo.bar:3128',
                        'http://hostname': 'foo.bar:4012'}`. The proxies are
                        used on each request.
                    token (`str` or `bool`, *optional*):
                        The token to use as HTTP bearer authorization for remote
                        files. If `True`, will use the token generated when
                        running `transformers-cli login` (stored in
                        `~/.huggingface`).
                    cache_dir (`Union[str, os.PathLike]`, *optional*):
                        Path to a directory in which a downloaded pretrained
                        model configuration should be cached if the standard
                        cache should not be used.
                    local_files_only(`bool`, *optional*, defaults to `False`):
                        Whether to only look at local files (i.e., do not try to
                        download the model).
                    model_kwargs (`Dict`, *optional*):
                        model_kwargs will be passed to the model during
                        initialization

                <Tip>

                Passing `token=True` is required when you want to use a
                private model.

                </Tip>
        """

        model_id = pretrained_model_name_or_path

        config_file: Optional[str] = None
        if os.path.isdir(model_id):
            if CONFIG_NAME in os.listdir(model_id):
                config_file = os.path.join(model_id, CONFIG_NAME)
            else:
                logger.warning(f"{CONFIG_NAME} not found in {Path(model_id).resolve()}")
        else:
            try:
                config_file = hf_hub_download(
                    repo_id=model_id,
                    filename=CONFIG_NAME,
                    revision=revision,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
                    token=token,
                    local_files_only=local_files_only,
                )
            except requests.exceptions.RequestException:
                logger.warning(f"{CONFIG_NAME} not found in HuggingFace Hub")

        if config_file is not None:
            with open(config_file, "r", encoding="utf-8") as f:
                config = json.load(f)
            model_kwargs.update({"config": config})

        return cls._from_pretrained(
            model_id,
            revision,
            cache_dir,
            force_download,
            proxies,
            resume_download,
            local_files_only,
            token,
            **model_kwargs,
        )

    @validate_hf_hub_args
    def push_to_hub(
            self,
            # NOTE: deprecated signature that will change in 0.12
            *,
            repo_path_or_name: Optional[str] = None,
            repo_url: Optional[str] = None,
            commit_message: str = "Add model",
            organization: Optional[str] = None,
            private: bool = False,
            api_endpoint: Optional[str] = None,
            token: Optional[str] = None,
            git_user: Optional[str] = None,
            git_email: Optional[str] = None,
            config: Optional[dict] = None,
            skip_lfs_files: bool = False,
            # NOTE: New arguments since 0.9
            repo_id: Optional[str] = None,  # optional only until 0.12
            branch: Optional[str] = None,
            create_pr: Optional[bool] = None,
            allow_patterns: Optional[Union[List[str], str]] = None,
            ignore_patterns: Optional[Union[List[str], str]] = None,
            # TODO (release 0.12): signature must be the following
            # repo_id: str,
            # *,
            # commit_message: str = "Add model",
            # private: bool = False,
            # api_endpoint: Optional[str] = None,
            # token: Optional[str] = None,
            # branch: Optional[str] = None,
            # create_pr: Optional[bool] = None,
            # config: Optional[dict] = None,
            # allow_patterns: Optional[Union[List[str], str]] = None,
            # ignore_patterns: Optional[Union[List[str], str]] = None,
    ) -> str:
        """
        Upload model checkpoint to the Hub.

        Use `allow_patterns` and `ignore_patterns` to precisely filter which files
        should be pushed to the hub. See [`upload_folder`] reference for more details.

        Parameters:
            repo_id (`str`, *optional*):
                Repository name to which push.
            commit_message (`str`, *optional*):
                Message to commit while pushing.
            private (`bool`, *optional*, defaults to `False`):
                Whether the repository created should be private.
            api_endpoint (`str`, *optional*):
                The API endpoint to use when pushing the model to the hub.
            token (`str`, *optional*):
                The token to use as HTTP bearer authorization for remote files.
                If not set, will use the token set when logging in with
                `transformers-cli login` (stored in `~/.huggingface`).
            branch (`str`, *optional*):
                The git branch on which to push the model. This defaults to
                the default branch as specified in your repository, which
                defaults to `"main"`.
            create_pr (`boolean`, *optional*):
                Whether or not to create a Pull Request from `branch` with that commit.
                Defaults to `False`.
            config (`dict`, *optional*):
                Configuration object to be saved alongside the model weights.
            allow_patterns (`List[str]` or `str`, *optional*):
                If provided, only files matching at least one pattern are pushed.
            ignore_patterns (`List[str]` or `str`, *optional*):
                If provided, files matching any of the patterns are not pushed.

        Returns:
            The url of the commit of your model in the given repository.
        """
        # If the repo id is set, it means we use the new version using HTTP endpoint
        # (introduced in v0.9).
        if repo_id is not None:
            api = HfApi(endpoint=api_endpoint)
            api.create_repo(
                repo_id=repo_id,
                repo_type="model",
                token=token,
                private=private,
                exist_ok=True,
            )

            # Push the files to the repo in a single commit
            with tempfile.TemporaryDirectory() as tmp:
                saved_path = Path(tmp) / repo_id
                self.save_pretrained(saved_path, config=config)
                return api.upload_folder(
                    repo_id=repo_id,
                    repo_type="model",
                    token=token,
                    folder_path=saved_path,
                    commit_message=commit_message,
                    revision=branch,
                    create_pr=create_pr,
                    allow_patterns=allow_patterns,
                    ignore_patterns=ignore_patterns,
                )

        # If the repo id is None, it means we use the deprecated version using Git
        # TODO: remove code between here and `return repo.git_push()` in release 0.12
        if repo_path_or_name is None and repo_url is None:
            raise ValueError(
                "You need to specify a `repo_path_or_name` or a `repo_url`."
            )

        if token is None and repo_url is None:
            token = HfFolder.get_token()
            if token is None:
                raise ValueError(
                    "You must login to the Hugging Face hub on this computer by typing"
                    " `huggingface-cli login` and entering your credentials to use"
                    " `token=True`. Alternatively, you can pass your own token"
                    " as the `token` argument."
                )
        elif isinstance(token, str):
            token = token
        else:
            token = None

        if repo_path_or_name is None:
            assert repo_url is not None, "A `None` repo URL would have raised above"
            repo_path_or_name = repo_url.split("/")[-1]

        # If no URL is passed and there's no path to a directory containing files, create a repo
        if repo_url is None and not os.path.exists(repo_path_or_name):
            repo_id = Path(repo_path_or_name).name
            if organization:
                repo_id = f"{organization}/{repo_id}"
            repo_url = HfApi(endpoint=api_endpoint).create_repo(
                repo_id=repo_id,
                token=token,
                private=private,
                repo_type=None,
                exist_ok=True,
            )

        repo = Repository(
            repo_path_or_name,
            clone_from=repo_url,
            token=token,
            git_user=git_user,
            git_email=git_email,
            skip_lfs_files=skip_lfs_files,
        )
        repo.git_pull(rebase=True)

        # Save the files in the cloned repo
        self.save_pretrained(repo_path_or_name, config=config)

        # Commit and push!
        repo.git_add(auto_lfs_track=True)
        repo.git_commit(commit_message)
        return repo.git_push()

    def _save_pretrained(self, save_directory):
        """
        Overwrite this method if you wish to save specific layers instead of the
        complete model.
        """
        path = os.path.join(save_directory, PYTORCH_WEIGHTS_NAME)
        model_to_save = self.module if hasattr(self, "module") else self
        torch.save(model_to_save.state_dict(), path)

    @classmethod
    def _from_pretrained(
            cls,
            model_id,
            revision,
            cache_dir,
            force_download,
            proxies,
            resume_download,
            local_files_only,
            token,
            map_location="cpu",
            strict=False,
            **model_kwargs,
    ):
        """
        Overwrite this method to initialize your model in a different way.
        """
        map_location = torch.device(map_location)

        if os.path.isdir(model_id):
            print("Loading weights from local directory")
            model_file = os.path.join(model_id, PYTORCH_WEIGHTS_NAME)
        else:
            model_file = hf_hub_download(
                repo_id=model_id,
                filename=PYTORCH_WEIGHTS_NAME,
                revision=revision,
                cache_dir=cache_dir,
                force_download=force_download,
                proxies=proxies,
                resume_download=resume_download,
                token=token,
                local_files_only=local_files_only,
            )
        model = cls(**model_kwargs)

        state_dict = torch.load(model_file, map_location=map_location)
        model.load_state_dict(state_dict, strict=strict)
        model.eval()

        return model

    @add_start_docstrings_to_model_forward(BART_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
    @add_end_docstrings(BART_GENERATION_EXAMPLE)
    def forward(
            self,
            input_ids: torch.LongTensor = None,
            attention_mask: Optional[torch.Tensor] = None,
            decoder_input_ids: Optional[torch.LongTensor] = None,
            decoder_attention_mask: Optional[torch.LongTensor] = None,
            head_mask: Optional[torch.Tensor] = None,
            decoder_head_mask: Optional[torch.Tensor] = None,
            cross_attn_head_mask: Optional[torch.Tensor] = None,
            encoder_outputs: Optional[List[torch.FloatTensor]] = None,
            past_key_values: Optional[List[torch.FloatTensor]] = None,
            inputs_embeds: Optional[torch.FloatTensor] = None,
            decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
            labels: Optional[torch.LongTensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    ) -> Union[Tuple, Seq2SeqLMOutput]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
        Returns:
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if labels is not None:
            if use_cache:
                logger.warning("The `use_cache` argument is changed to `False` since `labels` is provided.")
            use_cache = False
            if decoder_input_ids is None and decoder_inputs_embeds is None:
                decoder_input_ids = shift_tokens_right(
                    labels, self.config.pad_token_id, self.config.decoder_start_token_id
                )

        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            encoder_outputs=encoder_outputs,
            decoder_attention_mask=decoder_attention_mask,
            head_mask=head_mask,
            decoder_head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            decoder_inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        lm_logits = self.lm_head(outputs[0]) + self.final_logits_bias

        masked_lm_loss = None
        if labels is not None:
            loss_fct = CrossEntropyLoss()
            masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1))

        if not return_dict:
            output = (lm_logits,) + outputs[1:]
            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

        return Seq2SeqLMOutput(
            loss=masked_lm_loss,
            logits=lm_logits,
            past_key_values=outputs.past_key_values,
            decoder_hidden_states=outputs.decoder_hidden_states,
            decoder_attentions=outputs.decoder_attentions,
            cross_attentions=outputs.cross_attentions,
            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
            encoder_hidden_states=outputs.encoder_hidden_states,
            encoder_attentions=outputs.encoder_attentions,
        )

class BartDecoderPlus(BartDecoder):
    def __init__(self,keyBart:BartForConditionalGeneration,adapter_hid_dim: int) -> None:
        super().__init__(keyBart.get_decoder().config)
        self.decoder = keyBart.model.decoder
        self.adapters = nn.ModuleList([Adapter(self.decoder.config.d_model,adapter_hid_dim) for _ in range(len(self.decoder.layers))])
        self.config = self.decoder.config
        self.dropout = self.decoder.dropout
        self.layerdrop = self.decoder.layerdrop
        self.padding_idx = self.decoder.padding_idx
        self.max_target_positions = self.decoder.max_target_positions
        self.embed_scale = self.decoder.embed_scale
        self.embed_tokens = self.decoder.embed_tokens
        self.embed_positions = self.decoder.embed_positions
        self.layers = self.decoder.layers
        self.layernorm_embedding = self.decoder.layernorm_embedding
        self.gradient_checkpointing = self.decoder.gradient_checkpointing

    def forward(
            self,
            input_ids: torch.LongTensor = None,
            attention_mask: Optional[torch.Tensor] = None,
            encoder_hidden_states: Optional[torch.FloatTensor] = None,
            encoder_attention_mask: Optional[torch.LongTensor] = None,
            head_mask: Optional[torch.Tensor] = None,
            cross_attn_head_mask: Optional[torch.Tensor] = None,
            past_key_values: Optional[List[torch.FloatTensor]] = None,
            inputs_embeds: Optional[torch.FloatTensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # retrieve input_ids and inputs_embeds
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
        elif input_ids is not None:
            input = input_ids
            input_shape = input.shape
            input_ids = input_ids.view(-1, input_shape[-1])
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            input = inputs_embeds[:, :, -1]
        else:
            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")

        # past_key_values_length
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if inputs_embeds is None:
            inputs_embeds = self.decoder.embed_tokens(input) * self.decoder.embed_scale

        attention_mask = self.decoder._prepare_decoder_attention_mask(
            attention_mask, input_shape, inputs_embeds, past_key_values_length
        )

        # expand encoder attention mask
        if encoder_hidden_states is not None and encoder_attention_mask is not None:
            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
            encoder_attention_mask = _expand_mask(encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])

        # embed positions
        positions = self.decoder.embed_positions(input, past_key_values_length)

        hidden_states = inputs_embeds + positions
        hidden_states = self.decoder.layernorm_embedding(hidden_states)

        hidden_states = nn.functional.dropout(hidden_states, p=self.decoder.dropout, training=self.decoder.training)

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
        next_decoder_cache = () if use_cache else None

        # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired
        for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
            if attn_mask is not None:
                if attn_mask.size()[0] != (len(self.decoder.layers)):
                    raise ValueError(
                        f"The `{mask_name}` should be specified for {len(self.decoder.layers)} layers, but it is for"
                        f" {head_mask.size()[0]}."
                    )

        for idx, decoder_layer in enumerate(self.decoder.layers):
            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
            dropout_probability = random.uniform(0, 1)
            if self.decoder.training and (dropout_probability < self.decoder.layerdrop):
                continue

            past_key_value = past_key_values[idx] if past_key_values is not None else None

            if self.decoder.gradient_checkpointing and self.decoder.training:

                if use_cache:
                    logger.warning(
                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                    )
                    use_cache = False

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        # None for past_key_value
                        return module(*inputs, output_attentions, use_cache)

                    return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(decoder_layer),
                    hidden_states,
                    attention_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    head_mask[idx] if head_mask is not None else None,
                    cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
                    None,
                )
            else:

                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    encoder_hidden_states=encoder_hidden_states,
                    encoder_attention_mask=encoder_attention_mask,
                    layer_head_mask=(head_mask[idx] if head_mask is not None else None),
                    cross_attn_layer_head_mask=(
                        cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None
                    ),
                    past_key_value=past_key_value,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                )
            hidden_states = layer_outputs[0]

            ######################### new #################################
            hidden_states = self.adapters[idx](hidden_states)
            ######################### new #################################

            if use_cache:
                next_decoder_cache += (layer_outputs[3 if output_attentions else 1],)

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

                if encoder_hidden_states is not None:
                    all_cross_attentions += (layer_outputs[2],)

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = next_decoder_cache if use_cache else None
        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
            cross_attentions=all_cross_attentions,
        )

class BartPlus(BartModel):
    def __init__(self,keyBart: BartForConditionalGeneration, adapter_hid_dim: int ) -> None:
        super().__init__(keyBart.model.config)
        self.config = keyBart.model.config

        self.shared = keyBart.model.shared
        self.encoder = keyBart.model.encoder
        self.decoder = BartDecoderPlus(keyBart,adapter_hid_dim=adapter_hid_dim)

    def forward(
            self,
            input_ids: torch.LongTensor = None,
            attention_mask: Optional[torch.Tensor] = None,
            decoder_input_ids: Optional[torch.LongTensor] = None,
            decoder_attention_mask: Optional[torch.LongTensor] = None,
            head_mask: Optional[torch.Tensor] = None,
            decoder_head_mask: Optional[torch.Tensor] = None,
            cross_attn_head_mask: Optional[torch.Tensor] = None,
            encoder_outputs: Optional[List[torch.FloatTensor]] = None,
            past_key_values: Optional[List[torch.FloatTensor]] = None,
            inputs_embeds: Optional[torch.FloatTensor] = None,
            decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    ) -> Union[Tuple, Seq2SeqModelOutput]:

        # different to other models, Bart automatically creates decoder_input_ids from
        # input_ids if no decoder_input_ids are provided
        if decoder_input_ids is None and decoder_inputs_embeds is None:
            if input_ids is None:
                raise ValueError(
                    "If no `decoder_input_ids` or `decoder_inputs_embeds` are "
                    "passed, `input_ids` cannot be `None`. Please pass either "
                    "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`."
                )

            decoder_input_ids = shift_tokens_right(
                input_ids, self.config.pad_token_id, self.config.decoder_start_token_id
            )

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if encoder_outputs is None:
            encoder_outputs = self.encoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
                head_mask=head_mask,
                inputs_embeds=inputs_embeds,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
        # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
        elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
            encoder_outputs = BaseModelOutput(
                last_hidden_state=encoder_outputs[0],
                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
            )

        # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=encoder_outputs[0],
            encoder_attention_mask=attention_mask,
            head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        if not return_dict:
            return decoder_outputs + encoder_outputs

        return Seq2SeqModelOutput(
            last_hidden_state=decoder_outputs.last_hidden_state,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
        )


import argparse
import imp
import os
import sys
import logging
from xml.etree.ElementInclude import default_loader
logging.basicConfig(level=logging.INFO)
import shutil
import random
import numpy as np
import pandas as pd
from datetime import datetime
from collections import defaultdict

import multiprocessing
from functools import partial

from openfold.utils import data_utils

def cdrsfasta_from_truncated_pdbs(
    data_dir,
    output_dir,
    save_single=True,
    max_workers=16,
):
    job_args = [x for x in os.listdir(data_dir) if x.endswith(".pdb")]
    job_args = zip(job_args, list(range(len(job_args))))

    with multiprocessing.Pool(max_workers) as p:
        func = partial(data_utils.pdb2cdrsfasta, data_dir=data_dir)
        if save_single:
            fastas = []
            for (ret, basename) in p.starmap(func, job_args):
                fastas.extend(ret)
            trunc_fasta_path = os.path.join(output_dir, "cdrs.fasta")
            with open(trunc_fasta_path, "w") as f:
                f.write('\n'.join(fastas))
        else:
            for (ret, basename) in p.starmap(func, job_args):
                trunc_fasta_path = os.path.join(output_dir, f"{basename}_cdrs.fasta")
                with open(trunc_fasta_path, "w") as f:
                    f.write('\n'.join(ret))

def copy_pdb_to_new_dir(
    pdb_id,
    old_pdb_dir,
    new_pdb_dir,
    old_fasta_dir=None,
    new_fasta_dir=None,
):
    """copy pdb and fasta into a new directory"""
    old_pdb_path = os.path.join(old_pdb_dir, pdb_id + '.pdb')
    new_pdb_path = os.path.join(new_pdb_dir, pdb_id + '.pdb')
    shutil.copy2(old_pdb_path, new_pdb_path)

    if old_fasta_dir is not None and new_fasta_dir is not None:
        old_fasta_path = os.path.join(old_fasta_dir, pdb_id + '.fasta')
        new_fasta_path = os.path.join(new_fasta_dir, pdb_id + '.fasta')
        shutil.copy2(old_fasta_path, new_fasta_path)

def get_cluster_rep(cluster_res):
    cluster_dict = defaultdict(list)
    with open(cluster_res, 'r') as f:
        for line in f:
            k, v = line.strip().split('\t')
            k = k.rsplit('_', 1)[0]
            v = v.rsplit('_', 1)[0]
            cluster_dict[k].append(v)

    cluster_rep = sorted(list(cluster_dict.keys()))
    logging.info(
        f"get {len(cluster_rep)} cluster representatives\n"
    )
    return cluster_rep

def main_get_cdrs(args):

    os.makedirs(args.output_dir, exist_ok=True)

    # 1.
    cdrsfasta_from_truncated_pdbs(args.input_dir, args.output_dir)

def main_copy_pdbs(args):

    os.makedirs(args.output_dir, exist_ok=True)

    # 2.
    cluster_rep = get_cluster_rep(args.cluster)
    logging.info("copying pdbs...")
    for pdb in cluster_rep:
        copy_pdb_to_new_dir(
            pdb, args.input_dir, args.output_dir
        )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_dir", type=str, default=None,
        help="Path to a directory containing truncated pdbs"
    )
    parser.add_argument(
        "--output_dir", type=str, default=None,
        help="Path to a directory containing truncated fastas"
    )
    parser.add_argument(
        "--cluster", type=str, default=None,
        help="Path to a directory containing truncated fastas"
    )
    args = parser.parse_args()

    if args.cluster is None:
        main_get_cdrs(args)
    else:
        main_copy_pdbs(args)


import abc
import json
import re

# Base classes

class BaseIdl:
    """Abstract base class, used for JSON serialization."""
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def json_serializable(self):
        """Returns a JSON serializable form of the object.

        This should be a dictionary, with keys scoped names of the form
        Class::key, where the scope is the class name.
        This is so we produce identical output to the Perl code, which uses
        the Perl module JSON.pm, which uses this format.
        """
        pass

class TypedObject:
    """Object with a type, such as an Attribute or Operation (return value).

    The type can be an actual type, or can be a typedef, which must be resolved
    before passing data to the code generator.
    """
    __metaclass__ = abc.ABCMeta
    data_type = None
    extended_attributes = None

    def resolve_typedefs(self, typedefs):
        """Resolve typedefs to actual types in the object."""
        additional_extended_attributes = {}
        # Convert string representation to and from an IdlType object
        # to handle parsing
        data_type_object = IdlType.from_string(self.data_type)
        base_type = data_type_object.base_type
        if base_type in typedefs:
            replacement_type = typedefs[base_type]
            data_type_object.base_type = replacement_type.data_type
            additional_extended_attributes = replacement_type.extended_attributes
        self.data_type = str(data_type_object)
        self.extended_attributes.update(additional_extended_attributes)

# IDL classes

class IdlDefinitions(BaseIdl):
    def __init__(self, callback_functions=None, enumerations=None, exceptions=None, file_name=None, interfaces=None, typedefs=None):
        self.callback_functions = callback_functions or {}
        self.enumerations = enumerations or {}
        self.exceptions = exceptions or {}
        self.file_name = file_name or None
        self.interfaces = interfaces or {}
        # Typedefs are not exposed by bindings; resolve typedefs with the
        # actual types and then discard the Typedefs.
        # http://www.w3.org/TR/WebIDL/#idl-typedefs
        if typedefs:
            self.resolve_typedefs(typedefs)

    def resolve_typedefs(self, typedefs):
        for callback_function in self.callback_functions.itervalues():
            callback_function.resolve_typedefs(typedefs)
        for exception in self.exceptions.itervalues():
            exception.resolve_typedefs(typedefs)
        for interface in self.interfaces.itervalues():
            interface.resolve_typedefs(typedefs)

    def json_serializable(self):
        return {
                'idlDocument::callbackFunctions': self.callback_functions.values(),
                'idlDocument::enumerations': self.enumerations.values(),
                'idlDocument::fileName': self.file_name,
                # Perl treats exceptions as a kind of interface
                'idlDocument::interfaces': sorted(self.exceptions.values() + self.interfaces.values()),
                }

    def to_json(self, debug=False):
        """Returns a JSON string representing the Definitions.

        The JSON output should be identical with the output of the Perl parser,
        specifically the function serializeJSON in deprecated_idl_serializer.pm,
        which takes a Perl object created by deprecated_idl_parser.pm.
        """
        # Sort so order consistent, allowing comparison of output
        if debug:
            # indent turns on pretty-printing for legibility
            return json.dumps(self, cls=IdlEncoder, sort_keys=True, indent=4)
        # Use compact separators so output identical to Perl
        return json.dumps(self, cls=IdlEncoder, sort_keys=True, separators=(',', ':'))

class IdlCallbackFunction(BaseIdl, TypedObject):
    def __init__(self, name=None, data_type=None, arguments=None):
        self.data_type = data_type
        self.name = name
        self.arguments = arguments or []

    def resolve_typedefs(self, typedefs):
        TypedObject.resolve_typedefs(self, typedefs)
        for argument in self.arguments:
            argument.resolve_typedefs(typedefs)
        raise ValueError('Typedefs in callback functions are untested!')

    def json_serializable(self):
        return {
            'callbackFunction::name': self.name,
            'callbackFunction::type': self.data_type,
            'callbackFunction::parameters': self.arguments,
            }

class IdlEnum(BaseIdl):
    def __init__(self, name=None, values=None):
        self.name = name
        self.values = values or []

    def json_serializable(self):
        return {
            'domEnum::name': self.name,
            'domEnum::values': self.values,
            }

class IdlInterface(BaseIdl):
    def __init__(self, attributes=None, constants=None, constructors=None, custom_constructors=None, extended_attributes=None, operations=None, is_callback=False, is_partial=False, name=None, parent=None):
        self.attributes = attributes or []
        self.constants = constants or []
        self.constructors = constructors or []
        self.custom_constructors = custom_constructors or []
        self.extended_attributes = extended_attributes or {}
        self.operations = operations or []
        self.is_callback = is_callback
        self.is_partial = is_partial
        self.name = name
        self.parent = parent

    def resolve_typedefs(self, typedefs):
        for attribute in self.attributes:
            attribute.resolve_typedefs(typedefs)
        for constant in self.constants:
            constant.resolve_typedefs(typedefs)
        for constructor in self.constructors:
            constructor.resolve_typedefs(typedefs)
        for custom_constructor in self.custom_constructors:
            custom_constructor.resolve_typedefs(typedefs)
        for operation in self.operations:
            operation.resolve_typedefs(typedefs)

    def json_serializable(self):
        return {
            'domInterface::attributes': self.attributes,
            'domInterface::constants': self.constants,
            'domInterface::constructors': self.constructors,
            'domInterface::customConstructors': self.custom_constructors,
            'domInterface::extendedAttributes': none_to_value_is_missing(self.extended_attributes),
            'domInterface::functions': self.operations,
            'domInterface::isException': None,
            'domInterface::isCallback': boolean_to_perl(false_to_none(self.is_callback)),
            'domInterface::isPartial': false_to_none(self.is_partial),
            'domInterface::name': self.name,
            'domInterface::parent': self.parent,
            }

class IdlException(BaseIdl):
    def __init__(self, name=None, constants=None, operations=None, attributes=None, extended_attributes=None):
        self.attributes = attributes or []
        self.constants = constants or []
        self.extended_attributes = extended_attributes or {}
        self.operations = operations or []
        self.name = name

    def resolve_typedefs(self, typedefs):
        for constant in self.constants:
            constant.resolve_typedefs(typedefs)
        for attribute in self.attributes:
            attribute.resolve_typedefs(typedefs)
        for operations in self.operations:
            operations.resolve_typedefs(typedefs)

    def json_serializable(self):
        return {
            # Perl code treats Exceptions as a kind of Interface
            'domInterface::name': self.name,
            'domInterface::attributes': self.attributes,
            'domInterface::constants': self.constants,
            'domInterface::extendedAttributes': none_to_value_is_missing(self.extended_attributes),
            'domInterface::functions': self.operations,
            # These values don't vary for exceptions
            'domInterface::constructors': [],
            'domInterface::customConstructors': [],
            'domInterface::isException': 1,
            'domInterface::isCallback': None,
            'domInterface::isPartial': None,
            'domInterface::parent': None,
            }

class IdlAttribute(BaseIdl, TypedObject):
    def __init__(self, data_type=None, extended_attributes=None, getter_exceptions=None, is_nullable=False, is_static=False, is_read_only=False, name=None, setter_exceptions=None):
        self.data_type = data_type
        self.extended_attributes = extended_attributes or {}
        self.getter_exceptions = getter_exceptions or []
        self.is_nullable = is_nullable
        self.is_static = is_static
        self.is_read_only = is_read_only
        self.name = name
        self.setter_exceptions = setter_exceptions or []

    def json_serializable(self):
        return {
            'domAttribute::extendedAttributes': none_to_value_is_missing(self.extended_attributes),
            'domAttribute::getterExceptions': self.getter_exceptions,
            'domAttribute::isNullable': boolean_to_perl_quoted(false_to_none(self.is_nullable)),
            'domAttribute::isReadOnly': boolean_to_perl(false_to_none(self.is_read_only)),
            'domAttribute::isStatic': boolean_to_perl(false_to_none(self.is_static)),
            'domAttribute::name': self.name,
            'domAttribute::setterExceptions': self.setter_exceptions,
            'domAttribute::type': self.data_type,
            }

class IdlConstant(BaseIdl, TypedObject):
    def __init__(self, name=None, data_type=None, value=None, extended_attributes=None):
        self.data_type = data_type
        self.extended_attributes = extended_attributes or {}
        self.name = name
        self.value = value

    def json_serializable(self):
        return {
            'domConstant::extendedAttributes': none_to_value_is_missing(self.extended_attributes),
            'domConstant::name': self.name,
            'domConstant::type': self.data_type,
            'domConstant::value': self.value,
            }

class IdlOperation(BaseIdl, TypedObject):
    def __init__(self, is_static=False, name=None, data_type=None, extended_attributes=None, specials=None, arguments=None, overloaded_index=None):
        self.is_static = is_static
        self.name = name or ''
        self.data_type = data_type
        self.extended_attributes = extended_attributes or {}
        self.specials = specials or []
        self.arguments = arguments or []
        self.overloaded_index = overloaded_index

    def resolve_typedefs(self, typedefs):
        TypedObject.resolve_typedefs(self, typedefs)
        for argument in self.arguments:
            argument.resolve_typedefs(typedefs)

    def json_serializable(self):
        return {
            'domFunction::extendedAttributes': none_to_value_is_missing(self.extended_attributes),
            'domFunction::isStatic': boolean_to_perl(false_to_none(self.is_static)),
            'domFunction::name': self.name,
            'domFunction::overloadedIndex': self.overloaded_index,
            'domFunction::parameters': self.arguments,
            'domFunction::specials': self.specials,
            'domFunction::type': self.data_type,
            }

class IdlArgument(BaseIdl, TypedObject):
    def __init__(self, name=None, data_type=None, extended_attributes=None, is_optional=False, is_nullable=None, is_variadic=False):
        self.data_type = data_type
        self.extended_attributes = extended_attributes or {}
        # FIXME: boolean values are inconsistent.
        # The below hack is so that generated JSON is identical to
        # Perl-generated JSON, where the exact values depend on the code path.
        # False and None (Perl: 0 and undef) are semantically interchangeable,
        # but yield different JSON.
        # Once Perl removed, have all default to False.
        if is_optional is None:
            is_optional = False
            if is_variadic is None:
                is_variadic = False
        self.is_nullable = is_nullable  # (T?)
        self.is_optional = is_optional  # (optional T)
        self.is_variadic = is_variadic  # (T...)
        self.name = name

    def json_serializable(self):
        return {
            'domParameter::extendedAttributes': none_to_value_is_missing(self.extended_attributes),
            'domParameter::isNullable': boolean_to_perl_quoted(self.is_nullable),
            'domParameter::isOptional': boolean_to_perl(self.is_optional),
            'domParameter::isVariadic': boolean_to_perl(self.is_variadic),
            'domParameter::name': self.name,
            'domParameter::type': self.data_type,
            }

# Type classes

class IdlType:
    # FIXME: replace type strings with these objects,
    # so don't need to parse everywhere types are used.
    # Types are stored internally as strings, not objects,
    # e.g., as 'sequence<Foo>' or 'Foo[]',
    # hence need to parse the string whenever a type is used.
    # FIXME: incorporate Nullable, Variadic, etc.
    # FIXME: properly should nest types
    # Formally types are nested, e.g., short?[] vs. short[]?,
    # but in practice these complex types aren't used and can treat
    # as orthogonal properties.
    def __init__(self, base_type, is_array=False, is_sequence=False):
        if is_array and is_sequence:
            raise ValueError('Array of Sequences are not allowed.')
        self.base_type = base_type
        self.is_array = is_array
        self.is_sequence = is_sequence

    def __str__(self):
        type_string = self.base_type
        if self.is_array:
            return type_string + '[]'
        if self.is_sequence:
            return 'sequence<%s>' % type_string
        return type_string

    @classmethod
    def from_string(cls, type_string):
        sequence_re = r'^sequence<([^>]*)>$'
        if type_string.endswith('[]'):
            type_string = type_string[:-2]
            sequence_match = re.match(sequence_re, type_string)
            if sequence_match:
                raise ValueError('Array of Sequences are not allowed.')
            return cls(type_string, is_array=True)
        sequence_match = re.match(sequence_re, type_string)
        if sequence_match:
            base_type = sequence_match.group(1)
            return cls(base_type, is_sequence=True)
        return cls(type_string)

class IdlTypedef:
    # Internal to IDL parsing: typedefs are all translated during IdlObject
    # construction, and the typedefs themselves not stored in the object.
    def __init__(self, extended_attributes=None, data_type=None):
        self.extended_attributes = extended_attributes or {}
        self.data_type = data_type

class IdlUnionType(BaseIdl):
    def __init__(self, union_member_types=None):
        self.union_member_types = union_member_types or []

    def json_serializable(self):
        return {
            'UnionType::unionMemberTypes': self.union_member_types,
            }

# Perl JSON compatibility functions

def none_to_value_is_missing(extended_attributes):
    # Perl IDL Parser uses 'VALUE_IS_MISSING' for null values in
    # extended attributes, so add this as a filter when exporting to JSON.
    new_extended_attributes = {}
    for key, value in extended_attributes.iteritems():
        if value is None:
            new_extended_attributes[key] = 'VALUE_IS_MISSING'
        else:
            new_extended_attributes[key] = value
    return new_extended_attributes

def boolean_to_perl(value):
    # Perl stores booleans as 1, 0, or undefined (JSON null);
    # convert to this format.
    if value is None:
        return None
    return int(value)

def boolean_to_perl_quoted(value):
    # Bug-for-bug compatibility with Perl.
    # The value of isNullable is quoted ('1', '0', or undefined), rather than
    # an integer, so add quotes.
    if value is None:
        return None
    return str(int(value))

def false_to_none(value):
    # The Perl parser generally uses undefined (Python None) rather than False
    # for boolean flags, because the value is simply left undefined, rather than
    # explicitly set to False.
    if value is False:
        return None
    return value

# JSON export

class IdlEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, BaseIdl):
            return obj.json_serializable()
        return json.JSONEncoder.default(self, obj)


import pytest
import torch
from deepspeed.ops import op_builder

inference_module = None

def run_quantize_ds(activations, num_groups, q_bits, is_symmetric_quant):
    global inference_module
    if inference_module is None:
        inference_module = op_builder.QuantizerBuilder().load()

    return inference_module.quantize(
        activations,
        num_groups,
        q_bits,
        inference_module.Symmetric
        if is_symmetric_quant else inference_module.Asymmetric)

def get_q_props(q_bits):
    q_range = 2**q_bits
    q_min = -(2**(q_bits - 1))
    q_max = (2**(q_bits - 1) - 1)

    q_min = torch.IntTensor([q_min]).to(device='cuda')
    q_max = torch.IntTensor([q_max]).to(device='cuda')
    return q_range, q_max, q_min

def get_scale_zero_point(q_bits,
                         is_symmetric_quant,
                         max,
                         min,
                         absmax,
                         scales=None,
                         zero_points=None):

    q_range, q_max, q_min = get_q_props(q_bits)

    if is_symmetric_quant:
        scale = torch.empty_like(absmax)
        for i, x in enumerate(absmax):
            scale[i] = torch.ones_like(x) if x == 0 else q_range / (2 * x)
        zero_point = torch.zeros(scale.shape, dtype=torch.float32, device='cuda')
    else:
        scale = torch.empty_like(max)
        for i, x in enumerate(max):
            scale[i] = torch.ones_like(x) if max[i] == min[i] else q_range / (max[i] -
                                                                              min[i])
        zero_point = q_min - (min * scale)

    return scale, zero_point

def int4x2to2xint4(int4X2tensor):
    high = int4X2tensor >> 4
    low = (int4X2tensor << 4) >> 4
    return torch.stack((high, low), dim=-1).flatten()

def run_float_quantize(q_bits, is_symmetric_quant, activations_ref, num_groups):

    # Reference implementation
    # https://pytorch.org/docs/stable/quantization-support.html

    activations_ref = activations_ref.reshape(num_groups, -1).to(dtype=torch.float32)

    max_abs_activations_ref = torch.amax(torch.abs(activations_ref),
                                         dim=-1).view(num_groups,
                                                      -1)
    max_activations_ref = torch.amax(activations_ref, dim=-1).view(num_groups, -1)
    min_activations_ref = torch.amin(activations_ref, dim=-1).view(num_groups, -1)

    _, q_max, q_min = get_q_props(q_bits)

    scale, zero_point = get_scale_zero_point(q_bits, is_symmetric_quant, max_activations_ref, min_activations_ref, max_abs_activations_ref)

    data_f = activations_ref * scale

    if not is_symmetric_quant:
        data_f = data_f + zero_point

    data_i32 = torch.round(data_f).to(dtype=torch.int32)

    data_i32 = torch.minimum(torch.maximum(data_i32,
                                           q_min.expand_as(data_i32)),
                             q_max.expand_as(data_i32))
    data_i8 = data_i32.to(dtype=torch.int8)

    scales = (1.0 / scale).reshape(-1, 1)
    offsets = zero_point.reshape(-1, 1)
    params = torch.cat((scales, offsets), dim=-1)

    return data_i8, params

@pytest.mark.inference
@pytest.mark.parametrize("num_groups", [1, 13, 512])
@pytest.mark.parametrize("num_elems",
                         [8,
                          16,
                          32,
                          64,
                          128,
                          256,
                          4096,
                          8192,
                          12288,
                          16384])
@pytest.mark.parametrize("is_symmetric_quant", [True, False])
@pytest.mark.parametrize("q_bits", [4, 8])
@pytest.mark.parametrize("directed_case", ["all_zeros", None])
def test_float_quantize(num_elems,
                        num_groups,
                        is_symmetric_quant,
                        q_bits,
                        directed_case):

    if directed_case == "all_zeros":
        activations_ds = torch.zeros((num_groups,
                                      num_elems),
                                     dtype=torch.float16,
                                     device='cuda')
    else:
        activations_ds = torch.randn((num_groups,
                                      num_elems),
                                     dtype=torch.float16,
                                     device='cuda')
    activations_ref = activations_ds.clone().detach()

    ref_out_tensor, ref_params = run_float_quantize(q_bits, is_symmetric_quant, activations_ref, num_groups)

    ds_out_tensor, ds_out_params = run_quantize_ds(activations_ds, num_groups, q_bits, is_symmetric_quant)

    if (q_bits == 4):
        ds_out_tensor = int4x2to2xint4(ds_out_tensor)

    # Allow a max difference of 1 to account for differences in rounding in pytorch implementation
    assert (torch.all(
        torch.lt(torch.abs(ds_out_tensor.flatten() - ref_out_tensor.flatten()),
                 2)))
    if is_symmetric_quant:
        assert (torch.allclose(ds_out_params.flatten(), ref_params[:, 0].flatten()))
    else:
        assert (torch.allclose(ds_out_params[:,
                                             0].flatten(),
                               ref_params[:,
                                          0].flatten()))
        assert (torch.allclose(ds_out_params[:,
                                             1].flatten(),
                               ref_params[:,
                                          1].flatten(),
                               atol=5e-5,
                               rtol=5e-5))

def run_integer_quantize_ds(activations, num_groups, q_bits):
    global inference_module
    if inference_module is None:
        inference_module = op_builder.QuantizerBuilder().load()

    return inference_module.quantize(activations,
                                     num_groups,
                                     q_bits,
                                     inference_module.IntegerSymmetric)

def run_integer_quantize(q_bits, activations_ref, num_groups):

    activations_ref = activations_ref.reshape(num_groups, -1).to(dtype=torch.float32)

    max_abs_activations_ref = torch.amax(torch.abs(activations_ref),
                                         dim=-1).view(num_groups,
                                                      -1)

    _, q_max, q_min = get_q_props(q_bits)
    print(max_abs_activations_ref)
    max_abs_activations_ref = (max_abs_activations_ref + 1).to(torch.int8).to(
        torch.float32)
    print(max_abs_activations_ref)
    numerator = activations_ref * q_max
    print(numerator.dtype)
    denominator = max_abs_activations_ref

    data_f = numerator / denominator

    data_i32 = torch.round(data_f).to(dtype=torch.int32)
    data_i32 = torch.minimum(torch.maximum(data_i32,
                                           q_min.expand_as(data_i32)),
                             q_max.expand_as(data_i32))
    data_i8 = data_i32.to(dtype=torch.int8)

    return data_i8, max_abs_activations_ref.to(torch.int32)

@pytest.mark.inference
@pytest.mark.parametrize("num_groups", [1, 2, 4, 8, 16, 32, 64, 512])
@pytest.mark.parametrize("num_elems", [4096, 8192, 12288, 16384])
@pytest.mark.parametrize("q_bits", [4, 8])
def test_integer_quantize(num_elems, num_groups, q_bits):

    activations_ds = torch.ones((num_groups,
                                 num_elems),
                                dtype=torch.float16,
                                device='cuda') * 0.35
    activations_ref = activations_ds.clone().detach()

    ref_out_tensor, ref_params = run_integer_quantize(q_bits, activations_ref, num_groups)

    ds_out_tensor, ds_out_params = run_integer_quantize_ds(activations_ds, num_groups, q_bits)

    if (q_bits == 4):
        ds_out_tensor = int4x2to2xint4(ds_out_tensor)

    # Allow a max difference of 1 to account for differences in rounding in pytorch implementation
    assert (torch.all(
        torch.lt(torch.abs(ds_out_tensor.flatten() - ref_out_tensor.flatten()),
                 2)))
    assert (torch.allclose(ds_out_params.flatten(), ref_params.flatten()))


#!/usr/bin/python

import sys

from PyQt5.QtCore import Qt, QMimeData
from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QTextEdit
from PyQt5.QtWidgets import QPushButton, QVBoxLayout, QGridLayout, QFileDialog
from PyQt5.QtGui import QTextCursor, QPainter

class HexEditWidget(QTextEdit):
    def __init__(self, parent):
        QTextEdit.__init__(self, parent)

        self.setCursorWidth(8)
        self.setLineWrapMode(QTextEdit.NoWrap);

        # set the font
        font = self.font()
        font.setFamily("Courier")
        font.setPointSize(10)
        self.setFont(font)

        self.bytesPerLine = 16
        self.addressWidth = 10

        self.addrColor = Qt.red
        self.addrBackground = Qt.white
        self.dataColor = Qt.black
        self.dataBackground = Qt.white
        self.asciiColor = Qt.black
        self.asciiBackground = Qt.lightGray

        self.cursorPositionChanged.connect(self.showPosition)

    def paintEvent(self, event):
        QTextEdit.paintEvent(self, event)

        viewPort = self.viewport()

        cursor = self.textCursor()
        columnNumber = cursor.positionInBlock()
        rect = self.cursorRect()

        p = QPainter(viewPort)

        byteIndex = (columnNumber - 10) // 3

        # 4 pixels offset to the left, 58 columns to the ascii area
        p.setPen(Qt.red)
        p.drawRect( 4 + (58 + byteIndex) * rect.size().width(),
                   #rect.left() + 384,
                   rect.top(),# columnNumber * (rect.size().width()),
                    rect.size().width() - 1, rect.size().height())

    def keyPressEvent(self, event):
        # Ensure that shortcuts like CTRL-C are working
        if event.modifiers() & Qt.ControlModifier:
            if event.key() == Qt.Key_C:
                QTextEdit.keyPressEvent(self, event)

        else:
            # Filter keys (TODO)
            text = event.text()
            if text:
                print(event.text())
            else:
                print("KEY: {}".format(event.key()))

    def keyReleaseEvent(self, event):
        pass

    def mousePressEvent(self, event):
        QTextEdit.mousePressEvent(self, event)

        cursor = self.cursorForPosition(event.pos())
        columnNumber = cursor.positionInBlock()

        # Limit cursor position to the editable area
        # - TODO: remove hard coded numbers
        if columnNumber < self.addressWidth:
            columnNumber = self.addressWidth
        if columnNumber > 73:
            columnNumber = 73;

        # Not all of the characters displayed in the hex dump area can be edited
        if columnNumber < 58:
            bytePos = columnNumber - self.addressWidth
            byteIdx = bytePos % 3   # 0, 1, 2 => "7F "
            if byteIdx == 2:
                columnNumber -= 1

        block = cursor.block()      # The current block of the cursor
        cursor.setPosition(block.position() + columnNumber)

        self.setTextCursor(cursor)

        self.viewport().update()    # necessary to get the right coordinates in the paint event

    def getCurrentAddr(self, crsr):
        '''Returns the address of the byte where the cursor is located'''

        lineNumber = crsr.blockNumber()
        columnNumber = crsr.positionInBlock()
        if columnNumber < 58:   # data area
            addr = self.bytesPerLine * lineNumber + (columnNumber - self.addressWidth) // 3
        else:
            addr = self.bytesPerLine * lineNumber + (columnNumber - 58)
        return addr

    def getDataPositionForAddr(self, addr):
        '''returns the cursor position within the data area for the given address'''
        lineLength = self.addressWidth + 3* self.bytesPerLine + 16 + 1  # addr: data ascii\n
        blockNo = addr // self.bytesPerLine
        column = (addr % self.bytesPerLine) * 3 + self.addressWidth
        result = blockNo * lineLength + column
        return result

    def getByte(self, addr):
        crsr = self.textCursor()

        byteStart = self.getDataPositionForAddr(addr)
        byteEnd = byteStart + 2
        crsr.setPosition(byteStart)
        crsr.setPosition(byteEnd, QTextCursor.KeepAnchor)
        data = int(crsr.selectedText(), 16)
        return data

    def showPosition(self):
        '''cursorPositionChanged() slot'''

        crsr = self.textCursor()
        addr = self.getCurrentAddr(crsr)
        data = self.getByte(addr)

        # TODO: result be invalid if improper data is available
        if (data & 0b11111110) == 0b11111100:
            count = 6
            unicode = (data & 0b00000001)
        elif (data & 0b11111100) == 0b11111000:
            count = 5
            unicode = (data & 0b00000011)
        elif (data & 0b11111000) == 0b11110000:
            count = 4
            unicode = (data & 0b00000111)
        elif (data & 0b11110000) == 0b11100000:
            count = 3
            unicode = (data & 0b00001111)
        elif (data & 0b11100000) == 0b11000000:
            count = 2
            unicode = (data & 0b00011111)
        elif (data & 0b10000000) == 0b00000000:
            count = 1
            unicode = data
        else:
            count = 0

        for x in range(1, count):
            nextByte = self.getByte(addr + x)

            # shift bits by 6 to get space for the next bunch of bits
            unicode = unicode * 0x40

            # add the next bunch of bits to the unicode value
            unicode = unicode + (nextByte & 0b00111111)

        if count == 0:
            print('Byte address: 0x{0:08X} = {1:02X} = \\u (N/A)'.format(addr, data))
        else:
            print('Byte address: 0x{0:08X} = {1:02X} = \\u{2:2X}'.format(addr, data, unicode))

    def addLine(self, addr, hexDump, asciiDump):
        self.textCursor().movePosition(QTextCursor.End)

        self.setTextColor(self.addrColor)
        self.setTextBackgroundColor(self.addrBackground)
        self.insertPlainText("{0:08X}: ".format(addr))

        self.setTextColor(self.dataColor)
        self.setTextBackgroundColor(self.dataBackground)
        self.insertPlainText("{0:48}".format(hexDump))

        self.setTextColor(self.asciiColor)
        self.setTextBackgroundColor(self.asciiBackground)
        self.insertPlainText("{0}\n".format(asciiDump))

    def setContents(self, contents):
        self.blockSignals(True)

        addr = 0
        hexDump = ""
        asciiDump = ""
        column = 0

        self.clear()
        for c in contents:
            hexDump = hexDump + "{0:02X} ".format(c)
            asciiDump = asciiDump + ( chr(c) if c > 31 and c < 128 else '.')
            column += 1
            if (column % 16) == 0:
                self.addLine(addr, hexDump, asciiDump)

                asciiDump = ""
                hexDump = ""
                column = 0
                addr += 16
        if column > 0:
            self.addLine(addr, hexDump, asciiDump)

        self.blockSignals(False)

        crsr = self.textCursor()
        crsr.setPosition(10)
        self.setTextCursor(crsr)

class MainWindow(QMainWindow):
    def __init__(self):

        QMainWindow.__init__(self)
        self.resize(800, 600)

        # Create the main content widget
        mainWidget = QWidget(self)
        self.setCentralWidget(mainWidget)

        # Create a text component at the top area of the main widget
        self.output = HexEditWidget(mainWidget)

        mainLayout = QVBoxLayout(mainWidget)
        mainLayout.addWidget(self.output)

        # Create buttons in a grid layout below the top area
        buttonWidget = QWidget(mainWidget)
        self.buttonLayout = QGridLayout(buttonWidget)
        mainLayout.addWidget(buttonWidget)

        theButton = QPushButton("Load...")
        theButton.clicked.connect(self.chooseFile)
        self.buttonLayout.addWidget(theButton, 0, 0)

        statusBar = self.statusBar();

        self.loadFile("__pycache__/helloPyside.cpython-33.pyc")

    def chooseFile(self):
        fileName = QFileDialog.getOpenFileName(self, "Open file ...");
        fileName = fileName[0]
        if fileName:
            self.loadFile(fileName)

    def loadFile(self, fileName):
        file = open(fileName, 'rb')
        contents = file.read()
        file.close()
        self.output.setContents(contents)

class HexDump():
    def __init__(self):
        pass

    def main(self, args):
        # Create a Qt application
        app = QApplication(sys.argv)

        # Create and whow the main window
        wnd = MainWindow()
        wnd.show()

        # Run the application
        app.exec_()

if __name__ == "__main__":
    app = HexDump()
    app.main(sys.argv)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Usage:
# PyStandMate.py --help
# PyStandMate.py --response-file PyStandMate.rsp
# PyStandMate.py --bitness 32 --compiler MSVC --python-version 3.8.10 --console
# PyStandMate.py --compiler MSVC --python-version 3.8.10 --package PyQt5 --pip-index-url https://pypi.doubanio.com/simple
# PyStandMate.py --bitness 64 --compiler MSVC --python-version 3.8.10 --package PyQt5 --pip-index-url https://pypi.doubanio.com/simple
# PyStandMate.py --bitness 64 --compiler MSVC --python-version 3.10.0 --package PyQt6 --pip-index-url https://pypi.doubanio.com/simple

import argparse
import collections
import os
from pathlib import Path
from pprint import pprint
import re
import shutil
import subprocess
import sys
import urllib.parse
import urllib.request
import zipfile

EmbedPython = collections.namedtuple("EmbedPython", ["url", "filename"])

DOWNLOAD_DIR = "download"
BUILD_DIR = "build"
PUBLISH_DIR = "publish"

DEFAULT_USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6) AppleWebKit/603.3.8 (KHTML, like Gecko) Version/10.1.2 Safari/603.3.8"

class LoadFromFile(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        with values as f:
            # parse arguments in the file and store them in the target namespace
            parser.parse_args(f.read().split(), namespace)

def fetch_page(url, encoding="utf-8"):
    request = urllib.request.Request(url, headers={"User-Agent": DEFAULT_USER_AGENT})
    response = urllib.request.urlopen(request)

    if encoding:
        charset = response.headers.get_content_charset(failobj=encoding)
        for line in response:
            yield line.decode(charset)
    else:
        yield from response

def fetch_page_contents(url, encoding="utf-8"):
    return "".join(fetch_page(url, encoding))

def find_urls(s):
    return re.findall(r'href=[\'"]?([^\'" >]+)', s)

def download_pystand(version, target_dir):
    if version.startswith("v") or version.startswith("V"):
        version = version[1:]

    filename = f"PyStand-v{version}-exe.zip"
    pystand_path = target_dir / filename

    if not pystand_path.is_file():
        if not target_dir.is_dir():
            print(f"Create directory {target_dir}...")
            target_dir.mkdir()

        dload_url = f"https://github.com/skywind3000/PyStand/releases/download/{version}/{filename}"
        print(f"Download {dload_url} -> {filename}...")
        urllib.request.urlretrieve(dload_url, pystand_path)

    return pystand_path

def get_pystand_subdir(compiler, bitness, is_console):
    subsystem = "CLI" if is_console else "GUI"

    if compiler == "MSVC":
        arch = "Win32" if bitness == 32 else "x64"
    else:
        arch = "mingw32" if bitness == 32 else "mingw64"

    return f"PyStand-{arch}-{subsystem}"

def get_pystand_publish_subdir(version, compiler, bitness, is_console):
    if version.startswith("v") or version.startswith("V"):
        version = version[1:]

    subdir = get_pystand_subdir(compiler, bitness, is_console)
    subdir = subdir.replace("PyStand", f"PyStand-v{version}")

    return subdir

def get_embed_python_versions():
    page_contents = fetch_page_contents("https://www.python.org/downloads/windows/")
    dload_urls = find_urls(page_contents)

    result = collections.OrderedDict()

    for url in dload_urls:
        url_parts = urllib.parse.urlparse(url)
        if Path(url_parts.path).match("*embed*.zip"):
            version = url_parts.path.split("-")[1]
            embed_python = EmbedPython(url, Path(url_parts.path).name)
            if not version in result:
                result[version] = [embed_python]
            else:
                result[version].append(embed_python)

    return result

def download_embed_python(version, bitness, target_dir):
    arch = "win32" if bitness == 32 else "amd64"
    filename = f"python-{version}-embed-{arch}.zip"
    embed_python_path = target_dir / filename

    if not embed_python_path.is_file():
        print(f"{filename} doesn't exist, will download it first.")

        print("Get available Windows embeddable Python packages...")
        embed_python_versions = get_embed_python_versions()

        if version not in embed_python_versions:
            print(f"Couldn't find embeddable Python package of version {version}.")
            print("Available versions:")
            pprint(tuple(embed_python_versions.keys()))
            sys.exit(1)

        if not target_dir.is_dir():
            print(f"Create directory {target_dir}...")
            target_dir.mkdir()

        embed_python_list = embed_python_versions[version]
        for embed_python in embed_python_list:
            if arch in embed_python.filename:
                print(f"Download {embed_python.url} -> {embed_python.filename}...")
                urllib.request.urlretrieve(embed_python.url, embed_python_path)
                break

    if not embed_python_path.is_file():
        print(f"Couldn't find a suitable version of embeddable Python.")
        sys.exit(2)

    return embed_python_path

def install_pip(embed_python_dir, get_pip_path):
    pth_files = tuple(embed_python_dir.glob("*._pth"))
    if pth_files:
        if len(pth_files) != 1:
            print("There are more than one ._pth files:")
            pprint(pth_files)
            sys.exit(3)

        pth_file = pth_files[0]

        print("Uncomment import site...")

        with open(pth_file, "r") as fp:
            pth_file_content = fp.read()

        pth_file_content = pth_file_content.replace("#import site", "import site")
        with open(pth_file, "w") as fp:
            fp.write(pth_file_content)

    print("Install pip...")
    python = embed_python_dir / "python.exe"
    # subprocess.run([python, get_pip_path])
    subprocess.run(["cmd", "/C", python, get_pip_path], check=True)

def install_package(embed_python_dir, package, pip_index_url):
    python = embed_python_dir / "python.exe"

    args = ["cmd", "/C", python, "-m", "pip", "install", package]
    if pip_index_url:
        args.extend(["-i", pip_index_url])

    subprocess.run(args, check=True)

def install_requirements(embed_python_dir, requirements_file, pip_index_url):
    python = embed_python_dir / "python.exe"

    args = ["cmd", "/C", python, "-m", "pip", "install", "-r", requirements_file]
    if pip_index_url:
        args.extend(["-i", pip_index_url])

    subprocess.run(args, check=True)

def install_packages(embed_python_dir, packages, pip_index_url):
    for package in packages:
        if "requirements.txt" in Path(package).name:
            install_requirements(embed_python_dir, package, pip_index_url)
        else:
            install_package(embed_python_dir, package, pip_index_url)

def main():
    script_path = Path(sys.argv[0])
    script_dir = script_path.parent

    parser = argparse.ArgumentParser(
        prog=script_path.stem,
        description="PyStand packaging mate.",
    )

    parser.add_argument("--pystand-version", default="1.0.11", help="PyStand version")
    parser.add_argument(
        "--bitness", type=int, choices=(32, 64), default=32, help="Bitness"
    )
    parser.add_argument(
        "--compiler", choices=("MSVC", "GCC"), default="MSVC", help="Compiler"
    )
    parser.add_argument(
        "--console", action="store_true", help="Use PyStand CLI instead of GUI"
    )
    parser.add_argument(
        "--pystand-int", default="PyStand.int", help="PyStand.int script"
    )
    parser.add_argument("--python-version", default="3.8.10", help="Python version")
    parser.add_argument(
        "--package", nargs="+", help="A list of 3rd-party packages to be installed"
    )
    parser.add_argument(
        "--pip-index-url",
        help="Base url of Python package index",
    )
    parser.add_argument(
        "--response-file",
        type=open,
        action=LoadFromFile,
        help="Read options stored in a response file",
    )

    args = parser.parse_args()

    # 1. Download PyStand
    pystand_path = download_pystand(args.pystand_version, script_dir / DOWNLOAD_DIR)

    # 2. Extract PyStand
    pystand_dir = script_dir / BUILD_DIR / pystand_path.stem
    print(f"Extract {pystand_path.name} -> {pystand_dir.name}...")
    with zipfile.ZipFile(pystand_path, "r") as zip_ref:
        zip_ref.extractall(pystand_dir)

    # 3. Download Python
    embed_python_path = download_embed_python(
        args.python_version, args.bitness, script_dir / DOWNLOAD_DIR
    )

    # 4. Extract Python
    embed_python_dir = script_dir / BUILD_DIR / embed_python_path.stem
    if embed_python_dir.is_dir():
        print("Remove build directory...")
        shutil.rmtree(embed_python_dir)
    print(f"Extract {embed_python_path.name} -> {embed_python_dir.name}...")
    with zipfile.ZipFile(embed_python_path, "r") as zip_ref:
        zip_ref.extractall(embed_python_dir)

    # 5. Put together
    pystand_publish_subdir = get_pystand_publish_subdir(
        args.pystand_version, args.compiler, args.bitness, args.console
    )
    pystand_publish_dir = script_dir / PUBLISH_DIR / pystand_publish_subdir
    if pystand_publish_dir.is_dir():
        print(f"Remove publish directory...")
        shutil.rmtree(pystand_publish_dir)
    if not pystand_publish_dir.is_dir():
        print(f"Create directory {PUBLISH_DIR}{os.sep}{pystand_publish_subdir}...")
        pystand_publish_dir.mkdir(parents=True)
    # 5.1 Copy PyStand
    pystand_subdir = get_pystand_subdir(args.compiler, args.bitness, args.console)
    pystand_src_path = pystand_dir / pystand_subdir / "PyStand.exe"
    pystand_dst_path = pystand_publish_dir / "PyStand.exe"
    print(f"Copy PyStand...")
    shutil.copy2(pystand_src_path, pystand_dst_path)
    # 5.2 Copy Python
    print(f"Copy Python...")
    runtime_dir = pystand_publish_dir / "runtime"
    shutil.copytree(embed_python_dir, runtime_dir, dirs_exist_ok=True)
    # 5.3 Copy PyStand.int
    pystand_int_path = Path(args.pystand_int)
    if pystand_int_path.is_file():
        print(f"Copy {pystand_int_path.name} -> PyStand.int...")
        shutil.copy2(pystand_int_path, pystand_publish_dir / "PyStand.int")
    else:
        print(f"{pystand_int_path} is not a regular file.")

    if not args.package:
        sys.exit(0)

    # 6. Download & install pip
    get_pip_path = script_dir / DOWNLOAD_DIR / "get-pip.py"
    if not get_pip_path.is_file():
        print("Download get-pip.py...")
        urllib.request.urlretrieve("https://bootstrap.pypa.io/get-pip.py", get_pip_path)

    target_get_pip_path = embed_python_dir / "get-pip.py"
    shutil.copyfile(get_pip_path, target_get_pip_path)
    install_pip(embed_python_dir, target_get_pip_path)

    # 7. Memorize files and directories that are created by installing pip and setuptools.
    pip_facilities = [target_get_pip_path]
    site_packages_dir = embed_python_dir / "Lib" / "site-packages"
    scripts_dir = embed_python_dir / "Scripts"
    pip_facilities.extend(tuple(site_packages_dir.iterdir()))
    pip_facilities.extend(tuple(scripts_dir.iterdir()))
    # pprint(pip_facilities)

    # 8. Install packages
    print("Install packages...")
    install_packages(embed_python_dir, args.package, args.pip_index_url)

    # 9. Remove pip and setuptools.
    print("Remove pip and setuptools...")
    for facility in pip_facilities:
        if facility.is_file():
            print(f"Remove file {facility.name}...")
            facility.unlink(missing_ok=True)
        elif facility.is_dir():
            print(f"Remove directory {facility.name}...")
            shutil.rmtree(facility)

    # 10. Copy site-packages.
    if site_packages_dir.is_dir():
        # Remove .dist-info folders.
        for dist_info_dir in site_packages_dir.glob("*.dist-info"):
            if dist_info_dir.is_dir():
                print(f"Remove directory {dist_info_dir.name}...")
                shutil.rmtree(dist_info_dir)

        print("Copy installed packages...")
        shutil.copytree(
            site_packages_dir, pystand_publish_dir / "site-packages", dirs_exist_ok=True
        )

if __name__ == "__main__":
    main()

# Format code:
# pip install black
# black PyStandMate.py

# References:
# [Regular expression to extract URL from an HTML link](https://stackoverflow.com/questions/499345/regular-expression-to-extract-url-from-an-html-link)
# [Unzipping files in Python](https://stackoverflow.com/questions/3451111/unzipping-files-in-python)
# https://gist.github.com/myd7349/9f7c6334e67d1aee68a722a15df4a62a
# [Replace string within file contents](https://stackoverflow.com/questions/4128144/replace-string-within-file-contents)
# https://docs.python.org/3.11/library/pathlib.html
# [Could not find a version that satisfies the requirement setuptools](https://github.com/pypa/pip/issues/7730)
# [How to run a pip install command from a subproces.run()](https://stackoverflow.com/questions/69345839/how-to-run-a-pip-install-command-from-a-subproces-run)
# [How can I Install a Python module within code?](https://stackoverflow.com/questions/12332975/how-can-i-install-a-python-module-within-code)
# [How to run `pip` in a virtualenv with subprocess.check_call()?](https://stackoverflow.com/questions/28574058/how-to-run-pip-in-a-virtualenv-with-subprocess-check-call)
# [Python: Platform independent way to modify PATH environment variable](https://stackoverflow.com/questions/1681208/python-platform-independent-way-to-modify-path-environment-variable)
# [Check if Python Package is installed](https://stackoverflow.com/questions/1051254/check-if-python-package-is-installed)
# [how to get argparse to read arguments from a file with an option rather than prefix](https://stackoverflow.com/questions/27433316/how-to-get-argparse-to-read-arguments-from-a-file-with-an-option-rather-than-pre)
# [Find default pip index-url](https://stackoverflow.com/questions/50100576/find-default-pip-index-url)
# [pyqt5 安装 sipbuild](https://www.cnblogs.com/hany-postq473111315/p/15402473.html)
# [[PyQt] Building PyQt from source with sip v5](https://www.riverbankcomputing.com/pipermail/pyqt/2019-October/042281.html)

# Issues:
# 1. PyStandMate.py --package parse
# 2. PyStandMate.py --bitness 32 --compiler MSVC --package PyQt6 --pip-index-url https://pypi.tuna.tsinghua.edu.cn/simple
#    > ModuleNotFoundError: No module named 'sipbuild'
# 3. PyStandMate.py --bitness 32 --compiler MSVC --package sip PyQt6 --pip-index-url https://pypi.tuna.tsinghua.edu.cn/simple
#    > ModuleNotFoundError: No module named 'pyqtbuild'
# 4. PyStandMate.py --bitness 32 --compiler MSVC --package sip pyqt-builder PyQt6 --pip-index-url https://pypi.tuna.tsinghua.edu.cn/simple