from typing import List, Dict, Union, Optional

import logging
import numpy as np
from pathlib import Path
from tqdm.auto import tqdm

import torch
from torch.nn import DataParallel
from torch.utils.data.sampler import SequentialSampler

from haystack.schema import Document
from haystack.document_stores import BaseDocumentStore
from haystack.nodes.retriever.base import BaseRetriever
from haystack.modeling.model.tokenization import Tokenizer
from haystack.modeling.model.language_model import LanguageModel
from haystack.modeling.model.biadaptive_hash_model import BiAdaptiveHashModel
from haystack.modeling.model.prediction_head import TextSimilarityHead, BinarySimilarityHead
from haystack.modeling.data_handler.processor import TextSimilarityProcessor
from haystack.modeling.data_handler.data_silo import DataSilo
from haystack.modeling.data_handler.dataloader import NamedDataLoader
from haystack.modeling.model.optimization import initialize_optimizer
from haystack.modeling.training.base import Trainer, EarlyStopping
from haystack.modeling.utils import initialize_device_settings
from copy import deepcopy
import math

logger = logging.getLogger(__name__)

class BinaryPassageRetriever(BaseRetriever):
    """
    Retriever that uses a bi-encoder (one transformer for query, one transformer for passage).
    See the original paper for more details:
    Karpukhin, Vladimir, et al. (2020): "Dense Passage Retrieval for Open-Domain Question Answering."
    (https://arxiv.org/abs/2004.04906).
    """

    def __init__(
        self,
        document_store: BaseDocumentStore,
        query_embedding_model: Union[Path, str] = "facebook/dpr-question_encoder-single-nq-base",
        passage_embedding_model: Union[Path, str] = "facebook/dpr-ctx_encoder-single-nq-base",
        model_version: Optional[str] = None,
        max_seq_len_query: int = 64,
        max_seq_len_passage: int = 256,
        top_k: int = 10,
        use_gpu: bool = True,
        batch_size: int = 16,
        embed_title: bool = True,
        use_fast_tokenizers: bool = True,
        infer_tokenizer_classes: bool = False,
        similarity_function: str = "dot_product",
        global_loss_buffer_size: int = 150000,
        progress_bar: bool = True,
        devices: Optional[List[Union[int, str, torch.device]]] = None,
        use_auth_token: Optional[Union[str, bool]] = None,
        hashnet_gamma: float=0.1,
        candidates: int=1000,
    ):
        """
        Init the Retriever incl. the two encoder models from a local or remote model checkpoint.
        The checkpoint format matches huggingface transformers' model format

        **Example:**

                ```python
                |    # remote model from FAIR
                |    DensePassageRetriever(document_store=your_doc_store,
                |                          query_embedding_model="facebook/dpr-question_encoder-single-nq-base",
                |                          passage_embedding_model="facebook/dpr-ctx_encoder-single-nq-base")
                |    # or from local path
                |    DensePassageRetriever(document_store=your_doc_store,
                |                          query_embedding_model="model_directory/question-encoder",
                |                          passage_embedding_model="model_directory/context-encoder")
                ```

        :param document_store: An instance of DocumentStore from which to retrieve documents.
        :param query_embedding_model: Local path or remote name of question encoder checkpoint. The format equals the
                                      one used by hugging-face transformers' modelhub models
                                      Currently available remote names: ``"facebook/dpr-question_encoder-single-nq-base"``
        :param passage_embedding_model: Local path or remote name of passage encoder checkpoint. The format equals the
                                        one used by hugging-face transformers' modelhub models
                                        Currently available remote names: ``"facebook/dpr-ctx_encoder-single-nq-base"``
        :param model_version: The version of model to use from the HuggingFace model hub. Can be tag name, branch name, or commit hash.
        :param max_seq_len_query: Longest length of each query sequence. Maximum number of tokens for the query text. Longer ones will be cut down."
        :param max_seq_len_passage: Longest length of each passage/context sequence. Maximum number of tokens for the passage text. Longer ones will be cut down."
        :param top_k: How many documents to return per query.
        :param use_gpu: Whether to use all available GPUs or the CPU. Falls back on CPU if no GPU is available.
        :param batch_size: Number of questions or passages to encode at once. In case of multiple gpus, this will be the total batch size.
        :param embed_title: Whether to concatenate title and passage to a text pair that is then used to create the embedding.
                            This is the approach used in the original paper and is likely to improve performance if your
                            titles contain meaningful information for retrieval (topic, entities etc.) .
                            The title is expected to be present in doc.meta["name"] and can be supplied in the documents
                            before writing them to the DocumentStore like this:
                            {"text": "my text", "meta": {"name": "my title"}}.
        :param use_fast_tokenizers: Whether to use fast Rust tokenizers
        :param infer_tokenizer_classes: Whether to infer tokenizer class from the model config / name.
                                        If `False`, the class always loads `DPRQuestionEncoderTokenizer` and `DPRContextEncoderTokenizer`.
        :param similarity_function: Which function to apply for calculating the similarity of query and passage embeddings during training.
                                    Options: `dot_product` (Default) or `cosine`
        :param global_loss_buffer_size: Buffer size for all_gather() in DDP.
                                        Increase if errors like "encoded data exceeds max_size ..." come up
        :param progress_bar: Whether to show a tqdm progress bar or not.
                             Can be helpful to disable in production deployments to keep the logs clean.
        :param devices: List of GPU devices to limit inference to certain GPUs and not use all available ones (e.g. ["cuda:0"]).
                        As multi-GPU training is currently not implemented for DPR, training will only use the first device provided in this list.
        :param use_auth_token:  API token used to download private models from Huggingface. If this parameter is set to `True`,
                                the local token will be used, which must be previously created via `transformer-cli login`.
                                Additional information can be found here https://huggingface.co/transformers/main_classes/model.html#transformers.PreTrainedModel.from_pretrained
        """
        super().__init__()

        if devices is not None:
            self.devices = devices
        else:
            self.devices, _ = initialize_device_settings(use_cuda=use_gpu, multi_gpu=True)

        if batch_size < len(self.devices):
            logger.warning("Batch size is less than the number of devices. All gpus will not be utilized.")

        self.document_store = document_store
        self.batch_size = batch_size
        self.progress_bar = progress_bar
        self.top_k = top_k
        self.hashnet_gamma = hashnet_gamma
        self.candidates = candidates

        if document_store is None:
            logger.warning(
                "DensePassageRetriever initialized without a document store. "
                "This is fine if you are performing DPR training. "
                "Otherwise, please provide a document store in the constructor."
            )

        self.infer_tokenizer_classes = infer_tokenizer_classes
        tokenizers_default_classes = {"query": "DPRQuestionEncoderTokenizer", "passage": "DPRContextEncoderTokenizer"}
        if self.infer_tokenizer_classes:
            tokenizers_default_classes["query"] = None  # type: ignore
            tokenizers_default_classes["passage"] = None  # type: ignore

        # Init & Load Encoders
        self.query_tokenizer = Tokenizer.load(
            pretrained_model_name_or_path=query_embedding_model,
            revision=model_version,
            do_lower_case=True,
            use_fast=use_fast_tokenizers,
            tokenizer_class=tokenizers_default_classes["query"],
            use_auth_token=use_auth_token,
        )
        self.query_encoder = LanguageModel.load(
            pretrained_model_name_or_path=query_embedding_model,
            revision=model_version,
            language_model_class="DPRQuestionEncoder",
            use_auth_token=use_auth_token,
        )
        self.passage_tokenizer = Tokenizer.load(
            pretrained_model_name_or_path=passage_embedding_model,
            revision=model_version,
            do_lower_case=True,
            use_fast=use_fast_tokenizers,
            tokenizer_class=tokenizers_default_classes["passage"],
            use_auth_token=use_auth_token,
        )
        self.passage_encoder = LanguageModel.load(
            pretrained_model_name_or_path=passage_embedding_model,
            revision=model_version,
            language_model_class="DPRContextEncoder",
            use_auth_token=use_auth_token,
        )

        self.processor = TextSimilarityProcessor(
            query_tokenizer=self.query_tokenizer,
            passage_tokenizer=self.passage_tokenizer,
            max_seq_len_passage=max_seq_len_passage,
            max_seq_len_query=max_seq_len_query,
            label_list=["hard_negative", "positive"],
            metric="text_similarity_metric",
            embed_title=embed_title,
            num_hard_negatives=0,
            num_positives=1,
        )
        prediction_head = BinarySimilarityHead(
            similarity_function=similarity_function, global_loss_buffer_size=global_loss_buffer_size
        )
        self.model = BiAdaptiveHashModel(
            language_model1=self.query_encoder,
            language_model2=self.passage_encoder,
            prediction_heads=[prediction_head],
            embeds_dropout_prob=0.1,
            lm1_output_types=["per_sequence"],
            lm2_output_types=["per_sequence"],
            device=str(self.devices[0]),
            hashnet_gamma=self.hashnet_gamma,
        )

        self.model.connect_heads_with_processor(self.processor.tasks, require_labels=False)

        if len(self.devices) > 1:
            self.model = DataParallel(self.model, device_ids=self.devices)

    def retrieve(
        self,
        query: str,
        filters: dict = None,
        top_k: Optional[int] = None,
        index: str = None,
        headers: Optional[Dict[str, str]] = None,
    ) -> List[Document]: # TODO - should change so it uses query binary and dense representation. Generation and re-ranking
        """
        Scan through documents in DocumentStore and return a small number documents
        that are most relevant to the query.

        :param query: The query
        :param filters: A dictionary where the keys specify a metadata field and the value is a list of accepted values for that field
        :param top_k: How many documents to return per query.
        :param index: The name of the index in the DocumentStore from which to retrieve documents
        """
        if top_k is None:
            top_k = self.top_k
        if not self.document_store:
            logger.error("Cannot perform retrieve() since DensePassageRetriever initialized with document_store=None")
            return []
        if index is None:
            index = self.document_store.index
        # query_embeddings, query_emb_binary = self.embed_queries(texts=[query])
        query_embeddings, query_emb_binary = self.embed_one_query(query=[{"query": query}])

        if self.candidates>len(self.document_store.all_documents):
            self.candidates=len(self.document_store.all_documents)

        bin_query_embeddings = np.packbits(np.where(query_emb_binary > 0, 1, 0), axis=-1).reshape(1, -1)

        raw_index = self.document_store.faiss_indexes[index]
        _, ids_arr = raw_index.search(bin_query_embeddings, self.candidates)

        passage_embeddings = np.vstack(
            [np.unpackbits(raw_index.reconstruct(int(id_))) for id_ in ids_arr.reshape(-1)]
        )

        passage_embeddings = passage_embeddings.reshape(
            query_embeddings.shape[0], passage_embeddings.shape[0], query_embeddings.shape[1]
        )
        passage_embeddings = passage_embeddings.astype(np.float32)

        passage_embeddings = passage_embeddings * 2 - 1 # Turns 1, 0 into 1, -1

        scores_arr = np.einsum("ijk,ik->ij", passage_embeddings, query_embeddings)
        sorted_indices = np.argsort(-scores_arr, axis=1)[0][:top_k]
        documents = [self.document_store.all_documents[ids_arr[0][i]] for i in sorted_indices]
        # documents = self.document_store.query_by_embedding(
        #     query_emb=query_emb_binary[0], top_k=top_k, filters=filters, index=index, headers=headers,
        #     return_embedding=False
        # )
        return documents

    def embed_one_query(self, query):
        dataset, tensor_names, _, baskets = self.processor.dataset_from_dicts(
            query, indices=[i for i in range(len(query))], return_baskets=True
        )

        data_loader = NamedDataLoader(
            dataset=dataset, sampler=SequentialSampler(dataset), batch_size=self.batch_size, tensor_names=tensor_names
        )

        # self.model.eval() # Changed to eval() method to avoid doing it every query

        for batch in data_loader:
            batch = {key: batch[key].to(self.devices[0]) for key in batch}
            with torch.no_grad():
                binary_query, binary_passage, dense_query = self.model.forward(**batch)[0]
                return dense_query.cpu().numpy(), binary_query.cpu().numpy()
        return None

    def _get_predictions(self, dicts):
        """
        Feed a preprocessed dataset to the model and get the actual predictions (forward pass + formatting).

        :param dicts: list of dictionaries
        examples:[{'query': "where is florida?"}, {'query': "who wrote lord of the rings?"}, ...]
                [{'passages': [{
                    "title": 'Big Little Lies (TV series)',
                    "text": 'series garnered several accolades. It received..',
                    "label": 'positive',
                    "external_id": '18768923'},
                    {"title": 'Framlingham Castle',
                    "text": 'Castle on the Hill "Castle on the Hill" is a song by English..',
                    "label": 'positive',
                    "external_id": '19930582'}, ...]
        :return: dictionary of embeddings for "passages" and "query"
        """
        dataset, tensor_names, _, baskets = self.processor.dataset_from_dicts(
            dicts, indices=[i for i in range(len(dicts))], return_baskets=True
        )

        data_loader = NamedDataLoader(
            dataset=dataset, sampler=SequentialSampler(dataset), batch_size=self.batch_size, tensor_names=tensor_names
        )
        all_embeddings = {"query": [], "passages": [], "dense_query": []}
        self.model.eval()

        # When running evaluations etc., we don't want a progress bar for every single query
        if len(dataset) == 1:
            disable_tqdm = True
        else:
            disable_tqdm = not self.progress_bar

        with tqdm(
            total=len(data_loader) * self.batch_size,
            unit=" Docs",
            desc=f"Create embeddings",
            position=0,
            leave=False,
            disable=disable_tqdm,
        ) as progress_bar:
            for batch in data_loader:
                batch = {key: batch[key].to(self.devices[0]) for key in batch}

                # get logits
                with torch.no_grad():
                    binary_query, binary_passage, dense_query = self.model.forward(**batch)[0]
                    if binary_query is not None:
                        all_embeddings["query"].append(binary_query.cpu().numpy())
                    if binary_passage is not None:
                        all_embeddings["passages"].append(binary_passage.cpu().numpy())
                    if dense_query is not None:
                        all_embeddings["dense_query"].append(dense_query.cpu().numpy())
                progress_bar.update(self.batch_size)

        if all_embeddings["passages"]:
            all_embeddings["passages"] = np.concatenate(all_embeddings["passages"])
        if all_embeddings["query"]:
            all_embeddings["query"] = np.concatenate(all_embeddings["query"])
        if all_embeddings["dense_query"]:
            all_embeddings["dense_query"] = np.concatenate(all_embeddings["dense_query"])
        return all_embeddings

    def embed_queries(self, texts: List[str]) -> List[np.ndarray]:
        """
        Create embeddings for a list of queries using the query encoder

        :param texts: Queries to embed
        :return: Embeddings, one per input queries
        """
        queries = [{"query": q} for q in texts]
        result = self._get_predictions(queries)
        return result["dense_query"], result["query"]

    def embed_documents(self, docs: List[Document]) -> List[np.ndarray]:
        """
        Create embeddings for a list of documents using the passage encoder

        :param docs: List of Document objects used to represent documents / passages in a standardized way within Haystack.
        :return: Embeddings of documents / passages shape (batch_size, embedding_dim)
        """
        if self.processor.num_hard_negatives != 0:
            logger.warning(
                f"'num_hard_negatives' is set to {self.processor.num_hard_negatives}, but inference does "
                f"not require any hard negatives. Setting num_hard_negatives to 0."
            )
            self.processor.num_hard_negatives = 0

        passages = [
            {
                "passages": [
                    {
                        "title": d.meta["name"] if d.meta and "name" in d.meta else "",
                        "text": d.content,
                        "label": d.meta["label"] if d.meta and "label" in d.meta else "positive",
                        "external_id": d.id,
                    }
                ]
            }
            for d in docs
        ]
        embeddings = self._get_predictions(passages)["passages"]

        return embeddings

    def train(
        self,
        data_dir: str,
        train_filename: str,
        dev_filename: str = None,
        test_filename: str = None,
        max_samples: int = None,
        max_processes: int = 128,
        multiprocessing_strategy: Optional[str] = None,
        early_stopping: Optional[EarlyStopping] = None,
        logging_wandb: bool=False,
        checkpoint_every: Optional[int] = None,
        checkpoint_root_dir: Optional[Path] = None,
        from_epoch: int = 0,
        from_step: int = 0,
        checkpoint_on_sigterm: bool = False,
        checkpoints_to_keep: int = 5,
        dev_split: float = 0,
        batch_size: int = 2,
        embed_title: bool = True,
        num_hard_negatives: int = 1,
        num_positives: int = 1,
        n_epochs: int = 3,
        evaluate_every: int = 1000,
        n_gpu: int = 1,
        learning_rate: float = 1e-5,
        epsilon: float = 1e-08,
        weight_decay: float = 0.0,
        num_warmup_steps: int = 100,
        grad_acc_steps: int = 1,
        use_amp: str = None,
        optimizer_name: str = "AdamW",
        optimizer_correct_bias: bool = True,
        save_dir: str = "../saved_models/dpr",
        query_encoder_save_dir: str = "query_encoder",
        passage_encoder_save_dir: str = "passage_encoder",
    ):
        """
        train a DensePassageRetrieval model
        :param data_dir: Directory where training file, dev file and test file are present
        :param train_filename: training filename
        :param dev_filename: development set filename, file to be used by model in eval step of training
        :param test_filename: test set filename, file to be used by model in test step after training
        :param max_samples: maximum number of input samples to convert. Can be used for debugging a smaller dataset.
        :param max_processes: the maximum number of processes to spawn in the multiprocessing.Pool used in DataSilo.
                              It can be set to 1 to disable the use of multiprocessing or make debugging easier.
        :param multiprocessing_strategy: Set the multiprocessing sharing strategy, this can be one of file_descriptor/file_system depending on your OS.
                                         If your system has low limits for the number of open file descriptors, and you can’t raise them,
                                         you should use the file_system strategy.
        :param dev_split: The proportion of the train set that will sliced. Only works if dev_filename is set to None
        :param batch_size: total number of samples in 1 batch of data
        :param embed_title: whether to concatenate passage title with each passage. The default setting in official DPR embeds passage title with the corresponding passage
        :param num_hard_negatives: number of hard negative passages(passages which are very similar(high score by BM25) to query but do not contain the answer
        :param num_positives: number of positive passages
        :param n_epochs: number of epochs to train the model on
        :param evaluate_every: number of training steps after evaluation is run
        :param n_gpu: number of gpus to train on
        :param learning_rate: learning rate of optimizer
        :param epsilon: epsilon parameter of optimizer
        :param weight_decay: weight decay parameter of optimizer
        :param grad_acc_steps: number of steps to accumulate gradient over before back-propagation is done
        :param use_amp: Whether to use automatic mixed precision (AMP) or not. The options are:
                    "O0" (FP32)
                    "O1" (Mixed Precision)
                    "O2" (Almost FP16)
                    "O3" (Pure FP16).
                    For more information, refer to: https://nvidia.github.io/apex/amp.html
        :param optimizer_name: what optimizer to use (default: AdamW)
        :param num_warmup_steps: number of warmup steps
        :param optimizer_correct_bias: Whether to correct bias in optimizer
        :param save_dir: directory where models are saved
        :param query_encoder_save_dir: directory inside save_dir where query_encoder model files are saved
        :param passage_encoder_save_dir: directory inside save_dir where passage_encoder model files are saved
        """
        self.processor.embed_title = embed_title
        self.processor.data_dir = Path(data_dir)
        self.processor.train_filename = train_filename
        self.processor.dev_filename = dev_filename
        self.processor.test_filename = test_filename
        self.processor.max_samples = max_samples
        self.processor.dev_split = dev_split
        self.processor.num_hard_negatives = num_hard_negatives
        self.processor.num_positives = num_positives

        if isinstance(self.model, DataParallel):
            self.model.module.connect_heads_with_processor(self.processor.tasks, require_labels=True)
        else:
            self.model.connect_heads_with_processor(self.processor.tasks, require_labels=True)

        data_silo = DataSilo(
            processor=self.processor,
            batch_size=batch_size,
            distributed=False,
            max_processes=max_processes,
            multiprocessing_strategy=multiprocessing_strategy,
        )

        # 5. Create an optimizer
        self.model, optimizer, lr_schedule = initialize_optimizer(
            model=self.model,
            learning_rate=learning_rate,
            optimizer_opts={
                "name": optimizer_name,
                "correct_bias": optimizer_correct_bias,
                "weight_decay": weight_decay,
                "eps": epsilon,
            },
            schedule_opts={"name": "LinearWarmup", "num_warmup_steps": num_warmup_steps},
            n_batches=len(data_silo.loaders["train"]),
            n_epochs=n_epochs,
            grad_acc_steps=grad_acc_steps,
            device=self.devices[0],  # Only use first device while multi-gpu training is not implemented
            use_amp=use_amp,
        )

        self.model.prediction_heads[0].n_passages = num_positives+num_hard_negatives

        # 6. Feed everything to the Trainer, which keeps care of growing our model and evaluates it from time to time
        trainer = Trainer(
            model=self.model,
            optimizer=optimizer,
            data_silo=data_silo,
            epochs=n_epochs,
            n_gpu=n_gpu,
            lr_schedule=lr_schedule,
            evaluate_every=evaluate_every,
            device=self.devices[0],  # Only use first device while multi-gpu training is not implemented
            use_amp=use_amp,
            early_stopping=early_stopping,
            logging_wandb=logging_wandb,
            eval_report=False,
            checkpoint_every=checkpoint_every,
            checkpoint_root_dir=checkpoint_root_dir,
            from_epoch=from_epoch,
            from_step=from_step,
            checkpoint_on_sigterm=checkpoint_on_sigterm,
            checkpoints_to_keep=checkpoints_to_keep,
            query_encoder_save_dir=query_encoder_save_dir,
            passage_encoder_save_dir=passage_encoder_save_dir,
            retriever=self
        )

        # 7. Let it grow! Watch the tracked metrics live on the public mlflow server: https://public-mlflow.deepset.ai
        trainer.train()

        self.model.save(Path(save_dir), lm1_name=query_encoder_save_dir, lm2_name=passage_encoder_save_dir)
        self.query_tokenizer.save_pretrained(f"{save_dir}/{query_encoder_save_dir}")
        self.passage_tokenizer.save_pretrained(f"{save_dir}/{passage_encoder_save_dir}")

        if len(self.devices) > 1 and not isinstance(self.model, DataParallel):
            self.model = DataParallel(self.model, device_ids=self.devices)

    def save(
        self,
        save_dir: Union[Path, str],
        query_encoder_dir: str = "query_encoder",
        passage_encoder_dir: str = "passage_encoder",
    ):
        """
        Save DensePassageRetriever to the specified directory.

        :param save_dir: Directory to save to.
        :param query_encoder_dir: Directory in save_dir that contains query encoder model.
        :param passage_encoder_dir: Directory in save_dir that contains passage encoder model.
        :return: None
        """
        save_dir = Path(save_dir)
        self.model.save(save_dir, lm1_name=query_encoder_dir, lm2_name=passage_encoder_dir)
        save_dir = str(save_dir)
        self.query_tokenizer.save_pretrained(save_dir + f"/{query_encoder_dir}")
        self.passage_tokenizer.save_pretrained(save_dir + f"/{passage_encoder_dir}")

    @classmethod
    def load(
        cls,
        load_dir: Union[Path, str],
        document_store: BaseDocumentStore,
        max_seq_len_query: int = 64,
        max_seq_len_passage: int = 256,
        use_gpu: bool = True,
        batch_size: int = 16,
        embed_title: bool = True,
        use_fast_tokenizers: bool = True,
        similarity_function: str = "dot_product",
        query_encoder_dir: str = "query_encoder",
        passage_encoder_dir: str = "passage_encoder",
        infer_tokenizer_classes: bool = False,
    ):
        """
        Load DensePassageRetriever from the specified directory.
        """
        load_dir = Path(load_dir)
        dpr = cls(
            document_store=document_store,
            query_embedding_model=Path(load_dir) / query_encoder_dir,
            passage_embedding_model=Path(load_dir) / passage_encoder_dir,
            max_seq_len_query=max_seq_len_query,
            max_seq_len_passage=max_seq_len_passage,
            use_gpu=use_gpu,
            batch_size=batch_size,
            embed_title=embed_title,
            use_fast_tokenizers=use_fast_tokenizers,
            similarity_function=similarity_function,
            infer_tokenizer_classes=infer_tokenizer_classes,
        )
        logger.info(f"DPR model loaded from {load_dir}")

        return dpr


"""
This module contains routines to create
a ST dataset and some statistics. The dataset
will contain several files with the ST data in different
formats
"""
import sys
import os
import numpy as np
from collections import defaultdict
import pandas as pd
from stpipeline.common.clustering import *
from stpipeline.common.unique_events_parser import parse_unique_events
import logging
import sys

def computeUniqueUMIs(transcripts, umi_counting_offset, umi_allowed_mismatches, group_umi_func):
    """
    Helper function to compute unique transcripts UMIs from
    a given list of transcripts. The function using an offset (genomic coordinates)
    where all UMIs will be grouped together by a grouping function and with a certain
    number of mismatches allowed (hamming distance)
    """
    # Sort transcripts by strand and start position
    sorted_transcripts = sorted(transcripts, key = lambda x: (x[5], x[1]))
    # Group transcripts by strand and start-position allowing an offset
    # And then performs the UMI clustering in each group to finally
    # compute the gene count as the sum of the unique UMIs for each group (strand,start,offset)
    grouped_transcripts = defaultdict(list)
    # TODO A probably better approach is to get the mean of all the start positions
    # and then make mean +- 300bp (user defined) a group to account for the library
    # size variability and then group the rest of transcripts normally by (strand, start, position).
    unique_transcripts = list()
    num_transcripts = len(transcripts)
    for i in range(num_transcripts - 1):
        current = sorted_transcripts[i]
        nextone = sorted_transcripts[i + 1]
        grouped_transcripts[current[6]].append(current)
        if abs(current[1] - nextone[1]) > umi_counting_offset or current[5] != nextone[5]:
            # A new group has been reached (strand, start-pos, offset)
            # Compute unique UMIs by hamming distance
            unique_umis = group_umi_func(list(grouped_transcripts.keys()), umi_allowed_mismatches)
            # Choose 1 random transcript for the clustered transcripts (by UMI)
            unique_transcripts += [random.choice(grouped_transcripts[u_umi]) for u_umi in unique_umis]
            # Reset the container
            grouped_transcripts = defaultdict(list)
    # We process the last one and more transcripts if they were not processed
    lastone = sorted_transcripts[num_transcripts - 1]
    grouped_transcripts[lastone[6]].append(lastone)
    unique_umis = group_umi_func(list(grouped_transcripts.keys()), umi_allowed_mismatches)
    unique_transcripts += [random.choice(grouped_transcripts[u_umi]) for u_umi in unique_umis]
    return unique_transcripts

def createDataset(input_file,
                  qa_stats,
                  gff_filename=None,
                  umi_cluster_algorithm="hierarchical",
                  umi_allowed_mismatches=1,
                  umi_counting_offset=250,
                  diable_umi=False,
                  output_folder=None,
                  output_template=None,
                  verbose=True):
    """
    The functions parses the reads in BAM format
    that have been annotated and demultiplexed (containing spatial barcode).
    It then groups them by gene-barcode to count reads accounting for duplicates
    using the UMIs (clustering them suing the strand and start position).
    It outputs the records in a matrix of counts in TSV format and BED format and it also
    writes out some statistics.
    :param input_file: the file with the annotated-demultiplexed records in BAM format
    :param qa_stats: the Stats object to add some stats (THIS IS PASSED BY REFERENCE)
    :param umi_cluster_algorithm: the clustering algorithm to cluster UMIs
    :param umi_allowed_mismatches: the number of miss matches allowed to remove
                                  duplicates by UMIs
    :param umi_counting_offset: the number of bases allowed as offset (start position) when counting UMIs
    :param diable_umi: when True the UMI filtering step will not be performed
    :param output_folder: path to place the output files
    :param output_template: the name of the dataset
    :param verbose: True if we can to collect the stats in the logger
    :type input_file: str
    :type umi_cluster_algorithm: str
    :type umi_allowed_mismatches: boolean
    :type umi_counting_offset: integer
    :type diable_umi: bool
    :type output_folder: str
    :type output_template: str
    :type verbose: bool
    :raises: RuntimeError,ValueError,OSError,CalledProcessError
    """
    logger = logging.getLogger("STPipeline")

    if not os.path.isfile(input_file):
        error = "Error creating dataset, input file not present {}\n".format(input_file)
        logger.error(error)
        raise RuntimeError(error)

    if output_template:
        filenameDataFrame = "{}_stdata.tsv".format(output_template)
        filenameReadsBED = "{}_reads.bed".format(output_template)
    else:
        filenameDataFrame = "stdata.tsv"
        filenameReadsBED = "reads.bed"

    # Some counters
    total_record = 0
    discarded_reads = 0

    # Obtain the clustering function
    if umi_cluster_algorithm == "naive":
        group_umi_func = countUMINaive
    elif umi_cluster_algorithm == "hierarchical":
        group_umi_func = countUMIHierarchical
    elif umi_cluster_algorithm == "Adjacent":
        group_umi_func = dedup_adj
    elif umi_cluster_algorithm == "AdjacentBi":
        group_umi_func = dedup_dir_adj
    elif umi_cluster_algorithm == "Affinity":
        group_umi_func = affinity_umi_removal
    else:
        error = "Error creating dataset.\n" \
        "Incorrect clustering algorithm {}".format(umi_cluster_algorithm)
        logger.error(error)
        raise RuntimeError(error)

    # Containers needed to create the data frame
    list_row_values = list()
    list_indexes = list()

    # Parse unique events to generate the unique counts and the BED file
    unique_events = parse_unique_events(input_file, gff_filename)
    with open(os.path.join(output_folder, filenameReadsBED), "w") as reads_handler:
        # this is the generator returning a dictionary with spots for each gene
        for gene, spots in unique_events:
            transcript_counts_by_spot = {}
            for spot_coordinates, reads in list(spots.items()):
                (x,y) = spot_coordinates
                # Re-compute the read count accounting for duplicates using the UMIs
                # Transcripts is the list of transcripts (chrom, start, end, clear_name, mapping_quality, strand, UMI)
                # First:
                # Get the original number of transcripts (reads)
                reads = list(reads)
                read_count = len(reads)
                if not diable_umi:
                    # Compute unique transcripts (based on UMI, strand and start position +- threshold)
                    unique_transcripts = computeUniqueUMIs(reads, umi_counting_offset,
                                                           umi_allowed_mismatches, group_umi_func)
                else:
                    unique_transcripts = reads
                # The new transcript count
                transcript_count = len(unique_transcripts)
                assert transcript_count > 0 and transcript_count <= read_count
                # Update the discarded reads count
                discarded_reads += (read_count - transcript_count)
                # Update read counts in the container (replace the list
                # of transcripts for a number so it can be exported as a data frame)
                transcript_counts_by_spot["{0}x{1}".format(x, y)] = transcript_count
                # Write every unique transcript to the BED output (adding spot coordinate and gene name)
                for read in unique_transcripts:
                    reads_handler.write("{0}\t{1}\t{2}\t{3}\t{4}\t{5}\t{6}\t{7}\t{8}\n".format(read[0],
                                                                                read[1],
                                                                                read[2],
                                                                                read[3],
                                                                                read[4],
                                                                                read[5],
                                                                                gene,
                                                                                x,y))
                # keep a counter of the number of unique events (spot - gene) processed
                total_record += 1

            # Add spot and dict [gene] -> count to containers
            list_indexes.append(gene)
            list_row_values.append(transcript_counts_by_spot)

    if total_record == 0:
        error = "Error creating dataset, input file did not contain any transcript\n"
        logger.error(error)
        raise RuntimeError(error)

    # Create the data frame
    counts_table = pd.DataFrame(list_row_values, index=list_indexes)
    counts_table.fillna(0, inplace=True)
    counts_table=counts_table.T # Transpose the dictionary to still get the spots as rows and genes as columns in the final tsv

    # Compute some statistics
    total_barcodes = len(counts_table.index)
    total_transcripts = np.sum(counts_table.values, dtype=np.int32)
    number_genes = len(counts_table.columns)
    aggregated_spot_counts = counts_table.sum(axis=1).values
    aggregated_gene_counts = (counts_table != 0).sum(axis=1).values
    max_genes_feature = aggregated_gene_counts.max()
    min_genes_feature = aggregated_gene_counts.min()
    max_reads_feature = aggregated_spot_counts.max()
    min_reads_feature = aggregated_spot_counts.min()
    average_reads_feature = np.mean(aggregated_spot_counts)
    average_genes_feature = np.mean(aggregated_gene_counts)
    std_reads_feature = np.std(aggregated_spot_counts)
    std_genes_feature = np.std(aggregated_gene_counts)

    # Print some statistics
    if verbose:
        logger.info("Number of reads present: {}".format(total_transcripts))
        logger.info("Number of unique events (gene-spot) present: {}".format(total_record))
        logger.info("Number of unique genes present: {}".format(number_genes))
        logger.info("Max number of genes over all spots: {}".format(max_genes_feature))
        logger.info("Min number of genes over all spots: {}".format(min_genes_feature))
        logger.info("Max number of reads over all spots: {}".format(max_reads_feature))
        logger.info("Min number of reads over all spots: {}".format(min_reads_feature))
        logger.info("Average number genes per spot: {}".format(average_genes_feature))
        logger.info("Average number reads per spot: {}".format(average_reads_feature))
        logger.info("Std. number genes per spot: {}".format(std_genes_feature))
        logger.info("Std. number reads per spot: {}".format(std_reads_feature))
        logger.info("Number of discarded reads (possible duplicates): {}".format(discarded_reads))

    # Update the QA object
    qa_stats.reads_after_duplicates_removal = int(total_transcripts)
    qa_stats.barcodes_found = total_barcodes
    qa_stats.genes_found = number_genes
    qa_stats.duplicates_found = discarded_reads
    qa_stats.max_genes_feature = max_genes_feature
    qa_stats.min_genes_feature = min_genes_feature
    qa_stats.max_reads_feature = max_reads_feature
    qa_stats.min_reads_feature = min_reads_feature
    qa_stats.average_gene_feature = average_genes_feature
    qa_stats.average_reads_feature = average_reads_feature

    # Write data frame to file
    counts_table.to_csv(os.path.join(output_folder, filenameDataFrame), sep="\t", na_rep=0)


# Copyright (c) 2019 Georgia Tech Robot Learning Lab
# Licensed under the MIT License.

import numpy as np
import tensorflow as tf
import rl.core.utils.tf_utils as U
from rl.core.online_learners import online_optimizer as OO
from rl.core.online_learners import base_algorithms as BA
from rl.core.function_approximators.policies import Policy

class rlOnlineOptimizer(object):
    pass

def _rlOnlineOptimizerDecorator(cls):
    """
    A decorator for reusing OO.OnlineOptimizers while making them use policy as the variable.
    This would create a new class that is both the base class in OO.OnlineOptimizer and rlOnlineOptimizer.
    """
    assert issubclass(cls, OO.OnlineOptimizer)
    class decorated_cls(cls, rlOnlineOptimizer):
        """ self._policy is only used as an interface variable for other rl
        codes, which can be modified on-the-fly, NOT the state of the optimzier
        """

        def __init__(self, policy, base_alg, p=0.0, **kwargs):
            assert isinstance(policy, Policy) and isinstance(base_alg, BA.BaseAlgorithm)
            super().__init__(base_alg, p=p, **kwargs)
            self._policy = policy

        def update(self, *args, **kwargs):
            super().update(*args, **kwargs)
            self._policy.variable = self.x

        @property
        def x(self):
            return super().x

        @x.setter
        def x(self, val):
            cls.x.fset(self, val)
            self._policy.variable=self.x

        def assign(self, policy):
            # NOTE this does not change the state of the optimizer
            self._policy.assign(policy)

    # to make them look the same as intended
    decorated_cls.__name__ = cls.__name__
    decorated_cls.__qualname__ = cls.__qualname__
    return decorated_cls

class rlPiccoloFisher(object):
    pass

def _rlPiccoloFisherDecorator(cls):
    """
    A class decorator for using Fisher information matrix (and KL divergence)
    as the regularization in a Piccolo class that has been decorated as an
    rlOnlineOptimizer.
    """
    assert issubclass(cls, OO.Piccolo)
    if not issubclass(cls, rlOnlineOptimizer):
        cls = _rlOnlineOptimizerDecorator(cls)

    class Reg(object):
        # regularization based on KL divergence between policies
        def __init__(self, refpol, varpol, default_damping=0.1):
            """
            refpol:
                reference point to compute gradient.
            varpol:
                variable policy, which has the variables to optimize over.
            """
            assert isinstance(refpol, Policy) and isinstance(varpol, Policy)
            self._damping0 = default_damping
            self.refpol = refpol
            self.varpol = varpol
            self.obs = None
            self.damping = None

        def kl(self, x):
            self.varpol.variable = x
            return self.refpol.kl(self.varpol, self.obs, reversesd=False)

        def fvp(self, g):
            return self.refpol.fvp(self.obs, g) + self.damping * g

        @property
        def std(self):
            return self.refpol.std if hasattr(self.refpol, 'std') else 1.0

        def assign(self, reg):
            assert type(self) == type(reg)
            self.refpol.assign(reg.refpol)
            self.obs = np.copy(reg.obs)
            self.damping = np.copy(reg.damping)
            self._damping0 = np.copy(reg._damping0)

        def update(self, obs):
            self.refpol.assign(self.varpol)
            self.obs = np.copy(obs)
            self.damping = self._damping0  # /np.mean(self.std**2.0)

        @property
        def initialized(self):
            return not self.obs is None

    # We use Fisher information matrix and KL divergence to define the mvp (and
    # dist_fun) used in SecondOrderUpdate
    class decorated_cls(cls, rlPiccoloFisher):

        def __init__(self, policy, base_alg, p=0.0, use_shift=True, damping=0.1, **kwargs):
            assert isinstance(base_alg, BA.SecondOrderUpdate)
            super().__init__(policy, base_alg, p=p, **kwargs)
            # It uses reg_new and reg_old to define the regularization
            # reg_old is intended for memory
            self._reg_old = Reg(self._policy.copy('reg_pol_old'), self._policy, default_damping=damping)
            self._reg_new = Reg(self._policy.copy('reg_pol_new'), self._policy, default_damping=damping)
            self._reg_swp = Reg(self._policy.copy('reg_pol_swp'), self._policy, default_damping=damping)

            # Whether to overide the Fisher with the on-policy one in the
            # Prediction Step. It is not ideal in theory but improves numerical
            # stability.
            self._use_shift = use_shift

        @property
        def _fvp(self):
            return lambda g: (self._reg_new.fvp(g) + self._reg_old.fvp(g)) / 2.0

        @property
        def _kl(self):
            return self._reg_new.kl

        # overwrite _predict and _correct to provde the required mvp and dist_fun
        def _predict(self, g_hat, ro=None, **kwargs):
            obs = self._reg_new.obs if ro is None else ro.obs
            self._reg_swp.update(obs)  # save reg info in swp.
            if self._use_shift:  # may improve numericacl stability but break the theory
                self._reg_old.update(obs)
                self._reg_new.update(obs)
            else:  # use the previous regs, except for initialization
                if not self._reg_old.initialized or not self._reg_new.initialized:
                    self._reg_old.update(obs)
                    self._reg_new.update(obs)

            return super()._predict(g_hat, mvp=self._fvp, **kwargs)

        def _correct(self, g, ro=None, **kwargs):
            # **kwargs are not used
            assert ro is not None
            # update Fisher information matrices
            self._reg_new.update(ro.obs)
            if not self._reg_swp.initialized:  # when no predict has been called.
                self._reg_old.update(ro.obs)
            else:
                self._reg_old.assign(self._reg_swp)

            # call the usual piccolo update
            if isinstance(self._base_alg, BA.TrustRegionSecondOrderUpdate):
                return super()._correct(g, mvp=self._fvp, dist_fun=self._kl, **kwargs)
            elif isinstance(self._base_alg, BA.AdaptiveSecondOrderUpdate):
                return super()._correct(g, mvp=self._fvp, **kwargs)

    # to make them look the same as intended
    decorated_cls.__name__ = cls.__name__
    decorated_cls.__qualname__ = cls.__qualname__
    return decorated_cls

def _rlPiccoloOptDecorator(cls):
    assert issubclass(cls, OO.PiccoloOpt)
    if not issubclass(cls, rlOnlineOptimizer):
        cls = _rlOnlineOptimizerDecorator(cls)

    class decorated_cls(cls):
        def _correct(self, g_hat, grad_hat=None, loss_hat=None, callback=None,
                     warm_start=True, stop_std_grad=False, **kwargs):
            return super()._correct(g_hat, grad_hat=grad_hat, loss_hat=loss_hat, callback=callback, warm_start=warm_start, **kwargs)

        def _predict(self, g_hat, grad_hat=None, loss_hat=None, callback=None,
                     warm_start=True, stop_std_grad=False, **kwargs):
            # grad_hat, loss_hat, callback are functions that take no arguments
            # but inheritently depend on the policy (i.e. self._policy).
            assert (grad_hat is not None) and (loss_hat is not None)

            # prevent std from degenerating during multi-step updates
            if stop_std_grad:
                self._policy.stop_std_grad(True)

                def set_policy(x):
                    std = self._policy.std
                    self._policy.variable = x
                    self._policy.std = std

                def get_projection(x0):
                    # constraint the std as given by x0
                    self._policy.variable = x0
                    std0 = np.copy(self._policy.std)

                    def projection(x):
                        self._policy.variable = x
                        self._policy.std = std0
                        self._policy.stop_std_grad(False)  # turn it back on
                        return self._policy.variable
                    return projection
            else:
                def set_policy(x):
                    self._policy.variable = x
                get_projection = None

            # define the problem
            def _grad_hat(x):
                set_policy(x)
                return grad_hat()

            def _loss_hat(x):
                set_policy(x)
                return loss_hat()
            if callback is not None:
                def _callback(x):
                    self._policy.variable = x
                    return callback()
            else:
                _callback = None
            # call the core library
            g_hat = super()._predict(g_hat, grad_hat=_grad_hat, loss_hat=_loss_hat,
                                     callback=_callback, warm_start=warm_start,
                                     get_projection=get_projection, **kwargs)
            self._policy.stop_std_grad(False)  # turn it back on
            return g_hat

    # to make them look the same as intended
    decorated_cls.__name__ = cls.__name__
    decorated_cls.__qualname__ = cls.__qualname__
    return decorated_cls

# =====================================================================================
# Here we define all rl-oriented classes for policy optimization
@_rlOnlineOptimizerDecorator
class BasicOnlineOptimizer(OO.BasicOnlineOptimizer):
    pass

@_rlOnlineOptimizerDecorator
class Piccolo(OO.Piccolo):
    pass

@_rlPiccoloOptDecorator
@_rlOnlineOptimizerDecorator
class PiccoloOptBasic(OO.PiccoloOptBasic):
    pass

@_rlPiccoloFisherDecorator
class PiccoloFisherReg(Piccolo):
    pass

@_rlPiccoloFisherDecorator
class PiccoloOptBasicFisherReg(PiccoloOptBasic):
    pass


from mmcv import visualization
from numpy.lib.arraysetops import isin
import torch
from torch._C import DeviceObjType
import torch.nn as nn
import numpy as np

from mmdet.models import DETECTORS
from .depth_net import PackNetSlim01
from .sf_net import SFNet
from .lidar_loss.temporal_spatial_loss import warp_ref_image_temporal, PhotometricLoss, \
    calc_scene_flow_consistency_loss, warp_ref_image_spatial, calc_depth_consistency_loss, \
    calc_invdepth_consistency_loss, calc_stereo_rgb_loss

from .utils import get_depth_metrics, get_motion_metrics, remap_invdepth_color,\
    get_smooth_loss, group_smoothness,  sparsity_loss, get_motion_metrics, \
    flow2rgb

def scale_intrinsics(K, x_scale, y_scale):
    """
    Scale intrinsics given x_scale and y_scale factors
    x_scale = x_tgt_size / x_prev_size

    """
    K[..., 0, 0] *= x_scale
    K[..., 1, 1] *= y_scale
    K[..., 0, 2] = (K[..., 0, 2] + 0.5) * x_scale - 0.5
    K[..., 1, 2] = (K[..., 1, 2] + 0.5) * y_scale - 0.5
    return K

@DETECTORS.register_module()
class SpatialTempNet(nn.Module):
    _left_swap = [i for i in range(-1,5)]
    _right_swap = [i % 6 for i in range(1,7)]

    def __init__(self, depth_net_cfg, sf_net_cfg=None,
                scale_depth=False,
                scale_depth_for_temp=True,
                depth_supervision_ratio=1.0,
                depth_smoothing=1e-2,
                motion_smoothing=1e-3,
                motion_sparse=1e-3,
                sf_consis=1.0,
                depth_consis=1.0,
                rgb_consis=1.0,
                stereo_rgb_consis=1.0,
                loss_decay=1.0,
                **kwargs):

        super(SpatialTempNet, self).__init__()
        self.depth_net = PackNetSlim01(**depth_net_cfg)
        self.sf_net = SFNet()
        self.photometric_loss = PhotometricLoss(alpha=0.1, clip_loss=0.5, C1=1e-4, C2=9e-4)
        self.scale_depth = scale_depth
        self.scale_depth_for_temp = scale_depth_for_temp
        # < 0 means no supervised depth loss
        self.depth_supervision_ratio= depth_supervision_ratio
        self.depth_smoothing = depth_smoothing
        self.motion_smoothing = motion_smoothing
        self.motion_sparse = motion_sparse
        self.sf_consis = sf_consis
        self.depth_consis = depth_consis
        self.rgb_consis = rgb_consis
        self.stereo_rgb_consis = stereo_rgb_consis
        self.loss_decay = loss_decay

    def forward_train(self, data, **kwargs):
        # prev: 0-5; now: 6-11; next: 12:17
        # img0, img1, img2, ... img17
        img_keys = ['img{}'.format(i) for i in range(18)]

        now_depth_keys = ['depth_map{}'.format(i) for i in range(6, 12)]
        now_sf_keys = ['sf_map{}'.format(i) for i in range(6, 12)]

        # Imgs!
        # shape [18, N, 3, H, W]
        imgs = [data[tmp_key] for tmp_key in img_keys]

        prev_imgs, now_imgs, next_imgs = imgs[0:6], imgs[6:12], imgs[12:]
        # shape [N*6, 3, H, W]
        prev_imgs = torch.cat(prev_imgs, dim=0)
        now_imgs = torch.cat(now_imgs, dim=0)
        next_imgs = torch.cat(next_imgs, dim=0)

        # Depths GT!
        # shape [N*6, H, W]
        now_depth = torch.cat([data[tmp_key] for tmp_key in now_depth_keys], dim=0)
        # shape [N*6, 1, H, W]
        now_depth = now_depth.unsqueeze(dim=1)

        # SceneFlow GT!
        # shape [N*6, H, W, 4]
        now_sf = torch.cat([data[tmp_key] for tmp_key in now_sf_keys], dim=0)
        # shape [N*6, 4, H, W]
        #from IPython import embed
        #embed()
        now_sf = now_sf.permute(0, 3, 1, 2)

        # Depth Pred!
        # shape [N*6, 1, H, W]

        all_imgs = torch.cat([prev_imgs, now_imgs, next_imgs], dim=0)
        all_disp_preds = self.depth_net.get_pred(all_imgs)

        # SceneFlow Pred!
        # shape [N*6, 3, H, W]
        # now2next_sf_pred = self.sf_net(now_imgs, next_imgs)
        # now2prev_sf_pred = self.sf_net(now_imgs, prev_imgs)

        # [N, 18, 3, 3]
        camera_intrinsic = data['cam_intrinsic']
        #img_x, img_y = now_imgs.size(-2), now_imgs.size(-1)
        img_x, img_y = now_imgs.size(-1), now_imgs.size(-2)
        #print(now_imgs.shape)
        scale_x = img_x / 1600
        scale_y = img_y / 900
        camera_intrinsic = scale_intrinsics(camera_intrinsic, scale_x, scale_y)
        prev_cam_intrin = camera_intrinsic[:, 0:6, ...].view(-1, 3, 3)
        now_cam_intrin = camera_intrinsic[:, 6:12, ...].view(-1, 3, 3)
        next_cam_intrin = camera_intrinsic[:, 12:18, ...].view(-1, 3, 3)
        # [N, 18, 4, 4]
        cam_pose = data['cam_pose']

        # [N, 6, 4, 4] => [N*6, 4, 4]
        prev_cam_pose = cam_pose[:, 0:6, ...].view(-1, 4, 4)
        now_cam_pose = cam_pose[:, 6:12, ...].view(-1, 4, 4)
        next_cam_pose = cam_pose[:, 12:18, ...].view(-1, 4, 4)

        mask = now_depth > 0
        sparsity = torch.sum(mask) * 1.0 / torch.numel(mask)

        B, _, H, W = now_imgs.shape
        loss_list = []
        log_vars_list = []
        now2next_sf_pred = None
        now2prev_sf_pred = None
        for pred_idx, disp_raw in enumerate(all_disp_preds):
            #from IPython import embed
            #embed()
            # note: following loss can only be computed with batch size as one
            # now_inv_depth: [N*6, 1, H, W]
            prev_disp, now_disp, next_disp = torch.split(disp_raw, 6, dim=0)
            _, _, Hp, Wp = now_disp.shape
            if Hp != H:
                now_inv_depth = nn.functional.interpolate(now_disp, size=[H, W])
                # print(Hp, Wp, H, W)
                prev_disp = nn.functional.interpolate(prev_disp, size=[H, W])
                next_disp = nn.functional.interpolate(next_disp, size=[H, W])
            else:
                now_inv_depth = now_disp

            depth_pred = 1.0 / torch.clamp(now_inv_depth, min=1e-6)
            prev_depth_pred = 1.0 / torch.clamp(prev_disp, min=1e-6)
            next_depth_pred = 1.0 / torch.clamp(next_disp, min=1e-6)
            if self.scale_depth_for_temp:
                depth_scale = torch.median(now_depth[mask]) / (torch.median(depth_pred[mask]) + 1e-4)
                depth_scale = depth_scale.detach()
            else:
                depth_scale = 1.0

            # scaling the depth prediction!
            if now2next_sf_pred is None:
                prev_imgs_depth = torch.cat([prev_imgs, prev_depth_pred], dim=1)
                next_imgs_depth = torch.cat([next_imgs, next_depth_pred], dim=1)

                rec_prev_imgs_no_sf = warp_ref_image_temporal(now_inv_depth / depth_scale, prev_imgs_depth,
                                                now_cam_intrin, prev_cam_intrin,
                                                now_cam_pose, prev_cam_pose,
                                                #now2prev_sf_pred)
                                                None)
                                                #
                rec_next_imgs_no_sf = warp_ref_image_temporal(now_inv_depth / depth_scale, next_imgs_depth,
                                                now_cam_intrin, next_cam_intrin,
                                                now_cam_pose, next_cam_pose,
                                                #now2next_sf_pred)
                                                None)
                                                #now2next_sf_pred)

                now_imgs_depth = torch.cat([now_imgs, depth_pred.detach()], dim=1)
                now2next_inp = torch.cat([now_imgs_depth.detach(), rec_next_imgs_no_sf.detach(), next_imgs_depth], dim=1)
                now2prev_inp = torch.cat([prev_imgs_depth.detach(), rec_prev_imgs_no_sf.detach(), prev_imgs_depth], dim=1)
                now2next_sf_pred = self.sf_net(now2next_inp)
                now2prev_sf_pred = self.sf_net(now2prev_inp)

            rec_prev_imgs = warp_ref_image_temporal(now_inv_depth.detach() / depth_scale, prev_imgs,
                                            now_cam_intrin, prev_cam_intrin,
                                            now_cam_pose, prev_cam_pose,
                                            now2prev_sf_pred)
                                            #None)
                                            #
            rec_next_imgs = warp_ref_image_temporal(now_inv_depth.detach() / depth_scale, next_imgs,
                                            now_cam_intrin, next_cam_intrin,
                                            now_cam_pose, next_cam_pose,
                                            now2next_sf_pred)
                                            #None)
                                            #now2next_sf_pred)
            #rec_left_imgs = warp_ref_image_temporal(now_inv_depth, now_imgs)

            # [N, 1, H, W]

            #temp_rec_loss = self.photometric_loss([rec_prev_imgs], [now_imgs])[0] + \
            #                self.photometric_loss([rec_next_imgs], [now_imgs])[0]
            #temp_rec_loss = temp_rec_loss.mean()

            temp_rec_loss_maps = [torch.mean(torch.abs(rec_prev_imgs - now_imgs), dim=1, keepdim=True),
                                torch.mean(torch.abs(rec_next_imgs - now_imgs), dim=1, keepdim=True),
                                torch.mean(torch.abs(prev_imgs - now_imgs), dim=1, keepdim=True),
                                torch.mean(torch.abs(next_imgs - now_imgs), dim=1, keepdim=True), ]
            temp_rec_loss_map = torch.cat(temp_rec_loss_maps, dim=1)
            temp_rec_loss_map, _ = torch.min(temp_rec_loss_map, dim=1)

            temp_rec_loss = torch.mean(temp_rec_loss_map)

            stereo_left_loss, left_rec_imgs = calc_stereo_rgb_loss(now_inv_depth, now_imgs[self._left_swap, ...], now_imgs,
                                                now_cam_intrin, now_cam_intrin[self._left_swap, ...],
                                                now_cam_pose, now_cam_pose[self._left_swap, ], )
            stereo_right_loss, right_rec_imgs = calc_stereo_rgb_loss(now_inv_depth, now_imgs[self._right_swap, ...], now_imgs,
                                                now_cam_intrin, now_cam_intrin[self._right_swap, ...],
                                                now_cam_pose, now_cam_pose[self._right_swap, ], )

            stereo_rgb_loss = stereo_left_loss + stereo_right_loss

            consis_sf_loss, sf_valid_mask_ratio = calc_scene_flow_consistency_loss([now2prev_sf_pred, now2next_sf_pred], now_cam_intrin, now_cam_pose, now_inv_depth)

            # scale the depth.  depth-consis tend to make depth prediction smaller
            consis_dep_loss, valid_mask_ratio = calc_depth_consistency_loss(depth_pred, now_cam_intrin, now_cam_pose)

            # TODO: scale or not?
            depth_smoothing_loss = get_smooth_loss(depth_pred, now_imgs)

            motion_smoothing_loss = group_smoothness(now2prev_sf_pred) + \
                                    group_smoothness(now2next_sf_pred)
            motion_sparse_loss = sparsity_loss(now2prev_sf_pred) + \
                                 sparsity_loss(now2next_sf_pred)

            #from IPython import embed
            #embed()
            loss = self.rgb_consis * temp_rec_loss
            loss = loss + self.stereo_rgb_consis * stereo_rgb_loss
            loss = loss + self.sf_consis * consis_sf_loss
            loss = loss + self.depth_consis * consis_dep_loss
            loss = loss + self.depth_smoothing * depth_smoothing_loss
            loss = loss + self.motion_smoothing * motion_smoothing_loss
            loss = loss + self.motion_sparse * motion_sparse_loss
            if self.depth_supervision_ratio > 0:

                # TODO: smooth L1;  imbalance L1
                depth_sup_loss = torch.abs((now_depth - depth_pred)) * mask
                depth_sup_loss = torch.sum(depth_sup_loss) / torch.sum(mask)

                loss = loss + self.depth_supervision_ratio * depth_sup_loss

            with torch.no_grad():
                metrics = get_depth_metrics(depth_pred, now_depth, mask, scale=self.scale_depth)
                # abs_diff, abs_rel, sq_rel, rmse, rmse_log
                metrics = [m.item() for m in metrics]
                # depth_scale = median(gt_depth) / median(pred_depth)
                abs_diff, abs_rel, sq_rel, rmse, rmse_log, depth_scale = metrics

                # add visualization map
                if pred_idx == 0:
                    #std = torch.tensor([58.395, 57.12, 57.375]).cuda().view(1, -1, 1, 1)
                    #mean = torch.tensor([123.675, 116.28, 103.53]).cuda().view(1, -1, 1, 1)
                    #img = now_imgs * std + mean
                    #img = img / 255.0
                    img = now_imgs

                    # scale the depth for visualization
                    inv_depth_pred_img0 = (now_inv_depth[0]).clamp(min=1e-5)

                    cmap_inv_depth_pred_img0 = remap_invdepth_color(inv_depth_pred_img0)

                    visualization_map = {
                        'inv_depth_pred': cmap_inv_depth_pred_img0.transpose(2,0,1),
                        'img': img[0].detach().cpu().numpy(),
                        'prev_img': prev_imgs[0].detach().cpu().numpy(),
                        'next_img': next_imgs[0].detach().cpu().numpy(),
                        'left_img': now_imgs[self._left_swap, ][0].detach().cpu().numpy(),
                        'right_img': now_imgs[self._right_swap][0].detach().cpu().numpy(),
                        'rec_from_prev': rec_prev_imgs[0].detach().cpu().numpy(),
                        'rec_from_next': rec_next_imgs[0].detach().cpu().numpy(),
                        'rec_from_left': left_rec_imgs[0].detach().cpu().numpy(),
                        'rec_from_right': right_rec_imgs[0].detach().cpu().numpy(),
                    }

                log_vars = {'loss': loss.item(),
                            'temp_rec_loss': temp_rec_loss.item(),
                            'stereo_rgb_loss': stereo_rgb_loss.item(),
                            'depth_consis_loss': consis_dep_loss.item(),
                            'sceneflow_consis_loss': consis_sf_loss.item(),
                            'depth_smoothing': depth_smoothing_loss.item(),
                            'motion_smoothing': motion_smoothing_loss.item(),
                            'motion_sparse_loss': motion_sparse_loss.item(),
                            'valid_mask_ratio': valid_mask_ratio,
                            'sparsity': sparsity.item(),
                            'abs_diff': abs_diff, 'abs_rel': abs_rel,
                            'sq_rel': sq_rel, 'rmse': rmse,
                            'rmse_log': rmse_log, 'depth_scale': depth_scale
                            }
            #print(log_vars, '---', pred_idx)
            if self.depth_supervision_ratio > 0:
                log_vars['depth_sup_loss'] = depth_sup_loss.item()

            loss_list.append(loss)
            log_vars_list.append(log_vars)

        #print(log_vars)
        loss_ratio = 1.0
        final_loss = 0
        for temp_loss in loss_list:
            final_loss = final_loss + temp_loss * loss_ratio
            loss_ratio = loss_ratio * self.loss_decay

        final_log_vars = {}
        for log_vars in log_vars_list:
            for name, value in log_vars.items():
                if name not in final_log_vars:
                    final_log_vars[name] = [value]
                else:
                    final_log_vars[name].append(value)

        for name, value_list in final_log_vars.items():
            final_log_vars[name] = sum(value_list) / len(value_list)
            #print('mean-', len(value_list))

        with torch.no_grad():
            fused_motion = (now2next_sf_pred - now2prev_sf_pred) / 2
            motion_metrics = get_motion_metrics(fused_motion,  # time_delat: 0.5s / 0.1s for key/non-key
                                                now_sf/10.0, scale=self.scale_depth)
            motion_metrics = [m.item() for m in motion_metrics]
            epe, epe_rel, motion_scale = motion_metrics
            final_log_vars['epe'] = epe
            final_log_vars['epe_rel'] = epe_rel
            final_log_vars['motion_scale'] = motion_scale

            visual_motion = fused_motion[0]
            visual_motion = flow2rgb(visual_motion.detach().cpu().numpy().transpose(1,2,0)).transpose(2, 0, 1)
            cat_a, cat_b = visualization_map['img'], visualization_map['inv_depth_pred']
            pred_visual = np.concatenate([cat_a, cat_b, visual_motion], axis=-1)

            time_warp_visual = np.concatenate([cat_a,
                                visualization_map['rec_from_prev'],
                                visualization_map['rec_from_next']], axis=-1)
            spatial_warp_visual = np.concatenate([visualization_map['rec_from_left'],
                                                cat_a,
                                                visualization_map['rec_from_right'], ], axis=-1)
            image_prev_now_next = np.concatenate([visualization_map['prev_img'],
                                        cat_a,
                                        visualization_map['prev_img'],], axis=-1)
            image_left_now_right = np.concatenate([visualization_map['left_img'],
                                        cat_a,
                                        visualization_map['right_img'],], axis=-1)
            final_visualization_map = {'img-depth-motion': pred_visual,
                                        'time_warp': time_warp_visual,
                                        'spatial_warp': spatial_warp_visual,
                                        'image_time': image_prev_now_next,
                                        'image_sptial': image_left_now_right, }

        return final_loss, final_log_vars, final_visualization_map

    def forward(self, return_loss=True, rescale=False, **data):
        if not return_loss:
            loss, log_vars, visualization_map = self.forward_train(data)
            ret_keys = ['abs_diff', 'abs_rel',
                        'sq_rel', 'rmse',
                        'rmse_log','loss']

            values = [log_vars[t_key] for t_key in ret_keys]
            return [values]

        return self.forward_train(data)

    def train_step(self, data, optimzier):

        loss, log_vars, visualization_map = self.forward_train(data)

        # 'pred', 'data', 'label', 'depth_at_gt' is used for visualization only!
        outputs = {'loss':loss, 'log_vars':log_vars,
                    'num_samples':data['img1'].size(0)}

        for visual_name, visual_map in visualization_map.items():
            outputs[visual_name] = visual_map
        #print('output', outputs)
        return outputs

    def val_step(self, data, optimizer):

        return self.train_step(self, data, optimizer)

    def eval_forward(self, data):
        img_keys = ['img{}'.format(i) for i in range(18)]

        now_depth_keys = ['depth_map{}'.format(i) for i in range(6, 12)]
        now_sf_keys = ['sf_map{}'.format(i) for i in range(6, 12)]

        # Imgs!
        # shape [18, N, 3, H, W]
        imgs = [data[tmp_key] for tmp_key in img_keys]

        prev_imgs, now_imgs, next_imgs = imgs[0:6], imgs[6:12], imgs[12:]
        # shape [N*6, 3, H, W]
        prev_imgs = torch.cat(prev_imgs, dim=0)
        now_imgs = torch.cat(now_imgs, dim=0)
        next_imgs = torch.cat(next_imgs, dim=0)

        now_inv_depth_preds = self.depth_net.get_pred(now_imgs)

        camera_intrinsic = data['cam_intrinsic']
        #img_x, img_y = now_imgs.size(-2), now_imgs.size(-1)
        img_x, img_y = now_imgs.size(-1), now_imgs.size(-2)
        #print(now_imgs.shape)
        scale_x = img_x / 1600
        scale_y = img_y / 900
        camera_intrinsic = scale_intrinsics(camera_intrinsic, scale_x, scale_y)
        now_cam_intrin = camera_intrinsic[:, 6:12, ...].view(-1, 3, 3)
        # [N, 18, 4, 4]
        cam_pose = data['cam_pose']
        # [N, 6, 4, 4] => [N*6, 4, 4]
        now_cam_pose = cam_pose[:, 6:12, ...].view(-1, 4, 4)

        outputs = {'inv_depth_preds': now_inv_depth_preds,
                    'cam_intrinsic': now_cam_intrin,
                    'cam_pose': now_cam_pose,
                    'imgs': now_imgs}
        return outputs


from __future__ import print_function
import os
import json
import pickle
import numpy as np
import utils
import torch
import random
import torchvision.models as models
from torch.utils.data import Dataset
from torchvision import datasets, transforms

class FeatureDataset(Dataset):
    def __init__(self,
                 name,
                 dataset_path='./data'):
        super(FeatureDataset, self).__init__()
        self.keys = ['vital', 'demo', 'gt']
        self.set_features(name, dataset_path)
        self.tensorize()

    def set_features(self, name, dataset_path):
        print('loading features from files')
        self.data = {}
        for key in self.keys:
            self.data[key] = np.load(os.path.join(dataset_path, name+'_'+key+'.npy'))
            if key == 'gt':
                self.data[key] = np.expand_dims(self.data[key], axis=-1)

    def tensorize(self):
        self.tensor = {}
        for key in self.keys:
            self.tensor[key] = torch.from_numpy(self.data[key])

    def __getitem__(self, index):
        vital = self.tensor['vital'][index]
        #demo = self.tensor['demo'][index]
        gt = self.tensor['gt'][index]
        return vital, gt

    def __len__(self):
        return len(self.data['gt'])


# Train for prediction of both semantic maps and instances for the content of a given vessel mask in an image. Should run out of the box with the example set.
#...............................Imports..................................................................
import os
import numpy as np
import FCN_NetModel as NET_FCN # The net Class
import torch
import Reader as LabPicsReader
from scipy.optimize import linear_sum_assignment
import json
import Visuallization as vis
import ClassesGroups
##################################Input paramaters#########################################################################################
#.................................Main Input folder...........................................................................................
ChemLabPicsDir="ExampleTrain/" # Input folder for labpic chemistry training
MedLabPicsDir="ExampleTrain/" # Input folder for labpics medical  training
#......................................Main input parameters..................................................................................

MinSize=250 # min image height width
MaxSize=1200 #  Maxn image height widt
MaxPixels=800*800*3# max number of pixels in a batch reduce to solve out of memoru problems
MinMaskSize=1000 # Min size of vessel mask in pixels, smaller vessels will be ignored
TrainingMode=True # Train or test mode
IgnoreParts=True # Dont train on vessel parts
IgnoreSurfacePhase=False # Don train on materials that are just stuck on the glassware surface (and dont cover volume)
IgnoreScattered=True # Ignore material phases which are scattered in  sparse droplet/particles
MaxBatchSize=6# max images in class
Learning_Rate=1e-5 # initial learning rate
Weight_Decay=1e-5# Weight for the weight decay loss function
TrainedModelWeightDir="logs/" # Folder where trained model weight and information will be stored"
TrainLossTxtFile=TrainedModelWeightDir+"TrainLoss.txt" #Where train losses will be writen
MAX_ITERATION = int(100000010) # Max  number of training iteration
#********************************Classes to use in training the semantic segmentation the net will only produce maps for this clases************************************************************************************************************
MaterialClasses=ClassesGroups.MaterialClass
ClassToUse=ClassesGroups.VesselContentClasses
#****************************************************************************************************************************************************

if not os.path.exists(TrainedModelWeightDir):
    os.mkdir(TrainedModelWeightDir)
Trained_model_path="" # Path of trained model weights If you want to return to trained model, else should be =""

#=========================Load net weights and parameters from previous runs if exist====================================================================================================================
InitStep=1
if os.path.exists(TrainedModelWeightDir + "/Defult.torch"):
    Trained_model_path=TrainedModelWeightDir + "/Defult.torch"
if os.path.exists(TrainedModelWeightDir+"/Learning_Rate.npy"):
    Learning_Rate=np.load(TrainedModelWeightDir+"/Learning_Rate.npy")
if os.path.exists(TrainedModelWeightDir+"/itr.npy"): InitStep=int(np.load(TrainedModelWeightDir+"/itr.npy"))

#---------------------Create and Initiate net and create optimizer------------------------------------------------------------------------------------
Net=NET_FCN.Net(ClassList=ClassToUse) # Create neural net

if Trained_model_path!="": # Optional initiate full net by loading a previosly trained net weights if exist
    Net.load_state_dict(torch.load(Trained_model_path))
Net=Net.cuda() # Train on cuda
optimizer=torch.optim.Adam(params=Net.parameters(),lr=Learning_Rate,weight_decay=Weight_Decay) # Create adam optimizer

#----------------------------------------Create readers for data set--------------------------------------------------------------------------------------------------------------

ChemReader=LabPicsReader.Reader(ChemLabPicsDir,MaxBatchSize,MinSize,MaxSize,MaxPixels,MinMaskSize, TrainingMode,IgnoreParts,IgnoreSurfacePhase,IgnoreScattered,ClassToUse=ClassToUse)
MedReader=LabPicsReader.Reader(MedLabPicsDir,MaxBatchSize,MinSize,MaxSize,MaxPixels,MinMaskSize, TrainingMode,IgnoreParts,IgnoreSurfacePhase,IgnoreScattered,ClassToUse=ClassToUse)

#--------------------------- Create logs files for saving loss during training----------------------------------------------------------------------------------------------------------
if not os.path.exists(TrainedModelWeightDir): os.makedirs(TrainedModelWeightDir) # Create folder for trained weight
f = open(TrainLossTxtFile, "w+")# Training loss log file
f.write("Iteration\tloss\t Learning Rate=")
f.close()
#-------------------statitics  paramters for tracking loss--------------------------------------------------------------------------------
PrevAvgInsLoss=0 # Previously average instance loss
PrevAvgSemLoss=0 # previously average semantic loss
AVGInsLoss=0  #  average instance loss
AVGSemLoss=0  # average semantic loss
AVGCatLoss={} # Average loss per category  for semantic segmentation
for nm in ClassToUse:
    AVGCatLoss[nm]=0
#..............Start Training loop: Main Training....................................................................
print("Start Training")
for itr in range(InitStep,MAX_ITERATION): # Main training loop
    fr = 1 / np.min([itr - InitStep + 1, 2000])
    print(itr)
    Mode=""
#-----------------------------Read data----------------------------------------------------------------
    if np.random.rand() < 0.4: # Read instance data
            Mode = "Instance"
            if np.random.rand()<0.62:
                Imgs, GTMasks,WeightMaps, ROI, InsData = ChemReader.LoadBatchInstance()# Read from labpic chemistry
            else:
                Imgs, GTMasks, WeightMaps, ROI, InsData = MedReader.LoadBatchInstance()# Read from labpic medical
    else: # Read semantic data
            Mode = "Semantic"
            if np.random.rand()<0.62:
                Imgs, GTMasks,WeightMaps, ROI = ChemReader.LoadBatchSemantic() # Read from labpic chemistry
            else:
                Imgs, GTMasks, WeightMaps, ROI = MedReader.LoadBatchSemantic() # Read from labpic medical
    #----------------Display readed data----------------------
    # for i in range(Imgs.shape[0]):
    #     I = Imgs[i].copy()
    #     I1 = I.copy()
    #     I1[:, :, 0][ROI[i] > 0] = 0
    #     I1[:, :, 1][ROI[i] > 0] = 0
    #     vis.show(np.concatenate([I, I1], axis=1), "Vessel  ")
    #     for nm in GTMasks[i]:
    #         if nm == 'Vessel': continue
    #         I2 = I.copy()
    #         I2[:, :, 0][GTMasks[i][nm] > 0] = 0
    #         I2[:, :, 1][GTMasks[i][nm] > 0] = 0
    #       #  print(InsData[i][nm])
    #       #  vis.show(np.concatenate([I, I1, I2, vis.GreyScaleToRGB(WeightMaps[i][nm] * 255)], axis=1),str(InsData[i][nm]))\
    #         vis.show(np.concatenate([I, I1, I2, vis.GreyScaleToRGB(WeightMaps[i][nm] * 255)], axis=1),nm+str(np.unique(WeightMaps[i][nm])))
    #-----------------------------------------------------------------------------------------------------------------
    # for i in range(Imgs.shape[0]):
    #     I = Imgs[i].copy()
    #     I1 = I.copy()
    #     I1[:, :, 0][ROI[i] > 0] = 0
    #     I1[:, :, 1][ROI[i] > 0] = 0
    #     vis.show(np.concatenate([I, I1], axis=1), "Vessel  ")
    #     for nm in GTMasks:
    #         if GTMasks[nm][i].sum()==0: continue
    #         I2 = I.copy()
    #         I2[:, :, 0][GTMasks[nm][i] > 0] = 0
    #         I2[:, :, 1][GTMasks[nm][i] > 0] = 0
    #       #  print(InsData[i][nm])
    #       #  vis.show(np.concatenate([I, I1, I2, vis.GreyScaleToRGB(WeightMaps[i][nm] * 255)], axis=1),str(InsData[i][nm]))\
    #         vis.show(np.concatenate([I, I1, I2, vis.GreyScaleToRGB(WeightMaps[nm][i] * 255)], axis=1),nm+str(np.unique(WeightMaps[nm][i])))
    #     continue
    #----------------------------------------------
    #print("RUN PREDICITION")
#-------------------------------------Run Prediction--------------------------------------------------------------------------------
    ProbInst, LbInst, ProbSemantic, LbSemantic = Net.forward(Images=Imgs,ROI=ROI,PredictSemantic=(Mode == "Semantic"),PredictInstance = (Mode == "Instance")) # Run net inference and get prediction
    Net.zero_grad()
    Loss=0
    batchSize = Imgs.shape[0]
#******************Instance Segmentation loss find best matching GT /Predicted) segments (hungarian matching)**************************************************************************
    if (Mode == "Instance"):
        for iii in range(len(GTMasks)): # Note the matching  and loss is done now for every image of the batch indepenently. This is not very effiecnt p
            LbSum = {}
            GTMasksPT={}
            WeightMapsPT={}
            Cost = np.zeros([len(LbInst), len(GTMasks[iii])], dtype=np.float32) # Create correlation matrix of every IOU between every GT segment and predicted segment
            for ff,nm in enumerate(GTMasks[iii]):
                GTMasksPT[ff]=torch.autograd.Variable(torch.from_numpy((GTMasks[iii][nm]).astype(np.float32)).cuda(), requires_grad=False)
                WeightMapsPT[ff]=torch.autograd.Variable(torch.from_numpy((WeightMaps[iii][nm]).astype(np.float32)).cuda(), requires_grad=False)
                GTMasksSum=GTMasksPT[ff].sum()
                for ip in range(len(LbInst)):
                    if not ip in LbSum:
                        LbSum[ip]=LbInst[ip][iii].sum()
                    inter=(LbInst[ip][iii]*GTMasksPT[ff]).sum()
                    iou=inter/(GTMasksSum+LbSum[ip]-inter+0.0000001) # Calculate IOU between predicted and GT segment
                    Cost[ip,ff]=iou.data.cpu().numpy()
    #--------------------------Find match and calculate loss-------------------------------------------------------------
            row_ind, col_ind = linear_sum_assignment(-Cost) # Hungarian matching find the best matching prediction to GT segment, based on IOU matching
            for i in range(len(row_ind)): # Caclulate crossentropy  loss between  matching predicted and GT masks
                Loss -= torch.mean(GTMasksPT[col_ind[i]]* torch.log(ProbInst[row_ind[i]][iii][1]+0.000001)*WeightMapsPT[col_ind[i]])
                Loss -= torch.mean((1-GTMasksPT[col_ind[i]])* torch.log(ProbInst[row_ind[i]][iii][0]+0.000001))#
            for i in range(len(ProbInst)): # For unmatched set GT to empty segment
                if i not in (row_ind):
                    Loss -= torch.mean(torch.log(ProbInst[i][iii][0] + 0.000001))
        Loss /= batchSize
        AVGInsLoss = (1 - fr) * AVGInsLoss + fr * Loss.data.cpu().numpy() # Add to average loss for statitics
#*********************************Calculate semantic loss******************************************************************************************
    if Mode == "Semantic":
        Loss = 0
        LossByCat = {}
        lossCount = 0
        #  print("lll===="+str(len(list(OutLbDict))))
        #  print("-2-")
        for nm in ProbSemantic: # Go over all classes and match semantic maps
            lossCount += 1
            GT = torch.autograd.Variable(torch.from_numpy(GTMasks[nm].astype(np.float32)).cuda(), requires_grad=False)
            WeightMapsPT = torch.autograd.Variable(torch.from_numpy(WeightMaps[nm].astype(np.float32)).cuda(), requires_grad=False)
            LossByCat[nm] = -torch.mean((GT * torch.log(ProbSemantic[nm][:, 1, :, :] + 0.0000001) * WeightMapsPT + (1 - GT) * torch.log(ProbSemantic[nm][:, 0, :, :] + 0.0000001)))
            Loss += LossByCat[nm]
            AVGCatLoss[nm] = AVGCatLoss[nm] * (1 - fr) + fr * float(LossByCat[nm].data.cpu().numpy())  # Intiate runing average loss
            AVGSemLoss = (1 - fr) * AVGSemLoss + fr * Loss.data.cpu().numpy() # average semantic segmentation loss for statitics
#-------------------------------------------------------------------------------------------------------------------

    Loss.backward() # Backpropogate loss
    optimizer.step() # Apply gradient descent change to weight

####################################Saving and displaying###########################################################################

# --------------Save trained model------------------------------------------------------------------------------------------------------------------------------------------
    if itr % 1000 == 0:# and itr>0: #Save model weight once every 1000 steps temporary
        print("Saving Model to file in "+TrainedModelWeightDir+"/Defult.torch")
        torch.save(Net.state_dict(), TrainedModelWeightDir + "/Defult.torch")
        torch.save(Net.state_dict(), TrainedModelWeightDir + "/DefultBack.torch")
        print("model saved")
        np.save(TrainedModelWeightDir+"/Learning_Rate.npy",Learning_Rate)
        np.save(TrainedModelWeightDir+"/itr.npy",itr)
    if itr % 60000 == 0 and itr>0: #Save model weight once every 60k steps permenant
        print("Saving Model to file in "+TrainedModelWeightDir+"/"+ str(itr) + ".torch")
        torch.save(Net.state_dict(), TrainedModelWeightDir + "/" + str(itr) + ".torch")
        print("model saved")
#......................Write and display train loss..........................................................................
    if itr % 10==0: # Display train loss

        txt="\n"+str(itr)+"\t Semantic Loss "+str(AVGSemLoss)+"\t Ins Loss "+str(AVGInsLoss)+"\t Learning Rate "+str(Learning_Rate) +"\n"
        for nm in ClassToUse:
            txt+=nm+") "+str(AVGCatLoss[nm])+"  "
        print(txt)
        #Write train loss to file
        with open(TrainLossTxtFile, "a") as f:
            f.write(txt)
            f.close()
#----------------Update learning rate fractal manner-------------------------------------------------------------------------------
    if itr%10000==0:
        if PrevAvgInsLoss*0.95<AVGInsLoss or PrevAvgSemLoss*0.95<AVGSemLoss: # if average loss as not decrease in the last 10k step reduce learning rate
            Learning_Rate*=0.9 # Reduce learing rate
            if Learning_Rate<=4e-7: # If learning rate is to small increase it back up
                Learning_Rate=5e-6
            print("Learning Rate="+str(Learning_Rate))
            print("======================================================================================================================")
            optimizer = torch.optim.Adam(params=Net.parameters(), lr=Learning_Rate,weight_decay=Weight_Decay)  # Create adam optimizer with new loss
            torch.cuda.empty_cache()  # Empty cuda memory to avoid memory leaks
        PrevAvgInsLoss=AVGInsLoss+0.0000000001
        PrevAvgSemLoss= AVGSemLoss+0.0000001


# Copyright (c) 2019 Alibaba Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import numpy as np
import tensorflow as tf
import sys

from easy_rl.agents import AgentBase
from easy_rl import models
from easy_rl.models import PGModel
from easy_rl.utils.buffer import TrajectoryBuffer
from easy_rl.utils.utils import compute_targets

class A3CAgent(AgentBase):
    """A3C where workers interact with the environments and update the model parameters (on the servers) simutaneously without strict consistency considerations.

    See http://arxiv.org/abs/1602.01783 for more details.
    """

    _agent_name = "A3C"
    _valid_model_classes = [PGModel]
    _default_model_class_names = ["A3C"]

    def _init(self, model_config, ckpt_dir, custom_model):

        model_config = model_config or {}

        if custom_model is not None:
            assert np.any([
                issubclass(custom_model, e) for e in self._valid_model_classes
            ])
            model_class = custom_model
        else:
            model_name = model_config.get("type", self._default_model_class)
            assert model_name in self._valid_model_class_names, "{} does NOT support {} model.".format(
                self._agent_name, model_name)
            model_class = models.models[model_name]
        with self.distributed_handler.device:
            self.model = model_class(
                self.executor.ob_ph_spec,
                self.executor.action_ph_spec,
                model_config=model_config)
            self._behavior_model = self.model

            self._global_num_sampled = tf.get_variable(
                name="global_num_sampled", dtype=tf.int64, shape=())

        self._update_global_num_sampled = tf.assign_add(
            self._global_num_sampled,
            np.int64(self.config["sample_batch_size"]),
            use_locking=True)

        global_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)

        self.executor.setup(
            self.distributed_handler.master,
            self.distributed_handler.is_chief,
            self.model.global_step,
            ckpt_dir,
            self.model.summary_ops,
            global_vars=global_vars,
            local_vars=None)

        # transitions are stored in chronological order
        self._buffer = TrajectoryBuffer()

    def send_experience(self, obs, actions, rewards, dones, value_preds):
        """Store experience in the buffer

        Postprocess the collected transitions and store the them in the buffer.
        `obs`, `actions`, `rewards`, `next_obs`, `dones` are basic fields.

        Arguments:
            obs (list): the observed states.
            actions (list): the actions taken in response to the states.
            rewards (list): collected from the environment.
            next_obs (list): the obsered next states.
            dones (list): bool flags indicating whether the episode finishes.
            weights is the initializd priorities of samples.
            kwargs (dict) contains the additional algorithm-specific field(s).
        """

        # there is no need to compute_targets for v-trace model
        advantages, value_targets = compute_targets(
            rewards,
            dones,
            value_preds,
            gamma=self.config.get('gamma', 0.95),
            lambda_=self.config.get('lambda_', 1.0),
            use_gae=self.config.get('use_gae', True))
        traj_len = len(advantages)

        kwargs_actually_used = dict()
        kwargs_actually_used["advantages"] = advantages
        kwargs_actually_used["value_targets"] = value_targets
        self._buffer.add(
            obs=np.asarray(obs)[:traj_len],
            actions=np.asarray(actions)[:traj_len],
            rewards=np.asarray(rewards)[:traj_len],
            dones=np.asarray(dones)[:traj_len],
            **dict((k, v) for k, v in kwargs_actually_used.items()))

        if len(self._buffer) >= self.config["batch_size"]:
            self._ready_to_receive = True

        self._ready_to_send = False

        self.executor.run(self._update_global_num_sampled, {})

        # clear the lists
        del obs[:len(obs) - 1]
        del actions[:len(actions) - 1]
        del rewards[:len(rewards) - 1]
        del dones[:len(dones) - 1]
        del value_preds[:len(value_preds) - 1]

    def receive_experience(self):
        """Acquire training examples from the buffer.

        Returns:
            batch_data (dict): contains fields of data for making an update.
        """

        self._ready_to_receive = False
        return self._buffer.sample(self.config["batch_size"])

    def _init_act_count(self):

        # In order to get the same size of value_preds of next_obs, one more timestep is needed in the first iteration.
        self._act_count = -1

    def join(self):
        """Call `server.join()` if the agent object serves as a parameter server.
        """

        self.distributed_handler.server.join()

    def should_stop(self):
        """Judge whether the agent should stop working.
        """
        if self.distributed_handler.job_name == "worker":
            global_timesteps = self.executor.run(self._global_num_sampled, {})
            if global_timesteps > self.config.get("scheduled_timesteps",
                                                  1000000):
                return True
        return False


"""
	验证码识别测试模块
	create by Ian in 2018-8-7 13:20:46
"""
import os
import tensorflow as tf
from PIL import Image
from nets import nets_factory
import numpy as np
import matplotlib.pyplot as plt

# 不同字符数量
CHAR_SET_LEN = 36
# 图片高度
IMAGE_HEIGHT = 60
# 图片宽度
IMAGE_WIDTH = 160
# 批次
BATCH_SIZE = 1
# tfrecord文件存放路径
TFRECORD_FILE = "./tmp/test.tfrecords"

# placeholder
x = tf.placeholder(tf.float32, [None, 224, 224])

# 从tfrecord读出数据
def read_and_decode(filename):
    # 根据文件名生成一个队列
    filename_queue = tf.train.string_input_producer([filename])
    reader = tf.TFRecordReader()
    # 返回文件名和文件
    _, serialized_example = reader.read(filename_queue)
    features = tf.parse_single_example(serialized_example,
                                       features={
                                           'image' : tf.FixedLenFeature([], tf.string),
                                           'label0': tf.FixedLenFeature([], tf.int64),
                                           'label1': tf.FixedLenFeature([], tf.int64),
                                           'label2': tf.FixedLenFeature([], tf.int64),
                                           'label3': tf.FixedLenFeature([], tf.int64),
                                       })
    # 获取图片数据
    image = tf.decode_raw(features['image'], tf.uint8)
    # 没有经过预处理的灰度图
    image_raw = tf.reshape(image, [224, 224])
    # tf.train.shuffle_batch必须确定shape
    image = tf.reshape(image, [224, 224])
    # 图片预处理
    image = tf.cast(image, tf.float32) / 255.0
    image = tf.subtract(image, 0.5)
    image = tf.multiply(image, 2.0)
    # 获取label
    label0 = tf.cast(features['label0'], tf.int32)
    label1 = tf.cast(features['label1'], tf.int32)
    label2 = tf.cast(features['label2'], tf.int32)
    label3 = tf.cast(features['label3'], tf.int32)

    return image, image_raw, label0, label1, label2, label3

# 获取图片数据和标签
image, image_raw, label0, label1, label2, label3 = read_and_decode(TFRECORD_FILE)

#使用shuffle_batch可以随机打乱
image_batch, image_raw_batch, label_batch0, label_batch1, label_batch2, label_batch3 = tf.train.shuffle_batch(
        [image, image_raw, label0, label1, label2, label3], batch_size = BATCH_SIZE,
        capacity = 50000, min_after_dequeue=10000, num_threads=1)

#定义网络结构
train_network_fn = nets_factory.get_network_fn(
    'alexnet_v2',
    num_classes=CHAR_SET_LEN,
    weight_decay=0.0005,
    is_training=False)

with tf.Session() as sess:
    # inputs: a tensor of size [batch_size, height, width, channels]
    X = tf.reshape(x, [BATCH_SIZE, 224, 224, 1])
    # 数据输入网络得到输出值
    logits0,logits1,logits2,logits3,end_points = train_network_fn(X)

    # 预测值
    predict0 = tf.reshape(logits0, [-1, CHAR_SET_LEN])
    predict0 = tf.argmax(predict0, 1)

    predict1 = tf.reshape(logits1, [-1, CHAR_SET_LEN])
    predict1 = tf.argmax(predict1, 1)

    predict2 = tf.reshape(logits2, [-1, CHAR_SET_LEN])
    predict2 = tf.argmax(predict2, 1)

    predict3 = tf.reshape(logits3, [-1, CHAR_SET_LEN])
    predict3 = tf.argmax(predict3, 1)

    # 初始化
    sess.run(tf.global_variables_initializer())
    # 载入训练好的模型
    saver = tf.train.Saver()
    saver.restore(sess,'./model/crack_captcha.model-6000')

    # 创建一个协调器，管理线程
    coord = tf.train.Coordinator()
    # 启动QueueRunner, 此时文件名队列已经进队
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    for i in range(3):
        # 获取一个批次的数据和标签
        b_image, b_image_raw, b_label0, b_label1 ,b_label2 ,b_label3 = sess.run([image_batch,
                                                                    image_raw_batch,
                                                                    label_batch0,
                                                                    label_batch1,
                                                                    label_batch2,
                                                                    label_batch3])
        # 显示图片
        img=Image.fromarray(b_image_raw[0],'L')
        plt.imshow(img)
        plt.axis('off')
        plt.show()
        # 打印标签
        print('label:',b_label0, b_label1 ,b_label2 ,b_label3)
        # 预测
        label0,label1,label2,label3 = sess.run([predict0,predict1,predict2,predict3], feed_dict={x: b_image})
        # 打印预测值
        print('predict:',label0,label1,label2,label3)

    # 通知其他线程关闭
    coord.request_stop()