"""
Transfer Learning
"""

import subprocess
import time

from kili.client import Kili

SECONDS_TO_WAIT = 1
LABEL_FIELDS = ["isLatestLabelForUser", "labelType", "jsonResponse", "createdAt"]
FIELDS = ["id", "content", "externalId"] + [f"labels.{x}" for x in LABEL_FIELDS]


def get_labels_of_types(asset, label_types):
    """
    Extracts the latest labels from an asset

    Args:
        asset: the asset to extract the labels from
        label_types: type of label, either DEFAULT or REVIEW
    """
    labels = [
        label
        for label in asset["labels"]
        if label["labelType"] in label_types and label["isLatestLabelForUser"]
    ]
    return sorted(labels, key=lambda label: label["createdAt"])


class TransferLearning:
    """
    TransferLearning
    """

    # pylint: disable=too-many-instance-attributes,too-many-arguments

    def __init__(
        self,
        api_key,
        api_endpoint,
        project_id,
        number_of_inferences,
        minimum_number_of_assets_to_launch_training=10,
    ):
        self.kili = Kili(api_key=api_key, api_endpoint=api_endpoint)
        self.api_key = api_key
        self.project_id = project_id
        self.current_inference_number = 0
        self.current_training_number = 0
        self.last_training_number = -1
        self.assets_seen_in_training = []
        self.minimum_number_of_assets_to_launch_training = (
            minimum_number_of_assets_to_launch_training
        )
        self.number_of_inferences = number_of_inferences

    def get_assets_to_train(self):
        """
        Collects the assets to train on
        """
        assets = self.kili.assets(project_id=self.project_id, fields=FIELDS)
        assets_to_train = []
        for asset in assets:
            default_labels = get_labels_of_types(asset, ["DEFAULT"])
            review_labels = get_labels_of_types(asset, ["REVIEWED"])
            if len(review_labels) > 0:
                asset["labels"] = [review_labels[-1]]
                assets_to_train.append(asset)
            elif len(default_labels) == 1:
                asset["labels"] = [default_labels[-1]]
                assets_to_train.append(asset)
            elif len(review_labels) == 0 and len(default_labels) > 0:
                print(f'Asset {asset["id"]} has several labels: it should be reviewed')
            else:
                continue
        return assets_to_train

    @staticmethod
    def train(assets_to_train):
        """
        Prints that training is ongoing
        """
        print(
            f"Launch training for {len(assets_to_train)} assets:"
            f' {[asset["id"] for asset in assets_to_train]}'
        )

    def launch_train(self):
        """
        Launches the training
        """
        print("Launching train")
        time.sleep(SECONDS_TO_WAIT)
        assets_to_train = self.get_assets_to_train()
        if len(self.assets_seen_in_training) == 0:
            filtered_assets_to_train = assets_to_train
        else:
            filtered_assets_to_train = [
                asset
                for asset in assets_to_train
                if all(asset["id"] not in training for training in self.assets_seen_in_training)
            ]
        if len(filtered_assets_to_train) >= self.minimum_number_of_assets_to_launch_training:
            print("Starting training")
            TransferLearning.train(filtered_assets_to_train)
            self.current_training_number += 1
            self.assets_seen_in_training.append([asset["id"] for asset in filtered_assets_to_train])
        else:
            print("Not enough labeled assets to start training")

    def get_assets_to_predict(self):
        """
        Collects the assets to predict
        """
        assets = self.kili.assets(project_id=self.project_id, fields=FIELDS)
        assets_to_predict = []
        for asset in assets:
            labels = get_labels_of_types(asset, ["DEFAULT", "REVIEWED"])
            if len(labels) == 0:
                assets_to_predict.append(asset)
        return assets_to_predict

    @staticmethod
    def predict(assets_to_predict):
        """
        Prints that prediction is ongoing
        """
        print(
            f"Launch inference for {len(assets_to_predict)} assets:"
            f' {[asset["id"] for asset in assets_to_predict]}'
        )

    def launch_predict(self):
        """
        Launches the prediction
        """
        print("Launching prediction")
        time.sleep(SECONDS_TO_WAIT)
        if self.current_training_number == self.last_training_number:
            print("Inference will not be launched for now...")
            return
        assets_to_predict = self.get_assets_to_predict()
        if len(assets_to_predict) > 0:
            current_training_number = self.current_training_number
            TransferLearning.predict(assets_to_predict)
            self.last_training_number = current_training_number
            self.current_inference_number += 1

    @staticmethod
    def launch_tensorboard():
        """
        Launches the tensorboard
        """
        print("Starting Tensorboard...")
        with subprocess.Popen(["tensorboard", "--logdir=runs"]) as proc:
            print(proc.stdout.read())

    def launch(self):
        """
        Launches the tensorboard, training and prediction
        """
        TransferLearning.launch_tensorboard()
        while self.current_inference_number < self.number_of_inferences:
            print(f"Iteration inference {self.current_inference_number}")
            self.launch_train()
            self.launch_predict()
